From b1597cc45df301246f8dede93a51429f35813e03 Mon Sep 17 00:00:00 2001 From: Ruslan Iushchenko Date: Thu, 15 Oct 2026 09:04:02 +0200 Subject: [PATCH 1/5] When getting Kafka offsets from data written to the output merge updated offsets with previous offsets to ensure all partitions are present. --- .../absa/pramen/api/offset/DataOffset.scala | 1 + .../core/bookkeeper/OffsetManagerUtils.scala | 39 ++++++++++- .../model/OffsetRecordConverter.scala | 1 + .../pramen/core/metastore/MetastoreImpl.scala | 2 +- .../pipeline/IncrementalIngestionJob.scala | 19 +++++- .../bookkeeper/OffsetManagerUtilsSuite.scala | 67 +++++++++++++++++-- 6 files changed, 118 insertions(+), 11 deletions(-) diff --git a/pramen/api/src/main/scala/za/co/absa/pramen/api/offset/DataOffset.scala b/pramen/api/src/main/scala/za/co/absa/pramen/api/offset/DataOffset.scala index c5f1f410a..7e6b324d1 100644 --- a/pramen/api/src/main/scala/za/co/absa/pramen/api/offset/DataOffset.scala +++ b/pramen/api/src/main/scala/za/co/absa/pramen/api/offset/DataOffset.scala @@ -49,6 +49,7 @@ object DataOffset { infoDate: LocalDate, batchId: Long, offsetType: OffsetType, + minOffsetOpt: Option[OffsetValue], createdAt: Instant ) extends DataOffset { diff --git a/pramen/core/src/main/scala/za/co/absa/pramen/core/bookkeeper/OffsetManagerUtils.scala b/pramen/core/src/main/scala/za/co/absa/pramen/core/bookkeeper/OffsetManagerUtils.scala index 20d097a7e..b0462ec06 100644 --- a/pramen/core/src/main/scala/za/co/absa/pramen/core/bookkeeper/OffsetManagerUtils.scala +++ b/pramen/core/src/main/scala/za/co/absa/pramen/core/bookkeeper/OffsetManagerUtils.scala @@ -23,7 +23,35 @@ import za.co.absa.pramen.api.offset.{KafkaPartition, OffsetType, OffsetValue} import za.co.absa.pramen.api.sql.SqlGeneratorBase object OffsetManagerUtils { - def getMinMaxValueFromData(df: DataFrame, offsetColumn: String, offsetType: OffsetType): Option[(OffsetValue, OffsetValue)] = { + + /** + * Merges two Kafka partition sets in such a way that 'primaryValue' is considered a maximum offset + * for a given increment ingestion job. But if data ingested does not cover all partitions, such partitions + * might be missing in the list. In this case, values from the 'baseValue' are used. + */ + def mergeKafkaPartitions(primaryValue: OffsetValue.KafkaValue, baseValue: OffsetValue.KafkaValue): OffsetValue.KafkaValue = { + val partitions1 = primaryValue.value.map(v => (v.partition, v.offset)).toMap + val partitions2 = baseValue.value.map(v => (v.partition, v.offset)).toMap + val maxPartition = Math.max(partitions1.keys.max, partitions2.keys.max) + + val newPartitions = for (i <- 0 to maxPartition) yield { + if (partitions1.contains(i)) + KafkaPartition(i, partitions1(i)) + else if (partitions2.contains(i)) + KafkaPartition(i, partitions2(i)) + else + KafkaPartition(i, 0) + } + + OffsetValue.KafkaValue(newPartitions) + } + + + /** + * + */ + def getMinMaxValueFromData(df: DataFrame, offsetColumn: String, offsetType: OffsetType, currentMaxOffset: Option[OffsetValue]): Option[(OffsetValue, OffsetValue)] = { + if (df.isEmpty) { None } else { @@ -37,10 +65,15 @@ object OffsetManagerUtils { max(col(kafkaOffsetFieldName)).as("max_offset") ).orderBy(kafkaPartitionFieldName) - val minValue = OffsetValue.KafkaValue(aggregatedDf.collect().map(row => KafkaPartition(row.getAs[Int](0), row.getAs[Long](1))).toSeq) + val minValue: OffsetValue.KafkaValue = OffsetValue.KafkaValue(aggregatedDf.collect().map(row => KafkaPartition(row.getAs[Int](0), row.getAs[Long](1))).toSeq) val maxValue = OffsetValue.KafkaValue(aggregatedDf.collect().map(row => KafkaPartition(row.getAs[Int](0), row.getAs[Long](2))).toSeq) - Some(minValue, maxValue) + val (minValueMerged, maxValueMerged) = currentMaxOffset match { + case Some(value: OffsetValue.KafkaValue) => (mergeKafkaPartitions(minValue, value), mergeKafkaPartitions(maxValue, value)) + case _ => (minValue, maxValue) + } + + Some(minValueMerged, maxValueMerged) } else { val row = df.agg(min(offsetType.getSparkCol(col(offsetColumn)).cast(StringType)), max(offsetType.getSparkCol(col(offsetColumn))).cast(StringType)) diff --git a/pramen/core/src/main/scala/za/co/absa/pramen/core/bookkeeper/model/OffsetRecordConverter.scala b/pramen/core/src/main/scala/za/co/absa/pramen/core/bookkeeper/model/OffsetRecordConverter.scala index d6ae4db23..9bda69b04 100644 --- a/pramen/core/src/main/scala/za/co/absa/pramen/core/bookkeeper/model/OffsetRecordConverter.scala +++ b/pramen/core/src/main/scala/za/co/absa/pramen/core/bookkeeper/model/OffsetRecordConverter.scala @@ -43,6 +43,7 @@ object OffsetRecordConverter { LocalDate.parse(r.infoDate), r.batchId, dataType, + minOffsetOpt, Instant.ofEpochMilli(r.createdAtMilli) ) } diff --git a/pramen/core/src/main/scala/za/co/absa/pramen/core/metastore/MetastoreImpl.scala b/pramen/core/src/main/scala/za/co/absa/pramen/core/metastore/MetastoreImpl.scala index 3c72cf477..179cf7a2d 100644 --- a/pramen/core/src/main/scala/za/co/absa/pramen/core/metastore/MetastoreImpl.scala +++ b/pramen/core/src/main/scala/za/co/absa/pramen/core/metastore/MetastoreImpl.scala @@ -403,6 +403,6 @@ object MetastoreImpl { private[core] def getMinMaxOffsetFromMetastoreDf(df: DataFrame, batchIdColumn: String, currentMax: Option[OffsetValue]): Option[(OffsetValue, OffsetValue)] = { val offsetType = if (df.schema.fields.find(_.name == batchIdColumn).get.dataType == StringType) OffsetType.StringType else OffsetType.IntegralType - OffsetManagerUtils.getMinMaxValueFromData(df, batchIdColumn, offsetType) + OffsetManagerUtils.getMinMaxValueFromData(df, batchIdColumn, offsetType, currentMax) } } diff --git a/pramen/core/src/main/scala/za/co/absa/pramen/core/pipeline/IncrementalIngestionJob.scala b/pramen/core/src/main/scala/za/co/absa/pramen/core/pipeline/IncrementalIngestionJob.scala index 6b0d283b8..74eb707bf 100644 --- a/pramen/core/src/main/scala/za/co/absa/pramen/core/pipeline/IncrementalIngestionJob.scala +++ b/pramen/core/src/main/scala/za/co/absa/pramen/core/pipeline/IncrementalIngestionJob.scala @@ -121,7 +121,6 @@ class IncrementalIngestionJob(operationDef: OperationDef, throw new IllegalStateException(s"No offsets for '${outputTable.name}' for '$infoDate'. Cannot rerun.") } case (true, false) => - val om = bookkeeper.getOffsetManager val infoDateLatestOffset = om.getMaxInfoDateAndOffset(outputTable.name, Some(infoDate)) infoDateLatestOffset match { case Some(maxOffset) => @@ -158,6 +157,18 @@ class IncrementalIngestionJob(operationDef: OperationDef, validateOffsetColumn(df, offsetInfo) val req = om.startWriteOffsets(outputTable.name, infoDate, offsetInfo.offsetType) + val currentMaxOffsetOpt = if (offsetInfo.offsetType == OffsetType.KafkaType) { + val isReRun = runReason == TaskRunReason.Rerun + + if (isReRun) { + om.getMaxInfoDateAndOffset(outputTable.name, Option(infoDate)).map(_.maximumOffset) + } else { + om.getMaxInfoDateAndOffset(outputTable.name, None).map(_.maximumOffset) + } + } else { + None + } + val stats = try { val saveMode = if (isRerun) SaveMode.Overwrite else SaveMode.Append val statsToReturn = metastore.saveTable(outputTable.name, infoDate, dfToSave, inputRecordCount, saveModeOverride = Some(saveMode)) @@ -167,7 +178,7 @@ class IncrementalIngestionJob(operationDef: OperationDef, else metastore.getBatch(outputTable.name, infoDate, None) - OffsetManagerUtils.getMinMaxValueFromData(updatedDf, offsetInfo.offsetColumn, offsetInfo.offsetType) match { + OffsetManagerUtils.getMinMaxValueFromData(updatedDf, offsetInfo.offsetColumn, offsetInfo.offsetType, currentMaxOffsetOpt) match { case Some((minOffset, maxOffset)) => if (isRerun) { om.commitRerun(req, minOffset, maxOffset) @@ -278,7 +289,9 @@ class IncrementalIngestionJob(operationDef: OperationDef, throw new IllegalArgumentException(s"Offset column '${offsetInfo.offsetColumn}' not found in the output table '${outputTable.name}'. Cannot update uncommitted offsets.") } - val (newMinOffset, newMaxOffset) = OffsetManagerUtils.getMinMaxValueFromData(df, offsetInfo.offsetColumn, offsetInfo.offsetType).get + val minOffsets = uncommittedOffsets.flatMap(_.minOffsetOpt) + val maxOfMinOffsetsOpt = if (minOffsets.nonEmpty) Some(minOffsets.max) else None + val (newMinOffset, newMaxOffset) = OffsetManagerUtils.getMinMaxValueFromData(df, offsetInfo.offsetColumn, offsetInfo.offsetType, maxOfMinOffsetsOpt).get log.warn(s"Fixing uncommitted offsets. New offset to commit for ${outputTable.name} at $infoDate: " + s"min offset: ${newMinOffset.valueString}, max offset: ${newMaxOffset.valueString}.") diff --git a/pramen/core/src/test/scala/za/co/absa/pramen/core/tests/bookkeeper/OffsetManagerUtilsSuite.scala b/pramen/core/src/test/scala/za/co/absa/pramen/core/tests/bookkeeper/OffsetManagerUtilsSuite.scala index dc25c1807..245bcefda 100644 --- a/pramen/core/src/test/scala/za/co/absa/pramen/core/tests/bookkeeper/OffsetManagerUtilsSuite.scala +++ b/pramen/core/src/test/scala/za/co/absa/pramen/core/tests/bookkeeper/OffsetManagerUtilsSuite.scala @@ -29,11 +29,70 @@ class OffsetManagerUtilsSuite extends AnyWordSpec with SparkTestBase { import spark.implicits._ + "mergeKafkaPartitions" should { + val base = OffsetValue.KafkaValue(Seq( + KafkaPartition(0, 11), + KafkaPartition(1, 21), + KafkaPartition(2, 31), + KafkaPartition(3, 41), + KafkaPartition(4, 51), + )) + + "return correct offsets when primary offset is complete" in { + val primary = OffsetValue.KafkaValue(Seq( + KafkaPartition(0, 10), + KafkaPartition(1, 20), + KafkaPartition(2, 30), + KafkaPartition(3, 40), + KafkaPartition(4, 50), + )) + + val mered = OffsetManagerUtils.mergeKafkaPartitions(primary, base) + + assert(mered.value == primary.value) + } + + "return correct offsets when primary offset has gaps" in { + val primary = OffsetValue.KafkaValue(Seq( + KafkaPartition(0, 10), + KafkaPartition(2, 30), + KafkaPartition(4, 50), + )) + + val mered = OffsetManagerUtils.mergeKafkaPartitions(primary, base) + + assert(mered.value == Seq( + KafkaPartition(0, 10), + KafkaPartition(1, 21), + KafkaPartition(2, 30), + KafkaPartition(3, 41), + KafkaPartition(4, 50), + )) + } + + "return correct offsets when primary offset has missing first and last partitions" in { + val primary = OffsetValue.KafkaValue(Seq( + KafkaPartition(1, 20), + KafkaPartition(3, 40), + )) + + val mered = OffsetManagerUtils.mergeKafkaPartitions(primary, base) + + assert(mered.value == Seq( + KafkaPartition(0, 11), + KafkaPartition(1, 20), + KafkaPartition(2, 31), + KafkaPartition(3, 40), + KafkaPartition(4, 51), + )) + } + } + "getMinMaxValueFromData" should { "work for an integral data type" in { val df = List(("A", 1), ("B", 2), ("C", 3)).toDF("a", "offset") - val (minValue, maxValue) = OffsetManagerUtils.getMinMaxValueFromData(df, "offset", OffsetType.IntegralType).get + val (minValue, maxValue) = OffsetManagerUtils.getMinMaxValueFromData(df, "offset", OffsetType.IntegralType, None).get assert(minValue == OffsetValue.IntegralValue(1)) assert(maxValue == OffsetValue.IntegralValue(3)) @@ -42,7 +101,7 @@ class OffsetManagerUtilsSuite extends AnyWordSpec with SparkTestBase { "work for an string data type" in { val df = List(("A", 3), ("B", 1), ("C", 2)).toDF("offset", "b") - val (minValue, maxValue) = OffsetManagerUtils.getMinMaxValueFromData(df, "offset", OffsetType.StringType).get + val (minValue, maxValue) = OffsetManagerUtils.getMinMaxValueFromData(df, "offset", OffsetType.StringType, None).get assert(minValue == OffsetValue.StringValue("A")) assert(maxValue == OffsetValue.StringValue("C")) @@ -63,7 +122,7 @@ class OffsetManagerUtilsSuite extends AnyWordSpec with SparkTestBase { col("kafka_partition").as("partition") )) - val (minValue, maxValue) = OffsetManagerUtils.getMinMaxValueFromData(df, "kafka", OffsetType.KafkaType).get + val (minValue, maxValue) = OffsetManagerUtils.getMinMaxValueFromData(df, "kafka", OffsetType.KafkaType, None).get assert(minValue == OffsetValue.KafkaValue(Seq(KafkaPartition(0, 1), KafkaPartition(1, 3), KafkaPartition(2, 1)))) assert(maxValue == OffsetValue.KafkaValue(Seq(KafkaPartition(0, 2), KafkaPartition(1, 4), KafkaPartition(2, 3)))) @@ -74,7 +133,7 @@ class OffsetManagerUtilsSuite extends AnyWordSpec with SparkTestBase { val df = List(("A", baseTime), ("B", baseTime + 1000), ("C", baseTime + 1500)).toDF("a", "offset") .withColumn("offset", (col("offset") / 1000).cast(TimestampType)) - val (minValue, maxValue) = OffsetManagerUtils.getMinMaxValueFromData(df, "offset", OffsetType.DateTimeType).get + val (minValue, maxValue) = OffsetManagerUtils.getMinMaxValueFromData(df, "offset", OffsetType.DateTimeType, None).get assert(minValue == OffsetValue.DateTimeValue(Instant.ofEpochMilli(baseTime))) assert(maxValue == OffsetValue.DateTimeValue(Instant.ofEpochMilli(baseTime + 1500))) From 07b2ec55ef8e41172cf5d4d9e5f70efb47e65d5f Mon Sep 17 00:00:00 2001 From: Ruslan Iushchenko Date: Thu, 15 Oct 2026 11:02:22 +0200 Subject: [PATCH 2/5] Fix compilation for Scala 2.11 --- .../tests/bookkeeper/OffsetManagerUtilsSuite.scala | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/pramen/core/src/test/scala/za/co/absa/pramen/core/tests/bookkeeper/OffsetManagerUtilsSuite.scala b/pramen/core/src/test/scala/za/co/absa/pramen/core/tests/bookkeeper/OffsetManagerUtilsSuite.scala index 245bcefda..8bebc59c8 100644 --- a/pramen/core/src/test/scala/za/co/absa/pramen/core/tests/bookkeeper/OffsetManagerUtilsSuite.scala +++ b/pramen/core/src/test/scala/za/co/absa/pramen/core/tests/bookkeeper/OffsetManagerUtilsSuite.scala @@ -35,7 +35,7 @@ class OffsetManagerUtilsSuite extends AnyWordSpec with SparkTestBase { KafkaPartition(1, 21), KafkaPartition(2, 31), KafkaPartition(3, 41), - KafkaPartition(4, 51), + KafkaPartition(4, 51) )) "return correct offsets when primary offset is complete" in { @@ -44,7 +44,7 @@ class OffsetManagerUtilsSuite extends AnyWordSpec with SparkTestBase { KafkaPartition(1, 20), KafkaPartition(2, 30), KafkaPartition(3, 40), - KafkaPartition(4, 50), + KafkaPartition(4, 50) )) val mered = OffsetManagerUtils.mergeKafkaPartitions(primary, base) @@ -56,7 +56,7 @@ class OffsetManagerUtilsSuite extends AnyWordSpec with SparkTestBase { val primary = OffsetValue.KafkaValue(Seq( KafkaPartition(0, 10), KafkaPartition(2, 30), - KafkaPartition(4, 50), + KafkaPartition(4, 50) )) val mered = OffsetManagerUtils.mergeKafkaPartitions(primary, base) @@ -66,14 +66,14 @@ class OffsetManagerUtilsSuite extends AnyWordSpec with SparkTestBase { KafkaPartition(1, 21), KafkaPartition(2, 30), KafkaPartition(3, 41), - KafkaPartition(4, 50), + KafkaPartition(4, 50) )) } "return correct offsets when primary offset has missing first and last partitions" in { val primary = OffsetValue.KafkaValue(Seq( KafkaPartition(1, 20), - KafkaPartition(3, 40), + KafkaPartition(3, 40) )) val mered = OffsetManagerUtils.mergeKafkaPartitions(primary, base) @@ -83,7 +83,7 @@ class OffsetManagerUtilsSuite extends AnyWordSpec with SparkTestBase { KafkaPartition(1, 20), KafkaPartition(2, 31), KafkaPartition(3, 40), - KafkaPartition(4, 51), + KafkaPartition(4, 51) )) } } From 93965c31979e3e9efa26d88198d91f43f345c007 Mon Sep 17 00:00:00 2001 From: Ruslan Iushchenko Date: Thu, 15 Oct 2026 12:24:37 +0200 Subject: [PATCH 3/5] Fix offset merging to take into account that we store kafka offset as an inclusive interval. Fix uncommitted offsets processing for Kafka source --- .../core/bookkeeper/OffsetManager.scala | 2 +- .../core/bookkeeper/OffsetManagerCached.scala | 4 +- .../bookkeeper/OffsetManagerDynamoDb.scala | 5 +- .../core/bookkeeper/OffsetManagerJdbc.scala | 5 +- .../core/bookkeeper/OffsetManagerUtils.scala | 14 +++-- .../pipeline/IncrementalIngestionJob.scala | 5 +- .../IncrementalPipelineLongFixture.scala | 12 ++-- .../bookkeeper/OffsetManagerJdbcSuite.scala | 52 ++++++++--------- .../bookkeeper/OffsetManagerUtilsSuite.scala | 21 +++++-- .../extras/source/KafkaAvroSource.scala | 56 +++++++++++-------- 10 files changed, 102 insertions(+), 74 deletions(-) diff --git a/pramen/core/src/main/scala/za/co/absa/pramen/core/bookkeeper/OffsetManager.scala b/pramen/core/src/main/scala/za/co/absa/pramen/core/bookkeeper/OffsetManager.scala index 39addcd27..bd4c6b45e 100644 --- a/pramen/core/src/main/scala/za/co/absa/pramen/core/bookkeeper/OffsetManager.scala +++ b/pramen/core/src/main/scala/za/co/absa/pramen/core/bookkeeper/OffsetManager.scala @@ -64,7 +64,7 @@ trait OffsetManager extends AutoCloseable { * Starts an uncommitted offset for an incremental ingestion for a day. * This can only be done for the latest information date. */ - def startWriteOffsets(table: String, infoDate: LocalDate, offsetType: OffsetType): DataOffsetRequest + def startWriteOffsets(table: String, infoDate: LocalDate, offsetType: OffsetType, minOffset: Option[OffsetValue]): DataOffsetRequest /** * Commits changes to the table. If maxOffset is diff --git a/pramen/core/src/main/scala/za/co/absa/pramen/core/bookkeeper/OffsetManagerCached.scala b/pramen/core/src/main/scala/za/co/absa/pramen/core/bookkeeper/OffsetManagerCached.scala index 52f411ab0..201402e80 100644 --- a/pramen/core/src/main/scala/za/co/absa/pramen/core/bookkeeper/OffsetManagerCached.scala +++ b/pramen/core/src/main/scala/za/co/absa/pramen/core/bookkeeper/OffsetManagerCached.scala @@ -57,8 +57,8 @@ class OffsetManagerCached(offsetManager: OffsetManager) extends OffsetManager { } } - override def startWriteOffsets(table: String, infoDate: LocalDate, offsetType: OffsetType): DataOffsetRequest = { - offsetManager.startWriteOffsets(table, infoDate, offsetType) + override def startWriteOffsets(table: String, infoDate: LocalDate, offsetType: OffsetType, minOffset: Option[OffsetValue]): DataOffsetRequest = { + offsetManager.startWriteOffsets(table, infoDate, offsetType, minOffset) } override def commitOffsets(request: DataOffsetRequest, minOffset: OffsetValue, maxOffset: OffsetValue): Unit = { diff --git a/pramen/core/src/main/scala/za/co/absa/pramen/core/bookkeeper/OffsetManagerDynamoDb.scala b/pramen/core/src/main/scala/za/co/absa/pramen/core/bookkeeper/OffsetManagerDynamoDb.scala index 36627538d..0c4b083dc 100644 --- a/pramen/core/src/main/scala/za/co/absa/pramen/core/bookkeeper/OffsetManagerDynamoDb.scala +++ b/pramen/core/src/main/scala/za/co/absa/pramen/core/bookkeeper/OffsetManagerDynamoDb.scala @@ -132,10 +132,11 @@ class OffsetManagerDynamoDb( } } - override def startWriteOffsets(table: String, infoDate: LocalDate, offsetType: OffsetType): DataOffsetRequest = { + override def startWriteOffsets(table: String, infoDate: LocalDate, offsetType: OffsetType, minOffset: Option[OffsetValue]): DataOffsetRequest = { val createdAt = Instant.now() val createdAtMilli = createdAt.toEpochMilli val compositeKey = s"${infoDate.toString}#${createdAtMilli}" + val minOffsetStr = minOffset.map(_.valueString).getOrElse("") try { val putRequest = PutItemRequest.builder() @@ -145,7 +146,7 @@ class OffsetManagerDynamoDb( ATTR_COMPOSITE_KEY -> AttributeValue.builder().s(compositeKey).build(), ATTR_INFO_DATE -> AttributeValue.builder().s(infoDate.toString).build(), ATTR_DATA_TYPE -> AttributeValue.builder().s(offsetType.dataTypeString).build(), - ATTR_MIN_OFFSET -> AttributeValue.builder().s("").build(), + ATTR_MIN_OFFSET -> AttributeValue.builder().s(minOffsetStr).build(), ATTR_MAX_OFFSET -> AttributeValue.builder().s("").build(), ATTR_BATCH_ID -> AttributeValue.builder().n(batchId.toString).build(), ATTR_CREATED_AT -> AttributeValue.builder().n(createdAtMilli.toString).build() diff --git a/pramen/core/src/main/scala/za/co/absa/pramen/core/bookkeeper/OffsetManagerJdbc.scala b/pramen/core/src/main/scala/za/co/absa/pramen/core/bookkeeper/OffsetManagerJdbc.scala index b1e164860..c386b9612 100644 --- a/pramen/core/src/main/scala/za/co/absa/pramen/core/bookkeeper/OffsetManagerJdbc.scala +++ b/pramen/core/src/main/scala/za/co/absa/pramen/core/bookkeeper/OffsetManagerJdbc.scala @@ -72,10 +72,11 @@ class OffsetManagerJdbc(db: Database, slickProfile: JdbcProfile, offsetTable: Of } } - override def startWriteOffsets(table: String, infoDate: LocalDate, offsetType: OffsetType): DataOffsetRequest = { + override def startWriteOffsets(table: String, infoDate: LocalDate, offsetType: OffsetType, minOffset: Option[OffsetValue]): DataOffsetRequest = { val createdAt = Instant.now() + val minOffsetStr = minOffset.map(_.valueString).getOrElse("") - val record = OffsetRecord(table, infoDate.toString, offsetType.dataTypeString, "", "", batchId, createdAt.toEpochMilli, None) + val record = OffsetRecord(table, infoDate.toString, offsetType.dataTypeString, minOffsetStr, "", batchId, createdAt.toEpochMilli, None) slickUtils.ensureDbConnected(db) db.run( diff --git a/pramen/core/src/main/scala/za/co/absa/pramen/core/bookkeeper/OffsetManagerUtils.scala b/pramen/core/src/main/scala/za/co/absa/pramen/core/bookkeeper/OffsetManagerUtils.scala index b0462ec06..49974046c 100644 --- a/pramen/core/src/main/scala/za/co/absa/pramen/core/bookkeeper/OffsetManagerUtils.scala +++ b/pramen/core/src/main/scala/za/co/absa/pramen/core/bookkeeper/OffsetManagerUtils.scala @@ -29,10 +29,16 @@ object OffsetManagerUtils { * for a given increment ingestion job. But if data ingested does not cover all partitions, such partitions * might be missing in the list. In this case, values from the 'baseValue' are used. */ - def mergeKafkaPartitions(primaryValue: OffsetValue.KafkaValue, baseValue: OffsetValue.KafkaValue): OffsetValue.KafkaValue = { + def mergeKafkaPartitions(primaryValue: OffsetValue.KafkaValue, baseValue: OffsetValue.KafkaValue, asMinimumOffset: Boolean): OffsetValue.KafkaValue = { val partitions1 = primaryValue.value.map(v => (v.partition, v.offset)).toMap - val partitions2 = baseValue.value.map(v => (v.partition, v.offset)).toMap - val maxPartition = Math.max(partitions1.keys.max, partitions2.keys.max) + val partitions02 = baseValue.value.map(v => (v.partition, v.offset)).toMap + val maxPartition = Math.max(partitions1.keys.max, partitions02.keys.max) + + val partitions2 = if (asMinimumOffset) { + partitions02.map(o => (o._1, o._2 + 1)) + } else { + partitions02 + } val newPartitions = for (i <- 0 to maxPartition) yield { if (partitions1.contains(i)) @@ -69,7 +75,7 @@ object OffsetManagerUtils { val maxValue = OffsetValue.KafkaValue(aggregatedDf.collect().map(row => KafkaPartition(row.getAs[Int](0), row.getAs[Long](2))).toSeq) val (minValueMerged, maxValueMerged) = currentMaxOffset match { - case Some(value: OffsetValue.KafkaValue) => (mergeKafkaPartitions(minValue, value), mergeKafkaPartitions(maxValue, value)) + case Some(value: OffsetValue.KafkaValue) => (mergeKafkaPartitions(minValue, value, asMinimumOffset = true), mergeKafkaPartitions(maxValue, value, asMinimumOffset = false)) case _ => (minValue, maxValue) } diff --git a/pramen/core/src/main/scala/za/co/absa/pramen/core/pipeline/IncrementalIngestionJob.scala b/pramen/core/src/main/scala/za/co/absa/pramen/core/pipeline/IncrementalIngestionJob.scala index 74eb707bf..36b303987 100644 --- a/pramen/core/src/main/scala/za/co/absa/pramen/core/pipeline/IncrementalIngestionJob.scala +++ b/pramen/core/src/main/scala/za/co/absa/pramen/core/pipeline/IncrementalIngestionJob.scala @@ -155,7 +155,6 @@ class IncrementalIngestionJob(operationDef: OperationDef, ) validateOffsetColumn(df, offsetInfo) - val req = om.startWriteOffsets(outputTable.name, infoDate, offsetInfo.offsetType) val currentMaxOffsetOpt = if (offsetInfo.offsetType == OffsetType.KafkaType) { val isReRun = runReason == TaskRunReason.Rerun @@ -169,6 +168,8 @@ class IncrementalIngestionJob(operationDef: OperationDef, None } + val req = om.startWriteOffsets(outputTable.name, infoDate, offsetInfo.offsetType, currentMaxOffsetOpt) + val stats = try { val saveMode = if (isRerun) SaveMode.Overwrite else SaveMode.Append val statsToReturn = metastore.saveTable(outputTable.name, infoDate, dfToSave, inputRecordCount, saveModeOverride = Some(saveMode)) @@ -296,7 +297,7 @@ class IncrementalIngestionJob(operationDef: OperationDef, log.warn(s"Fixing uncommitted offsets. New offset to commit for ${outputTable.name} at $infoDate: " + s"min offset: ${newMinOffset.valueString}, max offset: ${newMaxOffset.valueString}.") - val req = om.startWriteOffsets(outputTable.name, infoDate, offsetInfo.offsetType) + val req = om.startWriteOffsets(outputTable.name, infoDate, offsetInfo.offsetType, None) om.commitOffsets(req, newMinOffset, newMaxOffset) uncommittedOffsets.foreach { of => diff --git a/pramen/core/src/test/scala/za/co/absa/pramen/core/integration/IncrementalPipelineLongFixture.scala b/pramen/core/src/test/scala/za/co/absa/pramen/core/integration/IncrementalPipelineLongFixture.scala index 31dd24d32..41fd5ef4f 100644 --- a/pramen/core/src/test/scala/za/co/absa/pramen/core/integration/IncrementalPipelineLongFixture.scala +++ b/pramen/core/src/test/scala/za/co/absa/pramen/core/integration/IncrementalPipelineLongFixture.scala @@ -251,7 +251,7 @@ class IncrementalPipelineLongFixture extends AnyWordSpec // Adding an uncommitted offset for 2021-02-17 val om = new OffsetManagerJdbc(pramenDb.slickDb, pramenDb.slickProfile, pramenDb.offsetTable, 123L) - om.startWriteOffsets("table1", infoDate.minusDays(1), OffsetType.IntegralType) + om.startWriteOffsets("table1", infoDate.minusDays(1), OffsetType.IntegralType, None) Thread.sleep(10) val table1Path = new Path(tempDir, "table1") @@ -699,7 +699,7 @@ class IncrementalPipelineLongFixture extends AnyWordSpec def testOffsetOnlyDealWithUncommittedOffsetsWithNoPath(metastoreFormat: String): Assertion = { val om = new OffsetManagerJdbc(pramenDb.slickDb, pramenDb.slickProfile, pramenDb.offsetTable, 123L) - om.startWriteOffsets("table1", infoDate, OffsetType.IntegralType) + om.startWriteOffsets("table1", infoDate, OffsetType.IntegralType, None) Thread.sleep(10) @@ -738,7 +738,7 @@ class IncrementalPipelineLongFixture extends AnyWordSpec def testOffsetOnlyDealWithUncommittedOffsetsWithNoData(metastoreFormat: String): Assertion = { val om = new OffsetManagerJdbc(pramenDb.slickDb, pramenDb.slickProfile, pramenDb.offsetTable, 123L) - om.startWriteOffsets("table1", infoDate, OffsetType.IntegralType) + om.startWriteOffsets("table1", infoDate, OffsetType.IntegralType, None) Thread.sleep(10) @@ -789,12 +789,12 @@ class IncrementalPipelineLongFixture extends AnyWordSpec def testOffsetOnlyDealWithUncommittedOffsetsWithData(metastoreFormat: String): Assertion = { val om1 = new OffsetManagerJdbc(pramenDb.slickDb, pramenDb.slickProfile, pramenDb.offsetTable, 123L) - om1.startWriteOffsets("table1", infoDate, OffsetType.IntegralType) + om1.startWriteOffsets("table1", infoDate, OffsetType.IntegralType, None) Thread.sleep(10) val om2 = new OffsetManagerJdbc(pramenDb.slickDb, pramenDb.slickProfile, pramenDb.offsetTable, 123L) - om2.startWriteOffsets("table1", infoDate, OffsetType.IntegralType) + om2.startWriteOffsets("table1", infoDate, OffsetType.IntegralType, None) Thread.sleep(10) @@ -867,7 +867,7 @@ class IncrementalPipelineLongFixture extends AnyWordSpec def testOffsetOnlyFailWhenInputTableDoestHaveOffsetField(metastoreFormat: String): Assertion = { val om1 = new OffsetManagerJdbc(pramenDb.slickDb, pramenDb.slickProfile, pramenDb.offsetTable, 123L) - om1.startWriteOffsets("table1", infoDate, OffsetType.IntegralType) + om1.startWriteOffsets("table1", infoDate, OffsetType.IntegralType, None) Thread.sleep(10) diff --git a/pramen/core/src/test/scala/za/co/absa/pramen/core/tests/bookkeeper/OffsetManagerJdbcSuite.scala b/pramen/core/src/test/scala/za/co/absa/pramen/core/tests/bookkeeper/OffsetManagerJdbcSuite.scala index 4b759cb3f..5111c0c3a 100644 --- a/pramen/core/src/test/scala/za/co/absa/pramen/core/tests/bookkeeper/OffsetManagerJdbcSuite.scala +++ b/pramen/core/src/test/scala/za/co/absa/pramen/core/tests/bookkeeper/OffsetManagerJdbcSuite.scala @@ -65,7 +65,7 @@ class OffsetManagerJdbcSuite extends AnyWordSpec with RelationalDbFixture with B val nextHour = now.plusSeconds(3600) val om = getOffsetManager - om.startWriteOffsets("table1", infoDate, OffsetType.IntegralType) + om.startWriteOffsets("table1", infoDate, OffsetType.IntegralType, None) val actualNonEmpty = om.getOffsets("table1", infoDate) val actualEmpty1 = om.getOffsets("table1", infoDate.plusDays(1)) @@ -87,7 +87,7 @@ class OffsetManagerJdbcSuite extends AnyWordSpec with RelationalDbFixture with B "return committed offsets" in { val om = getOffsetManager - val transactionReference = om.startWriteOffsets("table1", infoDate, OffsetType.IntegralType) + val transactionReference = om.startWriteOffsets("table1", infoDate, OffsetType.IntegralType, None) Thread.sleep(10) om.commitOffsets(transactionReference, OffsetValue.IntegralValue(1), OffsetValue.IntegralValue(100)) @@ -137,15 +137,15 @@ class OffsetManagerJdbcSuite extends AnyWordSpec with RelationalDbFixture with B "return the maximum info date and offsets " in { val om = getOffsetManager - val t1 = om.startWriteOffsets("table1", infoDate, OffsetType.IntegralType) + val t1 = om.startWriteOffsets("table1", infoDate, OffsetType.IntegralType, None) om.commitOffsets(t1, OffsetValue.IntegralValue(1), OffsetValue.IntegralValue(100)) Thread.sleep(10) - val t2 = om.startWriteOffsets("table1", infoDate, OffsetType.IntegralType) + val t2 = om.startWriteOffsets("table1", infoDate, OffsetType.IntegralType, None) om.commitOffsets(t2, OffsetValue.IntegralValue(101), OffsetValue.IntegralValue(200)) Thread.sleep(10) - val t3 = om.startWriteOffsets("table1", infoDate.plusDays(1), OffsetType.IntegralType) + val t3 = om.startWriteOffsets("table1", infoDate.plusDays(1), OffsetType.IntegralType, None) om.commitOffsets(t3, OffsetValue.IntegralValue(201), OffsetValue.IntegralValue(300)) val summaryMultiDay = om.getMaxInfoDateAndOffset("table1", None) @@ -175,7 +175,7 @@ class OffsetManagerJdbcSuite extends AnyWordSpec with RelationalDbFixture with B "work for empty offsets" in { val om = getOffsetManager - val t1 = om.startWriteOffsets("table1", infoDate, OffsetType.IntegralType) + val t1 = om.startWriteOffsets("table1", infoDate, OffsetType.IntegralType, None) om.commitOffsets(t1, OffsetValue.IntegralValue(10), OffsetValue.IntegralValue(100)) val offsets = om.getOffsets("table1", infoDate) @@ -192,11 +192,11 @@ class OffsetManagerJdbcSuite extends AnyWordSpec with RelationalDbFixture with B "work for non-empty offsets" in { val om = getOffsetManager - val t1 = om.startWriteOffsets("table1", infoDate, OffsetType.IntegralType) + val t1 = om.startWriteOffsets("table1", infoDate, OffsetType.IntegralType, None) om.commitOffsets(t1, OffsetValue.IntegralValue(10), OffsetValue.IntegralValue(100)) Thread.sleep(10) - val t2 = om.startWriteOffsets("table1", infoDate, OffsetType.IntegralType) + val t2 = om.startWriteOffsets("table1", infoDate, OffsetType.IntegralType, None) om.commitOffsets(t2, OffsetValue.IntegralValue(101), OffsetValue.IntegralValue(200)) val offsets = om.getOffsets("table1", infoDate) @@ -223,15 +223,15 @@ class OffsetManagerJdbcSuite extends AnyWordSpec with RelationalDbFixture with B "work with reruns that do not change offsets" in { val om = getOffsetManager - val t1 = om.startWriteOffsets("table1", infoDate, OffsetType.IntegralType) + val t1 = om.startWriteOffsets("table1", infoDate, OffsetType.IntegralType, None) om.commitOffsets(t1, OffsetValue.IntegralValue(0), OffsetValue.IntegralValue(100)) Thread.sleep(10) - val t2 = om.startWriteOffsets("table1", infoDate, OffsetType.IntegralType) + val t2 = om.startWriteOffsets("table1", infoDate, OffsetType.IntegralType, None) om.commitOffsets(t2, OffsetValue.IntegralValue(101), OffsetValue.IntegralValue(200)) Thread.sleep(10) - val t3 = om.startWriteOffsets("table1", infoDate, OffsetType.IntegralType) + val t3 = om.startWriteOffsets("table1", infoDate, OffsetType.IntegralType, None) om.commitRerun(t3, OffsetValue.IntegralValue(0), OffsetValue.IntegralValue(200)) val offsets = om.getOffsets("table1", infoDate) @@ -250,15 +250,15 @@ class OffsetManagerJdbcSuite extends AnyWordSpec with RelationalDbFixture with B "work with reruns that do change offsets" in { val om = getOffsetManager - val t1 = om.startWriteOffsets("table1", infoDate.minusDays(1), OffsetType.IntegralType) + val t1 = om.startWriteOffsets("table1", infoDate.minusDays(1), OffsetType.IntegralType, None) om.commitOffsets(t1, OffsetValue.IntegralValue(0), OffsetValue.IntegralValue(100)) Thread.sleep(10) - val t2 = om.startWriteOffsets("table1", infoDate, OffsetType.IntegralType) + val t2 = om.startWriteOffsets("table1", infoDate, OffsetType.IntegralType, None) om.commitOffsets(t2, OffsetValue.IntegralValue(101), OffsetValue.IntegralValue(200)) Thread.sleep(10) - val t3 = om.startWriteOffsets("table1", infoDate, OffsetType.IntegralType) + val t3 = om.startWriteOffsets("table1", infoDate, OffsetType.IntegralType, None) om.commitRerun(t3, OffsetValue.IntegralValue(201), OffsetValue.IntegralValue(300)) val offsets = om.getOffsets("table1", infoDate) @@ -277,15 +277,15 @@ class OffsetManagerJdbcSuite extends AnyWordSpec with RelationalDbFixture with B "work with reruns that deletes all data and no previous offsets" in { val om = getOffsetManager - val t1 = om.startWriteOffsets("table1", infoDate, OffsetType.IntegralType) + val t1 = om.startWriteOffsets("table1", infoDate, OffsetType.IntegralType, None) om.commitOffsets(t1, OffsetValue.IntegralValue(0), OffsetValue.IntegralValue(100)) Thread.sleep(10) - val t2 = om.startWriteOffsets("table1", infoDate, OffsetType.IntegralType) + val t2 = om.startWriteOffsets("table1", infoDate, OffsetType.IntegralType, None) om.commitOffsets(t2, OffsetValue.IntegralValue(101), OffsetValue.IntegralValue(200)) Thread.sleep(10) - val t3 = om.startWriteOffsets("table1", infoDate, OffsetType.IntegralType) + val t3 = om.startWriteOffsets("table1", infoDate, OffsetType.IntegralType, None) om.commitRerun(t3, OffsetValue.IntegralValue(-1), OffsetValue.IntegralValue(-1)) val offsets = om.getOffsets("table1", infoDate) @@ -304,15 +304,15 @@ class OffsetManagerJdbcSuite extends AnyWordSpec with RelationalDbFixture with B "work with reruns that deletes all data and there are previous offsets" in { val om = getOffsetManager - val t1 = om.startWriteOffsets("table1", infoDate.minusDays(1), OffsetType.IntegralType) + val t1 = om.startWriteOffsets("table1", infoDate.minusDays(1), OffsetType.IntegralType, None) om.commitOffsets(t1, OffsetValue.IntegralValue(Long.MinValue), OffsetValue.IntegralValue(0)) Thread.sleep(10) - val t2 = om.startWriteOffsets("table1", infoDate, OffsetType.IntegralType) + val t2 = om.startWriteOffsets("table1", infoDate, OffsetType.IntegralType, None) om.commitOffsets(t2, OffsetValue.IntegralValue(1), OffsetValue.IntegralValue(100)) Thread.sleep(10) - val t3 = om.startWriteOffsets("table1", infoDate, OffsetType.IntegralType) + val t3 = om.startWriteOffsets("table1", infoDate, OffsetType.IntegralType, None) om.commitRerun(t3, OffsetValue.IntegralValue(0), OffsetValue.IntegralValue(0)) val offsets = om.getOffsets("table1", infoDate) @@ -331,15 +331,15 @@ class OffsetManagerJdbcSuite extends AnyWordSpec with RelationalDbFixture with B "throw an exception when offsets are incorrect" in { val om = getOffsetManager - val t1 = om.startWriteOffsets("table1", infoDate, OffsetType.IntegralType) + val t1 = om.startWriteOffsets("table1", infoDate, OffsetType.IntegralType, None) om.commitOffsets(t1, OffsetValue.IntegralValue(Long.MinValue), OffsetValue.IntegralValue(0)) Thread.sleep(10) - val t2 = om.startWriteOffsets("table1", infoDate, OffsetType.IntegralType) + val t2 = om.startWriteOffsets("table1", infoDate, OffsetType.IntegralType, None) om.commitOffsets(t2, OffsetValue.IntegralValue(1), OffsetValue.IntegralValue(100)) Thread.sleep(10) - val t3 = om.startWriteOffsets("table1", infoDate, OffsetType.IntegralType) + val t3 = om.startWriteOffsets("table1", infoDate, OffsetType.IntegralType, None) val ex = intercept[IllegalArgumentException] { om.commitRerun(t3, OffsetValue.IntegralValue(1), OffsetValue.IntegralValue(-1)) } @@ -380,7 +380,7 @@ class OffsetManagerJdbcSuite extends AnyWordSpec with RelationalDbFixture with B "work for empty offsets" in { val om = getOffsetManager - val t1 = om.startWriteOffsets("table1", infoDate, OffsetType.IntegralType) + val t1 = om.startWriteOffsets("table1", infoDate, OffsetType.IntegralType, None) om.rollbackOffsets(t1) val offsets = om.getOffsets("table1", infoDate) @@ -391,11 +391,11 @@ class OffsetManagerJdbcSuite extends AnyWordSpec with RelationalDbFixture with B "work for non-empty offsets" in { val om = getOffsetManager - val t1 = om.startWriteOffsets("table1", infoDate, OffsetType.IntegralType) + val t1 = om.startWriteOffsets("table1", infoDate, OffsetType.IntegralType, None) om.commitOffsets(t1, OffsetValue.IntegralValue(0), OffsetValue.IntegralValue(100)) Thread.sleep(10) - val t2 = om.startWriteOffsets("table1", infoDate, OffsetType.IntegralType) + val t2 = om.startWriteOffsets("table1", infoDate, OffsetType.IntegralType, None) om.rollbackOffsets(t2) val offsets = om.getOffsets("table1", infoDate) diff --git a/pramen/core/src/test/scala/za/co/absa/pramen/core/tests/bookkeeper/OffsetManagerUtilsSuite.scala b/pramen/core/src/test/scala/za/co/absa/pramen/core/tests/bookkeeper/OffsetManagerUtilsSuite.scala index 8bebc59c8..c044f4889 100644 --- a/pramen/core/src/test/scala/za/co/absa/pramen/core/tests/bookkeeper/OffsetManagerUtilsSuite.scala +++ b/pramen/core/src/test/scala/za/co/absa/pramen/core/tests/bookkeeper/OffsetManagerUtilsSuite.scala @@ -47,9 +47,11 @@ class OffsetManagerUtilsSuite extends AnyWordSpec with SparkTestBase { KafkaPartition(4, 50) )) - val mered = OffsetManagerUtils.mergeKafkaPartitions(primary, base) + val mered1 = OffsetManagerUtils.mergeKafkaPartitions(primary, base, asMinimumOffset = false) + val mered2 = OffsetManagerUtils.mergeKafkaPartitions(primary, base, asMinimumOffset = true) - assert(mered.value == primary.value) + assert(mered1.value == primary.value) + assert(mered2.value == primary.value) } "return correct offsets when primary offset has gaps" in { @@ -59,7 +61,7 @@ class OffsetManagerUtilsSuite extends AnyWordSpec with SparkTestBase { KafkaPartition(4, 50) )) - val mered = OffsetManagerUtils.mergeKafkaPartitions(primary, base) + val mered = OffsetManagerUtils.mergeKafkaPartitions(primary, base, asMinimumOffset = false) assert(mered.value == Seq( KafkaPartition(0, 10), @@ -76,15 +78,24 @@ class OffsetManagerUtilsSuite extends AnyWordSpec with SparkTestBase { KafkaPartition(3, 40) )) - val mered = OffsetManagerUtils.mergeKafkaPartitions(primary, base) + val mered1 = OffsetManagerUtils.mergeKafkaPartitions(primary, base, asMinimumOffset = false) + val mered2 = OffsetManagerUtils.mergeKafkaPartitions(primary, base, asMinimumOffset = true) - assert(mered.value == Seq( + assert(mered1.value == Seq( KafkaPartition(0, 11), KafkaPartition(1, 20), KafkaPartition(2, 31), KafkaPartition(3, 40), KafkaPartition(4, 51) )) + + assert(mered2.value == Seq( + KafkaPartition(0, 12), + KafkaPartition(1, 20), + KafkaPartition(2, 32), + KafkaPartition(3, 40), + KafkaPartition(4, 52) + )) } } diff --git a/pramen/extras/src/main/scala/za/co/absa/pramen/extras/source/KafkaAvroSource.scala b/pramen/extras/src/main/scala/za/co/absa/pramen/extras/source/KafkaAvroSource.scala index 119873b08..955781a14 100644 --- a/pramen/extras/src/main/scala/za/co/absa/pramen/extras/source/KafkaAvroSource.scala +++ b/pramen/extras/src/main/scala/za/co/absa/pramen/extras/source/KafkaAvroSource.scala @@ -163,32 +163,40 @@ class KafkaAvroSource(sourceConfig: Config, .get("kafka.endingOffsets") .map(v => s"{$q$topic$q: $v}") - val startingOffsets = offsetFromOpt match { - case Some(offset) => - // The starting offset is inclusive in Spark. - if (offset.dataType != OffsetType.KafkaType) - throw new IllegalArgumentException(s"KafkaAvroSource supports only 'kafka' offsets, got ${offset.dataType.dataTypeString}") - - // If 'to' part of the interval is defined, use the closed interval min <= offset <= max, otherwise min < offset. - val offsetFrom = if (offsetToOpt.isDefined) - offset - else - offset.asInstanceOf[KafkaValue].increment - - Map("startingOffsets" -> s"{$q$topic$q: ${offsetFrom.valueString}}") - case None => - Map("startingOffsets" -> proposedStartingOffset.getOrElse("earliest")) + val startingOffsets = proposedStartingOffset match { + case Some(offsets) => Map("startingOffsets" -> s"{$q$topic$q: $offsets}") + case None => offsetFromOpt match { + case Some(offset) => + // The starting offset is inclusive in Spark. + if (offset.dataType != OffsetType.KafkaType) + throw new IllegalArgumentException(s"KafkaAvroSource supports only 'kafka' offsets, got ${offset.dataType.dataTypeString}") + + // If 'to' part of the interval is defined, use the closed interval min <= offset <= max, otherwise min < offset. + val offsetFrom = if (offsetToOpt.isDefined) + offset + else + offset.asInstanceOf[KafkaValue].increment + + Map("startingOffsets" -> s"{$q$topic$q: ${offsetFrom.valueString}}") + case None => + Map("startingOffsets" -> "earliest") + } } - val endingOffsets = offsetToOpt match { - case Some(offset) => - // The ending offset is exclusive in Spark. - if (offset.dataType != OffsetType.KafkaType) - throw new IllegalArgumentException(s"KafkaAvroSource supports only 'kafka' offsets, got ${offset.dataType.dataTypeString}") - val offsetTo = offset.asInstanceOf[KafkaValue].increment - Map("endingOffsets" -> s"{$q$topic$q: ${offsetTo.valueString}}") - case None => - Map("endingOffsets" -> proposedEndingOffset.getOrElse("latest")) + val endingOffsets = proposedEndingOffset match { + case Some(offsets) => + Map("endingOffsets" -> s"{$q$topic$q: $offsets}") + case None => + offsetToOpt match { + case Some(offset) => + // The ending offset is exclusive in Spark. + if (offset.dataType != OffsetType.KafkaType) + throw new IllegalArgumentException(s"KafkaAvroSource supports only 'kafka' offsets, got ${offset.dataType.dataTypeString}") + val offsetTo = offset.asInstanceOf[KafkaValue].increment + Map("endingOffsets" -> s"{$q$topic$q: ${offsetTo.valueString}}") + case None => + Map("endingOffsets" -> "latest") + } } val kafkaOptions = kafkaAvroConfig.extraOptions ++ startingOffsets ++ endingOffsets + From 0ed9f83abefc4f09aaf297f5b434ceb7137163a6 Mon Sep 17 00:00:00 2001 From: Ruslan Iushchenko Date: Thu, 15 Oct 2026 14:50:00 +0200 Subject: [PATCH 4/5] Fix reruns for Kafka incremental ingestion. --- .../scala/za/co/absa/pramen/api/offset/OffsetValue.scala | 4 ++++ .../absa/pramen/core/bookkeeper/OffsetManagerJdbc.scala | 8 +++++++- 2 files changed, 11 insertions(+), 1 deletion(-) diff --git a/pramen/api/src/main/scala/za/co/absa/pramen/api/offset/OffsetValue.scala b/pramen/api/src/main/scala/za/co/absa/pramen/api/offset/OffsetValue.scala index d62adc2cf..5e6e62fd1 100644 --- a/pramen/api/src/main/scala/za/co/absa/pramen/api/offset/OffsetValue.scala +++ b/pramen/api/src/main/scala/za/co/absa/pramen/api/offset/OffsetValue.scala @@ -122,6 +122,10 @@ object OffsetValue { def increment: OffsetValue = { KafkaValue(value.map(p => p.copy(offset = p.offset + 1))) } + + def decrement: OffsetValue = { + KafkaValue(value.map(p => p.copy(offset = p.offset - 1))) + } } diff --git a/pramen/core/src/main/scala/za/co/absa/pramen/core/bookkeeper/OffsetManagerJdbc.scala b/pramen/core/src/main/scala/za/co/absa/pramen/core/bookkeeper/OffsetManagerJdbc.scala index c386b9612..1f9ad6856 100644 --- a/pramen/core/src/main/scala/za/co/absa/pramen/core/bookkeeper/OffsetManagerJdbc.scala +++ b/pramen/core/src/main/scala/za/co/absa/pramen/core/bookkeeper/OffsetManagerJdbc.scala @@ -99,7 +99,13 @@ class OffsetManagerJdbc(db: Database, slickProfile: JdbcProfile, offsetTable: Of } override def commitRerun(request: DataOffsetRequest, minOffset: OffsetValue, maxOffset: OffsetValue): Unit = { - if (minOffset.compareTo(maxOffset) > 0) { + val minOffsetToCheck = if (minOffset.dataType == OffsetType.KafkaType) { + minOffset.asInstanceOf[OffsetValue.KafkaValue].decrement + } else { + minOffset + } + + if (minOffsetToCheck.compareTo(maxOffset) > 0) { throw new IllegalArgumentException(s"minOffset is greater than maxOffset: ${minOffset.valueString} > ${maxOffset.valueString}") } From b4820e7ccf14785b063c4b83f6ad73950ecc3fc1 Mon Sep 17 00:00:00 2001 From: Ruslan Iushchenko Date: Fri, 16 Oct 2026 11:27:00 +0200 Subject: [PATCH 5/5] Fix Kafka offset handling for reruns, empty partition merges, and proposed offset formatting --- .../pramen/core/bookkeeper/OffsetManagerDynamoDb.scala | 8 +++++++- .../absa/pramen/core/bookkeeper/OffsetManagerUtils.scala | 4 +++- .../za/co/absa/pramen/extras/source/KafkaAvroSource.scala | 4 ++-- 3 files changed, 12 insertions(+), 4 deletions(-) diff --git a/pramen/core/src/main/scala/za/co/absa/pramen/core/bookkeeper/OffsetManagerDynamoDb.scala b/pramen/core/src/main/scala/za/co/absa/pramen/core/bookkeeper/OffsetManagerDynamoDb.scala index 0c4b083dc..7ca8009a3 100644 --- a/pramen/core/src/main/scala/za/co/absa/pramen/core/bookkeeper/OffsetManagerDynamoDb.scala +++ b/pramen/core/src/main/scala/za/co/absa/pramen/core/bookkeeper/OffsetManagerDynamoDb.scala @@ -189,7 +189,13 @@ class OffsetManagerDynamoDb( } override def commitRerun(request: DataOffsetRequest, minOffset: OffsetValue, maxOffset: OffsetValue): Unit = { - if (minOffset.compareTo(maxOffset) > 0) { + val minOffsetToCheck = if (minOffset.dataType == OffsetType.KafkaType) { + minOffset.asInstanceOf[OffsetValue.KafkaValue].decrement + } else { + minOffset + } + + if (minOffsetToCheck.compareTo(maxOffset) > 0) { throw new IllegalArgumentException(s"minOffset is greater than maxOffset: ${minOffset.valueString} > ${maxOffset.valueString}") } diff --git a/pramen/core/src/main/scala/za/co/absa/pramen/core/bookkeeper/OffsetManagerUtils.scala b/pramen/core/src/main/scala/za/co/absa/pramen/core/bookkeeper/OffsetManagerUtils.scala index 49974046c..8f4709a69 100644 --- a/pramen/core/src/main/scala/za/co/absa/pramen/core/bookkeeper/OffsetManagerUtils.scala +++ b/pramen/core/src/main/scala/za/co/absa/pramen/core/bookkeeper/OffsetManagerUtils.scala @@ -32,7 +32,9 @@ object OffsetManagerUtils { def mergeKafkaPartitions(primaryValue: OffsetValue.KafkaValue, baseValue: OffsetValue.KafkaValue, asMinimumOffset: Boolean): OffsetValue.KafkaValue = { val partitions1 = primaryValue.value.map(v => (v.partition, v.offset)).toMap val partitions02 = baseValue.value.map(v => (v.partition, v.offset)).toMap - val maxPartition = Math.max(partitions1.keys.max, partitions02.keys.max) + val allPartitions = partitions1.keys ++ partitions02.keys + if (allPartitions.isEmpty) return OffsetValue.KafkaValue(Seq.empty) + val maxPartition = allPartitions.max val partitions2 = if (asMinimumOffset) { partitions02.map(o => (o._1, o._2 + 1)) diff --git a/pramen/extras/src/main/scala/za/co/absa/pramen/extras/source/KafkaAvroSource.scala b/pramen/extras/src/main/scala/za/co/absa/pramen/extras/source/KafkaAvroSource.scala index 955781a14..1c4a7e0b1 100644 --- a/pramen/extras/src/main/scala/za/co/absa/pramen/extras/source/KafkaAvroSource.scala +++ b/pramen/extras/src/main/scala/za/co/absa/pramen/extras/source/KafkaAvroSource.scala @@ -164,7 +164,7 @@ class KafkaAvroSource(sourceConfig: Config, .map(v => s"{$q$topic$q: $v}") val startingOffsets = proposedStartingOffset match { - case Some(offsets) => Map("startingOffsets" -> s"{$q$topic$q: $offsets}") + case Some(offsets) => Map("startingOffsets" -> offsets) case None => offsetFromOpt match { case Some(offset) => // The starting offset is inclusive in Spark. @@ -185,7 +185,7 @@ class KafkaAvroSource(sourceConfig: Config, val endingOffsets = proposedEndingOffset match { case Some(offsets) => - Map("endingOffsets" -> s"{$q$topic$q: $offsets}") + Map("endingOffsets" -> offsets) case None => offsetToOpt match { case Some(offset) =>