batch-level semantics, remove client validation, update test titles

Author: TheovanKraay

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

@@ -1037,26 +1037,28 @@ private object TransactionalBulkWriter {
 
   private val bulkProcessingThresholds = new CosmosBulkExecutionThresholdsState()
 
-  private val maxPendingOperationsPerJVM: Int = DefaultMaxPendingOperationPerCore * SparkUtils.getNumberOfHostCPUCores
+  // For batch-level backpressure: use a conservative estimate of concurrent batches across the JVM
+  // Each batch can have up to 100 operations, so this sizing provides headroom for the scheduler queues
+  private val maxPendingBatchesPerJVM: Int = DefaultMaxPendingOperationPerCore * SparkUtils.getNumberOfHostCPUCores / 50
 
   // Custom bounded elastic scheduler to consume input flux
   val bulkWriterRequestsBoundedElastic: Scheduler = Schedulers.newBoundedElastic(
     Schedulers.DEFAULT_BOUNDED_ELASTIC_SIZE,
-    Schedulers.DEFAULT_BOUNDED_ELASTIC_QUEUESIZE + 2 * maxPendingOperationsPerJVM,
+    Schedulers.DEFAULT_BOUNDED_ELASTIC_QUEUESIZE + 2 * maxPendingBatchesPerJVM,
     BULK_WRITER_REQUESTS_BOUNDED_ELASTIC_THREAD_NAME,
     TTL_FOR_SCHEDULER_WORKER_IN_SECONDS, true)
 
   // Custom bounded elastic scheduler to consume input flux
   val bulkWriterInputBoundedElastic: Scheduler = Schedulers.newBoundedElastic(
     Schedulers.DEFAULT_BOUNDED_ELASTIC_SIZE,
-    Schedulers.DEFAULT_BOUNDED_ELASTIC_QUEUESIZE + 2 * maxPendingOperationsPerJVM,
+    Schedulers.DEFAULT_BOUNDED_ELASTIC_QUEUESIZE + 2 * maxPendingBatchesPerJVM,
     BULK_WRITER_INPUT_BOUNDED_ELASTIC_THREAD_NAME,
     TTL_FOR_SCHEDULER_WORKER_IN_SECONDS, true)
 
   // Custom bounded elastic scheduler to switch off IO thread to process response.
   val bulkWriterResponsesBoundedElastic: Scheduler = Schedulers.newBoundedElastic(
     Schedulers.DEFAULT_BOUNDED_ELASTIC_SIZE,
-    Schedulers.DEFAULT_BOUNDED_ELASTIC_QUEUESIZE + maxPendingOperationsPerJVM,
+    Schedulers.DEFAULT_BOUNDED_ELASTIC_QUEUESIZE + maxPendingBatchesPerJVM,
     BULK_WRITER_RESPONSES_BOUNDED_ELASTIC_THREAD_NAME,
     TTL_FOR_SCHEDULER_WORKER_IN_SECONDS, true)
 

sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/TransactionalBatchITest.scala

@@ -424,11 +424,11 @@ class TransactionalBatchITest extends IntegrationSpec
     rootCause.getMessage should include("Batch request has more operations than what is supported")
   }
 
