improvePerfDuringPartitionPlanning (#46727)

* improve perf for partition planning

---------

Co-authored-by: annie-mac <[email protected]>

Author: xinlian12

sdk/cosmos/azure-cosmos-spark_3-3_2-12/CHANGELOG.md

@@ -9,6 +9,7 @@
 #### Bugs Fixed
 
 #### Other Changes
+* Added improvement to reduce partition planning time for large containers. - See [46727](https://github.com/Azure/azure-sdk-for-java/pull/46727)
 
 ### 4.39.0 (2025-09-05)
 

sdk/cosmos/azure-cosmos-spark_3-3_2-12/src/main/scala/com/azure/cosmos/spark/ChangeFeedMicroBatchStream.scala

@@ -2,9 +2,9 @@
 // Licensed under the MIT License.
 package com.azure.cosmos.spark
 
+import com.azure.cosmos.changeFeedMetrics.{ChangeFeedMetricsListener, ChangeFeedMetricsTracker}
 import com.azure.cosmos.implementation.guava25.collect.{HashBiMap, Maps}
 import com.azure.cosmos.implementation.{SparkBridgeImplementationInternal, UUIDs}
-import com.azure.cosmos.changeFeedMetrics.{ChangeFeedMetricsListener, ChangeFeedMetricsTracker}
 import com.azure.cosmos.spark.CosmosPredicates.{assertNotNull, assertNotNullOrEmpty, assertOnSparkDriver}
 import com.azure.cosmos.spark.diagnostics.{DiagnosticsContext, LoggerHelper}
 import org.apache.spark.broadcast.Broadcast
@@ -115,6 +115,7 @@ private class ChangeFeedMicroBatchStream
 
     assert(end.inputPartitions.isDefined, "Argument 'endOffset.inputPartitions' must not be null or empty.")
 
+    val parsedStartChangeFeedState = SparkBridgeImplementationInternal.parseChangeFeedState(start.changeFeedState)
     end
       .inputPartitions
       .get
@@ -123,7 +124,7 @@ private class ChangeFeedMicroBatchStream
           partition
            .withContinuationState(
              SparkBridgeImplementationInternal
-              .extractChangeFeedStateForRange(start.changeFeedState, partition.feedRange),
+              .extractChangeFeedStateForRange(parsedStartChangeFeedState, partition.feedRange),
              clearEndLsn = false)
            .withIndex(index)
       })

sdk/cosmos/azure-cosmos-spark_3-4_2-12/CHANGELOG.md

@@ -9,6 +9,7 @@
 #### Bugs Fixed
 
 #### Other Changes
+* Added improvement to reduce partition planning time for large containers. - See [46727](https://github.com/Azure/azure-sdk-for-java/pull/46727)
 
 ### 4.39.0 (2025-09-05)
 

sdk/cosmos/azure-cosmos-spark_3-4_2-12/src/main/scala/com/azure/cosmos/spark/ChangeFeedMicroBatchStream.scala

@@ -116,6 +116,7 @@ private class ChangeFeedMicroBatchStream
 
     assert(end.inputPartitions.isDefined, "Argument 'endOffset.inputPartitions' must not be null or empty.")
 
+    val parsedStartChangeFeedState = SparkBridgeImplementationInternal.parseChangeFeedState(start.changeFeedState)
     end
       .inputPartitions
       .get
@@ -124,7 +125,7 @@ private class ChangeFeedMicroBatchStream
         partition
          .withContinuationState(
            SparkBridgeImplementationInternal
-            .extractChangeFeedStateForRange(start.changeFeedState, partition.feedRange),
+            .extractChangeFeedStateForRange(parsedStartChangeFeedState, partition.feedRange),
            clearEndLsn = false)
          .withIndex(index)
       })

sdk/cosmos/azure-cosmos-spark_3-5_2-12/CHANGELOG.md

@@ -9,6 +9,7 @@
 #### Bugs Fixed
 
 #### Other Changes
