Fixes bug preventing usage of feedRangeFilter with change feed (Azure#29338)

* Fixed bug preventing usage of feedRangeFilter with change feed

* Fixed FeedRange split when FeedRange is smaller than requested split count

Author: FabianMeiswinkel

sdk/cosmos/azure-cosmos-spark_3_2-12/src/main/scala/com/azure/cosmos/spark/CosmosPartitionPlanner.scala

@@ -396,25 +396,40 @@ private object CosmosPartitionPlanner extends BasicLoggingTrait {
         Some(range)
       })
 
+    val result = new ArrayBuffer[PartitionMetadata]
     orderedRanges
-      .map(range => {
-        while (!SparkBridgeImplementationInternal.doRangesOverlap(range, startTokens(startTokensIndex)._1)) {
+      .foreach(range => {
+        logInfo(s"merging range $range")
+        val initialStartTokensIndex = startTokensIndex
+        val initialLatestTokensIndex = latestTokensIndex
+        while (startTokensIndex < startTokens.length &&
+          !SparkBridgeImplementationInternal.doRangesOverlap(range, startTokens(startTokensIndex)._1)) {
+
           startTokensIndex += 1
-          if (startTokensIndex >= startTokens.length) {
-            throw new IllegalStateException(s"No overlapping start token found for range '$range'.")
-          }
         }
 
-        while (!SparkBridgeImplementationInternal.doRangesOverlap(range, latestTokens(latestTokensIndex).feedRange)) {
+        while (startTokensIndex < startTokens.length &&
+          latestTokensIndex < latestTokens.length &&
+          !SparkBridgeImplementationInternal.doRangesOverlap(range, latestTokens(latestTokensIndex).feedRange)) {
+
           latestTokensIndex += 1
-          if (latestTokensIndex >= latestTokens.length) {
-            throw new IllegalStateException(s"No overlapping latest token found for range '$range'.")
-          }
         }
 
-        val startLsn: Long = startTokens(startTokensIndex)._2
-        latestTokens(latestTokensIndex).cloneForSubRange(range, startLsn)
+        if (startTokensIndex < startTokens.length &&
+          latestTokensIndex < latestTokens.length) {
+
+          val startLsn: Long = startTokens(startTokensIndex)._2
+          val newToken = latestTokens(latestTokensIndex).cloneForSubRange(range, startLsn)
+          result.append(newToken)
+        } else {
+          startTokensIndex = initialStartTokensIndex
+          latestTokensIndex = initialLatestTokensIndex
+        }
       })
+
+    assert(result.size > 0)
+
+    result.toArray
   }
   // scalastyle:on method.length
 

sdk/cosmos/azure-cosmos-spark_3_2-12/src/test/scala/com/azure/cosmos/spark/SparkE2EStructuredStreamingITest.scala

@@ -4,8 +4,8 @@ package com.azure.cosmos.spark
 
 import com.azure.cosmos.CosmosAsyncContainer
 import com.azure.cosmos.implementation.{TestConfigurations, Utils}
-import com.azure.cosmos.models.{PartitionKey, ThroughputProperties}
-import org.apache.spark.sql.SparkSession
+import com.azure.cosmos.models.{ModelBridgeInternal, PartitionKey, ThroughputProperties}
+import org.apache.spark.sql.{DataFrame, SparkSession}
 import org.apache.spark.sql.streaming.{StreamingQueryListener, Trigger}
 import org.apache.spark.sql.streaming.StreamingQueryListener.{QueryProgressEvent, QueryStartedEvent, QueryTerminatedEvent}
 
@@ -663,6 +663,100 @@ class SparkE2EStructuredStreamingITest
     targetContainer.delete().block()
   }
 
+  "spark change feed micro batch (incremental)" can
+    "filter by feedRange (Restrictive partitioning strategy)" taggedAs(Retryable) in {
+
+    runChangeFeedFeedRangeFilterTest("Restrictive")
+  }
+
+  "spark change feed micro batch (incremental)" can
+    "filter by feedRange (Default partitioning strategy)" taggedAs(Retryable) in {
+
+    runChangeFeedFeedRangeFilterTest("Default")
+  }
+
+  "spark change feed micro batch (incremental)" can
+    "filter by feedRange (Aggressive partitioning strategy)" taggedAs(Retryable) in {
+
+    runChangeFeedFeedRangeFilterTest("Aggressive")
+  }
+
+  private[this] def runChangeFeedFeedRangeFilterTest(partitioningStrategy: String): Unit = {
+    val processedRecordCount = new AtomicLong(0)
+    val forEachBatchRecordCount = new AtomicLong(0)
+    var spark = this.createSparkSession(processedRecordCount)
+    val cosmosEndpoint = TestConfigurations.HOST
+    val cosmosMasterKey = TestConfigurations.MASTER_KEY
+    val sourceContainer = cosmosClient.getDatabase(cosmosDatabase).getContainer(cosmosContainer)
+    val testId = UUID.randomUUID().toString
+
+    // Initially ingest 100 records
+    var lastId = ""
+    for (i <- 0 until 20) {
+      lastId = this.ingestTestDocument(sourceContainer, i)
+    }
+
+    Thread.sleep(2100)
+
+    val pkDefinition = sourceContainer.read().block().getProperties.getPartitionKeyDefinition
+    val pkDefinitionJson = ModelBridgeInternal.getJsonSerializable(pkDefinition).toJson
+
+    val feedRangeFilter = new GetFeedRangeForPartitionKeyValue().call(pkDefinitionJson, lastId)
+
+    val changeFeedCfg = Map(
+      "spark.cosmos.accountEndpoint" -> cosmosEndpoint,
+      "spark.cosmos.accountKey" -> cosmosMasterKey,
+      "spark.cosmos.database" -> cosmosDatabase,
+      "spark.cosmos.container" -> cosmosContainer,
+      "spark.cosmos.read.inferSchema.enabled" -> "false",
+      "spark.cosmos.read.partitioning.strategy" -> partitioningStrategy,
+      "spark.cosmos.partitioning.feedRangeFilter" -> feedRangeFilter
+    )
+
+    val changeFeedDF = spark
+      .readStream
+      .format("cosmos.oltp.changeFeed")
+      .options(changeFeedCfg)
+      .load()
+
+    val microBatchQuery = changeFeedDF
+      .writeStream
+      .foreachBatch { (batchDF: DataFrame, batchId: Long) =>
+        batchDF.persist()
+        val recordCount = batchDF.count()
+        forEachBatchRecordCount.addAndGet(recordCount)
+        println(s"BatchId: $batchId, Document count: $recordCount")
+        batchDF.unpersist()
+        ()
+      }
+      .trigger(Trigger.ProcessingTime("500 milliseconds"))
+      .queryName(testId)
+      .option("checkpointLocation", s"/tmp/$testId/")
+      .start()
+
+    Thread.sleep(5000)
+
+    microBatchQuery.lastProgress should not be null
+    microBatchQuery.lastProgress.sources should not be null
+    microBatchQuery.lastProgress.sources should not be null
+    microBatchQuery.lastProgress.sources(0).endOffset should not be null
+    getPartitionCountInOffset(microBatchQuery.lastProgress.sources(0).endOffset) >= 1 shouldEqual true
+
+    microBatchQuery.stop()
+
+    var sourceCount: Long = getRecordCountOfContainer(sourceContainer)
+    logInfo(s"RecordCount in source container after first execution: $sourceCount")
+
+    forEachBatchRecordCount.get() shouldEqual 1L
+    processedRecordCount.get() shouldEqual 1L
+    sourceCount shouldEqual 20L
+
+    // close and recreate spark session to validate
+    // that it is possible to recover the previous query
+    // from the commit log
+    spark.close()
+  }
+
   private[this] def ingestTestDocument
   (
     container: CosmosAsyncContainer,
@@ -690,7 +784,7 @@ class SparkE2EStructuredStreamingITest
       override def onQueryStarted(queryStarted: QueryStartedEvent): Unit = {}
       override def onQueryTerminated(queryTerminated: QueryTerminatedEvent): Unit = {}
       override def onQueryProgress(queryProgress: QueryProgressEvent): Unit = {
-        processedRecordCount.addAndGet(queryProgress.progress.sink.numOutputRows)
+        processedRecordCount.addAndGet(queryProgress.progress.numInputRows)
       }
     })
 

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/feedranges/FeedRangeInternal.java

@@ -265,31 +265,35 @@ static List<FeedRangeEpkImpl> trySplitWithHashV1(
         String minRange = effectiveRange.getMin();
         long diff = max - min;
         List<FeedRangeEpkImpl> splitFeedRanges = new ArrayList<>(targetedSplitCount);
-        for (int i = 1; i < targetedSplitCount; i++) {
-            long splitPoint = min + (i * (diff / targetedSplitCount));
-            String maxRange = PartitionKeyInternalHelper.toHexEncodedBinaryString(
-                new NumberPartitionKeyComponent[] {
-                    new NumberPartitionKeyComponent(splitPoint)
-                });
+        if (diff < targetedSplitCount) {
+            splitFeedRanges.add(new FeedRangeEpkImpl(effectiveRange));
+        } else {
+            for (int i = 1; i < targetedSplitCount; i++) {
+                long splitPoint = min + (i * (diff / targetedSplitCount));
+                String maxRange = PartitionKeyInternalHelper.toHexEncodedBinaryString(
+                    new NumberPartitionKeyComponent[] {
+                        new NumberPartitionKeyComponent(splitPoint)
+                    });
+                splitFeedRanges.add(
+                    new FeedRangeEpkImpl(
+                        new Range<>(
+                            minRange,
+                            maxRange,
+                            i > 1 || effectiveRange.isMinInclusive(),
+                            false)));
+
+                minRange = maxRange;
+            }
+
             splitFeedRanges.add(
                 new FeedRangeEpkImpl(
                     new Range<>(
                         minRange,
-                        maxRange,
-                        i > 1 || effectiveRange.isMinInclusive(),
-                        false)));
-
-            minRange = maxRange;
+                        effectiveRange.getMax(),
+                        true,
+                        effectiveRange.isMaxInclusive())));
         }
 
-        splitFeedRanges.add(
-            new FeedRangeEpkImpl(
-                new Range<>(
-                    minRange,
-                    effectiveRange.getMax(),
-                    true,
-                    effectiveRange.isMaxInclusive())));
-
         return splitFeedRanges;
     }