-  "Transactional Batch with Hierarchical Partition Keys" should "create items atomically with PermId and SourceId" in {
+  "Transactional Batch with Hierarchical Partition Keys" should "create items atomically with testPrimaryKey and testSecondaryKey" in {
     val cosmosEndpoint = TestConfigurations.HOST
     val cosmosMasterKey = TestConfigurations.MASTER_KEY
 
-    // Create container with hierarchical partition keys (PermId, SourceId)
+    // Create container with hierarchical partition keys (testPrimaryKey, testSecondaryKey)
     val containerName = s"test-hpk-${UUID.randomUUID()}"
     val containerProperties = new com.azure.cosmos.models.CosmosContainerProperties(
       containerName,
@@ -494,7 +494,7 @@ class TransactionalBatchITest extends IntegrationSpec
     val cosmosEndpoint = TestConfigurations.HOST
     val cosmosMasterKey = TestConfigurations.MASTER_KEY
 
-    // Create container with hierarchical partition keys (PermId, SourceId)
+    // Create container with hierarchical partition keys (testPrimaryKey, testSecondaryKey)
     val containerName = s"test-hpk-temporal-${UUID.randomUUID()}"
     val containerProperties = new com.azure.cosmos.models.CosmosContainerProperties(
       containerName,
@@ -570,11 +570,11 @@ class TransactionalBatchITest extends IntegrationSpec
     }
   }
 
-  it should "handle operations across multiple PermId/SourceId combinations" in {
+  it should "handle operations across multiple testPrimaryKey/testSecondaryKey combinations" in {
     val cosmosEndpoint = TestConfigurations.HOST
     val cosmosMasterKey = TestConfigurations.MASTER_KEY
 
-    // Create container with hierarchical partition keys (PermId, SourceId)
+    // Create container with hierarchical partition keys (testPrimaryKey, testSecondaryKey)
     val containerName = s"test-hpk-multi-${UUID.randomUUID()}"
     val containerProperties = new com.azure.cosmos.models.CosmosContainerProperties(
       containerName,
@@ -633,7 +633,7 @@ class TransactionalBatchITest extends IntegrationSpec
     val cosmosEndpoint = TestConfigurations.HOST
     val cosmosMasterKey = TestConfigurations.MASTER_KEY
 
-    // Create container with hierarchical partition keys (PermId, SourceId)
+    // Create container with hierarchical partition keys (testPrimaryKey, testSecondaryKey)
     val containerName = s"test-hpk-limit-${UUID.randomUUID()}"
     val containerProperties = new com.azure.cosmos.models.CosmosContainerProperties(
       containerName,

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/TransactionalBulkExecutor.java

@@ -95,7 +95,6 @@ public final class TransactionalBulkExecutor<TContext> implements Disposable {
         ImplementationBridgeHelpers.CosmosBatchResponseHelper.getCosmosBatchResponseAccessor();
 
     private final CosmosAsyncContainer container;
-    private final int maxMicroBatchPayloadSizeInBytes;
     private final AsyncDocumentClient docClientWrapper;
     private final String operationContextText;
     private final OperationContextAndListenerTuple operationListener;
@@ -126,7 +125,6 @@ public TransactionalBulkExecutor(CosmosAsyncContainer container,
         checkNotNull(inputOperations, "expected non-null inputOperations");
         checkNotNull(cosmosBulkOptions, "expected non-null bulkOptions");
 
-        this.maxMicroBatchPayloadSizeInBytes = cosmosBulkOptions.getMaxMicroBatchPayloadSizeInBytes();
         this.cosmosBulkExecutionOptions = cosmosBulkOptions;
         this.container = container;
         this.bulkSpanName = "transactionalBatch." + this.container.getId();
@@ -286,31 +284,6 @@ private Flux<CosmosBulkOperationResponse<TContext>> executeTransactionalBatch(Co
         String batchTrackingId = UUIDs.nonBlockingRandomUUID().toString();
         PartitionKey partitionKey = cosmosBatch.getPartitionKeyValue();
 
-        // Validate batch payload size before execution
-        int totalSerializedLength = this.calculateTotalSerializedLength(operations);
-        if (totalSerializedLength > this.maxMicroBatchPayloadSizeInBytes) {
-            String errorMessage = String.format(
-                "Transactional batch exceeds maximum payload size: %d bytes (limit: %d bytes), PK: %s",
-                totalSerializedLength,
-                this.maxMicroBatchPayloadSizeInBytes,
-                partitionKey);
-            logger.error("{}, Context: {}", errorMessage, this.operationContextText);
-            
-            // Return error responses for all operations in the batch
-            CosmosException payloadSizeException = BridgeInternal.createCosmosException(
-                HttpConstants.StatusCodes.REQUEST_ENTITY_TOO_LARGE,
-                errorMessage);
-            
-            return Flux.fromIterable(operations)
-                .map(operation -> {
-                    TContext actualContext = this.getActualContext(operation);
-                    return ModelBridgeInternal.createCosmosBulkOperationResponse(
-                        operation,
-                        payloadSizeException,
-                        actualContext);
-                });
-        }
-        
         logDebugOrWarning(
             "Executing transactional batch - {} operations, PK: {}, TrackingId: {}, Context: {}",
             operations.size(),
@@ -427,16 +400,6 @@ private Flux<CosmosBulkOperationResponse<TContext>> executeTransactionalBatch(Co
             });
     }
 
-    private int calculateTotalSerializedLength(List<CosmosItemOperation> operations) {
-        int totalLength = 0;
-        for (CosmosItemOperation operation : operations) {
-            if (operation instanceof CosmosItemOperationBase) {
-                totalLength += ((CosmosItemOperationBase) operation).getSerializedLength(this.effectiveItemSerializer);
-            }
-        }
-        return totalLength;
-    }
-
     private TContext getActualContext(CosmosItemOperation itemOperation) {
         ItemBulkOperation<?, ?> itemBulkOperation = null;