Cosmos Spark transactional batch support

Author: TheovanKraay

sdk/cosmos/azure-cosmos-spark_3-5/src/main/scala/com/azure/cosmos/spark/ItemsWriterBuilder.scala

@@ -2,11 +2,14 @@
 // Licensed under the MIT License.
 package com.azure.cosmos.spark
 
+import com.azure.cosmos.{CosmosAsyncClient, ReadConsistencyStrategy, SparkBridgeInternal}
 import com.azure.cosmos.spark.diagnostics.LoggerHelper
 import org.apache.spark.broadcast.Broadcast
+import org.apache.spark.sql.connector.distributions.{Distribution, Distributions}
+import org.apache.spark.sql.connector.expressions.{Expression, Expressions, NullOrdering, SortDirection, SortOrder}
 import org.apache.spark.sql.connector.metric.CustomMetric
 import org.apache.spark.sql.connector.write.streaming.StreamingWrite
-import org.apache.spark.sql.connector.write.{BatchWrite, Write, WriteBuilder}
+import org.apache.spark.sql.connector.write.{BatchWrite, RequiresDistributionAndOrdering, Write, WriteBuilder}
 import org.apache.spark.sql.types.StructType
 import org.apache.spark.sql.util.CaseInsensitiveStringMap
 
@@ -46,7 +49,7 @@ private class ItemsWriterBuilder
       diagnosticsConfig,
       sparkEnvironmentInfo)
 
-  private class CosmosWrite extends Write {
+  private class CosmosWrite extends Write with RequiresDistributionAndOrdering {
 
     private[this] val supportedCosmosMetrics: Array[CustomMetric] = {
       Array(
@@ -56,6 +59,15 @@ private class ItemsWriterBuilder
       )
     }
 
+    private[this] val writeConfig = CosmosWriteConfig.parseWriteConfig(
+      userConfig.asCaseSensitiveMap().asScala.toMap,
+      inputSchema
+    )
+
+    private[this] val containerConfig = CosmosContainerConfig.parseCosmosContainerConfig(
+      userConfig.asCaseSensitiveMap().asScala.toMap
+    )
+
     override def toBatch(): BatchWrite =
       new ItemsBatchWriter(
         userConfig.asCaseSensitiveMap().asScala.toMap,
@@ -73,5 +85,84 @@ private class ItemsWriterBuilder
         sparkEnvironmentInfo)
 
     override def supportedCustomMetrics(): Array[CustomMetric] = supportedCosmosMetrics
+
+    override def requiredDistribution(): Distribution = {
+      if (writeConfig.bulkEnabled && writeConfig.bulkEnableTransactions) {
+        // For transactional writes, partition by all partition key columns
+        val partitionKeyPaths = getPartitionKeyColumnNames()
+        if (partitionKeyPaths.nonEmpty) {
+          // Use public Expressions.column() factory - returns NamedReference
+          val clustering = partitionKeyPaths.map(path => Expressions.column(path): Expression).toArray
+          Distributions.clustered(clustering)
+        } else {
+          Distributions.unspecified()
+        }
+      } else {
+        Distributions.unspecified()
+      }
+    }
+
+    override def requiredOrdering(): Array[SortOrder] = {
+      if (writeConfig.bulkEnabled && writeConfig.bulkEnableTransactions) {
+        // For transactional writes, order by all partition key columns (ascending)
+        val partitionKeyPaths = getPartitionKeyColumnNames()
+        if (partitionKeyPaths.nonEmpty) {
+          partitionKeyPaths.map { path =>
+            // Use public Expressions.sort() factory for creating SortOrder
+            Expressions.sort(
+              Expressions.column(path),
+              SortDirection.ASCENDING,
+              NullOrdering.NULLS_FIRST
+            )
+          }.toArray
+        } else {
+          Array.empty[SortOrder]
+        }
+      } else {
+        Array.empty[SortOrder]
+      }
+    }
+
+    private def getPartitionKeyColumnNames(): Seq[String] = {
+      try {
+        // Need to create a temporary container client to get partition key definition
+        val clientCacheItem = CosmosClientCache(
+          CosmosClientConfiguration(
+            userConfig.asCaseSensitiveMap().asScala.toMap,
+            ReadConsistencyStrategy.EVENTUAL,
+            sparkEnvironmentInfo
+          ),
+          Some(cosmosClientStateHandles.value.cosmosClientMetadataCaches),
+          "ItemsWriterBuilder-PKLookup"
+        )
+
+        val container = ThroughputControlHelper.getContainer(
+          userConfig.asCaseSensitiveMap().asScala.toMap,
+          containerConfig,
+          clientCacheItem,
+          None
+        )
+
+        val containerProperties = SparkBridgeInternal.getContainerPropertiesFromCollectionCache(container)
+        val partitionKeyDefinition = containerProperties.getPartitionKeyDefinition
+
+        // Release the client
+        clientCacheItem.close()
+
+        if (partitionKeyDefinition != null && partitionKeyDefinition.getPaths != null) {
+          val paths = partitionKeyDefinition.getPaths.asScala
+          paths.map(path => {
+            // Remove leading '/' from partition key path (e.g., "/pk" -> "pk")
+            if (path.startsWith("/")) path.substring(1) else path
+          }).toSeq
+        } else {
+          Seq.empty[String]
+        }
+      } catch {
+        case ex: Exception =>
+          log.logWarning(s"Failed to get partition key definition for transactional writes: ${ex.getMessage}")
+          Seq.empty[String]
+      }
+    }
   }
 }

sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosConfig.scala

