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/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/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..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 @@ -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() @@ -188,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/OffsetManagerJdbc.scala b/pramen/core/src/main/scala/za/co/absa/pramen/core/bookkeeper/OffsetManagerJdbc.scala index b1e164860..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 @@ -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( @@ -98,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}") } 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..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 @@ -23,7 +23,43 @@ 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, 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 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)) + } else { + partitions02 + } + + 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 +73,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, asMinimumOffset = true), mergeKafkaPartitions(maxValue, value, asMinimumOffset = false)) + 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..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 @@ -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) => @@ -156,7 +155,20 @@ 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 req = om.startWriteOffsets(outputTable.name, infoDate, offsetInfo.offsetType, currentMaxOffsetOpt) val stats = try { val saveMode = if (isRerun) SaveMode.Overwrite else SaveMode.Append @@ -167,7 +179,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,12 +290,14 @@ 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}.") - 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 dc25c1807..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 @@ -29,11 +29,81 @@ 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 mered1 = OffsetManagerUtils.mergeKafkaPartitions(primary, base, asMinimumOffset = false) + val mered2 = OffsetManagerUtils.mergeKafkaPartitions(primary, base, asMinimumOffset = true) + + assert(mered1.value == primary.value) + assert(mered2.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, asMinimumOffset = false) + + 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 mered1 = OffsetManagerUtils.mergeKafkaPartitions(primary, base, asMinimumOffset = false) + val mered2 = OffsetManagerUtils.mergeKafkaPartitions(primary, base, asMinimumOffset = true) + + 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) + )) + } + } + "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 +112,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 +133,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 +144,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))) 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..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 @@ -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" -> 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" -> 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 +