+* Added improvement to reduce partition planning time for large containers. - See [46727](https://github.com/Azure/azure-sdk-for-java/pull/46727)
 
 ### 4.39.0 (2025-09-05)
 

sdk/cosmos/azure-cosmos-spark_3-5_2-12/src/main/scala/com/azure/cosmos/spark/ChangeFeedMicroBatchStream.scala

@@ -117,6 +117,7 @@ private class ChangeFeedMicroBatchStream
 
     assert(end.inputPartitions.isDefined, "Argument 'endOffset.inputPartitions' must not be null or empty.")
 
+    val parsedStartChangeFeedState = SparkBridgeImplementationInternal.parseChangeFeedState(start.changeFeedState)
     end
       .inputPartitions
       .get
@@ -125,7 +126,7 @@ private class ChangeFeedMicroBatchStream
         partition
          .withContinuationState(
            SparkBridgeImplementationInternal
-            .extractChangeFeedStateForRange(start.changeFeedState, partition.feedRange),
+            .extractChangeFeedStateForRange(parsedStartChangeFeedState, partition.feedRange),
            clearEndLsn = false)
          .withIndex(index)
       })

sdk/cosmos/azure-cosmos-spark_3_2-12/src/main/scala/com/azure/cosmos/implementation/SparkBridgeImplementationInternal.scala

@@ -172,16 +172,19 @@ private[cosmos] object SparkBridgeImplementationInternal extends BasicLoggingTra
 
   def extractChangeFeedStateForRange
   (
-    stateJsonBase64: String,
+    changeFeedState: ChangeFeedState,
     feedRange: NormalizedRange
   ): String = {
-    assert(!Strings.isNullOrWhiteSpace(stateJsonBase64), s"Argument 'stateJsonBase64' must not be null or empty.")
-    ChangeFeedState
-      .fromString(stateJsonBase64)
+    changeFeedState
       .extractForEffectiveRange(toCosmosRange(feedRange))
       .toString
   }
 
+  def parseChangeFeedState(changeFeedStateJsonString: String): ChangeFeedState = {
+    assert(!Strings.isNullOrWhiteSpace(changeFeedStateJsonString), s"Argument 'changeFeedStateJsonString' must not be null or empty.")
+    ChangeFeedState.fromString(changeFeedStateJsonString)
+  }
+
   def toFeedRange(range: NormalizedRange): FeedRange = {
     new FeedRangeEpkImpl(toCosmosRange(range))
   }

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

@@ -185,13 +185,14 @@ private class ChangeFeedBatch
 
       // Latest offset above has the EndLsn specified based on the point-in-time latest offset
       // For batch mode instead we need to reset it so that the change feed will get fully drained
+      val parsedInitialOffset = SparkBridgeImplementationInternal.parseChangeFeedState(initialOffsetJson)
       val inputPartitions = latestOffset
         .inputPartitions
         .get
         .map(partition => partition
           .withContinuationState(
             SparkBridgeImplementationInternal
-              .extractChangeFeedStateForRange(initialOffsetJson, partition.feedRange),
+              .extractChangeFeedStateForRange(parsedInitialOffset, partition.feedRange),
             clearEndLsn = !hasBatchCheckpointLocation))
         .map(_.asInstanceOf[InputPartition])
 

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

@@ -80,11 +80,11 @@ class CreateSpark2ContinuationsFromChangeFeedOffset extends UDF2[Map[String, Str
     pkRanges
       .foreach(pkRange => {
         val normalizedRange = rangeToNormalizedRange(pkRange.toRange)
-
+        val parsedChangeFeedState = SparkBridgeImplementationInternal.parseChangeFeedState(offset.changeFeedState)
         val effectiveChangeFeedState = ChangeFeedState
           .fromString(
             SparkBridgeImplementationInternal
-              .extractChangeFeedStateForRange(offset.changeFeedState, normalizedRange)
+              .extractChangeFeedStateForRange(parsedChangeFeedState, normalizedRange)
           )
 
         val containerResourceId = effectiveChangeFeedState.getContainerRid