@@ -113,6 +113,7 @@ private[spark] object CosmosConfigNames {
   val ClientTelemetryEnabled = "spark.cosmos.clientTelemetry.enabled" // keep this to avoid breaking changes
   val ClientTelemetryEndpoint = "spark.cosmos.clientTelemetry.endpoint" // keep this to avoid breaking changes
   val WriteBulkEnabled = "spark.cosmos.write.bulk.enabled"
+  val WriteBulkEnableTransactions = "spark.cosmos.write.bulk.enableTransactions"
   val WriteBulkMaxPendingOperations = "spark.cosmos.write.bulk.maxPendingOperations"
   val WriteBulkMaxBatchSize = "spark.cosmos.write.bulk.maxBatchSize"
   val WriteBulkMinTargetBatchSize = "spark.cosmos.write.bulk.minTargetBatchSize"
@@ -242,6 +243,7 @@ private[spark] object CosmosConfigNames {
     ClientTelemetryEnabled,
     ClientTelemetryEndpoint,
     WriteBulkEnabled,
+    WriteBulkEnableTransactions,
     WriteBulkMaxPendingOperations,
     WriteBulkMaxConcurrentPartitions,
     WriteBulkPayloadSizeInBytes,
@@ -1462,6 +1464,7 @@ private case class CosmosPatchConfigs(columnConfigsMap: TrieMap[String, CosmosPa
 private case class CosmosWriteConfig(itemWriteStrategy: ItemWriteStrategy,
                                      maxRetryCount: Int,
                                      bulkEnabled: Boolean,
+                                     bulkEnableTransactions: Boolean = false,
                                      bulkMaxPendingOperations: Option[Int] = None,
                                      pointMaxConcurrency: Option[Int] = None,
                                      maxConcurrentCosmosPartitions: Option[Int] = None,
@@ -1486,6 +1489,12 @@ private object CosmosWriteConfig {
     parseFromStringFunction = bulkEnabledAsString => bulkEnabledAsString.toBoolean,
     helpMessage = "Cosmos DB Item Write bulk enabled")
 
+  private val bulkEnableTransactions = CosmosConfigEntry[Boolean](key = CosmosConfigNames.WriteBulkEnableTransactions,
+    defaultValue = Option.apply(false),
+    mandatory = false,
+    parseFromStringFunction = enableTransactionsAsString => enableTransactionsAsString.toBoolean,
+    helpMessage = "Cosmos DB Item Write enable transactional batch - requires bulk write to be enabled and Spark 3.5+")
+
   private val microBatchPayloadSizeInBytes = CosmosConfigEntry[Int](key = CosmosConfigNames.WriteBulkPayloadSizeInBytes,
     defaultValue = Option.apply(BatchRequestResponseConstants.DEFAULT_MAX_DIRECT_MODE_BATCH_REQUEST_BODY_SIZE_IN_BYTES),
     mandatory = false,
@@ -1758,6 +1767,7 @@ private object CosmosWriteConfig {
     val itemWriteStrategyOpt = CosmosConfigEntry.parse(cfg, itemWriteStrategy)
     val maxRetryCountOpt = CosmosConfigEntry.parse(cfg, maxRetryCount)
     val bulkEnabledOpt = CosmosConfigEntry.parse(cfg, bulkEnabled)
+    val bulkEnableTransactionsOpt = CosmosConfigEntry.parse(cfg, bulkEnableTransactions)
     var patchConfigsOpt = Option.empty[CosmosPatchConfigs]
     val throughputControlConfigOpt = CosmosThroughputControlConfig.parseThroughputControlConfig(cfg)
     val microBatchPayloadSizeInBytesOpt = CosmosConfigEntry.parse(cfg, microBatchPayloadSizeInBytes)
@@ -1788,6 +1798,7 @@ private object CosmosWriteConfig {
       itemWriteStrategyOpt.get,
       maxRetryCountOpt.get,
       bulkEnabled = bulkEnabledOpt.get,
+      bulkEnableTransactions = bulkEnableTransactionsOpt.getOrElse(false),
       bulkMaxPendingOperations = CosmosConfigEntry.parse(cfg, bulkMaxPendingOperations),
       pointMaxConcurrency = CosmosConfigEntry.parse(cfg, pointWriteConcurrency),
       maxConcurrentCosmosPartitions = CosmosConfigEntry.parse(cfg, bulkMaxConcurrentPartitions),

sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosWriterBase.scala

@@ -62,14 +62,25 @@ private abstract class CosmosWriterBase(
 
   private val writer: AtomicReference[AsyncItemWriter] = new AtomicReference(
     if (cosmosWriteConfig.bulkEnabled) {
-      new BulkWriter(
-        container,
-        cosmosTargetContainerConfig,
-        partitionKeyDefinition,
-        cosmosWriteConfig,
-        diagnosticsConfig,
-        getOutputMetricsPublisher(),
-        commitAttempt.getAndIncrement())
+      if (cosmosWriteConfig.bulkEnableTransactions) {
+        new TransactionalBulkWriter(
+          container,
+          cosmosTargetContainerConfig,
+          partitionKeyDefinition,
+          cosmosWriteConfig,
+          diagnosticsConfig,
+          getOutputMetricsPublisher(),
+          commitAttempt.getAndIncrement())
+      } else {
+        new BulkWriter(
+          container,
+          cosmosTargetContainerConfig,
+          partitionKeyDefinition,
+          cosmosWriteConfig,
+          diagnosticsConfig,
+          getOutputMetricsPublisher(),
+          commitAttempt.getAndIncrement())
+      }
     } else {
       new PointWriter(
         container,