change

Author: annie-mac

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

@@ -2,7 +2,9 @@
 // Licensed under the MIT License.
 package com.azure.cosmos.spark
 
+import com.azure.cosmos.changeFeedMetrics.{ChangeFeedMetricsListener, ChangeFeedMetricsTracker}
 import com.azure.cosmos.implementation.SparkBridgeImplementationInternal
+import com.azure.cosmos.implementation.guava25.collect.{HashBiMap, Maps}
 import com.azure.cosmos.spark.CosmosPredicates.{assertNotNull, assertNotNullOrEmpty, assertOnSparkDriver}
 import com.azure.cosmos.spark.diagnostics.{DiagnosticsContext, LoggerHelper}
 import org.apache.spark.broadcast.Broadcast
@@ -13,6 +15,12 @@ import org.apache.spark.sql.types.StructType
 
 import java.time.Duration
 import java.util.UUID
+import java.util.concurrent.ConcurrentHashMap
+import java.util.concurrent.atomic.AtomicLong
+
+// scalastyle:off underscore.import
+import scala.collection.JavaConverters._
+// scalastyle:on underscore.import
 
 // scala style rule flaky - even complaining on partial log messages
 // scalastyle:off multiple.string.literals
@@ -59,6 +67,12 @@ private class ChangeFeedMicroBatchStream
 
   private var latestOffsetSnapshot: Option[ChangeFeedOffset] = None
 
+  private val partitionIndex = new AtomicLong(0)
+  private val partitionIndexMap = Maps.synchronizedBiMap(HashBiMap.create[NormalizedRange, Long]())
+  private val partitionMetricsMap = new ConcurrentHashMap[NormalizedRange, ChangeFeedMetricsTracker]()
+
+  session.sparkContext.addSparkListener(new ChangeFeedMetricsListener(partitionIndexMap, partitionMetricsMap))
+
   override def latestOffset(): Offset = {
     // For Spark data streams implementing SupportsAdmissionControl trait
     // latestOffset(Offset, ReadLimit) is called instead
@@ -101,11 +115,15 @@ private class ChangeFeedMicroBatchStream
     end
       .inputPartitions
       .get
-      .map(partition => partition
-        .withContinuationState(
-          SparkBridgeImplementationInternal
+      .map(partition => {
+        val index = partitionIndexMap.asScala.getOrElseUpdate(partition.feedRange, partitionIndex.incrementAndGet())
+        partition
+         .withContinuationState(
+           SparkBridgeImplementationInternal
             .extractChangeFeedStateForRange(start.changeFeedState, partition.feedRange),
-          clearEndLsn = false))
+           clearEndLsn = false)
+         .withIndex(index)
+      })
   }
 
   /**

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

@@ -550,37 +550,37 @@ private object CosmosPartitionPlanner extends BasicLoggingTrait {
                                        metadata: PartitionMetadata,
                                        isChangeFeed: Boolean,
                                        partitionMetricsMap: Option[ConcurrentHashMap[NormalizedRange, ChangeFeedMetricsTracker]] = None): Long = {
-    var normalizedChangesPerLsnOpt: Option[Double] = None
+    var changesPerLsnFromMetricsOpt: Option[Double] = None
     if (isChangeFeed) {
       partitionMetricsMap match {
         case Some(metricsMap) =>
           if (metricsMap.containsKey(metadata.feedRange)) {
-            normalizedChangesPerLsnOpt = metricsMap.get(metadata.feedRange).getWeightedAvgChangesPerLsn
+            changesPerLsnFromMetricsOpt = metricsMap.get(metadata.feedRange).getWeightedAvgChangesPerLsn
           }
         case None =>
       }
     }
 
-    metadata.getWeightedLsnGap(normalizedChangesPerLsnOpt)
+    metadata.getWeightedLsnGap(changesPerLsnFromMetricsOpt)
   }
 
   private[spark] def getAvgChangesPerLsn(
                                                 metadata: PartitionMetadata,
                                                 isChangeFeed: Boolean,
                                                 partitionMetricsMap: Option[ConcurrentHashMap[NormalizedRange, ChangeFeedMetricsTracker]] = None): Double = {
-    var normalizedChangesPerLsnOpt: Option[Double] = None
+    var changesPerLsnFromMetricsOpt: Option[Double] = None
     if (isChangeFeed) {
       partitionMetricsMap match {
         case Some(matricsMap) =>
           if (matricsMap.containsKey(metadata.feedRange)) {
-            normalizedChangesPerLsnOpt = partitionMetricsMap.get.get(metadata.feedRange).getWeightedAvgChangesPerLsn
+            changesPerLsnFromMetricsOpt = partitionMetricsMap.get.get(metadata.feedRange).getWeightedAvgChangesPerLsn
           }
         case None =>
       }
     }
 
-    if (normalizedChangesPerLsnOpt.isDefined) {
-      normalizedChangesPerLsnOpt.get
+    if (changesPerLsnFromMetricsOpt.isDefined) {
+      changesPerLsnFromMetricsOpt.get
     } else {
       metadata.getAvgChangesPerLsn
     }

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

@@ -114,18 +114,18 @@ private[cosmos] case class PartitionMetadata
     )
   }
 
-  def getWeightedLsnGap(normalizedChangesPerLsn: Option[Double] = None): Long = {
+  def getWeightedLsnGap(changesPerLsnFromMetrics: Option[Double] = None): Long = {
     val progressFactor = math.max(this.getAndValidateLatestLsn - this.startLsn, 0)
     if (progressFactor == 0) {
       0
     } else {
-      val effectiveNormalizedChangesPerLsn =
-        normalizedChangesPerLsn match {
+      val effectiveChangesPerLsn =
+        changesPerLsnFromMetrics match {
           case Some(changesPerLsn) => changesPerLsn
           case None => this.getAvgChangesPerLsn
         }
 
-      val weightedGap: Double = progressFactor * effectiveNormalizedChangesPerLsn
+      val weightedGap: Double = progressFactor * effectiveChangesPerLsn
       // Any double less than 1 gets rounded to 0 when toLong is invoked
       weightedGap.toLong.max(1)
     }