From 9770833eb593e21165051daa0cbb2ed07fb57470 Mon Sep 17 00:00:00 2001 From: Fabian Meiswinkel Date: Mon, 13 Apr 2026 22:49:34 +0000 Subject: [PATCH 01/64] Adding readManyByPartitionKey API --- .../azure/cosmos/spark/CosmosConstants.scala | 1 + .../cosmos/spark/CosmosItemsDataSource.scala | 109 ++++- .../spark/CosmosPartitionKeyHelper.scala | 46 +++ .../CosmosReadManyByPartitionKeyReader.scala | 153 +++++++ ...tionReaderWithReadManyByPartitionKey.scala | 233 +++++++++++ .../udf/GetCosmosPartitionKeyValue.scala | 25 ++ .../spark/CosmosPartitionKeyHelperSpec.scala | 81 ++++ ...eaderWithReadManyByPartitionKeyITest.scala | 158 ++++++++ .../cosmos/ReadManyByPartitionKeyTest.java | 381 ++++++++++++++++++ ...ReadManyByPartitionKeyQueryHelperTest.java | 349 ++++++++++++++++ .../azure/cosmos/CosmosAsyncContainer.java | 111 +++++ .../com/azure/cosmos/CosmosContainer.java | 67 +++ .../implementation/AsyncDocumentClient.java | 21 + .../azure/cosmos/implementation/Configs.java | 19 + .../ReadManyByPartitionKeyQueryHelper.java | 199 +++++++++ .../implementation/RxDocumentClientImpl.java | 227 +++++++++++ .../DocumentQueryExecutionContextFactory.java | 11 + .../docs/readManyByPartitionKey-design.md | 133 ++++++ 18 files changed, 2323 insertions(+), 1 deletion(-) create mode 100644 sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosPartitionKeyHelper.scala create mode 100644 sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosReadManyByPartitionKeyReader.scala create mode 100644 sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/ItemsPartitionReaderWithReadManyByPartitionKey.scala create mode 100644 sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/udf/GetCosmosPartitionKeyValue.scala create mode 100644 sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/CosmosPartitionKeyHelperSpec.scala create mode 100644 sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/ItemsPartitionReaderWithReadManyByPartitionKeyITest.scala create mode 100644 sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/ReadManyByPartitionKeyTest.java create mode 100644 sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelperTest.java create mode 100644 sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelper.java create mode 100644 sdk/cosmos/docs/readManyByPartitionKey-design.md diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosConstants.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosConstants.scala index 9ece47416526..00761f23d399 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosConstants.scala +++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosConstants.scala @@ -45,6 +45,7 @@ private[cosmos] object CosmosConstants { val Id = "id" val ETag = "_etag" val ItemIdentity = "_itemIdentity" + val PartitionKeyIdentity = "_partitionKeyIdentity" } object StatusCodes { diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosItemsDataSource.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosItemsDataSource.scala index a35cff27af68..2fbc036724e7 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosItemsDataSource.scala +++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosItemsDataSource.scala @@ -2,9 +2,10 @@ // Licensed under the MIT License. package com.azure.cosmos.spark -import com.azure.cosmos.models.{CosmosItemIdentity, PartitionKey} +import com.azure.cosmos.models.{CosmosItemIdentity, PartitionKey, PartitionKeyBuilder} import com.azure.cosmos.spark.CosmosPredicates.assertOnSparkDriver import com.azure.cosmos.spark.diagnostics.BasicLoggingTrait +import com.azure.cosmos.{SparkBridgeInternal} import org.apache.spark.sql.{DataFrame, Row, SparkSession} import java.util @@ -112,4 +113,110 @@ object CosmosItemsDataSource { readManyReader.readMany(df.rdd, readManyFilterExtraction) } + + def readManyByPartitionKey(df: DataFrame, userConfig: java.util.Map[String, String]): DataFrame = { + readManyByPartitionKey(df, userConfig, null) + } + + def readManyByPartitionKey( + df: DataFrame, + userConfig: java.util.Map[String, String], + userProvidedSchema: StructType): DataFrame = { + + val readManyReader = new CosmosReadManyByPartitionKeyReader( + userProvidedSchema, + userConfig.asScala.toMap) + + // Option 1: Look for the _partitionKeyIdentity column (produced by GetCosmosPartitionKeyValue UDF) + val pkIdentityFieldExtraction = df + .schema + .find(field => field.name.equals(CosmosConstants.Properties.PartitionKeyIdentity) && field.dataType.equals(StringType)) + .map(field => (row: Row) => + CosmosPartitionKeyHelper.tryParsePartitionKey(row.getString(row.fieldIndex(field.name))).get) + + // Option 2: Detect PK columns by matching the container's partition key paths against the DataFrame schema + val pkColumnExtraction: Option[Row => PartitionKey] = if (pkIdentityFieldExtraction.isDefined) { + None // no need to resolve PK paths - _partitionKeyIdentity column takes precedence + } else { + val effectiveConfig = CosmosConfig.getEffectiveConfig( + databaseName = None, + containerName = None, + userConfig.asScala.toMap) + val readConfig = CosmosReadConfig.parseCosmosReadConfig(effectiveConfig) + val containerConfig = CosmosContainerConfig.parseCosmosContainerConfig(effectiveConfig) + val sparkEnvironmentInfo = CosmosClientConfiguration.getSparkEnvironmentInfo(None) + val calledFrom = s"CosmosItemsDataSource.readManyByPartitionKey" + + val pkPathsOpt = Loan( + List[Option[CosmosClientCacheItem]]( + Some( + CosmosClientCache( + CosmosClientConfiguration( + effectiveConfig, + readConsistencyStrategy = readConfig.readConsistencyStrategy, + sparkEnvironmentInfo), + None, + calledFrom)), + ThroughputControlHelper.getThroughputControlClientCacheItem( + effectiveConfig, + calledFrom, + None, + sparkEnvironmentInfo) + )) + .to(clientCacheItems => { + val container = + ThroughputControlHelper.getContainer( + effectiveConfig, + containerConfig, + clientCacheItems(0).get, + clientCacheItems(1)) + + val pkDefinition = SparkBridgeInternal + .getContainerPropertiesFromCollectionCache(container) + .getPartitionKeyDefinition + + pkDefinition.getPaths.asScala.map(_.stripPrefix("/")).toList + }) + + // Check if ALL PK path columns exist in the DataFrame schema + val dfFieldNames = df.schema.fieldNames.toSet + val allPkColumnsPresent = pkPathsOpt.forall(path => dfFieldNames.contains(path)) + + if (allPkColumnsPresent && pkPathsOpt.nonEmpty) { + val pkPaths = pkPathsOpt + Some((row: Row) => { + if (pkPaths.size == 1) { + // Single partition key + new PartitionKey(row.getAs[Any](pkPaths.head)) + } else { + // Hierarchical partition key — build level by level + val builder = new PartitionKeyBuilder() + for (path <- pkPaths) { + val value = row.getAs[Any](path) + value match { + case s: String => builder.add(s) + case n: Number => builder.add(n.doubleValue()) + case b: Boolean => builder.add(b) + case null => builder.addNoneValue() + case other => builder.add(other.toString) + } + } + builder.build() + } + }) + } else { + None + } + } + + val pkExtraction = pkIdentityFieldExtraction + .orElse(pkColumnExtraction) + .getOrElse( + throw new IllegalArgumentException( + "Cannot determine partition key extraction from the input DataFrame. " + + "Either add a '_partitionKeyIdentity' column (using the GetCosmosPartitionKeyValue UDF) " + + "or ensure the DataFrame contains columns matching the container's partition key paths.")) + + readManyReader.readManyByPartitionKey(df.rdd, pkExtraction) + } } diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosPartitionKeyHelper.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosPartitionKeyHelper.scala new file mode 100644 index 000000000000..616e1893b343 --- /dev/null +++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosPartitionKeyHelper.scala @@ -0,0 +1,46 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.azure.cosmos.spark + +import com.azure.cosmos.implementation.routing.PartitionKeyInternal +import com.azure.cosmos.implementation.{ImplementationBridgeHelpers, Utils} +import com.azure.cosmos.models.PartitionKey +import com.azure.cosmos.spark.diagnostics.BasicLoggingTrait + +import java.util + +// scalastyle:off underscore.import +import scala.collection.JavaConverters._ +// scalastyle:on underscore.import + +private[spark] object CosmosPartitionKeyHelper extends BasicLoggingTrait { + // pattern will be recognized + // pk(partitionKeyValue) + // + // (?i) : The whole matching is case-insensitive + // pk[(](.*)[)]: partitionKey Value + private val cosmosPartitionKeyStringRegx = """(?i)pk[(](.*)[)]""".r + private val objectMapper = Utils.getSimpleObjectMapper + + def getCosmosPartitionKeyValueString(partitionKeyValue: List[Object]): String = { + s"pk(${objectMapper.writeValueAsString(partitionKeyValue.asJava)})" + } + + def tryParsePartitionKey(cosmosPartitionKeyString: String): Option[PartitionKey] = { + cosmosPartitionKeyString match { + case cosmosPartitionKeyStringRegx(pkValue) => + val partitionKeyValue = Utils.parse(pkValue, classOf[Object]) + partitionKeyValue match { + case arrayList: util.ArrayList[Object] => + Some( + ImplementationBridgeHelpers + .PartitionKeyHelper + .getPartitionKeyAccessor + .toPartitionKey(PartitionKeyInternal.fromObjectArray(arrayList.toArray, false))) + case _ => Some(new PartitionKey(partitionKeyValue)) + } + case _ => None + } + } +} diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosReadManyByPartitionKeyReader.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosReadManyByPartitionKeyReader.scala new file mode 100644 index 000000000000..1d324d4855ff --- /dev/null +++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosReadManyByPartitionKeyReader.scala @@ -0,0 +1,153 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. +package com.azure.cosmos.spark + +import com.azure.cosmos.{CosmosException, ReadConsistencyStrategy} +import com.azure.cosmos.implementation.{CosmosClientMetadataCachesSnapshot, UUIDs} +import com.azure.cosmos.models.PartitionKey +import com.azure.cosmos.spark.CosmosPredicates.assertOnSparkDriver +import com.azure.cosmos.spark.diagnostics.{BasicLoggingTrait, DiagnosticsContext} +import com.fasterxml.jackson.databind.node.ObjectNode +import org.apache.spark.TaskContext +import org.apache.spark.broadcast.Broadcast +import org.apache.spark.rdd.RDD +import org.apache.spark.sql.{DataFrame, Row, SparkSession} +import org.apache.spark.sql.types.StructType + +import java.util.UUID + +private[spark] class CosmosReadManyByPartitionKeyReader( + val userProvidedSchema: StructType, + val userConfig: Map[String, String] + ) extends BasicLoggingTrait with Serializable { + val effectiveUserConfig: Map[String, String] = CosmosConfig.getEffectiveConfig( + databaseName = None, + containerName = None, + userConfig) + + val clientConfig: CosmosAccountConfig = CosmosAccountConfig.parseCosmosAccountConfig(effectiveUserConfig) + val readConfig: CosmosReadConfig = CosmosReadConfig.parseCosmosReadConfig(effectiveUserConfig) + val cosmosContainerConfig: CosmosContainerConfig = + CosmosContainerConfig.parseCosmosContainerConfig(effectiveUserConfig) + //scalastyle:off multiple.string.literals + val tableName: String = s"com.azure.cosmos.spark.items.${clientConfig.accountName}." + + s"${cosmosContainerConfig.database}.${cosmosContainerConfig.container}" + private lazy val sparkSession = { + assertOnSparkDriver() + SparkSession.active + } + val sparkEnvironmentInfo: String = CosmosClientConfiguration.getSparkEnvironmentInfo(Some(sparkSession)) + logTrace(s"Instantiated ${this.getClass.getSimpleName} for $tableName") + + private[spark] def initializeAndBroadcastCosmosClientStatesForContainer(): Broadcast[CosmosClientMetadataCachesSnapshots] = { + val calledFrom = s"CosmosReadManyByPartitionKeyReader($tableName).initializeAndBroadcastCosmosClientStateForContainer" + Loan( + List[Option[CosmosClientCacheItem]]( + Some( + CosmosClientCache( + CosmosClientConfiguration( + effectiveUserConfig, + readConsistencyStrategy = readConfig.readConsistencyStrategy, + sparkEnvironmentInfo), + None, + calledFrom)), + ThroughputControlHelper.getThroughputControlClientCacheItem( + effectiveUserConfig, + calledFrom, + None, + sparkEnvironmentInfo) + )) + .to(clientCacheItems => { + val container = + ThroughputControlHelper.getContainer( + effectiveUserConfig, + cosmosContainerConfig, + clientCacheItems(0).get, + clientCacheItems(1)) + try { + container.readItem( + UUIDs.nonBlockingRandomUUID().toString, + new PartitionKey(UUIDs.nonBlockingRandomUUID().toString), + classOf[ObjectNode]) + .block() + } catch { + case _: CosmosException => None + } + + val state = new CosmosClientMetadataCachesSnapshot() + state.serialize(clientCacheItems(0).get.cosmosClient) + + var throughputControlState: Option[CosmosClientMetadataCachesSnapshot] = None + if (clientCacheItems(1).isDefined) { + throughputControlState = Some(new CosmosClientMetadataCachesSnapshot()) + throughputControlState.get.serialize(clientCacheItems(1).get.cosmosClient) + } + + val metadataSnapshots = CosmosClientMetadataCachesSnapshots(state, throughputControlState) + sparkSession.sparkContext.broadcast(metadataSnapshots) + }) + } + + def readManyByPartitionKey(inputRdd: RDD[Row], pkExtraction: Row => PartitionKey): DataFrame = { + val correlationActivityId = UUIDs.nonBlockingRandomUUID() + val calledFrom = s"CosmosReadManyByPartitionKeyReader.readManyByPartitionKey($correlationActivityId)" + val schema = Loan( + List[Option[CosmosClientCacheItem]]( + Some(CosmosClientCache( + CosmosClientConfiguration( + effectiveUserConfig, + readConsistencyStrategy = readConfig.readConsistencyStrategy, + sparkEnvironmentInfo), + None, + calledFrom + )), + ThroughputControlHelper.getThroughputControlClientCacheItem( + effectiveUserConfig, + calledFrom, + None, + sparkEnvironmentInfo) + )) + .to(clientCacheItems => Option.apply(userProvidedSchema).getOrElse( + CosmosTableSchemaInferrer.inferSchema( + clientCacheItems(0).get, + clientCacheItems(1), + effectiveUserConfig, + ItemsTable.defaultSchemaForInferenceDisabled))) + + val clientStates = initializeAndBroadcastCosmosClientStatesForContainer + + sparkSession.sqlContext.createDataFrame( + inputRdd.mapPartitionsWithIndex( + (partitionIndex: Int, rowIterator: Iterator[Row]) => { + val pkIterator: Iterator[PartitionKey] = rowIterator + .map(row => pkExtraction.apply(row)) + + logInfo(s"Creating an ItemsPartitionReaderWithReadManyByPartitionKey for Activity $correlationActivityId to read for " + + s"input partition [$partitionIndex] ${tableName}") + + val reader = new ItemsPartitionReaderWithReadManyByPartitionKey( + effectiveUserConfig, + CosmosReadManyHelper.FullRangeFeedRange, + schema, + DiagnosticsContext(correlationActivityId, partitionIndex.toString), + clientStates, + DiagnosticsConfig.parseDiagnosticsConfig(effectiveUserConfig), + sparkEnvironmentInfo, + TaskContext.get, + pkIterator) + + new Iterator[Row] { + override def hasNext: Boolean = reader.next() + + override def next(): Row = reader.getCurrentRow() + } + }, + preservesPartitioning = true + ), + schema) + } +} + +private object CosmosReadManyByPartitionKeyHelper { + val FullRangeFeedRange: NormalizedRange = NormalizedRange("", "FF") +} diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/ItemsPartitionReaderWithReadManyByPartitionKey.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/ItemsPartitionReaderWithReadManyByPartitionKey.scala new file mode 100644 index 000000000000..68e41cad3ec5 --- /dev/null +++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/ItemsPartitionReaderWithReadManyByPartitionKey.scala @@ -0,0 +1,233 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.azure.cosmos.spark + +import com.azure.cosmos.{CosmosAsyncContainer, CosmosEndToEndOperationLatencyPolicyConfigBuilder, CosmosItemSerializerNoExceptionWrapping, SparkBridgeInternal} +import com.azure.cosmos.implementation.spark.OperationContextAndListenerTuple +import com.azure.cosmos.implementation.{ImplementationBridgeHelpers, ObjectNodeMap, SparkRowItem, Utils} +import com.azure.cosmos.models.{CosmosReadManyRequestOptions, ModelBridgeInternal, PartitionKey, PartitionKeyDefinition} +import com.azure.cosmos.spark.BulkWriter.getThreadInfo +import com.azure.cosmos.spark.CosmosTableSchemaInferrer.IdAttributeName +import com.azure.cosmos.spark.diagnostics.{DetailedFeedDiagnosticsProvider, DiagnosticsContext, DiagnosticsLoader, LoggerHelper, SparkTaskContext} +import com.fasterxml.jackson.databind.node.ObjectNode +import org.apache.spark.TaskContext +import org.apache.spark.broadcast.Broadcast +import org.apache.spark.sql.Row +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.catalyst.encoders.ExpressionEncoder +import org.apache.spark.sql.connector.read.PartitionReader +import org.apache.spark.sql.types.StructType + +import java.util + +// scalastyle:off underscore.import +import scala.collection.JavaConverters._ +// scalastyle:on underscore.import + +private[spark] case class ItemsPartitionReaderWithReadManyByPartitionKey +( + config: Map[String, String], + feedRange: NormalizedRange, + readSchema: StructType, + diagnosticsContext: DiagnosticsContext, + cosmosClientStateHandles: Broadcast[CosmosClientMetadataCachesSnapshots], + diagnosticsConfig: DiagnosticsConfig, + sparkEnvironmentInfo: String, + taskContext: TaskContext, + readManyPartitionKeys: Iterator[PartitionKey] +) + extends PartitionReader[InternalRow] { + + private lazy val log = LoggerHelper.getLogger(diagnosticsConfig, this.getClass) + + private val readManyOptions = new CosmosReadManyRequestOptions() + private val readManyOptionsImpl = ImplementationBridgeHelpers + .CosmosReadManyRequestOptionsHelper + .getCosmosReadManyRequestOptionsAccessor + .getImpl(readManyOptions) + + private val readConfig = CosmosReadConfig.parseCosmosReadConfig(config) + ThroughputControlHelper.populateThroughputControlGroupName(readManyOptionsImpl, readConfig.throughputControlConfig) + + private val operationContext = { + assert(taskContext != null) + + SparkTaskContext(diagnosticsContext.correlationActivityId, + taskContext.stageId(), + taskContext.partitionId(), + taskContext.taskAttemptId(), + feedRange.toString) + } + + private val operationContextAndListenerTuple: Option[OperationContextAndListenerTuple] = { + if (diagnosticsConfig.mode.isDefined) { + val listener = + DiagnosticsLoader.getDiagnosticsProvider(diagnosticsConfig).getLogger(this.getClass) + + val ctxAndListener = new OperationContextAndListenerTuple(operationContext, listener) + + readManyOptionsImpl + .setOperationContextAndListenerTuple(ctxAndListener) + + Some(ctxAndListener) + } else { + None + } + } + + log.logTrace(s"Instantiated ${this.getClass.getSimpleName}, Context: ${operationContext.toString} $getThreadInfo") + + private val containerTargetConfig = CosmosContainerConfig.parseCosmosContainerConfig(config) + + log.logInfo(s"Using ReadManyByPartitionKey from feed range $feedRange of " + + s"container ${containerTargetConfig.database}.${containerTargetConfig.container} - " + + s"correlationActivityId ${diagnosticsContext.correlationActivityId}, " + + s"Context: ${operationContext.toString} $getThreadInfo") + + private val clientCacheItem = CosmosClientCache( + CosmosClientConfiguration(config, readConfig.readConsistencyStrategy, sparkEnvironmentInfo), + Some(cosmosClientStateHandles.value.cosmosClientMetadataCaches), + s"ItemsPartitionReaderWithReadManyByPartitionKey($feedRange, ${containerTargetConfig.database}.${containerTargetConfig.container})" + ) + + private val throughputControlClientCacheItemOpt = + ThroughputControlHelper.getThroughputControlClientCacheItem( + config, + clientCacheItem.context, + Some(cosmosClientStateHandles), + sparkEnvironmentInfo) + + private val cosmosAsyncContainer = + ThroughputControlHelper.getContainer( + config, + containerTargetConfig, + clientCacheItem, + throughputControlClientCacheItemOpt) + + private val partitionKeyDefinition: PartitionKeyDefinition = { + TransientErrorsRetryPolicy.executeWithRetry(() => { + SparkBridgeInternal + .getContainerPropertiesFromCollectionCache(cosmosAsyncContainer).getPartitionKeyDefinition + }) + } + + private val cosmosSerializationConfig = CosmosSerializationConfig.parseSerializationConfig(config) + private val cosmosRowConverter = CosmosRowConverter.get(cosmosSerializationConfig) + + readManyOptionsImpl + .setCustomItemSerializer( + new CosmosItemSerializerNoExceptionWrapping { + override def serialize[T](item: T): util.Map[String, AnyRef] = ??? + + override def deserialize[T](jsonNodeMap: util.Map[String, AnyRef], classType: Class[T]): T = { + if (jsonNodeMap == null) { + throw new IllegalStateException("The 'jsonNodeMap' should never be null here.") + } + + if (classType != classOf[SparkRowItem]) { + throw new IllegalStateException("The 'classType' must be 'classOf[SparkRowItem])' here.") + } + + val objectNode: ObjectNode = jsonNodeMap match { + case map: ObjectNodeMap => + map.getObjectNode + case _ => + Utils.getSimpleObjectMapper.convertValue(jsonNodeMap, classOf[ObjectNode]) + } + + val partitionKey = PartitionKeyHelper.getPartitionKeyPath(objectNode, partitionKeyDefinition) + + val row = cosmosRowConverter.fromObjectNodeToRow(readSchema, + objectNode, + readConfig.schemaConversionMode) + + SparkRowItem(row, getPartitionKeyForFeedDiagnostics(partitionKey)).asInstanceOf[T] + } + } + ) + + // Collect all PK values upfront — readManyByPartitionKey needs the full list to + // group by physical partition and issue parallel queries. + // Deduplicate by PK string representation — safe because the list size is bounded + // by the per-call limit of the readManyByPartitionKey API. + private lazy val pkList = { + val seen = new java.util.LinkedHashMap[String, PartitionKey]() + readManyPartitionKeys.foreach(pk => seen.putIfAbsent(pk.toString, pk)) + new java.util.ArrayList[PartitionKey](seen.values()) + } + + private val endToEndTimeoutPolicy = + new CosmosEndToEndOperationLatencyPolicyConfigBuilder( + java.time.Duration.ofSeconds(CosmosConstants.readOperationEndToEndTimeoutInSeconds)) + .enable(true) + .build + + // Single iterator over all PKs — the SDK handles per-physical-partition batching + // internally to avoid oversized SQL queries. + private lazy val iterator = new TransientIOErrorsRetryingIterator[SparkRowItem]( + continuationToken => { + val options = new CosmosReadManyRequestOptions() + val optionsImpl = ImplementationBridgeHelpers + .CosmosReadManyRequestOptionsHelper + .getCosmosReadManyRequestOptionsAccessor + .getImpl(options) + + ThroughputControlHelper.populateThroughputControlGroupName(optionsImpl, readConfig.throughputControlConfig) + + if (operationContextAndListenerTuple.isDefined) { + optionsImpl.setOperationContextAndListenerTuple(operationContextAndListenerTuple.get) + } + + optionsImpl.setCustomItemSerializer(readManyOptionsImpl.getCustomItemSerializer) + + if (pkList.isEmpty) { + cosmosAsyncContainer.readManyByPartitionKey( + new java.util.ArrayList[PartitionKey](), options, classOf[SparkRowItem]) + } else { + readConfig.customQuery match { + case Some(query) => + cosmosAsyncContainer.readManyByPartitionKey(pkList, query.toSqlQuerySpec, options, classOf[SparkRowItem]) + case None => + cosmosAsyncContainer.readManyByPartitionKey(pkList, options, classOf[SparkRowItem]) + } + } + }, + readConfig.maxItemCount, + readConfig.prefetchBufferSize, + operationContextAndListenerTuple, + None + ) + + private val rowSerializer: ExpressionEncoder.Serializer[Row] = RowSerializerPool.getOrCreateSerializer(readSchema) + + private def shouldLogDetailedFeedDiagnostics(): Boolean = { + diagnosticsConfig.mode.isDefined && + diagnosticsConfig.mode.get.equalsIgnoreCase(classOf[DetailedFeedDiagnosticsProvider].getName) + } + + private def getPartitionKeyForFeedDiagnostics(pkValue: PartitionKey): Option[PartitionKey] = { + if (shouldLogDetailedFeedDiagnostics()) { + Some(pkValue) + } else { + None + } + } + + override def next(): Boolean = iterator.hasNext + + override def get(): InternalRow = { + cosmosRowConverter.fromRowToInternalRow(iterator.next().row, rowSerializer) + } + + def getCurrentRow(): Row = iterator.next().row + + override def close(): Unit = { + this.iterator.close() + RowSerializerPool.returnSerializerToPool(readSchema, rowSerializer) + clientCacheItem.close() + if (throughputControlClientCacheItemOpt.isDefined) { + throughputControlClientCacheItemOpt.get.close() + } + } +} diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/udf/GetCosmosPartitionKeyValue.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/udf/GetCosmosPartitionKeyValue.scala new file mode 100644 index 000000000000..a58d5b723b8b --- /dev/null +++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/udf/GetCosmosPartitionKeyValue.scala @@ -0,0 +1,25 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.azure.cosmos.spark.udf + +import com.azure.cosmos.spark.CosmosPartitionKeyHelper +import com.azure.cosmos.spark.CosmosPredicates.requireNotNull +import org.apache.spark.sql.api.java.UDF1 + +@SerialVersionUID(1L) +class GetCosmosPartitionKeyValue extends UDF1[Object, String] { + override def call + ( + partitionKeyValue: Object + ): String = { + requireNotNull(partitionKeyValue, "partitionKeyValue") + + partitionKeyValue match { + // for subpartitions case - Seq covers both WrappedArray (Scala 2.12) and ArraySeq (Scala 2.13) + case seq: Seq[Any] => + CosmosPartitionKeyHelper.getCosmosPartitionKeyValueString(seq.map(_.asInstanceOf[Object]).toList) + case _ => CosmosPartitionKeyHelper.getCosmosPartitionKeyValueString(List(partitionKeyValue)) + } + } +} diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/CosmosPartitionKeyHelperSpec.scala b/sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/CosmosPartitionKeyHelperSpec.scala new file mode 100644 index 000000000000..d127710da287 --- /dev/null +++ b/sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/CosmosPartitionKeyHelperSpec.scala @@ -0,0 +1,81 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.azure.cosmos.spark + +import com.azure.cosmos.models.{PartitionKey, PartitionKeyBuilder} + +class CosmosPartitionKeyHelperSpec extends UnitSpec { + //scalastyle:off multiple.string.literals + //scalastyle:off magic.number + + it should "return the correct partition key value string for single PK" in { + val pkString = CosmosPartitionKeyHelper.getCosmosPartitionKeyValueString(List("pk1")) + pkString shouldEqual "pk([\"pk1\"])" + } + + it should "return the correct partition key value string for HPK" in { + val pkString = CosmosPartitionKeyHelper.getCosmosPartitionKeyValueString(List("city1", "zip1")) + pkString shouldEqual "pk([\"city1\",\"zip1\"])" + } + + it should "return the correct partition key value string for 3-level HPK" in { + val pkString = CosmosPartitionKeyHelper.getCosmosPartitionKeyValueString(List("a", "b", "c")) + pkString shouldEqual "pk([\"a\",\"b\",\"c\"])" + } + + it should "parse valid single PK string" in { + val pk = CosmosPartitionKeyHelper.tryParsePartitionKey("pk([\"myPkValue\"])") + pk.isDefined shouldBe true + pk.get shouldEqual new PartitionKey("myPkValue") + } + + it should "parse valid HPK string" in { + val pk = CosmosPartitionKeyHelper.tryParsePartitionKey("pk([\"city1\",\"zip1\"])") + pk.isDefined shouldBe true + val expected = new PartitionKeyBuilder().add("city1").add("zip1").build() + pk.get shouldEqual expected + } + + it should "parse valid 3-level HPK string" in { + val pk = CosmosPartitionKeyHelper.tryParsePartitionKey("pk([\"a\",\"b\",\"c\"])") + pk.isDefined shouldBe true + val expected = new PartitionKeyBuilder().add("a").add("b").add("c").build() + pk.get shouldEqual expected + } + + it should "roundtrip single PK" in { + val original = "pk([\"roundtrip\"])" + val parsed = CosmosPartitionKeyHelper.tryParsePartitionKey(original) + parsed.isDefined shouldBe true + val serialized = CosmosPartitionKeyHelper.getCosmosPartitionKeyValueString(List("roundtrip")) + serialized shouldEqual original + } + + it should "roundtrip HPK" in { + val original = "pk([\"city\",\"zip\"])" + val parsed = CosmosPartitionKeyHelper.tryParsePartitionKey(original) + parsed.isDefined shouldBe true + val serialized = CosmosPartitionKeyHelper.getCosmosPartitionKeyValueString(List("city", "zip")) + serialized shouldEqual original + } + + it should "return None for malformed string" in { + val pk = CosmosPartitionKeyHelper.tryParsePartitionKey("invalid_format") + pk.isDefined shouldBe false + } + + it should "return None for missing pk prefix" in { + val pk = CosmosPartitionKeyHelper.tryParsePartitionKey("[\"value\"]") + pk.isDefined shouldBe false + } + + it should "be case-insensitive for parsing" in { + val pk = CosmosPartitionKeyHelper.tryParsePartitionKey("PK([\"value\"])") + pk.isDefined shouldBe true + pk.get shouldEqual new PartitionKey("value") + } + + //scalastyle:on multiple.string.literals + //scalastyle:on magic.number +} diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/ItemsPartitionReaderWithReadManyByPartitionKeyITest.scala b/sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/ItemsPartitionReaderWithReadManyByPartitionKeyITest.scala new file mode 100644 index 000000000000..5c2d7b59836d --- /dev/null +++ b/sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/ItemsPartitionReaderWithReadManyByPartitionKeyITest.scala @@ -0,0 +1,158 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.azure.cosmos.spark + +import com.azure.cosmos.implementation.{CosmosClientMetadataCachesSnapshot, TestConfigurations, Utils} +import com.azure.cosmos.models.PartitionKey +import com.azure.cosmos.spark.diagnostics.DiagnosticsContext +import com.fasterxml.jackson.databind.node.ObjectNode +import org.apache.spark.MockTaskContext +import org.apache.spark.broadcast.Broadcast +import org.apache.spark.sql.types.{StringType, StructField, StructType} + +import java.util.UUID +import scala.collection.mutable.ListBuffer + +class ItemsPartitionReaderWithReadManyByPartitionKeyITest + extends IntegrationSpec + with Spark + with AutoCleanableCosmosContainersWithPkAsPartitionKey { + private val idProperty = "id" + private val pkProperty = "pk" + + //scalastyle:off multiple.string.literals + //scalastyle:off magic.number + + it should "be able to retrieve all items for given partition keys" in { + val container = cosmosClient.getDatabase(cosmosDatabase).getContainer(cosmosContainersWithPkAsPartitionKey) + + // Create items with known PK values + val partitionKeyDefinition = container.read().block().getProperties.getPartitionKeyDefinition + val allItemsByPk = scala.collection.mutable.Map[String, ListBuffer[ObjectNode]]() + val pkValues = List("pkA", "pkB", "pkC") + + for (pk <- pkValues) { + allItemsByPk(pk) = ListBuffer[ObjectNode]() + for (_ <- 1 to 5) { + val objectNode = Utils.getSimpleObjectMapper.createObjectNode() + objectNode.put(idProperty, UUID.randomUUID().toString) + objectNode.put(pkProperty, pk) + container.createItem(objectNode).block() + allItemsByPk(pk) += objectNode + } + } + + val config = Map( + "spark.cosmos.accountEndpoint" -> TestConfigurations.HOST, + "spark.cosmos.accountKey" -> TestConfigurations.MASTER_KEY, + "spark.cosmos.database" -> cosmosDatabase, + "spark.cosmos.container" -> cosmosContainersWithPkAsPartitionKey, + "spark.cosmos.read.inferSchema.enabled" -> "true", + "spark.cosmos.applicationName" -> "ReadManyByPKTest" + ) + + val readSchema = StructType(Seq( + StructField(idProperty, StringType, false), + StructField(pkProperty, StringType, false) + )) + + val diagnosticsContext = DiagnosticsContext(UUID.randomUUID(), "") + val diagnosticsConfig = DiagnosticsConfig.parseDiagnosticsConfig(config) + val cosmosClientMetadataCachesSnapshots = getCosmosClientMetadataCachesSnapshots() + + // Read items for pkA and pkB (not pkC) + val targetPks = List("pkA", "pkB") + val pkIterator = targetPks.map(pk => new PartitionKey(pk)).iterator + + val reader = ItemsPartitionReaderWithReadManyByPartitionKey( + config, + NormalizedRange("", "FF"), + readSchema, + diagnosticsContext, + cosmosClientMetadataCachesSnapshots, + diagnosticsConfig, + "", + MockTaskContext.mockTaskContext(), + pkIterator + ) + + val cosmosRowConverter = CosmosRowConverter.get(CosmosSerializationConfig.parseSerializationConfig(config)) + val itemsReadFromReader = ListBuffer[ObjectNode]() + while (reader.next()) { + itemsReadFromReader += cosmosRowConverter.fromInternalRowToObjectNode(reader.get(), readSchema) + } + + // Should have 10 items (5 for pkA + 5 for pkB) + itemsReadFromReader.size shouldEqual 10 + + // All items should be from pkA or pkB + itemsReadFromReader.foreach(item => { + val pk = item.get(pkProperty).asText() + targetPks should contain(pk) + }) + + // Validate all expected IDs are present + val expectedIds = (allItemsByPk("pkA") ++ allItemsByPk("pkB")).map(_.get(idProperty).asText()).toSet + val actualIds = itemsReadFromReader.map(_.get(idProperty).asText()).toSet + actualIds shouldEqual expectedIds + + reader.close() + } + + it should "return empty results for non-existent partition keys" in { + val config = Map( + "spark.cosmos.accountEndpoint" -> TestConfigurations.HOST, + "spark.cosmos.accountKey" -> TestConfigurations.MASTER_KEY, + "spark.cosmos.database" -> cosmosDatabase, + "spark.cosmos.container" -> cosmosContainersWithPkAsPartitionKey, + "spark.cosmos.read.inferSchema.enabled" -> "true", + "spark.cosmos.applicationName" -> "ReadManyByPKEmptyTest" + ) + + val readSchema = StructType(Seq( + StructField(idProperty, StringType, false), + StructField(pkProperty, StringType, false) + )) + + val diagnosticsContext = DiagnosticsContext(UUID.randomUUID(), "") + val diagnosticsConfig = DiagnosticsConfig.parseDiagnosticsConfig(config) + val cosmosClientMetadataCachesSnapshots = getCosmosClientMetadataCachesSnapshots() + + val pkIterator = List(new PartitionKey("nonExistentPk")).iterator + + val reader = ItemsPartitionReaderWithReadManyByPartitionKey( + config, + NormalizedRange("", "FF"), + readSchema, + diagnosticsContext, + cosmosClientMetadataCachesSnapshots, + diagnosticsConfig, + "", + MockTaskContext.mockTaskContext(), + pkIterator + ) + + val itemsReadFromReader = ListBuffer[ObjectNode]() + val cosmosRowConverter = CosmosRowConverter.get(CosmosSerializationConfig.parseSerializationConfig(config)) + while (reader.next()) { + itemsReadFromReader += cosmosRowConverter.fromInternalRowToObjectNode(reader.get(), readSchema) + } + + itemsReadFromReader.size shouldEqual 0 + reader.close() + } + + private def getCosmosClientMetadataCachesSnapshots(): Broadcast[CosmosClientMetadataCachesSnapshots] = { + val cosmosClientMetadataCachesSnapshot = new CosmosClientMetadataCachesSnapshot() + cosmosClientMetadataCachesSnapshot.serialize(cosmosClient) + + spark.sparkContext.broadcast( + CosmosClientMetadataCachesSnapshots( + cosmosClientMetadataCachesSnapshot, + Option.empty[CosmosClientMetadataCachesSnapshot])) + } + + //scalastyle:on multiple.string.literals + //scalastyle:on magic.number +} diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/ReadManyByPartitionKeyTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/ReadManyByPartitionKeyTest.java new file mode 100644 index 000000000000..8fa657b62105 --- /dev/null +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/ReadManyByPartitionKeyTest.java @@ -0,0 +1,381 @@ +/* + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. + */ + +package com.azure.cosmos; + +import com.azure.cosmos.models.CosmosContainerProperties; +import com.azure.cosmos.models.CosmosItemRequestOptions; +import com.azure.cosmos.models.FeedResponse; +import com.azure.cosmos.models.PartitionKey; +import com.azure.cosmos.models.PartitionKeyBuilder; +import com.azure.cosmos.models.PartitionKeyDefinition; +import com.azure.cosmos.models.PartitionKeyDefinitionVersion; +import com.azure.cosmos.models.PartitionKind; +import com.azure.cosmos.models.SqlParameter; +import com.azure.cosmos.models.SqlQuerySpec; +import com.azure.cosmos.rx.TestSuiteBase; +import com.azure.cosmos.util.CosmosPagedIterable; +import com.fasterxml.jackson.databind.node.ObjectNode; +import org.testng.annotations.AfterClass; +import org.testng.annotations.BeforeClass; +import org.testng.annotations.Factory; +import org.testng.annotations.Test; + +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.List; +import java.util.UUID; +import java.util.stream.Collectors; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.fail; + +public class ReadManyByPartitionKeyTest extends TestSuiteBase { + + private String preExistingDatabaseId = CosmosDatabaseForTest.generateId(); + private CosmosClient client; + private CosmosDatabase createdDatabase; + + // Single PK container (/mypk) + private CosmosContainer singlePkContainer; + + // HPK container (/city, /zipcode, /areaCode) + private CosmosContainer multiHashContainer; + + @Factory(dataProvider = "clientBuilders") + public ReadManyByPartitionKeyTest(CosmosClientBuilder clientBuilder) { + super(clientBuilder); + } + + @BeforeClass(groups = {"emulator"}, timeOut = SETUP_TIMEOUT) + public void before_ReadManyByPartitionKeyTest() { + client = getClientBuilder().buildClient(); + createdDatabase = createSyncDatabase(client, preExistingDatabaseId); + + // Single PK container + String singlePkContainerName = UUID.randomUUID().toString(); + CosmosContainerProperties singlePkProps = new CosmosContainerProperties(singlePkContainerName, "/mypk"); + createdDatabase.createContainer(singlePkProps); + singlePkContainer = createdDatabase.getContainer(singlePkContainerName); + + // HPK container + String multiHashContainerName = UUID.randomUUID().toString(); + PartitionKeyDefinition hpkDef = new PartitionKeyDefinition(); + hpkDef.setKind(PartitionKind.MULTI_HASH); + hpkDef.setVersion(PartitionKeyDefinitionVersion.V2); + ArrayList paths = new ArrayList<>(); + paths.add("/city"); + paths.add("/zipcode"); + paths.add("/areaCode"); + hpkDef.setPaths(paths); + + CosmosContainerProperties hpkProps = new CosmosContainerProperties(multiHashContainerName, hpkDef); + createdDatabase.createContainer(hpkProps); + multiHashContainer = createdDatabase.getContainer(multiHashContainerName); + } + + @AfterClass(groups = {"emulator"}, timeOut = SHUTDOWN_TIMEOUT, alwaysRun = true) + public void afterClass() { + safeDeleteSyncDatabase(createdDatabase); + safeCloseSyncClient(client); + } + + //region Single PK tests + + @Test(groups = {"emulator"}, timeOut = TIMEOUT) + public void singlePk_readManyByPartitionKey_basic() { + // Create items with different PKs + List items = createSinglePkItems("pk1", 3); + items.addAll(createSinglePkItems("pk2", 2)); + items.addAll(createSinglePkItems("pk3", 4)); + + // Read by 2 partition keys + List pkValues = Arrays.asList( + new PartitionKey("pk1"), + new PartitionKey("pk2")); + + CosmosPagedIterable results = singlePkContainer.readManyByPartitionKey(pkValues, ObjectNode.class); + List resultList = results.stream().collect(Collectors.toList()); + + assertThat(resultList).hasSize(5); // 3 + 2 + resultList.forEach(item -> { + String pk = item.get("mypk").asText(); + assertThat(pk).isIn("pk1", "pk2"); + }); + + // Cleanup + cleanupContainer(singlePkContainer); + } + + @Test(groups = {"emulator"}, timeOut = TIMEOUT) + public void singlePk_readManyByPartitionKey_withProjection() { + List items = createSinglePkItems("pkProj", 2); + + List pkValues = Collections.singletonList(new PartitionKey("pkProj")); + SqlQuerySpec customQuery = new SqlQuerySpec("SELECT c.id, c.mypk FROM c"); + + CosmosPagedIterable results = singlePkContainer.readManyByPartitionKey( + pkValues, customQuery, null, ObjectNode.class); + List resultList = results.stream().collect(Collectors.toList()); + + assertThat(resultList).hasSize(2); + // Should only have id and mypk fields (plus system properties) + resultList.forEach(item -> { + assertThat(item.has("id")).isTrue(); + assertThat(item.has("mypk")).isTrue(); + }); + + cleanupContainer(singlePkContainer); + } + + @Test(groups = {"emulator"}, timeOut = TIMEOUT) + public void singlePk_readManyByPartitionKey_withAdditionalFilter() { + // Create items with different "status" values + createSinglePkItemsWithStatus("pkFilter", "active", 3); + createSinglePkItemsWithStatus("pkFilter", "inactive", 2); + + List pkValues = Collections.singletonList(new PartitionKey("pkFilter")); + SqlQuerySpec customQuery = new SqlQuerySpec( + "SELECT * FROM c WHERE c.status = @status", + Arrays.asList(new SqlParameter("@status", "active"))); + + CosmosPagedIterable results = singlePkContainer.readManyByPartitionKey( + pkValues, customQuery, null, ObjectNode.class); + List resultList = results.stream().collect(Collectors.toList()); + + assertThat(resultList).hasSize(3); + resultList.forEach(item -> { + assertThat(item.get("status").asText()).isEqualTo("active"); + }); + + cleanupContainer(singlePkContainer); + } + + @Test(groups = {"emulator"}, timeOut = TIMEOUT) + public void singlePk_readManyByPartitionKey_emptyResults() { + List pkValues = Collections.singletonList(new PartitionKey("nonExistent")); + + CosmosPagedIterable results = singlePkContainer.readManyByPartitionKey(pkValues, ObjectNode.class); + List resultList = results.stream().collect(Collectors.toList()); + + assertThat(resultList).isEmpty(); + } + + //endregion + + //region HPK tests + + @Test(groups = {"emulator"}, timeOut = TIMEOUT) + public void hpk_readManyByPartitionKey_fullPk() { + createHpkItems(); + + // Read by full PKs + List pkValues = Arrays.asList( + new PartitionKeyBuilder().add("Redmond").add("98053").add(1).build(), + new PartitionKeyBuilder().add("Pittsburgh").add("15232").add(2).build()); + + CosmosPagedIterable results = multiHashContainer.readManyByPartitionKey(pkValues, ObjectNode.class); + List resultList = results.stream().collect(Collectors.toList()); + + // Redmond/98053/1 has 2 items, Pittsburgh/15232/2 has 1 item + assertThat(resultList).hasSize(3); + + cleanupContainer(multiHashContainer); + } + + @Test(groups = {"emulator"}, timeOut = TIMEOUT) + public void hpk_readManyByPartitionKey_partialPk_singleLevel() { + createHpkItems(); + + // Read by partial PK (only city) + List pkValues = Collections.singletonList( + new PartitionKeyBuilder().add("Redmond").build()); + + CosmosPagedIterable results = multiHashContainer.readManyByPartitionKey(pkValues, ObjectNode.class); + List resultList = results.stream().collect(Collectors.toList()); + + // Redmond has 3 items total (2 with 98053/1 and 1 with 12345/1) + assertThat(resultList).hasSize(3); + resultList.forEach(item -> { + assertThat(item.get("city").asText()).isEqualTo("Redmond"); + }); + + cleanupContainer(multiHashContainer); + } + + @Test(groups = {"emulator"}, timeOut = TIMEOUT) + public void hpk_readManyByPartitionKey_partialPk_twoLevels() { + createHpkItems(); + + // Read by partial PK (city + zipcode) + List pkValues = Collections.singletonList( + new PartitionKeyBuilder().add("Redmond").add("98053").build()); + + CosmosPagedIterable results = multiHashContainer.readManyByPartitionKey(pkValues, ObjectNode.class); + List resultList = results.stream().collect(Collectors.toList()); + + // Redmond/98053 has 2 items + assertThat(resultList).hasSize(2); + resultList.forEach(item -> { + assertThat(item.get("city").asText()).isEqualTo("Redmond"); + assertThat(item.get("zipcode").asText()).isEqualTo("98053"); + }); + + cleanupContainer(multiHashContainer); + } + + @Test(groups = {"emulator"}, timeOut = TIMEOUT) + public void hpk_readManyByPartitionKey_withProjection() { + createHpkItems(); + + List pkValues = Collections.singletonList( + new PartitionKeyBuilder().add("Redmond").add("98053").add(1).build()); + + SqlQuerySpec customQuery = new SqlQuerySpec("SELECT c.id, c.city FROM c"); + + CosmosPagedIterable results = multiHashContainer.readManyByPartitionKey( + pkValues, customQuery, null, ObjectNode.class); + List resultList = results.stream().collect(Collectors.toList()); + + assertThat(resultList).hasSize(2); + + cleanupContainer(multiHashContainer); + } + + //endregion + + //region Negative/validation tests + + @Test(groups = {"emulator"}, timeOut = TIMEOUT) + public void rejectsAggregateQuery() { + List pkValues = Collections.singletonList(new PartitionKey("pk1")); + SqlQuerySpec aggregateQuery = new SqlQuerySpec("SELECT COUNT(1) FROM c"); + + try { + singlePkContainer.readManyByPartitionKey(pkValues, aggregateQuery, null, ObjectNode.class) + .stream().collect(Collectors.toList()); + fail("Should have thrown IllegalArgumentException for aggregate query"); + } catch (IllegalArgumentException e) { + assertThat(e.getMessage()).contains("aggregates"); + } + } + + @Test(groups = {"emulator"}, timeOut = TIMEOUT) + public void rejectsOrderByQuery() { + List pkValues = Collections.singletonList(new PartitionKey("pk1")); + SqlQuerySpec orderByQuery = new SqlQuerySpec("SELECT * FROM c ORDER BY c.id"); + + try { + singlePkContainer.readManyByPartitionKey(pkValues, orderByQuery, null, ObjectNode.class) + .stream().collect(Collectors.toList()); + fail("Should have thrown IllegalArgumentException for ORDER BY query"); + } catch (IllegalArgumentException e) { + assertThat(e.getMessage()).contains("ORDER BY"); + } + } + + @Test(groups = {"emulator"}, timeOut = TIMEOUT) + public void rejectsDistinctQuery() { + List pkValues = Collections.singletonList(new PartitionKey("pk1")); + SqlQuerySpec distinctQuery = new SqlQuerySpec("SELECT DISTINCT c.mypk FROM c"); + + try { + singlePkContainer.readManyByPartitionKey(pkValues, distinctQuery, null, ObjectNode.class) + .stream().collect(Collectors.toList()); + fail("Should have thrown IllegalArgumentException for DISTINCT query"); + } catch (IllegalArgumentException e) { + assertThat(e.getMessage()).contains("DISTINCT"); + } + } + + @Test(groups = {"emulator"}, timeOut = TIMEOUT) + public void rejectsGroupByQuery() { + List pkValues = Collections.singletonList(new PartitionKey("pk1")); + SqlQuerySpec groupByQuery = new SqlQuerySpec("SELECT c.mypk, COUNT(1) as cnt FROM c GROUP BY c.mypk"); + + try { + singlePkContainer.readManyByPartitionKey(pkValues, groupByQuery, null, ObjectNode.class) + .stream().collect(Collectors.toList()); + fail("Should have thrown IllegalArgumentException for GROUP BY query"); + } catch (IllegalArgumentException e) { + assertThat(e.getMessage()).contains("GROUP BY"); + } + } + + @Test(groups = {"emulator"}, timeOut = TIMEOUT, expectedExceptions = IllegalArgumentException.class) + public void rejectsNullPartitionKeyList() { + singlePkContainer.readManyByPartitionKey((List) null, ObjectNode.class); + } + + @Test(groups = {"emulator"}, timeOut = TIMEOUT, expectedExceptions = IllegalArgumentException.class) + public void rejectsEmptyPartitionKeyList() { + singlePkContainer.readManyByPartitionKey(new ArrayList<>(), ObjectNode.class) + .stream().collect(Collectors.toList()); + } + + //endregion + + //region helper methods + + private List createSinglePkItems(String pkValue, int count) { + List items = new ArrayList<>(); + for (int i = 0; i < count; i++) { + ObjectNode item = com.azure.cosmos.implementation.Utils.getSimpleObjectMapper().createObjectNode(); + item.put("id", UUID.randomUUID().toString()); + item.put("mypk", pkValue); + singlePkContainer.createItem(item); + items.add(item); + } + return items; + } + + private List createSinglePkItemsWithStatus(String pkValue, String status, int count) { + List items = new ArrayList<>(); + for (int i = 0; i < count; i++) { + ObjectNode item = com.azure.cosmos.implementation.Utils.getSimpleObjectMapper().createObjectNode(); + item.put("id", UUID.randomUUID().toString()); + item.put("mypk", pkValue); + item.put("status", status); + singlePkContainer.createItem(item); + items.add(item); + } + return items; + } + + private void createHpkItems() { + // Same data as CosmosMultiHashTest.createItems() + createHpkItem("Redmond", "98053", 1); + createHpkItem("Redmond", "98053", 1); + createHpkItem("Pittsburgh", "15232", 2); + createHpkItem("Stonybrook", "11790", 3); + createHpkItem("Stonybrook", "11794", 3); + createHpkItem("Stonybrook", "11791", 3); + createHpkItem("Redmond", "12345", 1); + } + + private void createHpkItem(String city, String zipcode, int areaCode) { + ObjectNode item = com.azure.cosmos.implementation.Utils.getSimpleObjectMapper().createObjectNode(); + item.put("id", UUID.randomUUID().toString()); + item.put("city", city); + item.put("zipcode", zipcode); + item.put("areaCode", areaCode); + multiHashContainer.createItem(item); + } + + private void cleanupContainer(CosmosContainer container) { + CosmosPagedIterable allItems = container.queryItems( + "SELECT * FROM c", new com.azure.cosmos.models.CosmosQueryRequestOptions(), ObjectNode.class); + allItems.forEach(item -> { + try { + container.deleteItem(item, new CosmosItemRequestOptions()); + } catch (CosmosException e) { + // ignore cleanup failures + } + }); + } + + //endregion +} diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelperTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelperTest.java new file mode 100644 index 000000000000..9c68b3b126e4 --- /dev/null +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelperTest.java @@ -0,0 +1,349 @@ +/* + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. + */ + +package com.azure.cosmos.implementation; + +import com.azure.cosmos.models.PartitionKey; +import com.azure.cosmos.models.PartitionKeyBuilder; +import com.azure.cosmos.models.PartitionKeyDefinition; +import com.azure.cosmos.models.PartitionKeyDefinitionVersion; +import com.azure.cosmos.models.PartitionKind; +import com.azure.cosmos.models.SqlParameter; +import com.azure.cosmos.models.SqlQuerySpec; +import org.testng.annotations.Test; + +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.List; +import java.util.stream.Collectors; + +import static org.assertj.core.api.Assertions.assertThat; + +public class ReadManyByPartitionKeyQueryHelperTest { + + //region Single PK (HASH) tests + + @Test(groups = { "unit" }) + public void singlePk_defaultQuery_singleValue() { + PartitionKeyDefinition pkDef = createSinglePkDefinition("/mypk"); + List selectors = createSelectors(pkDef); + List pkValues = Collections.singletonList(new PartitionKey("pk1")); + + SqlQuerySpec result = ReadManyByPartitionKeyQueryHelper.createReadManyByPkQuerySpec( + "SELECT * FROM c", new ArrayList<>(), pkValues, selectors, pkDef); + + assertThat(result.getQueryText()).contains("SELECT * FROM c WHERE"); + assertThat(result.getQueryText()).contains("IN ("); + assertThat(result.getQueryText()).contains("@pkParam0"); + assertThat(result.getParameters()).hasSize(1); + assertThat(result.getParameters().get(0).getValue(Object.class)).isEqualTo("pk1"); + } + + @Test(groups = { "unit" }) + public void singlePk_defaultQuery_multipleValues() { + PartitionKeyDefinition pkDef = createSinglePkDefinition("/mypk"); + List selectors = createSelectors(pkDef); + List pkValues = Arrays.asList( + new PartitionKey("pk1"), + new PartitionKey("pk2"), + new PartitionKey("pk3")); + + SqlQuerySpec result = ReadManyByPartitionKeyQueryHelper.createReadManyByPkQuerySpec( + "SELECT * FROM c", new ArrayList<>(), pkValues, selectors, pkDef); + + assertThat(result.getQueryText()).contains("IN ("); + assertThat(result.getQueryText()).contains("@pkParam0"); + assertThat(result.getQueryText()).contains("@pkParam1"); + assertThat(result.getQueryText()).contains("@pkParam2"); + assertThat(result.getParameters()).hasSize(3); + } + + @Test(groups = { "unit" }) + public void singlePk_customQuery_noWhere() { + PartitionKeyDefinition pkDef = createSinglePkDefinition("/mypk"); + List selectors = createSelectors(pkDef); + List pkValues = Arrays.asList(new PartitionKey("pk1"), new PartitionKey("pk2")); + + SqlQuerySpec result = ReadManyByPartitionKeyQueryHelper.createReadManyByPkQuerySpec( + "SELECT c.name, c.age FROM c", new ArrayList<>(), pkValues, selectors, pkDef); + + assertThat(result.getQueryText()).startsWith("SELECT c.name, c.age FROM c WHERE"); + assertThat(result.getQueryText()).contains("IN ("); + } + + @Test(groups = { "unit" }) + public void singlePk_customQuery_withExistingWhere() { + PartitionKeyDefinition pkDef = createSinglePkDefinition("/mypk"); + List selectors = createSelectors(pkDef); + List pkValues = Collections.singletonList(new PartitionKey("pk1")); + + List baseParams = new ArrayList<>(); + baseParams.add(new SqlParameter("@minAge", 18)); + + SqlQuerySpec result = ReadManyByPartitionKeyQueryHelper.createReadManyByPkQuerySpec( + "SELECT * FROM c WHERE c.age > @minAge", baseParams, pkValues, selectors, pkDef); + + // Should AND the PK filter to the existing WHERE clause + assertThat(result.getQueryText()).contains("WHERE (c.age > @minAge) AND ("); + assertThat(result.getQueryText()).contains("IN ("); + assertThat(result.getParameters()).hasSize(2); // @minAge + @pkParam1 + assertThat(result.getParameters().get(0).getName()).isEqualTo("@minAge"); + } + + //endregion + + //region HPK (MULTI_HASH) tests + + @Test(groups = { "unit" }) + public void hpk_fullPk_defaultQuery() { + PartitionKeyDefinition pkDef = createMultiHashPkDefinition("/city", "/zipcode"); + List selectors = createSelectors(pkDef); + + PartitionKey pk = new PartitionKeyBuilder().add("Redmond").add("98052").build(); + List pkValues = Collections.singletonList(pk); + + SqlQuerySpec result = ReadManyByPartitionKeyQueryHelper.createReadManyByPkQuerySpec( + "SELECT * FROM c", new ArrayList<>(), pkValues, selectors, pkDef); + + assertThat(result.getQueryText()).contains("SELECT * FROM c WHERE"); + // Should use OR/AND pattern, not IN + assertThat(result.getQueryText()).doesNotContain("IN ("); + assertThat(result.getQueryText()).contains("c[\"city\"] = @pkParam0"); + assertThat(result.getQueryText()).contains("AND"); + assertThat(result.getQueryText()).contains("c[\"zipcode\"] = @pkParam1"); + assertThat(result.getParameters()).hasSize(2); + assertThat(result.getParameters().get(0).getValue(Object.class)).isEqualTo("Redmond"); + assertThat(result.getParameters().get(1).getValue(Object.class)).isEqualTo("98052"); + } + + @Test(groups = { "unit" }) + public void hpk_fullPk_multipleValues() { + PartitionKeyDefinition pkDef = createMultiHashPkDefinition("/city", "/zipcode"); + List selectors = createSelectors(pkDef); + + List pkValues = Arrays.asList( + new PartitionKeyBuilder().add("Redmond").add("98052").build(), + new PartitionKeyBuilder().add("Seattle").add("98101").build()); + + SqlQuerySpec result = ReadManyByPartitionKeyQueryHelper.createReadManyByPkQuerySpec( + "SELECT * FROM c", new ArrayList<>(), pkValues, selectors, pkDef); + + assertThat(result.getQueryText()).contains("OR"); + assertThat(result.getQueryText()).contains("c[\"city\"] = @pkParam0"); + assertThat(result.getQueryText()).contains("c[\"zipcode\"] = @pkParam1"); + assertThat(result.getQueryText()).contains("c[\"city\"] = @pkParam2"); + assertThat(result.getQueryText()).contains("c[\"zipcode\"] = @pkParam3"); + assertThat(result.getParameters()).hasSize(4); + } + + @Test(groups = { "unit" }) + public void hpk_partialPk_singleLevel() { + PartitionKeyDefinition pkDef = createMultiHashPkDefinition("/city", "/zipcode", "/areaCode"); + List selectors = createSelectors(pkDef); + + // Partial PK — only first level + PartitionKey partialPk = new PartitionKeyBuilder().add("Redmond").build(); + List pkValues = Collections.singletonList(partialPk); + + SqlQuerySpec result = ReadManyByPartitionKeyQueryHelper.createReadManyByPkQuerySpec( + "SELECT * FROM c", new ArrayList<>(), pkValues, selectors, pkDef); + + assertThat(result.getQueryText()).contains("c[\"city\"] = @pkParam0"); + // Should NOT include zipcode or areaCode since it's partial + assertThat(result.getQueryText()).doesNotContain("zipcode"); + assertThat(result.getQueryText()).doesNotContain("areaCode"); + assertThat(result.getParameters()).hasSize(1); + } + + @Test(groups = { "unit" }) + public void hpk_partialPk_twoLevels() { + PartitionKeyDefinition pkDef = createMultiHashPkDefinition("/city", "/zipcode", "/areaCode"); + List selectors = createSelectors(pkDef); + + // Partial PK — first two levels + PartitionKey partialPk = new PartitionKeyBuilder().add("Redmond").add("98052").build(); + List pkValues = Collections.singletonList(partialPk); + + SqlQuerySpec result = ReadManyByPartitionKeyQueryHelper.createReadManyByPkQuerySpec( + "SELECT * FROM c", new ArrayList<>(), pkValues, selectors, pkDef); + + assertThat(result.getQueryText()).contains("c[\"city\"] = @pkParam0"); + assertThat(result.getQueryText()).contains("c[\"zipcode\"] = @pkParam1"); + assertThat(result.getQueryText()).doesNotContain("areaCode"); + assertThat(result.getParameters()).hasSize(2); + } + + @Test(groups = { "unit" }) + public void hpk_customQuery_withWhere() { + PartitionKeyDefinition pkDef = createMultiHashPkDefinition("/city", "/zipcode"); + List selectors = createSelectors(pkDef); + + List baseParams = new ArrayList<>(); + baseParams.add(new SqlParameter("@status", "active")); + + PartitionKey pk = new PartitionKeyBuilder().add("Redmond").add("98052").build(); + List pkValues = Collections.singletonList(pk); + + SqlQuerySpec result = ReadManyByPartitionKeyQueryHelper.createReadManyByPkQuerySpec( + "SELECT c.name FROM c WHERE c.status = @status", baseParams, pkValues, selectors, pkDef); + + assertThat(result.getQueryText()).contains("WHERE (c.status = @status) AND ("); + assertThat(result.getQueryText()).contains("c[\"city\"] = @pkParam1"); + assertThat(result.getParameters()).hasSize(3); // @status + 2 pk params + } + + //endregion + + //region findTopLevelWhereIndex tests + + @Test(groups = { "unit" }) + public void findWhere_simpleQuery() { + int idx = ReadManyByPartitionKeyQueryHelper.findTopLevelWhereIndex("SELECT * FROM C WHERE C.ID = 1"); + assertThat(idx).isEqualTo(16); + } + + @Test(groups = { "unit" }) + public void findWhere_noWhere() { + int idx = ReadManyByPartitionKeyQueryHelper.findTopLevelWhereIndex("SELECT * FROM C"); + assertThat(idx).isEqualTo(-1); + } + + @Test(groups = { "unit" }) + public void findWhere_whereInSubquery() { + // WHERE inside parentheses (subquery) should be ignored + int idx = ReadManyByPartitionKeyQueryHelper.findTopLevelWhereIndex( + "SELECT * FROM C WHERE EXISTS(SELECT VALUE T FROM T IN C.TAGS WHERE T = 'FOO')"); + // Should find the outer WHERE, not the inner one + assertThat(idx).isEqualTo(16); + } + + @Test(groups = { "unit" }) + public void findWhere_caseInsensitive() { + int idx = ReadManyByPartitionKeyQueryHelper.findTopLevelWhereIndex("SELECT * FROM C WHERE C.X = 1"); + assertThat(idx).isGreaterThan(0); + } + + @Test(groups = { "unit" }) + public void findWhere_whereNotKeyword() { + // "ELSEWHERE" should not match + int idx = ReadManyByPartitionKeyQueryHelper.findTopLevelWhereIndex("SELECT * FROM ELSEWHERE"); + assertThat(idx).isEqualTo(-1); + } + + //endregion + + //region Custom alias tests + + @Test(groups = { "unit" }) + public void singlePk_customAlias() { + PartitionKeyDefinition pkDef = createSinglePkDefinition("/mypk"); + List selectors = createSelectors(pkDef); + List pkValues = Arrays.asList(new PartitionKey("pk1"), new PartitionKey("pk2")); + + SqlQuerySpec result = ReadManyByPartitionKeyQueryHelper.createReadManyByPkQuerySpec( + "SELECT x.id, x.mypk FROM x", new ArrayList<>(), pkValues, selectors, pkDef); + + assertThat(result.getQueryText()).startsWith("SELECT x.id, x.mypk FROM x WHERE"); + assertThat(result.getQueryText()).contains("x[\"mypk\"] IN ("); + assertThat(result.getQueryText()).doesNotContain("c[\"mypk\"]"); + } + + @Test(groups = { "unit" }) + public void singlePk_customAlias_withWhere() { + PartitionKeyDefinition pkDef = createSinglePkDefinition("/mypk"); + List selectors = createSelectors(pkDef); + List pkValues = Collections.singletonList(new PartitionKey("pk1")); + + List baseParams = new ArrayList<>(); + baseParams.add(new SqlParameter("@cat", "HelloWorld")); + + SqlQuerySpec result = ReadManyByPartitionKeyQueryHelper.createReadManyByPkQuerySpec( + "SELECT x.id, x.mypk FROM x WHERE x.category = @cat", baseParams, pkValues, selectors, pkDef); + + assertThat(result.getQueryText()).contains("WHERE (x.category = @cat) AND (x[\"mypk\"] IN ("); + } + + @Test(groups = { "unit" }) + public void hpk_customAlias() { + PartitionKeyDefinition pkDef = createMultiHashPkDefinition("/city", "/zipcode"); + List selectors = createSelectors(pkDef); + + PartitionKey pk = new PartitionKeyBuilder().add("Redmond").add("98052").build(); + List pkValues = Collections.singletonList(pk); + + SqlQuerySpec result = ReadManyByPartitionKeyQueryHelper.createReadManyByPkQuerySpec( + "SELECT r.name FROM root r", new ArrayList<>(), pkValues, selectors, pkDef); + + assertThat(result.getQueryText()).contains("r[\"city\"] = @pkParam0"); + assertThat(result.getQueryText()).contains("r[\"zipcode\"] = @pkParam1"); + assertThat(result.getQueryText()).doesNotContain("c[\""); + } + + //endregion + + //region extractTableAlias tests + + @Test(groups = { "unit" }) + public void extractAlias_defaultC() { + assertThat(ReadManyByPartitionKeyQueryHelper.extractTableAlias("SELECT * FROM c")).isEqualTo("c"); + } + + @Test(groups = { "unit" }) + public void extractAlias_customX() { + assertThat(ReadManyByPartitionKeyQueryHelper.extractTableAlias("SELECT x.id FROM x WHERE x.age > 5")).isEqualTo("x"); + } + + @Test(groups = { "unit" }) + public void extractAlias_rootWithAlias() { + assertThat(ReadManyByPartitionKeyQueryHelper.extractTableAlias("SELECT r.name FROM root r")).isEqualTo("r"); + } + + @Test(groups = { "unit" }) + public void extractAlias_rootNoAlias() { + assertThat(ReadManyByPartitionKeyQueryHelper.extractTableAlias("SELECT * FROM root")).isEqualTo("root"); + } + + @Test(groups = { "unit" }) + public void extractAlias_containerWithWhere() { + assertThat(ReadManyByPartitionKeyQueryHelper.extractTableAlias("SELECT * FROM items WHERE items.status = 'active'")).isEqualTo("items"); + } + + @Test(groups = { "unit" }) + public void extractAlias_caseInsensitive() { + assertThat(ReadManyByPartitionKeyQueryHelper.extractTableAlias("select * from MyContainer where MyContainer.id = '1'")).isEqualTo("MyContainer"); + } + + //endregion + + //region helpers + + private PartitionKeyDefinition createSinglePkDefinition(String path) { + PartitionKeyDefinition pkDef = new PartitionKeyDefinition(); + pkDef.setKind(PartitionKind.HASH); + pkDef.setVersion(PartitionKeyDefinitionVersion.V2); + pkDef.setPaths(Collections.singletonList(path)); + return pkDef; + } + + private PartitionKeyDefinition createMultiHashPkDefinition(String... paths) { + PartitionKeyDefinition pkDef = new PartitionKeyDefinition(); + pkDef.setKind(PartitionKind.MULTI_HASH); + pkDef.setVersion(PartitionKeyDefinitionVersion.V2); + pkDef.setPaths(Arrays.asList(paths)); + return pkDef; + } + + private List createSelectors(PartitionKeyDefinition pkDef) { + return pkDef.getPaths() + .stream() + .map(pathPart -> pathPart.substring(1)) // skip starting / + .map(pathPart -> pathPart.replace("\"", "\\")) + .map(part -> "[\"" + part + "\"]") + .collect(Collectors.toList()); + } + + //endregion +} diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java index b3888f1bad3a..06d6b7d167bd 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java @@ -1577,6 +1577,117 @@ private Mono> readManyInternal( context); } + /** + * Reads many documents matching the provided partition key values. + * Unlike {@link #readMany(List, Class)} this method does not require item ids - it queries + * all documents matching the provided partition key values. Uses {@code SELECT * FROM c} + * as the base query. + * + * @param the type parameter + * @param partitionKeys list of partition key values to read documents for + * @param classType class type + * @return a {@link CosmosPagedFlux} containing one or several feed response pages + */ + public CosmosPagedFlux readManyByPartitionKey( + List partitionKeys, + Class classType) { + + return this.readManyByPartitionKey(partitionKeys, null, null, classType); + } + + /** + * Reads many documents matching the provided partition key values. + * Unlike {@link #readMany(List, Class)} this method does not require item ids - it queries + * all documents matching the provided partition key values. Uses {@code SELECT * FROM c} + * as the base query. + * + * @param the type parameter + * @param partitionKeys list of partition key values to read documents for + * @param requestOptions the optional request options + * @param classType class type + * @return a {@link CosmosPagedFlux} containing one or several feed response pages + */ + public CosmosPagedFlux readManyByPartitionKey( + List partitionKeys, + CosmosReadManyRequestOptions requestOptions, + Class classType) { + + return this.readManyByPartitionKey(partitionKeys, null, requestOptions, classType); + } + + /** + * Reads many documents matching the provided partition key values with a custom query. + * The custom query can be used to apply projections (e.g. {@code SELECT c.name, c.age FROM c}) + * and/or additional filters (e.g. {@code SELECT * FROM c WHERE c.status = 'active'}). + * The SDK will automatically append partition key filtering to the custom query. + *

+ * The custom query must be a simple streamable query — aggregates, ORDER BY, DISTINCT, + * GROUP BY, DCOUNT, vector search, and full-text search are not supported and will be + * rejected. + *

+ * Partial hierarchical partition keys are supported and will fan out to multiple + * physical partitions. + * + * @param the type parameter + * @param partitionKeys list of partition key values to read documents for + * @param customQuery optional custom query for projections/additional filters (null means SELECT * FROM c) + * @param requestOptions the optional request options + * @param classType class type + * @return a {@link CosmosPagedFlux} containing one or several feed response pages + */ + public CosmosPagedFlux readManyByPartitionKey( + List partitionKeys, + SqlQuerySpec customQuery, + CosmosReadManyRequestOptions requestOptions, + Class classType) { + + return UtilBridgeInternal.createCosmosPagedFlux( + readManyByPartitionKeyInternalFunc(partitionKeys, customQuery, requestOptions, classType)); + } + + private Function>> readManyByPartitionKeyInternalFunc( + List partitionKeys, + SqlQuerySpec customQuery, + CosmosReadManyRequestOptions requestOptions, + Class classType) { + + CosmosAsyncClient client = this.getDatabase().getClient(); + + return (pagedFluxOptions -> { + CosmosQueryRequestOptions queryRequestOptions = requestOptions == null + ? new CosmosQueryRequestOptions() + : queryOptionsAccessor.clone(readManyOptionsAccessor.getImpl(requestOptions)); + queryRequestOptions.setMaxDegreeOfParallelism(-1); + queryRequestOptions.setQueryName("readManyByPartitionKey"); + CosmosQueryRequestOptionsBase cosmosQueryRequestOptionsImpl = queryOptionsAccessor.getImpl(queryRequestOptions); + applyPolicies(OperationType.Query, ResourceType.Document, cosmosQueryRequestOptionsImpl, this.readManyItemsSpanName); + + QueryFeedOperationState state = new QueryFeedOperationState( + client, + this.readManyItemsSpanName, + database.getId(), + this.getId(), + ResourceType.Document, + OperationType.Query, + queryOptionsAccessor.getQueryNameOrDefault(queryRequestOptions, this.readManyItemsSpanName), + queryRequestOptions, + pagedFluxOptions + ); + + pagedFluxOptions.setFeedOperationState(state); + + return CosmosBridgeInternal + .getAsyncDocumentClient(this.getDatabase()) + .readManyByPartitionKey( + partitionKeys, + customQuery, + BridgeInternal.getLink(this), + state, + classType) + .map(response -> prepareFeedResponse(response, false)); + }); + } + /** * Reads all the items of a logical partition * diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosContainer.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosContainer.java index 04a6060c1927..0bd8be5850c0 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosContainer.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosContainer.java @@ -540,6 +540,73 @@ public FeedResponse readMany( classType)); } + /** + * Reads many documents matching the provided partition key values. + * Unlike {@link #readMany(List, Class)} this method does not require item ids - it queries + * all documents matching the provided partition key values. Uses {@code SELECT * FROM c} + * as the base query. + * + * @param the type parameter + * @param partitionKeys list of partition key values to read documents for + * @param classType class type + * @return a {@link CosmosPagedIterable} containing the results + */ + public CosmosPagedIterable readManyByPartitionKey( + List partitionKeys, + Class classType) { + + return getCosmosPagedIterable(this.asyncContainer.readManyByPartitionKey(partitionKeys, classType)); + } + + /** + * Reads many documents matching the provided partition key values. + * Unlike {@link #readMany(List, Class)} this method does not require item ids - it queries + * all documents matching the provided partition key values. Uses {@code SELECT * FROM c} + * as the base query. + * + * @param the type parameter + * @param partitionKeys list of partition key values to read documents for + * @param requestOptions the optional request options + * @param classType class type + * @return a {@link CosmosPagedIterable} containing the results + */ + public CosmosPagedIterable readManyByPartitionKey( + List partitionKeys, + CosmosReadManyRequestOptions requestOptions, + Class classType) { + + return getCosmosPagedIterable(this.asyncContainer.readManyByPartitionKey(partitionKeys, requestOptions, classType)); + } + + /** + * Reads many documents matching the provided partition key values with a custom query. + * The custom query can be used to apply projections (e.g. {@code SELECT c.name, c.age FROM c}) + * and/or additional filters (e.g. {@code SELECT * FROM c WHERE c.status = 'active'}). + * The SDK will automatically append partition key filtering to the custom query. + *

+ * The custom query must be a simple streamable query — aggregates, ORDER BY, DISTINCT, + * GROUP BY, DCOUNT, vector search, and full-text search are not supported and will be + * rejected. + *

+ * Partial hierarchical partition keys are supported and will fan out to multiple + * physical partitions. + * + * @param the type parameter + * @param partitionKeys list of partition key values to read documents for + * @param customQuery optional custom query for projections/additional filters (null means SELECT * FROM c) + * @param requestOptions the optional request options + * @param classType class type + * @return a {@link CosmosPagedIterable} containing the results + */ + public CosmosPagedIterable readManyByPartitionKey( + List partitionKeys, + SqlQuerySpec customQuery, + CosmosReadManyRequestOptions requestOptions, + Class classType) { + + return getCosmosPagedIterable(this.asyncContainer.readManyByPartitionKey(partitionKeys, customQuery, requestOptions, classType)); + } + /** * Reads all the items of a logical partition returning the results as {@link CosmosPagedIterable}. * diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/AsyncDocumentClient.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/AsyncDocumentClient.java index 49e1fdf57f64..26f1ff64ea5d 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/AsyncDocumentClient.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/AsyncDocumentClient.java @@ -1584,6 +1584,27 @@ Mono> readMany( QueryFeedOperationState state, Class klass); + /** + * Reads many documents by partition key values. + * Unlike {@link #readMany(List, String, QueryFeedOperationState, Class)} this method does not require + * item ids - it queries all documents matching the provided partition key values. + * Partial hierarchical partition keys are supported and will fan out to multiple physical partitions. + * + * @param partitionKeys list of partition key values to read documents for + * @param customQuery optional custom query (for projections/additional filters) - null means SELECT * FROM c + * @param collectionLink link for the documentcollection/container to be queried + * @param state the query operation state + * @param klass class type + * @param the type parameter + * @return a Flux with feed response pages of documents + */ + Flux> readManyByPartitionKey( + List partitionKeys, + SqlQuerySpec customQuery, + String collectionLink, + QueryFeedOperationState state, + Class klass); + /** * Read all documents of a certain logical partition. *

diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/Configs.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/Configs.java index 337055c6947f..162b0740f408 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/Configs.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/Configs.java @@ -248,6 +248,11 @@ public class Configs { public static final String MIN_TARGET_BULK_MICRO_BATCH_SIZE_VARIABLE = "COSMOS_MIN_TARGET_BULK_MICRO_BATCH_SIZE"; public static final int DEFAULT_MIN_TARGET_BULK_MICRO_BATCH_SIZE = 1; + // readManyByPartitionKey: max number of PK values per query per physical partition + private static final String READ_MANY_BY_PK_MAX_BATCH_SIZE = "COSMOS.READ_MANY_BY_PK_MAX_BATCH_SIZE"; + private static final String READ_MANY_BY_PK_MAX_BATCH_SIZE_VARIABLE = "COSMOS_READ_MANY_BY_PK_MAX_BATCH_SIZE"; + private static final int DEFAULT_READ_MANY_BY_PK_MAX_BATCH_SIZE = 1000; + public static final String MAX_BULK_MICRO_BATCH_CONCURRENCY = "COSMOS.MAX_BULK_MICRO_BATCH_CONCURRENCY"; public static final String MAX_BULK_MICRO_BATCH_CONCURRENCY_VARIABLE = "COSMOS_MAX_BULK_MICRO_BATCH_CONCURRENCY"; public static final int DEFAULT_MAX_BULK_MICRO_BATCH_CONCURRENCY = 1; @@ -816,6 +821,20 @@ public static int getMinTargetBulkMicroBatchSize() { return DEFAULT_MIN_TARGET_BULK_MICRO_BATCH_SIZE; } + public static int getReadManyByPkMaxBatchSize() { + String valueFromSystemProperty = System.getProperty(READ_MANY_BY_PK_MAX_BATCH_SIZE); + if (valueFromSystemProperty != null && !valueFromSystemProperty.isEmpty()) { + return Math.max(1, Integer.parseInt(valueFromSystemProperty)); + } + + String valueFromEnvVariable = System.getenv(READ_MANY_BY_PK_MAX_BATCH_SIZE_VARIABLE); + if (valueFromEnvVariable != null && !valueFromEnvVariable.isEmpty()) { + return Math.max(1, Integer.parseInt(valueFromEnvVariable)); + } + + return DEFAULT_READ_MANY_BY_PK_MAX_BATCH_SIZE; + } + public static int getMaxBulkMicroBatchConcurrency() { String valueFromSystemProperty = System.getProperty(MAX_BULK_MICRO_BATCH_CONCURRENCY); if (valueFromSystemProperty != null && !valueFromSystemProperty.isEmpty()) { diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelper.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelper.java new file mode 100644 index 000000000000..fbe329bff844 --- /dev/null +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelper.java @@ -0,0 +1,199 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. +package com.azure.cosmos.implementation; + +import com.azure.cosmos.BridgeInternal; +import com.azure.cosmos.implementation.routing.PartitionKeyInternal; +import com.azure.cosmos.models.ModelBridgeInternal; +import com.azure.cosmos.models.PartitionKey; +import com.azure.cosmos.models.PartitionKeyDefinition; +import com.azure.cosmos.models.PartitionKind; +import com.azure.cosmos.models.SqlParameter; +import com.azure.cosmos.models.SqlQuerySpec; + +import java.util.ArrayList; +import java.util.List; + +/** + * Helper for constructing SqlQuerySpec instances for readManyByPartitionKey operations. + * This class is not intended to be used directly by end-users. + */ +public class ReadManyByPartitionKeyQueryHelper { + + private static final String DEFAULT_TABLE_ALIAS = "c"; + + public static SqlQuerySpec createReadManyByPkQuerySpec( + String baseQueryText, + List baseParameters, + List pkValues, + List partitionKeySelectors, + PartitionKeyDefinition pkDefinition) { + + // Extract the table alias from the FROM clause (e.g. "FROM x" → "x", "FROM c" → "c") + String tableAlias = extractTableAlias(baseQueryText); + + StringBuilder pkFilter = new StringBuilder(); + List parameters = new ArrayList<>(baseParameters); + int paramCount = baseParameters.size(); + + boolean isSinglePathPk = partitionKeySelectors.size() == 1; + + if (isSinglePathPk && pkDefinition.getKind() != PartitionKind.MULTI_HASH) { + // Single PK path — use IN clause: alias["pkPath"] IN (@pk0, @pk1, ...) + pkFilter.append(" "); + pkFilter.append(tableAlias); + pkFilter.append(partitionKeySelectors.get(0)); + pkFilter.append(" IN ( "); + for (int i = 0; i < pkValues.size(); i++) { + PartitionKeyInternal pkInternal = BridgeInternal.getPartitionKeyInternal(pkValues.get(i)); + Object[] pkComponents = pkInternal.toObjectArray(); + String pkParamName = "@pkParam" + paramCount; + parameters.add(new SqlParameter(pkParamName, pkComponents[0])); + paramCount++; + + pkFilter.append(pkParamName); + if (i < pkValues.size() - 1) { + pkFilter.append(", "); + } + } + pkFilter.append(" )"); + } else { + // Multiple PK paths (HPK) or MULTI_HASH — use OR of AND clauses + pkFilter.append(" "); + for (int i = 0; i < pkValues.size(); i++) { + PartitionKeyInternal pkInternal = BridgeInternal.getPartitionKeyInternal(pkValues.get(i)); + Object[] pkComponents = pkInternal.toObjectArray(); + + pkFilter.append("("); + for (int j = 0; j < pkComponents.length; j++) { + String pkParamName = "@pkParam" + paramCount; + parameters.add(new SqlParameter(pkParamName, pkComponents[j])); + paramCount++; + + if (j > 0) { + pkFilter.append(" AND "); + } + pkFilter.append(tableAlias); + pkFilter.append(partitionKeySelectors.get(j)); + pkFilter.append(" = "); + pkFilter.append(pkParamName); + } + pkFilter.append(")"); + + if (i < pkValues.size() - 1) { + pkFilter.append(" OR "); + } + } + } + + // Compose final query: handle existing WHERE clause in base query + String finalQuery; + int whereIndex = findTopLevelWhereIndex(baseQueryText); + if (whereIndex >= 0) { + // Base query has WHERE — AND our PK filter + String beforeWhere = baseQueryText.substring(0, whereIndex); + String afterWhere = baseQueryText.substring(whereIndex + 5); // skip "WHERE" + finalQuery = beforeWhere + "WHERE (" + afterWhere.trim() + ") AND (" + pkFilter.toString().trim() + ")"; + } else { + // No WHERE — add one + finalQuery = baseQueryText + " WHERE" + pkFilter.toString(); + } + + return new SqlQuerySpec(finalQuery, parameters); + } + + /** + * Extracts the table/collection alias from a SQL query's FROM clause. + * Handles: "SELECT * FROM c", "SELECT x.id FROM x WHERE ...", "SELECT * FROM root r", etc. + * Returns the alias used after FROM (last token before WHERE or end of FROM clause). + */ + static String extractTableAlias(String queryText) { + String upper = queryText.toUpperCase(); + int fromIndex = findTopLevelKeywordIndex(upper, "FROM"); + if (fromIndex < 0) { + return DEFAULT_TABLE_ALIAS; + } + + // Start scanning after "FROM" + int afterFrom = fromIndex + 4; + // Skip whitespace + while (afterFrom < queryText.length() && Character.isWhitespace(queryText.charAt(afterFrom))) { + afterFrom++; + } + + // Collect the container name token (could be "root", "c", etc.) + int tokenStart = afterFrom; + while (afterFrom < queryText.length() + && !Character.isWhitespace(queryText.charAt(afterFrom)) + && queryText.charAt(afterFrom) != '(' + && queryText.charAt(afterFrom) != ')') { + afterFrom++; + } + String containerName = queryText.substring(tokenStart, afterFrom); + + // Skip whitespace after container name + while (afterFrom < queryText.length() && Character.isWhitespace(queryText.charAt(afterFrom))) { + afterFrom++; + } + + // Check if there's an alias after the container name (before WHERE or end) + if (afterFrom < queryText.length()) { + char nextChar = Character.toUpperCase(queryText.charAt(afterFrom)); + // If the next token is a keyword (WHERE, ORDER, GROUP, JOIN) or end, containerName IS the alias + if (nextChar == 'W' || nextChar == 'O' || nextChar == 'G' || nextChar == 'J') { + // Check if it's actually a keyword + String remaining = upper.substring(afterFrom); + if (remaining.startsWith("WHERE") || remaining.startsWith("ORDER") + || remaining.startsWith("GROUP") || remaining.startsWith("JOIN")) { + return containerName; + } + } + // Otherwise the next token is the alias ("FROM root r" → alias is "r") + int aliasStart = afterFrom; + while (afterFrom < queryText.length() + && !Character.isWhitespace(queryText.charAt(afterFrom)) + && queryText.charAt(afterFrom) != '(' + && queryText.charAt(afterFrom) != ')') { + afterFrom++; + } + if (afterFrom > aliasStart) { + return queryText.substring(aliasStart, afterFrom); + } + } + + return containerName; + } + + /** + * Finds the index of a top-level SQL keyword in the query text (case-insensitive), + * ignoring occurrences inside parentheses. + */ + static int findTopLevelKeywordIndex(String queryText, String keyword) { + String queryTextUpper = queryText.toUpperCase(); + String keywordUpper = keyword.toUpperCase(); + int depth = 0; + int keyLen = keywordUpper.length(); + for (int i = 0; i <= queryTextUpper.length() - keyLen; i++) { + char ch = queryTextUpper.charAt(i); + if (ch == '(') { + depth++; + } else if (ch == ')') { + depth--; + } else if (depth == 0 && ch == keywordUpper.charAt(0) + && queryTextUpper.startsWith(keywordUpper, i) + && (i == 0 || !Character.isLetterOrDigit(queryTextUpper.charAt(i - 1))) + && (i + keyLen >= queryTextUpper.length() || !Character.isLetterOrDigit(queryTextUpper.charAt(i + keyLen)))) { + return i; + } + } + return -1; + } + + /** + * Finds the index of the top-level WHERE keyword in the query text, + * ignoring WHERE that appears inside parentheses (subqueries). + */ + public static int findTopLevelWhereIndex(String queryTextUpper) { + return findTopLevelKeywordIndex(queryTextUpper, "WHERE"); + } +} diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java index 5555b3da671c..17461fa56531 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java @@ -4367,6 +4367,233 @@ private Mono>> readMany( ); } + @Override + public Flux> readManyByPartitionKey( + List partitionKeys, + SqlQuerySpec customQuery, + String collectionLink, + QueryFeedOperationState state, + Class klass) { + + checkNotNull(partitionKeys, "Argument 'partitionKeys' must not be null."); + checkArgument(!partitionKeys.isEmpty(), "Argument 'partitionKeys' must not be empty."); + + final ScopedDiagnosticsFactory diagnosticsFactory = new ScopedDiagnosticsFactory(this, true); + state.registerDiagnosticsFactory( + () -> {}, // we never want to reset in readManyByPartitionKey + (ctx) -> diagnosticsFactory.merge(ctx) + ); + + String resourceLink = parentResourceLinkToQueryLink(collectionLink, ResourceType.Document); + RxDocumentServiceRequest request = RxDocumentServiceRequest.create(diagnosticsFactory, + OperationType.Query, + ResourceType.Document, + collectionLink, null + ); + + Mono> collectionObs = + collectionCache.resolveCollectionAsync(null, request); + + return collectionObs + .flatMapMany(documentCollectionResourceResponse -> { + final DocumentCollection collection = documentCollectionResourceResponse.v; + if (collection == null) { + return Flux.error(new IllegalStateException("Collection cannot be null")); + } + + final PartitionKeyDefinition pkDefinition = collection.getPartitionKey(); + + Mono> valueHolderMono = partitionKeyRangeCache + .tryLookupAsync( + BridgeInternal.getMetaDataDiagnosticContext(request.requestContext.cosmosDiagnostics), + collection.getResourceId(), + null, + null); + + // Validate custom query if provided + Mono queryValidationMono; + if (customQuery != null) { + queryValidationMono = validateCustomQueryForReadManyByPartitionKey( + customQuery, resourceLink, state.getQueryOptions()); + } else { + queryValidationMono = Mono.empty(); + } + + return Mono.zip(valueHolderMono, queryValidationMono.then(Mono.just(true))) + .flatMapMany(tuple -> { + CollectionRoutingMap routingMap = tuple.getT1().v; + if (routingMap == null) { + return Flux.error(new IllegalStateException("Failed to get routing map.")); + } + + Map> partitionRangePkMap = + groupPartitionKeysByPhysicalPartition(partitionKeys, pkDefinition, routingMap); + + List partitionKeySelectors = createPkSelectors(pkDefinition); + + String baseQueryText; + List baseParameters; + if (customQuery != null) { + baseQueryText = customQuery.getQueryText(); + baseParameters = customQuery.getParameters() != null + ? new ArrayList<>(customQuery.getParameters()) + : new ArrayList<>(); + } else { + baseQueryText = "SELECT * FROM c"; + baseParameters = new ArrayList<>(); + } + + // Build per-physical-partition batched queries. + // Each physical partition may have many PKs — split into batches + // to avoid oversized SQL queries. Batch size is configurable via + // system property COSMOS.READ_MANY_BY_PK_MAX_BATCH_SIZE (default 1000). + int maxPksPerPartitionQuery = Configs.getReadManyByPkMaxBatchSize(); + + // Build batches per partition as a list of lists (one inner list per partition). + // Then interleave in round-robin order so that concurrent execution + // prefers different physical partitions over multiple batches of the same partition. + List>> batchesPerPartition = new ArrayList<>(); + int maxBatchesPerPartition = 0; + + for (Map.Entry> entry : partitionRangePkMap.entrySet()) { + List allPks = entry.getValue(); + if (allPks.isEmpty()) { + continue; + } + List> partitionBatches = new ArrayList<>(); + for (int i = 0; i < allPks.size(); i += maxPksPerPartitionQuery) { + List batch = allPks.subList( + i, Math.min(i + maxPksPerPartitionQuery, allPks.size())); + SqlQuerySpec querySpec = ReadManyByPartitionKeyQueryHelper + .createReadManyByPkQuerySpec( + baseQueryText, baseParameters, batch, + partitionKeySelectors, pkDefinition); + partitionBatches.add( + Collections.singletonMap(entry.getKey(), querySpec)); + } + batchesPerPartition.add(partitionBatches); + maxBatchesPerPartition = Math.max(maxBatchesPerPartition, partitionBatches.size()); + } + + if (batchesPerPartition.isEmpty()) { + return Flux.empty(); + } + + // Round-robin interleave: [batch0-p1, batch0-p2, ..., batch0-pN, batch1-p1, batch1-p2, ...] + // This ensures that with bounded concurrency, different partitions are + // preferred over sequential batches of the same partition. + List> interleavedBatches = new ArrayList<>(); + for (int batchIdx = 0; batchIdx < maxBatchesPerPartition; batchIdx++) { + for (List> partitionBatches : batchesPerPartition) { + if (batchIdx < partitionBatches.size()) { + interleavedBatches.add(partitionBatches.get(batchIdx)); + } + } + } + + // Execute all batches with bounded concurrency. + List>> queryFluxes = interleavedBatches + .stream() + .map(batchMap -> queryForReadMany( + diagnosticsFactory, + resourceLink, + new SqlQuerySpec(DUMMY_SQL_QUERY), + state.getQueryOptions(), + klass, + ResourceType.Document, + collection, + Collections.unmodifiableMap(batchMap))) + .collect(Collectors.toList()); + + int fluxConcurrency = Math.min(queryFluxes.size(), + Math.max(Configs.getCPUCnt(), 1)); + + return Flux.mergeSequential(queryFluxes, fluxConcurrency, 1); + }); + }); + } + + private Mono validateCustomQueryForReadManyByPartitionKey( + SqlQuerySpec customQuery, + String resourceLink, + CosmosQueryRequestOptions queryRequestOptions) { + + IDocumentQueryClient queryClient = documentQueryClientImpl( + RxDocumentClientImpl.this, getOperationContextAndListenerTuple(queryRequestOptions)); + + return DocumentQueryExecutionContextFactory + .fetchQueryPlanForValidation(this, queryClient, customQuery, resourceLink, queryRequestOptions) + .flatMap(queryPlan -> { + QueryInfo queryInfo = queryPlan.getQueryInfo(); + + if (queryInfo.hasAggregates()) { + return Mono.error(new IllegalArgumentException( + "Custom query for readMany by partition key must not contain aggregates.")); + } + if (queryInfo.hasOrderBy()) { + return Mono.error(new IllegalArgumentException( + "Custom query for readMany by partition key must not contain ORDER BY.")); + } + if (queryInfo.hasDistinct()) { + return Mono.error(new IllegalArgumentException( + "Custom query for readMany by partition key must not contain DISTINCT.")); + } + if (queryInfo.hasGroupBy()) { + return Mono.error(new IllegalArgumentException( + "Custom query for readMany by partition key must not contain GROUP BY.")); + } + if (queryInfo.hasDCount()) { + return Mono.error(new IllegalArgumentException( + "Custom query for readMany by partition key must not contain DCOUNT.")); + } + if (queryInfo.hasNonStreamingOrderBy()) { + return Mono.error(new IllegalArgumentException( + "Custom query for readMany by partition key must not contain non-streaming ORDER BY.")); + } + if (queryPlan.hasHybridSearchQueryInfo()) { + return Mono.error(new IllegalArgumentException( + "Custom query for readMany by partition key must not contain hybrid/vector/full-text search.")); + } + + return Mono.empty(); + }); + } + + private Map> groupPartitionKeysByPhysicalPartition( + List partitionKeys, + PartitionKeyDefinition pkDefinition, + CollectionRoutingMap routingMap) { + + Map> partitionRangePkMap = new HashMap<>(); + + for (PartitionKey pk : partitionKeys) { + PartitionKeyInternal pkInternal = BridgeInternal.getPartitionKeyInternal(pk); + int componentCount = pkInternal.getComponents().size(); + int definedPathCount = pkDefinition.getPaths().size(); + + List targetRanges; + + if (pkDefinition.getKind() == PartitionKind.MULTI_HASH && componentCount < definedPathCount) { + // Partial HPK — compute EPK prefix range and find all overlapping physical partitions + Range epkRange = PartitionKeyInternalHelper.getEPKRangeForPrefixPartitionKey( + pkInternal, pkDefinition); + targetRanges = routingMap.getOverlappingRanges(epkRange); + } else { + // Full PK — maps to exactly one physical partition + String effectivePartitionKeyString = PartitionKeyInternalHelper + .getEffectivePartitionKeyString(pkInternal, pkDefinition); + PartitionKeyRange range = routingMap.getRangeByEffectivePartitionKey(effectivePartitionKeyString); + targetRanges = Collections.singletonList(range); + } + + for (PartitionKeyRange range : targetRanges) { + partitionRangePkMap.computeIfAbsent(range, k -> new ArrayList<>()).add(pk); + } + } + + return partitionRangePkMap; + } + private Map getRangeQueryMap( Map> partitionRangeItemKeyMap, PartitionKeyDefinition partitionKeyDefinition) { diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/DocumentQueryExecutionContextFactory.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/DocumentQueryExecutionContextFactory.java index e142f35339dd..0e8aa78ec190 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/DocumentQueryExecutionContextFactory.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/DocumentQueryExecutionContextFactory.java @@ -324,6 +324,17 @@ private static List getFeedRangeEpks(List> range return feedRanges; } + public static Mono fetchQueryPlanForValidation( + DiagnosticsClientContext diagnosticsClientContext, + IDocumentQueryClient queryClient, + SqlQuerySpec sqlQuerySpec, + String resourceLink, + CosmosQueryRequestOptions queryRequestOptions) { + + return QueryPlanRetriever.getQueryPlanThroughGatewayAsync( + diagnosticsClientContext, queryClient, sqlQuerySpec, resourceLink, queryRequestOptions); + } + public static Flux> createDocumentQueryExecutionContextAsync( DiagnosticsClientContext diagnosticsClientContext, IDocumentQueryClient client, diff --git a/sdk/cosmos/docs/readManyByPartitionKey-design.md b/sdk/cosmos/docs/readManyByPartitionKey-design.md new file mode 100644 index 000000000000..f53cde1db37a --- /dev/null +++ b/sdk/cosmos/docs/readManyByPartitionKey-design.md @@ -0,0 +1,133 @@ +# readMany by Partition Key — Design & Implementation Plan + +## Overview + +New `readMany` overloads on `CosmosAsyncContainer` / `CosmosContainer` that accept a +`List` (without item-id). The SDK splits the PK values by physical +partition, generates a streaming query per physical partition, and returns results as +`CosmosPagedFlux` / `CosmosPagedIterable`. + +An optional `SqlQuerySpec` parameter lets callers supply a custom query for projections +and additional filters. The SDK appends the auto-generated PK WHERE clause to it. + +## Decisions + +| Topic | Decision | +|---|---| +| API name | `readMany` — new overload distinguished by `List` parameter | +| Return type | `CosmosPagedFlux` (async) / `CosmosPagedIterable` (sync) | +| Custom query format | `SqlQuerySpec` — full query with parameters; SDK ANDs the PK filter | +| Partial HPK | Supported from the start; prefix PKs fan out via `getOverlappingRanges` | +| PK deduplication | Done at Spark layer only, not in the SDK | +| Spark UDF | New `GetCosmosPartitionKeyValue` UDF | +| Custom query validation | Gateway query plan; reject aggregates/ORDER BY/DISTINCT/GROUP BY/DCount/vector/fulltext | +| Max PK list size | Enforced per invocation (same effective cap as existing readMany) | + +## Phase 1 — SDK Core (`azure-cosmos`) + +### Step 1: New public overloads in CosmosAsyncContainer + +```java + CosmosPagedFlux readMany(List partitionKeys, Class classType) + CosmosPagedFlux readMany(List partitionKeys, + CosmosReadManyRequestOptions requestOptions, + Class classType) + CosmosPagedFlux readMany(List partitionKeys, + SqlQuerySpec customQuery, + CosmosReadManyRequestOptions requestOptions, + Class classType) +``` + +All delegate to a private `readManyByPartitionKeyInternal(...)`. + +### Step 2: Sync wrappers in CosmosContainer + +Same signatures returning `CosmosPagedIterable`, delegating to the async container. + +### Step 3: Internal orchestration (RxDocumentClientImpl) + +1. Resolve collection metadata + PK definition from cache. +2. Fetch routing map from `partitionKeyRangeCache`. +3. For each `PartitionKey`: + - Compute effective partition key (EPK). + - Full PK → `getRangeByEffectivePartitionKey()` (single range). + - Partial HPK → compute EPK prefix range → `getOverlappingRanges()` (multiple ranges). + **Note:** partial HPK intentionally fans out to multiple physical partitions. +4. Group PK values by `PartitionKeyRange`. +5. If custom `SqlQuerySpec` provided → validate via query plan (Step 4). +6. Per physical partition → build `SqlQuerySpec` with PK WHERE clause (Step 5). +7. Execute queries via `createReadManyQueryAsync()`. +8. Return results as `CosmosPagedFlux`. + +### Step 4: Custom query validation + +One-time call per invocation (existing query plan caching applies): + +- `QueryPlanRetriever.getQueryPlanThroughGatewayAsync()` for the user query. +- Reject (`IllegalArgumentException`) if: + - `queryInfo.hasAggregates()` + - `queryInfo.hasOrderBy()` + - `queryInfo.hasDistinct()` + - `queryInfo.hasGroupBy()` + - `queryInfo.hasDCount()` + - `queryInfo.hasNonStreamingOrderBy()` + - `partitionedQueryExecutionInfo.hasHybridSearchQueryInfo()` + +### Step 5: Query construction + +**Single PK (HASH):** +```sql +{baseQuery} WHERE c["{pkPath}"] IN (@pk0, @pk1, @pk2) +``` + +**Full HPK (MULTI_HASH):** +```sql +{baseQuery} WHERE (c["{path1}"] = @p0l1 AND c["{path2}"] = @p0l2) + OR (c["{path1}"] = @p1l1 AND c["{path2}"] = @p1l2) +``` + +**Partial HPK (prefix-only):** +```sql +{baseQuery} WHERE (c["{path1}"] = @p0l1) + OR (c["{path1}"] = @p1l1) +``` + +If the base query already has a WHERE clause: +```sql +{selectAndFrom} WHERE ({existingWhere}) AND ({pkFilter}) +``` + +### Step 6: Bridge / accessor wiring + +Expose internal method through `ImplementationBridgeHelpers`. + +## Phase 2 — Spark Connector (`azure-cosmos-spark_3`) + +### Step 7: New UDF — `GetCosmosPartitionKeyValue` + +- Input: partition key column(s) as array. +- Output: serialized PK string. + +### Step 8: PK-only serialization helper + +`CosmosPartitionKeyHelper`: +- `getCosmosPartitionKeyValueString(pkValues)` — serialize. +- `tryParsePartitionKey(serialized)` — deserialize. + +### Step 9: `CosmosItemsDataSource.readManyByPartitionKey` + +Static entry points, deduplicates PKs at Spark level, delegates to reader. + +### Step 10: `CosmosReadManyByPartitionKeyReader` + +Per-Spark-partition execution, analogous to `CosmosReadManyReader`. + +### Step 11: `ItemsPartitionReaderWithReadManyByPartitionKey` + +Calls new SDK API with `Iterator[PartitionKey]`, iterates `CosmosPagedFlux` pages. + +## Phase 3 — Testing + +- Unit tests: query construction (single PK, HPK, partial HPK, custom query composition). +- Unit tests: query plan rejection (aggregates, ORDER BY, DISTINCT, etc.). +- Integration tests: end-to-end SDK + Spark UDF. From 9a5b3e96e7e6d1ad83990bfea9244a8284ccbb81 Mon Sep 17 00:00:00 2001 From: Fabian Meiswinkel Date: Tue, 14 Apr 2026 16:49:23 +0200 Subject: [PATCH 02/64] Update sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/ItemsPartitionReaderWithReadManyByPartitionKey.scala Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- .../ItemsPartitionReaderWithReadManyByPartitionKey.scala | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/ItemsPartitionReaderWithReadManyByPartitionKey.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/ItemsPartitionReaderWithReadManyByPartitionKey.scala index 68e41cad3ec5..9df4b79fb238 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/ItemsPartitionReaderWithReadManyByPartitionKey.scala +++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/ItemsPartitionReaderWithReadManyByPartitionKey.scala @@ -118,7 +118,14 @@ private[spark] case class ItemsPartitionReaderWithReadManyByPartitionKey readManyOptionsImpl .setCustomItemSerializer( new CosmosItemSerializerNoExceptionWrapping { - override def serialize[T](item: T): util.Map[String, AnyRef] = ??? + override def serialize[T](item: T): util.Map[String, AnyRef] = { + throw new UnsupportedOperationException( + s"Serialization is not supported by the custom item serializer in " + + s"ItemsPartitionReaderWithReadManyByPartitionKey; this serializer is intended " + + s"for deserializing read-many responses into SparkRowItem only. " + + s"Unexpected item type: ${if (item == null) "null" else item.getClass.getName}" + ) + } override def deserialize[T](jsonNodeMap: util.Map[String, AnyRef], classType: Class[T]): T = { if (jsonNodeMap == null) { From a8720c3c9f27a4f1861ea2e59d5016541d90904b Mon Sep 17 00:00:00 2001 From: Fabian Meiswinkel Date: Tue, 14 Apr 2026 16:50:35 +0200 Subject: [PATCH 03/64] Update sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/ItemsPartitionReaderWithReadManyByPartitionKey.scala Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- ...tionReaderWithReadManyByPartitionKey.scala | 89 ++++++++++++------- 1 file changed, 57 insertions(+), 32 deletions(-) diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/ItemsPartitionReaderWithReadManyByPartitionKey.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/ItemsPartitionReaderWithReadManyByPartitionKey.scala index 9df4b79fb238..dddbb23a0050 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/ItemsPartitionReaderWithReadManyByPartitionKey.scala +++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/ItemsPartitionReaderWithReadManyByPartitionKey.scala @@ -170,41 +170,66 @@ private[spark] case class ItemsPartitionReaderWithReadManyByPartitionKey .enable(true) .build + private trait CloseableSparkRowItemIterator { + def hasNext: Boolean + def next(): SparkRowItem + def close(): Unit + } + + private object EmptySparkRowItemIterator extends CloseableSparkRowItemIterator { + override def hasNext: Boolean = false + + override def next(): SparkRowItem = { + throw new java.util.NoSuchElementException("No items available for empty partition-key list.") + } + + override def close(): Unit = {} + } + // Single iterator over all PKs — the SDK handles per-physical-partition batching // internally to avoid oversized SQL queries. - private lazy val iterator = new TransientIOErrorsRetryingIterator[SparkRowItem]( - continuationToken => { - val options = new CosmosReadManyRequestOptions() - val optionsImpl = ImplementationBridgeHelpers - .CosmosReadManyRequestOptionsHelper - .getCosmosReadManyRequestOptionsAccessor - .getImpl(options) - - ThroughputControlHelper.populateThroughputControlGroupName(optionsImpl, readConfig.throughputControlConfig) - - if (operationContextAndListenerTuple.isDefined) { - optionsImpl.setOperationContextAndListenerTuple(operationContextAndListenerTuple.get) - } - - optionsImpl.setCustomItemSerializer(readManyOptionsImpl.getCustomItemSerializer) - - if (pkList.isEmpty) { - cosmosAsyncContainer.readManyByPartitionKey( - new java.util.ArrayList[PartitionKey](), options, classOf[SparkRowItem]) - } else { - readConfig.customQuery match { - case Some(query) => - cosmosAsyncContainer.readManyByPartitionKey(pkList, query.toSqlQuerySpec, options, classOf[SparkRowItem]) - case None => - cosmosAsyncContainer.readManyByPartitionKey(pkList, options, classOf[SparkRowItem]) - } + // Short-circuit empty PK lists locally because the SDK rejects empty partition-key lists. + private lazy val iterator: CloseableSparkRowItemIterator = + if (pkList.isEmpty) { + EmptySparkRowItemIterator + } else { + new CloseableSparkRowItemIterator { + private val delegate = new TransientIOErrorsRetryingIterator[SparkRowItem]( + continuationToken => { + val options = new CosmosReadManyRequestOptions() + val optionsImpl = ImplementationBridgeHelpers + .CosmosReadManyRequestOptionsHelper + .getCosmosReadManyRequestOptionsAccessor + .getImpl(options) + + ThroughputControlHelper.populateThroughputControlGroupName(optionsImpl, readConfig.throughputControlConfig) + + if (operationContextAndListenerTuple.isDefined) { + optionsImpl.setOperationContextAndListenerTuple(operationContextAndListenerTuple.get) + } + + optionsImpl.setCustomItemSerializer(readManyOptionsImpl.getCustomItemSerializer) + + readConfig.customQuery match { + case Some(query) => + cosmosAsyncContainer.readManyByPartitionKey(pkList, query.toSqlQuerySpec, options, classOf[SparkRowItem]) + case None => + cosmosAsyncContainer.readManyByPartitionKey(pkList, options, classOf[SparkRowItem]) + } + }, + readConfig.maxItemCount, + readConfig.prefetchBufferSize, + operationContextAndListenerTuple, + None + ) + + override def hasNext: Boolean = delegate.hasNext + + override def next(): SparkRowItem = delegate.next() + + override def close(): Unit = delegate.close() } - }, - readConfig.maxItemCount, - readConfig.prefetchBufferSize, - operationContextAndListenerTuple, - None - ) + } private val rowSerializer: ExpressionEncoder.Serializer[Row] = RowSerializerPool.getOrCreateSerializer(readSchema) From d499da76fb44640c0405a1ac71195aa18781c27c Mon Sep 17 00:00:00 2001 From: Fabian Meiswinkel Date: Tue, 14 Apr 2026 16:51:24 +0200 Subject: [PATCH 04/64] Update sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- .../com/azure/cosmos/implementation/RxDocumentClientImpl.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java index 17461fa56531..db667ddeb420 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java @@ -4508,7 +4508,7 @@ public Flux> readManyByPartitionKey( int fluxConcurrency = Math.min(queryFluxes.size(), Math.max(Configs.getCPUCnt(), 1)); - return Flux.mergeSequential(queryFluxes, fluxConcurrency, 1); + return Flux.merge(queryFluxes, fluxConcurrency, 1); }); }); } From c3c542a33a79c189caff739f4412969c93c25b38 Mon Sep 17 00:00:00 2001 From: Fabian Meiswinkel Date: Tue, 14 Apr 2026 16:52:35 +0200 Subject: [PATCH 05/64] Update sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelper.java Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- .../cosmos/implementation/ReadManyByPartitionKeyQueryHelper.java | 1 - 1 file changed, 1 deletion(-) diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelper.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelper.java index fbe329bff844..b8f9fa4ecc81 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelper.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelper.java @@ -4,7 +4,6 @@ import com.azure.cosmos.BridgeInternal; import com.azure.cosmos.implementation.routing.PartitionKeyInternal; -import com.azure.cosmos.models.ModelBridgeInternal; import com.azure.cosmos.models.PartitionKey; import com.azure.cosmos.models.PartitionKeyDefinition; import com.azure.cosmos.models.PartitionKind; From 4416354e03e8973ae6a9195b05d2859301bd1561 Mon Sep 17 00:00:00 2001 From: Fabian Meiswinkel Date: Tue, 14 Apr 2026 20:50:32 +0000 Subject: [PATCH 06/64] =?UTF-8?q?=C2=B4Fixing=20code=20review=20comments?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- ...ReadManyByPartitionKeyQueryHelperTest.java | 34 +++++++++---------- .../ReadManyByPartitionKeyQueryHelper.java | 10 +++--- .../implementation/RxDocumentClientImpl.java | 2 +- 3 files changed, 24 insertions(+), 22 deletions(-) diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelperTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelperTest.java index 9c68b3b126e4..4f82db50bb60 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelperTest.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelperTest.java @@ -37,7 +37,7 @@ public void singlePk_defaultQuery_singleValue() { assertThat(result.getQueryText()).contains("SELECT * FROM c WHERE"); assertThat(result.getQueryText()).contains("IN ("); - assertThat(result.getQueryText()).contains("@pkParam0"); + assertThat(result.getQueryText()).contains("@__rmPk_0"); assertThat(result.getParameters()).hasSize(1); assertThat(result.getParameters().get(0).getValue(Object.class)).isEqualTo("pk1"); } @@ -55,9 +55,9 @@ public void singlePk_defaultQuery_multipleValues() { "SELECT * FROM c", new ArrayList<>(), pkValues, selectors, pkDef); assertThat(result.getQueryText()).contains("IN ("); - assertThat(result.getQueryText()).contains("@pkParam0"); - assertThat(result.getQueryText()).contains("@pkParam1"); - assertThat(result.getQueryText()).contains("@pkParam2"); + assertThat(result.getQueryText()).contains("@__rmPk_0"); + assertThat(result.getQueryText()).contains("@__rmPk_1"); + assertThat(result.getQueryText()).contains("@__rmPk_2"); assertThat(result.getParameters()).hasSize(3); } @@ -89,7 +89,7 @@ public void singlePk_customQuery_withExistingWhere() { // Should AND the PK filter to the existing WHERE clause assertThat(result.getQueryText()).contains("WHERE (c.age > @minAge) AND ("); assertThat(result.getQueryText()).contains("IN ("); - assertThat(result.getParameters()).hasSize(2); // @minAge + @pkParam1 + assertThat(result.getParameters()).hasSize(2); // @minAge + @__rmPk_0 assertThat(result.getParameters().get(0).getName()).isEqualTo("@minAge"); } @@ -111,9 +111,9 @@ public void hpk_fullPk_defaultQuery() { assertThat(result.getQueryText()).contains("SELECT * FROM c WHERE"); // Should use OR/AND pattern, not IN assertThat(result.getQueryText()).doesNotContain("IN ("); - assertThat(result.getQueryText()).contains("c[\"city\"] = @pkParam0"); + assertThat(result.getQueryText()).contains("c[\"city\"] = @__rmPk_0"); assertThat(result.getQueryText()).contains("AND"); - assertThat(result.getQueryText()).contains("c[\"zipcode\"] = @pkParam1"); + assertThat(result.getQueryText()).contains("c[\"zipcode\"] = @__rmPk_1"); assertThat(result.getParameters()).hasSize(2); assertThat(result.getParameters().get(0).getValue(Object.class)).isEqualTo("Redmond"); assertThat(result.getParameters().get(1).getValue(Object.class)).isEqualTo("98052"); @@ -132,10 +132,10 @@ public void hpk_fullPk_multipleValues() { "SELECT * FROM c", new ArrayList<>(), pkValues, selectors, pkDef); assertThat(result.getQueryText()).contains("OR"); - assertThat(result.getQueryText()).contains("c[\"city\"] = @pkParam0"); - assertThat(result.getQueryText()).contains("c[\"zipcode\"] = @pkParam1"); - assertThat(result.getQueryText()).contains("c[\"city\"] = @pkParam2"); - assertThat(result.getQueryText()).contains("c[\"zipcode\"] = @pkParam3"); + assertThat(result.getQueryText()).contains("c[\"city\"] = @__rmPk_0"); + assertThat(result.getQueryText()).contains("c[\"zipcode\"] = @__rmPk_1"); + assertThat(result.getQueryText()).contains("c[\"city\"] = @__rmPk_2"); + assertThat(result.getQueryText()).contains("c[\"zipcode\"] = @__rmPk_3"); assertThat(result.getParameters()).hasSize(4); } @@ -151,7 +151,7 @@ public void hpk_partialPk_singleLevel() { SqlQuerySpec result = ReadManyByPartitionKeyQueryHelper.createReadManyByPkQuerySpec( "SELECT * FROM c", new ArrayList<>(), pkValues, selectors, pkDef); - assertThat(result.getQueryText()).contains("c[\"city\"] = @pkParam0"); + assertThat(result.getQueryText()).contains("c[\"city\"] = @__rmPk_0"); // Should NOT include zipcode or areaCode since it's partial assertThat(result.getQueryText()).doesNotContain("zipcode"); assertThat(result.getQueryText()).doesNotContain("areaCode"); @@ -170,8 +170,8 @@ public void hpk_partialPk_twoLevels() { SqlQuerySpec result = ReadManyByPartitionKeyQueryHelper.createReadManyByPkQuerySpec( "SELECT * FROM c", new ArrayList<>(), pkValues, selectors, pkDef); - assertThat(result.getQueryText()).contains("c[\"city\"] = @pkParam0"); - assertThat(result.getQueryText()).contains("c[\"zipcode\"] = @pkParam1"); + assertThat(result.getQueryText()).contains("c[\"city\"] = @__rmPk_0"); + assertThat(result.getQueryText()).contains("c[\"zipcode\"] = @__rmPk_1"); assertThat(result.getQueryText()).doesNotContain("areaCode"); assertThat(result.getParameters()).hasSize(2); } @@ -191,7 +191,7 @@ public void hpk_customQuery_withWhere() { "SELECT c.name FROM c WHERE c.status = @status", baseParams, pkValues, selectors, pkDef); assertThat(result.getQueryText()).contains("WHERE (c.status = @status) AND ("); - assertThat(result.getQueryText()).contains("c[\"city\"] = @pkParam1"); + assertThat(result.getQueryText()).contains("c[\"city\"] = @__rmPk_0"); assertThat(result.getParameters()).hasSize(3); // @status + 2 pk params } @@ -277,8 +277,8 @@ public void hpk_customAlias() { SqlQuerySpec result = ReadManyByPartitionKeyQueryHelper.createReadManyByPkQuerySpec( "SELECT r.name FROM root r", new ArrayList<>(), pkValues, selectors, pkDef); - assertThat(result.getQueryText()).contains("r[\"city\"] = @pkParam0"); - assertThat(result.getQueryText()).contains("r[\"zipcode\"] = @pkParam1"); + assertThat(result.getQueryText()).contains("r[\"city\"] = @__rmPk_0"); + assertThat(result.getQueryText()).contains("r[\"zipcode\"] = @__rmPk_1"); assertThat(result.getQueryText()).doesNotContain("c[\""); } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelper.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelper.java index b8f9fa4ecc81..36880eb99b59 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelper.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelper.java @@ -20,6 +20,8 @@ public class ReadManyByPartitionKeyQueryHelper { private static final String DEFAULT_TABLE_ALIAS = "c"; + // Internal parameter prefix — uses double-underscore to avoid collisions with user-provided parameters + private static final String PK_PARAM_PREFIX = "@__rmPk_"; public static SqlQuerySpec createReadManyByPkQuerySpec( String baseQueryText, @@ -33,12 +35,12 @@ public static SqlQuerySpec createReadManyByPkQuerySpec( StringBuilder pkFilter = new StringBuilder(); List parameters = new ArrayList<>(baseParameters); - int paramCount = baseParameters.size(); + int paramCount = 0; boolean isSinglePathPk = partitionKeySelectors.size() == 1; if (isSinglePathPk && pkDefinition.getKind() != PartitionKind.MULTI_HASH) { - // Single PK path — use IN clause: alias["pkPath"] IN (@pk0, @pk1, ...) + // Single PK path — use IN clause: alias["pkPath"] IN (@__rmPk_0, @__rmPk_1, ...) pkFilter.append(" "); pkFilter.append(tableAlias); pkFilter.append(partitionKeySelectors.get(0)); @@ -46,7 +48,7 @@ public static SqlQuerySpec createReadManyByPkQuerySpec( for (int i = 0; i < pkValues.size(); i++) { PartitionKeyInternal pkInternal = BridgeInternal.getPartitionKeyInternal(pkValues.get(i)); Object[] pkComponents = pkInternal.toObjectArray(); - String pkParamName = "@pkParam" + paramCount; + String pkParamName = PK_PARAM_PREFIX + paramCount; parameters.add(new SqlParameter(pkParamName, pkComponents[0])); paramCount++; @@ -65,7 +67,7 @@ public static SqlQuerySpec createReadManyByPkQuerySpec( pkFilter.append("("); for (int j = 0; j < pkComponents.length; j++) { - String pkParamName = "@pkParam" + paramCount; + String pkParamName = PK_PARAM_PREFIX + paramCount; parameters.add(new SqlParameter(pkParamName, pkComponents[j])); paramCount++; diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java index db667ddeb420..5c276fb18e7e 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java @@ -4508,7 +4508,7 @@ public Flux> readManyByPartitionKey( int fluxConcurrency = Math.min(queryFluxes.size(), Math.max(Configs.getCPUCnt(), 1)); - return Flux.merge(queryFluxes, fluxConcurrency, 1); + return Flux.merge(Flux.fromIterable(queryFluxes), fluxConcurrency, 1); }); }); } From 588a7550c545717f35bdb76e5e0fead1d38ac3b8 Mon Sep 17 00:00:00 2001 From: Fabian Meiswinkel Date: Wed, 15 Apr 2026 09:49:34 +0000 Subject: [PATCH 07/64] Update CosmosAsyncContainer.java --- .../main/java/com/azure/cosmos/CosmosAsyncContainer.java | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java index b0258ec1b561..42e5dfeb0f17 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java @@ -1680,10 +1680,10 @@ private Function>> readManyByPa return (pagedFluxOptions -> { CosmosQueryRequestOptions queryRequestOptions = requestOptions == null ? new CosmosQueryRequestOptions() - : queryOptionsAccessor.clone(readManyOptionsAccessor.getImpl(requestOptions)); + : queryOptionsAccessor().clone(readManyOptionsAccessor().getImpl(requestOptions)); queryRequestOptions.setMaxDegreeOfParallelism(-1); queryRequestOptions.setQueryName("readManyByPartitionKey"); - CosmosQueryRequestOptionsBase cosmosQueryRequestOptionsImpl = queryOptionsAccessor.getImpl(queryRequestOptions); + CosmosQueryRequestOptionsBase cosmosQueryRequestOptionsImpl = queryOptionsAccessor().getImpl(queryRequestOptions); applyPolicies(OperationType.Query, ResourceType.Document, cosmosQueryRequestOptionsImpl, this.readManyItemsSpanName); QueryFeedOperationState state = new QueryFeedOperationState( @@ -1693,7 +1693,7 @@ private Function>> readManyByPa this.getId(), ResourceType.Document, OperationType.Query, - queryOptionsAccessor.getQueryNameOrDefault(queryRequestOptions, this.readManyItemsSpanName), + queryOptionsAccessor().getQueryNameOrDefault(queryRequestOptions, this.readManyItemsSpanName), queryRequestOptions, pagedFluxOptions ); From f5485527a9c45892c58945cee66d02f5bab463b4 Mon Sep 17 00:00:00 2001 From: Fabian Meiswinkel Date: Wed, 15 Apr 2026 12:04:49 +0000 Subject: [PATCH 08/64] Update ReadManyByPartitionKeyTest.java --- .../azure/cosmos/ReadManyByPartitionKeyTest.java | 16 +++++++++++++++- 1 file changed, 15 insertions(+), 1 deletion(-) diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/ReadManyByPartitionKeyTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/ReadManyByPartitionKeyTest.java index 8fa657b62105..a5cb16dc2715 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/ReadManyByPartitionKeyTest.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/ReadManyByPartitionKeyTest.java @@ -294,7 +294,7 @@ public void rejectsDistinctQuery() { @Test(groups = {"emulator"}, timeOut = TIMEOUT) public void rejectsGroupByQuery() { List pkValues = Collections.singletonList(new PartitionKey("pk1")); - SqlQuerySpec groupByQuery = new SqlQuerySpec("SELECT c.mypk, COUNT(1) as cnt FROM c GROUP BY c.mypk"); + SqlQuerySpec groupByQuery = new SqlQuerySpec("SELECT c.mypk FROM c GROUP BY c.mypk"); try { singlePkContainer.readManyByPartitionKey(pkValues, groupByQuery, null, ObjectNode.class) @@ -305,6 +305,20 @@ public void rejectsGroupByQuery() { } } + @Test(groups = {"emulator"}, timeOut = TIMEOUT) + public void rejectsGroupByWithAggregateQuery() { + List pkValues = Collections.singletonList(new PartitionKey("pk1")); + SqlQuerySpec groupByWithAggregateQuery = new SqlQuerySpec("SELECT c.mypk, COUNT(1) as cnt FROM c GROUP BY c.mypk"); + + try { + singlePkContainer.readManyByPartitionKey(pkValues, groupByWithAggregateQuery, null, ObjectNode.class) + .stream().collect(Collectors.toList()); + fail("Should have thrown IllegalArgumentException for GROUP BY with aggregate query"); + } catch (IllegalArgumentException e) { + assertThat(e.getMessage()).contains("aggregates"); + } + } + @Test(groups = {"emulator"}, timeOut = TIMEOUT, expectedExceptions = IllegalArgumentException.class) public void rejectsNullPartitionKeyList() { singlePkContainer.readManyByPartitionKey((List) null, ObjectNode.class); From f68cf02ff71cb784ccf6a5668589d1b64d4831a9 Mon Sep 17 00:00:00 2001 From: Fabian Meiswinkel Date: Wed, 15 Apr 2026 20:08:33 +0000 Subject: [PATCH 09/64] Fixing test issues --- .../java/com/azure/cosmos/ReadManyByPartitionKeyTest.java | 2 +- .../azure/cosmos/implementation/RxDocumentClientImpl.java | 8 ++++---- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/ReadManyByPartitionKeyTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/ReadManyByPartitionKeyTest.java index a5cb16dc2715..da76d795f981 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/ReadManyByPartitionKeyTest.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/ReadManyByPartitionKeyTest.java @@ -315,7 +315,7 @@ public void rejectsGroupByWithAggregateQuery() { .stream().collect(Collectors.toList()); fail("Should have thrown IllegalArgumentException for GROUP BY with aggregate query"); } catch (IllegalArgumentException e) { - assertThat(e.getMessage()).contains("aggregates"); + assertThat(e.getMessage()).contains("GROUP BY"); } } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java index dbd0d5553278..26986fb1dd8c 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java @@ -4524,6 +4524,10 @@ private Mono validateCustomQueryForReadManyByPartitionKey( .flatMap(queryPlan -> { QueryInfo queryInfo = queryPlan.getQueryInfo(); + if (queryInfo.hasGroupBy()) { + return Mono.error(new IllegalArgumentException( + "Custom query for readMany by partition key must not contain GROUP BY.")); + } if (queryInfo.hasAggregates()) { return Mono.error(new IllegalArgumentException( "Custom query for readMany by partition key must not contain aggregates.")); @@ -4536,10 +4540,6 @@ private Mono validateCustomQueryForReadManyByPartitionKey( return Mono.error(new IllegalArgumentException( "Custom query for readMany by partition key must not contain DISTINCT.")); } - if (queryInfo.hasGroupBy()) { - return Mono.error(new IllegalArgumentException( - "Custom query for readMany by partition key must not contain GROUP BY.")); - } if (queryInfo.hasDCount()) { return Mono.error(new IllegalArgumentException( "Custom query for readMany by partition key must not contain DCOUNT.")); From 8b6c4b168ea2c4b5e538524a9b8419513dce9693 Mon Sep 17 00:00:00 2001 From: Fabian Meiswinkel Date: Wed, 15 Apr 2026 22:36:29 +0000 Subject: [PATCH 10/64] Update CosmosAsyncContainer.java --- .../main/java/com/azure/cosmos/CosmosAsyncContainer.java | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java index 42e5dfeb0f17..79579aaf1715 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java @@ -1665,6 +1665,13 @@ public CosmosPagedFlux readManyByPartitionKey( CosmosReadManyRequestOptions requestOptions, Class classType) { + if (partitionKeys == null) { + throw new IllegalArgumentException("Argument 'partitionKeys' must not be null."); + } + if (partitionKeys.isEmpty()) { + throw new IllegalArgumentException("Argument 'partitionKeys' must not be empty."); + } + return UtilBridgeInternal.createCosmosPagedFlux( readManyByPartitionKeyInternalFunc(partitionKeys, customQuery, requestOptions, classType)); } From 56b067a93391c061108a3b8ba8bf1d5db3de4a91 Mon Sep 17 00:00:00 2001 From: Fabian Meiswinkel Date: Thu, 16 Apr 2026 00:49:47 +0000 Subject: [PATCH 11/64] Reacted to code review feedback --- .../cosmos/spark/CosmosItemsDataSource.scala | 8 +- .../spark/CosmosPartitionKeyHelper.scala | 7 +- .../CosmosReadManyByPartitionKeyReader.scala | 3 - ...tionReaderWithReadManyByPartitionKey.scala | 18 +-- .../spark/CosmosPartitionKeyHelperSpec.scala | 12 ++ .../cosmos/ReadManyByPartitionKeyTest.java | 67 ++++++++++ ...ReadManyByPartitionKeyQueryHelperTest.java | 77 +++++++++++ .../azure/cosmos/CosmosAsyncContainer.java | 8 +- .../ReadManyByPartitionKeyQueryHelper.java | 19 ++- .../implementation/RxDocumentClientImpl.java | 7 +- .../docs/readManyByPartitionKey-design.md | 124 +++++++++++------- 11 files changed, 269 insertions(+), 81 deletions(-) diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosItemsDataSource.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosItemsDataSource.scala index 2fbc036724e7..6257e96e81e5 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosItemsDataSource.scala +++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosItemsDataSource.scala @@ -147,7 +147,7 @@ object CosmosItemsDataSource { val sparkEnvironmentInfo = CosmosClientConfiguration.getSparkEnvironmentInfo(None) val calledFrom = s"CosmosItemsDataSource.readManyByPartitionKey" - val pkPathsOpt = Loan( + val pkPaths = Loan( List[Option[CosmosClientCacheItem]]( Some( CosmosClientCache( @@ -180,10 +180,10 @@ object CosmosItemsDataSource { // Check if ALL PK path columns exist in the DataFrame schema val dfFieldNames = df.schema.fieldNames.toSet - val allPkColumnsPresent = pkPathsOpt.forall(path => dfFieldNames.contains(path)) + val allPkColumnsPresent = pkPaths.forall(path => dfFieldNames.contains(path)) - if (allPkColumnsPresent && pkPathsOpt.nonEmpty) { - val pkPaths = pkPathsOpt + if (allPkColumnsPresent && pkPaths.nonEmpty) { + // pkPaths already defined above Some((row: Row) => { if (pkPaths.size == 1) { // Single partition key diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosPartitionKeyHelper.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosPartitionKeyHelper.scala index 616e1893b343..27776f5c3de6 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosPartitionKeyHelper.scala +++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosPartitionKeyHelper.scala @@ -30,15 +30,14 @@ private[spark] object CosmosPartitionKeyHelper extends BasicLoggingTrait { def tryParsePartitionKey(cosmosPartitionKeyString: String): Option[PartitionKey] = { cosmosPartitionKeyString match { case cosmosPartitionKeyStringRegx(pkValue) => - val partitionKeyValue = Utils.parse(pkValue, classOf[Object]) - partitionKeyValue match { - case arrayList: util.ArrayList[Object] => + scala.util.Try(Utils.parse(pkValue, classOf[Object])).toOption.flatMap { + case arrayList: util.ArrayList[Object @unchecked] => Some( ImplementationBridgeHelpers .PartitionKeyHelper .getPartitionKeyAccessor .toPartitionKey(PartitionKeyInternal.fromObjectArray(arrayList.toArray, false))) - case _ => Some(new PartitionKey(partitionKeyValue)) + case other => Some(new PartitionKey(other)) } case _ => None } diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosReadManyByPartitionKeyReader.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosReadManyByPartitionKeyReader.scala index 1d324d4855ff..91f3a56bc664 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosReadManyByPartitionKeyReader.scala +++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosReadManyByPartitionKeyReader.scala @@ -148,6 +148,3 @@ private[spark] class CosmosReadManyByPartitionKeyReader( } } -private object CosmosReadManyByPartitionKeyHelper { - val FullRangeFeedRange: NormalizedRange = NormalizedRange("", "FF") -} diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/ItemsPartitionReaderWithReadManyByPartitionKey.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/ItemsPartitionReaderWithReadManyByPartitionKey.scala index dddbb23a0050..da3b81d951ae 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/ItemsPartitionReaderWithReadManyByPartitionKey.scala +++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/ItemsPartitionReaderWithReadManyByPartitionKey.scala @@ -196,25 +196,11 @@ private[spark] case class ItemsPartitionReaderWithReadManyByPartitionKey new CloseableSparkRowItemIterator { private val delegate = new TransientIOErrorsRetryingIterator[SparkRowItem]( continuationToken => { - val options = new CosmosReadManyRequestOptions() - val optionsImpl = ImplementationBridgeHelpers - .CosmosReadManyRequestOptionsHelper - .getCosmosReadManyRequestOptionsAccessor - .getImpl(options) - - ThroughputControlHelper.populateThroughputControlGroupName(optionsImpl, readConfig.throughputControlConfig) - - if (operationContextAndListenerTuple.isDefined) { - optionsImpl.setOperationContextAndListenerTuple(operationContextAndListenerTuple.get) - } - - optionsImpl.setCustomItemSerializer(readManyOptionsImpl.getCustomItemSerializer) - readConfig.customQuery match { case Some(query) => - cosmosAsyncContainer.readManyByPartitionKey(pkList, query.toSqlQuerySpec, options, classOf[SparkRowItem]) + cosmosAsyncContainer.readManyByPartitionKey(pkList, query.toSqlQuerySpec, readManyOptions, classOf[SparkRowItem]) case None => - cosmosAsyncContainer.readManyByPartitionKey(pkList, options, classOf[SparkRowItem]) + cosmosAsyncContainer.readManyByPartitionKey(pkList, readManyOptions, classOf[SparkRowItem]) } }, readConfig.maxItemCount, diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/CosmosPartitionKeyHelperSpec.scala b/sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/CosmosPartitionKeyHelperSpec.scala index d127710da287..182f0c3cc3d3 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/CosmosPartitionKeyHelperSpec.scala +++ b/sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/CosmosPartitionKeyHelperSpec.scala @@ -76,6 +76,18 @@ class CosmosPartitionKeyHelperSpec extends UnitSpec { pk.get shouldEqual new PartitionKey("value") } + + it should "return None for malformed JSON inside pk() wrapper" in { + // Invalid JSON that would cause JsonProcessingException + val pk = CosmosPartitionKeyHelper.tryParsePartitionKey("pk({invalid json})") + pk.isDefined shouldBe false + } + + it should "return None for truncated JSON inside pk() wrapper" in { + val pk = CosmosPartitionKeyHelper.tryParsePartitionKey("pk(["unterminated)") + pk.isDefined shouldBe false + } + //scalastyle:on multiple.string.literals //scalastyle:on magic.number } diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/ReadManyByPartitionKeyTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/ReadManyByPartitionKeyTest.java index da76d795f981..2c26d564ed24 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/ReadManyByPartitionKeyTest.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/ReadManyByPartitionKeyTest.java @@ -332,6 +332,73 @@ public void rejectsEmptyPartitionKeyList() { //endregion + + //region Batch size tests (#10) + + @Test(groups = {"emulator"}, timeOut = TIMEOUT) + public void singlePk_readManyByPartitionKey_withSmallBatchSize() { + // Temporarily set batch size to 2 to exercise the batching/interleaving logic + String originalValue = System.getProperty("COSMOS.READ_MANY_BY_PK_MAX_BATCH_SIZE"); + try { + System.setProperty("COSMOS.READ_MANY_BY_PK_MAX_BATCH_SIZE", "2"); + + // Create items across 4 PKs (more than the batch size of 2) + List items = createSinglePkItems("batchPk1", 2); + items.addAll(createSinglePkItems("batchPk2", 2)); + items.addAll(createSinglePkItems("batchPk3", 2)); + items.addAll(createSinglePkItems("batchPk4", 2)); + + // Read all 4 PKs — should be split into batches of 2 + List pkValues = Arrays.asList( + new PartitionKey("batchPk1"), + new PartitionKey("batchPk2"), + new PartitionKey("batchPk3"), + new PartitionKey("batchPk4")); + + CosmosPagedIterable results = singlePkContainer.readManyByPartitionKey(pkValues, ObjectNode.class); + List resultList = results.stream().collect(Collectors.toList()); + + assertThat(resultList).hasSize(8); // 2 items per PK * 4 PKs + resultList.forEach(item -> { + String pk = item.get("mypk").asText(); + assertThat(pk).isIn("batchPk1", "batchPk2", "batchPk3", "batchPk4"); + }); + + cleanupContainer(singlePkContainer); + } finally { + if (originalValue != null) { + System.setProperty("COSMOS.READ_MANY_BY_PK_MAX_BATCH_SIZE", originalValue); + } else { + System.clearProperty("COSMOS.READ_MANY_BY_PK_MAX_BATCH_SIZE"); + } + } + } + + //endregion + + //region Custom serializer regression tests (#5) + + @Test(groups = {"emulator"}, timeOut = TIMEOUT) + public void singlePk_readManyByPartitionKey_withRequestOptions() { + // This test ensures that request options (like throughput control settings) + // are properly propagated through the readManyByPartitionKey path. + // It acts as a regression test for the redundant options construction bug. + List items = createSinglePkItems("pkOpts", 3); + + List pkValues = Collections.singletonList(new PartitionKey("pkOpts")); + com.azure.cosmos.models.CosmosReadManyRequestOptions options = new com.azure.cosmos.models.CosmosReadManyRequestOptions(); + + CosmosPagedIterable results = singlePkContainer.readManyByPartitionKey( + pkValues, options, ObjectNode.class); + List resultList = results.stream().collect(Collectors.toList()); + + assertThat(resultList).hasSize(3); + + cleanupContainer(singlePkContainer); + } + + //endregion + //region helper methods private List createSinglePkItems(String pkValue, int count) { diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelperTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelperTest.java index 4f82db50bb60..95c109ba025f 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelperTest.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelperTest.java @@ -318,6 +318,83 @@ public void extractAlias_caseInsensitive() { //endregion + + //region String literal handling tests (#1) + + @Test(groups = { "unit" }) + public void findWhere_ignoresWhereInsideStringLiteral() { + // WHERE inside a string literal should be ignored + int idx = ReadManyByPartitionKeyQueryHelper.findTopLevelWhereIndex( + "SELECT * FROM c WHERE c.msg = 'use WHERE clause here'"); + // Should find the outer WHERE at position 16, not the one inside the string + assertThat(idx).isEqualTo(16); + } + + @Test(groups = { "unit" }) + public void findWhere_ignoresParenthesesInsideStringLiteral() { + // Parentheses inside string literal should not affect depth tracking + int idx = ReadManyByPartitionKeyQueryHelper.findTopLevelWhereIndex( + "SELECT * FROM c WHERE c.name = 'foo(bar)' AND c.x = 1"); + assertThat(idx).isEqualTo(16); + } + + @Test(groups = { "unit" }) + public void findWhere_handlesUnbalancedParenInStringLiteral() { + // Unbalanced paren inside string literal must not corrupt depth + int idx = ReadManyByPartitionKeyQueryHelper.findTopLevelWhereIndex( + "SELECT * FROM c WHERE c.val = 'open(' AND c.active = true"); + assertThat(idx).isEqualTo(16); + } + + @Test(groups = { "unit" }) + public void findWhere_handlesStringLiteralBeforeWhere() { + // String literal in SELECT before WHERE + int idx = ReadManyByPartitionKeyQueryHelper.findTopLevelWhereIndex( + "SELECT 'WHERE' as label FROM c WHERE c.id = '1'"); + // The WHERE inside quotes should be ignored; the real WHERE is further along + assertThat(idx).isGreaterThan(30); + } + + @Test(groups = { "unit" }) + public void singlePk_customQuery_withStringLiteralContainingParens() { + PartitionKeyDefinition pkDef = createSinglePkDefinition("/mypk"); + List selectors = createSelectors(pkDef); + List pkValues = Collections.singletonList(new PartitionKey("pk1")); + + List baseParams = new ArrayList<>(); + baseParams.add(new SqlParameter("@msg", "hello")); + + SqlQuerySpec result = ReadManyByPartitionKeyQueryHelper.createReadManyByPkQuerySpec( + "SELECT * FROM c WHERE c.msg = 'test(value)WHERE'", baseParams, pkValues, selectors, pkDef); + + // Should correctly AND the PK filter to the real WHERE clause + assertThat(result.getQueryText()).contains("WHERE (c.msg = 'test(value)WHERE') AND ("); + } + + //endregion + + //region OFFSET/LIMIT/HAVING alias detection tests (#9) + + @Test(groups = { "unit" }) + public void extractAlias_containerWithOffset() { + assertThat(ReadManyByPartitionKeyQueryHelper.extractTableAlias( + "SELECT * FROM c OFFSET 10 LIMIT 5")).isEqualTo("c"); + } + + @Test(groups = { "unit" }) + public void extractAlias_containerWithLimit() { + assertThat(ReadManyByPartitionKeyQueryHelper.extractTableAlias( + "SELECT * FROM c LIMIT 10")).isEqualTo("c"); + } + + @Test(groups = { "unit" }) + public void extractAlias_containerWithHaving() { + assertThat(ReadManyByPartitionKeyQueryHelper.extractTableAlias( + "SELECT c.status, COUNT(1) FROM c GROUP BY c.status HAVING COUNT(1) > 1")).isEqualTo("c"); + } + + //endregion + //region helpers private PartitionKeyDefinition createSinglePkDefinition(String path) { diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java index 79579aaf1715..025a957a4606 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java @@ -165,6 +165,7 @@ private static ImplementationBridgeHelpers.CosmosBatchRequestOptionsHelper.Cosmo private final String createItemSpanName; private final String readAllItemsSpanName; private final String readManyItemsSpanName; + private final String readManyByPartitionKeyItemsSpanName; private final String readAllItemsOfLogicalPartitionSpanName; private final String queryItemsSpanName; private final String queryChangeFeedSpanName; @@ -198,6 +199,7 @@ protected CosmosAsyncContainer(CosmosAsyncContainer toBeWrappedContainer) { this.createItemSpanName = "createItem." + this.id; this.readAllItemsSpanName = "readAllItems." + this.id; this.readManyItemsSpanName = "readManyItems." + this.id; + this.readManyByPartitionKeyItemsSpanName = "readManyByPartitionKeyItems." + this.id; this.readAllItemsOfLogicalPartitionSpanName = "readAllItemsOfLogicalPartition." + this.id; this.queryItemsSpanName = "queryItems." + this.id; this.queryChangeFeedSpanName = "queryChangeFeed." + this.id; @@ -1691,16 +1693,16 @@ private Function>> readManyByPa queryRequestOptions.setMaxDegreeOfParallelism(-1); queryRequestOptions.setQueryName("readManyByPartitionKey"); CosmosQueryRequestOptionsBase cosmosQueryRequestOptionsImpl = queryOptionsAccessor().getImpl(queryRequestOptions); - applyPolicies(OperationType.Query, ResourceType.Document, cosmosQueryRequestOptionsImpl, this.readManyItemsSpanName); + applyPolicies(OperationType.Query, ResourceType.Document, cosmosQueryRequestOptionsImpl, this.readManyByPartitionKeyItemsSpanName); QueryFeedOperationState state = new QueryFeedOperationState( client, - this.readManyItemsSpanName, + this.readManyByPartitionKeyItemsSpanName, database.getId(), this.getId(), ResourceType.Document, OperationType.Query, - queryOptionsAccessor().getQueryNameOrDefault(queryRequestOptions, this.readManyItemsSpanName), + queryOptionsAccessor().getQueryNameOrDefault(queryRequestOptions, this.readManyByPartitionKeyItemsSpanName), queryRequestOptions, pagedFluxOptions ); diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelper.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelper.java index 36880eb99b59..d538542df1ba 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelper.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelper.java @@ -140,12 +140,15 @@ static String extractTableAlias(String queryText) { // Check if there's an alias after the container name (before WHERE or end) if (afterFrom < queryText.length()) { char nextChar = Character.toUpperCase(queryText.charAt(afterFrom)); - // If the next token is a keyword (WHERE, ORDER, GROUP, JOIN) or end, containerName IS the alias - if (nextChar == 'W' || nextChar == 'O' || nextChar == 'G' || nextChar == 'J') { + // If the next token is a keyword (WHERE, ORDER, GROUP, JOIN, OFFSET, LIMIT, HAVING) or end, containerName IS the alias + if (nextChar == 'W' || nextChar == 'O' || nextChar == 'G' || nextChar == 'J' + || nextChar == 'L' || nextChar == 'H') { // Check if it's actually a keyword String remaining = upper.substring(afterFrom); if (remaining.startsWith("WHERE") || remaining.startsWith("ORDER") - || remaining.startsWith("GROUP") || remaining.startsWith("JOIN")) { + || remaining.startsWith("GROUP") || remaining.startsWith("JOIN") + || remaining.startsWith("OFFSET") || remaining.startsWith("LIMIT") + || remaining.startsWith("HAVING")) { return containerName; } } @@ -167,7 +170,7 @@ static String extractTableAlias(String queryText) { /** * Finds the index of a top-level SQL keyword in the query text (case-insensitive), - * ignoring occurrences inside parentheses. + * ignoring occurrences inside parentheses or string literals. */ static int findTopLevelKeywordIndex(String queryText, String keyword) { String queryTextUpper = queryText.toUpperCase(); @@ -176,6 +179,14 @@ static int findTopLevelKeywordIndex(String queryText, String keyword) { int keyLen = keywordUpper.length(); for (int i = 0; i <= queryTextUpper.length() - keyLen; i++) { char ch = queryTextUpper.charAt(i); + // Skip string literals enclosed in single quotes + if (queryText.charAt(i) == '\'') { + i++; + while (i < queryText.length() && queryText.charAt(i) != '\'') { + i++; + } + continue; + } if (ch == '(') { depth++; } else if (ch == ')') { diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java index 26986fb1dd8c..b4c72532280c 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java @@ -4417,9 +4417,10 @@ public Flux> readManyByPartitionKey( queryValidationMono = Mono.empty(); } - return Mono.zip(valueHolderMono, queryValidationMono.then(Mono.just(true))) - .flatMapMany(tuple -> { - CollectionRoutingMap routingMap = tuple.getT1().v; + return valueHolderMono + .delayUntil(ignored -> queryValidationMono) + .flatMapMany(routingMapHolder -> { + CollectionRoutingMap routingMap = routingMapHolder.v; if (routingMap == null) { return Flux.error(new IllegalStateException("Failed to get routing map.")); } diff --git a/sdk/cosmos/docs/readManyByPartitionKey-design.md b/sdk/cosmos/docs/readManyByPartitionKey-design.md index f53cde1db37a..95d7624f0c8b 100644 --- a/sdk/cosmos/docs/readManyByPartitionKey-design.md +++ b/sdk/cosmos/docs/readManyByPartitionKey-design.md @@ -1,10 +1,10 @@ -# readMany by Partition Key — Design & Implementation Plan +# readManyByPartitionKey — Design & Implementation ## Overview -New `readMany` overloads on `CosmosAsyncContainer` / `CosmosContainer` that accept a +New `readManyByPartitionKey` methods on `CosmosAsyncContainer` / `CosmosContainer` that accept a `List` (without item-id). The SDK splits the PK values by physical -partition, generates a streaming query per physical partition, and returns results as +partition, generates batched streaming queries per physical partition, and returns results as `CosmosPagedFlux` / `CosmosPagedIterable`. An optional `SqlQuerySpec` parameter lets callers supply a custom query for projections @@ -14,31 +14,36 @@ and additional filters. The SDK appends the auto-generated PK WHERE clause to it | Topic | Decision | |---|---| -| API name | `readMany` — new overload distinguished by `List` parameter | +| API name | `readManyByPartitionKey` — distinct name to avoid ambiguity with existing `readMany(List)` | | Return type | `CosmosPagedFlux` (async) / `CosmosPagedIterable` (sync) | | Custom query format | `SqlQuerySpec` — full query with parameters; SDK ANDs the PK filter | | Partial HPK | Supported from the start; prefix PKs fan out via `getOverlappingRanges` | | PK deduplication | Done at Spark layer only, not in the SDK | | Spark UDF | New `GetCosmosPartitionKeyValue` UDF | -| Custom query validation | Gateway query plan; reject aggregates/ORDER BY/DISTINCT/GROUP BY/DCount/vector/fulltext | -| Max PK list size | Enforced per invocation (same effective cap as existing readMany) | +| Custom query validation | Gateway query plan; reject aggregates/ORDER BY/DISTINCT/GROUP BY/DCount/non-streaming ORDER BY/vector/fulltext | +| PK list size | No hard upper-bound enforced; SDK batches internally per physical partition (default 1000 PKs per batch, configurable via `COSMOS.READ_MANY_BY_PK_MAX_BATCH_SIZE`) | +| Eager validation | Null and empty PK list rejected eagerly (not lazily in reactive chain) | +| Telemetry | Separate span name `readManyByPartitionKeyItems.` (distinct from existing `readManyItems`) | +| Query construction | Table alias auto-detected from FROM clause; string literals and subqueries handled correctly | ## Phase 1 — SDK Core (`azure-cosmos`) ### Step 1: New public overloads in CosmosAsyncContainer ```java - CosmosPagedFlux readMany(List partitionKeys, Class classType) - CosmosPagedFlux readMany(List partitionKeys, - CosmosReadManyRequestOptions requestOptions, - Class classType) - CosmosPagedFlux readMany(List partitionKeys, - SqlQuerySpec customQuery, - CosmosReadManyRequestOptions requestOptions, - Class classType) + CosmosPagedFlux readManyByPartitionKey(List partitionKeys, Class classType) + CosmosPagedFlux readManyByPartitionKey(List partitionKeys, + CosmosReadManyRequestOptions requestOptions, + Class classType) + CosmosPagedFlux readManyByPartitionKey(List partitionKeys, + SqlQuerySpec customQuery, + CosmosReadManyRequestOptions requestOptions, + Class classType) ``` -All delegate to a private `readManyByPartitionKeyInternal(...)`. +All delegate to a private `readManyByPartitionKeyInternalFunc(...)`. + +**Eager validation:** The 4-arg method validates `partitionKeys` is non-null and non-empty before constructing the reactive pipeline, throwing `IllegalArgumentException` synchronously. ### Step 2: Sync wrappers in CosmosContainer @@ -47,49 +52,56 @@ Same signatures returning `CosmosPagedIterable`, delegating to the async cont ### Step 3: Internal orchestration (RxDocumentClientImpl) 1. Resolve collection metadata + PK definition from cache. -2. Fetch routing map from `partitionKeyRangeCache`. +2. Fetch routing map from `partitionKeyRangeCache` **in parallel with** custom query validation (Step 4). 3. For each `PartitionKey`: - Compute effective partition key (EPK). - Full PK → `getRangeByEffectivePartitionKey()` (single range). - Partial HPK → compute EPK prefix range → `getOverlappingRanges()` (multiple ranges). **Note:** partial HPK intentionally fans out to multiple physical partitions. 4. Group PK values by `PartitionKeyRange`. -5. If custom `SqlQuerySpec` provided → validate via query plan (Step 4). -6. Per physical partition → build `SqlQuerySpec` with PK WHERE clause (Step 5). -7. Execute queries via `createReadManyQueryAsync()`. -8. Return results as `CosmosPagedFlux`. +5. Per physical partition → split PKs into batches of `maxPksPerPartitionQuery` (configurable, default 1000). +6. Per batch → build `SqlQuerySpec` with PK WHERE clause (Step 5). +7. Interleave batches across physical partitions in round-robin order so that bounded concurrency prefers different physical partitions over sequential batches of the same partition. +8. Execute queries via `queryForReadMany()` with bounded concurrency (`Math.min(batchCount, cpuCount)`). +9. Return results as `CosmosPagedFlux`. ### Step 4: Custom query validation -One-time call per invocation (existing query plan caching applies): +One-time call per invocation (existing query plan caching applies). Runs **in parallel** with routing map lookup to minimize latency: - `QueryPlanRetriever.getQueryPlanThroughGatewayAsync()` for the user query. - Reject (`IllegalArgumentException`) if: + - `queryInfo.hasGroupBy()` — checked first (takes precedence over aggregates since `hasAggregates()` also returns true for GROUP BY queries) - `queryInfo.hasAggregates()` - `queryInfo.hasOrderBy()` - `queryInfo.hasDistinct()` - - `queryInfo.hasGroupBy()` - `queryInfo.hasDCount()` - `queryInfo.hasNonStreamingOrderBy()` - `partitionedQueryExecutionInfo.hasHybridSearchQueryInfo()` ### Step 5: Query construction +Query construction is implemented in `ReadManyByPartitionKeyQueryHelper`. The helper: +- Extracts the table alias from the FROM clause (handles `FROM c`, `FROM root r`, `FROM x WHERE ...`) +- Handles string literals in queries (parens/keywords inside `'...'` are correctly skipped) +- Recognizes SQL keywords: WHERE, ORDER, GROUP, JOIN, OFFSET, LIMIT, HAVING +- Uses parameterized queries (`@__rmPk_` prefix) to prevent SQL injection + **Single PK (HASH):** ```sql -{baseQuery} WHERE c["{pkPath}"] IN (@pk0, @pk1, @pk2) +{baseQuery} WHERE {alias}["{pkPath}"] IN (@__rmPk_0, @__rmPk_1, @__rmPk_2) ``` **Full HPK (MULTI_HASH):** ```sql -{baseQuery} WHERE (c["{path1}"] = @p0l1 AND c["{path2}"] = @p0l2) - OR (c["{path1}"] = @p1l1 AND c["{path2}"] = @p1l2) +{baseQuery} WHERE ({alias}["{path1}"] = @__rmPk_0 AND {alias}["{path2}"] = @__rmPk_1) + OR ({alias}["{path1}"] = @__rmPk_2 AND {alias}["{path2}"] = @__rmPk_3) ``` **Partial HPK (prefix-only):** ```sql -{baseQuery} WHERE (c["{path1}"] = @p0l1) - OR (c["{path1}"] = @p1l1) +{baseQuery} WHERE ({alias}["{path1}"] = @__rmPk_0) + OR ({alias}["{path1}"] = @__rmPk_1) ``` If the base query already has a WHERE clause: @@ -97,37 +109,61 @@ If the base query already has a WHERE clause: {selectAndFrom} WHERE ({existingWhere}) AND ({pkFilter}) ``` -### Step 6: Bridge / accessor wiring +### Step 6: Interface wiring + +New method `readManyByPartitionKey` added directly to `AsyncDocumentClient` interface, implemented in `RxDocumentClientImpl`. New `fetchQueryPlanForValidation` static method added to `DocumentQueryExecutionContextFactory` for custom query validation. -Expose internal method through `ImplementationBridgeHelpers`. +### Step 7: Configuration + +New configurable batch size via system property `COSMOS.READ_MANY_BY_PK_MAX_BATCH_SIZE` or environment variable `COSMOS_READ_MANY_BY_PK_MAX_BATCH_SIZE` (default: 1000, minimum: 1). Follows existing `Configs` patterns. ## Phase 2 — Spark Connector (`azure-cosmos-spark_3`) -### Step 7: New UDF — `GetCosmosPartitionKeyValue` +### Step 8: New UDF — `GetCosmosPartitionKeyValue` -- Input: partition key column(s) as array. -- Output: serialized PK string. +- Input: partition key value (single value or Seq for hierarchical PKs). +- Output: serialized PK string in format `pk([...json...])`. +- **Null handling:** Throws on null input (Scala convention; callers should filter nulls upstream). -### Step 8: PK-only serialization helper +### Step 9: PK-only serialization helper `CosmosPartitionKeyHelper`: -- `getCosmosPartitionKeyValueString(pkValues)` — serialize. -- `tryParsePartitionKey(serialized)` — deserialize. +- `getCosmosPartitionKeyValueString(pkValues: List[Object]): String` — serialize to `pk([...])` format. +- `tryParsePartitionKey(serialized: String): Option[PartitionKey]` — deserialize; returns `None` for malformed input including invalid JSON (wrapped in `scala.util.Try`). + +### Step 10: `CosmosItemsDataSource.readManyByPartitionKey` -### Step 9: `CosmosItemsDataSource.readManyByPartitionKey` +Static entry points that accept a DataFrame and Cosmos config. PK extraction supports two modes: +1. **UDF-produced column**: DataFrame contains `_partitionKeyIdentity` column (from `GetCosmosPartitionKeyValue` UDF). +2. **Schema-matched columns**: DataFrame columns match the container's PK paths. -Static entry points, deduplicates PKs at Spark level, delegates to reader. +Falls back with `IllegalArgumentException` if neither mode is possible. -### Step 10: `CosmosReadManyByPartitionKeyReader` +### Step 11: `CosmosReadManyByPartitionKeyReader` -Per-Spark-partition execution, analogous to `CosmosReadManyReader`. +Orchestrator that resolves schema, initializes and broadcasts client state to executors, then maps each Spark partition to an `ItemsPartitionReaderWithReadManyByPartitionKey`. -### Step 11: `ItemsPartitionReaderWithReadManyByPartitionKey` +### Step 12: `ItemsPartitionReaderWithReadManyByPartitionKey` -Calls new SDK API with `Iterator[PartitionKey]`, iterates `CosmosPagedFlux` pages. +Spark `PartitionReader[InternalRow]` that: +- Deduplicates PKs via `LinkedHashMap` (by PK string representation). +- Passes the pre-built `CosmosReadManyRequestOptions` (with throughput control, diagnostics, custom serializer) to the SDK. +- Uses `TransientIOErrorsRetryingIterator` for retry handling. +- Short-circuits empty PK lists to avoid SDK rejection. ## Phase 3 — Testing -- Unit tests: query construction (single PK, HPK, partial HPK, custom query composition). -- Unit tests: query plan rejection (aggregates, ORDER BY, DISTINCT, etc.). -- Integration tests: end-to-end SDK + Spark UDF. +### Unit tests +- Query construction: single PK, HPK full/partial, custom query composition, table alias detection. +- Query plan rejection: aggregates, ORDER BY, DISTINCT, GROUP BY (with and without aggregates), DCOUNT. +- String literal handling: WHERE/parentheses inside string constants. +- Keyword detection: WHERE, ORDER, GROUP, JOIN, OFFSET, LIMIT, HAVING. +- PK serialization/deserialization roundtrip (including malformed JSON handling). +- `findTopLevelWhereIndex` edge cases: subqueries, string literals, case insensitivity. + +### Integration tests +- End-to-end SDK: single PK basic, projections, filters, empty results, HPK full/partial, request options propagation. +- Batch size validation: temporarily lowered batch size to exercise batching/interleaving logic. +- Null/empty PK list rejection (eager validation). +- Spark connector: `ItemsPartitionReaderWithReadManyByPartitionKey` with known PK values and non-existent PKs. +- `CosmosPartitionKeyHelper`: single/HPK roundtrip, case insensitivity, malformed input. From d9504c91f343d6b6c2e970586ce3c5f994c44dd8 Mon Sep 17 00:00:00 2001 From: Fabian Meiswinkel Date: Thu, 16 Apr 2026 10:56:51 +0000 Subject: [PATCH 12/64] Fix build issues --- .../azure/cosmos/spark/CosmosPartitionKeyHelperSpec.scala | 2 +- sdk/cosmos/cspell.yaml | 6 ++++++ 2 files changed, 7 insertions(+), 1 deletion(-) create mode 100644 sdk/cosmos/cspell.yaml diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/CosmosPartitionKeyHelperSpec.scala b/sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/CosmosPartitionKeyHelperSpec.scala index 182f0c3cc3d3..6528d44f5ebe 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/CosmosPartitionKeyHelperSpec.scala +++ b/sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/CosmosPartitionKeyHelperSpec.scala @@ -84,7 +84,7 @@ class CosmosPartitionKeyHelperSpec extends UnitSpec { } it should "return None for truncated JSON inside pk() wrapper" in { - val pk = CosmosPartitionKeyHelper.tryParsePartitionKey("pk(["unterminated)") + val pk = CosmosPartitionKeyHelper.tryParsePartitionKey("pk([\"unterminated)") pk.isDefined shouldBe false } diff --git a/sdk/cosmos/cspell.yaml b/sdk/cosmos/cspell.yaml new file mode 100644 index 000000000000..94a4002c2c9c --- /dev/null +++ b/sdk/cosmos/cspell.yaml @@ -0,0 +1,6 @@ +import: + - ../../.vscode/cspell.json +overrides: + - filename: "**/sdk/cosmos/*" + words: + - DCOUNT From 681830e2d4a134c1f72976cf6cdea22929f6a69e Mon Sep 17 00:00:00 2001 From: Fabian Meiswinkel Date: Thu, 16 Apr 2026 21:48:17 +0000 Subject: [PATCH 13/64] Fixing changelog --- sdk/cosmos/azure-cosmos-spark_3-3_2-12/CHANGELOG.md | 1 + sdk/cosmos/azure-cosmos-spark_3-4_2-12/CHANGELOG.md | 1 + sdk/cosmos/azure-cosmos-spark_3-5_2-12/CHANGELOG.md | 1 + sdk/cosmos/azure-cosmos-spark_3-5_2-13/CHANGELOG.md | 1 + sdk/cosmos/azure-cosmos-spark_4-0_2-13/CHANGELOG.md | 1 + sdk/cosmos/azure-cosmos/CHANGELOG.md | 1 + 6 files changed, 6 insertions(+) diff --git a/sdk/cosmos/azure-cosmos-spark_3-3_2-12/CHANGELOG.md b/sdk/cosmos/azure-cosmos-spark_3-3_2-12/CHANGELOG.md index cbf97c610f9f..fe114462019c 100644 --- a/sdk/cosmos/azure-cosmos-spark_3-3_2-12/CHANGELOG.md +++ b/sdk/cosmos/azure-cosmos-spark_3-3_2-12/CHANGELOG.md @@ -3,6 +3,7 @@ ### 4.47.0-beta.1 (Unreleased) #### Features Added +* Added new `CosmosItemsDataSource.readManyByPartitionKey` Spark function to execute bulk queries by a list of pk-values with better efficiency. See [PR 48801](https://github.com/Azure/azure-sdk-for-java/pull/48801) #### Breaking Changes diff --git a/sdk/cosmos/azure-cosmos-spark_3-4_2-12/CHANGELOG.md b/sdk/cosmos/azure-cosmos-spark_3-4_2-12/CHANGELOG.md index c9097e749f03..3b2c7ce36db1 100644 --- a/sdk/cosmos/azure-cosmos-spark_3-4_2-12/CHANGELOG.md +++ b/sdk/cosmos/azure-cosmos-spark_3-4_2-12/CHANGELOG.md @@ -3,6 +3,7 @@ ### 4.47.0-beta.1 (Unreleased) #### Features Added +* Added new `CosmosItemsDataSource.readManyByPartitionKey` Spark function to execute bulk queries by a list of pk-values with better efficiency. See [PR 48801](https://github.com/Azure/azure-sdk-for-java/pull/48801) #### Breaking Changes diff --git a/sdk/cosmos/azure-cosmos-spark_3-5_2-12/CHANGELOG.md b/sdk/cosmos/azure-cosmos-spark_3-5_2-12/CHANGELOG.md index f5eac38bdb71..2240a48b1654 100644 --- a/sdk/cosmos/azure-cosmos-spark_3-5_2-12/CHANGELOG.md +++ b/sdk/cosmos/azure-cosmos-spark_3-5_2-12/CHANGELOG.md @@ -3,6 +3,7 @@ ### 4.47.0-beta.1 (Unreleased) #### Features Added +* Added new `CosmosItemsDataSource.readManyByPartitionKey` Spark function to execute bulk queries by a list of pk-values with better efficiency. See [PR 48801](https://github.com/Azure/azure-sdk-for-java/pull/48801) #### Breaking Changes diff --git a/sdk/cosmos/azure-cosmos-spark_3-5_2-13/CHANGELOG.md b/sdk/cosmos/azure-cosmos-spark_3-5_2-13/CHANGELOG.md index 919d7fbfa325..20a3e3a61bd7 100644 --- a/sdk/cosmos/azure-cosmos-spark_3-5_2-13/CHANGELOG.md +++ b/sdk/cosmos/azure-cosmos-spark_3-5_2-13/CHANGELOG.md @@ -3,6 +3,7 @@ ### 4.47.0-beta.1 (Unreleased) #### Features Added +* Added new `CosmosItemsDataSource.readManyByPartitionKey` Spark function to execute bulk queries by a list of pk-values with better efficiency. See [PR 48801](https://github.com/Azure/azure-sdk-for-java/pull/48801) #### Breaking Changes diff --git a/sdk/cosmos/azure-cosmos-spark_4-0_2-13/CHANGELOG.md b/sdk/cosmos/azure-cosmos-spark_4-0_2-13/CHANGELOG.md index 3972ae6aeb98..d8368be6a0da 100644 --- a/sdk/cosmos/azure-cosmos-spark_4-0_2-13/CHANGELOG.md +++ b/sdk/cosmos/azure-cosmos-spark_4-0_2-13/CHANGELOG.md @@ -3,6 +3,7 @@ ### 4.47.0-beta.1 (Unreleased) #### Features Added +* Added new `CosmosItemsDataSource.readManyByPartitionKey` Spark function to execute bulk queries by a list of pk-values with better efficiency. See [PR 48801](https://github.com/Azure/azure-sdk-for-java/pull/48801) #### Breaking Changes diff --git a/sdk/cosmos/azure-cosmos/CHANGELOG.md b/sdk/cosmos/azure-cosmos/CHANGELOG.md index e8ea564fab7b..faf661ddd80f 100644 --- a/sdk/cosmos/azure-cosmos/CHANGELOG.md +++ b/sdk/cosmos/azure-cosmos/CHANGELOG.md @@ -3,6 +3,7 @@ ### 4.80.0-beta.1 (Unreleased) #### Features Added +* Added new `readManyByPartitioNKey` to bulk query by a list of pk-values with better efficiency. See [PR 48801](https://github.com/Azure/azure-sdk-for-java/pull/48801) #### Breaking Changes From 0b8905dbb011e5eae5720098ca0ddd2de9d4757a Mon Sep 17 00:00:00 2001 From: Fabian Meiswinkel Date: Thu, 16 Apr 2026 23:46:18 +0000 Subject: [PATCH 14/64] Addressing code review comments --- .../com/azure/cosmos/spark/CosmosConfig.scala | 22 +++++++++- .../cosmos/spark/CosmosItemsDataSource.scala | 5 ++- ...tionReaderWithReadManyByPartitionKey.scala | 2 + .../azure/cosmos/spark/CosmosConfigSpec.scala | 42 +++++++++++++++++++ .../spark/CosmosPartitionKeyHelperSpec.scala | 11 +++++ 5 files changed, 79 insertions(+), 3 deletions(-) diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosConfig.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosConfig.scala index 951f4735444d..e1b8f0b51f8a 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosConfig.scala +++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosConfig.scala @@ -92,6 +92,7 @@ private[spark] object CosmosConfigNames { val ReadPartitioningFeedRangeFilter = "spark.cosmos.partitioning.feedRangeFilter" val ReadRuntimeFilteringEnabled = "spark.cosmos.read.runtimeFiltering.enabled" val ReadManyFilteringEnabled = "spark.cosmos.read.readManyFiltering.enabled" + val ReadManyByPkNullHandling = "spark.cosmos.read.readManyByPk.nullHandling" val ViewsRepositoryPath = "spark.cosmos.views.repositoryPath" val DiagnosticsMode = "spark.cosmos.diagnostics" val DiagnosticsSamplingMaxCount = "spark.cosmos.diagnostics.sampling.maxCount" @@ -226,6 +227,7 @@ private[spark] object CosmosConfigNames { ReadPartitioningFeedRangeFilter, ReadRuntimeFilteringEnabled, ReadManyFilteringEnabled, + ReadManyByPkNullHandling, ViewsRepositoryPath, DiagnosticsMode, DiagnosticsSamplingIntervalInSeconds, @@ -1042,7 +1044,8 @@ private case class CosmosReadConfig(readConsistencyStrategy: ReadConsistencyStra throughputControlConfig: Option[CosmosThroughputControlConfig] = None, runtimeFilteringEnabled: Boolean, readManyFilteringConfig: CosmosReadManyFilteringConfig, - responseContinuationTokenLimitInKb: Option[Int] = None) + responseContinuationTokenLimitInKb: Option[Int] = None, + readManyByPkTreatNullAsNone: Boolean = false) private object SchemaConversionModes extends Enumeration { type SchemaConversionMode = Value @@ -1136,6 +1139,18 @@ private object CosmosReadConfig { helpMessage = " Indicates whether dynamic partition pruning filters will be pushed down when applicable." ) + private val ReadManyByPkNullHandling = CosmosConfigEntry[String]( + key = CosmosConfigNames.ReadManyByPkNullHandling, + mandatory = false, + defaultValue = Some("Null"), + parseFromStringFunction = value => value, + helpMessage = "Determines how null values in hierarchical partition key components are treated " + + "for readManyByPartitionKey. 'Null' (default) maps null to a JSON null value via addNullValue(), " + + "which is appropriate when the document field exists with an explicit null value. " + + "'None' maps null to PartitionKey.NONE via addNoneValue(), which should only be used when the " + + "partition key path does not exist at all in the document." + ) + def parseCosmosReadConfig(cfg: Map[String, String]): CosmosReadConfig = { val forceEventualConsistency = CosmosConfigEntry.parse(cfg, ForceEventualConsistency) val readConsistencyStrategyOverride = CosmosConfigEntry.parse(cfg, ReadConsistencyStrategyOverride) @@ -1158,6 +1173,8 @@ private object CosmosReadConfig { val throughputControlConfigOpt = CosmosThroughputControlConfig.parseThroughputControlConfig(cfg) val runtimeFilteringEnabled = CosmosConfigEntry.parse(cfg, ReadRuntimeFilteringEnabled) val readManyFilteringConfig = CosmosReadManyFilteringConfig.parseCosmosReadManyFilterConfig(cfg) + val readManyByPkNullHandling = CosmosConfigEntry.parse(cfg, ReadManyByPkNullHandling) + val readManyByPkTreatNullAsNone = readManyByPkNullHandling.getOrElse("Null").equalsIgnoreCase("None") val effectiveReadConsistencyStrategy = if (readConsistencyStrategyOverride.getOrElse(ReadConsistencyStrategy.DEFAULT) != ReadConsistencyStrategy.DEFAULT) { readConsistencyStrategyOverride.get @@ -1189,7 +1206,8 @@ private object CosmosReadConfig { throughputControlConfigOpt, runtimeFilteringEnabled.get, readManyFilteringConfig, - responseContinuationTokenLimitInKb) + responseContinuationTokenLimitInKb, + readManyByPkTreatNullAsNone) } } diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosItemsDataSource.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosItemsDataSource.scala index 6257e96e81e5..ac2299929f6a 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosItemsDataSource.scala +++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosItemsDataSource.scala @@ -146,6 +146,7 @@ object CosmosItemsDataSource { val containerConfig = CosmosContainerConfig.parseCosmosContainerConfig(effectiveConfig) val sparkEnvironmentInfo = CosmosClientConfiguration.getSparkEnvironmentInfo(None) val calledFrom = s"CosmosItemsDataSource.readManyByPartitionKey" + val treatNullAsNone = readConfig.readManyByPkTreatNullAsNone val pkPaths = Loan( List[Option[CosmosClientCacheItem]]( @@ -197,7 +198,9 @@ object CosmosItemsDataSource { case s: String => builder.add(s) case n: Number => builder.add(n.doubleValue()) case b: Boolean => builder.add(b) - case null => builder.addNoneValue() + case null => + if (treatNullAsNone) builder.addNoneValue() + else builder.addNullValue() case other => builder.add(other.toString) } } diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/ItemsPartitionReaderWithReadManyByPartitionKey.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/ItemsPartitionReaderWithReadManyByPartitionKey.scala index da3b81d951ae..73477b3a488d 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/ItemsPartitionReaderWithReadManyByPartitionKey.scala +++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/ItemsPartitionReaderWithReadManyByPartitionKey.scala @@ -170,6 +170,8 @@ private[spark] case class ItemsPartitionReaderWithReadManyByPartitionKey .enable(true) .build + readManyOptionsImpl.setCosmosEndToEndOperationLatencyPolicyConfig(endToEndTimeoutPolicy) + private trait CloseableSparkRowItemIterator { def hasNext: Boolean def next(): SparkRowItem diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/CosmosConfigSpec.scala b/sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/CosmosConfigSpec.scala index 17f75e45a746..17a298d62131 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/CosmosConfigSpec.scala +++ b/sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/CosmosConfigSpec.scala @@ -457,6 +457,7 @@ class CosmosConfigSpec extends UnitSpec with BasicLoggingTrait { config.runtimeFilteringEnabled shouldBe true config.readManyFilteringConfig.readManyFilteringEnabled shouldBe false config.readManyFilteringConfig.readManyFilterProperty shouldEqual "_itemIdentity" + config.readManyByPkTreatNullAsNone shouldBe false userConfig = Map( "spark.cosmos.read.forceEventualConsistency" -> "false", @@ -630,6 +631,47 @@ class CosmosConfigSpec extends UnitSpec with BasicLoggingTrait { config.customQuery.get.queryText shouldBe queryText } + it should "parse readManyByPk nullHandling configuration" in { + // Default (not specified) should treat null as JSON null (addNullValue) + var userConfig = Map( + "spark.cosmos.read.forceEventualConsistency" -> "false" + ) + var config = CosmosReadConfig.parseCosmosReadConfig(userConfig) + config.readManyByPkTreatNullAsNone shouldBe false + + // Explicit "Null" should treat null as JSON null (addNullValue) + userConfig = Map( + "spark.cosmos.read.forceEventualConsistency" -> "false", + "spark.cosmos.read.readManyByPk.nullHandling" -> "Null" + ) + config = CosmosReadConfig.parseCosmosReadConfig(userConfig) + config.readManyByPkTreatNullAsNone shouldBe false + + // Case-insensitive "null" + userConfig = Map( + "spark.cosmos.read.forceEventualConsistency" -> "false", + "spark.cosmos.read.readManyByPk.nullHandling" -> "null" + ) + config = CosmosReadConfig.parseCosmosReadConfig(userConfig) + config.readManyByPkTreatNullAsNone shouldBe false + + // "None" should treat null as PartitionKey.NONE (addNoneValue) + userConfig = Map( + "spark.cosmos.read.forceEventualConsistency" -> "false", + "spark.cosmos.read.readManyByPk.nullHandling" -> "None" + ) + config = CosmosReadConfig.parseCosmosReadConfig(userConfig) + config.readManyByPkTreatNullAsNone shouldBe true + + // Case-insensitive "none" + userConfig = Map( + "spark.cosmos.read.forceEventualConsistency" -> "false", + "spark.cosmos.read.readManyByPk.nullHandling" -> "none" + ) + config = CosmosReadConfig.parseCosmosReadConfig(userConfig) + config.readManyByPkTreatNullAsNone shouldBe true + } + it should "throw on invalid read configuration" in { val userConfig = Map( "spark.cosmos.read.schemaConversionMode" -> "not a valid value" diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/CosmosPartitionKeyHelperSpec.scala b/sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/CosmosPartitionKeyHelperSpec.scala index 6528d44f5ebe..1ac40e395847 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/CosmosPartitionKeyHelperSpec.scala +++ b/sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/CosmosPartitionKeyHelperSpec.scala @@ -88,6 +88,17 @@ class CosmosPartitionKeyHelperSpec extends UnitSpec { pk.isDefined shouldBe false } + it should "produce different partition keys for addNullValue vs addNoneValue in HPK" in { + // addNullValue represents an explicit JSON null for a field that exists with value null + val pkWithNull = new PartitionKeyBuilder().add("Redmond").addNullValue().build() + + // addNoneValue represents PartitionKey.NONE, meaning the field is absent/undefined + val pkWithNone = new PartitionKeyBuilder().add("Redmond").addNoneValue().build() + + // These MUST produce different partition key hashes and route to different physical partitions + pkWithNull should not equal pkWithNone + } + //scalastyle:on multiple.string.literals //scalastyle:on magic.number } From 22abc780ed8b81da8665297b59d25abad896d97f Mon Sep 17 00:00:00 2001 From: Fabian Meiswinkel Date: Thu, 16 Apr 2026 23:54:19 +0000 Subject: [PATCH 15/64] Addressing code review feedback --- ...tionReaderWithReadManyByPartitionKey.scala | 25 +-- ...tryingReadManyByPartitionKeyIterator.scala | 175 ++++++++++++++++++ sdk/cosmos/azure-cosmos/CHANGELOG.md | 2 +- 3 files changed, 186 insertions(+), 16 deletions(-) create mode 100644 sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/TransientIOErrorsRetryingReadManyByPartitionKeyIterator.scala diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/ItemsPartitionReaderWithReadManyByPartitionKey.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/ItemsPartitionReaderWithReadManyByPartitionKey.scala index 73477b3a488d..8d4952b2144c 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/ItemsPartitionReaderWithReadManyByPartitionKey.scala +++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/ItemsPartitionReaderWithReadManyByPartitionKey.scala @@ -6,7 +6,7 @@ package com.azure.cosmos.spark import com.azure.cosmos.{CosmosAsyncContainer, CosmosEndToEndOperationLatencyPolicyConfigBuilder, CosmosItemSerializerNoExceptionWrapping, SparkBridgeInternal} import com.azure.cosmos.implementation.spark.OperationContextAndListenerTuple import com.azure.cosmos.implementation.{ImplementationBridgeHelpers, ObjectNodeMap, SparkRowItem, Utils} -import com.azure.cosmos.models.{CosmosReadManyRequestOptions, ModelBridgeInternal, PartitionKey, PartitionKeyDefinition} +import com.azure.cosmos.models.{CosmosReadManyRequestOptions, ModelBridgeInternal, PartitionKey, PartitionKeyDefinition, SqlQuerySpec} import com.azure.cosmos.spark.BulkWriter.getThreadInfo import com.azure.cosmos.spark.CosmosTableSchemaInferrer.IdAttributeName import com.azure.cosmos.spark.diagnostics.{DetailedFeedDiagnosticsProvider, DiagnosticsContext, DiagnosticsLoader, LoggerHelper, SparkTaskContext} @@ -188,27 +188,22 @@ private[spark] case class ItemsPartitionReaderWithReadManyByPartitionKey override def close(): Unit = {} } - // Single iterator over all PKs — the SDK handles per-physical-partition batching - // internally to avoid oversized SQL queries. - // Short-circuit empty PK lists locally because the SDK rejects empty partition-key lists. + // Batch partition keys and retry each batch independently on transient I/O errors. + // This avoids the continuation-token problem with TransientIOErrorsRetryingIterator + // where a retry would re-read all data from scratch, causing silent data duplication. private lazy val iterator: CloseableSparkRowItemIterator = if (pkList.isEmpty) { EmptySparkRowItemIterator } else { new CloseableSparkRowItemIterator { - private val delegate = new TransientIOErrorsRetryingIterator[SparkRowItem]( - continuationToken => { - readConfig.customQuery match { - case Some(query) => - cosmosAsyncContainer.readManyByPartitionKey(pkList, query.toSqlQuerySpec, readManyOptions, classOf[SparkRowItem]) - case None => - cosmosAsyncContainer.readManyByPartitionKey(pkList, readManyOptions, classOf[SparkRowItem]) - } - }, + private val delegate = new TransientIOErrorsRetryingReadManyByPartitionKeyIterator[SparkRowItem]( + cosmosAsyncContainer, + pkList, + readConfig.customQuery.map(_.toSqlQuerySpec), + readManyOptions, readConfig.maxItemCount, - readConfig.prefetchBufferSize, operationContextAndListenerTuple, - None + classOf[SparkRowItem] ) override def hasNext: Boolean = delegate.hasNext diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/TransientIOErrorsRetryingReadManyByPartitionKeyIterator.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/TransientIOErrorsRetryingReadManyByPartitionKeyIterator.scala new file mode 100644 index 000000000000..dfdc380b09ce --- /dev/null +++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/TransientIOErrorsRetryingReadManyByPartitionKeyIterator.scala @@ -0,0 +1,175 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.azure.cosmos.spark + +import com.azure.cosmos.CosmosAsyncContainer +import com.azure.cosmos.implementation.OperationCancelledException +import com.azure.cosmos.implementation.spark.OperationContextAndListenerTuple +import com.azure.cosmos.models.{CosmosReadManyRequestOptions, PartitionKey, SqlQuerySpec} +import com.azure.cosmos.spark.diagnostics.BasicLoggingTrait + +import java.util.concurrent.{ExecutorService, SynchronousQueue, ThreadPoolExecutor, TimeUnit, TimeoutException} +import scala.concurrent.{Await, ExecutionContext, Future} + +// scalastyle:off underscore.import +import scala.collection.JavaConverters._ +// scalastyle:on underscore.import + +/** + * Retry-safe iterator for readManyByPartitionKey that batches partition keys and retries + * each batch independently on transient I/O errors. This avoids the continuation-token problem + * where TransientIOErrorsRetryingIterator would re-read all data from scratch on retry, + * causing silent data duplication. + */ +private[spark] class TransientIOErrorsRetryingReadManyByPartitionKeyIterator[TSparkRow] +( + val container: CosmosAsyncContainer, + val partitionKeys: java.util.List[PartitionKey], + val customQuery: Option[SqlQuerySpec], + val queryOptions: CosmosReadManyRequestOptions, + val pageSize: Int, + val operationContextAndListener: Option[OperationContextAndListenerTuple], + val classType: Class[TSparkRow] +) extends BufferedIterator[TSparkRow] with BasicLoggingTrait with AutoCloseable { + + private val maxPageRetrievalTimeout = scala.concurrent.duration.FiniteDuration( + 5 + CosmosConstants.readOperationEndToEndTimeoutInSeconds, + scala.concurrent.duration.SECONDS) + + private[spark] var currentItemIterator: Option[BufferedIterator[TSparkRow]] = None + private val pkBatchIterator = partitionKeys.asScala.iterator.grouped(pageSize) + + override def hasNext: Boolean = { + if (hasBufferedNext) { + true + } else { + hasNextInternal + } + } + + private def hasNextInternal: Boolean = { + var returnValue: Option[Boolean] = None + + while (returnValue.isEmpty) { + if (pkBatchIterator.hasNext) { + val pkBatch = pkBatchIterator.next().toList + returnValue = + TransientErrorsRetryPolicy.executeWithRetry( + () => hasNextInternalCore(pkBatch), + statusResetFuncBetweenRetry = Some(() => { currentItemIterator = None }) + ) + } else { + returnValue = Some(false) + } + } + + returnValue.get + } + + private def hasNextInternalCore(pkBatch: List[PartitionKey]): Option[Boolean] = { + val pkJavaList = new java.util.ArrayList[PartitionKey](pkBatch.asJava) + val results = try { + Await.result( + Future { + val flux = customQuery match { + case Some(query) => + container.readManyByPartitionKey(pkJavaList, query, queryOptions, classType) + case None => + container.readManyByPartitionKey(pkJavaList, queryOptions, classType) + } + + // Collect all pages for this batch into a single list + flux.collectList().block() + }(TransientIOErrorsRetryingReadManyByPartitionKeyIterator.executionContext), + maxPageRetrievalTimeout) + } catch { + case endToEndTimeoutException: OperationCancelledException => + val operationContextString = operationContextAndListener match { + case Some(o) => if (o.getOperationContext != null) { + o.getOperationContext.toString + } else { + "n/a" + } + case None => "n/a" + } + + val message = s"End-to-end timeout hit when trying to retrieve readManyByPartitionKey batch. " + + s"Batch size: ${pkBatch.size}, Context: $operationContextString" + + logError(message, throwable = endToEndTimeoutException) + + throw endToEndTimeoutException + case timeoutException: TimeoutException => + val operationContextString = operationContextAndListener match { + case Some(o) => if (o.getOperationContext != null) { + o.getOperationContext.toString + } else { + "n/a" + } + case None => "n/a" + } + + val message = s"Attempting to retrieve readManyByPartitionKey batch timed out. " + + s"Batch size: ${pkBatch.size}, Context: $operationContextString" + + logError(message, timeoutException) + + val exception = new OperationCancelledException( + message, + null + ) + exception.setStackTrace(timeoutException.getStackTrace) + throw exception + + case other: Throwable => throw other + } + + val iteratorCandidate = results.iterator().asScala.buffered + + if (iteratorCandidate.hasNext) { + currentItemIterator = Some(iteratorCandidate) + Some(true) + } else { + None + } + } + + private def hasBufferedNext: Boolean = { + currentItemIterator match { + case Some(iterator) => if (iterator.hasNext) { + true + } else { + currentItemIterator = None + false + } + case None => false + } + } + + override def next(): TSparkRow = { + currentItemIterator.get.next() + } + + override def head(): TSparkRow = { + currentItemIterator.get.head + } + + override def close(): Unit = {} +} + +private object TransientIOErrorsRetryingReadManyByPartitionKeyIterator extends BasicLoggingTrait { + private val maxConcurrency = SparkUtils.getNumberOfHostCPUCores + + val executorService: ExecutorService = new ThreadPoolExecutor( + maxConcurrency, + maxConcurrency, + 0L, + TimeUnit.MILLISECONDS, + new SynchronousQueue(), + SparkUtils.daemonThreadFactory(), + new ThreadPoolExecutor.CallerRunsPolicy() + ) + + val executionContext: ExecutionContext = ExecutionContext.fromExecutorService(executorService) +} diff --git a/sdk/cosmos/azure-cosmos/CHANGELOG.md b/sdk/cosmos/azure-cosmos/CHANGELOG.md index faf661ddd80f..904c01c3238f 100644 --- a/sdk/cosmos/azure-cosmos/CHANGELOG.md +++ b/sdk/cosmos/azure-cosmos/CHANGELOG.md @@ -3,7 +3,7 @@ ### 4.80.0-beta.1 (Unreleased) #### Features Added -* Added new `readManyByPartitioNKey` to bulk query by a list of pk-values with better efficiency. See [PR 48801](https://github.com/Azure/azure-sdk-for-java/pull/48801) +* Added new `readManyByPartitionKey` to bulk query by a list of pk-values with better efficiency. See [PR 48801](https://github.com/Azure/azure-sdk-for-java/pull/48801) #### Breaking Changes From 662b1a4b90ee6954d7467c4b707b340a2d6b446d Mon Sep 17 00:00:00 2001 From: Fabian Meiswinkel Date: Fri, 17 Apr 2026 00:01:35 +0000 Subject: [PATCH 16/64] Update CosmosItemsDataSource.scala --- .../cosmos/spark/CosmosItemsDataSource.scala | 31 ++++++++++++------- 1 file changed, 20 insertions(+), 11 deletions(-) diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosItemsDataSource.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosItemsDataSource.scala index ac2299929f6a..aac5a53d8a8d 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosItemsDataSource.scala +++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosItemsDataSource.scala @@ -188,21 +188,12 @@ object CosmosItemsDataSource { Some((row: Row) => { if (pkPaths.size == 1) { // Single partition key - new PartitionKey(row.getAs[Any](pkPaths.head)) + buildPartitionKey(row.getAs[Any](pkPaths.head), treatNullAsNone) } else { // Hierarchical partition key — build level by level val builder = new PartitionKeyBuilder() for (path <- pkPaths) { - val value = row.getAs[Any](path) - value match { - case s: String => builder.add(s) - case n: Number => builder.add(n.doubleValue()) - case b: Boolean => builder.add(b) - case null => - if (treatNullAsNone) builder.addNoneValue() - else builder.addNullValue() - case other => builder.add(other.toString) - } + addPartitionKeyComponent(builder, row.getAs[Any](path), treatNullAsNone) } builder.build() } @@ -222,4 +213,22 @@ object CosmosItemsDataSource { readManyReader.readManyByPartitionKey(df.rdd, pkExtraction) } + + private def addPartitionKeyComponent(builder: PartitionKeyBuilder, value: Any, treatNullAsNone: Boolean): Unit = { + value match { + case s: String => builder.add(s) + case n: Number => builder.add(n.doubleValue()) + case b: Boolean => builder.add(b) + case null => + if (treatNullAsNone) builder.addNoneValue() + else builder.addNullValue() + case other => builder.add(other.toString) + } + } + + private def buildPartitionKey(value: Any, treatNullAsNone: Boolean): PartitionKey = { + val builder = new PartitionKeyBuilder() + addPartitionKeyComponent(builder, value, treatNullAsNone) + builder.build() + } } From c764de9de02caa44d307a31b5cb5f8a6755ca6d4 Mon Sep 17 00:00:00 2001 From: Fabian Meiswinkel Date: Fri, 17 Apr 2026 00:03:16 +0000 Subject: [PATCH 17/64] Update CosmosItemsDataSource.scala --- .../com/azure/cosmos/spark/CosmosItemsDataSource.scala | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosItemsDataSource.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosItemsDataSource.scala index aac5a53d8a8d..86ef865bcb83 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosItemsDataSource.scala +++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosItemsDataSource.scala @@ -131,8 +131,13 @@ object CosmosItemsDataSource { val pkIdentityFieldExtraction = df .schema .find(field => field.name.equals(CosmosConstants.Properties.PartitionKeyIdentity) && field.dataType.equals(StringType)) - .map(field => (row: Row) => - CosmosPartitionKeyHelper.tryParsePartitionKey(row.getString(row.fieldIndex(field.name))).get) + .map(field => (row: Row) => { + val rawValue = row.getString(row.fieldIndex(field.name)) + CosmosPartitionKeyHelper.tryParsePartitionKey(rawValue) + .getOrElse(throw new IllegalArgumentException( + s"Invalid _partitionKeyIdentity value in row: '$rawValue'. " + + "Expected format: pk([...json...])")) + }) // Option 2: Detect PK columns by matching the container's partition key paths against the DataFrame schema val pkColumnExtraction: Option[Row => PartitionKey] = if (pkIdentityFieldExtraction.isDefined) { From 080ce4a22931755e02810d48cf6d2fd1d38da719 Mon Sep 17 00:00:00 2001 From: Fabian Meiswinkel Date: Fri, 17 Apr 2026 11:20:49 +0000 Subject: [PATCH 18/64] Update RxDocumentClientImpl.java --- .../implementation/RxDocumentClientImpl.java | 56 +++++++++++++++++++ 1 file changed, 56 insertions(+) diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java index b4c72532280c..e5d8248c1578 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java @@ -4382,6 +4382,62 @@ public Flux> readManyByPartitionKey( (ctx) -> diagnosticsFactory.merge(ctx) ); + StaleResourceRetryPolicy staleResourceRetryPolicy = new StaleResourceRetryPolicy( + this.collectionCache, + null, + collectionLink, + queryOptionsAccessor().getProperties(state.getQueryOptions()), + queryOptionsAccessor().getHeaders(state.getQueryOptions()), + this.sessionContainer, + diagnosticsFactory, + ResourceType.Document + ); + + return ObservableHelper + .fluxInlineIfPossibleAsObs( + () -> readManyByPartitionKey( + partitionKeys, customQuery, collectionLink, state, diagnosticsFactory, klass), + staleResourceRetryPolicy + ) + .onErrorMap(throwable -> { + if (throwable instanceof CosmosException) { + CosmosException cosmosException = (CosmosException) throwable; + CosmosDiagnostics diagnostics = cosmosException.getDiagnostics(); + if (diagnostics != null) { + state.mergeDiagnosticsContext(); + CosmosDiagnosticsContext ctx = state.getDiagnosticsContextSnapshot(); + if (ctx != null) { + ctxAccessor().recordOperation( + ctx, + cosmosException.getStatusCode(), + cosmosException.getSubStatusCode(), + 0, + cosmosException.getRequestCharge(), + diagnostics, + throwable + ); + diagAccessor() + .setDiagnosticsContext( + diagnostics, + state.getDiagnosticsContextSnapshot()); + } + } + + return cosmosException; + } + + return throwable; + }); + } + + private Flux> readManyByPartitionKey( + List partitionKeys, + SqlQuerySpec customQuery, + String collectionLink, + QueryFeedOperationState state, + ScopedDiagnosticsFactory diagnosticsFactory, + Class klass) { + String resourceLink = parentResourceLinkToQueryLink(collectionLink, ResourceType.Document); RxDocumentServiceRequest request = RxDocumentServiceRequest.create(diagnosticsFactory, OperationType.Query, From b01f8758eea8c870df2a746ddc0b350fe4e9cfd8 Mon Sep 17 00:00:00 2001 From: Fabian Meiswinkel Date: Fri, 17 Apr 2026 11:37:51 +0000 Subject: [PATCH 19/64] Fix readManyByPartitionKey retries --- ...tionReaderWithReadManyByPartitionKey.scala | 1 + ...tryingReadManyByPartitionKeyIterator.scala | 236 ++++++++++++------ 2 files changed, 161 insertions(+), 76 deletions(-) diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/ItemsPartitionReaderWithReadManyByPartitionKey.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/ItemsPartitionReaderWithReadManyByPartitionKey.scala index 8d4952b2144c..c67cc9c10be1 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/ItemsPartitionReaderWithReadManyByPartitionKey.scala +++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/ItemsPartitionReaderWithReadManyByPartitionKey.scala @@ -202,6 +202,7 @@ private[spark] case class ItemsPartitionReaderWithReadManyByPartitionKey readConfig.customQuery.map(_.toSqlQuerySpec), readManyOptions, readConfig.maxItemCount, + readConfig.prefetchBufferSize, operationContextAndListenerTuple, classOf[SparkRowItem] ) diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/TransientIOErrorsRetryingReadManyByPartitionKeyIterator.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/TransientIOErrorsRetryingReadManyByPartitionKeyIterator.scala index dfdc380b09ce..dcfdf4f93536 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/TransientIOErrorsRetryingReadManyByPartitionKeyIterator.scala +++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/TransientIOErrorsRetryingReadManyByPartitionKeyIterator.scala @@ -3,24 +3,29 @@ package com.azure.cosmos.spark -import com.azure.cosmos.CosmosAsyncContainer +import com.azure.cosmos.{CosmosAsyncContainer, CosmosException} import com.azure.cosmos.implementation.OperationCancelledException import com.azure.cosmos.implementation.spark.OperationContextAndListenerTuple -import com.azure.cosmos.models.{CosmosReadManyRequestOptions, PartitionKey, SqlQuerySpec} +import com.azure.cosmos.models.{CosmosReadManyRequestOptions, FeedResponse, PartitionKey, SqlQuerySpec} import com.azure.cosmos.spark.diagnostics.BasicLoggingTrait +import com.azure.cosmos.util.CosmosPagedIterable import java.util.concurrent.{ExecutorService, SynchronousQueue, ThreadPoolExecutor, TimeUnit, TimeoutException} +import java.util.concurrent.atomic.AtomicLong import scala.concurrent.{Await, ExecutionContext, Future} +import scala.util.Random +import scala.util.control.Breaks // scalastyle:off underscore.import import scala.collection.JavaConverters._ // scalastyle:on underscore.import /** - * Retry-safe iterator for readManyByPartitionKey that batches partition keys and retries - * each batch independently on transient I/O errors. This avoids the continuation-token problem - * where TransientIOErrorsRetryingIterator would re-read all data from scratch on retry, - * causing silent data duplication. + * Retry-safe iterator for readManyByPartitionKey that batches partition keys and lazily + * iterates pages within each batch via CosmosPagedIterable — consistent with how + * TransientIOErrorsRetryingIterator handles normal queries. On transient I/O errors the + * current batch's flux is recreated and pages already consumed are replayed, avoiding + * the memory overhead of collectList and matching the query iterator's structure. */ private[spark] class TransientIOErrorsRetryingReadManyByPartitionKeyIterator[TSparkRow] ( @@ -29,109 +34,139 @@ private[spark] class TransientIOErrorsRetryingReadManyByPartitionKeyIterator[TSp val customQuery: Option[SqlQuerySpec], val queryOptions: CosmosReadManyRequestOptions, val pageSize: Int, + val pagePrefetchBufferSize: Int, val operationContextAndListener: Option[OperationContextAndListenerTuple], val classType: Class[TSparkRow] ) extends BufferedIterator[TSparkRow] with BasicLoggingTrait with AutoCloseable { + private[spark] var maxRetryIntervalInMs = CosmosConstants.maxRetryIntervalForTransientFailuresInMs + private[spark] var maxRetryCount = CosmosConstants.maxRetryCountForTransientFailures + private val maxPageRetrievalTimeout = scala.concurrent.duration.FiniteDuration( 5 + CosmosConstants.readOperationEndToEndTimeoutInSeconds, scala.concurrent.duration.SECONDS) + private val rnd = Random + private val retryCount = new AtomicLong(0) + private lazy val operationContextString = operationContextAndListener match { + case Some(o) => if (o.getOperationContext != null) { + o.getOperationContext.toString + } else { + "n/a" + } + case None => "n/a" + } + + private[spark] var currentFeedResponseIterator: Option[BufferedIterator[FeedResponse[TSparkRow]]] = None private[spark] var currentItemIterator: Option[BufferedIterator[TSparkRow]] = None + private val pkBatchIterator = partitionKeys.asScala.iterator.grouped(pageSize) + // Track the current batch so we can replay it on retry + private var currentBatch: Option[java.util.List[PartitionKey]] = None override def hasNext: Boolean = { - if (hasBufferedNext) { - true - } else { - hasNextInternal - } + executeWithRetry("hasNextInternal", () => hasNextInternal) } private def hasNextInternal: Boolean = { var returnValue: Option[Boolean] = None while (returnValue.isEmpty) { - if (pkBatchIterator.hasNext) { - val pkBatch = pkBatchIterator.next().toList - returnValue = - TransientErrorsRetryPolicy.executeWithRetry( - () => hasNextInternalCore(pkBatch), - statusResetFuncBetweenRetry = Some(() => { currentItemIterator = None }) - ) - } else { - returnValue = Some(false) - } + returnValue = hasNextInternalCore } returnValue.get } - private def hasNextInternalCore(pkBatch: List[PartitionKey]): Option[Boolean] = { - val pkJavaList = new java.util.ArrayList[PartitionKey](pkBatch.asJava) - val results = try { - Await.result( - Future { - val flux = customQuery match { - case Some(query) => - container.readManyByPartitionKey(pkJavaList, query, queryOptions, classType) + private def hasNextInternalCore: Option[Boolean] = { + if (hasBufferedNext) { + Some(true) + } else { + val feedResponseIterator = currentFeedResponseIterator match { + case Some(existing) => existing + case None => + // Need a new feed response iterator — either for the current batch (on retry) + // or for the next batch + val batch = currentBatch match { + case Some(b) => b // retry of current batch case None => - container.readManyByPartitionKey(pkJavaList, queryOptions, classType) + if (pkBatchIterator.hasNext) { + val nextBatch = new java.util.ArrayList[PartitionKey](pkBatchIterator.next().toList.asJava) + currentBatch = Some(nextBatch) + nextBatch + } else { + return Some(false) // no more batches + } } - // Collect all pages for this batch into a single list - flux.collectList().block() - }(TransientIOErrorsRetryingReadManyByPartitionKeyIterator.executionContext), - maxPageRetrievalTimeout) - } catch { - case endToEndTimeoutException: OperationCancelledException => - val operationContextString = operationContextAndListener match { - case Some(o) => if (o.getOperationContext != null) { - o.getOperationContext.toString - } else { - "n/a" + val pagedFlux = customQuery match { + case Some(query) => + container.readManyByPartitionKey(batch, query, queryOptions, classType) + case None => + container.readManyByPartitionKey(batch, queryOptions, classType) } - case None => "n/a" - } - - val message = s"End-to-end timeout hit when trying to retrieve readManyByPartitionKey batch. " + - s"Batch size: ${pkBatch.size}, Context: $operationContextString" - - logError(message, throwable = endToEndTimeoutException) - throw endToEndTimeoutException - case timeoutException: TimeoutException => - val operationContextString = operationContextAndListener match { - case Some(o) => if (o.getOperationContext != null) { - o.getOperationContext.toString - } else { - "n/a" - } - case None => "n/a" - } + currentFeedResponseIterator = Some( + new CosmosPagedIterable[TSparkRow]( + pagedFlux, + pageSize, + pagePrefetchBufferSize + ) + .iterableByPage() + .iterator + .asScala + .buffered + ) - val message = s"Attempting to retrieve readManyByPartitionKey batch timed out. " + - s"Batch size: ${pkBatch.size}, Context: $operationContextString" + currentFeedResponseIterator.get + } - logError(message, timeoutException) + val hasNext: Boolean = try { + Await.result( + Future { + feedResponseIterator.hasNext + }(TransientIOErrorsRetryingReadManyByPartitionKeyIterator.executionContext), + maxPageRetrievalTimeout) + } catch { + case endToEndTimeoutException: OperationCancelledException => + val message = s"End-to-end timeout hit when trying to retrieve the next page. " + + s"Context: $operationContextString" + logError(message, throwable = endToEndTimeoutException) + throw endToEndTimeoutException - val exception = new OperationCancelledException( - message, - null - ) - exception.setStackTrace(timeoutException.getStackTrace) - throw exception + case timeoutException: TimeoutException => + val message = s"Attempting to retrieve the next page timed out. " + + s"Context: $operationContextString" + logError(message, timeoutException) + val exception = new OperationCancelledException(message, null) + exception.setStackTrace(timeoutException.getStackTrace) + throw exception - case other: Throwable => throw other - } + case other: Throwable => throw other + } - val iteratorCandidate = results.iterator().asScala.buffered + if (hasNext) { + val feedResponse = feedResponseIterator.next() + if (operationContextAndListener.isDefined) { + operationContextAndListener.get.getOperationListener.feedResponseProcessedListener( + operationContextAndListener.get.getOperationContext, + feedResponse) + } + val iteratorCandidate = feedResponse.getResults.iterator().asScala.buffered - if (iteratorCandidate.hasNext) { - currentItemIterator = Some(iteratorCandidate) - Some(true) - } else { - None + if (iteratorCandidate.hasNext) { + currentItemIterator = Some(iteratorCandidate) + Some(true) + } else { + // empty page interleaved — try again + None + } + } else { + // Current batch's flux is exhausted — move to next batch + currentBatch = None + currentFeedResponseIterator = None + None + } } } @@ -155,7 +190,56 @@ private[spark] class TransientIOErrorsRetryingReadManyByPartitionKeyIterator[TSp currentItemIterator.get.head } - override def close(): Unit = {} + private[spark] def executeWithRetry[T](methodName: String, func: () => T): T = { + val loop = new Breaks() + var returnValue: Option[T] = None + + loop.breakable { + while (true) { + val retryIntervalInMs = rnd.nextInt(maxRetryIntervalInMs) + + try { + returnValue = Some(func()) + retryCount.set(0) + loop.break + } + catch { + case cosmosException: CosmosException => + if (Exceptions.canBeTransientFailure(cosmosException.getStatusCode, cosmosException.getSubStatusCode)) { + val retryCountSnapshot = retryCount.incrementAndGet() + if (retryCountSnapshot > maxRetryCount) { + logError( + s"Too many transient failure retry attempts in " + + s"TransientIOErrorsRetryingReadManyByPartitionKeyIterator.$methodName", + cosmosException) + throw cosmosException + } else { + logWarning( + s"Transient failure handled in " + + s"TransientIOErrorsRetryingReadManyByPartitionKeyIterator.$methodName -" + + s" will be retried (attempt#$retryCountSnapshot) in ${retryIntervalInMs}ms", + cosmosException) + } + } else { + throw cosmosException + } + case other: Throwable => throw other + } + + // Reset iterators but keep currentBatch so the batch is replayed + currentItemIterator = None + currentFeedResponseIterator = None + Thread.sleep(retryIntervalInMs) + } + } + + returnValue.get + } + + override def close(): Unit = { + currentItemIterator = None + currentFeedResponseIterator = None + } } private object TransientIOErrorsRetryingReadManyByPartitionKeyIterator extends BasicLoggingTrait { From 7130d4aa35a228b116df9220394ab6a1ca569a9f Mon Sep 17 00:00:00 2001 From: Fabian Meiswinkel Date: Fri, 17 Apr 2026 11:48:54 +0000 Subject: [PATCH 20/64] Fix PK.None --- .../azure/cosmos/CosmosAsyncContainer.java | 6 ++ .../ReadManyByPartitionKeyQueryHelper.java | 97 ++++++++++++++----- .../implementation/RxDocumentClientImpl.java | 14 ++- 3 files changed, 88 insertions(+), 29 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java index 025a957a4606..4e234667c1c0 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java @@ -1673,6 +1673,12 @@ public CosmosPagedFlux readManyByPartitionKey( if (partitionKeys.isEmpty()) { throw new IllegalArgumentException("Argument 'partitionKeys' must not be empty."); } + for (PartitionKey pk : partitionKeys) { + if (pk == null) { + throw new IllegalArgumentException( + "Argument 'partitionKeys' must not contain null elements."); + } + } return UtilBridgeInternal.createCosmosPagedFlux( readManyByPartitionKeyInternalFunc(partitionKeys, customQuery, requestOptions, classType)); diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelper.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelper.java index d538542df1ba..4a6d2efdeee5 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelper.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelper.java @@ -40,24 +40,54 @@ public static SqlQuerySpec createReadManyByPkQuerySpec( boolean isSinglePathPk = partitionKeySelectors.size() == 1; if (isSinglePathPk && pkDefinition.getKind() != PartitionKind.MULTI_HASH) { - // Single PK path — use IN clause: alias["pkPath"] IN (@__rmPk_0, @__rmPk_1, ...) + // Single PK path — use IN clause for normal values, OR NOT IS_DEFINED for NONE + // First, separate NONE PKs from normal PKs + boolean hasNone = false; + List normalPkValues = new ArrayList<>(); + for (PartitionKey pk : pkValues) { + PartitionKeyInternal pkInternal = BridgeInternal.getPartitionKeyInternal(pk); + if (pkInternal.getComponents() == null) { + hasNone = true; + } else { + normalPkValues.add(pk); + } + } + pkFilter.append(" "); - pkFilter.append(tableAlias); - pkFilter.append(partitionKeySelectors.get(0)); - pkFilter.append(" IN ( "); - for (int i = 0; i < pkValues.size(); i++) { - PartitionKeyInternal pkInternal = BridgeInternal.getPartitionKeyInternal(pkValues.get(i)); - Object[] pkComponents = pkInternal.toObjectArray(); - String pkParamName = PK_PARAM_PREFIX + paramCount; - parameters.add(new SqlParameter(pkParamName, pkComponents[0])); - paramCount++; + boolean hasNormalValues = !normalPkValues.isEmpty(); + if (hasNormalValues && hasNone) { + pkFilter.append("("); + } + if (hasNormalValues) { + pkFilter.append(tableAlias); + pkFilter.append(partitionKeySelectors.get(0)); + pkFilter.append(" IN ( "); + for (int i = 0; i < normalPkValues.size(); i++) { + PartitionKeyInternal pkInternal = BridgeInternal.getPartitionKeyInternal(normalPkValues.get(i)); + Object[] pkComponents = pkInternal.toObjectArray(); + String pkParamName = PK_PARAM_PREFIX + paramCount; + parameters.add(new SqlParameter(pkParamName, pkComponents[0])); + paramCount++; - pkFilter.append(pkParamName); - if (i < pkValues.size() - 1) { - pkFilter.append(", "); + pkFilter.append(pkParamName); + if (i < normalPkValues.size() - 1) { + pkFilter.append(", "); + } } + pkFilter.append(" )"); + } + if (hasNone) { + if (hasNormalValues) { + pkFilter.append(" OR "); + } + pkFilter.append("NOT IS_DEFINED("); + pkFilter.append(tableAlias); + pkFilter.append(partitionKeySelectors.get(0)); + pkFilter.append(")"); + } + if (hasNormalValues && hasNone) { + pkFilter.append(")"); } - pkFilter.append(" )"); } else { // Multiple PK paths (HPK) or MULTI_HASH — use OR of AND clauses pkFilter.append(" "); @@ -65,21 +95,36 @@ public static SqlQuerySpec createReadManyByPkQuerySpec( PartitionKeyInternal pkInternal = BridgeInternal.getPartitionKeyInternal(pkValues.get(i)); Object[] pkComponents = pkInternal.toObjectArray(); - pkFilter.append("("); - for (int j = 0; j < pkComponents.length; j++) { - String pkParamName = PK_PARAM_PREFIX + paramCount; - parameters.add(new SqlParameter(pkParamName, pkComponents[j])); - paramCount++; + // PartitionKey.NONE — generate NOT IS_DEFINED for all PK paths + if (pkComponents == null) { + pkFilter.append("("); + for (int j = 0; j < partitionKeySelectors.size(); j++) { + if (j > 0) { + pkFilter.append(" AND "); + } + pkFilter.append("NOT IS_DEFINED("); + pkFilter.append(tableAlias); + pkFilter.append(partitionKeySelectors.get(j)); + pkFilter.append(")"); + } + pkFilter.append(")"); + } else { + pkFilter.append("("); + for (int j = 0; j < pkComponents.length; j++) { + String pkParamName = PK_PARAM_PREFIX + paramCount; + parameters.add(new SqlParameter(pkParamName, pkComponents[j])); + paramCount++; - if (j > 0) { - pkFilter.append(" AND "); + if (j > 0) { + pkFilter.append(" AND "); + } + pkFilter.append(tableAlias); + pkFilter.append(partitionKeySelectors.get(j)); + pkFilter.append(" = "); + pkFilter.append(pkParamName); } - pkFilter.append(tableAlias); - pkFilter.append(partitionKeySelectors.get(j)); - pkFilter.append(" = "); - pkFilter.append(pkParamName); + pkFilter.append(")"); } - pkFilter.append(")"); if (i < pkValues.size() - 1) { pkFilter.append(" OR "); diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java index e5d8248c1578..c70dedaa1f2c 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java @@ -4623,7 +4623,15 @@ private Map> groupPartitionKeysByPhysicalP for (PartitionKey pk : partitionKeys) { PartitionKeyInternal pkInternal = BridgeInternal.getPartitionKeyInternal(pk); - int componentCount = pkInternal.getComponents().size(); + + // PartitionKey.NONE wraps NonePartitionKey which has components = null. + // For routing purposes, treat NONE as UndefinedPartitionKey — documents ingested + // without a partition key path are stored with the undefined EPK. + PartitionKeyInternal effectivePkInternal = pkInternal.getComponents() == null + ? PartitionKeyInternal.UndefinedPartitionKey + : pkInternal; + + int componentCount = effectivePkInternal.getComponents().size(); int definedPathCount = pkDefinition.getPaths().size(); List targetRanges; @@ -4631,12 +4639,12 @@ private Map> groupPartitionKeysByPhysicalP if (pkDefinition.getKind() == PartitionKind.MULTI_HASH && componentCount < definedPathCount) { // Partial HPK — compute EPK prefix range and find all overlapping physical partitions Range epkRange = PartitionKeyInternalHelper.getEPKRangeForPrefixPartitionKey( - pkInternal, pkDefinition); + effectivePkInternal, pkDefinition); targetRanges = routingMap.getOverlappingRanges(epkRange); } else { // Full PK — maps to exactly one physical partition String effectivePartitionKeyString = PartitionKeyInternalHelper - .getEffectivePartitionKeyString(pkInternal, pkDefinition); + .getEffectivePartitionKeyString(effectivePkInternal, pkDefinition); PartitionKeyRange range = routingMap.getRangeByEffectivePartitionKey(effectivePartitionKeyString); targetRanges = Collections.singletonList(range); } From 93957f3a8442d730fe67fbc379ef5399f46f5665 Mon Sep 17 00:00:00 2001 From: Fabian Meiswinkel Date: Fri, 17 Apr 2026 11:56:12 +0000 Subject: [PATCH 21/64] Update ReadManyByPartitionKeyQueryHelper.java --- .../ReadManyByPartitionKeyQueryHelper.java | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelper.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelper.java index 4a6d2efdeee5..6d6cd084e01a 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelper.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelper.java @@ -224,10 +224,17 @@ static int findTopLevelKeywordIndex(String queryText, String keyword) { int keyLen = keywordUpper.length(); for (int i = 0; i <= queryTextUpper.length() - keyLen; i++) { char ch = queryTextUpper.charAt(i); - // Skip string literals enclosed in single quotes + // Skip string literals enclosed in single quotes (handle '' escape) if (queryText.charAt(i) == '\'') { i++; - while (i < queryText.length() && queryText.charAt(i) != '\'') { + while (i < queryText.length()) { + if (queryText.charAt(i) == '\'') { + if (i + 1 < queryText.length() && queryText.charAt(i + 1) == '\'') { + i += 2; // escaped quote — skip both + continue; + } + break; // end of string literal + } i++; } continue; From 9200f8fe53b95c51f4bec29e6afbc3b8480ec093 Mon Sep 17 00:00:00 2001 From: Fabian Meiswinkel Date: Fri, 17 Apr 2026 12:50:47 +0000 Subject: [PATCH 22/64] Fix code review feedback --- ...bianm_readManyByPK-vs-origin_main-full.txt | 3444 +++++++++++++++++ ...bianm_readManyByPK-vs-origin_main-stat.txt | 76 + .../azure-cosmos-spark_3-3_2-12/CHANGELOG.md | 2 +- .../azure-cosmos-spark_3-4_2-12/CHANGELOG.md | 2 +- .../azure-cosmos-spark_3-5_2-12/CHANGELOG.md | 2 +- .../azure-cosmos-spark_3-5_2-13/CHANGELOG.md | 2 +- .../com/azure/cosmos/spark/CosmosConfig.scala | 11 +- .../cosmos/spark/CosmosItemsDataSource.scala | 23 +- .../CosmosReadManyByPartitionKeyReader.scala | 7 +- ...tionReaderWithReadManyByPartitionKey.scala | 22 +- ...tryingReadManyByPartitionKeyIterator.scala | 127 +- .../udf/GetCosmosPartitionKeyValue.scala | 21 +- .../azure-cosmos-spark_4-0_2-13/CHANGELOG.md | 2 +- .../azure/cosmos/implementation/Configs.java | 10 +- .../ReadManyByPartitionKeyQueryHelper.java | 30 +- .../implementation/RxDocumentClientImpl.java | 16 +- 16 files changed, 3697 insertions(+), 100 deletions(-) create mode 100644 sdk/cosmos/.temp/branch-users_fabianm_readManyByPK-vs-origin_main-full.txt create mode 100644 sdk/cosmos/.temp/branch-users_fabianm_readManyByPK-vs-origin_main-stat.txt diff --git a/sdk/cosmos/.temp/branch-users_fabianm_readManyByPK-vs-origin_main-full.txt b/sdk/cosmos/.temp/branch-users_fabianm_readManyByPK-vs-origin_main-full.txt new file mode 100644 index 000000000000..5eee65a15a7b --- /dev/null +++ b/sdk/cosmos/.temp/branch-users_fabianm_readManyByPK-vs-origin_main-full.txt @@ -0,0 +1,3444 @@ +===== PR #LOCAL-users_fabianm_readManyByPK ===== +Title: Branch comparison users/fabianm/readManyByPK vs origin/main +Author: Fabian Meiswinkel +Status: DIVERGED (ahead 30, behind 4) +Branch: users/fabianm/readManyByPK -> origin/main +Head SHA: 93957f3a8442d730fe67fbc379ef5399f46f5665 +Merge Base: 20313f79ba8dd0dfa97862d0c31dd4b2e44ee671 +URL: N/A (local branch comparison) + +--- Description --- +Adds readManyByPartitionKey API (sync+async) and Spark connector support for PK-only reads, with query-plan-based validation +--- End Description --- + +===== Commits in PR ===== +9770833eb59 Adding readManyByPartitionKey API +ac287bcdf00 Merge branch 'main' of https://github.com/Azure/azure-sdk-for-java into users/fabianm/readManyByPK +9a5b3e96e7e Update sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/ItemsPartitionReaderWithReadManyByPartitionKey.scala +a8720c3c9f2 Update sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/ItemsPartitionReaderWithReadManyByPartitionKey.scala +d499da76fb4 Update sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java +c3c542a33a7 Update sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelper.java +4416354e03e ┬┤Fixing code review comments +3ab3f0d64f5 Merge branch 'main' of https://github.com/Azure/azure-sdk-for-java into users/fabianm/readManyByPK +588a7550c54 Update CosmosAsyncContainer.java +8c5cdb47b31 Merge branch 'main' into users/fabianm/readManyByPK +f5485527a9c Update ReadManyByPartitionKeyTest.java +f68cf02ff71 Fixing test issues +8b6c4b168ea Update CosmosAsyncContainer.java +8ba7f4db2da Merge branch 'main' into users/fabianm/readManyByPK +56b067a9339 Reacted to code review feedback +fa430e918fa Merge branch 'main' into users/fabianm/readManyByPK +d9504c91f34 Fix build issues +73151f09e5f Merge branch 'main' into users/fabianm/readManyByPK +681830e2d4a Fixing changelog +7f745e60641 Merge branch 'main' into users/fabianm/readManyByPK +0b8905dbb01 Addressing code review comments +22abc780ed8 Addressing code review feedback +662b1a4b90e Update CosmosItemsDataSource.scala +c764de9de02 Update CosmosItemsDataSource.scala +e1e6f5a6f73 Merge branch 'main' into users/fabianm/readManyByPK +080ce4a2293 Update RxDocumentClientImpl.java +516bbf3a95a Merge branch 'users/fabianm/readManyByPK' of https://github.com/Azure/azure-sdk-for-java into users/fabianm/readManyByPK +b01f8758eea Fix readManyByPartitionKey retries +7130d4aa35a Fix PK.None +93957f3a844 Update ReadManyByPartitionKeyQueryHelper.java + +===== Files Changed ===== + sdk/cosmos/azure-cosmos-spark_3-3_2-12/CHANGELOG.md (+1 -0) + sdk/cosmos/azure-cosmos-spark_3-4_2-12/CHANGELOG.md (+1 -0) + sdk/cosmos/azure-cosmos-spark_3-5_2-12/CHANGELOG.md (+1 -0) + sdk/cosmos/azure-cosmos-spark_3-5_2-13/CHANGELOG.md (+1 -0) + sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosConfig.scala (+20 -2) + sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosConstants.scala (+1 -0) + sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosItemsDataSource.scala (+125 -1) + sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosPartitionKeyHelper.scala (+45 -0) + sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosReadManyByPartitionKeyReader.scala (+150 -0) + sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/ItemsPartitionReaderWithReadManyByPartitionKey.scala (+249 -0) + sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/TransientIOErrorsRetryingReadManyByPartitionKeyIterator.scala (+259 -0) + sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/udf/GetCosmosPartitionKeyValue.scala (+25 -0) + sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/CosmosConfigSpec.scala (+42 -0) + sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/CosmosPartitionKeyHelperSpec.scala (+104 -0) + sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/ItemsPartitionReaderWithReadManyByPartitionKeyITest.scala (+158 -0) + sdk/cosmos/azure-cosmos-spark_4-0_2-13/CHANGELOG.md (+1 -0) + sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/ReadManyByPartitionKeyTest.java (+462 -0) + sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelperTest.java (+426 -0) + sdk/cosmos/azure-cosmos/CHANGELOG.md (+1 -0) + sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java (+126 -0) + sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosContainer.java (+67 -0) + sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/AsyncDocumentClient.java (+21 -0) + sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/Configs.java (+19 -0) + sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelper.java (+263 -0) + sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java (+292 -0) + sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/DocumentQueryExecutionContextFactory.java (+11 -0) + sdk/cosmos/cspell.yaml (+6 -0) + sdk/cosmos/docs/readManyByPartitionKey-design.md (+169 -0) + +===== Full Diff ===== +diff --git a/sdk/cosmos/azure-cosmos-spark_3-3_2-12/CHANGELOG.md b/sdk/cosmos/azure-cosmos-spark_3-3_2-12/CHANGELOG.md +index cbf97c610f9..fe114462019 100644 +--- a/sdk/cosmos/azure-cosmos-spark_3-3_2-12/CHANGELOG.md ++++ b/sdk/cosmos/azure-cosmos-spark_3-3_2-12/CHANGELOG.md +@@ -3,6 +3,7 @@ + ### 4.47.0-beta.1 (Unreleased) + + #### Features Added ++* Added new `CosmosItemsDataSource.readManyByPartitionKey` Spark function to execute bulk queries by a list of pk-values with better efficiency. See [PR 48801](https://github.com/Azure/azure-sdk-for-java/pull/48801) + + #### Breaking Changes + +diff --git a/sdk/cosmos/azure-cosmos-spark_3-4_2-12/CHANGELOG.md b/sdk/cosmos/azure-cosmos-spark_3-4_2-12/CHANGELOG.md +index c9097e749f0..3b2c7ce36db 100644 +--- a/sdk/cosmos/azure-cosmos-spark_3-4_2-12/CHANGELOG.md ++++ b/sdk/cosmos/azure-cosmos-spark_3-4_2-12/CHANGELOG.md +@@ -3,6 +3,7 @@ + ### 4.47.0-beta.1 (Unreleased) + + #### Features Added ++* Added new `CosmosItemsDataSource.readManyByPartitionKey` Spark function to execute bulk queries by a list of pk-values with better efficiency. See [PR 48801](https://github.com/Azure/azure-sdk-for-java/pull/48801) + + #### Breaking Changes + +diff --git a/sdk/cosmos/azure-cosmos-spark_3-5_2-12/CHANGELOG.md b/sdk/cosmos/azure-cosmos-spark_3-5_2-12/CHANGELOG.md +index f5eac38bdb7..2240a48b165 100644 +--- a/sdk/cosmos/azure-cosmos-spark_3-5_2-12/CHANGELOG.md ++++ b/sdk/cosmos/azure-cosmos-spark_3-5_2-12/CHANGELOG.md +@@ -3,6 +3,7 @@ + ### 4.47.0-beta.1 (Unreleased) + + #### Features Added ++* Added new `CosmosItemsDataSource.readManyByPartitionKey` Spark function to execute bulk queries by a list of pk-values with better efficiency. See [PR 48801](https://github.com/Azure/azure-sdk-for-java/pull/48801) + + #### Breaking Changes + +diff --git a/sdk/cosmos/azure-cosmos-spark_3-5_2-13/CHANGELOG.md b/sdk/cosmos/azure-cosmos-spark_3-5_2-13/CHANGELOG.md +index 919d7fbfa32..20a3e3a61bd 100644 +--- a/sdk/cosmos/azure-cosmos-spark_3-5_2-13/CHANGELOG.md ++++ b/sdk/cosmos/azure-cosmos-spark_3-5_2-13/CHANGELOG.md +@@ -3,6 +3,7 @@ + ### 4.47.0-beta.1 (Unreleased) + + #### Features Added ++* Added new `CosmosItemsDataSource.readManyByPartitionKey` Spark function to execute bulk queries by a list of pk-values with better efficiency. See [PR 48801](https://github.com/Azure/azure-sdk-for-java/pull/48801) + + #### Breaking Changes + +diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosConfig.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosConfig.scala +index 951f4735444..e1b8f0b51f8 100644 +--- a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosConfig.scala ++++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosConfig.scala +@@ -92,6 +92,7 @@ private[spark] object CosmosConfigNames { + val ReadPartitioningFeedRangeFilter = "spark.cosmos.partitioning.feedRangeFilter" + val ReadRuntimeFilteringEnabled = "spark.cosmos.read.runtimeFiltering.enabled" + val ReadManyFilteringEnabled = "spark.cosmos.read.readManyFiltering.enabled" ++ val ReadManyByPkNullHandling = "spark.cosmos.read.readManyByPk.nullHandling" + val ViewsRepositoryPath = "spark.cosmos.views.repositoryPath" + val DiagnosticsMode = "spark.cosmos.diagnostics" + val DiagnosticsSamplingMaxCount = "spark.cosmos.diagnostics.sampling.maxCount" +@@ -226,6 +227,7 @@ private[spark] object CosmosConfigNames { + ReadPartitioningFeedRangeFilter, + ReadRuntimeFilteringEnabled, + ReadManyFilteringEnabled, ++ ReadManyByPkNullHandling, + ViewsRepositoryPath, + DiagnosticsMode, + DiagnosticsSamplingIntervalInSeconds, +@@ -1042,7 +1044,8 @@ private case class CosmosReadConfig(readConsistencyStrategy: ReadConsistencyStra + throughputControlConfig: Option[CosmosThroughputControlConfig] = None, + runtimeFilteringEnabled: Boolean, + readManyFilteringConfig: CosmosReadManyFilteringConfig, +- responseContinuationTokenLimitInKb: Option[Int] = None) ++ responseContinuationTokenLimitInKb: Option[Int] = None, ++ readManyByPkTreatNullAsNone: Boolean = false) + + private object SchemaConversionModes extends Enumeration { + type SchemaConversionMode = Value +@@ -1136,6 +1139,18 @@ private object CosmosReadConfig { + helpMessage = " Indicates whether dynamic partition pruning filters will be pushed down when applicable." + ) + ++ private val ReadManyByPkNullHandling = CosmosConfigEntry[String]( ++ key = CosmosConfigNames.ReadManyByPkNullHandling, ++ mandatory = false, ++ defaultValue = Some("Null"), ++ parseFromStringFunction = value => value, ++ helpMessage = "Determines how null values in hierarchical partition key components are treated " + ++ "for readManyByPartitionKey. 'Null' (default) maps null to a JSON null value via addNullValue(), " + ++ "which is appropriate when the document field exists with an explicit null value. " + ++ "'None' maps null to PartitionKey.NONE via addNoneValue(), which should only be used when the " + ++ "partition key path does not exist at all in the document." ++ ) ++ + def parseCosmosReadConfig(cfg: Map[String, String]): CosmosReadConfig = { + val forceEventualConsistency = CosmosConfigEntry.parse(cfg, ForceEventualConsistency) + val readConsistencyStrategyOverride = CosmosConfigEntry.parse(cfg, ReadConsistencyStrategyOverride) +@@ -1158,6 +1173,8 @@ private object CosmosReadConfig { + val throughputControlConfigOpt = CosmosThroughputControlConfig.parseThroughputControlConfig(cfg) + val runtimeFilteringEnabled = CosmosConfigEntry.parse(cfg, ReadRuntimeFilteringEnabled) + val readManyFilteringConfig = CosmosReadManyFilteringConfig.parseCosmosReadManyFilterConfig(cfg) ++ val readManyByPkNullHandling = CosmosConfigEntry.parse(cfg, ReadManyByPkNullHandling) ++ val readManyByPkTreatNullAsNone = readManyByPkNullHandling.getOrElse("Null").equalsIgnoreCase("None") + + val effectiveReadConsistencyStrategy = if (readConsistencyStrategyOverride.getOrElse(ReadConsistencyStrategy.DEFAULT) != ReadConsistencyStrategy.DEFAULT) { + readConsistencyStrategyOverride.get +@@ -1189,7 +1206,8 @@ private object CosmosReadConfig { + throughputControlConfigOpt, + runtimeFilteringEnabled.get, + readManyFilteringConfig, +- responseContinuationTokenLimitInKb) ++ responseContinuationTokenLimitInKb, ++ readManyByPkTreatNullAsNone) + } + } + +diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosConstants.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosConstants.scala +index 9ece4741652..00761f23d39 100644 +--- a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosConstants.scala ++++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosConstants.scala +@@ -45,6 +45,7 @@ private[cosmos] object CosmosConstants { + val Id = "id" + val ETag = "_etag" + val ItemIdentity = "_itemIdentity" ++ val PartitionKeyIdentity = "_partitionKeyIdentity" + } + + object StatusCodes { +diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosItemsDataSource.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosItemsDataSource.scala +index a35cff27af6..86ef865bcb8 100644 +--- a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosItemsDataSource.scala ++++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosItemsDataSource.scala +@@ -2,9 +2,10 @@ + // Licensed under the MIT License. + package com.azure.cosmos.spark + +-import com.azure.cosmos.models.{CosmosItemIdentity, PartitionKey} ++import com.azure.cosmos.models.{CosmosItemIdentity, PartitionKey, PartitionKeyBuilder} + import com.azure.cosmos.spark.CosmosPredicates.assertOnSparkDriver + import com.azure.cosmos.spark.diagnostics.BasicLoggingTrait ++import com.azure.cosmos.{SparkBridgeInternal} + import org.apache.spark.sql.{DataFrame, Row, SparkSession} + + import java.util +@@ -112,4 +113,127 @@ object CosmosItemsDataSource { + + readManyReader.readMany(df.rdd, readManyFilterExtraction) + } ++ ++ def readManyByPartitionKey(df: DataFrame, userConfig: java.util.Map[String, String]): DataFrame = { ++ readManyByPartitionKey(df, userConfig, null) ++ } ++ ++ def readManyByPartitionKey( ++ df: DataFrame, ++ userConfig: java.util.Map[String, String], ++ userProvidedSchema: StructType): DataFrame = { ++ ++ val readManyReader = new CosmosReadManyByPartitionKeyReader( ++ userProvidedSchema, ++ userConfig.asScala.toMap) ++ ++ // Option 1: Look for the _partitionKeyIdentity column (produced by GetCosmosPartitionKeyValue UDF) ++ val pkIdentityFieldExtraction = df ++ .schema ++ .find(field => field.name.equals(CosmosConstants.Properties.PartitionKeyIdentity) && field.dataType.equals(StringType)) ++ .map(field => (row: Row) => { ++ val rawValue = row.getString(row.fieldIndex(field.name)) ++ CosmosPartitionKeyHelper.tryParsePartitionKey(rawValue) ++ .getOrElse(throw new IllegalArgumentException( ++ s"Invalid _partitionKeyIdentity value in row: '$rawValue'. " + ++ "Expected format: pk([...json...])")) ++ }) ++ ++ // Option 2: Detect PK columns by matching the container's partition key paths against the DataFrame schema ++ val pkColumnExtraction: Option[Row => PartitionKey] = if (pkIdentityFieldExtraction.isDefined) { ++ None // no need to resolve PK paths - _partitionKeyIdentity column takes precedence ++ } else { ++ val effectiveConfig = CosmosConfig.getEffectiveConfig( ++ databaseName = None, ++ containerName = None, ++ userConfig.asScala.toMap) ++ val readConfig = CosmosReadConfig.parseCosmosReadConfig(effectiveConfig) ++ val containerConfig = CosmosContainerConfig.parseCosmosContainerConfig(effectiveConfig) ++ val sparkEnvironmentInfo = CosmosClientConfiguration.getSparkEnvironmentInfo(None) ++ val calledFrom = s"CosmosItemsDataSource.readManyByPartitionKey" ++ val treatNullAsNone = readConfig.readManyByPkTreatNullAsNone ++ ++ val pkPaths = Loan( ++ List[Option[CosmosClientCacheItem]]( ++ Some( ++ CosmosClientCache( ++ CosmosClientConfiguration( ++ effectiveConfig, ++ readConsistencyStrategy = readConfig.readConsistencyStrategy, ++ sparkEnvironmentInfo), ++ None, ++ calledFrom)), ++ ThroughputControlHelper.getThroughputControlClientCacheItem( ++ effectiveConfig, ++ calledFrom, ++ None, ++ sparkEnvironmentInfo) ++ )) ++ .to(clientCacheItems => { ++ val container = ++ ThroughputControlHelper.getContainer( ++ effectiveConfig, ++ containerConfig, ++ clientCacheItems(0).get, ++ clientCacheItems(1)) ++ ++ val pkDefinition = SparkBridgeInternal ++ .getContainerPropertiesFromCollectionCache(container) ++ .getPartitionKeyDefinition ++ ++ pkDefinition.getPaths.asScala.map(_.stripPrefix("/")).toList ++ }) ++ ++ // Check if ALL PK path columns exist in the DataFrame schema ++ val dfFieldNames = df.schema.fieldNames.toSet ++ val allPkColumnsPresent = pkPaths.forall(path => dfFieldNames.contains(path)) ++ ++ if (allPkColumnsPresent && pkPaths.nonEmpty) { ++ // pkPaths already defined above ++ Some((row: Row) => { ++ if (pkPaths.size == 1) { ++ // Single partition key ++ buildPartitionKey(row.getAs[Any](pkPaths.head), treatNullAsNone) ++ } else { ++ // Hierarchical partition key ΓÇö build level by level ++ val builder = new PartitionKeyBuilder() ++ for (path <- pkPaths) { ++ addPartitionKeyComponent(builder, row.getAs[Any](path), treatNullAsNone) ++ } ++ builder.build() ++ } ++ }) ++ } else { ++ None ++ } ++ } ++ ++ val pkExtraction = pkIdentityFieldExtraction ++ .orElse(pkColumnExtraction) ++ .getOrElse( ++ throw new IllegalArgumentException( ++ "Cannot determine partition key extraction from the input DataFrame. " + ++ "Either add a '_partitionKeyIdentity' column (using the GetCosmosPartitionKeyValue UDF) " + ++ "or ensure the DataFrame contains columns matching the container's partition key paths.")) ++ ++ readManyReader.readManyByPartitionKey(df.rdd, pkExtraction) ++ } ++ ++ private def addPartitionKeyComponent(builder: PartitionKeyBuilder, value: Any, treatNullAsNone: Boolean): Unit = { ++ value match { ++ case s: String => builder.add(s) ++ case n: Number => builder.add(n.doubleValue()) ++ case b: Boolean => builder.add(b) ++ case null => ++ if (treatNullAsNone) builder.addNoneValue() ++ else builder.addNullValue() ++ case other => builder.add(other.toString) ++ } ++ } ++ ++ private def buildPartitionKey(value: Any, treatNullAsNone: Boolean): PartitionKey = { ++ val builder = new PartitionKeyBuilder() ++ addPartitionKeyComponent(builder, value, treatNullAsNone) ++ builder.build() ++ } + } +diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosPartitionKeyHelper.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosPartitionKeyHelper.scala +new file mode 100644 +index 00000000000..27776f5c3de +--- /dev/null ++++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosPartitionKeyHelper.scala +@@ -0,0 +1,45 @@ ++// Copyright (c) Microsoft Corporation. All rights reserved. ++// Licensed under the MIT License. ++ ++package com.azure.cosmos.spark ++ ++import com.azure.cosmos.implementation.routing.PartitionKeyInternal ++import com.azure.cosmos.implementation.{ImplementationBridgeHelpers, Utils} ++import com.azure.cosmos.models.PartitionKey ++import com.azure.cosmos.spark.diagnostics.BasicLoggingTrait ++ ++import java.util ++ ++// scalastyle:off underscore.import ++import scala.collection.JavaConverters._ ++// scalastyle:on underscore.import ++ ++private[spark] object CosmosPartitionKeyHelper extends BasicLoggingTrait { ++ // pattern will be recognized ++ // pk(partitionKeyValue) ++ // ++ // (?i) : The whole matching is case-insensitive ++ // pk[(](.*)[)]: partitionKey Value ++ private val cosmosPartitionKeyStringRegx = """(?i)pk[(](.*)[)]""".r ++ private val objectMapper = Utils.getSimpleObjectMapper ++ ++ def getCosmosPartitionKeyValueString(partitionKeyValue: List[Object]): String = { ++ s"pk(${objectMapper.writeValueAsString(partitionKeyValue.asJava)})" ++ } ++ ++ def tryParsePartitionKey(cosmosPartitionKeyString: String): Option[PartitionKey] = { ++ cosmosPartitionKeyString match { ++ case cosmosPartitionKeyStringRegx(pkValue) => ++ scala.util.Try(Utils.parse(pkValue, classOf[Object])).toOption.flatMap { ++ case arrayList: util.ArrayList[Object @unchecked] => ++ Some( ++ ImplementationBridgeHelpers ++ .PartitionKeyHelper ++ .getPartitionKeyAccessor ++ .toPartitionKey(PartitionKeyInternal.fromObjectArray(arrayList.toArray, false))) ++ case other => Some(new PartitionKey(other)) ++ } ++ case _ => None ++ } ++ } ++} +diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosReadManyByPartitionKeyReader.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosReadManyByPartitionKeyReader.scala +new file mode 100644 +index 00000000000..91f3a56bc66 +--- /dev/null ++++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosReadManyByPartitionKeyReader.scala +@@ -0,0 +1,150 @@ ++// Copyright (c) Microsoft Corporation. All rights reserved. ++// Licensed under the MIT License. ++package com.azure.cosmos.spark ++ ++import com.azure.cosmos.{CosmosException, ReadConsistencyStrategy} ++import com.azure.cosmos.implementation.{CosmosClientMetadataCachesSnapshot, UUIDs} ++import com.azure.cosmos.models.PartitionKey ++import com.azure.cosmos.spark.CosmosPredicates.assertOnSparkDriver ++import com.azure.cosmos.spark.diagnostics.{BasicLoggingTrait, DiagnosticsContext} ++import com.fasterxml.jackson.databind.node.ObjectNode ++import org.apache.spark.TaskContext ++import org.apache.spark.broadcast.Broadcast ++import org.apache.spark.rdd.RDD ++import org.apache.spark.sql.{DataFrame, Row, SparkSession} ++import org.apache.spark.sql.types.StructType ++ ++import java.util.UUID ++ ++private[spark] class CosmosReadManyByPartitionKeyReader( ++ val userProvidedSchema: StructType, ++ val userConfig: Map[String, String] ++ ) extends BasicLoggingTrait with Serializable { ++ val effectiveUserConfig: Map[String, String] = CosmosConfig.getEffectiveConfig( ++ databaseName = None, ++ containerName = None, ++ userConfig) ++ ++ val clientConfig: CosmosAccountConfig = CosmosAccountConfig.parseCosmosAccountConfig(effectiveUserConfig) ++ val readConfig: CosmosReadConfig = CosmosReadConfig.parseCosmosReadConfig(effectiveUserConfig) ++ val cosmosContainerConfig: CosmosContainerConfig = ++ CosmosContainerConfig.parseCosmosContainerConfig(effectiveUserConfig) ++ //scalastyle:off multiple.string.literals ++ val tableName: String = s"com.azure.cosmos.spark.items.${clientConfig.accountName}." + ++ s"${cosmosContainerConfig.database}.${cosmosContainerConfig.container}" ++ private lazy val sparkSession = { ++ assertOnSparkDriver() ++ SparkSession.active ++ } ++ val sparkEnvironmentInfo: String = CosmosClientConfiguration.getSparkEnvironmentInfo(Some(sparkSession)) ++ logTrace(s"Instantiated ${this.getClass.getSimpleName} for $tableName") ++ ++ private[spark] def initializeAndBroadcastCosmosClientStatesForContainer(): Broadcast[CosmosClientMetadataCachesSnapshots] = { ++ val calledFrom = s"CosmosReadManyByPartitionKeyReader($tableName).initializeAndBroadcastCosmosClientStateForContainer" ++ Loan( ++ List[Option[CosmosClientCacheItem]]( ++ Some( ++ CosmosClientCache( ++ CosmosClientConfiguration( ++ effectiveUserConfig, ++ readConsistencyStrategy = readConfig.readConsistencyStrategy, ++ sparkEnvironmentInfo), ++ None, ++ calledFrom)), ++ ThroughputControlHelper.getThroughputControlClientCacheItem( ++ effectiveUserConfig, ++ calledFrom, ++ None, ++ sparkEnvironmentInfo) ++ )) ++ .to(clientCacheItems => { ++ val container = ++ ThroughputControlHelper.getContainer( ++ effectiveUserConfig, ++ cosmosContainerConfig, ++ clientCacheItems(0).get, ++ clientCacheItems(1)) ++ try { ++ container.readItem( ++ UUIDs.nonBlockingRandomUUID().toString, ++ new PartitionKey(UUIDs.nonBlockingRandomUUID().toString), ++ classOf[ObjectNode]) ++ .block() ++ } catch { ++ case _: CosmosException => None ++ } ++ ++ val state = new CosmosClientMetadataCachesSnapshot() ++ state.serialize(clientCacheItems(0).get.cosmosClient) ++ ++ var throughputControlState: Option[CosmosClientMetadataCachesSnapshot] = None ++ if (clientCacheItems(1).isDefined) { ++ throughputControlState = Some(new CosmosClientMetadataCachesSnapshot()) ++ throughputControlState.get.serialize(clientCacheItems(1).get.cosmosClient) ++ } ++ ++ val metadataSnapshots = CosmosClientMetadataCachesSnapshots(state, throughputControlState) ++ sparkSession.sparkContext.broadcast(metadataSnapshots) ++ }) ++ } ++ ++ def readManyByPartitionKey(inputRdd: RDD[Row], pkExtraction: Row => PartitionKey): DataFrame = { ++ val correlationActivityId = UUIDs.nonBlockingRandomUUID() ++ val calledFrom = s"CosmosReadManyByPartitionKeyReader.readManyByPartitionKey($correlationActivityId)" ++ val schema = Loan( ++ List[Option[CosmosClientCacheItem]]( ++ Some(CosmosClientCache( ++ CosmosClientConfiguration( ++ effectiveUserConfig, ++ readConsistencyStrategy = readConfig.readConsistencyStrategy, ++ sparkEnvironmentInfo), ++ None, ++ calledFrom ++ )), ++ ThroughputControlHelper.getThroughputControlClientCacheItem( ++ effectiveUserConfig, ++ calledFrom, ++ None, ++ sparkEnvironmentInfo) ++ )) ++ .to(clientCacheItems => Option.apply(userProvidedSchema).getOrElse( ++ CosmosTableSchemaInferrer.inferSchema( ++ clientCacheItems(0).get, ++ clientCacheItems(1), ++ effectiveUserConfig, ++ ItemsTable.defaultSchemaForInferenceDisabled))) ++ ++ val clientStates = initializeAndBroadcastCosmosClientStatesForContainer ++ ++ sparkSession.sqlContext.createDataFrame( ++ inputRdd.mapPartitionsWithIndex( ++ (partitionIndex: Int, rowIterator: Iterator[Row]) => { ++ val pkIterator: Iterator[PartitionKey] = rowIterator ++ .map(row => pkExtraction.apply(row)) ++ ++ logInfo(s"Creating an ItemsPartitionReaderWithReadManyByPartitionKey for Activity $correlationActivityId to read for " ++ + s"input partition [$partitionIndex] ${tableName}") ++ ++ val reader = new ItemsPartitionReaderWithReadManyByPartitionKey( ++ effectiveUserConfig, ++ CosmosReadManyHelper.FullRangeFeedRange, ++ schema, ++ DiagnosticsContext(correlationActivityId, partitionIndex.toString), ++ clientStates, ++ DiagnosticsConfig.parseDiagnosticsConfig(effectiveUserConfig), ++ sparkEnvironmentInfo, ++ TaskContext.get, ++ pkIterator) ++ ++ new Iterator[Row] { ++ override def hasNext: Boolean = reader.next() ++ ++ override def next(): Row = reader.getCurrentRow() ++ } ++ }, ++ preservesPartitioning = true ++ ), ++ schema) ++ } ++} ++ +diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/ItemsPartitionReaderWithReadManyByPartitionKey.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/ItemsPartitionReaderWithReadManyByPartitionKey.scala +new file mode 100644 +index 00000000000..c67cc9c10be +--- /dev/null ++++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/ItemsPartitionReaderWithReadManyByPartitionKey.scala +@@ -0,0 +1,249 @@ ++// Copyright (c) Microsoft Corporation. All rights reserved. ++// Licensed under the MIT License. ++ ++package com.azure.cosmos.spark ++ ++import com.azure.cosmos.{CosmosAsyncContainer, CosmosEndToEndOperationLatencyPolicyConfigBuilder, CosmosItemSerializerNoExceptionWrapping, SparkBridgeInternal} ++import com.azure.cosmos.implementation.spark.OperationContextAndListenerTuple ++import com.azure.cosmos.implementation.{ImplementationBridgeHelpers, ObjectNodeMap, SparkRowItem, Utils} ++import com.azure.cosmos.models.{CosmosReadManyRequestOptions, ModelBridgeInternal, PartitionKey, PartitionKeyDefinition, SqlQuerySpec} ++import com.azure.cosmos.spark.BulkWriter.getThreadInfo ++import com.azure.cosmos.spark.CosmosTableSchemaInferrer.IdAttributeName ++import com.azure.cosmos.spark.diagnostics.{DetailedFeedDiagnosticsProvider, DiagnosticsContext, DiagnosticsLoader, LoggerHelper, SparkTaskContext} ++import com.fasterxml.jackson.databind.node.ObjectNode ++import org.apache.spark.TaskContext ++import org.apache.spark.broadcast.Broadcast ++import org.apache.spark.sql.Row ++import org.apache.spark.sql.catalyst.InternalRow ++import org.apache.spark.sql.catalyst.encoders.ExpressionEncoder ++import org.apache.spark.sql.connector.read.PartitionReader ++import org.apache.spark.sql.types.StructType ++ ++import java.util ++ ++// scalastyle:off underscore.import ++import scala.collection.JavaConverters._ ++// scalastyle:on underscore.import ++ ++private[spark] case class ItemsPartitionReaderWithReadManyByPartitionKey ++( ++ config: Map[String, String], ++ feedRange: NormalizedRange, ++ readSchema: StructType, ++ diagnosticsContext: DiagnosticsContext, ++ cosmosClientStateHandles: Broadcast[CosmosClientMetadataCachesSnapshots], ++ diagnosticsConfig: DiagnosticsConfig, ++ sparkEnvironmentInfo: String, ++ taskContext: TaskContext, ++ readManyPartitionKeys: Iterator[PartitionKey] ++) ++ extends PartitionReader[InternalRow] { ++ ++ private lazy val log = LoggerHelper.getLogger(diagnosticsConfig, this.getClass) ++ ++ private val readManyOptions = new CosmosReadManyRequestOptions() ++ private val readManyOptionsImpl = ImplementationBridgeHelpers ++ .CosmosReadManyRequestOptionsHelper ++ .getCosmosReadManyRequestOptionsAccessor ++ .getImpl(readManyOptions) ++ ++ private val readConfig = CosmosReadConfig.parseCosmosReadConfig(config) ++ ThroughputControlHelper.populateThroughputControlGroupName(readManyOptionsImpl, readConfig.throughputControlConfig) ++ ++ private val operationContext = { ++ assert(taskContext != null) ++ ++ SparkTaskContext(diagnosticsContext.correlationActivityId, ++ taskContext.stageId(), ++ taskContext.partitionId(), ++ taskContext.taskAttemptId(), ++ feedRange.toString) ++ } ++ ++ private val operationContextAndListenerTuple: Option[OperationContextAndListenerTuple] = { ++ if (diagnosticsConfig.mode.isDefined) { ++ val listener = ++ DiagnosticsLoader.getDiagnosticsProvider(diagnosticsConfig).getLogger(this.getClass) ++ ++ val ctxAndListener = new OperationContextAndListenerTuple(operationContext, listener) ++ ++ readManyOptionsImpl ++ .setOperationContextAndListenerTuple(ctxAndListener) ++ ++ Some(ctxAndListener) ++ } else { ++ None ++ } ++ } ++ ++ log.logTrace(s"Instantiated ${this.getClass.getSimpleName}, Context: ${operationContext.toString} $getThreadInfo") ++ ++ private val containerTargetConfig = CosmosContainerConfig.parseCosmosContainerConfig(config) ++ ++ log.logInfo(s"Using ReadManyByPartitionKey from feed range $feedRange of " + ++ s"container ${containerTargetConfig.database}.${containerTargetConfig.container} - " + ++ s"correlationActivityId ${diagnosticsContext.correlationActivityId}, " + ++ s"Context: ${operationContext.toString} $getThreadInfo") ++ ++ private val clientCacheItem = CosmosClientCache( ++ CosmosClientConfiguration(config, readConfig.readConsistencyStrategy, sparkEnvironmentInfo), ++ Some(cosmosClientStateHandles.value.cosmosClientMetadataCaches), ++ s"ItemsPartitionReaderWithReadManyByPartitionKey($feedRange, ${containerTargetConfig.database}.${containerTargetConfig.container})" ++ ) ++ ++ private val throughputControlClientCacheItemOpt = ++ ThroughputControlHelper.getThroughputControlClientCacheItem( ++ config, ++ clientCacheItem.context, ++ Some(cosmosClientStateHandles), ++ sparkEnvironmentInfo) ++ ++ private val cosmosAsyncContainer = ++ ThroughputControlHelper.getContainer( ++ config, ++ containerTargetConfig, ++ clientCacheItem, ++ throughputControlClientCacheItemOpt) ++ ++ private val partitionKeyDefinition: PartitionKeyDefinition = { ++ TransientErrorsRetryPolicy.executeWithRetry(() => { ++ SparkBridgeInternal ++ .getContainerPropertiesFromCollectionCache(cosmosAsyncContainer).getPartitionKeyDefinition ++ }) ++ } ++ ++ private val cosmosSerializationConfig = CosmosSerializationConfig.parseSerializationConfig(config) ++ private val cosmosRowConverter = CosmosRowConverter.get(cosmosSerializationConfig) ++ ++ readManyOptionsImpl ++ .setCustomItemSerializer( ++ new CosmosItemSerializerNoExceptionWrapping { ++ override def serialize[T](item: T): util.Map[String, AnyRef] = { ++ throw new UnsupportedOperationException( ++ s"Serialization is not supported by the custom item serializer in " + ++ s"ItemsPartitionReaderWithReadManyByPartitionKey; this serializer is intended " + ++ s"for deserializing read-many responses into SparkRowItem only. " + ++ s"Unexpected item type: ${if (item == null) "null" else item.getClass.getName}" ++ ) ++ } ++ ++ override def deserialize[T](jsonNodeMap: util.Map[String, AnyRef], classType: Class[T]): T = { ++ if (jsonNodeMap == null) { ++ throw new IllegalStateException("The 'jsonNodeMap' should never be null here.") ++ } ++ ++ if (classType != classOf[SparkRowItem]) { ++ throw new IllegalStateException("The 'classType' must be 'classOf[SparkRowItem])' here.") ++ } ++ ++ val objectNode: ObjectNode = jsonNodeMap match { ++ case map: ObjectNodeMap => ++ map.getObjectNode ++ case _ => ++ Utils.getSimpleObjectMapper.convertValue(jsonNodeMap, classOf[ObjectNode]) ++ } ++ ++ val partitionKey = PartitionKeyHelper.getPartitionKeyPath(objectNode, partitionKeyDefinition) ++ ++ val row = cosmosRowConverter.fromObjectNodeToRow(readSchema, ++ objectNode, ++ readConfig.schemaConversionMode) ++ ++ SparkRowItem(row, getPartitionKeyForFeedDiagnostics(partitionKey)).asInstanceOf[T] ++ } ++ } ++ ) ++ ++ // Collect all PK values upfront ΓÇö readManyByPartitionKey needs the full list to ++ // group by physical partition and issue parallel queries. ++ // Deduplicate by PK string representation ΓÇö safe because the list size is bounded ++ // by the per-call limit of the readManyByPartitionKey API. ++ private lazy val pkList = { ++ val seen = new java.util.LinkedHashMap[String, PartitionKey]() ++ readManyPartitionKeys.foreach(pk => seen.putIfAbsent(pk.toString, pk)) ++ new java.util.ArrayList[PartitionKey](seen.values()) ++ } ++ ++ private val endToEndTimeoutPolicy = ++ new CosmosEndToEndOperationLatencyPolicyConfigBuilder( ++ java.time.Duration.ofSeconds(CosmosConstants.readOperationEndToEndTimeoutInSeconds)) ++ .enable(true) ++ .build ++ ++ readManyOptionsImpl.setCosmosEndToEndOperationLatencyPolicyConfig(endToEndTimeoutPolicy) ++ ++ private trait CloseableSparkRowItemIterator { ++ def hasNext: Boolean ++ def next(): SparkRowItem ++ def close(): Unit ++ } ++ ++ private object EmptySparkRowItemIterator extends CloseableSparkRowItemIterator { ++ override def hasNext: Boolean = false ++ ++ override def next(): SparkRowItem = { ++ throw new java.util.NoSuchElementException("No items available for empty partition-key list.") ++ } ++ ++ override def close(): Unit = {} ++ } ++ ++ // Batch partition keys and retry each batch independently on transient I/O errors. ++ // This avoids the continuation-token problem with TransientIOErrorsRetryingIterator ++ // where a retry would re-read all data from scratch, causing silent data duplication. ++ private lazy val iterator: CloseableSparkRowItemIterator = ++ if (pkList.isEmpty) { ++ EmptySparkRowItemIterator ++ } else { ++ new CloseableSparkRowItemIterator { ++ private val delegate = new TransientIOErrorsRetryingReadManyByPartitionKeyIterator[SparkRowItem]( ++ cosmosAsyncContainer, ++ pkList, ++ readConfig.customQuery.map(_.toSqlQuerySpec), ++ readManyOptions, ++ readConfig.maxItemCount, ++ readConfig.prefetchBufferSize, ++ operationContextAndListenerTuple, ++ classOf[SparkRowItem] ++ ) ++ ++ override def hasNext: Boolean = delegate.hasNext ++ ++ override def next(): SparkRowItem = delegate.next() ++ ++ override def close(): Unit = delegate.close() ++ } ++ } ++ ++ private val rowSerializer: ExpressionEncoder.Serializer[Row] = RowSerializerPool.getOrCreateSerializer(readSchema) ++ ++ private def shouldLogDetailedFeedDiagnostics(): Boolean = { ++ diagnosticsConfig.mode.isDefined && ++ diagnosticsConfig.mode.get.equalsIgnoreCase(classOf[DetailedFeedDiagnosticsProvider].getName) ++ } ++ ++ private def getPartitionKeyForFeedDiagnostics(pkValue: PartitionKey): Option[PartitionKey] = { ++ if (shouldLogDetailedFeedDiagnostics()) { ++ Some(pkValue) ++ } else { ++ None ++ } ++ } ++ ++ override def next(): Boolean = iterator.hasNext ++ ++ override def get(): InternalRow = { ++ cosmosRowConverter.fromRowToInternalRow(iterator.next().row, rowSerializer) ++ } ++ ++ def getCurrentRow(): Row = iterator.next().row ++ ++ override def close(): Unit = { ++ this.iterator.close() ++ RowSerializerPool.returnSerializerToPool(readSchema, rowSerializer) ++ clientCacheItem.close() ++ if (throughputControlClientCacheItemOpt.isDefined) { ++ throughputControlClientCacheItemOpt.get.close() ++ } ++ } ++} +diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/TransientIOErrorsRetryingReadManyByPartitionKeyIterator.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/TransientIOErrorsRetryingReadManyByPartitionKeyIterator.scala +new file mode 100644 +index 00000000000..dcfdf4f9353 +--- /dev/null ++++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/TransientIOErrorsRetryingReadManyByPartitionKeyIterator.scala +@@ -0,0 +1,259 @@ ++// Copyright (c) Microsoft Corporation. All rights reserved. ++// Licensed under the MIT License. ++ ++package com.azure.cosmos.spark ++ ++import com.azure.cosmos.{CosmosAsyncContainer, CosmosException} ++import com.azure.cosmos.implementation.OperationCancelledException ++import com.azure.cosmos.implementation.spark.OperationContextAndListenerTuple ++import com.azure.cosmos.models.{CosmosReadManyRequestOptions, FeedResponse, PartitionKey, SqlQuerySpec} ++import com.azure.cosmos.spark.diagnostics.BasicLoggingTrait ++import com.azure.cosmos.util.CosmosPagedIterable ++ ++import java.util.concurrent.{ExecutorService, SynchronousQueue, ThreadPoolExecutor, TimeUnit, TimeoutException} ++import java.util.concurrent.atomic.AtomicLong ++import scala.concurrent.{Await, ExecutionContext, Future} ++import scala.util.Random ++import scala.util.control.Breaks ++ ++// scalastyle:off underscore.import ++import scala.collection.JavaConverters._ ++// scalastyle:on underscore.import ++ ++/** ++ * Retry-safe iterator for readManyByPartitionKey that batches partition keys and lazily ++ * iterates pages within each batch via CosmosPagedIterable ΓÇö consistent with how ++ * TransientIOErrorsRetryingIterator handles normal queries. On transient I/O errors the ++ * current batch's flux is recreated and pages already consumed are replayed, avoiding ++ * the memory overhead of collectList and matching the query iterator's structure. ++ */ ++private[spark] class TransientIOErrorsRetryingReadManyByPartitionKeyIterator[TSparkRow] ++( ++ val container: CosmosAsyncContainer, ++ val partitionKeys: java.util.List[PartitionKey], ++ val customQuery: Option[SqlQuerySpec], ++ val queryOptions: CosmosReadManyRequestOptions, ++ val pageSize: Int, ++ val pagePrefetchBufferSize: Int, ++ val operationContextAndListener: Option[OperationContextAndListenerTuple], ++ val classType: Class[TSparkRow] ++) extends BufferedIterator[TSparkRow] with BasicLoggingTrait with AutoCloseable { ++ ++ private[spark] var maxRetryIntervalInMs = CosmosConstants.maxRetryIntervalForTransientFailuresInMs ++ private[spark] var maxRetryCount = CosmosConstants.maxRetryCountForTransientFailures ++ ++ private val maxPageRetrievalTimeout = scala.concurrent.duration.FiniteDuration( ++ 5 + CosmosConstants.readOperationEndToEndTimeoutInSeconds, ++ scala.concurrent.duration.SECONDS) ++ ++ private val rnd = Random ++ private val retryCount = new AtomicLong(0) ++ private lazy val operationContextString = operationContextAndListener match { ++ case Some(o) => if (o.getOperationContext != null) { ++ o.getOperationContext.toString ++ } else { ++ "n/a" ++ } ++ case None => "n/a" ++ } ++ ++ private[spark] var currentFeedResponseIterator: Option[BufferedIterator[FeedResponse[TSparkRow]]] = None ++ private[spark] var currentItemIterator: Option[BufferedIterator[TSparkRow]] = None ++ ++ private val pkBatchIterator = partitionKeys.asScala.iterator.grouped(pageSize) ++ // Track the current batch so we can replay it on retry ++ private var currentBatch: Option[java.util.List[PartitionKey]] = None ++ ++ override def hasNext: Boolean = { ++ executeWithRetry("hasNextInternal", () => hasNextInternal) ++ } ++ ++ private def hasNextInternal: Boolean = { ++ var returnValue: Option[Boolean] = None ++ ++ while (returnValue.isEmpty) { ++ returnValue = hasNextInternalCore ++ } ++ ++ returnValue.get ++ } ++ ++ private def hasNextInternalCore: Option[Boolean] = { ++ if (hasBufferedNext) { ++ Some(true) ++ } else { ++ val feedResponseIterator = currentFeedResponseIterator match { ++ case Some(existing) => existing ++ case None => ++ // Need a new feed response iterator ΓÇö either for the current batch (on retry) ++ // or for the next batch ++ val batch = currentBatch match { ++ case Some(b) => b // retry of current batch ++ case None => ++ if (pkBatchIterator.hasNext) { ++ val nextBatch = new java.util.ArrayList[PartitionKey](pkBatchIterator.next().toList.asJava) ++ currentBatch = Some(nextBatch) ++ nextBatch ++ } else { ++ return Some(false) // no more batches ++ } ++ } ++ ++ val pagedFlux = customQuery match { ++ case Some(query) => ++ container.readManyByPartitionKey(batch, query, queryOptions, classType) ++ case None => ++ container.readManyByPartitionKey(batch, queryOptions, classType) ++ } ++ ++ currentFeedResponseIterator = Some( ++ new CosmosPagedIterable[TSparkRow]( ++ pagedFlux, ++ pageSize, ++ pagePrefetchBufferSize ++ ) ++ .iterableByPage() ++ .iterator ++ .asScala ++ .buffered ++ ) ++ ++ currentFeedResponseIterator.get ++ } ++ ++ val hasNext: Boolean = try { ++ Await.result( ++ Future { ++ feedResponseIterator.hasNext ++ }(TransientIOErrorsRetryingReadManyByPartitionKeyIterator.executionContext), ++ maxPageRetrievalTimeout) ++ } catch { ++ case endToEndTimeoutException: OperationCancelledException => ++ val message = s"End-to-end timeout hit when trying to retrieve the next page. " + ++ s"Context: $operationContextString" ++ logError(message, throwable = endToEndTimeoutException) ++ throw endToEndTimeoutException ++ ++ case timeoutException: TimeoutException => ++ val message = s"Attempting to retrieve the next page timed out. " + ++ s"Context: $operationContextString" ++ logError(message, timeoutException) ++ val exception = new OperationCancelledException(message, null) ++ exception.setStackTrace(timeoutException.getStackTrace) ++ throw exception ++ ++ case other: Throwable => throw other ++ } ++ ++ if (hasNext) { ++ val feedResponse = feedResponseIterator.next() ++ if (operationContextAndListener.isDefined) { ++ operationContextAndListener.get.getOperationListener.feedResponseProcessedListener( ++ operationContextAndListener.get.getOperationContext, ++ feedResponse) ++ } ++ val iteratorCandidate = feedResponse.getResults.iterator().asScala.buffered ++ ++ if (iteratorCandidate.hasNext) { ++ currentItemIterator = Some(iteratorCandidate) ++ Some(true) ++ } else { ++ // empty page interleaved ΓÇö try again ++ None ++ } ++ } else { ++ // Current batch's flux is exhausted ΓÇö move to next batch ++ currentBatch = None ++ currentFeedResponseIterator = None ++ None ++ } ++ } ++ } ++ ++ private def hasBufferedNext: Boolean = { ++ currentItemIterator match { ++ case Some(iterator) => if (iterator.hasNext) { ++ true ++ } else { ++ currentItemIterator = None ++ false ++ } ++ case None => false ++ } ++ } ++ ++ override def next(): TSparkRow = { ++ currentItemIterator.get.next() ++ } ++ ++ override def head(): TSparkRow = { ++ currentItemIterator.get.head ++ } ++ ++ private[spark] def executeWithRetry[T](methodName: String, func: () => T): T = { ++ val loop = new Breaks() ++ var returnValue: Option[T] = None ++ ++ loop.breakable { ++ while (true) { ++ val retryIntervalInMs = rnd.nextInt(maxRetryIntervalInMs) ++ ++ try { ++ returnValue = Some(func()) ++ retryCount.set(0) ++ loop.break ++ } ++ catch { ++ case cosmosException: CosmosException => ++ if (Exceptions.canBeTransientFailure(cosmosException.getStatusCode, cosmosException.getSubStatusCode)) { ++ val retryCountSnapshot = retryCount.incrementAndGet() ++ if (retryCountSnapshot > maxRetryCount) { ++ logError( ++ s"Too many transient failure retry attempts in " + ++ s"TransientIOErrorsRetryingReadManyByPartitionKeyIterator.$methodName", ++ cosmosException) ++ throw cosmosException ++ } else { ++ logWarning( ++ s"Transient failure handled in " + ++ s"TransientIOErrorsRetryingReadManyByPartitionKeyIterator.$methodName -" + ++ s" will be retried (attempt#$retryCountSnapshot) in ${retryIntervalInMs}ms", ++ cosmosException) ++ } ++ } else { ++ throw cosmosException ++ } ++ case other: Throwable => throw other ++ } ++ ++ // Reset iterators but keep currentBatch so the batch is replayed ++ currentItemIterator = None ++ currentFeedResponseIterator = None ++ Thread.sleep(retryIntervalInMs) ++ } ++ } ++ ++ returnValue.get ++ } ++ ++ override def close(): Unit = { ++ currentItemIterator = None ++ currentFeedResponseIterator = None ++ } ++} ++ ++private object TransientIOErrorsRetryingReadManyByPartitionKeyIterator extends BasicLoggingTrait { ++ private val maxConcurrency = SparkUtils.getNumberOfHostCPUCores ++ ++ val executorService: ExecutorService = new ThreadPoolExecutor( ++ maxConcurrency, ++ maxConcurrency, ++ 0L, ++ TimeUnit.MILLISECONDS, ++ new SynchronousQueue(), ++ SparkUtils.daemonThreadFactory(), ++ new ThreadPoolExecutor.CallerRunsPolicy() ++ ) ++ ++ val executionContext: ExecutionContext = ExecutionContext.fromExecutorService(executorService) ++} +diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/udf/GetCosmosPartitionKeyValue.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/udf/GetCosmosPartitionKeyValue.scala +new file mode 100644 +index 00000000000..a58d5b723b8 +--- /dev/null ++++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/udf/GetCosmosPartitionKeyValue.scala +@@ -0,0 +1,25 @@ ++// Copyright (c) Microsoft Corporation. All rights reserved. ++// Licensed under the MIT License. ++ ++package com.azure.cosmos.spark.udf ++ ++import com.azure.cosmos.spark.CosmosPartitionKeyHelper ++import com.azure.cosmos.spark.CosmosPredicates.requireNotNull ++import org.apache.spark.sql.api.java.UDF1 ++ ++@SerialVersionUID(1L) ++class GetCosmosPartitionKeyValue extends UDF1[Object, String] { ++ override def call ++ ( ++ partitionKeyValue: Object ++ ): String = { ++ requireNotNull(partitionKeyValue, "partitionKeyValue") ++ ++ partitionKeyValue match { ++ // for subpartitions case - Seq covers both WrappedArray (Scala 2.12) and ArraySeq (Scala 2.13) ++ case seq: Seq[Any] => ++ CosmosPartitionKeyHelper.getCosmosPartitionKeyValueString(seq.map(_.asInstanceOf[Object]).toList) ++ case _ => CosmosPartitionKeyHelper.getCosmosPartitionKeyValueString(List(partitionKeyValue)) ++ } ++ } ++} +diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/CosmosConfigSpec.scala b/sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/CosmosConfigSpec.scala +index 17f75e45a74..17a298d6213 100644 +--- a/sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/CosmosConfigSpec.scala ++++ b/sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/CosmosConfigSpec.scala +@@ -457,6 +457,7 @@ class CosmosConfigSpec extends UnitSpec with BasicLoggingTrait { + config.runtimeFilteringEnabled shouldBe true + config.readManyFilteringConfig.readManyFilteringEnabled shouldBe false + config.readManyFilteringConfig.readManyFilterProperty shouldEqual "_itemIdentity" ++ config.readManyByPkTreatNullAsNone shouldBe false + + userConfig = Map( + "spark.cosmos.read.forceEventualConsistency" -> "false", +@@ -630,6 +631,47 @@ class CosmosConfigSpec extends UnitSpec with BasicLoggingTrait { + config.customQuery.get.queryText shouldBe queryText + } + ++ it should "parse readManyByPk nullHandling configuration" in { ++ // Default (not specified) should treat null as JSON null (addNullValue) ++ var userConfig = Map( ++ "spark.cosmos.read.forceEventualConsistency" -> "false" ++ ) ++ var config = CosmosReadConfig.parseCosmosReadConfig(userConfig) ++ config.readManyByPkTreatNullAsNone shouldBe false ++ ++ // Explicit "Null" should treat null as JSON null (addNullValue) ++ userConfig = Map( ++ "spark.cosmos.read.forceEventualConsistency" -> "false", ++ "spark.cosmos.read.readManyByPk.nullHandling" -> "Null" ++ ) ++ config = CosmosReadConfig.parseCosmosReadConfig(userConfig) ++ config.readManyByPkTreatNullAsNone shouldBe false ++ ++ // Case-insensitive "null" ++ userConfig = Map( ++ "spark.cosmos.read.forceEventualConsistency" -> "false", ++ "spark.cosmos.read.readManyByPk.nullHandling" -> "null" ++ ) ++ config = CosmosReadConfig.parseCosmosReadConfig(userConfig) ++ config.readManyByPkTreatNullAsNone shouldBe false ++ ++ // "None" should treat null as PartitionKey.NONE (addNoneValue) ++ userConfig = Map( ++ "spark.cosmos.read.forceEventualConsistency" -> "false", ++ "spark.cosmos.read.readManyByPk.nullHandling" -> "None" ++ ) ++ config = CosmosReadConfig.parseCosmosReadConfig(userConfig) ++ config.readManyByPkTreatNullAsNone shouldBe true ++ ++ // Case-insensitive "none" ++ userConfig = Map( ++ "spark.cosmos.read.forceEventualConsistency" -> "false", ++ "spark.cosmos.read.readManyByPk.nullHandling" -> "none" ++ ) ++ config = CosmosReadConfig.parseCosmosReadConfig(userConfig) ++ config.readManyByPkTreatNullAsNone shouldBe true ++ } ++ + it should "throw on invalid read configuration" in { + val userConfig = Map( + "spark.cosmos.read.schemaConversionMode" -> "not a valid value" +diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/CosmosPartitionKeyHelperSpec.scala b/sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/CosmosPartitionKeyHelperSpec.scala +new file mode 100644 +index 00000000000..1ac40e39584 +--- /dev/null ++++ b/sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/CosmosPartitionKeyHelperSpec.scala +@@ -0,0 +1,104 @@ ++// Copyright (c) Microsoft Corporation. All rights reserved. ++// Licensed under the MIT License. ++ ++package com.azure.cosmos.spark ++ ++import com.azure.cosmos.models.{PartitionKey, PartitionKeyBuilder} ++ ++class CosmosPartitionKeyHelperSpec extends UnitSpec { ++ //scalastyle:off multiple.string.literals ++ //scalastyle:off magic.number ++ ++ it should "return the correct partition key value string for single PK" in { ++ val pkString = CosmosPartitionKeyHelper.getCosmosPartitionKeyValueString(List("pk1")) ++ pkString shouldEqual "pk([\"pk1\"])" ++ } ++ ++ it should "return the correct partition key value string for HPK" in { ++ val pkString = CosmosPartitionKeyHelper.getCosmosPartitionKeyValueString(List("city1", "zip1")) ++ pkString shouldEqual "pk([\"city1\",\"zip1\"])" ++ } ++ ++ it should "return the correct partition key value string for 3-level HPK" in { ++ val pkString = CosmosPartitionKeyHelper.getCosmosPartitionKeyValueString(List("a", "b", "c")) ++ pkString shouldEqual "pk([\"a\",\"b\",\"c\"])" ++ } ++ ++ it should "parse valid single PK string" in { ++ val pk = CosmosPartitionKeyHelper.tryParsePartitionKey("pk([\"myPkValue\"])") ++ pk.isDefined shouldBe true ++ pk.get shouldEqual new PartitionKey("myPkValue") ++ } ++ ++ it should "parse valid HPK string" in { ++ val pk = CosmosPartitionKeyHelper.tryParsePartitionKey("pk([\"city1\",\"zip1\"])") ++ pk.isDefined shouldBe true ++ val expected = new PartitionKeyBuilder().add("city1").add("zip1").build() ++ pk.get shouldEqual expected ++ } ++ ++ it should "parse valid 3-level HPK string" in { ++ val pk = CosmosPartitionKeyHelper.tryParsePartitionKey("pk([\"a\",\"b\",\"c\"])") ++ pk.isDefined shouldBe true ++ val expected = new PartitionKeyBuilder().add("a").add("b").add("c").build() ++ pk.get shouldEqual expected ++ } ++ ++ it should "roundtrip single PK" in { ++ val original = "pk([\"roundtrip\"])" ++ val parsed = CosmosPartitionKeyHelper.tryParsePartitionKey(original) ++ parsed.isDefined shouldBe true ++ val serialized = CosmosPartitionKeyHelper.getCosmosPartitionKeyValueString(List("roundtrip")) ++ serialized shouldEqual original ++ } ++ ++ it should "roundtrip HPK" in { ++ val original = "pk([\"city\",\"zip\"])" ++ val parsed = CosmosPartitionKeyHelper.tryParsePartitionKey(original) ++ parsed.isDefined shouldBe true ++ val serialized = CosmosPartitionKeyHelper.getCosmosPartitionKeyValueString(List("city", "zip")) ++ serialized shouldEqual original ++ } ++ ++ it should "return None for malformed string" in { ++ val pk = CosmosPartitionKeyHelper.tryParsePartitionKey("invalid_format") ++ pk.isDefined shouldBe false ++ } ++ ++ it should "return None for missing pk prefix" in { ++ val pk = CosmosPartitionKeyHelper.tryParsePartitionKey("[\"value\"]") ++ pk.isDefined shouldBe false ++ } ++ ++ it should "be case-insensitive for parsing" in { ++ val pk = CosmosPartitionKeyHelper.tryParsePartitionKey("PK([\"value\"])") ++ pk.isDefined shouldBe true ++ pk.get shouldEqual new PartitionKey("value") ++ } ++ ++ ++ it should "return None for malformed JSON inside pk() wrapper" in { ++ // Invalid JSON that would cause JsonProcessingException ++ val pk = CosmosPartitionKeyHelper.tryParsePartitionKey("pk({invalid json})") ++ pk.isDefined shouldBe false ++ } ++ ++ it should "return None for truncated JSON inside pk() wrapper" in { ++ val pk = CosmosPartitionKeyHelper.tryParsePartitionKey("pk([\"unterminated)") ++ pk.isDefined shouldBe false ++ } ++ ++ it should "produce different partition keys for addNullValue vs addNoneValue in HPK" in { ++ // addNullValue represents an explicit JSON null for a field that exists with value null ++ val pkWithNull = new PartitionKeyBuilder().add("Redmond").addNullValue().build() ++ ++ // addNoneValue represents PartitionKey.NONE, meaning the field is absent/undefined ++ val pkWithNone = new PartitionKeyBuilder().add("Redmond").addNoneValue().build() ++ ++ // These MUST produce different partition key hashes and route to different physical partitions ++ pkWithNull should not equal pkWithNone ++ } ++ ++ //scalastyle:on multiple.string.literals ++ //scalastyle:on magic.number ++} +diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/ItemsPartitionReaderWithReadManyByPartitionKeyITest.scala b/sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/ItemsPartitionReaderWithReadManyByPartitionKeyITest.scala +new file mode 100644 +index 00000000000..5c2d7b59836 +--- /dev/null ++++ b/sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/ItemsPartitionReaderWithReadManyByPartitionKeyITest.scala +@@ -0,0 +1,158 @@ ++// Copyright (c) Microsoft Corporation. All rights reserved. ++// Licensed under the MIT License. ++ ++package com.azure.cosmos.spark ++ ++import com.azure.cosmos.implementation.{CosmosClientMetadataCachesSnapshot, TestConfigurations, Utils} ++import com.azure.cosmos.models.PartitionKey ++import com.azure.cosmos.spark.diagnostics.DiagnosticsContext ++import com.fasterxml.jackson.databind.node.ObjectNode ++import org.apache.spark.MockTaskContext ++import org.apache.spark.broadcast.Broadcast ++import org.apache.spark.sql.types.{StringType, StructField, StructType} ++ ++import java.util.UUID ++import scala.collection.mutable.ListBuffer ++ ++class ItemsPartitionReaderWithReadManyByPartitionKeyITest ++ extends IntegrationSpec ++ with Spark ++ with AutoCleanableCosmosContainersWithPkAsPartitionKey { ++ private val idProperty = "id" ++ private val pkProperty = "pk" ++ ++ //scalastyle:off multiple.string.literals ++ //scalastyle:off magic.number ++ ++ it should "be able to retrieve all items for given partition keys" in { ++ val container = cosmosClient.getDatabase(cosmosDatabase).getContainer(cosmosContainersWithPkAsPartitionKey) ++ ++ // Create items with known PK values ++ val partitionKeyDefinition = container.read().block().getProperties.getPartitionKeyDefinition ++ val allItemsByPk = scala.collection.mutable.Map[String, ListBuffer[ObjectNode]]() ++ val pkValues = List("pkA", "pkB", "pkC") ++ ++ for (pk <- pkValues) { ++ allItemsByPk(pk) = ListBuffer[ObjectNode]() ++ for (_ <- 1 to 5) { ++ val objectNode = Utils.getSimpleObjectMapper.createObjectNode() ++ objectNode.put(idProperty, UUID.randomUUID().toString) ++ objectNode.put(pkProperty, pk) ++ container.createItem(objectNode).block() ++ allItemsByPk(pk) += objectNode ++ } ++ } ++ ++ val config = Map( ++ "spark.cosmos.accountEndpoint" -> TestConfigurations.HOST, ++ "spark.cosmos.accountKey" -> TestConfigurations.MASTER_KEY, ++ "spark.cosmos.database" -> cosmosDatabase, ++ "spark.cosmos.container" -> cosmosContainersWithPkAsPartitionKey, ++ "spark.cosmos.read.inferSchema.enabled" -> "true", ++ "spark.cosmos.applicationName" -> "ReadManyByPKTest" ++ ) ++ ++ val readSchema = StructType(Seq( ++ StructField(idProperty, StringType, false), ++ StructField(pkProperty, StringType, false) ++ )) ++ ++ val diagnosticsContext = DiagnosticsContext(UUID.randomUUID(), "") ++ val diagnosticsConfig = DiagnosticsConfig.parseDiagnosticsConfig(config) ++ val cosmosClientMetadataCachesSnapshots = getCosmosClientMetadataCachesSnapshots() ++ ++ // Read items for pkA and pkB (not pkC) ++ val targetPks = List("pkA", "pkB") ++ val pkIterator = targetPks.map(pk => new PartitionKey(pk)).iterator ++ ++ val reader = ItemsPartitionReaderWithReadManyByPartitionKey( ++ config, ++ NormalizedRange("", "FF"), ++ readSchema, ++ diagnosticsContext, ++ cosmosClientMetadataCachesSnapshots, ++ diagnosticsConfig, ++ "", ++ MockTaskContext.mockTaskContext(), ++ pkIterator ++ ) ++ ++ val cosmosRowConverter = CosmosRowConverter.get(CosmosSerializationConfig.parseSerializationConfig(config)) ++ val itemsReadFromReader = ListBuffer[ObjectNode]() ++ while (reader.next()) { ++ itemsReadFromReader += cosmosRowConverter.fromInternalRowToObjectNode(reader.get(), readSchema) ++ } ++ ++ // Should have 10 items (5 for pkA + 5 for pkB) ++ itemsReadFromReader.size shouldEqual 10 ++ ++ // All items should be from pkA or pkB ++ itemsReadFromReader.foreach(item => { ++ val pk = item.get(pkProperty).asText() ++ targetPks should contain(pk) ++ }) ++ ++ // Validate all expected IDs are present ++ val expectedIds = (allItemsByPk("pkA") ++ allItemsByPk("pkB")).map(_.get(idProperty).asText()).toSet ++ val actualIds = itemsReadFromReader.map(_.get(idProperty).asText()).toSet ++ actualIds shouldEqual expectedIds ++ ++ reader.close() ++ } ++ ++ it should "return empty results for non-existent partition keys" in { ++ val config = Map( ++ "spark.cosmos.accountEndpoint" -> TestConfigurations.HOST, ++ "spark.cosmos.accountKey" -> TestConfigurations.MASTER_KEY, ++ "spark.cosmos.database" -> cosmosDatabase, ++ "spark.cosmos.container" -> cosmosContainersWithPkAsPartitionKey, ++ "spark.cosmos.read.inferSchema.enabled" -> "true", ++ "spark.cosmos.applicationName" -> "ReadManyByPKEmptyTest" ++ ) ++ ++ val readSchema = StructType(Seq( ++ StructField(idProperty, StringType, false), ++ StructField(pkProperty, StringType, false) ++ )) ++ ++ val diagnosticsContext = DiagnosticsContext(UUID.randomUUID(), "") ++ val diagnosticsConfig = DiagnosticsConfig.parseDiagnosticsConfig(config) ++ val cosmosClientMetadataCachesSnapshots = getCosmosClientMetadataCachesSnapshots() ++ ++ val pkIterator = List(new PartitionKey("nonExistentPk")).iterator ++ ++ val reader = ItemsPartitionReaderWithReadManyByPartitionKey( ++ config, ++ NormalizedRange("", "FF"), ++ readSchema, ++ diagnosticsContext, ++ cosmosClientMetadataCachesSnapshots, ++ diagnosticsConfig, ++ "", ++ MockTaskContext.mockTaskContext(), ++ pkIterator ++ ) ++ ++ val itemsReadFromReader = ListBuffer[ObjectNode]() ++ val cosmosRowConverter = CosmosRowConverter.get(CosmosSerializationConfig.parseSerializationConfig(config)) ++ while (reader.next()) { ++ itemsReadFromReader += cosmosRowConverter.fromInternalRowToObjectNode(reader.get(), readSchema) ++ } ++ ++ itemsReadFromReader.size shouldEqual 0 ++ reader.close() ++ } ++ ++ private def getCosmosClientMetadataCachesSnapshots(): Broadcast[CosmosClientMetadataCachesSnapshots] = { ++ val cosmosClientMetadataCachesSnapshot = new CosmosClientMetadataCachesSnapshot() ++ cosmosClientMetadataCachesSnapshot.serialize(cosmosClient) ++ ++ spark.sparkContext.broadcast( ++ CosmosClientMetadataCachesSnapshots( ++ cosmosClientMetadataCachesSnapshot, ++ Option.empty[CosmosClientMetadataCachesSnapshot])) ++ } ++ ++ //scalastyle:on multiple.string.literals ++ //scalastyle:on magic.number ++} +diff --git a/sdk/cosmos/azure-cosmos-spark_4-0_2-13/CHANGELOG.md b/sdk/cosmos/azure-cosmos-spark_4-0_2-13/CHANGELOG.md +index 3972ae6aeb9..d8368be6a0d 100644 +--- a/sdk/cosmos/azure-cosmos-spark_4-0_2-13/CHANGELOG.md ++++ b/sdk/cosmos/azure-cosmos-spark_4-0_2-13/CHANGELOG.md +@@ -3,6 +3,7 @@ + ### 4.47.0-beta.1 (Unreleased) + + #### Features Added ++* Added new `CosmosItemsDataSource.readManyByPartitionKey` Spark function to execute bulk queries by a list of pk-values with better efficiency. See [PR 48801](https://github.com/Azure/azure-sdk-for-java/pull/48801) + + #### Breaking Changes + +diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/ReadManyByPartitionKeyTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/ReadManyByPartitionKeyTest.java +new file mode 100644 +index 00000000000..2c26d564ed2 +--- /dev/null ++++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/ReadManyByPartitionKeyTest.java +@@ -0,0 +1,462 @@ ++/* ++ * Copyright (c) Microsoft Corporation. All rights reserved. ++ * Licensed under the MIT License. ++ */ ++ ++package com.azure.cosmos; ++ ++import com.azure.cosmos.models.CosmosContainerProperties; ++import com.azure.cosmos.models.CosmosItemRequestOptions; ++import com.azure.cosmos.models.FeedResponse; ++import com.azure.cosmos.models.PartitionKey; ++import com.azure.cosmos.models.PartitionKeyBuilder; ++import com.azure.cosmos.models.PartitionKeyDefinition; ++import com.azure.cosmos.models.PartitionKeyDefinitionVersion; ++import com.azure.cosmos.models.PartitionKind; ++import com.azure.cosmos.models.SqlParameter; ++import com.azure.cosmos.models.SqlQuerySpec; ++import com.azure.cosmos.rx.TestSuiteBase; ++import com.azure.cosmos.util.CosmosPagedIterable; ++import com.fasterxml.jackson.databind.node.ObjectNode; ++import org.testng.annotations.AfterClass; ++import org.testng.annotations.BeforeClass; ++import org.testng.annotations.Factory; ++import org.testng.annotations.Test; ++ ++import java.util.ArrayList; ++import java.util.Arrays; ++import java.util.Collections; ++import java.util.List; ++import java.util.UUID; ++import java.util.stream.Collectors; ++ ++import static org.assertj.core.api.Assertions.assertThat; ++import static org.assertj.core.api.Assertions.fail; ++ ++public class ReadManyByPartitionKeyTest extends TestSuiteBase { ++ ++ private String preExistingDatabaseId = CosmosDatabaseForTest.generateId(); ++ private CosmosClient client; ++ private CosmosDatabase createdDatabase; ++ ++ // Single PK container (/mypk) ++ private CosmosContainer singlePkContainer; ++ ++ // HPK container (/city, /zipcode, /areaCode) ++ private CosmosContainer multiHashContainer; ++ ++ @Factory(dataProvider = "clientBuilders") ++ public ReadManyByPartitionKeyTest(CosmosClientBuilder clientBuilder) { ++ super(clientBuilder); ++ } ++ ++ @BeforeClass(groups = {"emulator"}, timeOut = SETUP_TIMEOUT) ++ public void before_ReadManyByPartitionKeyTest() { ++ client = getClientBuilder().buildClient(); ++ createdDatabase = createSyncDatabase(client, preExistingDatabaseId); ++ ++ // Single PK container ++ String singlePkContainerName = UUID.randomUUID().toString(); ++ CosmosContainerProperties singlePkProps = new CosmosContainerProperties(singlePkContainerName, "/mypk"); ++ createdDatabase.createContainer(singlePkProps); ++ singlePkContainer = createdDatabase.getContainer(singlePkContainerName); ++ ++ // HPK container ++ String multiHashContainerName = UUID.randomUUID().toString(); ++ PartitionKeyDefinition hpkDef = new PartitionKeyDefinition(); ++ hpkDef.setKind(PartitionKind.MULTI_HASH); ++ hpkDef.setVersion(PartitionKeyDefinitionVersion.V2); ++ ArrayList paths = new ArrayList<>(); ++ paths.add("/city"); ++ paths.add("/zipcode"); ++ paths.add("/areaCode"); ++ hpkDef.setPaths(paths); ++ ++ CosmosContainerProperties hpkProps = new CosmosContainerProperties(multiHashContainerName, hpkDef); ++ createdDatabase.createContainer(hpkProps); ++ multiHashContainer = createdDatabase.getContainer(multiHashContainerName); ++ } ++ ++ @AfterClass(groups = {"emulator"}, timeOut = SHUTDOWN_TIMEOUT, alwaysRun = true) ++ public void afterClass() { ++ safeDeleteSyncDatabase(createdDatabase); ++ safeCloseSyncClient(client); ++ } ++ ++ //region Single PK tests ++ ++ @Test(groups = {"emulator"}, timeOut = TIMEOUT) ++ public void singlePk_readManyByPartitionKey_basic() { ++ // Create items with different PKs ++ List items = createSinglePkItems("pk1", 3); ++ items.addAll(createSinglePkItems("pk2", 2)); ++ items.addAll(createSinglePkItems("pk3", 4)); ++ ++ // Read by 2 partition keys ++ List pkValues = Arrays.asList( ++ new PartitionKey("pk1"), ++ new PartitionKey("pk2")); ++ ++ CosmosPagedIterable results = singlePkContainer.readManyByPartitionKey(pkValues, ObjectNode.class); ++ List resultList = results.stream().collect(Collectors.toList()); ++ ++ assertThat(resultList).hasSize(5); // 3 + 2 ++ resultList.forEach(item -> { ++ String pk = item.get("mypk").asText(); ++ assertThat(pk).isIn("pk1", "pk2"); ++ }); ++ ++ // Cleanup ++ cleanupContainer(singlePkContainer); ++ } ++ ++ @Test(groups = {"emulator"}, timeOut = TIMEOUT) ++ public void singlePk_readManyByPartitionKey_withProjection() { ++ List items = createSinglePkItems("pkProj", 2); ++ ++ List pkValues = Collections.singletonList(new PartitionKey("pkProj")); ++ SqlQuerySpec customQuery = new SqlQuerySpec("SELECT c.id, c.mypk FROM c"); ++ ++ CosmosPagedIterable results = singlePkContainer.readManyByPartitionKey( ++ pkValues, customQuery, null, ObjectNode.class); ++ List resultList = results.stream().collect(Collectors.toList()); ++ ++ assertThat(resultList).hasSize(2); ++ // Should only have id and mypk fields (plus system properties) ++ resultList.forEach(item -> { ++ assertThat(item.has("id")).isTrue(); ++ assertThat(item.has("mypk")).isTrue(); ++ }); ++ ++ cleanupContainer(singlePkContainer); ++ } ++ ++ @Test(groups = {"emulator"}, timeOut = TIMEOUT) ++ public void singlePk_readManyByPartitionKey_withAdditionalFilter() { ++ // Create items with different "status" values ++ createSinglePkItemsWithStatus("pkFilter", "active", 3); ++ createSinglePkItemsWithStatus("pkFilter", "inactive", 2); ++ ++ List pkValues = Collections.singletonList(new PartitionKey("pkFilter")); ++ SqlQuerySpec customQuery = new SqlQuerySpec( ++ "SELECT * FROM c WHERE c.status = @status", ++ Arrays.asList(new SqlParameter("@status", "active"))); ++ ++ CosmosPagedIterable results = singlePkContainer.readManyByPartitionKey( ++ pkValues, customQuery, null, ObjectNode.class); ++ List resultList = results.stream().collect(Collectors.toList()); ++ ++ assertThat(resultList).hasSize(3); ++ resultList.forEach(item -> { ++ assertThat(item.get("status").asText()).isEqualTo("active"); ++ }); ++ ++ cleanupContainer(singlePkContainer); ++ } ++ ++ @Test(groups = {"emulator"}, timeOut = TIMEOUT) ++ public void singlePk_readManyByPartitionKey_emptyResults() { ++ List pkValues = Collections.singletonList(new PartitionKey("nonExistent")); ++ ++ CosmosPagedIterable results = singlePkContainer.readManyByPartitionKey(pkValues, ObjectNode.class); ++ List resultList = results.stream().collect(Collectors.toList()); ++ ++ assertThat(resultList).isEmpty(); ++ } ++ ++ //endregion ++ ++ //region HPK tests ++ ++ @Test(groups = {"emulator"}, timeOut = TIMEOUT) ++ public void hpk_readManyByPartitionKey_fullPk() { ++ createHpkItems(); ++ ++ // Read by full PKs ++ List pkValues = Arrays.asList( ++ new PartitionKeyBuilder().add("Redmond").add("98053").add(1).build(), ++ new PartitionKeyBuilder().add("Pittsburgh").add("15232").add(2).build()); ++ ++ CosmosPagedIterable results = multiHashContainer.readManyByPartitionKey(pkValues, ObjectNode.class); ++ List resultList = results.stream().collect(Collectors.toList()); ++ ++ // Redmond/98053/1 has 2 items, Pittsburgh/15232/2 has 1 item ++ assertThat(resultList).hasSize(3); ++ ++ cleanupContainer(multiHashContainer); ++ } ++ ++ @Test(groups = {"emulator"}, timeOut = TIMEOUT) ++ public void hpk_readManyByPartitionKey_partialPk_singleLevel() { ++ createHpkItems(); ++ ++ // Read by partial PK (only city) ++ List pkValues = Collections.singletonList( ++ new PartitionKeyBuilder().add("Redmond").build()); ++ ++ CosmosPagedIterable results = multiHashContainer.readManyByPartitionKey(pkValues, ObjectNode.class); ++ List resultList = results.stream().collect(Collectors.toList()); ++ ++ // Redmond has 3 items total (2 with 98053/1 and 1 with 12345/1) ++ assertThat(resultList).hasSize(3); ++ resultList.forEach(item -> { ++ assertThat(item.get("city").asText()).isEqualTo("Redmond"); ++ }); ++ ++ cleanupContainer(multiHashContainer); ++ } ++ ++ @Test(groups = {"emulator"}, timeOut = TIMEOUT) ++ public void hpk_readManyByPartitionKey_partialPk_twoLevels() { ++ createHpkItems(); ++ ++ // Read by partial PK (city + zipcode) ++ List pkValues = Collections.singletonList( ++ new PartitionKeyBuilder().add("Redmond").add("98053").build()); ++ ++ CosmosPagedIterable results = multiHashContainer.readManyByPartitionKey(pkValues, ObjectNode.class); ++ List resultList = results.stream().collect(Collectors.toList()); ++ ++ // Redmond/98053 has 2 items ++ assertThat(resultList).hasSize(2); ++ resultList.forEach(item -> { ++ assertThat(item.get("city").asText()).isEqualTo("Redmond"); ++ assertThat(item.get("zipcode").asText()).isEqualTo("98053"); ++ }); ++ ++ cleanupContainer(multiHashContainer); ++ } ++ ++ @Test(groups = {"emulator"}, timeOut = TIMEOUT) ++ public void hpk_readManyByPartitionKey_withProjection() { ++ createHpkItems(); ++ ++ List pkValues = Collections.singletonList( ++ new PartitionKeyBuilder().add("Redmond").add("98053").add(1).build()); ++ ++ SqlQuerySpec customQuery = new SqlQuerySpec("SELECT c.id, c.city FROM c"); ++ ++ CosmosPagedIterable results = multiHashContainer.readManyByPartitionKey( ++ pkValues, customQuery, null, ObjectNode.class); ++ List resultList = results.stream().collect(Collectors.toList()); ++ ++ assertThat(resultList).hasSize(2); ++ ++ cleanupContainer(multiHashContainer); ++ } ++ ++ //endregion ++ ++ //region Negative/validation tests ++ ++ @Test(groups = {"emulator"}, timeOut = TIMEOUT) ++ public void rejectsAggregateQuery() { ++ List pkValues = Collections.singletonList(new PartitionKey("pk1")); ++ SqlQuerySpec aggregateQuery = new SqlQuerySpec("SELECT COUNT(1) FROM c"); ++ ++ try { ++ singlePkContainer.readManyByPartitionKey(pkValues, aggregateQuery, null, ObjectNode.class) ++ .stream().collect(Collectors.toList()); ++ fail("Should have thrown IllegalArgumentException for aggregate query"); ++ } catch (IllegalArgumentException e) { ++ assertThat(e.getMessage()).contains("aggregates"); ++ } ++ } ++ ++ @Test(groups = {"emulator"}, timeOut = TIMEOUT) ++ public void rejectsOrderByQuery() { ++ List pkValues = Collections.singletonList(new PartitionKey("pk1")); ++ SqlQuerySpec orderByQuery = new SqlQuerySpec("SELECT * FROM c ORDER BY c.id"); ++ ++ try { ++ singlePkContainer.readManyByPartitionKey(pkValues, orderByQuery, null, ObjectNode.class) ++ .stream().collect(Collectors.toList()); ++ fail("Should have thrown IllegalArgumentException for ORDER BY query"); ++ } catch (IllegalArgumentException e) { ++ assertThat(e.getMessage()).contains("ORDER BY"); ++ } ++ } ++ ++ @Test(groups = {"emulator"}, timeOut = TIMEOUT) ++ public void rejectsDistinctQuery() { ++ List pkValues = Collections.singletonList(new PartitionKey("pk1")); ++ SqlQuerySpec distinctQuery = new SqlQuerySpec("SELECT DISTINCT c.mypk FROM c"); ++ ++ try { ++ singlePkContainer.readManyByPartitionKey(pkValues, distinctQuery, null, ObjectNode.class) ++ .stream().collect(Collectors.toList()); ++ fail("Should have thrown IllegalArgumentException for DISTINCT query"); ++ } catch (IllegalArgumentException e) { ++ assertThat(e.getMessage()).contains("DISTINCT"); ++ } ++ } ++ ++ @Test(groups = {"emulator"}, timeOut = TIMEOUT) ++ public void rejectsGroupByQuery() { ++ List pkValues = Collections.singletonList(new PartitionKey("pk1")); ++ SqlQuerySpec groupByQuery = new SqlQuerySpec("SELECT c.mypk FROM c GROUP BY c.mypk"); ++ ++ try { ++ singlePkContainer.readManyByPartitionKey(pkValues, groupByQuery, null, ObjectNode.class) ++ .stream().collect(Collectors.toList()); ++ fail("Should have thrown IllegalArgumentException for GROUP BY query"); ++ } catch (IllegalArgumentException e) { ++ assertThat(e.getMessage()).contains("GROUP BY"); ++ } ++ } ++ ++ @Test(groups = {"emulator"}, timeOut = TIMEOUT) ++ public void rejectsGroupByWithAggregateQuery() { ++ List pkValues = Collections.singletonList(new PartitionKey("pk1")); ++ SqlQuerySpec groupByWithAggregateQuery = new SqlQuerySpec("SELECT c.mypk, COUNT(1) as cnt FROM c GROUP BY c.mypk"); ++ ++ try { ++ singlePkContainer.readManyByPartitionKey(pkValues, groupByWithAggregateQuery, null, ObjectNode.class) ++ .stream().collect(Collectors.toList()); ++ fail("Should have thrown IllegalArgumentException for GROUP BY with aggregate query"); ++ } catch (IllegalArgumentException e) { ++ assertThat(e.getMessage()).contains("GROUP BY"); ++ } ++ } ++ ++ @Test(groups = {"emulator"}, timeOut = TIMEOUT, expectedExceptions = IllegalArgumentException.class) ++ public void rejectsNullPartitionKeyList() { ++ singlePkContainer.readManyByPartitionKey((List) null, ObjectNode.class); ++ } ++ ++ @Test(groups = {"emulator"}, timeOut = TIMEOUT, expectedExceptions = IllegalArgumentException.class) ++ public void rejectsEmptyPartitionKeyList() { ++ singlePkContainer.readManyByPartitionKey(new ArrayList<>(), ObjectNode.class) ++ .stream().collect(Collectors.toList()); ++ } ++ ++ //endregion ++ ++ ++ //region Batch size tests (#10) ++ ++ @Test(groups = {"emulator"}, timeOut = TIMEOUT) ++ public void singlePk_readManyByPartitionKey_withSmallBatchSize() { ++ // Temporarily set batch size to 2 to exercise the batching/interleaving logic ++ String originalValue = System.getProperty("COSMOS.READ_MANY_BY_PK_MAX_BATCH_SIZE"); ++ try { ++ System.setProperty("COSMOS.READ_MANY_BY_PK_MAX_BATCH_SIZE", "2"); ++ ++ // Create items across 4 PKs (more than the batch size of 2) ++ List items = createSinglePkItems("batchPk1", 2); ++ items.addAll(createSinglePkItems("batchPk2", 2)); ++ items.addAll(createSinglePkItems("batchPk3", 2)); ++ items.addAll(createSinglePkItems("batchPk4", 2)); ++ ++ // Read all 4 PKs ΓÇö should be split into batches of 2 ++ List pkValues = Arrays.asList( ++ new PartitionKey("batchPk1"), ++ new PartitionKey("batchPk2"), ++ new PartitionKey("batchPk3"), ++ new PartitionKey("batchPk4")); ++ ++ CosmosPagedIterable results = singlePkContainer.readManyByPartitionKey(pkValues, ObjectNode.class); ++ List resultList = results.stream().collect(Collectors.toList()); ++ ++ assertThat(resultList).hasSize(8); // 2 items per PK * 4 PKs ++ resultList.forEach(item -> { ++ String pk = item.get("mypk").asText(); ++ assertThat(pk).isIn("batchPk1", "batchPk2", "batchPk3", "batchPk4"); ++ }); ++ ++ cleanupContainer(singlePkContainer); ++ } finally { ++ if (originalValue != null) { ++ System.setProperty("COSMOS.READ_MANY_BY_PK_MAX_BATCH_SIZE", originalValue); ++ } else { ++ System.clearProperty("COSMOS.READ_MANY_BY_PK_MAX_BATCH_SIZE"); ++ } ++ } ++ } ++ ++ //endregion ++ ++ //region Custom serializer regression tests (#5) ++ ++ @Test(groups = {"emulator"}, timeOut = TIMEOUT) ++ public void singlePk_readManyByPartitionKey_withRequestOptions() { ++ // This test ensures that request options (like throughput control settings) ++ // are properly propagated through the readManyByPartitionKey path. ++ // It acts as a regression test for the redundant options construction bug. ++ List items = createSinglePkItems("pkOpts", 3); ++ ++ List pkValues = Collections.singletonList(new PartitionKey("pkOpts")); ++ com.azure.cosmos.models.CosmosReadManyRequestOptions options = new com.azure.cosmos.models.CosmosReadManyRequestOptions(); ++ ++ CosmosPagedIterable results = singlePkContainer.readManyByPartitionKey( ++ pkValues, options, ObjectNode.class); ++ List resultList = results.stream().collect(Collectors.toList()); ++ ++ assertThat(resultList).hasSize(3); ++ ++ cleanupContainer(singlePkContainer); ++ } ++ ++ //endregion ++ ++ //region helper methods ++ ++ private List createSinglePkItems(String pkValue, int count) { ++ List items = new ArrayList<>(); ++ for (int i = 0; i < count; i++) { ++ ObjectNode item = com.azure.cosmos.implementation.Utils.getSimpleObjectMapper().createObjectNode(); ++ item.put("id", UUID.randomUUID().toString()); ++ item.put("mypk", pkValue); ++ singlePkContainer.createItem(item); ++ items.add(item); ++ } ++ return items; ++ } ++ ++ private List createSinglePkItemsWithStatus(String pkValue, String status, int count) { ++ List items = new ArrayList<>(); ++ for (int i = 0; i < count; i++) { ++ ObjectNode item = com.azure.cosmos.implementation.Utils.getSimpleObjectMapper().createObjectNode(); ++ item.put("id", UUID.randomUUID().toString()); ++ item.put("mypk", pkValue); ++ item.put("status", status); ++ singlePkContainer.createItem(item); ++ items.add(item); ++ } ++ return items; ++ } ++ ++ private void createHpkItems() { ++ // Same data as CosmosMultiHashTest.createItems() ++ createHpkItem("Redmond", "98053", 1); ++ createHpkItem("Redmond", "98053", 1); ++ createHpkItem("Pittsburgh", "15232", 2); ++ createHpkItem("Stonybrook", "11790", 3); ++ createHpkItem("Stonybrook", "11794", 3); ++ createHpkItem("Stonybrook", "11791", 3); ++ createHpkItem("Redmond", "12345", 1); ++ } ++ ++ private void createHpkItem(String city, String zipcode, int areaCode) { ++ ObjectNode item = com.azure.cosmos.implementation.Utils.getSimpleObjectMapper().createObjectNode(); ++ item.put("id", UUID.randomUUID().toString()); ++ item.put("city", city); ++ item.put("zipcode", zipcode); ++ item.put("areaCode", areaCode); ++ multiHashContainer.createItem(item); ++ } ++ ++ private void cleanupContainer(CosmosContainer container) { ++ CosmosPagedIterable allItems = container.queryItems( ++ "SELECT * FROM c", new com.azure.cosmos.models.CosmosQueryRequestOptions(), ObjectNode.class); ++ allItems.forEach(item -> { ++ try { ++ container.deleteItem(item, new CosmosItemRequestOptions()); ++ } catch (CosmosException e) { ++ // ignore cleanup failures ++ } ++ }); ++ } ++ ++ //endregion ++} +diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelperTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelperTest.java +new file mode 100644 +index 00000000000..95c109ba025 +--- /dev/null ++++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelperTest.java +@@ -0,0 +1,426 @@ ++/* ++ * Copyright (c) Microsoft Corporation. All rights reserved. ++ * Licensed under the MIT License. ++ */ ++ ++package com.azure.cosmos.implementation; ++ ++import com.azure.cosmos.models.PartitionKey; ++import com.azure.cosmos.models.PartitionKeyBuilder; ++import com.azure.cosmos.models.PartitionKeyDefinition; ++import com.azure.cosmos.models.PartitionKeyDefinitionVersion; ++import com.azure.cosmos.models.PartitionKind; ++import com.azure.cosmos.models.SqlParameter; ++import com.azure.cosmos.models.SqlQuerySpec; ++import org.testng.annotations.Test; ++ ++import java.util.ArrayList; ++import java.util.Arrays; ++import java.util.Collections; ++import java.util.List; ++import java.util.stream.Collectors; ++ ++import static org.assertj.core.api.Assertions.assertThat; ++ ++public class ReadManyByPartitionKeyQueryHelperTest { ++ ++ //region Single PK (HASH) tests ++ ++ @Test(groups = { "unit" }) ++ public void singlePk_defaultQuery_singleValue() { ++ PartitionKeyDefinition pkDef = createSinglePkDefinition("/mypk"); ++ List selectors = createSelectors(pkDef); ++ List pkValues = Collections.singletonList(new PartitionKey("pk1")); ++ ++ SqlQuerySpec result = ReadManyByPartitionKeyQueryHelper.createReadManyByPkQuerySpec( ++ "SELECT * FROM c", new ArrayList<>(), pkValues, selectors, pkDef); ++ ++ assertThat(result.getQueryText()).contains("SELECT * FROM c WHERE"); ++ assertThat(result.getQueryText()).contains("IN ("); ++ assertThat(result.getQueryText()).contains("@__rmPk_0"); ++ assertThat(result.getParameters()).hasSize(1); ++ assertThat(result.getParameters().get(0).getValue(Object.class)).isEqualTo("pk1"); ++ } ++ ++ @Test(groups = { "unit" }) ++ public void singlePk_defaultQuery_multipleValues() { ++ PartitionKeyDefinition pkDef = createSinglePkDefinition("/mypk"); ++ List selectors = createSelectors(pkDef); ++ List pkValues = Arrays.asList( ++ new PartitionKey("pk1"), ++ new PartitionKey("pk2"), ++ new PartitionKey("pk3")); ++ ++ SqlQuerySpec result = ReadManyByPartitionKeyQueryHelper.createReadManyByPkQuerySpec( ++ "SELECT * FROM c", new ArrayList<>(), pkValues, selectors, pkDef); ++ ++ assertThat(result.getQueryText()).contains("IN ("); ++ assertThat(result.getQueryText()).contains("@__rmPk_0"); ++ assertThat(result.getQueryText()).contains("@__rmPk_1"); ++ assertThat(result.getQueryText()).contains("@__rmPk_2"); ++ assertThat(result.getParameters()).hasSize(3); ++ } ++ ++ @Test(groups = { "unit" }) ++ public void singlePk_customQuery_noWhere() { ++ PartitionKeyDefinition pkDef = createSinglePkDefinition("/mypk"); ++ List selectors = createSelectors(pkDef); ++ List pkValues = Arrays.asList(new PartitionKey("pk1"), new PartitionKey("pk2")); ++ ++ SqlQuerySpec result = ReadManyByPartitionKeyQueryHelper.createReadManyByPkQuerySpec( ++ "SELECT c.name, c.age FROM c", new ArrayList<>(), pkValues, selectors, pkDef); ++ ++ assertThat(result.getQueryText()).startsWith("SELECT c.name, c.age FROM c WHERE"); ++ assertThat(result.getQueryText()).contains("IN ("); ++ } ++ ++ @Test(groups = { "unit" }) ++ public void singlePk_customQuery_withExistingWhere() { ++ PartitionKeyDefinition pkDef = createSinglePkDefinition("/mypk"); ++ List selectors = createSelectors(pkDef); ++ List pkValues = Collections.singletonList(new PartitionKey("pk1")); ++ ++ List baseParams = new ArrayList<>(); ++ baseParams.add(new SqlParameter("@minAge", 18)); ++ ++ SqlQuerySpec result = ReadManyByPartitionKeyQueryHelper.createReadManyByPkQuerySpec( ++ "SELECT * FROM c WHERE c.age > @minAge", baseParams, pkValues, selectors, pkDef); ++ ++ // Should AND the PK filter to the existing WHERE clause ++ assertThat(result.getQueryText()).contains("WHERE (c.age > @minAge) AND ("); ++ assertThat(result.getQueryText()).contains("IN ("); ++ assertThat(result.getParameters()).hasSize(2); // @minAge + @__rmPk_0 ++ assertThat(result.getParameters().get(0).getName()).isEqualTo("@minAge"); ++ } ++ ++ //endregion ++ ++ //region HPK (MULTI_HASH) tests ++ ++ @Test(groups = { "unit" }) ++ public void hpk_fullPk_defaultQuery() { ++ PartitionKeyDefinition pkDef = createMultiHashPkDefinition("/city", "/zipcode"); ++ List selectors = createSelectors(pkDef); ++ ++ PartitionKey pk = new PartitionKeyBuilder().add("Redmond").add("98052").build(); ++ List pkValues = Collections.singletonList(pk); ++ ++ SqlQuerySpec result = ReadManyByPartitionKeyQueryHelper.createReadManyByPkQuerySpec( ++ "SELECT * FROM c", new ArrayList<>(), pkValues, selectors, pkDef); ++ ++ assertThat(result.getQueryText()).contains("SELECT * FROM c WHERE"); ++ // Should use OR/AND pattern, not IN ++ assertThat(result.getQueryText()).doesNotContain("IN ("); ++ assertThat(result.getQueryText()).contains("c[\"city\"] = @__rmPk_0"); ++ assertThat(result.getQueryText()).contains("AND"); ++ assertThat(result.getQueryText()).contains("c[\"zipcode\"] = @__rmPk_1"); ++ assertThat(result.getParameters()).hasSize(2); ++ assertThat(result.getParameters().get(0).getValue(Object.class)).isEqualTo("Redmond"); ++ assertThat(result.getParameters().get(1).getValue(Object.class)).isEqualTo("98052"); ++ } ++ ++ @Test(groups = { "unit" }) ++ public void hpk_fullPk_multipleValues() { ++ PartitionKeyDefinition pkDef = createMultiHashPkDefinition("/city", "/zipcode"); ++ List selectors = createSelectors(pkDef); ++ ++ List pkValues = Arrays.asList( ++ new PartitionKeyBuilder().add("Redmond").add("98052").build(), ++ new PartitionKeyBuilder().add("Seattle").add("98101").build()); ++ ++ SqlQuerySpec result = ReadManyByPartitionKeyQueryHelper.createReadManyByPkQuerySpec( ++ "SELECT * FROM c", new ArrayList<>(), pkValues, selectors, pkDef); ++ ++ assertThat(result.getQueryText()).contains("OR"); ++ assertThat(result.getQueryText()).contains("c[\"city\"] = @__rmPk_0"); ++ assertThat(result.getQueryText()).contains("c[\"zipcode\"] = @__rmPk_1"); ++ assertThat(result.getQueryText()).contains("c[\"city\"] = @__rmPk_2"); ++ assertThat(result.getQueryText()).contains("c[\"zipcode\"] = @__rmPk_3"); ++ assertThat(result.getParameters()).hasSize(4); ++ } ++ ++ @Test(groups = { "unit" }) ++ public void hpk_partialPk_singleLevel() { ++ PartitionKeyDefinition pkDef = createMultiHashPkDefinition("/city", "/zipcode", "/areaCode"); ++ List selectors = createSelectors(pkDef); ++ ++ // Partial PK ΓÇö only first level ++ PartitionKey partialPk = new PartitionKeyBuilder().add("Redmond").build(); ++ List pkValues = Collections.singletonList(partialPk); ++ ++ SqlQuerySpec result = ReadManyByPartitionKeyQueryHelper.createReadManyByPkQuerySpec( ++ "SELECT * FROM c", new ArrayList<>(), pkValues, selectors, pkDef); ++ ++ assertThat(result.getQueryText()).contains("c[\"city\"] = @__rmPk_0"); ++ // Should NOT include zipcode or areaCode since it's partial ++ assertThat(result.getQueryText()).doesNotContain("zipcode"); ++ assertThat(result.getQueryText()).doesNotContain("areaCode"); ++ assertThat(result.getParameters()).hasSize(1); ++ } ++ ++ @Test(groups = { "unit" }) ++ public void hpk_partialPk_twoLevels() { ++ PartitionKeyDefinition pkDef = createMultiHashPkDefinition("/city", "/zipcode", "/areaCode"); ++ List selectors = createSelectors(pkDef); ++ ++ // Partial PK ΓÇö first two levels ++ PartitionKey partialPk = new PartitionKeyBuilder().add("Redmond").add("98052").build(); ++ List pkValues = Collections.singletonList(partialPk); ++ ++ SqlQuerySpec result = ReadManyByPartitionKeyQueryHelper.createReadManyByPkQuerySpec( ++ "SELECT * FROM c", new ArrayList<>(), pkValues, selectors, pkDef); ++ ++ assertThat(result.getQueryText()).contains("c[\"city\"] = @__rmPk_0"); ++ assertThat(result.getQueryText()).contains("c[\"zipcode\"] = @__rmPk_1"); ++ assertThat(result.getQueryText()).doesNotContain("areaCode"); ++ assertThat(result.getParameters()).hasSize(2); ++ } ++ ++ @Test(groups = { "unit" }) ++ public void hpk_customQuery_withWhere() { ++ PartitionKeyDefinition pkDef = createMultiHashPkDefinition("/city", "/zipcode"); ++ List selectors = createSelectors(pkDef); ++ ++ List baseParams = new ArrayList<>(); ++ baseParams.add(new SqlParameter("@status", "active")); ++ ++ PartitionKey pk = new PartitionKeyBuilder().add("Redmond").add("98052").build(); ++ List pkValues = Collections.singletonList(pk); ++ ++ SqlQuerySpec result = ReadManyByPartitionKeyQueryHelper.createReadManyByPkQuerySpec( ++ "SELECT c.name FROM c WHERE c.status = @status", baseParams, pkValues, selectors, pkDef); ++ ++ assertThat(result.getQueryText()).contains("WHERE (c.status = @status) AND ("); ++ assertThat(result.getQueryText()).contains("c[\"city\"] = @__rmPk_0"); ++ assertThat(result.getParameters()).hasSize(3); // @status + 2 pk params ++ } ++ ++ //endregion ++ ++ //region findTopLevelWhereIndex tests ++ ++ @Test(groups = { "unit" }) ++ public void findWhere_simpleQuery() { ++ int idx = ReadManyByPartitionKeyQueryHelper.findTopLevelWhereIndex("SELECT * FROM C WHERE C.ID = 1"); ++ assertThat(idx).isEqualTo(16); ++ } ++ ++ @Test(groups = { "unit" }) ++ public void findWhere_noWhere() { ++ int idx = ReadManyByPartitionKeyQueryHelper.findTopLevelWhereIndex("SELECT * FROM C"); ++ assertThat(idx).isEqualTo(-1); ++ } ++ ++ @Test(groups = { "unit" }) ++ public void findWhere_whereInSubquery() { ++ // WHERE inside parentheses (subquery) should be ignored ++ int idx = ReadManyByPartitionKeyQueryHelper.findTopLevelWhereIndex( ++ "SELECT * FROM C WHERE EXISTS(SELECT VALUE T FROM T IN C.TAGS WHERE T = 'FOO')"); ++ // Should find the outer WHERE, not the inner one ++ assertThat(idx).isEqualTo(16); ++ } ++ ++ @Test(groups = { "unit" }) ++ public void findWhere_caseInsensitive() { ++ int idx = ReadManyByPartitionKeyQueryHelper.findTopLevelWhereIndex("SELECT * FROM C WHERE C.X = 1"); ++ assertThat(idx).isGreaterThan(0); ++ } ++ ++ @Test(groups = { "unit" }) ++ public void findWhere_whereNotKeyword() { ++ // "ELSEWHERE" should not match ++ int idx = ReadManyByPartitionKeyQueryHelper.findTopLevelWhereIndex("SELECT * FROM ELSEWHERE"); ++ assertThat(idx).isEqualTo(-1); ++ } ++ ++ //endregion ++ ++ //region Custom alias tests ++ ++ @Test(groups = { "unit" }) ++ public void singlePk_customAlias() { ++ PartitionKeyDefinition pkDef = createSinglePkDefinition("/mypk"); ++ List selectors = createSelectors(pkDef); ++ List pkValues = Arrays.asList(new PartitionKey("pk1"), new PartitionKey("pk2")); ++ ++ SqlQuerySpec result = ReadManyByPartitionKeyQueryHelper.createReadManyByPkQuerySpec( ++ "SELECT x.id, x.mypk FROM x", new ArrayList<>(), pkValues, selectors, pkDef); ++ ++ assertThat(result.getQueryText()).startsWith("SELECT x.id, x.mypk FROM x WHERE"); ++ assertThat(result.getQueryText()).contains("x[\"mypk\"] IN ("); ++ assertThat(result.getQueryText()).doesNotContain("c[\"mypk\"]"); ++ } ++ ++ @Test(groups = { "unit" }) ++ public void singlePk_customAlias_withWhere() { ++ PartitionKeyDefinition pkDef = createSinglePkDefinition("/mypk"); ++ List selectors = createSelectors(pkDef); ++ List pkValues = Collections.singletonList(new PartitionKey("pk1")); ++ ++ List baseParams = new ArrayList<>(); ++ baseParams.add(new SqlParameter("@cat", "HelloWorld")); ++ ++ SqlQuerySpec result = ReadManyByPartitionKeyQueryHelper.createReadManyByPkQuerySpec( ++ "SELECT x.id, x.mypk FROM x WHERE x.category = @cat", baseParams, pkValues, selectors, pkDef); ++ ++ assertThat(result.getQueryText()).contains("WHERE (x.category = @cat) AND (x[\"mypk\"] IN ("); ++ } ++ ++ @Test(groups = { "unit" }) ++ public void hpk_customAlias() { ++ PartitionKeyDefinition pkDef = createMultiHashPkDefinition("/city", "/zipcode"); ++ List selectors = createSelectors(pkDef); ++ ++ PartitionKey pk = new PartitionKeyBuilder().add("Redmond").add("98052").build(); ++ List pkValues = Collections.singletonList(pk); ++ ++ SqlQuerySpec result = ReadManyByPartitionKeyQueryHelper.createReadManyByPkQuerySpec( ++ "SELECT r.name FROM root r", new ArrayList<>(), pkValues, selectors, pkDef); ++ ++ assertThat(result.getQueryText()).contains("r[\"city\"] = @__rmPk_0"); ++ assertThat(result.getQueryText()).contains("r[\"zipcode\"] = @__rmPk_1"); ++ assertThat(result.getQueryText()).doesNotContain("c[\""); ++ } ++ ++ //endregion ++ ++ //region extractTableAlias tests ++ ++ @Test(groups = { "unit" }) ++ public void extractAlias_defaultC() { ++ assertThat(ReadManyByPartitionKeyQueryHelper.extractTableAlias("SELECT * FROM c")).isEqualTo("c"); ++ } ++ ++ @Test(groups = { "unit" }) ++ public void extractAlias_customX() { ++ assertThat(ReadManyByPartitionKeyQueryHelper.extractTableAlias("SELECT x.id FROM x WHERE x.age > 5")).isEqualTo("x"); ++ } ++ ++ @Test(groups = { "unit" }) ++ public void extractAlias_rootWithAlias() { ++ assertThat(ReadManyByPartitionKeyQueryHelper.extractTableAlias("SELECT r.name FROM root r")).isEqualTo("r"); ++ } ++ ++ @Test(groups = { "unit" }) ++ public void extractAlias_rootNoAlias() { ++ assertThat(ReadManyByPartitionKeyQueryHelper.extractTableAlias("SELECT * FROM root")).isEqualTo("root"); ++ } ++ ++ @Test(groups = { "unit" }) ++ public void extractAlias_containerWithWhere() { ++ assertThat(ReadManyByPartitionKeyQueryHelper.extractTableAlias("SELECT * FROM items WHERE items.status = 'active'")).isEqualTo("items"); ++ } ++ ++ @Test(groups = { "unit" }) ++ public void extractAlias_caseInsensitive() { ++ assertThat(ReadManyByPartitionKeyQueryHelper.extractTableAlias("select * from MyContainer where MyContainer.id = '1'")).isEqualTo("MyContainer"); ++ } ++ ++ //endregion ++ ++ ++ //region String literal handling tests (#1) ++ ++ @Test(groups = { "unit" }) ++ public void findWhere_ignoresWhereInsideStringLiteral() { ++ // WHERE inside a string literal should be ignored ++ int idx = ReadManyByPartitionKeyQueryHelper.findTopLevelWhereIndex( ++ "SELECT * FROM c WHERE c.msg = 'use WHERE clause here'"); ++ // Should find the outer WHERE at position 16, not the one inside the string ++ assertThat(idx).isEqualTo(16); ++ } ++ ++ @Test(groups = { "unit" }) ++ public void findWhere_ignoresParenthesesInsideStringLiteral() { ++ // Parentheses inside string literal should not affect depth tracking ++ int idx = ReadManyByPartitionKeyQueryHelper.findTopLevelWhereIndex( ++ "SELECT * FROM c WHERE c.name = 'foo(bar)' AND c.x = 1"); ++ assertThat(idx).isEqualTo(16); ++ } ++ ++ @Test(groups = { "unit" }) ++ public void findWhere_handlesUnbalancedParenInStringLiteral() { ++ // Unbalanced paren inside string literal must not corrupt depth ++ int idx = ReadManyByPartitionKeyQueryHelper.findTopLevelWhereIndex( ++ "SELECT * FROM c WHERE c.val = 'open(' AND c.active = true"); ++ assertThat(idx).isEqualTo(16); ++ } ++ ++ @Test(groups = { "unit" }) ++ public void findWhere_handlesStringLiteralBeforeWhere() { ++ // String literal in SELECT before WHERE ++ int idx = ReadManyByPartitionKeyQueryHelper.findTopLevelWhereIndex( ++ "SELECT 'WHERE' as label FROM c WHERE c.id = '1'"); ++ // The WHERE inside quotes should be ignored; the real WHERE is further along ++ assertThat(idx).isGreaterThan(30); ++ } ++ ++ @Test(groups = { "unit" }) ++ public void singlePk_customQuery_withStringLiteralContainingParens() { ++ PartitionKeyDefinition pkDef = createSinglePkDefinition("/mypk"); ++ List selectors = createSelectors(pkDef); ++ List pkValues = Collections.singletonList(new PartitionKey("pk1")); ++ ++ List baseParams = new ArrayList<>(); ++ baseParams.add(new SqlParameter("@msg", "hello")); ++ ++ SqlQuerySpec result = ReadManyByPartitionKeyQueryHelper.createReadManyByPkQuerySpec( ++ "SELECT * FROM c WHERE c.msg = 'test(value)WHERE'", baseParams, pkValues, selectors, pkDef); ++ ++ // Should correctly AND the PK filter to the real WHERE clause ++ assertThat(result.getQueryText()).contains("WHERE (c.msg = 'test(value)WHERE') AND ("); ++ } ++ ++ //endregion ++ ++ //region OFFSET/LIMIT/HAVING alias detection tests (#9) ++ ++ @Test(groups = { "unit" }) ++ public void extractAlias_containerWithOffset() { ++ assertThat(ReadManyByPartitionKeyQueryHelper.extractTableAlias( ++ "SELECT * FROM c OFFSET 10 LIMIT 5")).isEqualTo("c"); ++ } ++ ++ @Test(groups = { "unit" }) ++ public void extractAlias_containerWithLimit() { ++ assertThat(ReadManyByPartitionKeyQueryHelper.extractTableAlias( ++ "SELECT * FROM c LIMIT 10")).isEqualTo("c"); ++ } ++ ++ @Test(groups = { "unit" }) ++ public void extractAlias_containerWithHaving() { ++ assertThat(ReadManyByPartitionKeyQueryHelper.extractTableAlias( ++ "SELECT c.status, COUNT(1) FROM c GROUP BY c.status HAVING COUNT(1) > 1")).isEqualTo("c"); ++ } ++ ++ //endregion ++ ++ //region helpers ++ ++ private PartitionKeyDefinition createSinglePkDefinition(String path) { ++ PartitionKeyDefinition pkDef = new PartitionKeyDefinition(); ++ pkDef.setKind(PartitionKind.HASH); ++ pkDef.setVersion(PartitionKeyDefinitionVersion.V2); ++ pkDef.setPaths(Collections.singletonList(path)); ++ return pkDef; ++ } ++ ++ private PartitionKeyDefinition createMultiHashPkDefinition(String... paths) { ++ PartitionKeyDefinition pkDef = new PartitionKeyDefinition(); ++ pkDef.setKind(PartitionKind.MULTI_HASH); ++ pkDef.setVersion(PartitionKeyDefinitionVersion.V2); ++ pkDef.setPaths(Arrays.asList(paths)); ++ return pkDef; ++ } ++ ++ private List createSelectors(PartitionKeyDefinition pkDef) { ++ return pkDef.getPaths() ++ .stream() ++ .map(pathPart -> pathPart.substring(1)) // skip starting / ++ .map(pathPart -> pathPart.replace("\"", "\\")) ++ .map(part -> "[\"" + part + "\"]") ++ .collect(Collectors.toList()); ++ } ++ ++ //endregion ++} +diff --git a/sdk/cosmos/azure-cosmos/CHANGELOG.md b/sdk/cosmos/azure-cosmos/CHANGELOG.md +index e8ea564fab7..904c01c3238 100644 +--- a/sdk/cosmos/azure-cosmos/CHANGELOG.md ++++ b/sdk/cosmos/azure-cosmos/CHANGELOG.md +@@ -3,6 +3,7 @@ + ### 4.80.0-beta.1 (Unreleased) + + #### Features Added ++* Added new `readManyByPartitionKey` to bulk query by a list of pk-values with better efficiency. See [PR 48801](https://github.com/Azure/azure-sdk-for-java/pull/48801) + + #### Breaking Changes + +diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java +index ad871bb97c0..4e234667c1c 100644 +--- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java ++++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java +@@ -165,6 +165,7 @@ public class CosmosAsyncContainer { + private final String createItemSpanName; + private final String readAllItemsSpanName; + private final String readManyItemsSpanName; ++ private final String readManyByPartitionKeyItemsSpanName; + private final String readAllItemsOfLogicalPartitionSpanName; + private final String queryItemsSpanName; + private final String queryChangeFeedSpanName; +@@ -198,6 +199,7 @@ public class CosmosAsyncContainer { + this.createItemSpanName = "createItem." + this.id; + this.readAllItemsSpanName = "readAllItems." + this.id; + this.readManyItemsSpanName = "readManyItems." + this.id; ++ this.readManyByPartitionKeyItemsSpanName = "readManyByPartitionKeyItems." + this.id; + this.readAllItemsOfLogicalPartitionSpanName = "readAllItemsOfLogicalPartition." + this.id; + this.queryItemsSpanName = "queryItems." + this.id; + this.queryChangeFeedSpanName = "queryChangeFeed." + this.id; +@@ -1601,6 +1603,130 @@ public class CosmosAsyncContainer { + context); + } + ++ /** ++ * Reads many documents matching the provided partition key values. ++ * Unlike {@link #readMany(List, Class)} this method does not require item ids - it queries ++ * all documents matching the provided partition key values. Uses {@code SELECT * FROM c} ++ * as the base query. ++ * ++ * @param the type parameter ++ * @param partitionKeys list of partition key values to read documents for ++ * @param classType class type ++ * @return a {@link CosmosPagedFlux} containing one or several feed response pages ++ */ ++ public CosmosPagedFlux readManyByPartitionKey( ++ List partitionKeys, ++ Class classType) { ++ ++ return this.readManyByPartitionKey(partitionKeys, null, null, classType); ++ } ++ ++ /** ++ * Reads many documents matching the provided partition key values. ++ * Unlike {@link #readMany(List, Class)} this method does not require item ids - it queries ++ * all documents matching the provided partition key values. Uses {@code SELECT * FROM c} ++ * as the base query. ++ * ++ * @param the type parameter ++ * @param partitionKeys list of partition key values to read documents for ++ * @param requestOptions the optional request options ++ * @param classType class type ++ * @return a {@link CosmosPagedFlux} containing one or several feed response pages ++ */ ++ public CosmosPagedFlux readManyByPartitionKey( ++ List partitionKeys, ++ CosmosReadManyRequestOptions requestOptions, ++ Class classType) { ++ ++ return this.readManyByPartitionKey(partitionKeys, null, requestOptions, classType); ++ } ++ ++ /** ++ * Reads many documents matching the provided partition key values with a custom query. ++ * The custom query can be used to apply projections (e.g. {@code SELECT c.name, c.age FROM c}) ++ * and/or additional filters (e.g. {@code SELECT * FROM c WHERE c.status = 'active'}). ++ * The SDK will automatically append partition key filtering to the custom query. ++ *

++ * The custom query must be a simple streamable query ΓÇö aggregates, ORDER BY, DISTINCT, ++ * GROUP BY, DCOUNT, vector search, and full-text search are not supported and will be ++ * rejected. ++ *

++ * Partial hierarchical partition keys are supported and will fan out to multiple ++ * physical partitions. ++ * ++ * @param the type parameter ++ * @param partitionKeys list of partition key values to read documents for ++ * @param customQuery optional custom query for projections/additional filters (null means SELECT * FROM c) ++ * @param requestOptions the optional request options ++ * @param classType class type ++ * @return a {@link CosmosPagedFlux} containing one or several feed response pages ++ */ ++ public CosmosPagedFlux readManyByPartitionKey( ++ List partitionKeys, ++ SqlQuerySpec customQuery, ++ CosmosReadManyRequestOptions requestOptions, ++ Class classType) { ++ ++ if (partitionKeys == null) { ++ throw new IllegalArgumentException("Argument 'partitionKeys' must not be null."); ++ } ++ if (partitionKeys.isEmpty()) { ++ throw new IllegalArgumentException("Argument 'partitionKeys' must not be empty."); ++ } ++ for (PartitionKey pk : partitionKeys) { ++ if (pk == null) { ++ throw new IllegalArgumentException( ++ "Argument 'partitionKeys' must not contain null elements."); ++ } ++ } ++ ++ return UtilBridgeInternal.createCosmosPagedFlux( ++ readManyByPartitionKeyInternalFunc(partitionKeys, customQuery, requestOptions, classType)); ++ } ++ ++ private Function>> readManyByPartitionKeyInternalFunc( ++ List partitionKeys, ++ SqlQuerySpec customQuery, ++ CosmosReadManyRequestOptions requestOptions, ++ Class classType) { ++ ++ CosmosAsyncClient client = this.getDatabase().getClient(); ++ ++ return (pagedFluxOptions -> { ++ CosmosQueryRequestOptions queryRequestOptions = requestOptions == null ++ ? new CosmosQueryRequestOptions() ++ : queryOptionsAccessor().clone(readManyOptionsAccessor().getImpl(requestOptions)); ++ queryRequestOptions.setMaxDegreeOfParallelism(-1); ++ queryRequestOptions.setQueryName("readManyByPartitionKey"); ++ CosmosQueryRequestOptionsBase cosmosQueryRequestOptionsImpl = queryOptionsAccessor().getImpl(queryRequestOptions); ++ applyPolicies(OperationType.Query, ResourceType.Document, cosmosQueryRequestOptionsImpl, this.readManyByPartitionKeyItemsSpanName); ++ ++ QueryFeedOperationState state = new QueryFeedOperationState( ++ client, ++ this.readManyByPartitionKeyItemsSpanName, ++ database.getId(), ++ this.getId(), ++ ResourceType.Document, ++ OperationType.Query, ++ queryOptionsAccessor().getQueryNameOrDefault(queryRequestOptions, this.readManyByPartitionKeyItemsSpanName), ++ queryRequestOptions, ++ pagedFluxOptions ++ ); ++ ++ pagedFluxOptions.setFeedOperationState(state); ++ ++ return CosmosBridgeInternal ++ .getAsyncDocumentClient(this.getDatabase()) ++ .readManyByPartitionKey( ++ partitionKeys, ++ customQuery, ++ BridgeInternal.getLink(this), ++ state, ++ classType) ++ .map(response -> prepareFeedResponse(response, false)); ++ }); ++ } ++ + /** + * Reads all the items of a logical partition + * +diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosContainer.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosContainer.java +index 04a6060c192..0bd8be5850c 100644 +--- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosContainer.java ++++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosContainer.java +@@ -540,6 +540,73 @@ public class CosmosContainer { + classType)); + } + ++ /** ++ * Reads many documents matching the provided partition key values. ++ * Unlike {@link #readMany(List, Class)} this method does not require item ids - it queries ++ * all documents matching the provided partition key values. Uses {@code SELECT * FROM c} ++ * as the base query. ++ * ++ * @param the type parameter ++ * @param partitionKeys list of partition key values to read documents for ++ * @param classType class type ++ * @return a {@link CosmosPagedIterable} containing the results ++ */ ++ public CosmosPagedIterable readManyByPartitionKey( ++ List partitionKeys, ++ Class classType) { ++ ++ return getCosmosPagedIterable(this.asyncContainer.readManyByPartitionKey(partitionKeys, classType)); ++ } ++ ++ /** ++ * Reads many documents matching the provided partition key values. ++ * Unlike {@link #readMany(List, Class)} this method does not require item ids - it queries ++ * all documents matching the provided partition key values. Uses {@code SELECT * FROM c} ++ * as the base query. ++ * ++ * @param the type parameter ++ * @param partitionKeys list of partition key values to read documents for ++ * @param requestOptions the optional request options ++ * @param classType class type ++ * @return a {@link CosmosPagedIterable} containing the results ++ */ ++ public CosmosPagedIterable readManyByPartitionKey( ++ List partitionKeys, ++ CosmosReadManyRequestOptions requestOptions, ++ Class classType) { ++ ++ return getCosmosPagedIterable(this.asyncContainer.readManyByPartitionKey(partitionKeys, requestOptions, classType)); ++ } ++ ++ /** ++ * Reads many documents matching the provided partition key values with a custom query. ++ * The custom query can be used to apply projections (e.g. {@code SELECT c.name, c.age FROM c}) ++ * and/or additional filters (e.g. {@code SELECT * FROM c WHERE c.status = 'active'}). ++ * The SDK will automatically append partition key filtering to the custom query. ++ *

++ * The custom query must be a simple streamable query ΓÇö aggregates, ORDER BY, DISTINCT, ++ * GROUP BY, DCOUNT, vector search, and full-text search are not supported and will be ++ * rejected. ++ *

++ * Partial hierarchical partition keys are supported and will fan out to multiple ++ * physical partitions. ++ * ++ * @param the type parameter ++ * @param partitionKeys list of partition key values to read documents for ++ * @param customQuery optional custom query for projections/additional filters (null means SELECT * FROM c) ++ * @param requestOptions the optional request options ++ * @param classType class type ++ * @return a {@link CosmosPagedIterable} containing the results ++ */ ++ public CosmosPagedIterable readManyByPartitionKey( ++ List partitionKeys, ++ SqlQuerySpec customQuery, ++ CosmosReadManyRequestOptions requestOptions, ++ Class classType) { ++ ++ return getCosmosPagedIterable(this.asyncContainer.readManyByPartitionKey(partitionKeys, customQuery, requestOptions, classType)); ++ } ++ + /** + * Reads all the items of a logical partition returning the results as {@link CosmosPagedIterable}. + * +diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/AsyncDocumentClient.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/AsyncDocumentClient.java +index 945e768a82f..8e2499c9039 100644 +--- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/AsyncDocumentClient.java ++++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/AsyncDocumentClient.java +@@ -1584,6 +1584,27 @@ public interface AsyncDocumentClient { + QueryFeedOperationState state, + Class klass); + ++ /** ++ * Reads many documents by partition key values. ++ * Unlike {@link #readMany(List, String, QueryFeedOperationState, Class)} this method does not require ++ * item ids - it queries all documents matching the provided partition key values. ++ * Partial hierarchical partition keys are supported and will fan out to multiple physical partitions. ++ * ++ * @param partitionKeys list of partition key values to read documents for ++ * @param customQuery optional custom query (for projections/additional filters) - null means SELECT * FROM c ++ * @param collectionLink link for the documentcollection/container to be queried ++ * @param state the query operation state ++ * @param klass class type ++ * @param the type parameter ++ * @return a Flux with feed response pages of documents ++ */ ++ Flux> readManyByPartitionKey( ++ List partitionKeys, ++ SqlQuerySpec customQuery, ++ String collectionLink, ++ QueryFeedOperationState state, ++ Class klass); ++ + /** + * Read all documents of a certain logical partition. + *

+diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/Configs.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/Configs.java +index 337055c6947..162b0740f40 100644 +--- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/Configs.java ++++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/Configs.java +@@ -248,6 +248,11 @@ public class Configs { + public static final String MIN_TARGET_BULK_MICRO_BATCH_SIZE_VARIABLE = "COSMOS_MIN_TARGET_BULK_MICRO_BATCH_SIZE"; + public static final int DEFAULT_MIN_TARGET_BULK_MICRO_BATCH_SIZE = 1; + ++ // readManyByPartitionKey: max number of PK values per query per physical partition ++ private static final String READ_MANY_BY_PK_MAX_BATCH_SIZE = "COSMOS.READ_MANY_BY_PK_MAX_BATCH_SIZE"; ++ private static final String READ_MANY_BY_PK_MAX_BATCH_SIZE_VARIABLE = "COSMOS_READ_MANY_BY_PK_MAX_BATCH_SIZE"; ++ private static final int DEFAULT_READ_MANY_BY_PK_MAX_BATCH_SIZE = 1000; ++ + public static final String MAX_BULK_MICRO_BATCH_CONCURRENCY = "COSMOS.MAX_BULK_MICRO_BATCH_CONCURRENCY"; + public static final String MAX_BULK_MICRO_BATCH_CONCURRENCY_VARIABLE = "COSMOS_MAX_BULK_MICRO_BATCH_CONCURRENCY"; + public static final int DEFAULT_MAX_BULK_MICRO_BATCH_CONCURRENCY = 1; +@@ -816,6 +821,20 @@ public class Configs { + return DEFAULT_MIN_TARGET_BULK_MICRO_BATCH_SIZE; + } + ++ public static int getReadManyByPkMaxBatchSize() { ++ String valueFromSystemProperty = System.getProperty(READ_MANY_BY_PK_MAX_BATCH_SIZE); ++ if (valueFromSystemProperty != null && !valueFromSystemProperty.isEmpty()) { ++ return Math.max(1, Integer.parseInt(valueFromSystemProperty)); ++ } ++ ++ String valueFromEnvVariable = System.getenv(READ_MANY_BY_PK_MAX_BATCH_SIZE_VARIABLE); ++ if (valueFromEnvVariable != null && !valueFromEnvVariable.isEmpty()) { ++ return Math.max(1, Integer.parseInt(valueFromEnvVariable)); ++ } ++ ++ return DEFAULT_READ_MANY_BY_PK_MAX_BATCH_SIZE; ++ } ++ + public static int getMaxBulkMicroBatchConcurrency() { + String valueFromSystemProperty = System.getProperty(MAX_BULK_MICRO_BATCH_CONCURRENCY); + if (valueFromSystemProperty != null && !valueFromSystemProperty.isEmpty()) { +diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelper.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelper.java +new file mode 100644 +index 00000000000..6d6cd084e01 +--- /dev/null ++++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelper.java +@@ -0,0 +1,263 @@ ++// Copyright (c) Microsoft Corporation. All rights reserved. ++// Licensed under the MIT License. ++package com.azure.cosmos.implementation; ++ ++import com.azure.cosmos.BridgeInternal; ++import com.azure.cosmos.implementation.routing.PartitionKeyInternal; ++import com.azure.cosmos.models.PartitionKey; ++import com.azure.cosmos.models.PartitionKeyDefinition; ++import com.azure.cosmos.models.PartitionKind; ++import com.azure.cosmos.models.SqlParameter; ++import com.azure.cosmos.models.SqlQuerySpec; ++ ++import java.util.ArrayList; ++import java.util.List; ++ ++/** ++ * Helper for constructing SqlQuerySpec instances for readManyByPartitionKey operations. ++ * This class is not intended to be used directly by end-users. ++ */ ++public class ReadManyByPartitionKeyQueryHelper { ++ ++ private static final String DEFAULT_TABLE_ALIAS = "c"; ++ // Internal parameter prefix ΓÇö uses double-underscore to avoid collisions with user-provided parameters ++ private static final String PK_PARAM_PREFIX = "@__rmPk_"; ++ ++ public static SqlQuerySpec createReadManyByPkQuerySpec( ++ String baseQueryText, ++ List baseParameters, ++ List pkValues, ++ List partitionKeySelectors, ++ PartitionKeyDefinition pkDefinition) { ++ ++ // Extract the table alias from the FROM clause (e.g. "FROM x" ΓåÆ "x", "FROM c" ΓåÆ "c") ++ String tableAlias = extractTableAlias(baseQueryText); ++ ++ StringBuilder pkFilter = new StringBuilder(); ++ List parameters = new ArrayList<>(baseParameters); ++ int paramCount = 0; ++ ++ boolean isSinglePathPk = partitionKeySelectors.size() == 1; ++ ++ if (isSinglePathPk && pkDefinition.getKind() != PartitionKind.MULTI_HASH) { ++ // Single PK path ΓÇö use IN clause for normal values, OR NOT IS_DEFINED for NONE ++ // First, separate NONE PKs from normal PKs ++ boolean hasNone = false; ++ List normalPkValues = new ArrayList<>(); ++ for (PartitionKey pk : pkValues) { ++ PartitionKeyInternal pkInternal = BridgeInternal.getPartitionKeyInternal(pk); ++ if (pkInternal.getComponents() == null) { ++ hasNone = true; ++ } else { ++ normalPkValues.add(pk); ++ } ++ } ++ ++ pkFilter.append(" "); ++ boolean hasNormalValues = !normalPkValues.isEmpty(); ++ if (hasNormalValues && hasNone) { ++ pkFilter.append("("); ++ } ++ if (hasNormalValues) { ++ pkFilter.append(tableAlias); ++ pkFilter.append(partitionKeySelectors.get(0)); ++ pkFilter.append(" IN ( "); ++ for (int i = 0; i < normalPkValues.size(); i++) { ++ PartitionKeyInternal pkInternal = BridgeInternal.getPartitionKeyInternal(normalPkValues.get(i)); ++ Object[] pkComponents = pkInternal.toObjectArray(); ++ String pkParamName = PK_PARAM_PREFIX + paramCount; ++ parameters.add(new SqlParameter(pkParamName, pkComponents[0])); ++ paramCount++; ++ ++ pkFilter.append(pkParamName); ++ if (i < normalPkValues.size() - 1) { ++ pkFilter.append(", "); ++ } ++ } ++ pkFilter.append(" )"); ++ } ++ if (hasNone) { ++ if (hasNormalValues) { ++ pkFilter.append(" OR "); ++ } ++ pkFilter.append("NOT IS_DEFINED("); ++ pkFilter.append(tableAlias); ++ pkFilter.append(partitionKeySelectors.get(0)); ++ pkFilter.append(")"); ++ } ++ if (hasNormalValues && hasNone) { ++ pkFilter.append(")"); ++ } ++ } else { ++ // Multiple PK paths (HPK) or MULTI_HASH ΓÇö use OR of AND clauses ++ pkFilter.append(" "); ++ for (int i = 0; i < pkValues.size(); i++) { ++ PartitionKeyInternal pkInternal = BridgeInternal.getPartitionKeyInternal(pkValues.get(i)); ++ Object[] pkComponents = pkInternal.toObjectArray(); ++ ++ // PartitionKey.NONE ΓÇö generate NOT IS_DEFINED for all PK paths ++ if (pkComponents == null) { ++ pkFilter.append("("); ++ for (int j = 0; j < partitionKeySelectors.size(); j++) { ++ if (j > 0) { ++ pkFilter.append(" AND "); ++ } ++ pkFilter.append("NOT IS_DEFINED("); ++ pkFilter.append(tableAlias); ++ pkFilter.append(partitionKeySelectors.get(j)); ++ pkFilter.append(")"); ++ } ++ pkFilter.append(")"); ++ } else { ++ pkFilter.append("("); ++ for (int j = 0; j < pkComponents.length; j++) { ++ String pkParamName = PK_PARAM_PREFIX + paramCount; ++ parameters.add(new SqlParameter(pkParamName, pkComponents[j])); ++ paramCount++; ++ ++ if (j > 0) { ++ pkFilter.append(" AND "); ++ } ++ pkFilter.append(tableAlias); ++ pkFilter.append(partitionKeySelectors.get(j)); ++ pkFilter.append(" = "); ++ pkFilter.append(pkParamName); ++ } ++ pkFilter.append(")"); ++ } ++ ++ if (i < pkValues.size() - 1) { ++ pkFilter.append(" OR "); ++ } ++ } ++ } ++ ++ // Compose final query: handle existing WHERE clause in base query ++ String finalQuery; ++ int whereIndex = findTopLevelWhereIndex(baseQueryText); ++ if (whereIndex >= 0) { ++ // Base query has WHERE ΓÇö AND our PK filter ++ String beforeWhere = baseQueryText.substring(0, whereIndex); ++ String afterWhere = baseQueryText.substring(whereIndex + 5); // skip "WHERE" ++ finalQuery = beforeWhere + "WHERE (" + afterWhere.trim() + ") AND (" + pkFilter.toString().trim() + ")"; ++ } else { ++ // No WHERE ΓÇö add one ++ finalQuery = baseQueryText + " WHERE" + pkFilter.toString(); ++ } ++ ++ return new SqlQuerySpec(finalQuery, parameters); ++ } ++ ++ /** ++ * Extracts the table/collection alias from a SQL query's FROM clause. ++ * Handles: "SELECT * FROM c", "SELECT x.id FROM x WHERE ...", "SELECT * FROM root r", etc. ++ * Returns the alias used after FROM (last token before WHERE or end of FROM clause). ++ */ ++ static String extractTableAlias(String queryText) { ++ String upper = queryText.toUpperCase(); ++ int fromIndex = findTopLevelKeywordIndex(upper, "FROM"); ++ if (fromIndex < 0) { ++ return DEFAULT_TABLE_ALIAS; ++ } ++ ++ // Start scanning after "FROM" ++ int afterFrom = fromIndex + 4; ++ // Skip whitespace ++ while (afterFrom < queryText.length() && Character.isWhitespace(queryText.charAt(afterFrom))) { ++ afterFrom++; ++ } ++ ++ // Collect the container name token (could be "root", "c", etc.) ++ int tokenStart = afterFrom; ++ while (afterFrom < queryText.length() ++ && !Character.isWhitespace(queryText.charAt(afterFrom)) ++ && queryText.charAt(afterFrom) != '(' ++ && queryText.charAt(afterFrom) != ')') { ++ afterFrom++; ++ } ++ String containerName = queryText.substring(tokenStart, afterFrom); ++ ++ // Skip whitespace after container name ++ while (afterFrom < queryText.length() && Character.isWhitespace(queryText.charAt(afterFrom))) { ++ afterFrom++; ++ } ++ ++ // Check if there's an alias after the container name (before WHERE or end) ++ if (afterFrom < queryText.length()) { ++ char nextChar = Character.toUpperCase(queryText.charAt(afterFrom)); ++ // If the next token is a keyword (WHERE, ORDER, GROUP, JOIN, OFFSET, LIMIT, HAVING) or end, containerName IS the alias ++ if (nextChar == 'W' || nextChar == 'O' || nextChar == 'G' || nextChar == 'J' ++ || nextChar == 'L' || nextChar == 'H') { ++ // Check if it's actually a keyword ++ String remaining = upper.substring(afterFrom); ++ if (remaining.startsWith("WHERE") || remaining.startsWith("ORDER") ++ || remaining.startsWith("GROUP") || remaining.startsWith("JOIN") ++ || remaining.startsWith("OFFSET") || remaining.startsWith("LIMIT") ++ || remaining.startsWith("HAVING")) { ++ return containerName; ++ } ++ } ++ // Otherwise the next token is the alias ("FROM root r" ΓåÆ alias is "r") ++ int aliasStart = afterFrom; ++ while (afterFrom < queryText.length() ++ && !Character.isWhitespace(queryText.charAt(afterFrom)) ++ && queryText.charAt(afterFrom) != '(' ++ && queryText.charAt(afterFrom) != ')') { ++ afterFrom++; ++ } ++ if (afterFrom > aliasStart) { ++ return queryText.substring(aliasStart, afterFrom); ++ } ++ } ++ ++ return containerName; ++ } ++ ++ /** ++ * Finds the index of a top-level SQL keyword in the query text (case-insensitive), ++ * ignoring occurrences inside parentheses or string literals. ++ */ ++ static int findTopLevelKeywordIndex(String queryText, String keyword) { ++ String queryTextUpper = queryText.toUpperCase(); ++ String keywordUpper = keyword.toUpperCase(); ++ int depth = 0; ++ int keyLen = keywordUpper.length(); ++ for (int i = 0; i <= queryTextUpper.length() - keyLen; i++) { ++ char ch = queryTextUpper.charAt(i); ++ // Skip string literals enclosed in single quotes (handle '' escape) ++ if (queryText.charAt(i) == '\'') { ++ i++; ++ while (i < queryText.length()) { ++ if (queryText.charAt(i) == '\'') { ++ if (i + 1 < queryText.length() && queryText.charAt(i + 1) == '\'') { ++ i += 2; // escaped quote ΓÇö skip both ++ continue; ++ } ++ break; // end of string literal ++ } ++ i++; ++ } ++ continue; ++ } ++ if (ch == '(') { ++ depth++; ++ } else if (ch == ')') { ++ depth--; ++ } else if (depth == 0 && ch == keywordUpper.charAt(0) ++ && queryTextUpper.startsWith(keywordUpper, i) ++ && (i == 0 || !Character.isLetterOrDigit(queryTextUpper.charAt(i - 1))) ++ && (i + keyLen >= queryTextUpper.length() || !Character.isLetterOrDigit(queryTextUpper.charAt(i + keyLen)))) { ++ return i; ++ } ++ } ++ return -1; ++ } ++ ++ /** ++ * Finds the index of the top-level WHERE keyword in the query text, ++ * ignoring WHERE that appears inside parentheses (subqueries). ++ */ ++ public static int findTopLevelWhereIndex(String queryTextUpper) { ++ return findTopLevelKeywordIndex(queryTextUpper, "WHERE"); ++ } ++} +diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java +index 11121bca033..c70dedaa1f2 100644 +--- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java ++++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java +@@ -4365,6 +4365,298 @@ public class RxDocumentClientImpl implements AsyncDocumentClient, IAuthorization + ); + } + ++ @Override ++ public Flux> readManyByPartitionKey( ++ List partitionKeys, ++ SqlQuerySpec customQuery, ++ String collectionLink, ++ QueryFeedOperationState state, ++ Class klass) { ++ ++ checkNotNull(partitionKeys, "Argument 'partitionKeys' must not be null."); ++ checkArgument(!partitionKeys.isEmpty(), "Argument 'partitionKeys' must not be empty."); ++ ++ final ScopedDiagnosticsFactory diagnosticsFactory = new ScopedDiagnosticsFactory(this, true); ++ state.registerDiagnosticsFactory( ++ () -> {}, // we never want to reset in readManyByPartitionKey ++ (ctx) -> diagnosticsFactory.merge(ctx) ++ ); ++ ++ StaleResourceRetryPolicy staleResourceRetryPolicy = new StaleResourceRetryPolicy( ++ this.collectionCache, ++ null, ++ collectionLink, ++ queryOptionsAccessor().getProperties(state.getQueryOptions()), ++ queryOptionsAccessor().getHeaders(state.getQueryOptions()), ++ this.sessionContainer, ++ diagnosticsFactory, ++ ResourceType.Document ++ ); ++ ++ return ObservableHelper ++ .fluxInlineIfPossibleAsObs( ++ () -> readManyByPartitionKey( ++ partitionKeys, customQuery, collectionLink, state, diagnosticsFactory, klass), ++ staleResourceRetryPolicy ++ ) ++ .onErrorMap(throwable -> { ++ if (throwable instanceof CosmosException) { ++ CosmosException cosmosException = (CosmosException) throwable; ++ CosmosDiagnostics diagnostics = cosmosException.getDiagnostics(); ++ if (diagnostics != null) { ++ state.mergeDiagnosticsContext(); ++ CosmosDiagnosticsContext ctx = state.getDiagnosticsContextSnapshot(); ++ if (ctx != null) { ++ ctxAccessor().recordOperation( ++ ctx, ++ cosmosException.getStatusCode(), ++ cosmosException.getSubStatusCode(), ++ 0, ++ cosmosException.getRequestCharge(), ++ diagnostics, ++ throwable ++ ); ++ diagAccessor() ++ .setDiagnosticsContext( ++ diagnostics, ++ state.getDiagnosticsContextSnapshot()); ++ } ++ } ++ ++ return cosmosException; ++ } ++ ++ return throwable; ++ }); ++ } ++ ++ private Flux> readManyByPartitionKey( ++ List partitionKeys, ++ SqlQuerySpec customQuery, ++ String collectionLink, ++ QueryFeedOperationState state, ++ ScopedDiagnosticsFactory diagnosticsFactory, ++ Class klass) { ++ ++ String resourceLink = parentResourceLinkToQueryLink(collectionLink, ResourceType.Document); ++ RxDocumentServiceRequest request = RxDocumentServiceRequest.create(diagnosticsFactory, ++ OperationType.Query, ++ ResourceType.Document, ++ collectionLink, null ++ ); ++ ++ Mono> collectionObs = ++ collectionCache.resolveCollectionAsync(null, request); ++ ++ return collectionObs ++ .flatMapMany(documentCollectionResourceResponse -> { ++ final DocumentCollection collection = documentCollectionResourceResponse.v; ++ if (collection == null) { ++ return Flux.error(new IllegalStateException("Collection cannot be null")); ++ } ++ ++ final PartitionKeyDefinition pkDefinition = collection.getPartitionKey(); ++ ++ Mono> valueHolderMono = partitionKeyRangeCache ++ .tryLookupAsync( ++ BridgeInternal.getMetaDataDiagnosticContext(request.requestContext.cosmosDiagnostics), ++ collection.getResourceId(), ++ null, ++ null); ++ ++ // Validate custom query if provided ++ Mono queryValidationMono; ++ if (customQuery != null) { ++ queryValidationMono = validateCustomQueryForReadManyByPartitionKey( ++ customQuery, resourceLink, state.getQueryOptions()); ++ } else { ++ queryValidationMono = Mono.empty(); ++ } ++ ++ return valueHolderMono ++ .delayUntil(ignored -> queryValidationMono) ++ .flatMapMany(routingMapHolder -> { ++ CollectionRoutingMap routingMap = routingMapHolder.v; ++ if (routingMap == null) { ++ return Flux.error(new IllegalStateException("Failed to get routing map.")); ++ } ++ ++ Map> partitionRangePkMap = ++ groupPartitionKeysByPhysicalPartition(partitionKeys, pkDefinition, routingMap); ++ ++ List partitionKeySelectors = createPkSelectors(pkDefinition); ++ ++ String baseQueryText; ++ List baseParameters; ++ if (customQuery != null) { ++ baseQueryText = customQuery.getQueryText(); ++ baseParameters = customQuery.getParameters() != null ++ ? new ArrayList<>(customQuery.getParameters()) ++ : new ArrayList<>(); ++ } else { ++ baseQueryText = "SELECT * FROM c"; ++ baseParameters = new ArrayList<>(); ++ } ++ ++ // Build per-physical-partition batched queries. ++ // Each physical partition may have many PKs ΓÇö split into batches ++ // to avoid oversized SQL queries. Batch size is configurable via ++ // system property COSMOS.READ_MANY_BY_PK_MAX_BATCH_SIZE (default 1000). ++ int maxPksPerPartitionQuery = Configs.getReadManyByPkMaxBatchSize(); ++ ++ // Build batches per partition as a list of lists (one inner list per partition). ++ // Then interleave in round-robin order so that concurrent execution ++ // prefers different physical partitions over multiple batches of the same partition. ++ List>> batchesPerPartition = new ArrayList<>(); ++ int maxBatchesPerPartition = 0; ++ ++ for (Map.Entry> entry : partitionRangePkMap.entrySet()) { ++ List allPks = entry.getValue(); ++ if (allPks.isEmpty()) { ++ continue; ++ } ++ List> partitionBatches = new ArrayList<>(); ++ for (int i = 0; i < allPks.size(); i += maxPksPerPartitionQuery) { ++ List batch = allPks.subList( ++ i, Math.min(i + maxPksPerPartitionQuery, allPks.size())); ++ SqlQuerySpec querySpec = ReadManyByPartitionKeyQueryHelper ++ .createReadManyByPkQuerySpec( ++ baseQueryText, baseParameters, batch, ++ partitionKeySelectors, pkDefinition); ++ partitionBatches.add( ++ Collections.singletonMap(entry.getKey(), querySpec)); ++ } ++ batchesPerPartition.add(partitionBatches); ++ maxBatchesPerPartition = Math.max(maxBatchesPerPartition, partitionBatches.size()); ++ } ++ ++ if (batchesPerPartition.isEmpty()) { ++ return Flux.empty(); ++ } ++ ++ // Round-robin interleave: [batch0-p1, batch0-p2, ..., batch0-pN, batch1-p1, batch1-p2, ...] ++ // This ensures that with bounded concurrency, different partitions are ++ // preferred over sequential batches of the same partition. ++ List> interleavedBatches = new ArrayList<>(); ++ for (int batchIdx = 0; batchIdx < maxBatchesPerPartition; batchIdx++) { ++ for (List> partitionBatches : batchesPerPartition) { ++ if (batchIdx < partitionBatches.size()) { ++ interleavedBatches.add(partitionBatches.get(batchIdx)); ++ } ++ } ++ } ++ ++ // Execute all batches with bounded concurrency. ++ List>> queryFluxes = interleavedBatches ++ .stream() ++ .map(batchMap -> queryForReadMany( ++ diagnosticsFactory, ++ resourceLink, ++ new SqlQuerySpec(DUMMY_SQL_QUERY), ++ state.getQueryOptions(), ++ klass, ++ ResourceType.Document, ++ collection, ++ Collections.unmodifiableMap(batchMap))) ++ .collect(Collectors.toList()); ++ ++ int fluxConcurrency = Math.min(queryFluxes.size(), ++ Math.max(Configs.getCPUCnt(), 1)); ++ ++ return Flux.merge(Flux.fromIterable(queryFluxes), fluxConcurrency, 1); ++ }); ++ }); ++ } ++ ++ private Mono validateCustomQueryForReadManyByPartitionKey( ++ SqlQuerySpec customQuery, ++ String resourceLink, ++ CosmosQueryRequestOptions queryRequestOptions) { ++ ++ IDocumentQueryClient queryClient = documentQueryClientImpl( ++ RxDocumentClientImpl.this, getOperationContextAndListenerTuple(queryRequestOptions)); ++ ++ return DocumentQueryExecutionContextFactory ++ .fetchQueryPlanForValidation(this, queryClient, customQuery, resourceLink, queryRequestOptions) ++ .flatMap(queryPlan -> { ++ QueryInfo queryInfo = queryPlan.getQueryInfo(); ++ ++ if (queryInfo.hasGroupBy()) { ++ return Mono.error(new IllegalArgumentException( ++ "Custom query for readMany by partition key must not contain GROUP BY.")); ++ } ++ if (queryInfo.hasAggregates()) { ++ return Mono.error(new IllegalArgumentException( ++ "Custom query for readMany by partition key must not contain aggregates.")); ++ } ++ if (queryInfo.hasOrderBy()) { ++ return Mono.error(new IllegalArgumentException( ++ "Custom query for readMany by partition key must not contain ORDER BY.")); ++ } ++ if (queryInfo.hasDistinct()) { ++ return Mono.error(new IllegalArgumentException( ++ "Custom query for readMany by partition key must not contain DISTINCT.")); ++ } ++ if (queryInfo.hasDCount()) { ++ return Mono.error(new IllegalArgumentException( ++ "Custom query for readMany by partition key must not contain DCOUNT.")); ++ } ++ if (queryInfo.hasNonStreamingOrderBy()) { ++ return Mono.error(new IllegalArgumentException( ++ "Custom query for readMany by partition key must not contain non-streaming ORDER BY.")); ++ } ++ if (queryPlan.hasHybridSearchQueryInfo()) { ++ return Mono.error(new IllegalArgumentException( ++ "Custom query for readMany by partition key must not contain hybrid/vector/full-text search.")); ++ } ++ ++ return Mono.empty(); ++ }); ++ } ++ ++ private Map> groupPartitionKeysByPhysicalPartition( ++ List partitionKeys, ++ PartitionKeyDefinition pkDefinition, ++ CollectionRoutingMap routingMap) { ++ ++ Map> partitionRangePkMap = new HashMap<>(); ++ ++ for (PartitionKey pk : partitionKeys) { ++ PartitionKeyInternal pkInternal = BridgeInternal.getPartitionKeyInternal(pk); ++ ++ // PartitionKey.NONE wraps NonePartitionKey which has components = null. ++ // For routing purposes, treat NONE as UndefinedPartitionKey ΓÇö documents ingested ++ // without a partition key path are stored with the undefined EPK. ++ PartitionKeyInternal effectivePkInternal = pkInternal.getComponents() == null ++ ? PartitionKeyInternal.UndefinedPartitionKey ++ : pkInternal; ++ ++ int componentCount = effectivePkInternal.getComponents().size(); ++ int definedPathCount = pkDefinition.getPaths().size(); ++ ++ List targetRanges; ++ ++ if (pkDefinition.getKind() == PartitionKind.MULTI_HASH && componentCount < definedPathCount) { ++ // Partial HPK ΓÇö compute EPK prefix range and find all overlapping physical partitions ++ Range epkRange = PartitionKeyInternalHelper.getEPKRangeForPrefixPartitionKey( ++ effectivePkInternal, pkDefinition); ++ targetRanges = routingMap.getOverlappingRanges(epkRange); ++ } else { ++ // Full PK ΓÇö maps to exactly one physical partition ++ String effectivePartitionKeyString = PartitionKeyInternalHelper ++ .getEffectivePartitionKeyString(effectivePkInternal, pkDefinition); ++ PartitionKeyRange range = routingMap.getRangeByEffectivePartitionKey(effectivePartitionKeyString); ++ targetRanges = Collections.singletonList(range); ++ } ++ ++ for (PartitionKeyRange range : targetRanges) { ++ partitionRangePkMap.computeIfAbsent(range, k -> new ArrayList<>()).add(pk); ++ } ++ } ++ ++ return partitionRangePkMap; ++ } ++ + private Map getRangeQueryMap( + Map> partitionRangeItemKeyMap, + PartitionKeyDefinition partitionKeyDefinition) { +diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/DocumentQueryExecutionContextFactory.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/DocumentQueryExecutionContextFactory.java +index e62d8ed3d75..d8f9614343c 100644 +--- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/DocumentQueryExecutionContextFactory.java ++++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/DocumentQueryExecutionContextFactory.java +@@ -318,6 +318,17 @@ public class DocumentQueryExecutionContextFactory { + return feedRanges; + } + ++ public static Mono fetchQueryPlanForValidation( ++ DiagnosticsClientContext diagnosticsClientContext, ++ IDocumentQueryClient queryClient, ++ SqlQuerySpec sqlQuerySpec, ++ String resourceLink, ++ CosmosQueryRequestOptions queryRequestOptions) { ++ ++ return QueryPlanRetriever.getQueryPlanThroughGatewayAsync( ++ diagnosticsClientContext, queryClient, sqlQuerySpec, resourceLink, queryRequestOptions); ++ } ++ + public static Flux> createDocumentQueryExecutionContextAsync( + DiagnosticsClientContext diagnosticsClientContext, + IDocumentQueryClient client, +diff --git a/sdk/cosmos/cspell.yaml b/sdk/cosmos/cspell.yaml +new file mode 100644 +index 00000000000..94a4002c2c9 +--- /dev/null ++++ b/sdk/cosmos/cspell.yaml +@@ -0,0 +1,6 @@ ++import: ++ - ../../.vscode/cspell.json ++overrides: ++ - filename: "**/sdk/cosmos/*" ++ words: ++ - DCOUNT +diff --git a/sdk/cosmos/docs/readManyByPartitionKey-design.md b/sdk/cosmos/docs/readManyByPartitionKey-design.md +new file mode 100644 +index 00000000000..95d7624f0c8 +--- /dev/null ++++ b/sdk/cosmos/docs/readManyByPartitionKey-design.md +@@ -0,0 +1,169 @@ ++# readManyByPartitionKey ΓÇö Design & Implementation ++ ++## Overview ++ ++New `readManyByPartitionKey` methods on `CosmosAsyncContainer` / `CosmosContainer` that accept a ++`List` (without item-id). The SDK splits the PK values by physical ++partition, generates batched streaming queries per physical partition, and returns results as ++`CosmosPagedFlux` / `CosmosPagedIterable`. ++ ++An optional `SqlQuerySpec` parameter lets callers supply a custom query for projections ++and additional filters. The SDK appends the auto-generated PK WHERE clause to it. ++ ++## Decisions ++ ++| Topic | Decision | ++|---|---| ++| API name | `readManyByPartitionKey` ΓÇö distinct name to avoid ambiguity with existing `readMany(List)` | ++| Return type | `CosmosPagedFlux` (async) / `CosmosPagedIterable` (sync) | ++| Custom query format | `SqlQuerySpec` ΓÇö full query with parameters; SDK ANDs the PK filter | ++| Partial HPK | Supported from the start; prefix PKs fan out via `getOverlappingRanges` | ++| PK deduplication | Done at Spark layer only, not in the SDK | ++| Spark UDF | New `GetCosmosPartitionKeyValue` UDF | ++| Custom query validation | Gateway query plan; reject aggregates/ORDER BY/DISTINCT/GROUP BY/DCount/non-streaming ORDER BY/vector/fulltext | ++| PK list size | No hard upper-bound enforced; SDK batches internally per physical partition (default 1000 PKs per batch, configurable via `COSMOS.READ_MANY_BY_PK_MAX_BATCH_SIZE`) | ++| Eager validation | Null and empty PK list rejected eagerly (not lazily in reactive chain) | ++| Telemetry | Separate span name `readManyByPartitionKeyItems.` (distinct from existing `readManyItems`) | ++| Query construction | Table alias auto-detected from FROM clause; string literals and subqueries handled correctly | ++ ++## Phase 1 ΓÇö SDK Core (`azure-cosmos`) ++ ++### Step 1: New public overloads in CosmosAsyncContainer ++ ++```java ++ CosmosPagedFlux readManyByPartitionKey(List partitionKeys, Class classType) ++ CosmosPagedFlux readManyByPartitionKey(List partitionKeys, ++ CosmosReadManyRequestOptions requestOptions, ++ Class classType) ++ CosmosPagedFlux readManyByPartitionKey(List partitionKeys, ++ SqlQuerySpec customQuery, ++ CosmosReadManyRequestOptions requestOptions, ++ Class classType) ++``` ++ ++All delegate to a private `readManyByPartitionKeyInternalFunc(...)`. ++ ++**Eager validation:** The 4-arg method validates `partitionKeys` is non-null and non-empty before constructing the reactive pipeline, throwing `IllegalArgumentException` synchronously. ++ ++### Step 2: Sync wrappers in CosmosContainer ++ ++Same signatures returning `CosmosPagedIterable`, delegating to the async container. ++ ++### Step 3: Internal orchestration (RxDocumentClientImpl) ++ ++1. Resolve collection metadata + PK definition from cache. ++2. Fetch routing map from `partitionKeyRangeCache` **in parallel with** custom query validation (Step 4). ++3. For each `PartitionKey`: ++ - Compute effective partition key (EPK). ++ - Full PK ΓåÆ `getRangeByEffectivePartitionKey()` (single range). ++ - Partial HPK ΓåÆ compute EPK prefix range ΓåÆ `getOverlappingRanges()` (multiple ranges). ++ **Note:** partial HPK intentionally fans out to multiple physical partitions. ++4. Group PK values by `PartitionKeyRange`. ++5. Per physical partition ΓåÆ split PKs into batches of `maxPksPerPartitionQuery` (configurable, default 1000). ++6. Per batch ΓåÆ build `SqlQuerySpec` with PK WHERE clause (Step 5). ++7. Interleave batches across physical partitions in round-robin order so that bounded concurrency prefers different physical partitions over sequential batches of the same partition. ++8. Execute queries via `queryForReadMany()` with bounded concurrency (`Math.min(batchCount, cpuCount)`). ++9. Return results as `CosmosPagedFlux`. ++ ++### Step 4: Custom query validation ++ ++One-time call per invocation (existing query plan caching applies). Runs **in parallel** with routing map lookup to minimize latency: ++ ++- `QueryPlanRetriever.getQueryPlanThroughGatewayAsync()` for the user query. ++- Reject (`IllegalArgumentException`) if: ++ - `queryInfo.hasGroupBy()` ΓÇö checked first (takes precedence over aggregates since `hasAggregates()` also returns true for GROUP BY queries) ++ - `queryInfo.hasAggregates()` ++ - `queryInfo.hasOrderBy()` ++ - `queryInfo.hasDistinct()` ++ - `queryInfo.hasDCount()` ++ - `queryInfo.hasNonStreamingOrderBy()` ++ - `partitionedQueryExecutionInfo.hasHybridSearchQueryInfo()` ++ ++### Step 5: Query construction ++ ++Query construction is implemented in `ReadManyByPartitionKeyQueryHelper`. The helper: ++- Extracts the table alias from the FROM clause (handles `FROM c`, `FROM root r`, `FROM x WHERE ...`) ++- Handles string literals in queries (parens/keywords inside `'...'` are correctly skipped) ++- Recognizes SQL keywords: WHERE, ORDER, GROUP, JOIN, OFFSET, LIMIT, HAVING ++- Uses parameterized queries (`@__rmPk_` prefix) to prevent SQL injection ++ ++**Single PK (HASH):** ++```sql ++{baseQuery} WHERE {alias}["{pkPath}"] IN (@__rmPk_0, @__rmPk_1, @__rmPk_2) ++``` ++ ++**Full HPK (MULTI_HASH):** ++```sql ++{baseQuery} WHERE ({alias}["{path1}"] = @__rmPk_0 AND {alias}["{path2}"] = @__rmPk_1) ++ OR ({alias}["{path1}"] = @__rmPk_2 AND {alias}["{path2}"] = @__rmPk_3) ++``` ++ ++**Partial HPK (prefix-only):** ++```sql ++{baseQuery} WHERE ({alias}["{path1}"] = @__rmPk_0) ++ OR ({alias}["{path1}"] = @__rmPk_1) ++``` ++ ++If the base query already has a WHERE clause: ++```sql ++{selectAndFrom} WHERE ({existingWhere}) AND ({pkFilter}) ++``` ++ ++### Step 6: Interface wiring ++ ++New method `readManyByPartitionKey` added directly to `AsyncDocumentClient` interface, implemented in `RxDocumentClientImpl`. New `fetchQueryPlanForValidation` static method added to `DocumentQueryExecutionContextFactory` for custom query validation. ++ ++### Step 7: Configuration ++ ++New configurable batch size via system property `COSMOS.READ_MANY_BY_PK_MAX_BATCH_SIZE` or environment variable `COSMOS_READ_MANY_BY_PK_MAX_BATCH_SIZE` (default: 1000, minimum: 1). Follows existing `Configs` patterns. ++ ++## Phase 2 ΓÇö Spark Connector (`azure-cosmos-spark_3`) ++ ++### Step 8: New UDF ΓÇö `GetCosmosPartitionKeyValue` ++ ++- Input: partition key value (single value or Seq for hierarchical PKs). ++- Output: serialized PK string in format `pk([...json...])`. ++- **Null handling:** Throws on null input (Scala convention; callers should filter nulls upstream). ++ ++### Step 9: PK-only serialization helper ++ ++`CosmosPartitionKeyHelper`: ++- `getCosmosPartitionKeyValueString(pkValues: List[Object]): String` ΓÇö serialize to `pk([...])` format. ++- `tryParsePartitionKey(serialized: String): Option[PartitionKey]` ΓÇö deserialize; returns `None` for malformed input including invalid JSON (wrapped in `scala.util.Try`). ++ ++### Step 10: `CosmosItemsDataSource.readManyByPartitionKey` ++ ++Static entry points that accept a DataFrame and Cosmos config. PK extraction supports two modes: ++1. **UDF-produced column**: DataFrame contains `_partitionKeyIdentity` column (from `GetCosmosPartitionKeyValue` UDF). ++2. **Schema-matched columns**: DataFrame columns match the container's PK paths. ++ ++Falls back with `IllegalArgumentException` if neither mode is possible. ++ ++### Step 11: `CosmosReadManyByPartitionKeyReader` ++ ++Orchestrator that resolves schema, initializes and broadcasts client state to executors, then maps each Spark partition to an `ItemsPartitionReaderWithReadManyByPartitionKey`. ++ ++### Step 12: `ItemsPartitionReaderWithReadManyByPartitionKey` ++ ++Spark `PartitionReader[InternalRow]` that: ++- Deduplicates PKs via `LinkedHashMap` (by PK string representation). ++- Passes the pre-built `CosmosReadManyRequestOptions` (with throughput control, diagnostics, custom serializer) to the SDK. ++- Uses `TransientIOErrorsRetryingIterator` for retry handling. ++- Short-circuits empty PK lists to avoid SDK rejection. ++ ++## Phase 3 ΓÇö Testing ++ ++### Unit tests ++- Query construction: single PK, HPK full/partial, custom query composition, table alias detection. ++- Query plan rejection: aggregates, ORDER BY, DISTINCT, GROUP BY (with and without aggregates), DCOUNT. ++- String literal handling: WHERE/parentheses inside string constants. ++- Keyword detection: WHERE, ORDER, GROUP, JOIN, OFFSET, LIMIT, HAVING. ++- PK serialization/deserialization roundtrip (including malformed JSON handling). ++- `findTopLevelWhereIndex` edge cases: subqueries, string literals, case insensitivity. ++ ++### Integration tests ++- End-to-end SDK: single PK basic, projections, filters, empty results, HPK full/partial, request options propagation. ++- Batch size validation: temporarily lowered batch size to exercise batching/interleaving logic. ++- Null/empty PK list rejection (eager validation). ++- Spark connector: `ItemsPartitionReaderWithReadManyByPartitionKey` with known PK values and non-existent PKs. ++- `CosmosPartitionKeyHelper`: single/HPK roundtrip, case insensitivity, malformed input. + diff --git a/sdk/cosmos/.temp/branch-users_fabianm_readManyByPK-vs-origin_main-stat.txt b/sdk/cosmos/.temp/branch-users_fabianm_readManyByPK-vs-origin_main-stat.txt new file mode 100644 index 000000000000..360eb15accb5 --- /dev/null +++ b/sdk/cosmos/.temp/branch-users_fabianm_readManyByPK-vs-origin_main-stat.txt @@ -0,0 +1,76 @@ +===== PR #LOCAL-users_fabianm_readManyByPK ===== +Title: Branch comparison users/fabianm/readManyByPK vs origin/main +Author: Fabian Meiswinkel +Status: DIVERGED (ahead 30, behind 4) +Branch: users/fabianm/readManyByPK -> origin/main +Head SHA: 93957f3a8442d730fe67fbc379ef5399f46f5665 +Merge Base: 20313f79ba8dd0dfa97862d0c31dd4b2e44ee671 +URL: N/A (local branch comparison) + +--- Description --- +Adds readManyByPartitionKey API (sync+async) and Spark connector support for PK-only reads, with query-plan-based validation +--- End Description --- + +===== Commits in PR ===== +9770833eb59 Adding readManyByPartitionKey API +ac287bcdf00 Merge branch 'main' of https://github.com/Azure/azure-sdk-for-java into users/fabianm/readManyByPK +9a5b3e96e7e Update sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/ItemsPartitionReaderWithReadManyByPartitionKey.scala +a8720c3c9f2 Update sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/ItemsPartitionReaderWithReadManyByPartitionKey.scala +d499da76fb4 Update sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java +c3c542a33a7 Update sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelper.java +4416354e03e ┬┤Fixing code review comments +3ab3f0d64f5 Merge branch 'main' of https://github.com/Azure/azure-sdk-for-java into users/fabianm/readManyByPK +588a7550c54 Update CosmosAsyncContainer.java +8c5cdb47b31 Merge branch 'main' into users/fabianm/readManyByPK +f5485527a9c Update ReadManyByPartitionKeyTest.java +f68cf02ff71 Fixing test issues +8b6c4b168ea Update CosmosAsyncContainer.java +8ba7f4db2da Merge branch 'main' into users/fabianm/readManyByPK +56b067a9339 Reacted to code review feedback +fa430e918fa Merge branch 'main' into users/fabianm/readManyByPK +d9504c91f34 Fix build issues +73151f09e5f Merge branch 'main' into users/fabianm/readManyByPK +681830e2d4a Fixing changelog +7f745e60641 Merge branch 'main' into users/fabianm/readManyByPK +0b8905dbb01 Addressing code review comments +22abc780ed8 Addressing code review feedback +662b1a4b90e Update CosmosItemsDataSource.scala +c764de9de02 Update CosmosItemsDataSource.scala +e1e6f5a6f73 Merge branch 'main' into users/fabianm/readManyByPK +080ce4a2293 Update RxDocumentClientImpl.java +516bbf3a95a Merge branch 'users/fabianm/readManyByPK' of https://github.com/Azure/azure-sdk-for-java into users/fabianm/readManyByPK +b01f8758eea Fix readManyByPartitionKey retries +7130d4aa35a Fix PK.None +93957f3a844 Update ReadManyByPartitionKeyQueryHelper.java + +===== Files Changed ===== + sdk/cosmos/azure-cosmos-spark_3-3_2-12/CHANGELOG.md (+1 -0) + sdk/cosmos/azure-cosmos-spark_3-4_2-12/CHANGELOG.md (+1 -0) + sdk/cosmos/azure-cosmos-spark_3-5_2-12/CHANGELOG.md (+1 -0) + sdk/cosmos/azure-cosmos-spark_3-5_2-13/CHANGELOG.md (+1 -0) + sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosConfig.scala (+20 -2) + sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosConstants.scala (+1 -0) + sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosItemsDataSource.scala (+125 -1) + sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosPartitionKeyHelper.scala (+45 -0) + sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosReadManyByPartitionKeyReader.scala (+150 -0) + sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/ItemsPartitionReaderWithReadManyByPartitionKey.scala (+249 -0) + sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/TransientIOErrorsRetryingReadManyByPartitionKeyIterator.scala (+259 -0) + sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/udf/GetCosmosPartitionKeyValue.scala (+25 -0) + sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/CosmosConfigSpec.scala (+42 -0) + sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/CosmosPartitionKeyHelperSpec.scala (+104 -0) + sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/ItemsPartitionReaderWithReadManyByPartitionKeyITest.scala (+158 -0) + sdk/cosmos/azure-cosmos-spark_4-0_2-13/CHANGELOG.md (+1 -0) + sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/ReadManyByPartitionKeyTest.java (+462 -0) + sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelperTest.java (+426 -0) + sdk/cosmos/azure-cosmos/CHANGELOG.md (+1 -0) + sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java (+126 -0) + sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosContainer.java (+67 -0) + sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/AsyncDocumentClient.java (+21 -0) + sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/Configs.java (+19 -0) + sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelper.java (+263 -0) + sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java (+292 -0) + sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/DocumentQueryExecutionContextFactory.java (+11 -0) + sdk/cosmos/cspell.yaml (+6 -0) + sdk/cosmos/docs/readManyByPartitionKey-design.md (+169 -0) + + diff --git a/sdk/cosmos/azure-cosmos-spark_3-3_2-12/CHANGELOG.md b/sdk/cosmos/azure-cosmos-spark_3-3_2-12/CHANGELOG.md index fe114462019c..2b9b9cda9361 100644 --- a/sdk/cosmos/azure-cosmos-spark_3-3_2-12/CHANGELOG.md +++ b/sdk/cosmos/azure-cosmos-spark_3-3_2-12/CHANGELOG.md @@ -3,7 +3,7 @@ ### 4.47.0-beta.1 (Unreleased) #### Features Added -* Added new `CosmosItemsDataSource.readManyByPartitionKey` Spark function to execute bulk queries by a list of pk-values with better efficiency. See [PR 48801](https://github.com/Azure/azure-sdk-for-java/pull/48801) +* Added new `CosmosItemsDataSource.readManyByPartitionKey` Spark function to execute bulk queries by a list of pk-values with better efficiency. Configure null handling via `spark.cosmos.read.readManyByPk.nullHandling` - default `Null` treats a null PK column as JSON null (`addNullValue`), `None` treats it as `PartitionKey.NONE` (`addNoneValue` / `NOT IS_DEFINED`). These route to different physical partitions - picking the wrong mode silently returns zero rows. See [PR 48801](https://github.com/Azure/azure-sdk-for-java/pull/48801) #### Breaking Changes diff --git a/sdk/cosmos/azure-cosmos-spark_3-4_2-12/CHANGELOG.md b/sdk/cosmos/azure-cosmos-spark_3-4_2-12/CHANGELOG.md index 3b2c7ce36db1..771f5974be70 100644 --- a/sdk/cosmos/azure-cosmos-spark_3-4_2-12/CHANGELOG.md +++ b/sdk/cosmos/azure-cosmos-spark_3-4_2-12/CHANGELOG.md @@ -3,7 +3,7 @@ ### 4.47.0-beta.1 (Unreleased) #### Features Added -* Added new `CosmosItemsDataSource.readManyByPartitionKey` Spark function to execute bulk queries by a list of pk-values with better efficiency. See [PR 48801](https://github.com/Azure/azure-sdk-for-java/pull/48801) +* Added new `CosmosItemsDataSource.readManyByPartitionKey` Spark function to execute bulk queries by a list of pk-values with better efficiency. Configure null handling via `spark.cosmos.read.readManyByPk.nullHandling` - default `Null` treats a null PK column as JSON null (`addNullValue`), `None` treats it as `PartitionKey.NONE` (`addNoneValue` / `NOT IS_DEFINED`). These route to different physical partitions - picking the wrong mode silently returns zero rows. See [PR 48801](https://github.com/Azure/azure-sdk-for-java/pull/48801) #### Breaking Changes diff --git a/sdk/cosmos/azure-cosmos-spark_3-5_2-12/CHANGELOG.md b/sdk/cosmos/azure-cosmos-spark_3-5_2-12/CHANGELOG.md index 2240a48b1654..021d2e66929d 100644 --- a/sdk/cosmos/azure-cosmos-spark_3-5_2-12/CHANGELOG.md +++ b/sdk/cosmos/azure-cosmos-spark_3-5_2-12/CHANGELOG.md @@ -3,7 +3,7 @@ ### 4.47.0-beta.1 (Unreleased) #### Features Added -* Added new `CosmosItemsDataSource.readManyByPartitionKey` Spark function to execute bulk queries by a list of pk-values with better efficiency. See [PR 48801](https://github.com/Azure/azure-sdk-for-java/pull/48801) +* Added new `CosmosItemsDataSource.readManyByPartitionKey` Spark function to execute bulk queries by a list of pk-values with better efficiency. Configure null handling via `spark.cosmos.read.readManyByPk.nullHandling` - default `Null` treats a null PK column as JSON null (`addNullValue`), `None` treats it as `PartitionKey.NONE` (`addNoneValue` / `NOT IS_DEFINED`). These route to different physical partitions - picking the wrong mode silently returns zero rows. See [PR 48801](https://github.com/Azure/azure-sdk-for-java/pull/48801) #### Breaking Changes diff --git a/sdk/cosmos/azure-cosmos-spark_3-5_2-13/CHANGELOG.md b/sdk/cosmos/azure-cosmos-spark_3-5_2-13/CHANGELOG.md index 20a3e3a61bd7..f512a65dc491 100644 --- a/sdk/cosmos/azure-cosmos-spark_3-5_2-13/CHANGELOG.md +++ b/sdk/cosmos/azure-cosmos-spark_3-5_2-13/CHANGELOG.md @@ -3,7 +3,7 @@ ### 4.47.0-beta.1 (Unreleased) #### Features Added -* Added new `CosmosItemsDataSource.readManyByPartitionKey` Spark function to execute bulk queries by a list of pk-values with better efficiency. See [PR 48801](https://github.com/Azure/azure-sdk-for-java/pull/48801) +* Added new `CosmosItemsDataSource.readManyByPartitionKey` Spark function to execute bulk queries by a list of pk-values with better efficiency. Configure null handling via `spark.cosmos.read.readManyByPk.nullHandling` - default `Null` treats a null PK column as JSON null (`addNullValue`), `None` treats it as `PartitionKey.NONE` (`addNoneValue` / `NOT IS_DEFINED`). These route to different physical partitions - picking the wrong mode silently returns zero rows. See [PR 48801](https://github.com/Azure/azure-sdk-for-java/pull/48801) #### Breaking Changes diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosConfig.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosConfig.scala index e1b8f0b51f8a..4a483ef38a6b 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosConfig.scala +++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosConfig.scala @@ -1144,11 +1144,12 @@ private object CosmosReadConfig { mandatory = false, defaultValue = Some("Null"), parseFromStringFunction = value => value, - helpMessage = "Determines how null values in hierarchical partition key components are treated " + - "for readManyByPartitionKey. 'Null' (default) maps null to a JSON null value via addNullValue(), " + - "which is appropriate when the document field exists with an explicit null value. " + - "'None' maps null to PartitionKey.NONE via addNoneValue(), which should only be used when the " + - "partition key path does not exist at all in the document." + helpMessage = "Determines how null values in partition key columns are treated for " + + "readManyByPartitionKey. 'Null' (default) maps null to a JSON null via addNullValue(), which " + + "is appropriate when the document field exists with an explicit null value. 'None' maps null " + + "to PartitionKey.NONE via addNoneValue(), which should only be used when the partition key " + + "path does not exist at all in the document. These two semantics hash to DIFFERENT physical " + + "partitions - picking the wrong mode for your data will silently return zero rows." ) def parseCosmosReadConfig(cfg: Map[String, String]): CosmosReadConfig = { diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosItemsDataSource.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosItemsDataSource.scala index 86ef865bcb83..82e3158d765d 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosItemsDataSource.scala +++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosItemsDataSource.scala @@ -5,7 +5,7 @@ package com.azure.cosmos.spark import com.azure.cosmos.models.{CosmosItemIdentity, PartitionKey, PartitionKeyBuilder} import com.azure.cosmos.spark.CosmosPredicates.assertOnSparkDriver import com.azure.cosmos.spark.diagnostics.BasicLoggingTrait -import com.azure.cosmos.{SparkBridgeInternal} +import com.azure.cosmos.SparkBridgeInternal import org.apache.spark.sql.{DataFrame, Row, SparkSession} import java.util @@ -184,6 +184,15 @@ object CosmosItemsDataSource { pkDefinition.getPaths.asScala.map(_.stripPrefix("/")).toList }) + // Nested PK paths (containing /) cannot be resolved from top-level DataFrame columns. + // Surface an explicit error so users know to use the UDF-produced _partitionKeyIdentity column. + if (pkPaths.exists(_.contains("/"))) { + throw new IllegalArgumentException( + "Container has nested partition key path(s) " + pkPaths.mkString("[", ",", "]") + ". " + + "Nested paths cannot be resolved from DataFrame columns automatically - add a " + + "'_partitionKeyIdentity' column produced by the GetCosmosPartitionKeyValue UDF.") + } + // Check if ALL PK path columns exist in the DataFrame schema val dfFieldNames = df.schema.fieldNames.toSet val allPkColumnsPresent = pkPaths.forall(path => dfFieldNames.contains(path)) @@ -195,7 +204,7 @@ object CosmosItemsDataSource { // Single partition key buildPartitionKey(row.getAs[Any](pkPaths.head), treatNullAsNone) } else { - // Hierarchical partition key — build level by level + // Hierarchical partition key - build level by level val builder = new PartitionKeyBuilder() for (path <- pkPaths) { addPartitionKeyComponent(builder, row.getAs[Any](path), treatNullAsNone) @@ -227,7 +236,15 @@ object CosmosItemsDataSource { case null => if (treatNullAsNone) builder.addNoneValue() else builder.addNullValue() - case other => builder.add(other.toString) + case other => + // Reject unknown types rather than silently .toString-ing them - the document field + // was stored with its original type and a stringified value will never match. + // Supported types: String, Number (Byte/Short/Int/Long/Float/Double/BigDecimal), Boolean, null. + throw new IllegalArgumentException( + s"Unsupported partition key column type '${other.getClass.getName}' with value '$other'. " + + "Supported types are String, Number (integral or floating-point), Boolean, and null. " + + "For other source types, convert the column before calling readManyByPartitionKey or use " + + "the GetCosmosPartitionKeyValue UDF to produce a '_partitionKeyIdentity' column.") } } diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosReadManyByPartitionKeyReader.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosReadManyByPartitionKeyReader.scala index 91f3a56bc664..207f77eb4766 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosReadManyByPartitionKeyReader.scala +++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosReadManyByPartitionKeyReader.scala @@ -71,7 +71,12 @@ private[spark] class CosmosReadManyByPartitionKeyReader( classOf[ObjectNode]) .block() } catch { - case _: CosmosException => None + // The warm-up readItem is only used to hydrate the collection/routing-map caches. + // A 404 (item not found) is expected, but we log other CosmosExceptions at debug to + // aid diagnosis (auth failures, throttling, etc.) while not failing reader setup. + case ex: CosmosException => + logDebug(s"Warm-up readItem for metadata caches completed with exception: ${ex.getMessage}", ex) + None } val state = new CosmosClientMetadataCachesSnapshot() diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/ItemsPartitionReaderWithReadManyByPartitionKey.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/ItemsPartitionReaderWithReadManyByPartitionKey.scala index c67cc9c10be1..7fb6e0eb0aba 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/ItemsPartitionReaderWithReadManyByPartitionKey.scala +++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/ItemsPartitionReaderWithReadManyByPartitionKey.scala @@ -6,6 +6,7 @@ package com.azure.cosmos.spark import com.azure.cosmos.{CosmosAsyncContainer, CosmosEndToEndOperationLatencyPolicyConfigBuilder, CosmosItemSerializerNoExceptionWrapping, SparkBridgeInternal} import com.azure.cosmos.implementation.spark.OperationContextAndListenerTuple import com.azure.cosmos.implementation.{ImplementationBridgeHelpers, ObjectNodeMap, SparkRowItem, Utils} +import com.azure.cosmos.BridgeInternal import com.azure.cosmos.models.{CosmosReadManyRequestOptions, ModelBridgeInternal, PartitionKey, PartitionKeyDefinition, SqlQuerySpec} import com.azure.cosmos.spark.BulkWriter.getThreadInfo import com.azure.cosmos.spark.CosmosTableSchemaInferrer.IdAttributeName @@ -154,13 +155,17 @@ private[spark] case class ItemsPartitionReaderWithReadManyByPartitionKey } ) - // Collect all PK values upfront — readManyByPartitionKey needs the full list to - // group by physical partition and issue parallel queries. - // Deduplicate by PK string representation — safe because the list size is bounded - // by the per-call limit of the readManyByPartitionKey API. + // Collect all PK values upfront - readManyByPartitionKey needs the full list to + // group by physical partition (the SDK batches internally per physical partition). + // Deduplicate using the canonical PartitionKeyInternal JSON representation so that + // equivalent PKs built from different runtime types (Int vs Long vs Double) are + // collapsed, and distinct PKs that happen to toString() identically are not. private lazy val pkList = { val seen = new java.util.LinkedHashMap[String, PartitionKey]() - readManyPartitionKeys.foreach(pk => seen.putIfAbsent(pk.toString, pk)) + readManyPartitionKeys.foreach(pk => { + val key = BridgeInternal.getPartitionKeyInternal(pk).toJson + seen.putIfAbsent(key, pk) + }) new java.util.ArrayList[PartitionKey](seen.values()) } @@ -188,9 +193,10 @@ private[spark] case class ItemsPartitionReaderWithReadManyByPartitionKey override def close(): Unit = {} } - // Batch partition keys and retry each batch independently on transient I/O errors. - // This avoids the continuation-token problem with TransientIOErrorsRetryingIterator - // where a retry would re-read all data from scratch, causing silent data duplication. + // Pass the full PK list to the SDK (which batches per physical partition internally). + // On transient I/O failures the retry iterator tracks pages already emitted upstream + // and skips them on replay; if a failure occurs mid-page (after items from that page + // have been emitted) the task fails rather than risking row duplication. private lazy val iterator: CloseableSparkRowItemIterator = if (pkList.isEmpty) { EmptySparkRowItemIterator diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/TransientIOErrorsRetryingReadManyByPartitionKeyIterator.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/TransientIOErrorsRetryingReadManyByPartitionKeyIterator.scala index dcfdf4f93536..f82855218cce 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/TransientIOErrorsRetryingReadManyByPartitionKeyIterator.scala +++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/TransientIOErrorsRetryingReadManyByPartitionKeyIterator.scala @@ -16,16 +16,12 @@ import scala.concurrent.{Await, ExecutionContext, Future} import scala.util.Random import scala.util.control.Breaks -// scalastyle:off underscore.import -import scala.collection.JavaConverters._ -// scalastyle:on underscore.import - /** - * Retry-safe iterator for readManyByPartitionKey that batches partition keys and lazily - * iterates pages within each batch via CosmosPagedIterable — consistent with how - * TransientIOErrorsRetryingIterator handles normal queries. On transient I/O errors the - * current batch's flux is recreated and pages already consumed are replayed, avoiding - * the memory overhead of collectList and matching the query iterator's structure. + * Retry-safe iterator for readManyByPartitionKey. The full partition-key list is passed to the + * SDK in a single call - the SDK is responsible for fan-out and per-physical-partition batching + * (see Configs.getReadManyByPkMaxBatchSize()). This iterator therefore wraps a single + * CosmosPagedIterable and, on transient I/O failures, re-creates the underlying flux and + * skips the pages that were already emitted upstream so no row is delivered twice. */ private[spark] class TransientIOErrorsRetryingReadManyByPartitionKeyIterator[TSparkRow] ( @@ -57,13 +53,20 @@ private[spark] class TransientIOErrorsRetryingReadManyByPartitionKeyIterator[TSp case None => "n/a" } + // Number of pages that have been fully emitted upstream. On retry, we recreate the flux + // and skip this many pages before emitting any item, so already-delivered rows are not + // re-emitted. A page is "committed" as soon as we surface its first item to the caller - + // subsequent failures while still inside that page cannot be recovered from without + // risking duplication, so we fail fast in that case. + private var pagesCommitted: Long = 0 + // Whether the currently-buffered page has emitted at least one item. If true, we have + // passed the point of no return for this page: any transient failure here must surface, + // because we cannot partially-skip within a page on retry. + private var currentPagePartiallyConsumed: Boolean = false + private[spark] var currentFeedResponseIterator: Option[BufferedIterator[FeedResponse[TSparkRow]]] = None private[spark] var currentItemIterator: Option[BufferedIterator[TSparkRow]] = None - private val pkBatchIterator = partitionKeys.asScala.iterator.grouped(pageSize) - // Track the current batch so we can replay it on retry - private var currentBatch: Option[java.util.List[PartitionKey]] = None - override def hasNext: Boolean = { executeWithRetry("hasNextInternal", () => hasNextInternal) } @@ -85,38 +88,39 @@ private[spark] class TransientIOErrorsRetryingReadManyByPartitionKeyIterator[TSp val feedResponseIterator = currentFeedResponseIterator match { case Some(existing) => existing case None => - // Need a new feed response iterator — either for the current batch (on retry) - // or for the next batch - val batch = currentBatch match { - case Some(b) => b // retry of current batch - case None => - if (pkBatchIterator.hasNext) { - val nextBatch = new java.util.ArrayList[PartitionKey](pkBatchIterator.next().toList.asJava) - currentBatch = Some(nextBatch) - nextBatch - } else { - return Some(false) // no more batches - } - } - val pagedFlux = customQuery match { case Some(query) => - container.readManyByPartitionKey(batch, query, queryOptions, classType) + container.readManyByPartitionKey(partitionKeys, query, queryOptions, classType) case None => - container.readManyByPartitionKey(batch, queryOptions, classType) + container.readManyByPartitionKey(partitionKeys, queryOptions, classType) } - currentFeedResponseIterator = Some( - new CosmosPagedIterable[TSparkRow]( - pagedFlux, - pageSize, - pagePrefetchBufferSize - ) - .iterableByPage() - .iterator - .asScala - .buffered + val rawIterator = new CosmosPagedIterable[TSparkRow]( + pagedFlux, + pageSize, + pagePrefetchBufferSize ) + .iterableByPage() + .iterator + + // Skip pages already emitted upstream (replay-safe retry). + var skipped: Long = 0 + while (skipped < pagesCommitted && rawIterator.hasNext) { + rawIterator.next() + skipped += 1 + } + if (skipped < pagesCommitted) { + // The server returned fewer pages than before - cannot safely replay. + // Surface a clean error rather than silently emitting a truncated result. + throw new IllegalStateException( + s"readManyByPartitionKey retry replay failed: expected to skip $pagesCommitted " + + s"already-emitted pages but only $skipped were available. Context: $operationContextString") + } + + // scalastyle:off underscore.import + import scala.collection.JavaConverters._ + // scalastyle:on underscore.import + currentFeedResponseIterator = Some(rawIterator.asScala.buffered) currentFeedResponseIterator.get } @@ -152,20 +156,24 @@ private[spark] class TransientIOErrorsRetryingReadManyByPartitionKeyIterator[TSp operationContextAndListener.get.getOperationContext, feedResponse) } + // scalastyle:off underscore.import + import scala.collection.JavaConverters._ + // scalastyle:on underscore.import val iteratorCandidate = feedResponse.getResults.iterator().asScala.buffered if (iteratorCandidate.hasNext) { currentItemIterator = Some(iteratorCandidate) + currentPagePartiallyConsumed = false Some(true) } else { - // empty page interleaved — try again + // empty page - count it as committed (no items to replay) and try again + pagesCommitted += 1 None } } else { - // Current batch's flux is exhausted — move to next batch - currentBatch = None + // Flux exhausted currentFeedResponseIterator = None - None + Some(false) } } } @@ -175,6 +183,9 @@ private[spark] class TransientIOErrorsRetryingReadManyByPartitionKeyIterator[TSp case Some(iterator) => if (iterator.hasNext) { true } else { + // Entire page drained -> it is now committed for replay-skipping purposes. + pagesCommitted += 1 + currentPagePartiallyConsumed = false currentItemIterator = None false } @@ -183,11 +194,15 @@ private[spark] class TransientIOErrorsRetryingReadManyByPartitionKeyIterator[TSp } override def next(): TSparkRow = { - currentItemIterator.get.next() + executeWithRetry("next", () => { + val value = currentItemIterator.get.next() + currentPagePartiallyConsumed = true + value + }) } - override def head(): TSparkRow = { - currentItemIterator.get.head + override def head: TSparkRow = { + executeWithRetry("head", () => currentItemIterator.get.head) } private[spark] def executeWithRetry[T](methodName: String, func: () => T): T = { @@ -206,6 +221,17 @@ private[spark] class TransientIOErrorsRetryingReadManyByPartitionKeyIterator[TSp catch { case cosmosException: CosmosException => if (Exceptions.canBeTransientFailure(cosmosException.getStatusCode, cosmosException.getSubStatusCode)) { + if (currentPagePartiallyConsumed) { + // We have already emitted items from the current page upstream. Replaying + // the flux would re-skip only completed pages, not items within a page - + // which would cause silent duplication. Fail the task instead. + logError( + s"Transient failure in TransientIOErrorsRetryingReadManyByPartitionKeyIterator." + + s"$methodName after items from the current page were already emitted - " + + s"cannot safely retry without duplicating rows.", + cosmosException) + throw cosmosException + } val retryCountSnapshot = retryCount.incrementAndGet() if (retryCountSnapshot > maxRetryCount) { logError( @@ -217,7 +243,8 @@ private[spark] class TransientIOErrorsRetryingReadManyByPartitionKeyIterator[TSp logWarning( s"Transient failure handled in " + s"TransientIOErrorsRetryingReadManyByPartitionKeyIterator.$methodName -" + - s" will be retried (attempt#$retryCountSnapshot) in ${retryIntervalInMs}ms", + s" will be retried (attempt#$retryCountSnapshot) in ${retryIntervalInMs}ms " + + s"(pagesCommitted=$pagesCommitted)", cosmosException) } } else { @@ -226,7 +253,7 @@ private[spark] class TransientIOErrorsRetryingReadManyByPartitionKeyIterator[TSp case other: Throwable => throw other } - // Reset iterators but keep currentBatch so the batch is replayed + // Reset iterators; pagesCommitted is intentionally preserved so replay can skip them. currentItemIterator = None currentFeedResponseIterator = None Thread.sleep(retryIntervalInMs) @@ -236,6 +263,10 @@ private[spark] class TransientIOErrorsRetryingReadManyByPartitionKeyIterator[TSp returnValue.get } + // Clean up iterator references - the underlying Reactor subscription from + // CosmosPagedIterable.iterator will be cleaned up when the iterator is GC'd. + // This matches the behavior of TransientIOErrorsRetryingIterator; any still-prefetched + // pages are discarded with the iterator. override def close(): Unit = { currentItemIterator = None currentFeedResponseIterator = None diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/udf/GetCosmosPartitionKeyValue.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/udf/GetCosmosPartitionKeyValue.scala index a58d5b723b8b..fc038861a19a 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/udf/GetCosmosPartitionKeyValue.scala +++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/udf/GetCosmosPartitionKeyValue.scala @@ -4,22 +4,25 @@ package com.azure.cosmos.spark.udf import com.azure.cosmos.spark.CosmosPartitionKeyHelper -import com.azure.cosmos.spark.CosmosPredicates.requireNotNull import org.apache.spark.sql.api.java.UDF1 @SerialVersionUID(1L) class GetCosmosPartitionKeyValue extends UDF1[Object, String] { - override def call - ( - partitionKeyValue: Object - ): String = { - requireNotNull(partitionKeyValue, "partitionKeyValue") - + // Null is a valid partition-key value (JSON null). A null input is serialized as a + // single-level partition key with a JSON null component; parsing that string back via + // CosmosPartitionKeyHelper.tryParsePartitionKey yields a PartitionKey built with + // addNullValue(). If the caller instead wants PartitionKey.NONE semantics (absent PK + // field) they should filter the null row before calling this UDF and use the + // schema-matched readManyByPartitionKey path with readManyByPk.nullHandling=None. + override def call(partitionKeyValue: Object): String = { partitionKeyValue match { + case null => + CosmosPartitionKeyHelper.getCosmosPartitionKeyValueString(List(null)) // for subpartitions case - Seq covers both WrappedArray (Scala 2.12) and ArraySeq (Scala 2.13) case seq: Seq[Any] => CosmosPartitionKeyHelper.getCosmosPartitionKeyValueString(seq.map(_.asInstanceOf[Object]).toList) - case _ => CosmosPartitionKeyHelper.getCosmosPartitionKeyValueString(List(partitionKeyValue)) + case _ => + CosmosPartitionKeyHelper.getCosmosPartitionKeyValueString(List(partitionKeyValue)) } } -} +} \ No newline at end of file diff --git a/sdk/cosmos/azure-cosmos-spark_4-0_2-13/CHANGELOG.md b/sdk/cosmos/azure-cosmos-spark_4-0_2-13/CHANGELOG.md index d8368be6a0da..243863e8de65 100644 --- a/sdk/cosmos/azure-cosmos-spark_4-0_2-13/CHANGELOG.md +++ b/sdk/cosmos/azure-cosmos-spark_4-0_2-13/CHANGELOG.md @@ -3,7 +3,7 @@ ### 4.47.0-beta.1 (Unreleased) #### Features Added -* Added new `CosmosItemsDataSource.readManyByPartitionKey` Spark function to execute bulk queries by a list of pk-values with better efficiency. See [PR 48801](https://github.com/Azure/azure-sdk-for-java/pull/48801) +* Added new `CosmosItemsDataSource.readManyByPartitionKey` Spark function to execute bulk queries by a list of pk-values with better efficiency. Configure null handling via `spark.cosmos.read.readManyByPk.nullHandling` - default `Null` treats a null PK column as JSON null (`addNullValue`), `None` treats it as `PartitionKey.NONE` (`addNoneValue` / `NOT IS_DEFINED`). These route to different physical partitions - picking the wrong mode silently returns zero rows. See [PR 48801](https://github.com/Azure/azure-sdk-for-java/pull/48801) #### Breaking Changes diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/Configs.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/Configs.java index 162b0740f408..651019d91652 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/Configs.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/Configs.java @@ -58,7 +58,7 @@ public class Configs { private static final String NETTY_HTTP_CLIENT_METRICS_ENABLED = "COSMOS.NETTY_HTTP_CLIENT_METRICS_ENABLED"; private static final String NETTY_HTTP_CLIENT_METRICS_ENABLED_VARIABLE = "COSMOS_NETTY_HTTP_CLIENT_METRICS_ENABLED"; - // Thin client connect/acquire timeout — controls CONNECT_TIMEOUT_MILLIS for Gateway V2 data plane endpoints. + // Thin client connect/acquire timeout - controls CONNECT_TIMEOUT_MILLIS for Gateway V2 data plane endpoints. // Data plane requests are routed to the thin client regional endpoint (from RegionalRoutingContext) // which uses a non-443 port. These get a shorter 5s connect/acquire timeout. // Metadata requests target Gateway V1 endpoint (port 443) and retain the full 45s/60s timeout (unchanged). @@ -249,9 +249,9 @@ public class Configs { public static final int DEFAULT_MIN_TARGET_BULK_MICRO_BATCH_SIZE = 1; // readManyByPartitionKey: max number of PK values per query per physical partition - private static final String READ_MANY_BY_PK_MAX_BATCH_SIZE = "COSMOS.READ_MANY_BY_PK_MAX_BATCH_SIZE"; - private static final String READ_MANY_BY_PK_MAX_BATCH_SIZE_VARIABLE = "COSMOS_READ_MANY_BY_PK_MAX_BATCH_SIZE"; - private static final int DEFAULT_READ_MANY_BY_PK_MAX_BATCH_SIZE = 1000; + public static final String READ_MANY_BY_PK_MAX_BATCH_SIZE = "COSMOS.READ_MANY_BY_PK_MAX_BATCH_SIZE"; + public static final String READ_MANY_BY_PK_MAX_BATCH_SIZE_VARIABLE = "COSMOS_READ_MANY_BY_PK_MAX_BATCH_SIZE"; + public static final int DEFAULT_READ_MANY_BY_PK_MAX_BATCH_SIZE = 100; public static final String MAX_BULK_MICRO_BATCH_CONCURRENCY = "COSMOS.MAX_BULK_MICRO_BATCH_CONCURRENCY"; public static final String MAX_BULK_MICRO_BATCH_CONCURRENCY_VARIABLE = "COSMOS_MAX_BULK_MICRO_BATCH_CONCURRENCY"; @@ -683,7 +683,7 @@ public static int getThinClientConnectionTimeoutInMs() { } } - // Guard against invalid values — timeout must be at least 500ms + // Guard against invalid values - timeout must be at least 500ms if (value < 500) { logger.warn( "Invalid thin client connection timeout: {}ms. Must be >= 500. Falling back to default: {}ms.", diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelper.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelper.java index 6d6cd084e01a..24ac8ea19666 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelper.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelper.java @@ -20,7 +20,7 @@ public class ReadManyByPartitionKeyQueryHelper { private static final String DEFAULT_TABLE_ALIAS = "c"; - // Internal parameter prefix — uses double-underscore to avoid collisions with user-provided parameters + // Internal parameter prefix - uses double-underscore to avoid collisions with user-provided parameters private static final String PK_PARAM_PREFIX = "@__rmPk_"; public static SqlQuerySpec createReadManyByPkQuerySpec( @@ -30,7 +30,19 @@ public static SqlQuerySpec createReadManyByPkQuerySpec( List partitionKeySelectors, PartitionKeyDefinition pkDefinition) { - // Extract the table alias from the FROM clause (e.g. "FROM x" → "x", "FROM c" → "c") + // Guard against collisions with our internal parameter names - callers cannot realistically + // use the @__rmPk_ prefix for their own parameters, but if they do we surface a clear error + // rather than letting the server reject a SqlQuerySpec with duplicate parameter names. + for (SqlParameter baseParam : baseParameters) { + String name = baseParam.getName(); + if (name != null && name.startsWith(PK_PARAM_PREFIX)) { + throw new IllegalArgumentException( + "Custom query parameter name '" + name + "' collides with the reserved " + + "readManyByPartitionKey internal prefix '" + PK_PARAM_PREFIX + "'. Rename the parameter."); + } + } + + // Extract the table alias from the FROM clause (e.g. "FROM x" -> "x", "FROM c" -> "c") String tableAlias = extractTableAlias(baseQueryText); StringBuilder pkFilter = new StringBuilder(); @@ -40,7 +52,7 @@ public static SqlQuerySpec createReadManyByPkQuerySpec( boolean isSinglePathPk = partitionKeySelectors.size() == 1; if (isSinglePathPk && pkDefinition.getKind() != PartitionKind.MULTI_HASH) { - // Single PK path — use IN clause for normal values, OR NOT IS_DEFINED for NONE + // Single PK path - use IN clause for normal values, OR NOT IS_DEFINED for NONE // First, separate NONE PKs from normal PKs boolean hasNone = false; List normalPkValues = new ArrayList<>(); @@ -89,13 +101,13 @@ public static SqlQuerySpec createReadManyByPkQuerySpec( pkFilter.append(")"); } } else { - // Multiple PK paths (HPK) or MULTI_HASH — use OR of AND clauses + // Multiple PK paths (HPK) or MULTI_HASH - use OR of AND clauses pkFilter.append(" "); for (int i = 0; i < pkValues.size(); i++) { PartitionKeyInternal pkInternal = BridgeInternal.getPartitionKeyInternal(pkValues.get(i)); Object[] pkComponents = pkInternal.toObjectArray(); - // PartitionKey.NONE — generate NOT IS_DEFINED for all PK paths + // PartitionKey.NONE - generate NOT IS_DEFINED for all PK paths if (pkComponents == null) { pkFilter.append("("); for (int j = 0; j < partitionKeySelectors.size(); j++) { @@ -136,12 +148,12 @@ public static SqlQuerySpec createReadManyByPkQuerySpec( String finalQuery; int whereIndex = findTopLevelWhereIndex(baseQueryText); if (whereIndex >= 0) { - // Base query has WHERE — AND our PK filter + // Base query has WHERE - AND our PK filter String beforeWhere = baseQueryText.substring(0, whereIndex); String afterWhere = baseQueryText.substring(whereIndex + 5); // skip "WHERE" finalQuery = beforeWhere + "WHERE (" + afterWhere.trim() + ") AND (" + pkFilter.toString().trim() + ")"; } else { - // No WHERE — add one + // No WHERE - add one finalQuery = baseQueryText + " WHERE" + pkFilter.toString(); } @@ -197,7 +209,7 @@ static String extractTableAlias(String queryText) { return containerName; } } - // Otherwise the next token is the alias ("FROM root r" → alias is "r") + // Otherwise the next token is the alias ("FROM root r" -> alias is "r") int aliasStart = afterFrom; while (afterFrom < queryText.length() && !Character.isWhitespace(queryText.charAt(afterFrom)) @@ -230,7 +242,7 @@ static int findTopLevelKeywordIndex(String queryText, String keyword) { while (i < queryText.length()) { if (queryText.charAt(i) == '\'') { if (i + 1 < queryText.length() && queryText.charAt(i + 1) == '\'') { - i += 2; // escaped quote — skip both + i += 2; // escaped quote - skip both continue; } break; // end of string literal diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java index c70dedaa1f2c..3d9082a25bcc 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java @@ -4499,7 +4499,7 @@ private Flux> readManyByPartitionKey( } // Build per-physical-partition batched queries. - // Each physical partition may have many PKs — split into batches + // Each physical partition may have many PKs - split into batches // to avoid oversized SQL queries. Batch size is configurable via // system property COSMOS.READ_MANY_BY_PK_MAX_BATCH_SIZE (default 1000). int maxPksPerPartitionQuery = Configs.getReadManyByPkMaxBatchSize(); @@ -4534,9 +4534,11 @@ private Flux> readManyByPartitionKey( return Flux.empty(); } - // Round-robin interleave: [batch0-p1, batch0-p2, ..., batch0-pN, batch1-p1, batch1-p2, ...] - // This ensures that with bounded concurrency, different partitions are - // preferred over sequential batches of the same partition. + // Interleave batches across physical partitions so that the first batch for + // each partition is kicked off before the second batch of any partition. With bounded + // concurrency this spreads the initial wave of work across the cluster; note that + // skewed distributions (one partition with N batches, another with 1) will eventually + // fall back to sequential execution once the short partitions are drained. List> interleavedBatches = new ArrayList<>(); for (int batchIdx = 0; batchIdx < maxBatchesPerPartition; batchIdx++) { for (List> partitionBatches : batchesPerPartition) { @@ -4625,7 +4627,7 @@ private Map> groupPartitionKeysByPhysicalP PartitionKeyInternal pkInternal = BridgeInternal.getPartitionKeyInternal(pk); // PartitionKey.NONE wraps NonePartitionKey which has components = null. - // For routing purposes, treat NONE as UndefinedPartitionKey — documents ingested + // For routing purposes, treat NONE as UndefinedPartitionKey - documents ingested // without a partition key path are stored with the undefined EPK. PartitionKeyInternal effectivePkInternal = pkInternal.getComponents() == null ? PartitionKeyInternal.UndefinedPartitionKey @@ -4637,12 +4639,12 @@ private Map> groupPartitionKeysByPhysicalP List targetRanges; if (pkDefinition.getKind() == PartitionKind.MULTI_HASH && componentCount < definedPathCount) { - // Partial HPK — compute EPK prefix range and find all overlapping physical partitions + // Partial HPK - compute EPK prefix range and find all overlapping physical partitions Range epkRange = PartitionKeyInternalHelper.getEPKRangeForPrefixPartitionKey( effectivePkInternal, pkDefinition); targetRanges = routingMap.getOverlappingRanges(epkRange); } else { - // Full PK — maps to exactly one physical partition + // Full PK - maps to exactly one physical partition String effectivePartitionKeyString = PartitionKeyInternalHelper .getEffectivePartitionKeyString(effectivePkInternal, pkDefinition); PartitionKeyRange range = routingMap.getRangeByEffectivePartitionKey(effectivePartitionKeyString); From c96b6f6935913e3e3358fa74d591f5ebe1ff0663 Mon Sep 17 00:00:00 2001 From: Fabian Meiswinkel Date: Fri, 17 Apr 2026 13:38:57 +0000 Subject: [PATCH 23/64] Reacting to code review feedback --- sdk/cosmos/.gitignore | 2 + ...bianm_readManyByPK-vs-origin_main-full.txt | 3444 ----------------- ...bianm_readManyByPK-vs-origin_main-stat.txt | 76 - .../cosmos/spark/CosmosItemsDataSource.scala | 17 +- .../spark/CosmosPartitionKeyHelper.scala | 43 +- .../CosmosReadManyByPartitionKeyReader.scala | 8 + .../cosmos/ReadManyByPartitionKeyTest.java | 37 + .../azure/cosmos/CosmosAsyncContainer.java | 46 +- .../com/azure/cosmos/CosmosContainer.java | 29 +- .../ReadManyByPartitionKeyQueryHelper.java | 72 +- .../implementation/RxDocumentClientImpl.java | 8 +- 11 files changed, 218 insertions(+), 3564 deletions(-) delete mode 100644 sdk/cosmos/.temp/branch-users_fabianm_readManyByPK-vs-origin_main-full.txt delete mode 100644 sdk/cosmos/.temp/branch-users_fabianm_readManyByPK-vs-origin_main-stat.txt diff --git a/sdk/cosmos/.gitignore b/sdk/cosmos/.gitignore index 1ea74182f6bb..81d278c1728f 100644 --- a/sdk/cosmos/.gitignore +++ b/sdk/cosmos/.gitignore @@ -2,3 +2,5 @@ metastore_db/* spark-warehouse/* + +.temp/ diff --git a/sdk/cosmos/.temp/branch-users_fabianm_readManyByPK-vs-origin_main-full.txt b/sdk/cosmos/.temp/branch-users_fabianm_readManyByPK-vs-origin_main-full.txt deleted file mode 100644 index 5eee65a15a7b..000000000000 --- a/sdk/cosmos/.temp/branch-users_fabianm_readManyByPK-vs-origin_main-full.txt +++ /dev/null @@ -1,3444 +0,0 @@ -===== PR #LOCAL-users_fabianm_readManyByPK ===== -Title: Branch comparison users/fabianm/readManyByPK vs origin/main -Author: Fabian Meiswinkel -Status: DIVERGED (ahead 30, behind 4) -Branch: users/fabianm/readManyByPK -> origin/main -Head SHA: 93957f3a8442d730fe67fbc379ef5399f46f5665 -Merge Base: 20313f79ba8dd0dfa97862d0c31dd4b2e44ee671 -URL: N/A (local branch comparison) - ---- Description --- -Adds readManyByPartitionKey API (sync+async) and Spark connector support for PK-only reads, with query-plan-based validation ---- End Description --- - -===== Commits in PR ===== -9770833eb59 Adding readManyByPartitionKey API -ac287bcdf00 Merge branch 'main' of https://github.com/Azure/azure-sdk-for-java into users/fabianm/readManyByPK -9a5b3e96e7e Update sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/ItemsPartitionReaderWithReadManyByPartitionKey.scala -a8720c3c9f2 Update sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/ItemsPartitionReaderWithReadManyByPartitionKey.scala -d499da76fb4 Update sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java -c3c542a33a7 Update sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelper.java -4416354e03e ┬┤Fixing code review comments -3ab3f0d64f5 Merge branch 'main' of https://github.com/Azure/azure-sdk-for-java into users/fabianm/readManyByPK -588a7550c54 Update CosmosAsyncContainer.java -8c5cdb47b31 Merge branch 'main' into users/fabianm/readManyByPK -f5485527a9c Update ReadManyByPartitionKeyTest.java -f68cf02ff71 Fixing test issues -8b6c4b168ea Update CosmosAsyncContainer.java -8ba7f4db2da Merge branch 'main' into users/fabianm/readManyByPK -56b067a9339 Reacted to code review feedback -fa430e918fa Merge branch 'main' into users/fabianm/readManyByPK -d9504c91f34 Fix build issues -73151f09e5f Merge branch 'main' into users/fabianm/readManyByPK -681830e2d4a Fixing changelog -7f745e60641 Merge branch 'main' into users/fabianm/readManyByPK -0b8905dbb01 Addressing code review comments -22abc780ed8 Addressing code review feedback -662b1a4b90e Update CosmosItemsDataSource.scala -c764de9de02 Update CosmosItemsDataSource.scala -e1e6f5a6f73 Merge branch 'main' into users/fabianm/readManyByPK -080ce4a2293 Update RxDocumentClientImpl.java -516bbf3a95a Merge branch 'users/fabianm/readManyByPK' of https://github.com/Azure/azure-sdk-for-java into users/fabianm/readManyByPK -b01f8758eea Fix readManyByPartitionKey retries -7130d4aa35a Fix PK.None -93957f3a844 Update ReadManyByPartitionKeyQueryHelper.java - -===== Files Changed ===== - sdk/cosmos/azure-cosmos-spark_3-3_2-12/CHANGELOG.md (+1 -0) - sdk/cosmos/azure-cosmos-spark_3-4_2-12/CHANGELOG.md (+1 -0) - sdk/cosmos/azure-cosmos-spark_3-5_2-12/CHANGELOG.md (+1 -0) - sdk/cosmos/azure-cosmos-spark_3-5_2-13/CHANGELOG.md (+1 -0) - sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosConfig.scala (+20 -2) - sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosConstants.scala (+1 -0) - sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosItemsDataSource.scala (+125 -1) - sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosPartitionKeyHelper.scala (+45 -0) - sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosReadManyByPartitionKeyReader.scala (+150 -0) - sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/ItemsPartitionReaderWithReadManyByPartitionKey.scala (+249 -0) - sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/TransientIOErrorsRetryingReadManyByPartitionKeyIterator.scala (+259 -0) - sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/udf/GetCosmosPartitionKeyValue.scala (+25 -0) - sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/CosmosConfigSpec.scala (+42 -0) - sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/CosmosPartitionKeyHelperSpec.scala (+104 -0) - sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/ItemsPartitionReaderWithReadManyByPartitionKeyITest.scala (+158 -0) - sdk/cosmos/azure-cosmos-spark_4-0_2-13/CHANGELOG.md (+1 -0) - sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/ReadManyByPartitionKeyTest.java (+462 -0) - sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelperTest.java (+426 -0) - sdk/cosmos/azure-cosmos/CHANGELOG.md (+1 -0) - sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java (+126 -0) - sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosContainer.java (+67 -0) - sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/AsyncDocumentClient.java (+21 -0) - sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/Configs.java (+19 -0) - sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelper.java (+263 -0) - sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java (+292 -0) - sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/DocumentQueryExecutionContextFactory.java (+11 -0) - sdk/cosmos/cspell.yaml (+6 -0) - sdk/cosmos/docs/readManyByPartitionKey-design.md (+169 -0) - -===== Full Diff ===== -diff --git a/sdk/cosmos/azure-cosmos-spark_3-3_2-12/CHANGELOG.md b/sdk/cosmos/azure-cosmos-spark_3-3_2-12/CHANGELOG.md -index cbf97c610f9..fe114462019 100644 ---- a/sdk/cosmos/azure-cosmos-spark_3-3_2-12/CHANGELOG.md -+++ b/sdk/cosmos/azure-cosmos-spark_3-3_2-12/CHANGELOG.md -@@ -3,6 +3,7 @@ - ### 4.47.0-beta.1 (Unreleased) - - #### Features Added -+* Added new `CosmosItemsDataSource.readManyByPartitionKey` Spark function to execute bulk queries by a list of pk-values with better efficiency. See [PR 48801](https://github.com/Azure/azure-sdk-for-java/pull/48801) - - #### Breaking Changes - -diff --git a/sdk/cosmos/azure-cosmos-spark_3-4_2-12/CHANGELOG.md b/sdk/cosmos/azure-cosmos-spark_3-4_2-12/CHANGELOG.md -index c9097e749f0..3b2c7ce36db 100644 ---- a/sdk/cosmos/azure-cosmos-spark_3-4_2-12/CHANGELOG.md -+++ b/sdk/cosmos/azure-cosmos-spark_3-4_2-12/CHANGELOG.md -@@ -3,6 +3,7 @@ - ### 4.47.0-beta.1 (Unreleased) - - #### Features Added -+* Added new `CosmosItemsDataSource.readManyByPartitionKey` Spark function to execute bulk queries by a list of pk-values with better efficiency. See [PR 48801](https://github.com/Azure/azure-sdk-for-java/pull/48801) - - #### Breaking Changes - -diff --git a/sdk/cosmos/azure-cosmos-spark_3-5_2-12/CHANGELOG.md b/sdk/cosmos/azure-cosmos-spark_3-5_2-12/CHANGELOG.md -index f5eac38bdb7..2240a48b165 100644 ---- a/sdk/cosmos/azure-cosmos-spark_3-5_2-12/CHANGELOG.md -+++ b/sdk/cosmos/azure-cosmos-spark_3-5_2-12/CHANGELOG.md -@@ -3,6 +3,7 @@ - ### 4.47.0-beta.1 (Unreleased) - - #### Features Added -+* Added new `CosmosItemsDataSource.readManyByPartitionKey` Spark function to execute bulk queries by a list of pk-values with better efficiency. See [PR 48801](https://github.com/Azure/azure-sdk-for-java/pull/48801) - - #### Breaking Changes - -diff --git a/sdk/cosmos/azure-cosmos-spark_3-5_2-13/CHANGELOG.md b/sdk/cosmos/azure-cosmos-spark_3-5_2-13/CHANGELOG.md -index 919d7fbfa32..20a3e3a61bd 100644 ---- a/sdk/cosmos/azure-cosmos-spark_3-5_2-13/CHANGELOG.md -+++ b/sdk/cosmos/azure-cosmos-spark_3-5_2-13/CHANGELOG.md -@@ -3,6 +3,7 @@ - ### 4.47.0-beta.1 (Unreleased) - - #### Features Added -+* Added new `CosmosItemsDataSource.readManyByPartitionKey` Spark function to execute bulk queries by a list of pk-values with better efficiency. See [PR 48801](https://github.com/Azure/azure-sdk-for-java/pull/48801) - - #### Breaking Changes - -diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosConfig.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosConfig.scala -index 951f4735444..e1b8f0b51f8 100644 ---- a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosConfig.scala -+++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosConfig.scala -@@ -92,6 +92,7 @@ private[spark] object CosmosConfigNames { - val ReadPartitioningFeedRangeFilter = "spark.cosmos.partitioning.feedRangeFilter" - val ReadRuntimeFilteringEnabled = "spark.cosmos.read.runtimeFiltering.enabled" - val ReadManyFilteringEnabled = "spark.cosmos.read.readManyFiltering.enabled" -+ val ReadManyByPkNullHandling = "spark.cosmos.read.readManyByPk.nullHandling" - val ViewsRepositoryPath = "spark.cosmos.views.repositoryPath" - val DiagnosticsMode = "spark.cosmos.diagnostics" - val DiagnosticsSamplingMaxCount = "spark.cosmos.diagnostics.sampling.maxCount" -@@ -226,6 +227,7 @@ private[spark] object CosmosConfigNames { - ReadPartitioningFeedRangeFilter, - ReadRuntimeFilteringEnabled, - ReadManyFilteringEnabled, -+ ReadManyByPkNullHandling, - ViewsRepositoryPath, - DiagnosticsMode, - DiagnosticsSamplingIntervalInSeconds, -@@ -1042,7 +1044,8 @@ private case class CosmosReadConfig(readConsistencyStrategy: ReadConsistencyStra - throughputControlConfig: Option[CosmosThroughputControlConfig] = None, - runtimeFilteringEnabled: Boolean, - readManyFilteringConfig: CosmosReadManyFilteringConfig, -- responseContinuationTokenLimitInKb: Option[Int] = None) -+ responseContinuationTokenLimitInKb: Option[Int] = None, -+ readManyByPkTreatNullAsNone: Boolean = false) - - private object SchemaConversionModes extends Enumeration { - type SchemaConversionMode = Value -@@ -1136,6 +1139,18 @@ private object CosmosReadConfig { - helpMessage = " Indicates whether dynamic partition pruning filters will be pushed down when applicable." - ) - -+ private val ReadManyByPkNullHandling = CosmosConfigEntry[String]( -+ key = CosmosConfigNames.ReadManyByPkNullHandling, -+ mandatory = false, -+ defaultValue = Some("Null"), -+ parseFromStringFunction = value => value, -+ helpMessage = "Determines how null values in hierarchical partition key components are treated " + -+ "for readManyByPartitionKey. 'Null' (default) maps null to a JSON null value via addNullValue(), " + -+ "which is appropriate when the document field exists with an explicit null value. " + -+ "'None' maps null to PartitionKey.NONE via addNoneValue(), which should only be used when the " + -+ "partition key path does not exist at all in the document." -+ ) -+ - def parseCosmosReadConfig(cfg: Map[String, String]): CosmosReadConfig = { - val forceEventualConsistency = CosmosConfigEntry.parse(cfg, ForceEventualConsistency) - val readConsistencyStrategyOverride = CosmosConfigEntry.parse(cfg, ReadConsistencyStrategyOverride) -@@ -1158,6 +1173,8 @@ private object CosmosReadConfig { - val throughputControlConfigOpt = CosmosThroughputControlConfig.parseThroughputControlConfig(cfg) - val runtimeFilteringEnabled = CosmosConfigEntry.parse(cfg, ReadRuntimeFilteringEnabled) - val readManyFilteringConfig = CosmosReadManyFilteringConfig.parseCosmosReadManyFilterConfig(cfg) -+ val readManyByPkNullHandling = CosmosConfigEntry.parse(cfg, ReadManyByPkNullHandling) -+ val readManyByPkTreatNullAsNone = readManyByPkNullHandling.getOrElse("Null").equalsIgnoreCase("None") - - val effectiveReadConsistencyStrategy = if (readConsistencyStrategyOverride.getOrElse(ReadConsistencyStrategy.DEFAULT) != ReadConsistencyStrategy.DEFAULT) { - readConsistencyStrategyOverride.get -@@ -1189,7 +1206,8 @@ private object CosmosReadConfig { - throughputControlConfigOpt, - runtimeFilteringEnabled.get, - readManyFilteringConfig, -- responseContinuationTokenLimitInKb) -+ responseContinuationTokenLimitInKb, -+ readManyByPkTreatNullAsNone) - } - } - -diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosConstants.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosConstants.scala -index 9ece4741652..00761f23d39 100644 ---- a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosConstants.scala -+++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosConstants.scala -@@ -45,6 +45,7 @@ private[cosmos] object CosmosConstants { - val Id = "id" - val ETag = "_etag" - val ItemIdentity = "_itemIdentity" -+ val PartitionKeyIdentity = "_partitionKeyIdentity" - } - - object StatusCodes { -diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosItemsDataSource.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosItemsDataSource.scala -index a35cff27af6..86ef865bcb8 100644 ---- a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosItemsDataSource.scala -+++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosItemsDataSource.scala -@@ -2,9 +2,10 @@ - // Licensed under the MIT License. - package com.azure.cosmos.spark - --import com.azure.cosmos.models.{CosmosItemIdentity, PartitionKey} -+import com.azure.cosmos.models.{CosmosItemIdentity, PartitionKey, PartitionKeyBuilder} - import com.azure.cosmos.spark.CosmosPredicates.assertOnSparkDriver - import com.azure.cosmos.spark.diagnostics.BasicLoggingTrait -+import com.azure.cosmos.{SparkBridgeInternal} - import org.apache.spark.sql.{DataFrame, Row, SparkSession} - - import java.util -@@ -112,4 +113,127 @@ object CosmosItemsDataSource { - - readManyReader.readMany(df.rdd, readManyFilterExtraction) - } -+ -+ def readManyByPartitionKey(df: DataFrame, userConfig: java.util.Map[String, String]): DataFrame = { -+ readManyByPartitionKey(df, userConfig, null) -+ } -+ -+ def readManyByPartitionKey( -+ df: DataFrame, -+ userConfig: java.util.Map[String, String], -+ userProvidedSchema: StructType): DataFrame = { -+ -+ val readManyReader = new CosmosReadManyByPartitionKeyReader( -+ userProvidedSchema, -+ userConfig.asScala.toMap) -+ -+ // Option 1: Look for the _partitionKeyIdentity column (produced by GetCosmosPartitionKeyValue UDF) -+ val pkIdentityFieldExtraction = df -+ .schema -+ .find(field => field.name.equals(CosmosConstants.Properties.PartitionKeyIdentity) && field.dataType.equals(StringType)) -+ .map(field => (row: Row) => { -+ val rawValue = row.getString(row.fieldIndex(field.name)) -+ CosmosPartitionKeyHelper.tryParsePartitionKey(rawValue) -+ .getOrElse(throw new IllegalArgumentException( -+ s"Invalid _partitionKeyIdentity value in row: '$rawValue'. " + -+ "Expected format: pk([...json...])")) -+ }) -+ -+ // Option 2: Detect PK columns by matching the container's partition key paths against the DataFrame schema -+ val pkColumnExtraction: Option[Row => PartitionKey] = if (pkIdentityFieldExtraction.isDefined) { -+ None // no need to resolve PK paths - _partitionKeyIdentity column takes precedence -+ } else { -+ val effectiveConfig = CosmosConfig.getEffectiveConfig( -+ databaseName = None, -+ containerName = None, -+ userConfig.asScala.toMap) -+ val readConfig = CosmosReadConfig.parseCosmosReadConfig(effectiveConfig) -+ val containerConfig = CosmosContainerConfig.parseCosmosContainerConfig(effectiveConfig) -+ val sparkEnvironmentInfo = CosmosClientConfiguration.getSparkEnvironmentInfo(None) -+ val calledFrom = s"CosmosItemsDataSource.readManyByPartitionKey" -+ val treatNullAsNone = readConfig.readManyByPkTreatNullAsNone -+ -+ val pkPaths = Loan( -+ List[Option[CosmosClientCacheItem]]( -+ Some( -+ CosmosClientCache( -+ CosmosClientConfiguration( -+ effectiveConfig, -+ readConsistencyStrategy = readConfig.readConsistencyStrategy, -+ sparkEnvironmentInfo), -+ None, -+ calledFrom)), -+ ThroughputControlHelper.getThroughputControlClientCacheItem( -+ effectiveConfig, -+ calledFrom, -+ None, -+ sparkEnvironmentInfo) -+ )) -+ .to(clientCacheItems => { -+ val container = -+ ThroughputControlHelper.getContainer( -+ effectiveConfig, -+ containerConfig, -+ clientCacheItems(0).get, -+ clientCacheItems(1)) -+ -+ val pkDefinition = SparkBridgeInternal -+ .getContainerPropertiesFromCollectionCache(container) -+ .getPartitionKeyDefinition -+ -+ pkDefinition.getPaths.asScala.map(_.stripPrefix("/")).toList -+ }) -+ -+ // Check if ALL PK path columns exist in the DataFrame schema -+ val dfFieldNames = df.schema.fieldNames.toSet -+ val allPkColumnsPresent = pkPaths.forall(path => dfFieldNames.contains(path)) -+ -+ if (allPkColumnsPresent && pkPaths.nonEmpty) { -+ // pkPaths already defined above -+ Some((row: Row) => { -+ if (pkPaths.size == 1) { -+ // Single partition key -+ buildPartitionKey(row.getAs[Any](pkPaths.head), treatNullAsNone) -+ } else { -+ // Hierarchical partition key ΓÇö build level by level -+ val builder = new PartitionKeyBuilder() -+ for (path <- pkPaths) { -+ addPartitionKeyComponent(builder, row.getAs[Any](path), treatNullAsNone) -+ } -+ builder.build() -+ } -+ }) -+ } else { -+ None -+ } -+ } -+ -+ val pkExtraction = pkIdentityFieldExtraction -+ .orElse(pkColumnExtraction) -+ .getOrElse( -+ throw new IllegalArgumentException( -+ "Cannot determine partition key extraction from the input DataFrame. " + -+ "Either add a '_partitionKeyIdentity' column (using the GetCosmosPartitionKeyValue UDF) " + -+ "or ensure the DataFrame contains columns matching the container's partition key paths.")) -+ -+ readManyReader.readManyByPartitionKey(df.rdd, pkExtraction) -+ } -+ -+ private def addPartitionKeyComponent(builder: PartitionKeyBuilder, value: Any, treatNullAsNone: Boolean): Unit = { -+ value match { -+ case s: String => builder.add(s) -+ case n: Number => builder.add(n.doubleValue()) -+ case b: Boolean => builder.add(b) -+ case null => -+ if (treatNullAsNone) builder.addNoneValue() -+ else builder.addNullValue() -+ case other => builder.add(other.toString) -+ } -+ } -+ -+ private def buildPartitionKey(value: Any, treatNullAsNone: Boolean): PartitionKey = { -+ val builder = new PartitionKeyBuilder() -+ addPartitionKeyComponent(builder, value, treatNullAsNone) -+ builder.build() -+ } - } -diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosPartitionKeyHelper.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosPartitionKeyHelper.scala -new file mode 100644 -index 00000000000..27776f5c3de ---- /dev/null -+++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosPartitionKeyHelper.scala -@@ -0,0 +1,45 @@ -+// Copyright (c) Microsoft Corporation. All rights reserved. -+// Licensed under the MIT License. -+ -+package com.azure.cosmos.spark -+ -+import com.azure.cosmos.implementation.routing.PartitionKeyInternal -+import com.azure.cosmos.implementation.{ImplementationBridgeHelpers, Utils} -+import com.azure.cosmos.models.PartitionKey -+import com.azure.cosmos.spark.diagnostics.BasicLoggingTrait -+ -+import java.util -+ -+// scalastyle:off underscore.import -+import scala.collection.JavaConverters._ -+// scalastyle:on underscore.import -+ -+private[spark] object CosmosPartitionKeyHelper extends BasicLoggingTrait { -+ // pattern will be recognized -+ // pk(partitionKeyValue) -+ // -+ // (?i) : The whole matching is case-insensitive -+ // pk[(](.*)[)]: partitionKey Value -+ private val cosmosPartitionKeyStringRegx = """(?i)pk[(](.*)[)]""".r -+ private val objectMapper = Utils.getSimpleObjectMapper -+ -+ def getCosmosPartitionKeyValueString(partitionKeyValue: List[Object]): String = { -+ s"pk(${objectMapper.writeValueAsString(partitionKeyValue.asJava)})" -+ } -+ -+ def tryParsePartitionKey(cosmosPartitionKeyString: String): Option[PartitionKey] = { -+ cosmosPartitionKeyString match { -+ case cosmosPartitionKeyStringRegx(pkValue) => -+ scala.util.Try(Utils.parse(pkValue, classOf[Object])).toOption.flatMap { -+ case arrayList: util.ArrayList[Object @unchecked] => -+ Some( -+ ImplementationBridgeHelpers -+ .PartitionKeyHelper -+ .getPartitionKeyAccessor -+ .toPartitionKey(PartitionKeyInternal.fromObjectArray(arrayList.toArray, false))) -+ case other => Some(new PartitionKey(other)) -+ } -+ case _ => None -+ } -+ } -+} -diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosReadManyByPartitionKeyReader.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosReadManyByPartitionKeyReader.scala -new file mode 100644 -index 00000000000..91f3a56bc66 ---- /dev/null -+++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosReadManyByPartitionKeyReader.scala -@@ -0,0 +1,150 @@ -+// Copyright (c) Microsoft Corporation. All rights reserved. -+// Licensed under the MIT License. -+package com.azure.cosmos.spark -+ -+import com.azure.cosmos.{CosmosException, ReadConsistencyStrategy} -+import com.azure.cosmos.implementation.{CosmosClientMetadataCachesSnapshot, UUIDs} -+import com.azure.cosmos.models.PartitionKey -+import com.azure.cosmos.spark.CosmosPredicates.assertOnSparkDriver -+import com.azure.cosmos.spark.diagnostics.{BasicLoggingTrait, DiagnosticsContext} -+import com.fasterxml.jackson.databind.node.ObjectNode -+import org.apache.spark.TaskContext -+import org.apache.spark.broadcast.Broadcast -+import org.apache.spark.rdd.RDD -+import org.apache.spark.sql.{DataFrame, Row, SparkSession} -+import org.apache.spark.sql.types.StructType -+ -+import java.util.UUID -+ -+private[spark] class CosmosReadManyByPartitionKeyReader( -+ val userProvidedSchema: StructType, -+ val userConfig: Map[String, String] -+ ) extends BasicLoggingTrait with Serializable { -+ val effectiveUserConfig: Map[String, String] = CosmosConfig.getEffectiveConfig( -+ databaseName = None, -+ containerName = None, -+ userConfig) -+ -+ val clientConfig: CosmosAccountConfig = CosmosAccountConfig.parseCosmosAccountConfig(effectiveUserConfig) -+ val readConfig: CosmosReadConfig = CosmosReadConfig.parseCosmosReadConfig(effectiveUserConfig) -+ val cosmosContainerConfig: CosmosContainerConfig = -+ CosmosContainerConfig.parseCosmosContainerConfig(effectiveUserConfig) -+ //scalastyle:off multiple.string.literals -+ val tableName: String = s"com.azure.cosmos.spark.items.${clientConfig.accountName}." + -+ s"${cosmosContainerConfig.database}.${cosmosContainerConfig.container}" -+ private lazy val sparkSession = { -+ assertOnSparkDriver() -+ SparkSession.active -+ } -+ val sparkEnvironmentInfo: String = CosmosClientConfiguration.getSparkEnvironmentInfo(Some(sparkSession)) -+ logTrace(s"Instantiated ${this.getClass.getSimpleName} for $tableName") -+ -+ private[spark] def initializeAndBroadcastCosmosClientStatesForContainer(): Broadcast[CosmosClientMetadataCachesSnapshots] = { -+ val calledFrom = s"CosmosReadManyByPartitionKeyReader($tableName).initializeAndBroadcastCosmosClientStateForContainer" -+ Loan( -+ List[Option[CosmosClientCacheItem]]( -+ Some( -+ CosmosClientCache( -+ CosmosClientConfiguration( -+ effectiveUserConfig, -+ readConsistencyStrategy = readConfig.readConsistencyStrategy, -+ sparkEnvironmentInfo), -+ None, -+ calledFrom)), -+ ThroughputControlHelper.getThroughputControlClientCacheItem( -+ effectiveUserConfig, -+ calledFrom, -+ None, -+ sparkEnvironmentInfo) -+ )) -+ .to(clientCacheItems => { -+ val container = -+ ThroughputControlHelper.getContainer( -+ effectiveUserConfig, -+ cosmosContainerConfig, -+ clientCacheItems(0).get, -+ clientCacheItems(1)) -+ try { -+ container.readItem( -+ UUIDs.nonBlockingRandomUUID().toString, -+ new PartitionKey(UUIDs.nonBlockingRandomUUID().toString), -+ classOf[ObjectNode]) -+ .block() -+ } catch { -+ case _: CosmosException => None -+ } -+ -+ val state = new CosmosClientMetadataCachesSnapshot() -+ state.serialize(clientCacheItems(0).get.cosmosClient) -+ -+ var throughputControlState: Option[CosmosClientMetadataCachesSnapshot] = None -+ if (clientCacheItems(1).isDefined) { -+ throughputControlState = Some(new CosmosClientMetadataCachesSnapshot()) -+ throughputControlState.get.serialize(clientCacheItems(1).get.cosmosClient) -+ } -+ -+ val metadataSnapshots = CosmosClientMetadataCachesSnapshots(state, throughputControlState) -+ sparkSession.sparkContext.broadcast(metadataSnapshots) -+ }) -+ } -+ -+ def readManyByPartitionKey(inputRdd: RDD[Row], pkExtraction: Row => PartitionKey): DataFrame = { -+ val correlationActivityId = UUIDs.nonBlockingRandomUUID() -+ val calledFrom = s"CosmosReadManyByPartitionKeyReader.readManyByPartitionKey($correlationActivityId)" -+ val schema = Loan( -+ List[Option[CosmosClientCacheItem]]( -+ Some(CosmosClientCache( -+ CosmosClientConfiguration( -+ effectiveUserConfig, -+ readConsistencyStrategy = readConfig.readConsistencyStrategy, -+ sparkEnvironmentInfo), -+ None, -+ calledFrom -+ )), -+ ThroughputControlHelper.getThroughputControlClientCacheItem( -+ effectiveUserConfig, -+ calledFrom, -+ None, -+ sparkEnvironmentInfo) -+ )) -+ .to(clientCacheItems => Option.apply(userProvidedSchema).getOrElse( -+ CosmosTableSchemaInferrer.inferSchema( -+ clientCacheItems(0).get, -+ clientCacheItems(1), -+ effectiveUserConfig, -+ ItemsTable.defaultSchemaForInferenceDisabled))) -+ -+ val clientStates = initializeAndBroadcastCosmosClientStatesForContainer -+ -+ sparkSession.sqlContext.createDataFrame( -+ inputRdd.mapPartitionsWithIndex( -+ (partitionIndex: Int, rowIterator: Iterator[Row]) => { -+ val pkIterator: Iterator[PartitionKey] = rowIterator -+ .map(row => pkExtraction.apply(row)) -+ -+ logInfo(s"Creating an ItemsPartitionReaderWithReadManyByPartitionKey for Activity $correlationActivityId to read for " -+ + s"input partition [$partitionIndex] ${tableName}") -+ -+ val reader = new ItemsPartitionReaderWithReadManyByPartitionKey( -+ effectiveUserConfig, -+ CosmosReadManyHelper.FullRangeFeedRange, -+ schema, -+ DiagnosticsContext(correlationActivityId, partitionIndex.toString), -+ clientStates, -+ DiagnosticsConfig.parseDiagnosticsConfig(effectiveUserConfig), -+ sparkEnvironmentInfo, -+ TaskContext.get, -+ pkIterator) -+ -+ new Iterator[Row] { -+ override def hasNext: Boolean = reader.next() -+ -+ override def next(): Row = reader.getCurrentRow() -+ } -+ }, -+ preservesPartitioning = true -+ ), -+ schema) -+ } -+} -+ -diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/ItemsPartitionReaderWithReadManyByPartitionKey.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/ItemsPartitionReaderWithReadManyByPartitionKey.scala -new file mode 100644 -index 00000000000..c67cc9c10be ---- /dev/null -+++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/ItemsPartitionReaderWithReadManyByPartitionKey.scala -@@ -0,0 +1,249 @@ -+// Copyright (c) Microsoft Corporation. All rights reserved. -+// Licensed under the MIT License. -+ -+package com.azure.cosmos.spark -+ -+import com.azure.cosmos.{CosmosAsyncContainer, CosmosEndToEndOperationLatencyPolicyConfigBuilder, CosmosItemSerializerNoExceptionWrapping, SparkBridgeInternal} -+import com.azure.cosmos.implementation.spark.OperationContextAndListenerTuple -+import com.azure.cosmos.implementation.{ImplementationBridgeHelpers, ObjectNodeMap, SparkRowItem, Utils} -+import com.azure.cosmos.models.{CosmosReadManyRequestOptions, ModelBridgeInternal, PartitionKey, PartitionKeyDefinition, SqlQuerySpec} -+import com.azure.cosmos.spark.BulkWriter.getThreadInfo -+import com.azure.cosmos.spark.CosmosTableSchemaInferrer.IdAttributeName -+import com.azure.cosmos.spark.diagnostics.{DetailedFeedDiagnosticsProvider, DiagnosticsContext, DiagnosticsLoader, LoggerHelper, SparkTaskContext} -+import com.fasterxml.jackson.databind.node.ObjectNode -+import org.apache.spark.TaskContext -+import org.apache.spark.broadcast.Broadcast -+import org.apache.spark.sql.Row -+import org.apache.spark.sql.catalyst.InternalRow -+import org.apache.spark.sql.catalyst.encoders.ExpressionEncoder -+import org.apache.spark.sql.connector.read.PartitionReader -+import org.apache.spark.sql.types.StructType -+ -+import java.util -+ -+// scalastyle:off underscore.import -+import scala.collection.JavaConverters._ -+// scalastyle:on underscore.import -+ -+private[spark] case class ItemsPartitionReaderWithReadManyByPartitionKey -+( -+ config: Map[String, String], -+ feedRange: NormalizedRange, -+ readSchema: StructType, -+ diagnosticsContext: DiagnosticsContext, -+ cosmosClientStateHandles: Broadcast[CosmosClientMetadataCachesSnapshots], -+ diagnosticsConfig: DiagnosticsConfig, -+ sparkEnvironmentInfo: String, -+ taskContext: TaskContext, -+ readManyPartitionKeys: Iterator[PartitionKey] -+) -+ extends PartitionReader[InternalRow] { -+ -+ private lazy val log = LoggerHelper.getLogger(diagnosticsConfig, this.getClass) -+ -+ private val readManyOptions = new CosmosReadManyRequestOptions() -+ private val readManyOptionsImpl = ImplementationBridgeHelpers -+ .CosmosReadManyRequestOptionsHelper -+ .getCosmosReadManyRequestOptionsAccessor -+ .getImpl(readManyOptions) -+ -+ private val readConfig = CosmosReadConfig.parseCosmosReadConfig(config) -+ ThroughputControlHelper.populateThroughputControlGroupName(readManyOptionsImpl, readConfig.throughputControlConfig) -+ -+ private val operationContext = { -+ assert(taskContext != null) -+ -+ SparkTaskContext(diagnosticsContext.correlationActivityId, -+ taskContext.stageId(), -+ taskContext.partitionId(), -+ taskContext.taskAttemptId(), -+ feedRange.toString) -+ } -+ -+ private val operationContextAndListenerTuple: Option[OperationContextAndListenerTuple] = { -+ if (diagnosticsConfig.mode.isDefined) { -+ val listener = -+ DiagnosticsLoader.getDiagnosticsProvider(diagnosticsConfig).getLogger(this.getClass) -+ -+ val ctxAndListener = new OperationContextAndListenerTuple(operationContext, listener) -+ -+ readManyOptionsImpl -+ .setOperationContextAndListenerTuple(ctxAndListener) -+ -+ Some(ctxAndListener) -+ } else { -+ None -+ } -+ } -+ -+ log.logTrace(s"Instantiated ${this.getClass.getSimpleName}, Context: ${operationContext.toString} $getThreadInfo") -+ -+ private val containerTargetConfig = CosmosContainerConfig.parseCosmosContainerConfig(config) -+ -+ log.logInfo(s"Using ReadManyByPartitionKey from feed range $feedRange of " + -+ s"container ${containerTargetConfig.database}.${containerTargetConfig.container} - " + -+ s"correlationActivityId ${diagnosticsContext.correlationActivityId}, " + -+ s"Context: ${operationContext.toString} $getThreadInfo") -+ -+ private val clientCacheItem = CosmosClientCache( -+ CosmosClientConfiguration(config, readConfig.readConsistencyStrategy, sparkEnvironmentInfo), -+ Some(cosmosClientStateHandles.value.cosmosClientMetadataCaches), -+ s"ItemsPartitionReaderWithReadManyByPartitionKey($feedRange, ${containerTargetConfig.database}.${containerTargetConfig.container})" -+ ) -+ -+ private val throughputControlClientCacheItemOpt = -+ ThroughputControlHelper.getThroughputControlClientCacheItem( -+ config, -+ clientCacheItem.context, -+ Some(cosmosClientStateHandles), -+ sparkEnvironmentInfo) -+ -+ private val cosmosAsyncContainer = -+ ThroughputControlHelper.getContainer( -+ config, -+ containerTargetConfig, -+ clientCacheItem, -+ throughputControlClientCacheItemOpt) -+ -+ private val partitionKeyDefinition: PartitionKeyDefinition = { -+ TransientErrorsRetryPolicy.executeWithRetry(() => { -+ SparkBridgeInternal -+ .getContainerPropertiesFromCollectionCache(cosmosAsyncContainer).getPartitionKeyDefinition -+ }) -+ } -+ -+ private val cosmosSerializationConfig = CosmosSerializationConfig.parseSerializationConfig(config) -+ private val cosmosRowConverter = CosmosRowConverter.get(cosmosSerializationConfig) -+ -+ readManyOptionsImpl -+ .setCustomItemSerializer( -+ new CosmosItemSerializerNoExceptionWrapping { -+ override def serialize[T](item: T): util.Map[String, AnyRef] = { -+ throw new UnsupportedOperationException( -+ s"Serialization is not supported by the custom item serializer in " + -+ s"ItemsPartitionReaderWithReadManyByPartitionKey; this serializer is intended " + -+ s"for deserializing read-many responses into SparkRowItem only. " + -+ s"Unexpected item type: ${if (item == null) "null" else item.getClass.getName}" -+ ) -+ } -+ -+ override def deserialize[T](jsonNodeMap: util.Map[String, AnyRef], classType: Class[T]): T = { -+ if (jsonNodeMap == null) { -+ throw new IllegalStateException("The 'jsonNodeMap' should never be null here.") -+ } -+ -+ if (classType != classOf[SparkRowItem]) { -+ throw new IllegalStateException("The 'classType' must be 'classOf[SparkRowItem])' here.") -+ } -+ -+ val objectNode: ObjectNode = jsonNodeMap match { -+ case map: ObjectNodeMap => -+ map.getObjectNode -+ case _ => -+ Utils.getSimpleObjectMapper.convertValue(jsonNodeMap, classOf[ObjectNode]) -+ } -+ -+ val partitionKey = PartitionKeyHelper.getPartitionKeyPath(objectNode, partitionKeyDefinition) -+ -+ val row = cosmosRowConverter.fromObjectNodeToRow(readSchema, -+ objectNode, -+ readConfig.schemaConversionMode) -+ -+ SparkRowItem(row, getPartitionKeyForFeedDiagnostics(partitionKey)).asInstanceOf[T] -+ } -+ } -+ ) -+ -+ // Collect all PK values upfront ΓÇö readManyByPartitionKey needs the full list to -+ // group by physical partition and issue parallel queries. -+ // Deduplicate by PK string representation ΓÇö safe because the list size is bounded -+ // by the per-call limit of the readManyByPartitionKey API. -+ private lazy val pkList = { -+ val seen = new java.util.LinkedHashMap[String, PartitionKey]() -+ readManyPartitionKeys.foreach(pk => seen.putIfAbsent(pk.toString, pk)) -+ new java.util.ArrayList[PartitionKey](seen.values()) -+ } -+ -+ private val endToEndTimeoutPolicy = -+ new CosmosEndToEndOperationLatencyPolicyConfigBuilder( -+ java.time.Duration.ofSeconds(CosmosConstants.readOperationEndToEndTimeoutInSeconds)) -+ .enable(true) -+ .build -+ -+ readManyOptionsImpl.setCosmosEndToEndOperationLatencyPolicyConfig(endToEndTimeoutPolicy) -+ -+ private trait CloseableSparkRowItemIterator { -+ def hasNext: Boolean -+ def next(): SparkRowItem -+ def close(): Unit -+ } -+ -+ private object EmptySparkRowItemIterator extends CloseableSparkRowItemIterator { -+ override def hasNext: Boolean = false -+ -+ override def next(): SparkRowItem = { -+ throw new java.util.NoSuchElementException("No items available for empty partition-key list.") -+ } -+ -+ override def close(): Unit = {} -+ } -+ -+ // Batch partition keys and retry each batch independently on transient I/O errors. -+ // This avoids the continuation-token problem with TransientIOErrorsRetryingIterator -+ // where a retry would re-read all data from scratch, causing silent data duplication. -+ private lazy val iterator: CloseableSparkRowItemIterator = -+ if (pkList.isEmpty) { -+ EmptySparkRowItemIterator -+ } else { -+ new CloseableSparkRowItemIterator { -+ private val delegate = new TransientIOErrorsRetryingReadManyByPartitionKeyIterator[SparkRowItem]( -+ cosmosAsyncContainer, -+ pkList, -+ readConfig.customQuery.map(_.toSqlQuerySpec), -+ readManyOptions, -+ readConfig.maxItemCount, -+ readConfig.prefetchBufferSize, -+ operationContextAndListenerTuple, -+ classOf[SparkRowItem] -+ ) -+ -+ override def hasNext: Boolean = delegate.hasNext -+ -+ override def next(): SparkRowItem = delegate.next() -+ -+ override def close(): Unit = delegate.close() -+ } -+ } -+ -+ private val rowSerializer: ExpressionEncoder.Serializer[Row] = RowSerializerPool.getOrCreateSerializer(readSchema) -+ -+ private def shouldLogDetailedFeedDiagnostics(): Boolean = { -+ diagnosticsConfig.mode.isDefined && -+ diagnosticsConfig.mode.get.equalsIgnoreCase(classOf[DetailedFeedDiagnosticsProvider].getName) -+ } -+ -+ private def getPartitionKeyForFeedDiagnostics(pkValue: PartitionKey): Option[PartitionKey] = { -+ if (shouldLogDetailedFeedDiagnostics()) { -+ Some(pkValue) -+ } else { -+ None -+ } -+ } -+ -+ override def next(): Boolean = iterator.hasNext -+ -+ override def get(): InternalRow = { -+ cosmosRowConverter.fromRowToInternalRow(iterator.next().row, rowSerializer) -+ } -+ -+ def getCurrentRow(): Row = iterator.next().row -+ -+ override def close(): Unit = { -+ this.iterator.close() -+ RowSerializerPool.returnSerializerToPool(readSchema, rowSerializer) -+ clientCacheItem.close() -+ if (throughputControlClientCacheItemOpt.isDefined) { -+ throughputControlClientCacheItemOpt.get.close() -+ } -+ } -+} -diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/TransientIOErrorsRetryingReadManyByPartitionKeyIterator.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/TransientIOErrorsRetryingReadManyByPartitionKeyIterator.scala -new file mode 100644 -index 00000000000..dcfdf4f9353 ---- /dev/null -+++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/TransientIOErrorsRetryingReadManyByPartitionKeyIterator.scala -@@ -0,0 +1,259 @@ -+// Copyright (c) Microsoft Corporation. All rights reserved. -+// Licensed under the MIT License. -+ -+package com.azure.cosmos.spark -+ -+import com.azure.cosmos.{CosmosAsyncContainer, CosmosException} -+import com.azure.cosmos.implementation.OperationCancelledException -+import com.azure.cosmos.implementation.spark.OperationContextAndListenerTuple -+import com.azure.cosmos.models.{CosmosReadManyRequestOptions, FeedResponse, PartitionKey, SqlQuerySpec} -+import com.azure.cosmos.spark.diagnostics.BasicLoggingTrait -+import com.azure.cosmos.util.CosmosPagedIterable -+ -+import java.util.concurrent.{ExecutorService, SynchronousQueue, ThreadPoolExecutor, TimeUnit, TimeoutException} -+import java.util.concurrent.atomic.AtomicLong -+import scala.concurrent.{Await, ExecutionContext, Future} -+import scala.util.Random -+import scala.util.control.Breaks -+ -+// scalastyle:off underscore.import -+import scala.collection.JavaConverters._ -+// scalastyle:on underscore.import -+ -+/** -+ * Retry-safe iterator for readManyByPartitionKey that batches partition keys and lazily -+ * iterates pages within each batch via CosmosPagedIterable ΓÇö consistent with how -+ * TransientIOErrorsRetryingIterator handles normal queries. On transient I/O errors the -+ * current batch's flux is recreated and pages already consumed are replayed, avoiding -+ * the memory overhead of collectList and matching the query iterator's structure. -+ */ -+private[spark] class TransientIOErrorsRetryingReadManyByPartitionKeyIterator[TSparkRow] -+( -+ val container: CosmosAsyncContainer, -+ val partitionKeys: java.util.List[PartitionKey], -+ val customQuery: Option[SqlQuerySpec], -+ val queryOptions: CosmosReadManyRequestOptions, -+ val pageSize: Int, -+ val pagePrefetchBufferSize: Int, -+ val operationContextAndListener: Option[OperationContextAndListenerTuple], -+ val classType: Class[TSparkRow] -+) extends BufferedIterator[TSparkRow] with BasicLoggingTrait with AutoCloseable { -+ -+ private[spark] var maxRetryIntervalInMs = CosmosConstants.maxRetryIntervalForTransientFailuresInMs -+ private[spark] var maxRetryCount = CosmosConstants.maxRetryCountForTransientFailures -+ -+ private val maxPageRetrievalTimeout = scala.concurrent.duration.FiniteDuration( -+ 5 + CosmosConstants.readOperationEndToEndTimeoutInSeconds, -+ scala.concurrent.duration.SECONDS) -+ -+ private val rnd = Random -+ private val retryCount = new AtomicLong(0) -+ private lazy val operationContextString = operationContextAndListener match { -+ case Some(o) => if (o.getOperationContext != null) { -+ o.getOperationContext.toString -+ } else { -+ "n/a" -+ } -+ case None => "n/a" -+ } -+ -+ private[spark] var currentFeedResponseIterator: Option[BufferedIterator[FeedResponse[TSparkRow]]] = None -+ private[spark] var currentItemIterator: Option[BufferedIterator[TSparkRow]] = None -+ -+ private val pkBatchIterator = partitionKeys.asScala.iterator.grouped(pageSize) -+ // Track the current batch so we can replay it on retry -+ private var currentBatch: Option[java.util.List[PartitionKey]] = None -+ -+ override def hasNext: Boolean = { -+ executeWithRetry("hasNextInternal", () => hasNextInternal) -+ } -+ -+ private def hasNextInternal: Boolean = { -+ var returnValue: Option[Boolean] = None -+ -+ while (returnValue.isEmpty) { -+ returnValue = hasNextInternalCore -+ } -+ -+ returnValue.get -+ } -+ -+ private def hasNextInternalCore: Option[Boolean] = { -+ if (hasBufferedNext) { -+ Some(true) -+ } else { -+ val feedResponseIterator = currentFeedResponseIterator match { -+ case Some(existing) => existing -+ case None => -+ // Need a new feed response iterator ΓÇö either for the current batch (on retry) -+ // or for the next batch -+ val batch = currentBatch match { -+ case Some(b) => b // retry of current batch -+ case None => -+ if (pkBatchIterator.hasNext) { -+ val nextBatch = new java.util.ArrayList[PartitionKey](pkBatchIterator.next().toList.asJava) -+ currentBatch = Some(nextBatch) -+ nextBatch -+ } else { -+ return Some(false) // no more batches -+ } -+ } -+ -+ val pagedFlux = customQuery match { -+ case Some(query) => -+ container.readManyByPartitionKey(batch, query, queryOptions, classType) -+ case None => -+ container.readManyByPartitionKey(batch, queryOptions, classType) -+ } -+ -+ currentFeedResponseIterator = Some( -+ new CosmosPagedIterable[TSparkRow]( -+ pagedFlux, -+ pageSize, -+ pagePrefetchBufferSize -+ ) -+ .iterableByPage() -+ .iterator -+ .asScala -+ .buffered -+ ) -+ -+ currentFeedResponseIterator.get -+ } -+ -+ val hasNext: Boolean = try { -+ Await.result( -+ Future { -+ feedResponseIterator.hasNext -+ }(TransientIOErrorsRetryingReadManyByPartitionKeyIterator.executionContext), -+ maxPageRetrievalTimeout) -+ } catch { -+ case endToEndTimeoutException: OperationCancelledException => -+ val message = s"End-to-end timeout hit when trying to retrieve the next page. " + -+ s"Context: $operationContextString" -+ logError(message, throwable = endToEndTimeoutException) -+ throw endToEndTimeoutException -+ -+ case timeoutException: TimeoutException => -+ val message = s"Attempting to retrieve the next page timed out. " + -+ s"Context: $operationContextString" -+ logError(message, timeoutException) -+ val exception = new OperationCancelledException(message, null) -+ exception.setStackTrace(timeoutException.getStackTrace) -+ throw exception -+ -+ case other: Throwable => throw other -+ } -+ -+ if (hasNext) { -+ val feedResponse = feedResponseIterator.next() -+ if (operationContextAndListener.isDefined) { -+ operationContextAndListener.get.getOperationListener.feedResponseProcessedListener( -+ operationContextAndListener.get.getOperationContext, -+ feedResponse) -+ } -+ val iteratorCandidate = feedResponse.getResults.iterator().asScala.buffered -+ -+ if (iteratorCandidate.hasNext) { -+ currentItemIterator = Some(iteratorCandidate) -+ Some(true) -+ } else { -+ // empty page interleaved ΓÇö try again -+ None -+ } -+ } else { -+ // Current batch's flux is exhausted ΓÇö move to next batch -+ currentBatch = None -+ currentFeedResponseIterator = None -+ None -+ } -+ } -+ } -+ -+ private def hasBufferedNext: Boolean = { -+ currentItemIterator match { -+ case Some(iterator) => if (iterator.hasNext) { -+ true -+ } else { -+ currentItemIterator = None -+ false -+ } -+ case None => false -+ } -+ } -+ -+ override def next(): TSparkRow = { -+ currentItemIterator.get.next() -+ } -+ -+ override def head(): TSparkRow = { -+ currentItemIterator.get.head -+ } -+ -+ private[spark] def executeWithRetry[T](methodName: String, func: () => T): T = { -+ val loop = new Breaks() -+ var returnValue: Option[T] = None -+ -+ loop.breakable { -+ while (true) { -+ val retryIntervalInMs = rnd.nextInt(maxRetryIntervalInMs) -+ -+ try { -+ returnValue = Some(func()) -+ retryCount.set(0) -+ loop.break -+ } -+ catch { -+ case cosmosException: CosmosException => -+ if (Exceptions.canBeTransientFailure(cosmosException.getStatusCode, cosmosException.getSubStatusCode)) { -+ val retryCountSnapshot = retryCount.incrementAndGet() -+ if (retryCountSnapshot > maxRetryCount) { -+ logError( -+ s"Too many transient failure retry attempts in " + -+ s"TransientIOErrorsRetryingReadManyByPartitionKeyIterator.$methodName", -+ cosmosException) -+ throw cosmosException -+ } else { -+ logWarning( -+ s"Transient failure handled in " + -+ s"TransientIOErrorsRetryingReadManyByPartitionKeyIterator.$methodName -" + -+ s" will be retried (attempt#$retryCountSnapshot) in ${retryIntervalInMs}ms", -+ cosmosException) -+ } -+ } else { -+ throw cosmosException -+ } -+ case other: Throwable => throw other -+ } -+ -+ // Reset iterators but keep currentBatch so the batch is replayed -+ currentItemIterator = None -+ currentFeedResponseIterator = None -+ Thread.sleep(retryIntervalInMs) -+ } -+ } -+ -+ returnValue.get -+ } -+ -+ override def close(): Unit = { -+ currentItemIterator = None -+ currentFeedResponseIterator = None -+ } -+} -+ -+private object TransientIOErrorsRetryingReadManyByPartitionKeyIterator extends BasicLoggingTrait { -+ private val maxConcurrency = SparkUtils.getNumberOfHostCPUCores -+ -+ val executorService: ExecutorService = new ThreadPoolExecutor( -+ maxConcurrency, -+ maxConcurrency, -+ 0L, -+ TimeUnit.MILLISECONDS, -+ new SynchronousQueue(), -+ SparkUtils.daemonThreadFactory(), -+ new ThreadPoolExecutor.CallerRunsPolicy() -+ ) -+ -+ val executionContext: ExecutionContext = ExecutionContext.fromExecutorService(executorService) -+} -diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/udf/GetCosmosPartitionKeyValue.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/udf/GetCosmosPartitionKeyValue.scala -new file mode 100644 -index 00000000000..a58d5b723b8 ---- /dev/null -+++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/udf/GetCosmosPartitionKeyValue.scala -@@ -0,0 +1,25 @@ -+// Copyright (c) Microsoft Corporation. All rights reserved. -+// Licensed under the MIT License. -+ -+package com.azure.cosmos.spark.udf -+ -+import com.azure.cosmos.spark.CosmosPartitionKeyHelper -+import com.azure.cosmos.spark.CosmosPredicates.requireNotNull -+import org.apache.spark.sql.api.java.UDF1 -+ -+@SerialVersionUID(1L) -+class GetCosmosPartitionKeyValue extends UDF1[Object, String] { -+ override def call -+ ( -+ partitionKeyValue: Object -+ ): String = { -+ requireNotNull(partitionKeyValue, "partitionKeyValue") -+ -+ partitionKeyValue match { -+ // for subpartitions case - Seq covers both WrappedArray (Scala 2.12) and ArraySeq (Scala 2.13) -+ case seq: Seq[Any] => -+ CosmosPartitionKeyHelper.getCosmosPartitionKeyValueString(seq.map(_.asInstanceOf[Object]).toList) -+ case _ => CosmosPartitionKeyHelper.getCosmosPartitionKeyValueString(List(partitionKeyValue)) -+ } -+ } -+} -diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/CosmosConfigSpec.scala b/sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/CosmosConfigSpec.scala -index 17f75e45a74..17a298d6213 100644 ---- a/sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/CosmosConfigSpec.scala -+++ b/sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/CosmosConfigSpec.scala -@@ -457,6 +457,7 @@ class CosmosConfigSpec extends UnitSpec with BasicLoggingTrait { - config.runtimeFilteringEnabled shouldBe true - config.readManyFilteringConfig.readManyFilteringEnabled shouldBe false - config.readManyFilteringConfig.readManyFilterProperty shouldEqual "_itemIdentity" -+ config.readManyByPkTreatNullAsNone shouldBe false - - userConfig = Map( - "spark.cosmos.read.forceEventualConsistency" -> "false", -@@ -630,6 +631,47 @@ class CosmosConfigSpec extends UnitSpec with BasicLoggingTrait { - config.customQuery.get.queryText shouldBe queryText - } - -+ it should "parse readManyByPk nullHandling configuration" in { -+ // Default (not specified) should treat null as JSON null (addNullValue) -+ var userConfig = Map( -+ "spark.cosmos.read.forceEventualConsistency" -> "false" -+ ) -+ var config = CosmosReadConfig.parseCosmosReadConfig(userConfig) -+ config.readManyByPkTreatNullAsNone shouldBe false -+ -+ // Explicit "Null" should treat null as JSON null (addNullValue) -+ userConfig = Map( -+ "spark.cosmos.read.forceEventualConsistency" -> "false", -+ "spark.cosmos.read.readManyByPk.nullHandling" -> "Null" -+ ) -+ config = CosmosReadConfig.parseCosmosReadConfig(userConfig) -+ config.readManyByPkTreatNullAsNone shouldBe false -+ -+ // Case-insensitive "null" -+ userConfig = Map( -+ "spark.cosmos.read.forceEventualConsistency" -> "false", -+ "spark.cosmos.read.readManyByPk.nullHandling" -> "null" -+ ) -+ config = CosmosReadConfig.parseCosmosReadConfig(userConfig) -+ config.readManyByPkTreatNullAsNone shouldBe false -+ -+ // "None" should treat null as PartitionKey.NONE (addNoneValue) -+ userConfig = Map( -+ "spark.cosmos.read.forceEventualConsistency" -> "false", -+ "spark.cosmos.read.readManyByPk.nullHandling" -> "None" -+ ) -+ config = CosmosReadConfig.parseCosmosReadConfig(userConfig) -+ config.readManyByPkTreatNullAsNone shouldBe true -+ -+ // Case-insensitive "none" -+ userConfig = Map( -+ "spark.cosmos.read.forceEventualConsistency" -> "false", -+ "spark.cosmos.read.readManyByPk.nullHandling" -> "none" -+ ) -+ config = CosmosReadConfig.parseCosmosReadConfig(userConfig) -+ config.readManyByPkTreatNullAsNone shouldBe true -+ } -+ - it should "throw on invalid read configuration" in { - val userConfig = Map( - "spark.cosmos.read.schemaConversionMode" -> "not a valid value" -diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/CosmosPartitionKeyHelperSpec.scala b/sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/CosmosPartitionKeyHelperSpec.scala -new file mode 100644 -index 00000000000..1ac40e39584 ---- /dev/null -+++ b/sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/CosmosPartitionKeyHelperSpec.scala -@@ -0,0 +1,104 @@ -+// Copyright (c) Microsoft Corporation. All rights reserved. -+// Licensed under the MIT License. -+ -+package com.azure.cosmos.spark -+ -+import com.azure.cosmos.models.{PartitionKey, PartitionKeyBuilder} -+ -+class CosmosPartitionKeyHelperSpec extends UnitSpec { -+ //scalastyle:off multiple.string.literals -+ //scalastyle:off magic.number -+ -+ it should "return the correct partition key value string for single PK" in { -+ val pkString = CosmosPartitionKeyHelper.getCosmosPartitionKeyValueString(List("pk1")) -+ pkString shouldEqual "pk([\"pk1\"])" -+ } -+ -+ it should "return the correct partition key value string for HPK" in { -+ val pkString = CosmosPartitionKeyHelper.getCosmosPartitionKeyValueString(List("city1", "zip1")) -+ pkString shouldEqual "pk([\"city1\",\"zip1\"])" -+ } -+ -+ it should "return the correct partition key value string for 3-level HPK" in { -+ val pkString = CosmosPartitionKeyHelper.getCosmosPartitionKeyValueString(List("a", "b", "c")) -+ pkString shouldEqual "pk([\"a\",\"b\",\"c\"])" -+ } -+ -+ it should "parse valid single PK string" in { -+ val pk = CosmosPartitionKeyHelper.tryParsePartitionKey("pk([\"myPkValue\"])") -+ pk.isDefined shouldBe true -+ pk.get shouldEqual new PartitionKey("myPkValue") -+ } -+ -+ it should "parse valid HPK string" in { -+ val pk = CosmosPartitionKeyHelper.tryParsePartitionKey("pk([\"city1\",\"zip1\"])") -+ pk.isDefined shouldBe true -+ val expected = new PartitionKeyBuilder().add("city1").add("zip1").build() -+ pk.get shouldEqual expected -+ } -+ -+ it should "parse valid 3-level HPK string" in { -+ val pk = CosmosPartitionKeyHelper.tryParsePartitionKey("pk([\"a\",\"b\",\"c\"])") -+ pk.isDefined shouldBe true -+ val expected = new PartitionKeyBuilder().add("a").add("b").add("c").build() -+ pk.get shouldEqual expected -+ } -+ -+ it should "roundtrip single PK" in { -+ val original = "pk([\"roundtrip\"])" -+ val parsed = CosmosPartitionKeyHelper.tryParsePartitionKey(original) -+ parsed.isDefined shouldBe true -+ val serialized = CosmosPartitionKeyHelper.getCosmosPartitionKeyValueString(List("roundtrip")) -+ serialized shouldEqual original -+ } -+ -+ it should "roundtrip HPK" in { -+ val original = "pk([\"city\",\"zip\"])" -+ val parsed = CosmosPartitionKeyHelper.tryParsePartitionKey(original) -+ parsed.isDefined shouldBe true -+ val serialized = CosmosPartitionKeyHelper.getCosmosPartitionKeyValueString(List("city", "zip")) -+ serialized shouldEqual original -+ } -+ -+ it should "return None for malformed string" in { -+ val pk = CosmosPartitionKeyHelper.tryParsePartitionKey("invalid_format") -+ pk.isDefined shouldBe false -+ } -+ -+ it should "return None for missing pk prefix" in { -+ val pk = CosmosPartitionKeyHelper.tryParsePartitionKey("[\"value\"]") -+ pk.isDefined shouldBe false -+ } -+ -+ it should "be case-insensitive for parsing" in { -+ val pk = CosmosPartitionKeyHelper.tryParsePartitionKey("PK([\"value\"])") -+ pk.isDefined shouldBe true -+ pk.get shouldEqual new PartitionKey("value") -+ } -+ -+ -+ it should "return None for malformed JSON inside pk() wrapper" in { -+ // Invalid JSON that would cause JsonProcessingException -+ val pk = CosmosPartitionKeyHelper.tryParsePartitionKey("pk({invalid json})") -+ pk.isDefined shouldBe false -+ } -+ -+ it should "return None for truncated JSON inside pk() wrapper" in { -+ val pk = CosmosPartitionKeyHelper.tryParsePartitionKey("pk([\"unterminated)") -+ pk.isDefined shouldBe false -+ } -+ -+ it should "produce different partition keys for addNullValue vs addNoneValue in HPK" in { -+ // addNullValue represents an explicit JSON null for a field that exists with value null -+ val pkWithNull = new PartitionKeyBuilder().add("Redmond").addNullValue().build() -+ -+ // addNoneValue represents PartitionKey.NONE, meaning the field is absent/undefined -+ val pkWithNone = new PartitionKeyBuilder().add("Redmond").addNoneValue().build() -+ -+ // These MUST produce different partition key hashes and route to different physical partitions -+ pkWithNull should not equal pkWithNone -+ } -+ -+ //scalastyle:on multiple.string.literals -+ //scalastyle:on magic.number -+} -diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/ItemsPartitionReaderWithReadManyByPartitionKeyITest.scala b/sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/ItemsPartitionReaderWithReadManyByPartitionKeyITest.scala -new file mode 100644 -index 00000000000..5c2d7b59836 ---- /dev/null -+++ b/sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/ItemsPartitionReaderWithReadManyByPartitionKeyITest.scala -@@ -0,0 +1,158 @@ -+// Copyright (c) Microsoft Corporation. All rights reserved. -+// Licensed under the MIT License. -+ -+package com.azure.cosmos.spark -+ -+import com.azure.cosmos.implementation.{CosmosClientMetadataCachesSnapshot, TestConfigurations, Utils} -+import com.azure.cosmos.models.PartitionKey -+import com.azure.cosmos.spark.diagnostics.DiagnosticsContext -+import com.fasterxml.jackson.databind.node.ObjectNode -+import org.apache.spark.MockTaskContext -+import org.apache.spark.broadcast.Broadcast -+import org.apache.spark.sql.types.{StringType, StructField, StructType} -+ -+import java.util.UUID -+import scala.collection.mutable.ListBuffer -+ -+class ItemsPartitionReaderWithReadManyByPartitionKeyITest -+ extends IntegrationSpec -+ with Spark -+ with AutoCleanableCosmosContainersWithPkAsPartitionKey { -+ private val idProperty = "id" -+ private val pkProperty = "pk" -+ -+ //scalastyle:off multiple.string.literals -+ //scalastyle:off magic.number -+ -+ it should "be able to retrieve all items for given partition keys" in { -+ val container = cosmosClient.getDatabase(cosmosDatabase).getContainer(cosmosContainersWithPkAsPartitionKey) -+ -+ // Create items with known PK values -+ val partitionKeyDefinition = container.read().block().getProperties.getPartitionKeyDefinition -+ val allItemsByPk = scala.collection.mutable.Map[String, ListBuffer[ObjectNode]]() -+ val pkValues = List("pkA", "pkB", "pkC") -+ -+ for (pk <- pkValues) { -+ allItemsByPk(pk) = ListBuffer[ObjectNode]() -+ for (_ <- 1 to 5) { -+ val objectNode = Utils.getSimpleObjectMapper.createObjectNode() -+ objectNode.put(idProperty, UUID.randomUUID().toString) -+ objectNode.put(pkProperty, pk) -+ container.createItem(objectNode).block() -+ allItemsByPk(pk) += objectNode -+ } -+ } -+ -+ val config = Map( -+ "spark.cosmos.accountEndpoint" -> TestConfigurations.HOST, -+ "spark.cosmos.accountKey" -> TestConfigurations.MASTER_KEY, -+ "spark.cosmos.database" -> cosmosDatabase, -+ "spark.cosmos.container" -> cosmosContainersWithPkAsPartitionKey, -+ "spark.cosmos.read.inferSchema.enabled" -> "true", -+ "spark.cosmos.applicationName" -> "ReadManyByPKTest" -+ ) -+ -+ val readSchema = StructType(Seq( -+ StructField(idProperty, StringType, false), -+ StructField(pkProperty, StringType, false) -+ )) -+ -+ val diagnosticsContext = DiagnosticsContext(UUID.randomUUID(), "") -+ val diagnosticsConfig = DiagnosticsConfig.parseDiagnosticsConfig(config) -+ val cosmosClientMetadataCachesSnapshots = getCosmosClientMetadataCachesSnapshots() -+ -+ // Read items for pkA and pkB (not pkC) -+ val targetPks = List("pkA", "pkB") -+ val pkIterator = targetPks.map(pk => new PartitionKey(pk)).iterator -+ -+ val reader = ItemsPartitionReaderWithReadManyByPartitionKey( -+ config, -+ NormalizedRange("", "FF"), -+ readSchema, -+ diagnosticsContext, -+ cosmosClientMetadataCachesSnapshots, -+ diagnosticsConfig, -+ "", -+ MockTaskContext.mockTaskContext(), -+ pkIterator -+ ) -+ -+ val cosmosRowConverter = CosmosRowConverter.get(CosmosSerializationConfig.parseSerializationConfig(config)) -+ val itemsReadFromReader = ListBuffer[ObjectNode]() -+ while (reader.next()) { -+ itemsReadFromReader += cosmosRowConverter.fromInternalRowToObjectNode(reader.get(), readSchema) -+ } -+ -+ // Should have 10 items (5 for pkA + 5 for pkB) -+ itemsReadFromReader.size shouldEqual 10 -+ -+ // All items should be from pkA or pkB -+ itemsReadFromReader.foreach(item => { -+ val pk = item.get(pkProperty).asText() -+ targetPks should contain(pk) -+ }) -+ -+ // Validate all expected IDs are present -+ val expectedIds = (allItemsByPk("pkA") ++ allItemsByPk("pkB")).map(_.get(idProperty).asText()).toSet -+ val actualIds = itemsReadFromReader.map(_.get(idProperty).asText()).toSet -+ actualIds shouldEqual expectedIds -+ -+ reader.close() -+ } -+ -+ it should "return empty results for non-existent partition keys" in { -+ val config = Map( -+ "spark.cosmos.accountEndpoint" -> TestConfigurations.HOST, -+ "spark.cosmos.accountKey" -> TestConfigurations.MASTER_KEY, -+ "spark.cosmos.database" -> cosmosDatabase, -+ "spark.cosmos.container" -> cosmosContainersWithPkAsPartitionKey, -+ "spark.cosmos.read.inferSchema.enabled" -> "true", -+ "spark.cosmos.applicationName" -> "ReadManyByPKEmptyTest" -+ ) -+ -+ val readSchema = StructType(Seq( -+ StructField(idProperty, StringType, false), -+ StructField(pkProperty, StringType, false) -+ )) -+ -+ val diagnosticsContext = DiagnosticsContext(UUID.randomUUID(), "") -+ val diagnosticsConfig = DiagnosticsConfig.parseDiagnosticsConfig(config) -+ val cosmosClientMetadataCachesSnapshots = getCosmosClientMetadataCachesSnapshots() -+ -+ val pkIterator = List(new PartitionKey("nonExistentPk")).iterator -+ -+ val reader = ItemsPartitionReaderWithReadManyByPartitionKey( -+ config, -+ NormalizedRange("", "FF"), -+ readSchema, -+ diagnosticsContext, -+ cosmosClientMetadataCachesSnapshots, -+ diagnosticsConfig, -+ "", -+ MockTaskContext.mockTaskContext(), -+ pkIterator -+ ) -+ -+ val itemsReadFromReader = ListBuffer[ObjectNode]() -+ val cosmosRowConverter = CosmosRowConverter.get(CosmosSerializationConfig.parseSerializationConfig(config)) -+ while (reader.next()) { -+ itemsReadFromReader += cosmosRowConverter.fromInternalRowToObjectNode(reader.get(), readSchema) -+ } -+ -+ itemsReadFromReader.size shouldEqual 0 -+ reader.close() -+ } -+ -+ private def getCosmosClientMetadataCachesSnapshots(): Broadcast[CosmosClientMetadataCachesSnapshots] = { -+ val cosmosClientMetadataCachesSnapshot = new CosmosClientMetadataCachesSnapshot() -+ cosmosClientMetadataCachesSnapshot.serialize(cosmosClient) -+ -+ spark.sparkContext.broadcast( -+ CosmosClientMetadataCachesSnapshots( -+ cosmosClientMetadataCachesSnapshot, -+ Option.empty[CosmosClientMetadataCachesSnapshot])) -+ } -+ -+ //scalastyle:on multiple.string.literals -+ //scalastyle:on magic.number -+} -diff --git a/sdk/cosmos/azure-cosmos-spark_4-0_2-13/CHANGELOG.md b/sdk/cosmos/azure-cosmos-spark_4-0_2-13/CHANGELOG.md -index 3972ae6aeb9..d8368be6a0d 100644 ---- a/sdk/cosmos/azure-cosmos-spark_4-0_2-13/CHANGELOG.md -+++ b/sdk/cosmos/azure-cosmos-spark_4-0_2-13/CHANGELOG.md -@@ -3,6 +3,7 @@ - ### 4.47.0-beta.1 (Unreleased) - - #### Features Added -+* Added new `CosmosItemsDataSource.readManyByPartitionKey` Spark function to execute bulk queries by a list of pk-values with better efficiency. See [PR 48801](https://github.com/Azure/azure-sdk-for-java/pull/48801) - - #### Breaking Changes - -diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/ReadManyByPartitionKeyTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/ReadManyByPartitionKeyTest.java -new file mode 100644 -index 00000000000..2c26d564ed2 ---- /dev/null -+++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/ReadManyByPartitionKeyTest.java -@@ -0,0 +1,462 @@ -+/* -+ * Copyright (c) Microsoft Corporation. All rights reserved. -+ * Licensed under the MIT License. -+ */ -+ -+package com.azure.cosmos; -+ -+import com.azure.cosmos.models.CosmosContainerProperties; -+import com.azure.cosmos.models.CosmosItemRequestOptions; -+import com.azure.cosmos.models.FeedResponse; -+import com.azure.cosmos.models.PartitionKey; -+import com.azure.cosmos.models.PartitionKeyBuilder; -+import com.azure.cosmos.models.PartitionKeyDefinition; -+import com.azure.cosmos.models.PartitionKeyDefinitionVersion; -+import com.azure.cosmos.models.PartitionKind; -+import com.azure.cosmos.models.SqlParameter; -+import com.azure.cosmos.models.SqlQuerySpec; -+import com.azure.cosmos.rx.TestSuiteBase; -+import com.azure.cosmos.util.CosmosPagedIterable; -+import com.fasterxml.jackson.databind.node.ObjectNode; -+import org.testng.annotations.AfterClass; -+import org.testng.annotations.BeforeClass; -+import org.testng.annotations.Factory; -+import org.testng.annotations.Test; -+ -+import java.util.ArrayList; -+import java.util.Arrays; -+import java.util.Collections; -+import java.util.List; -+import java.util.UUID; -+import java.util.stream.Collectors; -+ -+import static org.assertj.core.api.Assertions.assertThat; -+import static org.assertj.core.api.Assertions.fail; -+ -+public class ReadManyByPartitionKeyTest extends TestSuiteBase { -+ -+ private String preExistingDatabaseId = CosmosDatabaseForTest.generateId(); -+ private CosmosClient client; -+ private CosmosDatabase createdDatabase; -+ -+ // Single PK container (/mypk) -+ private CosmosContainer singlePkContainer; -+ -+ // HPK container (/city, /zipcode, /areaCode) -+ private CosmosContainer multiHashContainer; -+ -+ @Factory(dataProvider = "clientBuilders") -+ public ReadManyByPartitionKeyTest(CosmosClientBuilder clientBuilder) { -+ super(clientBuilder); -+ } -+ -+ @BeforeClass(groups = {"emulator"}, timeOut = SETUP_TIMEOUT) -+ public void before_ReadManyByPartitionKeyTest() { -+ client = getClientBuilder().buildClient(); -+ createdDatabase = createSyncDatabase(client, preExistingDatabaseId); -+ -+ // Single PK container -+ String singlePkContainerName = UUID.randomUUID().toString(); -+ CosmosContainerProperties singlePkProps = new CosmosContainerProperties(singlePkContainerName, "/mypk"); -+ createdDatabase.createContainer(singlePkProps); -+ singlePkContainer = createdDatabase.getContainer(singlePkContainerName); -+ -+ // HPK container -+ String multiHashContainerName = UUID.randomUUID().toString(); -+ PartitionKeyDefinition hpkDef = new PartitionKeyDefinition(); -+ hpkDef.setKind(PartitionKind.MULTI_HASH); -+ hpkDef.setVersion(PartitionKeyDefinitionVersion.V2); -+ ArrayList paths = new ArrayList<>(); -+ paths.add("/city"); -+ paths.add("/zipcode"); -+ paths.add("/areaCode"); -+ hpkDef.setPaths(paths); -+ -+ CosmosContainerProperties hpkProps = new CosmosContainerProperties(multiHashContainerName, hpkDef); -+ createdDatabase.createContainer(hpkProps); -+ multiHashContainer = createdDatabase.getContainer(multiHashContainerName); -+ } -+ -+ @AfterClass(groups = {"emulator"}, timeOut = SHUTDOWN_TIMEOUT, alwaysRun = true) -+ public void afterClass() { -+ safeDeleteSyncDatabase(createdDatabase); -+ safeCloseSyncClient(client); -+ } -+ -+ //region Single PK tests -+ -+ @Test(groups = {"emulator"}, timeOut = TIMEOUT) -+ public void singlePk_readManyByPartitionKey_basic() { -+ // Create items with different PKs -+ List items = createSinglePkItems("pk1", 3); -+ items.addAll(createSinglePkItems("pk2", 2)); -+ items.addAll(createSinglePkItems("pk3", 4)); -+ -+ // Read by 2 partition keys -+ List pkValues = Arrays.asList( -+ new PartitionKey("pk1"), -+ new PartitionKey("pk2")); -+ -+ CosmosPagedIterable results = singlePkContainer.readManyByPartitionKey(pkValues, ObjectNode.class); -+ List resultList = results.stream().collect(Collectors.toList()); -+ -+ assertThat(resultList).hasSize(5); // 3 + 2 -+ resultList.forEach(item -> { -+ String pk = item.get("mypk").asText(); -+ assertThat(pk).isIn("pk1", "pk2"); -+ }); -+ -+ // Cleanup -+ cleanupContainer(singlePkContainer); -+ } -+ -+ @Test(groups = {"emulator"}, timeOut = TIMEOUT) -+ public void singlePk_readManyByPartitionKey_withProjection() { -+ List items = createSinglePkItems("pkProj", 2); -+ -+ List pkValues = Collections.singletonList(new PartitionKey("pkProj")); -+ SqlQuerySpec customQuery = new SqlQuerySpec("SELECT c.id, c.mypk FROM c"); -+ -+ CosmosPagedIterable results = singlePkContainer.readManyByPartitionKey( -+ pkValues, customQuery, null, ObjectNode.class); -+ List resultList = results.stream().collect(Collectors.toList()); -+ -+ assertThat(resultList).hasSize(2); -+ // Should only have id and mypk fields (plus system properties) -+ resultList.forEach(item -> { -+ assertThat(item.has("id")).isTrue(); -+ assertThat(item.has("mypk")).isTrue(); -+ }); -+ -+ cleanupContainer(singlePkContainer); -+ } -+ -+ @Test(groups = {"emulator"}, timeOut = TIMEOUT) -+ public void singlePk_readManyByPartitionKey_withAdditionalFilter() { -+ // Create items with different "status" values -+ createSinglePkItemsWithStatus("pkFilter", "active", 3); -+ createSinglePkItemsWithStatus("pkFilter", "inactive", 2); -+ -+ List pkValues = Collections.singletonList(new PartitionKey("pkFilter")); -+ SqlQuerySpec customQuery = new SqlQuerySpec( -+ "SELECT * FROM c WHERE c.status = @status", -+ Arrays.asList(new SqlParameter("@status", "active"))); -+ -+ CosmosPagedIterable results = singlePkContainer.readManyByPartitionKey( -+ pkValues, customQuery, null, ObjectNode.class); -+ List resultList = results.stream().collect(Collectors.toList()); -+ -+ assertThat(resultList).hasSize(3); -+ resultList.forEach(item -> { -+ assertThat(item.get("status").asText()).isEqualTo("active"); -+ }); -+ -+ cleanupContainer(singlePkContainer); -+ } -+ -+ @Test(groups = {"emulator"}, timeOut = TIMEOUT) -+ public void singlePk_readManyByPartitionKey_emptyResults() { -+ List pkValues = Collections.singletonList(new PartitionKey("nonExistent")); -+ -+ CosmosPagedIterable results = singlePkContainer.readManyByPartitionKey(pkValues, ObjectNode.class); -+ List resultList = results.stream().collect(Collectors.toList()); -+ -+ assertThat(resultList).isEmpty(); -+ } -+ -+ //endregion -+ -+ //region HPK tests -+ -+ @Test(groups = {"emulator"}, timeOut = TIMEOUT) -+ public void hpk_readManyByPartitionKey_fullPk() { -+ createHpkItems(); -+ -+ // Read by full PKs -+ List pkValues = Arrays.asList( -+ new PartitionKeyBuilder().add("Redmond").add("98053").add(1).build(), -+ new PartitionKeyBuilder().add("Pittsburgh").add("15232").add(2).build()); -+ -+ CosmosPagedIterable results = multiHashContainer.readManyByPartitionKey(pkValues, ObjectNode.class); -+ List resultList = results.stream().collect(Collectors.toList()); -+ -+ // Redmond/98053/1 has 2 items, Pittsburgh/15232/2 has 1 item -+ assertThat(resultList).hasSize(3); -+ -+ cleanupContainer(multiHashContainer); -+ } -+ -+ @Test(groups = {"emulator"}, timeOut = TIMEOUT) -+ public void hpk_readManyByPartitionKey_partialPk_singleLevel() { -+ createHpkItems(); -+ -+ // Read by partial PK (only city) -+ List pkValues = Collections.singletonList( -+ new PartitionKeyBuilder().add("Redmond").build()); -+ -+ CosmosPagedIterable results = multiHashContainer.readManyByPartitionKey(pkValues, ObjectNode.class); -+ List resultList = results.stream().collect(Collectors.toList()); -+ -+ // Redmond has 3 items total (2 with 98053/1 and 1 with 12345/1) -+ assertThat(resultList).hasSize(3); -+ resultList.forEach(item -> { -+ assertThat(item.get("city").asText()).isEqualTo("Redmond"); -+ }); -+ -+ cleanupContainer(multiHashContainer); -+ } -+ -+ @Test(groups = {"emulator"}, timeOut = TIMEOUT) -+ public void hpk_readManyByPartitionKey_partialPk_twoLevels() { -+ createHpkItems(); -+ -+ // Read by partial PK (city + zipcode) -+ List pkValues = Collections.singletonList( -+ new PartitionKeyBuilder().add("Redmond").add("98053").build()); -+ -+ CosmosPagedIterable results = multiHashContainer.readManyByPartitionKey(pkValues, ObjectNode.class); -+ List resultList = results.stream().collect(Collectors.toList()); -+ -+ // Redmond/98053 has 2 items -+ assertThat(resultList).hasSize(2); -+ resultList.forEach(item -> { -+ assertThat(item.get("city").asText()).isEqualTo("Redmond"); -+ assertThat(item.get("zipcode").asText()).isEqualTo("98053"); -+ }); -+ -+ cleanupContainer(multiHashContainer); -+ } -+ -+ @Test(groups = {"emulator"}, timeOut = TIMEOUT) -+ public void hpk_readManyByPartitionKey_withProjection() { -+ createHpkItems(); -+ -+ List pkValues = Collections.singletonList( -+ new PartitionKeyBuilder().add("Redmond").add("98053").add(1).build()); -+ -+ SqlQuerySpec customQuery = new SqlQuerySpec("SELECT c.id, c.city FROM c"); -+ -+ CosmosPagedIterable results = multiHashContainer.readManyByPartitionKey( -+ pkValues, customQuery, null, ObjectNode.class); -+ List resultList = results.stream().collect(Collectors.toList()); -+ -+ assertThat(resultList).hasSize(2); -+ -+ cleanupContainer(multiHashContainer); -+ } -+ -+ //endregion -+ -+ //region Negative/validation tests -+ -+ @Test(groups = {"emulator"}, timeOut = TIMEOUT) -+ public void rejectsAggregateQuery() { -+ List pkValues = Collections.singletonList(new PartitionKey("pk1")); -+ SqlQuerySpec aggregateQuery = new SqlQuerySpec("SELECT COUNT(1) FROM c"); -+ -+ try { -+ singlePkContainer.readManyByPartitionKey(pkValues, aggregateQuery, null, ObjectNode.class) -+ .stream().collect(Collectors.toList()); -+ fail("Should have thrown IllegalArgumentException for aggregate query"); -+ } catch (IllegalArgumentException e) { -+ assertThat(e.getMessage()).contains("aggregates"); -+ } -+ } -+ -+ @Test(groups = {"emulator"}, timeOut = TIMEOUT) -+ public void rejectsOrderByQuery() { -+ List pkValues = Collections.singletonList(new PartitionKey("pk1")); -+ SqlQuerySpec orderByQuery = new SqlQuerySpec("SELECT * FROM c ORDER BY c.id"); -+ -+ try { -+ singlePkContainer.readManyByPartitionKey(pkValues, orderByQuery, null, ObjectNode.class) -+ .stream().collect(Collectors.toList()); -+ fail("Should have thrown IllegalArgumentException for ORDER BY query"); -+ } catch (IllegalArgumentException e) { -+ assertThat(e.getMessage()).contains("ORDER BY"); -+ } -+ } -+ -+ @Test(groups = {"emulator"}, timeOut = TIMEOUT) -+ public void rejectsDistinctQuery() { -+ List pkValues = Collections.singletonList(new PartitionKey("pk1")); -+ SqlQuerySpec distinctQuery = new SqlQuerySpec("SELECT DISTINCT c.mypk FROM c"); -+ -+ try { -+ singlePkContainer.readManyByPartitionKey(pkValues, distinctQuery, null, ObjectNode.class) -+ .stream().collect(Collectors.toList()); -+ fail("Should have thrown IllegalArgumentException for DISTINCT query"); -+ } catch (IllegalArgumentException e) { -+ assertThat(e.getMessage()).contains("DISTINCT"); -+ } -+ } -+ -+ @Test(groups = {"emulator"}, timeOut = TIMEOUT) -+ public void rejectsGroupByQuery() { -+ List pkValues = Collections.singletonList(new PartitionKey("pk1")); -+ SqlQuerySpec groupByQuery = new SqlQuerySpec("SELECT c.mypk FROM c GROUP BY c.mypk"); -+ -+ try { -+ singlePkContainer.readManyByPartitionKey(pkValues, groupByQuery, null, ObjectNode.class) -+ .stream().collect(Collectors.toList()); -+ fail("Should have thrown IllegalArgumentException for GROUP BY query"); -+ } catch (IllegalArgumentException e) { -+ assertThat(e.getMessage()).contains("GROUP BY"); -+ } -+ } -+ -+ @Test(groups = {"emulator"}, timeOut = TIMEOUT) -+ public void rejectsGroupByWithAggregateQuery() { -+ List pkValues = Collections.singletonList(new PartitionKey("pk1")); -+ SqlQuerySpec groupByWithAggregateQuery = new SqlQuerySpec("SELECT c.mypk, COUNT(1) as cnt FROM c GROUP BY c.mypk"); -+ -+ try { -+ singlePkContainer.readManyByPartitionKey(pkValues, groupByWithAggregateQuery, null, ObjectNode.class) -+ .stream().collect(Collectors.toList()); -+ fail("Should have thrown IllegalArgumentException for GROUP BY with aggregate query"); -+ } catch (IllegalArgumentException e) { -+ assertThat(e.getMessage()).contains("GROUP BY"); -+ } -+ } -+ -+ @Test(groups = {"emulator"}, timeOut = TIMEOUT, expectedExceptions = IllegalArgumentException.class) -+ public void rejectsNullPartitionKeyList() { -+ singlePkContainer.readManyByPartitionKey((List) null, ObjectNode.class); -+ } -+ -+ @Test(groups = {"emulator"}, timeOut = TIMEOUT, expectedExceptions = IllegalArgumentException.class) -+ public void rejectsEmptyPartitionKeyList() { -+ singlePkContainer.readManyByPartitionKey(new ArrayList<>(), ObjectNode.class) -+ .stream().collect(Collectors.toList()); -+ } -+ -+ //endregion -+ -+ -+ //region Batch size tests (#10) -+ -+ @Test(groups = {"emulator"}, timeOut = TIMEOUT) -+ public void singlePk_readManyByPartitionKey_withSmallBatchSize() { -+ // Temporarily set batch size to 2 to exercise the batching/interleaving logic -+ String originalValue = System.getProperty("COSMOS.READ_MANY_BY_PK_MAX_BATCH_SIZE"); -+ try { -+ System.setProperty("COSMOS.READ_MANY_BY_PK_MAX_BATCH_SIZE", "2"); -+ -+ // Create items across 4 PKs (more than the batch size of 2) -+ List items = createSinglePkItems("batchPk1", 2); -+ items.addAll(createSinglePkItems("batchPk2", 2)); -+ items.addAll(createSinglePkItems("batchPk3", 2)); -+ items.addAll(createSinglePkItems("batchPk4", 2)); -+ -+ // Read all 4 PKs ΓÇö should be split into batches of 2 -+ List pkValues = Arrays.asList( -+ new PartitionKey("batchPk1"), -+ new PartitionKey("batchPk2"), -+ new PartitionKey("batchPk3"), -+ new PartitionKey("batchPk4")); -+ -+ CosmosPagedIterable results = singlePkContainer.readManyByPartitionKey(pkValues, ObjectNode.class); -+ List resultList = results.stream().collect(Collectors.toList()); -+ -+ assertThat(resultList).hasSize(8); // 2 items per PK * 4 PKs -+ resultList.forEach(item -> { -+ String pk = item.get("mypk").asText(); -+ assertThat(pk).isIn("batchPk1", "batchPk2", "batchPk3", "batchPk4"); -+ }); -+ -+ cleanupContainer(singlePkContainer); -+ } finally { -+ if (originalValue != null) { -+ System.setProperty("COSMOS.READ_MANY_BY_PK_MAX_BATCH_SIZE", originalValue); -+ } else { -+ System.clearProperty("COSMOS.READ_MANY_BY_PK_MAX_BATCH_SIZE"); -+ } -+ } -+ } -+ -+ //endregion -+ -+ //region Custom serializer regression tests (#5) -+ -+ @Test(groups = {"emulator"}, timeOut = TIMEOUT) -+ public void singlePk_readManyByPartitionKey_withRequestOptions() { -+ // This test ensures that request options (like throughput control settings) -+ // are properly propagated through the readManyByPartitionKey path. -+ // It acts as a regression test for the redundant options construction bug. -+ List items = createSinglePkItems("pkOpts", 3); -+ -+ List pkValues = Collections.singletonList(new PartitionKey("pkOpts")); -+ com.azure.cosmos.models.CosmosReadManyRequestOptions options = new com.azure.cosmos.models.CosmosReadManyRequestOptions(); -+ -+ CosmosPagedIterable results = singlePkContainer.readManyByPartitionKey( -+ pkValues, options, ObjectNode.class); -+ List resultList = results.stream().collect(Collectors.toList()); -+ -+ assertThat(resultList).hasSize(3); -+ -+ cleanupContainer(singlePkContainer); -+ } -+ -+ //endregion -+ -+ //region helper methods -+ -+ private List createSinglePkItems(String pkValue, int count) { -+ List items = new ArrayList<>(); -+ for (int i = 0; i < count; i++) { -+ ObjectNode item = com.azure.cosmos.implementation.Utils.getSimpleObjectMapper().createObjectNode(); -+ item.put("id", UUID.randomUUID().toString()); -+ item.put("mypk", pkValue); -+ singlePkContainer.createItem(item); -+ items.add(item); -+ } -+ return items; -+ } -+ -+ private List createSinglePkItemsWithStatus(String pkValue, String status, int count) { -+ List items = new ArrayList<>(); -+ for (int i = 0; i < count; i++) { -+ ObjectNode item = com.azure.cosmos.implementation.Utils.getSimpleObjectMapper().createObjectNode(); -+ item.put("id", UUID.randomUUID().toString()); -+ item.put("mypk", pkValue); -+ item.put("status", status); -+ singlePkContainer.createItem(item); -+ items.add(item); -+ } -+ return items; -+ } -+ -+ private void createHpkItems() { -+ // Same data as CosmosMultiHashTest.createItems() -+ createHpkItem("Redmond", "98053", 1); -+ createHpkItem("Redmond", "98053", 1); -+ createHpkItem("Pittsburgh", "15232", 2); -+ createHpkItem("Stonybrook", "11790", 3); -+ createHpkItem("Stonybrook", "11794", 3); -+ createHpkItem("Stonybrook", "11791", 3); -+ createHpkItem("Redmond", "12345", 1); -+ } -+ -+ private void createHpkItem(String city, String zipcode, int areaCode) { -+ ObjectNode item = com.azure.cosmos.implementation.Utils.getSimpleObjectMapper().createObjectNode(); -+ item.put("id", UUID.randomUUID().toString()); -+ item.put("city", city); -+ item.put("zipcode", zipcode); -+ item.put("areaCode", areaCode); -+ multiHashContainer.createItem(item); -+ } -+ -+ private void cleanupContainer(CosmosContainer container) { -+ CosmosPagedIterable allItems = container.queryItems( -+ "SELECT * FROM c", new com.azure.cosmos.models.CosmosQueryRequestOptions(), ObjectNode.class); -+ allItems.forEach(item -> { -+ try { -+ container.deleteItem(item, new CosmosItemRequestOptions()); -+ } catch (CosmosException e) { -+ // ignore cleanup failures -+ } -+ }); -+ } -+ -+ //endregion -+} -diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelperTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelperTest.java -new file mode 100644 -index 00000000000..95c109ba025 ---- /dev/null -+++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelperTest.java -@@ -0,0 +1,426 @@ -+/* -+ * Copyright (c) Microsoft Corporation. All rights reserved. -+ * Licensed under the MIT License. -+ */ -+ -+package com.azure.cosmos.implementation; -+ -+import com.azure.cosmos.models.PartitionKey; -+import com.azure.cosmos.models.PartitionKeyBuilder; -+import com.azure.cosmos.models.PartitionKeyDefinition; -+import com.azure.cosmos.models.PartitionKeyDefinitionVersion; -+import com.azure.cosmos.models.PartitionKind; -+import com.azure.cosmos.models.SqlParameter; -+import com.azure.cosmos.models.SqlQuerySpec; -+import org.testng.annotations.Test; -+ -+import java.util.ArrayList; -+import java.util.Arrays; -+import java.util.Collections; -+import java.util.List; -+import java.util.stream.Collectors; -+ -+import static org.assertj.core.api.Assertions.assertThat; -+ -+public class ReadManyByPartitionKeyQueryHelperTest { -+ -+ //region Single PK (HASH) tests -+ -+ @Test(groups = { "unit" }) -+ public void singlePk_defaultQuery_singleValue() { -+ PartitionKeyDefinition pkDef = createSinglePkDefinition("/mypk"); -+ List selectors = createSelectors(pkDef); -+ List pkValues = Collections.singletonList(new PartitionKey("pk1")); -+ -+ SqlQuerySpec result = ReadManyByPartitionKeyQueryHelper.createReadManyByPkQuerySpec( -+ "SELECT * FROM c", new ArrayList<>(), pkValues, selectors, pkDef); -+ -+ assertThat(result.getQueryText()).contains("SELECT * FROM c WHERE"); -+ assertThat(result.getQueryText()).contains("IN ("); -+ assertThat(result.getQueryText()).contains("@__rmPk_0"); -+ assertThat(result.getParameters()).hasSize(1); -+ assertThat(result.getParameters().get(0).getValue(Object.class)).isEqualTo("pk1"); -+ } -+ -+ @Test(groups = { "unit" }) -+ public void singlePk_defaultQuery_multipleValues() { -+ PartitionKeyDefinition pkDef = createSinglePkDefinition("/mypk"); -+ List selectors = createSelectors(pkDef); -+ List pkValues = Arrays.asList( -+ new PartitionKey("pk1"), -+ new PartitionKey("pk2"), -+ new PartitionKey("pk3")); -+ -+ SqlQuerySpec result = ReadManyByPartitionKeyQueryHelper.createReadManyByPkQuerySpec( -+ "SELECT * FROM c", new ArrayList<>(), pkValues, selectors, pkDef); -+ -+ assertThat(result.getQueryText()).contains("IN ("); -+ assertThat(result.getQueryText()).contains("@__rmPk_0"); -+ assertThat(result.getQueryText()).contains("@__rmPk_1"); -+ assertThat(result.getQueryText()).contains("@__rmPk_2"); -+ assertThat(result.getParameters()).hasSize(3); -+ } -+ -+ @Test(groups = { "unit" }) -+ public void singlePk_customQuery_noWhere() { -+ PartitionKeyDefinition pkDef = createSinglePkDefinition("/mypk"); -+ List selectors = createSelectors(pkDef); -+ List pkValues = Arrays.asList(new PartitionKey("pk1"), new PartitionKey("pk2")); -+ -+ SqlQuerySpec result = ReadManyByPartitionKeyQueryHelper.createReadManyByPkQuerySpec( -+ "SELECT c.name, c.age FROM c", new ArrayList<>(), pkValues, selectors, pkDef); -+ -+ assertThat(result.getQueryText()).startsWith("SELECT c.name, c.age FROM c WHERE"); -+ assertThat(result.getQueryText()).contains("IN ("); -+ } -+ -+ @Test(groups = { "unit" }) -+ public void singlePk_customQuery_withExistingWhere() { -+ PartitionKeyDefinition pkDef = createSinglePkDefinition("/mypk"); -+ List selectors = createSelectors(pkDef); -+ List pkValues = Collections.singletonList(new PartitionKey("pk1")); -+ -+ List baseParams = new ArrayList<>(); -+ baseParams.add(new SqlParameter("@minAge", 18)); -+ -+ SqlQuerySpec result = ReadManyByPartitionKeyQueryHelper.createReadManyByPkQuerySpec( -+ "SELECT * FROM c WHERE c.age > @minAge", baseParams, pkValues, selectors, pkDef); -+ -+ // Should AND the PK filter to the existing WHERE clause -+ assertThat(result.getQueryText()).contains("WHERE (c.age > @minAge) AND ("); -+ assertThat(result.getQueryText()).contains("IN ("); -+ assertThat(result.getParameters()).hasSize(2); // @minAge + @__rmPk_0 -+ assertThat(result.getParameters().get(0).getName()).isEqualTo("@minAge"); -+ } -+ -+ //endregion -+ -+ //region HPK (MULTI_HASH) tests -+ -+ @Test(groups = { "unit" }) -+ public void hpk_fullPk_defaultQuery() { -+ PartitionKeyDefinition pkDef = createMultiHashPkDefinition("/city", "/zipcode"); -+ List selectors = createSelectors(pkDef); -+ -+ PartitionKey pk = new PartitionKeyBuilder().add("Redmond").add("98052").build(); -+ List pkValues = Collections.singletonList(pk); -+ -+ SqlQuerySpec result = ReadManyByPartitionKeyQueryHelper.createReadManyByPkQuerySpec( -+ "SELECT * FROM c", new ArrayList<>(), pkValues, selectors, pkDef); -+ -+ assertThat(result.getQueryText()).contains("SELECT * FROM c WHERE"); -+ // Should use OR/AND pattern, not IN -+ assertThat(result.getQueryText()).doesNotContain("IN ("); -+ assertThat(result.getQueryText()).contains("c[\"city\"] = @__rmPk_0"); -+ assertThat(result.getQueryText()).contains("AND"); -+ assertThat(result.getQueryText()).contains("c[\"zipcode\"] = @__rmPk_1"); -+ assertThat(result.getParameters()).hasSize(2); -+ assertThat(result.getParameters().get(0).getValue(Object.class)).isEqualTo("Redmond"); -+ assertThat(result.getParameters().get(1).getValue(Object.class)).isEqualTo("98052"); -+ } -+ -+ @Test(groups = { "unit" }) -+ public void hpk_fullPk_multipleValues() { -+ PartitionKeyDefinition pkDef = createMultiHashPkDefinition("/city", "/zipcode"); -+ List selectors = createSelectors(pkDef); -+ -+ List pkValues = Arrays.asList( -+ new PartitionKeyBuilder().add("Redmond").add("98052").build(), -+ new PartitionKeyBuilder().add("Seattle").add("98101").build()); -+ -+ SqlQuerySpec result = ReadManyByPartitionKeyQueryHelper.createReadManyByPkQuerySpec( -+ "SELECT * FROM c", new ArrayList<>(), pkValues, selectors, pkDef); -+ -+ assertThat(result.getQueryText()).contains("OR"); -+ assertThat(result.getQueryText()).contains("c[\"city\"] = @__rmPk_0"); -+ assertThat(result.getQueryText()).contains("c[\"zipcode\"] = @__rmPk_1"); -+ assertThat(result.getQueryText()).contains("c[\"city\"] = @__rmPk_2"); -+ assertThat(result.getQueryText()).contains("c[\"zipcode\"] = @__rmPk_3"); -+ assertThat(result.getParameters()).hasSize(4); -+ } -+ -+ @Test(groups = { "unit" }) -+ public void hpk_partialPk_singleLevel() { -+ PartitionKeyDefinition pkDef = createMultiHashPkDefinition("/city", "/zipcode", "/areaCode"); -+ List selectors = createSelectors(pkDef); -+ -+ // Partial PK ΓÇö only first level -+ PartitionKey partialPk = new PartitionKeyBuilder().add("Redmond").build(); -+ List pkValues = Collections.singletonList(partialPk); -+ -+ SqlQuerySpec result = ReadManyByPartitionKeyQueryHelper.createReadManyByPkQuerySpec( -+ "SELECT * FROM c", new ArrayList<>(), pkValues, selectors, pkDef); -+ -+ assertThat(result.getQueryText()).contains("c[\"city\"] = @__rmPk_0"); -+ // Should NOT include zipcode or areaCode since it's partial -+ assertThat(result.getQueryText()).doesNotContain("zipcode"); -+ assertThat(result.getQueryText()).doesNotContain("areaCode"); -+ assertThat(result.getParameters()).hasSize(1); -+ } -+ -+ @Test(groups = { "unit" }) -+ public void hpk_partialPk_twoLevels() { -+ PartitionKeyDefinition pkDef = createMultiHashPkDefinition("/city", "/zipcode", "/areaCode"); -+ List selectors = createSelectors(pkDef); -+ -+ // Partial PK ΓÇö first two levels -+ PartitionKey partialPk = new PartitionKeyBuilder().add("Redmond").add("98052").build(); -+ List pkValues = Collections.singletonList(partialPk); -+ -+ SqlQuerySpec result = ReadManyByPartitionKeyQueryHelper.createReadManyByPkQuerySpec( -+ "SELECT * FROM c", new ArrayList<>(), pkValues, selectors, pkDef); -+ -+ assertThat(result.getQueryText()).contains("c[\"city\"] = @__rmPk_0"); -+ assertThat(result.getQueryText()).contains("c[\"zipcode\"] = @__rmPk_1"); -+ assertThat(result.getQueryText()).doesNotContain("areaCode"); -+ assertThat(result.getParameters()).hasSize(2); -+ } -+ -+ @Test(groups = { "unit" }) -+ public void hpk_customQuery_withWhere() { -+ PartitionKeyDefinition pkDef = createMultiHashPkDefinition("/city", "/zipcode"); -+ List selectors = createSelectors(pkDef); -+ -+ List baseParams = new ArrayList<>(); -+ baseParams.add(new SqlParameter("@status", "active")); -+ -+ PartitionKey pk = new PartitionKeyBuilder().add("Redmond").add("98052").build(); -+ List pkValues = Collections.singletonList(pk); -+ -+ SqlQuerySpec result = ReadManyByPartitionKeyQueryHelper.createReadManyByPkQuerySpec( -+ "SELECT c.name FROM c WHERE c.status = @status", baseParams, pkValues, selectors, pkDef); -+ -+ assertThat(result.getQueryText()).contains("WHERE (c.status = @status) AND ("); -+ assertThat(result.getQueryText()).contains("c[\"city\"] = @__rmPk_0"); -+ assertThat(result.getParameters()).hasSize(3); // @status + 2 pk params -+ } -+ -+ //endregion -+ -+ //region findTopLevelWhereIndex tests -+ -+ @Test(groups = { "unit" }) -+ public void findWhere_simpleQuery() { -+ int idx = ReadManyByPartitionKeyQueryHelper.findTopLevelWhereIndex("SELECT * FROM C WHERE C.ID = 1"); -+ assertThat(idx).isEqualTo(16); -+ } -+ -+ @Test(groups = { "unit" }) -+ public void findWhere_noWhere() { -+ int idx = ReadManyByPartitionKeyQueryHelper.findTopLevelWhereIndex("SELECT * FROM C"); -+ assertThat(idx).isEqualTo(-1); -+ } -+ -+ @Test(groups = { "unit" }) -+ public void findWhere_whereInSubquery() { -+ // WHERE inside parentheses (subquery) should be ignored -+ int idx = ReadManyByPartitionKeyQueryHelper.findTopLevelWhereIndex( -+ "SELECT * FROM C WHERE EXISTS(SELECT VALUE T FROM T IN C.TAGS WHERE T = 'FOO')"); -+ // Should find the outer WHERE, not the inner one -+ assertThat(idx).isEqualTo(16); -+ } -+ -+ @Test(groups = { "unit" }) -+ public void findWhere_caseInsensitive() { -+ int idx = ReadManyByPartitionKeyQueryHelper.findTopLevelWhereIndex("SELECT * FROM C WHERE C.X = 1"); -+ assertThat(idx).isGreaterThan(0); -+ } -+ -+ @Test(groups = { "unit" }) -+ public void findWhere_whereNotKeyword() { -+ // "ELSEWHERE" should not match -+ int idx = ReadManyByPartitionKeyQueryHelper.findTopLevelWhereIndex("SELECT * FROM ELSEWHERE"); -+ assertThat(idx).isEqualTo(-1); -+ } -+ -+ //endregion -+ -+ //region Custom alias tests -+ -+ @Test(groups = { "unit" }) -+ public void singlePk_customAlias() { -+ PartitionKeyDefinition pkDef = createSinglePkDefinition("/mypk"); -+ List selectors = createSelectors(pkDef); -+ List pkValues = Arrays.asList(new PartitionKey("pk1"), new PartitionKey("pk2")); -+ -+ SqlQuerySpec result = ReadManyByPartitionKeyQueryHelper.createReadManyByPkQuerySpec( -+ "SELECT x.id, x.mypk FROM x", new ArrayList<>(), pkValues, selectors, pkDef); -+ -+ assertThat(result.getQueryText()).startsWith("SELECT x.id, x.mypk FROM x WHERE"); -+ assertThat(result.getQueryText()).contains("x[\"mypk\"] IN ("); -+ assertThat(result.getQueryText()).doesNotContain("c[\"mypk\"]"); -+ } -+ -+ @Test(groups = { "unit" }) -+ public void singlePk_customAlias_withWhere() { -+ PartitionKeyDefinition pkDef = createSinglePkDefinition("/mypk"); -+ List selectors = createSelectors(pkDef); -+ List pkValues = Collections.singletonList(new PartitionKey("pk1")); -+ -+ List baseParams = new ArrayList<>(); -+ baseParams.add(new SqlParameter("@cat", "HelloWorld")); -+ -+ SqlQuerySpec result = ReadManyByPartitionKeyQueryHelper.createReadManyByPkQuerySpec( -+ "SELECT x.id, x.mypk FROM x WHERE x.category = @cat", baseParams, pkValues, selectors, pkDef); -+ -+ assertThat(result.getQueryText()).contains("WHERE (x.category = @cat) AND (x[\"mypk\"] IN ("); -+ } -+ -+ @Test(groups = { "unit" }) -+ public void hpk_customAlias() { -+ PartitionKeyDefinition pkDef = createMultiHashPkDefinition("/city", "/zipcode"); -+ List selectors = createSelectors(pkDef); -+ -+ PartitionKey pk = new PartitionKeyBuilder().add("Redmond").add("98052").build(); -+ List pkValues = Collections.singletonList(pk); -+ -+ SqlQuerySpec result = ReadManyByPartitionKeyQueryHelper.createReadManyByPkQuerySpec( -+ "SELECT r.name FROM root r", new ArrayList<>(), pkValues, selectors, pkDef); -+ -+ assertThat(result.getQueryText()).contains("r[\"city\"] = @__rmPk_0"); -+ assertThat(result.getQueryText()).contains("r[\"zipcode\"] = @__rmPk_1"); -+ assertThat(result.getQueryText()).doesNotContain("c[\""); -+ } -+ -+ //endregion -+ -+ //region extractTableAlias tests -+ -+ @Test(groups = { "unit" }) -+ public void extractAlias_defaultC() { -+ assertThat(ReadManyByPartitionKeyQueryHelper.extractTableAlias("SELECT * FROM c")).isEqualTo("c"); -+ } -+ -+ @Test(groups = { "unit" }) -+ public void extractAlias_customX() { -+ assertThat(ReadManyByPartitionKeyQueryHelper.extractTableAlias("SELECT x.id FROM x WHERE x.age > 5")).isEqualTo("x"); -+ } -+ -+ @Test(groups = { "unit" }) -+ public void extractAlias_rootWithAlias() { -+ assertThat(ReadManyByPartitionKeyQueryHelper.extractTableAlias("SELECT r.name FROM root r")).isEqualTo("r"); -+ } -+ -+ @Test(groups = { "unit" }) -+ public void extractAlias_rootNoAlias() { -+ assertThat(ReadManyByPartitionKeyQueryHelper.extractTableAlias("SELECT * FROM root")).isEqualTo("root"); -+ } -+ -+ @Test(groups = { "unit" }) -+ public void extractAlias_containerWithWhere() { -+ assertThat(ReadManyByPartitionKeyQueryHelper.extractTableAlias("SELECT * FROM items WHERE items.status = 'active'")).isEqualTo("items"); -+ } -+ -+ @Test(groups = { "unit" }) -+ public void extractAlias_caseInsensitive() { -+ assertThat(ReadManyByPartitionKeyQueryHelper.extractTableAlias("select * from MyContainer where MyContainer.id = '1'")).isEqualTo("MyContainer"); -+ } -+ -+ //endregion -+ -+ -+ //region String literal handling tests (#1) -+ -+ @Test(groups = { "unit" }) -+ public void findWhere_ignoresWhereInsideStringLiteral() { -+ // WHERE inside a string literal should be ignored -+ int idx = ReadManyByPartitionKeyQueryHelper.findTopLevelWhereIndex( -+ "SELECT * FROM c WHERE c.msg = 'use WHERE clause here'"); -+ // Should find the outer WHERE at position 16, not the one inside the string -+ assertThat(idx).isEqualTo(16); -+ } -+ -+ @Test(groups = { "unit" }) -+ public void findWhere_ignoresParenthesesInsideStringLiteral() { -+ // Parentheses inside string literal should not affect depth tracking -+ int idx = ReadManyByPartitionKeyQueryHelper.findTopLevelWhereIndex( -+ "SELECT * FROM c WHERE c.name = 'foo(bar)' AND c.x = 1"); -+ assertThat(idx).isEqualTo(16); -+ } -+ -+ @Test(groups = { "unit" }) -+ public void findWhere_handlesUnbalancedParenInStringLiteral() { -+ // Unbalanced paren inside string literal must not corrupt depth -+ int idx = ReadManyByPartitionKeyQueryHelper.findTopLevelWhereIndex( -+ "SELECT * FROM c WHERE c.val = 'open(' AND c.active = true"); -+ assertThat(idx).isEqualTo(16); -+ } -+ -+ @Test(groups = { "unit" }) -+ public void findWhere_handlesStringLiteralBeforeWhere() { -+ // String literal in SELECT before WHERE -+ int idx = ReadManyByPartitionKeyQueryHelper.findTopLevelWhereIndex( -+ "SELECT 'WHERE' as label FROM c WHERE c.id = '1'"); -+ // The WHERE inside quotes should be ignored; the real WHERE is further along -+ assertThat(idx).isGreaterThan(30); -+ } -+ -+ @Test(groups = { "unit" }) -+ public void singlePk_customQuery_withStringLiteralContainingParens() { -+ PartitionKeyDefinition pkDef = createSinglePkDefinition("/mypk"); -+ List selectors = createSelectors(pkDef); -+ List pkValues = Collections.singletonList(new PartitionKey("pk1")); -+ -+ List baseParams = new ArrayList<>(); -+ baseParams.add(new SqlParameter("@msg", "hello")); -+ -+ SqlQuerySpec result = ReadManyByPartitionKeyQueryHelper.createReadManyByPkQuerySpec( -+ "SELECT * FROM c WHERE c.msg = 'test(value)WHERE'", baseParams, pkValues, selectors, pkDef); -+ -+ // Should correctly AND the PK filter to the real WHERE clause -+ assertThat(result.getQueryText()).contains("WHERE (c.msg = 'test(value)WHERE') AND ("); -+ } -+ -+ //endregion -+ -+ //region OFFSET/LIMIT/HAVING alias detection tests (#9) -+ -+ @Test(groups = { "unit" }) -+ public void extractAlias_containerWithOffset() { -+ assertThat(ReadManyByPartitionKeyQueryHelper.extractTableAlias( -+ "SELECT * FROM c OFFSET 10 LIMIT 5")).isEqualTo("c"); -+ } -+ -+ @Test(groups = { "unit" }) -+ public void extractAlias_containerWithLimit() { -+ assertThat(ReadManyByPartitionKeyQueryHelper.extractTableAlias( -+ "SELECT * FROM c LIMIT 10")).isEqualTo("c"); -+ } -+ -+ @Test(groups = { "unit" }) -+ public void extractAlias_containerWithHaving() { -+ assertThat(ReadManyByPartitionKeyQueryHelper.extractTableAlias( -+ "SELECT c.status, COUNT(1) FROM c GROUP BY c.status HAVING COUNT(1) > 1")).isEqualTo("c"); -+ } -+ -+ //endregion -+ -+ //region helpers -+ -+ private PartitionKeyDefinition createSinglePkDefinition(String path) { -+ PartitionKeyDefinition pkDef = new PartitionKeyDefinition(); -+ pkDef.setKind(PartitionKind.HASH); -+ pkDef.setVersion(PartitionKeyDefinitionVersion.V2); -+ pkDef.setPaths(Collections.singletonList(path)); -+ return pkDef; -+ } -+ -+ private PartitionKeyDefinition createMultiHashPkDefinition(String... paths) { -+ PartitionKeyDefinition pkDef = new PartitionKeyDefinition(); -+ pkDef.setKind(PartitionKind.MULTI_HASH); -+ pkDef.setVersion(PartitionKeyDefinitionVersion.V2); -+ pkDef.setPaths(Arrays.asList(paths)); -+ return pkDef; -+ } -+ -+ private List createSelectors(PartitionKeyDefinition pkDef) { -+ return pkDef.getPaths() -+ .stream() -+ .map(pathPart -> pathPart.substring(1)) // skip starting / -+ .map(pathPart -> pathPart.replace("\"", "\\")) -+ .map(part -> "[\"" + part + "\"]") -+ .collect(Collectors.toList()); -+ } -+ -+ //endregion -+} -diff --git a/sdk/cosmos/azure-cosmos/CHANGELOG.md b/sdk/cosmos/azure-cosmos/CHANGELOG.md -index e8ea564fab7..904c01c3238 100644 ---- a/sdk/cosmos/azure-cosmos/CHANGELOG.md -+++ b/sdk/cosmos/azure-cosmos/CHANGELOG.md -@@ -3,6 +3,7 @@ - ### 4.80.0-beta.1 (Unreleased) - - #### Features Added -+* Added new `readManyByPartitionKey` to bulk query by a list of pk-values with better efficiency. See [PR 48801](https://github.com/Azure/azure-sdk-for-java/pull/48801) - - #### Breaking Changes - -diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java -index ad871bb97c0..4e234667c1c 100644 ---- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java -+++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java -@@ -165,6 +165,7 @@ public class CosmosAsyncContainer { - private final String createItemSpanName; - private final String readAllItemsSpanName; - private final String readManyItemsSpanName; -+ private final String readManyByPartitionKeyItemsSpanName; - private final String readAllItemsOfLogicalPartitionSpanName; - private final String queryItemsSpanName; - private final String queryChangeFeedSpanName; -@@ -198,6 +199,7 @@ public class CosmosAsyncContainer { - this.createItemSpanName = "createItem." + this.id; - this.readAllItemsSpanName = "readAllItems." + this.id; - this.readManyItemsSpanName = "readManyItems." + this.id; -+ this.readManyByPartitionKeyItemsSpanName = "readManyByPartitionKeyItems." + this.id; - this.readAllItemsOfLogicalPartitionSpanName = "readAllItemsOfLogicalPartition." + this.id; - this.queryItemsSpanName = "queryItems." + this.id; - this.queryChangeFeedSpanName = "queryChangeFeed." + this.id; -@@ -1601,6 +1603,130 @@ public class CosmosAsyncContainer { - context); - } - -+ /** -+ * Reads many documents matching the provided partition key values. -+ * Unlike {@link #readMany(List, Class)} this method does not require item ids - it queries -+ * all documents matching the provided partition key values. Uses {@code SELECT * FROM c} -+ * as the base query. -+ * -+ * @param the type parameter -+ * @param partitionKeys list of partition key values to read documents for -+ * @param classType class type -+ * @return a {@link CosmosPagedFlux} containing one or several feed response pages -+ */ -+ public CosmosPagedFlux readManyByPartitionKey( -+ List partitionKeys, -+ Class classType) { -+ -+ return this.readManyByPartitionKey(partitionKeys, null, null, classType); -+ } -+ -+ /** -+ * Reads many documents matching the provided partition key values. -+ * Unlike {@link #readMany(List, Class)} this method does not require item ids - it queries -+ * all documents matching the provided partition key values. Uses {@code SELECT * FROM c} -+ * as the base query. -+ * -+ * @param the type parameter -+ * @param partitionKeys list of partition key values to read documents for -+ * @param requestOptions the optional request options -+ * @param classType class type -+ * @return a {@link CosmosPagedFlux} containing one or several feed response pages -+ */ -+ public CosmosPagedFlux readManyByPartitionKey( -+ List partitionKeys, -+ CosmosReadManyRequestOptions requestOptions, -+ Class classType) { -+ -+ return this.readManyByPartitionKey(partitionKeys, null, requestOptions, classType); -+ } -+ -+ /** -+ * Reads many documents matching the provided partition key values with a custom query. -+ * The custom query can be used to apply projections (e.g. {@code SELECT c.name, c.age FROM c}) -+ * and/or additional filters (e.g. {@code SELECT * FROM c WHERE c.status = 'active'}). -+ * The SDK will automatically append partition key filtering to the custom query. -+ *

-+ * The custom query must be a simple streamable query ΓÇö aggregates, ORDER BY, DISTINCT, -+ * GROUP BY, DCOUNT, vector search, and full-text search are not supported and will be -+ * rejected. -+ *

-+ * Partial hierarchical partition keys are supported and will fan out to multiple -+ * physical partitions. -+ * -+ * @param the type parameter -+ * @param partitionKeys list of partition key values to read documents for -+ * @param customQuery optional custom query for projections/additional filters (null means SELECT * FROM c) -+ * @param requestOptions the optional request options -+ * @param classType class type -+ * @return a {@link CosmosPagedFlux} containing one or several feed response pages -+ */ -+ public CosmosPagedFlux readManyByPartitionKey( -+ List partitionKeys, -+ SqlQuerySpec customQuery, -+ CosmosReadManyRequestOptions requestOptions, -+ Class classType) { -+ -+ if (partitionKeys == null) { -+ throw new IllegalArgumentException("Argument 'partitionKeys' must not be null."); -+ } -+ if (partitionKeys.isEmpty()) { -+ throw new IllegalArgumentException("Argument 'partitionKeys' must not be empty."); -+ } -+ for (PartitionKey pk : partitionKeys) { -+ if (pk == null) { -+ throw new IllegalArgumentException( -+ "Argument 'partitionKeys' must not contain null elements."); -+ } -+ } -+ -+ return UtilBridgeInternal.createCosmosPagedFlux( -+ readManyByPartitionKeyInternalFunc(partitionKeys, customQuery, requestOptions, classType)); -+ } -+ -+ private Function>> readManyByPartitionKeyInternalFunc( -+ List partitionKeys, -+ SqlQuerySpec customQuery, -+ CosmosReadManyRequestOptions requestOptions, -+ Class classType) { -+ -+ CosmosAsyncClient client = this.getDatabase().getClient(); -+ -+ return (pagedFluxOptions -> { -+ CosmosQueryRequestOptions queryRequestOptions = requestOptions == null -+ ? new CosmosQueryRequestOptions() -+ : queryOptionsAccessor().clone(readManyOptionsAccessor().getImpl(requestOptions)); -+ queryRequestOptions.setMaxDegreeOfParallelism(-1); -+ queryRequestOptions.setQueryName("readManyByPartitionKey"); -+ CosmosQueryRequestOptionsBase cosmosQueryRequestOptionsImpl = queryOptionsAccessor().getImpl(queryRequestOptions); -+ applyPolicies(OperationType.Query, ResourceType.Document, cosmosQueryRequestOptionsImpl, this.readManyByPartitionKeyItemsSpanName); -+ -+ QueryFeedOperationState state = new QueryFeedOperationState( -+ client, -+ this.readManyByPartitionKeyItemsSpanName, -+ database.getId(), -+ this.getId(), -+ ResourceType.Document, -+ OperationType.Query, -+ queryOptionsAccessor().getQueryNameOrDefault(queryRequestOptions, this.readManyByPartitionKeyItemsSpanName), -+ queryRequestOptions, -+ pagedFluxOptions -+ ); -+ -+ pagedFluxOptions.setFeedOperationState(state); -+ -+ return CosmosBridgeInternal -+ .getAsyncDocumentClient(this.getDatabase()) -+ .readManyByPartitionKey( -+ partitionKeys, -+ customQuery, -+ BridgeInternal.getLink(this), -+ state, -+ classType) -+ .map(response -> prepareFeedResponse(response, false)); -+ }); -+ } -+ - /** - * Reads all the items of a logical partition - * -diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosContainer.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosContainer.java -index 04a6060c192..0bd8be5850c 100644 ---- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosContainer.java -+++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosContainer.java -@@ -540,6 +540,73 @@ public class CosmosContainer { - classType)); - } - -+ /** -+ * Reads many documents matching the provided partition key values. -+ * Unlike {@link #readMany(List, Class)} this method does not require item ids - it queries -+ * all documents matching the provided partition key values. Uses {@code SELECT * FROM c} -+ * as the base query. -+ * -+ * @param the type parameter -+ * @param partitionKeys list of partition key values to read documents for -+ * @param classType class type -+ * @return a {@link CosmosPagedIterable} containing the results -+ */ -+ public CosmosPagedIterable readManyByPartitionKey( -+ List partitionKeys, -+ Class classType) { -+ -+ return getCosmosPagedIterable(this.asyncContainer.readManyByPartitionKey(partitionKeys, classType)); -+ } -+ -+ /** -+ * Reads many documents matching the provided partition key values. -+ * Unlike {@link #readMany(List, Class)} this method does not require item ids - it queries -+ * all documents matching the provided partition key values. Uses {@code SELECT * FROM c} -+ * as the base query. -+ * -+ * @param the type parameter -+ * @param partitionKeys list of partition key values to read documents for -+ * @param requestOptions the optional request options -+ * @param classType class type -+ * @return a {@link CosmosPagedIterable} containing the results -+ */ -+ public CosmosPagedIterable readManyByPartitionKey( -+ List partitionKeys, -+ CosmosReadManyRequestOptions requestOptions, -+ Class classType) { -+ -+ return getCosmosPagedIterable(this.asyncContainer.readManyByPartitionKey(partitionKeys, requestOptions, classType)); -+ } -+ -+ /** -+ * Reads many documents matching the provided partition key values with a custom query. -+ * The custom query can be used to apply projections (e.g. {@code SELECT c.name, c.age FROM c}) -+ * and/or additional filters (e.g. {@code SELECT * FROM c WHERE c.status = 'active'}). -+ * The SDK will automatically append partition key filtering to the custom query. -+ *

-+ * The custom query must be a simple streamable query ΓÇö aggregates, ORDER BY, DISTINCT, -+ * GROUP BY, DCOUNT, vector search, and full-text search are not supported and will be -+ * rejected. -+ *

-+ * Partial hierarchical partition keys are supported and will fan out to multiple -+ * physical partitions. -+ * -+ * @param the type parameter -+ * @param partitionKeys list of partition key values to read documents for -+ * @param customQuery optional custom query for projections/additional filters (null means SELECT * FROM c) -+ * @param requestOptions the optional request options -+ * @param classType class type -+ * @return a {@link CosmosPagedIterable} containing the results -+ */ -+ public CosmosPagedIterable readManyByPartitionKey( -+ List partitionKeys, -+ SqlQuerySpec customQuery, -+ CosmosReadManyRequestOptions requestOptions, -+ Class classType) { -+ -+ return getCosmosPagedIterable(this.asyncContainer.readManyByPartitionKey(partitionKeys, customQuery, requestOptions, classType)); -+ } -+ - /** - * Reads all the items of a logical partition returning the results as {@link CosmosPagedIterable}. - * -diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/AsyncDocumentClient.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/AsyncDocumentClient.java -index 945e768a82f..8e2499c9039 100644 ---- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/AsyncDocumentClient.java -+++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/AsyncDocumentClient.java -@@ -1584,6 +1584,27 @@ public interface AsyncDocumentClient { - QueryFeedOperationState state, - Class klass); - -+ /** -+ * Reads many documents by partition key values. -+ * Unlike {@link #readMany(List, String, QueryFeedOperationState, Class)} this method does not require -+ * item ids - it queries all documents matching the provided partition key values. -+ * Partial hierarchical partition keys are supported and will fan out to multiple physical partitions. -+ * -+ * @param partitionKeys list of partition key values to read documents for -+ * @param customQuery optional custom query (for projections/additional filters) - null means SELECT * FROM c -+ * @param collectionLink link for the documentcollection/container to be queried -+ * @param state the query operation state -+ * @param klass class type -+ * @param the type parameter -+ * @return a Flux with feed response pages of documents -+ */ -+ Flux> readManyByPartitionKey( -+ List partitionKeys, -+ SqlQuerySpec customQuery, -+ String collectionLink, -+ QueryFeedOperationState state, -+ Class klass); -+ - /** - * Read all documents of a certain logical partition. - *

-diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/Configs.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/Configs.java -index 337055c6947..162b0740f40 100644 ---- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/Configs.java -+++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/Configs.java -@@ -248,6 +248,11 @@ public class Configs { - public static final String MIN_TARGET_BULK_MICRO_BATCH_SIZE_VARIABLE = "COSMOS_MIN_TARGET_BULK_MICRO_BATCH_SIZE"; - public static final int DEFAULT_MIN_TARGET_BULK_MICRO_BATCH_SIZE = 1; - -+ // readManyByPartitionKey: max number of PK values per query per physical partition -+ private static final String READ_MANY_BY_PK_MAX_BATCH_SIZE = "COSMOS.READ_MANY_BY_PK_MAX_BATCH_SIZE"; -+ private static final String READ_MANY_BY_PK_MAX_BATCH_SIZE_VARIABLE = "COSMOS_READ_MANY_BY_PK_MAX_BATCH_SIZE"; -+ private static final int DEFAULT_READ_MANY_BY_PK_MAX_BATCH_SIZE = 1000; -+ - public static final String MAX_BULK_MICRO_BATCH_CONCURRENCY = "COSMOS.MAX_BULK_MICRO_BATCH_CONCURRENCY"; - public static final String MAX_BULK_MICRO_BATCH_CONCURRENCY_VARIABLE = "COSMOS_MAX_BULK_MICRO_BATCH_CONCURRENCY"; - public static final int DEFAULT_MAX_BULK_MICRO_BATCH_CONCURRENCY = 1; -@@ -816,6 +821,20 @@ public class Configs { - return DEFAULT_MIN_TARGET_BULK_MICRO_BATCH_SIZE; - } - -+ public static int getReadManyByPkMaxBatchSize() { -+ String valueFromSystemProperty = System.getProperty(READ_MANY_BY_PK_MAX_BATCH_SIZE); -+ if (valueFromSystemProperty != null && !valueFromSystemProperty.isEmpty()) { -+ return Math.max(1, Integer.parseInt(valueFromSystemProperty)); -+ } -+ -+ String valueFromEnvVariable = System.getenv(READ_MANY_BY_PK_MAX_BATCH_SIZE_VARIABLE); -+ if (valueFromEnvVariable != null && !valueFromEnvVariable.isEmpty()) { -+ return Math.max(1, Integer.parseInt(valueFromEnvVariable)); -+ } -+ -+ return DEFAULT_READ_MANY_BY_PK_MAX_BATCH_SIZE; -+ } -+ - public static int getMaxBulkMicroBatchConcurrency() { - String valueFromSystemProperty = System.getProperty(MAX_BULK_MICRO_BATCH_CONCURRENCY); - if (valueFromSystemProperty != null && !valueFromSystemProperty.isEmpty()) { -diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelper.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelper.java -new file mode 100644 -index 00000000000..6d6cd084e01 ---- /dev/null -+++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelper.java -@@ -0,0 +1,263 @@ -+// Copyright (c) Microsoft Corporation. All rights reserved. -+// Licensed under the MIT License. -+package com.azure.cosmos.implementation; -+ -+import com.azure.cosmos.BridgeInternal; -+import com.azure.cosmos.implementation.routing.PartitionKeyInternal; -+import com.azure.cosmos.models.PartitionKey; -+import com.azure.cosmos.models.PartitionKeyDefinition; -+import com.azure.cosmos.models.PartitionKind; -+import com.azure.cosmos.models.SqlParameter; -+import com.azure.cosmos.models.SqlQuerySpec; -+ -+import java.util.ArrayList; -+import java.util.List; -+ -+/** -+ * Helper for constructing SqlQuerySpec instances for readManyByPartitionKey operations. -+ * This class is not intended to be used directly by end-users. -+ */ -+public class ReadManyByPartitionKeyQueryHelper { -+ -+ private static final String DEFAULT_TABLE_ALIAS = "c"; -+ // Internal parameter prefix ΓÇö uses double-underscore to avoid collisions with user-provided parameters -+ private static final String PK_PARAM_PREFIX = "@__rmPk_"; -+ -+ public static SqlQuerySpec createReadManyByPkQuerySpec( -+ String baseQueryText, -+ List baseParameters, -+ List pkValues, -+ List partitionKeySelectors, -+ PartitionKeyDefinition pkDefinition) { -+ -+ // Extract the table alias from the FROM clause (e.g. "FROM x" ΓåÆ "x", "FROM c" ΓåÆ "c") -+ String tableAlias = extractTableAlias(baseQueryText); -+ -+ StringBuilder pkFilter = new StringBuilder(); -+ List parameters = new ArrayList<>(baseParameters); -+ int paramCount = 0; -+ -+ boolean isSinglePathPk = partitionKeySelectors.size() == 1; -+ -+ if (isSinglePathPk && pkDefinition.getKind() != PartitionKind.MULTI_HASH) { -+ // Single PK path ΓÇö use IN clause for normal values, OR NOT IS_DEFINED for NONE -+ // First, separate NONE PKs from normal PKs -+ boolean hasNone = false; -+ List normalPkValues = new ArrayList<>(); -+ for (PartitionKey pk : pkValues) { -+ PartitionKeyInternal pkInternal = BridgeInternal.getPartitionKeyInternal(pk); -+ if (pkInternal.getComponents() == null) { -+ hasNone = true; -+ } else { -+ normalPkValues.add(pk); -+ } -+ } -+ -+ pkFilter.append(" "); -+ boolean hasNormalValues = !normalPkValues.isEmpty(); -+ if (hasNormalValues && hasNone) { -+ pkFilter.append("("); -+ } -+ if (hasNormalValues) { -+ pkFilter.append(tableAlias); -+ pkFilter.append(partitionKeySelectors.get(0)); -+ pkFilter.append(" IN ( "); -+ for (int i = 0; i < normalPkValues.size(); i++) { -+ PartitionKeyInternal pkInternal = BridgeInternal.getPartitionKeyInternal(normalPkValues.get(i)); -+ Object[] pkComponents = pkInternal.toObjectArray(); -+ String pkParamName = PK_PARAM_PREFIX + paramCount; -+ parameters.add(new SqlParameter(pkParamName, pkComponents[0])); -+ paramCount++; -+ -+ pkFilter.append(pkParamName); -+ if (i < normalPkValues.size() - 1) { -+ pkFilter.append(", "); -+ } -+ } -+ pkFilter.append(" )"); -+ } -+ if (hasNone) { -+ if (hasNormalValues) { -+ pkFilter.append(" OR "); -+ } -+ pkFilter.append("NOT IS_DEFINED("); -+ pkFilter.append(tableAlias); -+ pkFilter.append(partitionKeySelectors.get(0)); -+ pkFilter.append(")"); -+ } -+ if (hasNormalValues && hasNone) { -+ pkFilter.append(")"); -+ } -+ } else { -+ // Multiple PK paths (HPK) or MULTI_HASH ΓÇö use OR of AND clauses -+ pkFilter.append(" "); -+ for (int i = 0; i < pkValues.size(); i++) { -+ PartitionKeyInternal pkInternal = BridgeInternal.getPartitionKeyInternal(pkValues.get(i)); -+ Object[] pkComponents = pkInternal.toObjectArray(); -+ -+ // PartitionKey.NONE ΓÇö generate NOT IS_DEFINED for all PK paths -+ if (pkComponents == null) { -+ pkFilter.append("("); -+ for (int j = 0; j < partitionKeySelectors.size(); j++) { -+ if (j > 0) { -+ pkFilter.append(" AND "); -+ } -+ pkFilter.append("NOT IS_DEFINED("); -+ pkFilter.append(tableAlias); -+ pkFilter.append(partitionKeySelectors.get(j)); -+ pkFilter.append(")"); -+ } -+ pkFilter.append(")"); -+ } else { -+ pkFilter.append("("); -+ for (int j = 0; j < pkComponents.length; j++) { -+ String pkParamName = PK_PARAM_PREFIX + paramCount; -+ parameters.add(new SqlParameter(pkParamName, pkComponents[j])); -+ paramCount++; -+ -+ if (j > 0) { -+ pkFilter.append(" AND "); -+ } -+ pkFilter.append(tableAlias); -+ pkFilter.append(partitionKeySelectors.get(j)); -+ pkFilter.append(" = "); -+ pkFilter.append(pkParamName); -+ } -+ pkFilter.append(")"); -+ } -+ -+ if (i < pkValues.size() - 1) { -+ pkFilter.append(" OR "); -+ } -+ } -+ } -+ -+ // Compose final query: handle existing WHERE clause in base query -+ String finalQuery; -+ int whereIndex = findTopLevelWhereIndex(baseQueryText); -+ if (whereIndex >= 0) { -+ // Base query has WHERE ΓÇö AND our PK filter -+ String beforeWhere = baseQueryText.substring(0, whereIndex); -+ String afterWhere = baseQueryText.substring(whereIndex + 5); // skip "WHERE" -+ finalQuery = beforeWhere + "WHERE (" + afterWhere.trim() + ") AND (" + pkFilter.toString().trim() + ")"; -+ } else { -+ // No WHERE ΓÇö add one -+ finalQuery = baseQueryText + " WHERE" + pkFilter.toString(); -+ } -+ -+ return new SqlQuerySpec(finalQuery, parameters); -+ } -+ -+ /** -+ * Extracts the table/collection alias from a SQL query's FROM clause. -+ * Handles: "SELECT * FROM c", "SELECT x.id FROM x WHERE ...", "SELECT * FROM root r", etc. -+ * Returns the alias used after FROM (last token before WHERE or end of FROM clause). -+ */ -+ static String extractTableAlias(String queryText) { -+ String upper = queryText.toUpperCase(); -+ int fromIndex = findTopLevelKeywordIndex(upper, "FROM"); -+ if (fromIndex < 0) { -+ return DEFAULT_TABLE_ALIAS; -+ } -+ -+ // Start scanning after "FROM" -+ int afterFrom = fromIndex + 4; -+ // Skip whitespace -+ while (afterFrom < queryText.length() && Character.isWhitespace(queryText.charAt(afterFrom))) { -+ afterFrom++; -+ } -+ -+ // Collect the container name token (could be "root", "c", etc.) -+ int tokenStart = afterFrom; -+ while (afterFrom < queryText.length() -+ && !Character.isWhitespace(queryText.charAt(afterFrom)) -+ && queryText.charAt(afterFrom) != '(' -+ && queryText.charAt(afterFrom) != ')') { -+ afterFrom++; -+ } -+ String containerName = queryText.substring(tokenStart, afterFrom); -+ -+ // Skip whitespace after container name -+ while (afterFrom < queryText.length() && Character.isWhitespace(queryText.charAt(afterFrom))) { -+ afterFrom++; -+ } -+ -+ // Check if there's an alias after the container name (before WHERE or end) -+ if (afterFrom < queryText.length()) { -+ char nextChar = Character.toUpperCase(queryText.charAt(afterFrom)); -+ // If the next token is a keyword (WHERE, ORDER, GROUP, JOIN, OFFSET, LIMIT, HAVING) or end, containerName IS the alias -+ if (nextChar == 'W' || nextChar == 'O' || nextChar == 'G' || nextChar == 'J' -+ || nextChar == 'L' || nextChar == 'H') { -+ // Check if it's actually a keyword -+ String remaining = upper.substring(afterFrom); -+ if (remaining.startsWith("WHERE") || remaining.startsWith("ORDER") -+ || remaining.startsWith("GROUP") || remaining.startsWith("JOIN") -+ || remaining.startsWith("OFFSET") || remaining.startsWith("LIMIT") -+ || remaining.startsWith("HAVING")) { -+ return containerName; -+ } -+ } -+ // Otherwise the next token is the alias ("FROM root r" ΓåÆ alias is "r") -+ int aliasStart = afterFrom; -+ while (afterFrom < queryText.length() -+ && !Character.isWhitespace(queryText.charAt(afterFrom)) -+ && queryText.charAt(afterFrom) != '(' -+ && queryText.charAt(afterFrom) != ')') { -+ afterFrom++; -+ } -+ if (afterFrom > aliasStart) { -+ return queryText.substring(aliasStart, afterFrom); -+ } -+ } -+ -+ return containerName; -+ } -+ -+ /** -+ * Finds the index of a top-level SQL keyword in the query text (case-insensitive), -+ * ignoring occurrences inside parentheses or string literals. -+ */ -+ static int findTopLevelKeywordIndex(String queryText, String keyword) { -+ String queryTextUpper = queryText.toUpperCase(); -+ String keywordUpper = keyword.toUpperCase(); -+ int depth = 0; -+ int keyLen = keywordUpper.length(); -+ for (int i = 0; i <= queryTextUpper.length() - keyLen; i++) { -+ char ch = queryTextUpper.charAt(i); -+ // Skip string literals enclosed in single quotes (handle '' escape) -+ if (queryText.charAt(i) == '\'') { -+ i++; -+ while (i < queryText.length()) { -+ if (queryText.charAt(i) == '\'') { -+ if (i + 1 < queryText.length() && queryText.charAt(i + 1) == '\'') { -+ i += 2; // escaped quote ΓÇö skip both -+ continue; -+ } -+ break; // end of string literal -+ } -+ i++; -+ } -+ continue; -+ } -+ if (ch == '(') { -+ depth++; -+ } else if (ch == ')') { -+ depth--; -+ } else if (depth == 0 && ch == keywordUpper.charAt(0) -+ && queryTextUpper.startsWith(keywordUpper, i) -+ && (i == 0 || !Character.isLetterOrDigit(queryTextUpper.charAt(i - 1))) -+ && (i + keyLen >= queryTextUpper.length() || !Character.isLetterOrDigit(queryTextUpper.charAt(i + keyLen)))) { -+ return i; -+ } -+ } -+ return -1; -+ } -+ -+ /** -+ * Finds the index of the top-level WHERE keyword in the query text, -+ * ignoring WHERE that appears inside parentheses (subqueries). -+ */ -+ public static int findTopLevelWhereIndex(String queryTextUpper) { -+ return findTopLevelKeywordIndex(queryTextUpper, "WHERE"); -+ } -+} -diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java -index 11121bca033..c70dedaa1f2 100644 ---- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java -+++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java -@@ -4365,6 +4365,298 @@ public class RxDocumentClientImpl implements AsyncDocumentClient, IAuthorization - ); - } - -+ @Override -+ public Flux> readManyByPartitionKey( -+ List partitionKeys, -+ SqlQuerySpec customQuery, -+ String collectionLink, -+ QueryFeedOperationState state, -+ Class klass) { -+ -+ checkNotNull(partitionKeys, "Argument 'partitionKeys' must not be null."); -+ checkArgument(!partitionKeys.isEmpty(), "Argument 'partitionKeys' must not be empty."); -+ -+ final ScopedDiagnosticsFactory diagnosticsFactory = new ScopedDiagnosticsFactory(this, true); -+ state.registerDiagnosticsFactory( -+ () -> {}, // we never want to reset in readManyByPartitionKey -+ (ctx) -> diagnosticsFactory.merge(ctx) -+ ); -+ -+ StaleResourceRetryPolicy staleResourceRetryPolicy = new StaleResourceRetryPolicy( -+ this.collectionCache, -+ null, -+ collectionLink, -+ queryOptionsAccessor().getProperties(state.getQueryOptions()), -+ queryOptionsAccessor().getHeaders(state.getQueryOptions()), -+ this.sessionContainer, -+ diagnosticsFactory, -+ ResourceType.Document -+ ); -+ -+ return ObservableHelper -+ .fluxInlineIfPossibleAsObs( -+ () -> readManyByPartitionKey( -+ partitionKeys, customQuery, collectionLink, state, diagnosticsFactory, klass), -+ staleResourceRetryPolicy -+ ) -+ .onErrorMap(throwable -> { -+ if (throwable instanceof CosmosException) { -+ CosmosException cosmosException = (CosmosException) throwable; -+ CosmosDiagnostics diagnostics = cosmosException.getDiagnostics(); -+ if (diagnostics != null) { -+ state.mergeDiagnosticsContext(); -+ CosmosDiagnosticsContext ctx = state.getDiagnosticsContextSnapshot(); -+ if (ctx != null) { -+ ctxAccessor().recordOperation( -+ ctx, -+ cosmosException.getStatusCode(), -+ cosmosException.getSubStatusCode(), -+ 0, -+ cosmosException.getRequestCharge(), -+ diagnostics, -+ throwable -+ ); -+ diagAccessor() -+ .setDiagnosticsContext( -+ diagnostics, -+ state.getDiagnosticsContextSnapshot()); -+ } -+ } -+ -+ return cosmosException; -+ } -+ -+ return throwable; -+ }); -+ } -+ -+ private Flux> readManyByPartitionKey( -+ List partitionKeys, -+ SqlQuerySpec customQuery, -+ String collectionLink, -+ QueryFeedOperationState state, -+ ScopedDiagnosticsFactory diagnosticsFactory, -+ Class klass) { -+ -+ String resourceLink = parentResourceLinkToQueryLink(collectionLink, ResourceType.Document); -+ RxDocumentServiceRequest request = RxDocumentServiceRequest.create(diagnosticsFactory, -+ OperationType.Query, -+ ResourceType.Document, -+ collectionLink, null -+ ); -+ -+ Mono> collectionObs = -+ collectionCache.resolveCollectionAsync(null, request); -+ -+ return collectionObs -+ .flatMapMany(documentCollectionResourceResponse -> { -+ final DocumentCollection collection = documentCollectionResourceResponse.v; -+ if (collection == null) { -+ return Flux.error(new IllegalStateException("Collection cannot be null")); -+ } -+ -+ final PartitionKeyDefinition pkDefinition = collection.getPartitionKey(); -+ -+ Mono> valueHolderMono = partitionKeyRangeCache -+ .tryLookupAsync( -+ BridgeInternal.getMetaDataDiagnosticContext(request.requestContext.cosmosDiagnostics), -+ collection.getResourceId(), -+ null, -+ null); -+ -+ // Validate custom query if provided -+ Mono queryValidationMono; -+ if (customQuery != null) { -+ queryValidationMono = validateCustomQueryForReadManyByPartitionKey( -+ customQuery, resourceLink, state.getQueryOptions()); -+ } else { -+ queryValidationMono = Mono.empty(); -+ } -+ -+ return valueHolderMono -+ .delayUntil(ignored -> queryValidationMono) -+ .flatMapMany(routingMapHolder -> { -+ CollectionRoutingMap routingMap = routingMapHolder.v; -+ if (routingMap == null) { -+ return Flux.error(new IllegalStateException("Failed to get routing map.")); -+ } -+ -+ Map> partitionRangePkMap = -+ groupPartitionKeysByPhysicalPartition(partitionKeys, pkDefinition, routingMap); -+ -+ List partitionKeySelectors = createPkSelectors(pkDefinition); -+ -+ String baseQueryText; -+ List baseParameters; -+ if (customQuery != null) { -+ baseQueryText = customQuery.getQueryText(); -+ baseParameters = customQuery.getParameters() != null -+ ? new ArrayList<>(customQuery.getParameters()) -+ : new ArrayList<>(); -+ } else { -+ baseQueryText = "SELECT * FROM c"; -+ baseParameters = new ArrayList<>(); -+ } -+ -+ // Build per-physical-partition batched queries. -+ // Each physical partition may have many PKs ΓÇö split into batches -+ // to avoid oversized SQL queries. Batch size is configurable via -+ // system property COSMOS.READ_MANY_BY_PK_MAX_BATCH_SIZE (default 1000). -+ int maxPksPerPartitionQuery = Configs.getReadManyByPkMaxBatchSize(); -+ -+ // Build batches per partition as a list of lists (one inner list per partition). -+ // Then interleave in round-robin order so that concurrent execution -+ // prefers different physical partitions over multiple batches of the same partition. -+ List>> batchesPerPartition = new ArrayList<>(); -+ int maxBatchesPerPartition = 0; -+ -+ for (Map.Entry> entry : partitionRangePkMap.entrySet()) { -+ List allPks = entry.getValue(); -+ if (allPks.isEmpty()) { -+ continue; -+ } -+ List> partitionBatches = new ArrayList<>(); -+ for (int i = 0; i < allPks.size(); i += maxPksPerPartitionQuery) { -+ List batch = allPks.subList( -+ i, Math.min(i + maxPksPerPartitionQuery, allPks.size())); -+ SqlQuerySpec querySpec = ReadManyByPartitionKeyQueryHelper -+ .createReadManyByPkQuerySpec( -+ baseQueryText, baseParameters, batch, -+ partitionKeySelectors, pkDefinition); -+ partitionBatches.add( -+ Collections.singletonMap(entry.getKey(), querySpec)); -+ } -+ batchesPerPartition.add(partitionBatches); -+ maxBatchesPerPartition = Math.max(maxBatchesPerPartition, partitionBatches.size()); -+ } -+ -+ if (batchesPerPartition.isEmpty()) { -+ return Flux.empty(); -+ } -+ -+ // Round-robin interleave: [batch0-p1, batch0-p2, ..., batch0-pN, batch1-p1, batch1-p2, ...] -+ // This ensures that with bounded concurrency, different partitions are -+ // preferred over sequential batches of the same partition. -+ List> interleavedBatches = new ArrayList<>(); -+ for (int batchIdx = 0; batchIdx < maxBatchesPerPartition; batchIdx++) { -+ for (List> partitionBatches : batchesPerPartition) { -+ if (batchIdx < partitionBatches.size()) { -+ interleavedBatches.add(partitionBatches.get(batchIdx)); -+ } -+ } -+ } -+ -+ // Execute all batches with bounded concurrency. -+ List>> queryFluxes = interleavedBatches -+ .stream() -+ .map(batchMap -> queryForReadMany( -+ diagnosticsFactory, -+ resourceLink, -+ new SqlQuerySpec(DUMMY_SQL_QUERY), -+ state.getQueryOptions(), -+ klass, -+ ResourceType.Document, -+ collection, -+ Collections.unmodifiableMap(batchMap))) -+ .collect(Collectors.toList()); -+ -+ int fluxConcurrency = Math.min(queryFluxes.size(), -+ Math.max(Configs.getCPUCnt(), 1)); -+ -+ return Flux.merge(Flux.fromIterable(queryFluxes), fluxConcurrency, 1); -+ }); -+ }); -+ } -+ -+ private Mono validateCustomQueryForReadManyByPartitionKey( -+ SqlQuerySpec customQuery, -+ String resourceLink, -+ CosmosQueryRequestOptions queryRequestOptions) { -+ -+ IDocumentQueryClient queryClient = documentQueryClientImpl( -+ RxDocumentClientImpl.this, getOperationContextAndListenerTuple(queryRequestOptions)); -+ -+ return DocumentQueryExecutionContextFactory -+ .fetchQueryPlanForValidation(this, queryClient, customQuery, resourceLink, queryRequestOptions) -+ .flatMap(queryPlan -> { -+ QueryInfo queryInfo = queryPlan.getQueryInfo(); -+ -+ if (queryInfo.hasGroupBy()) { -+ return Mono.error(new IllegalArgumentException( -+ "Custom query for readMany by partition key must not contain GROUP BY.")); -+ } -+ if (queryInfo.hasAggregates()) { -+ return Mono.error(new IllegalArgumentException( -+ "Custom query for readMany by partition key must not contain aggregates.")); -+ } -+ if (queryInfo.hasOrderBy()) { -+ return Mono.error(new IllegalArgumentException( -+ "Custom query for readMany by partition key must not contain ORDER BY.")); -+ } -+ if (queryInfo.hasDistinct()) { -+ return Mono.error(new IllegalArgumentException( -+ "Custom query for readMany by partition key must not contain DISTINCT.")); -+ } -+ if (queryInfo.hasDCount()) { -+ return Mono.error(new IllegalArgumentException( -+ "Custom query for readMany by partition key must not contain DCOUNT.")); -+ } -+ if (queryInfo.hasNonStreamingOrderBy()) { -+ return Mono.error(new IllegalArgumentException( -+ "Custom query for readMany by partition key must not contain non-streaming ORDER BY.")); -+ } -+ if (queryPlan.hasHybridSearchQueryInfo()) { -+ return Mono.error(new IllegalArgumentException( -+ "Custom query for readMany by partition key must not contain hybrid/vector/full-text search.")); -+ } -+ -+ return Mono.empty(); -+ }); -+ } -+ -+ private Map> groupPartitionKeysByPhysicalPartition( -+ List partitionKeys, -+ PartitionKeyDefinition pkDefinition, -+ CollectionRoutingMap routingMap) { -+ -+ Map> partitionRangePkMap = new HashMap<>(); -+ -+ for (PartitionKey pk : partitionKeys) { -+ PartitionKeyInternal pkInternal = BridgeInternal.getPartitionKeyInternal(pk); -+ -+ // PartitionKey.NONE wraps NonePartitionKey which has components = null. -+ // For routing purposes, treat NONE as UndefinedPartitionKey ΓÇö documents ingested -+ // without a partition key path are stored with the undefined EPK. -+ PartitionKeyInternal effectivePkInternal = pkInternal.getComponents() == null -+ ? PartitionKeyInternal.UndefinedPartitionKey -+ : pkInternal; -+ -+ int componentCount = effectivePkInternal.getComponents().size(); -+ int definedPathCount = pkDefinition.getPaths().size(); -+ -+ List targetRanges; -+ -+ if (pkDefinition.getKind() == PartitionKind.MULTI_HASH && componentCount < definedPathCount) { -+ // Partial HPK ΓÇö compute EPK prefix range and find all overlapping physical partitions -+ Range epkRange = PartitionKeyInternalHelper.getEPKRangeForPrefixPartitionKey( -+ effectivePkInternal, pkDefinition); -+ targetRanges = routingMap.getOverlappingRanges(epkRange); -+ } else { -+ // Full PK ΓÇö maps to exactly one physical partition -+ String effectivePartitionKeyString = PartitionKeyInternalHelper -+ .getEffectivePartitionKeyString(effectivePkInternal, pkDefinition); -+ PartitionKeyRange range = routingMap.getRangeByEffectivePartitionKey(effectivePartitionKeyString); -+ targetRanges = Collections.singletonList(range); -+ } -+ -+ for (PartitionKeyRange range : targetRanges) { -+ partitionRangePkMap.computeIfAbsent(range, k -> new ArrayList<>()).add(pk); -+ } -+ } -+ -+ return partitionRangePkMap; -+ } -+ - private Map getRangeQueryMap( - Map> partitionRangeItemKeyMap, - PartitionKeyDefinition partitionKeyDefinition) { -diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/DocumentQueryExecutionContextFactory.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/DocumentQueryExecutionContextFactory.java -index e62d8ed3d75..d8f9614343c 100644 ---- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/DocumentQueryExecutionContextFactory.java -+++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/DocumentQueryExecutionContextFactory.java -@@ -318,6 +318,17 @@ public class DocumentQueryExecutionContextFactory { - return feedRanges; - } - -+ public static Mono fetchQueryPlanForValidation( -+ DiagnosticsClientContext diagnosticsClientContext, -+ IDocumentQueryClient queryClient, -+ SqlQuerySpec sqlQuerySpec, -+ String resourceLink, -+ CosmosQueryRequestOptions queryRequestOptions) { -+ -+ return QueryPlanRetriever.getQueryPlanThroughGatewayAsync( -+ diagnosticsClientContext, queryClient, sqlQuerySpec, resourceLink, queryRequestOptions); -+ } -+ - public static Flux> createDocumentQueryExecutionContextAsync( - DiagnosticsClientContext diagnosticsClientContext, - IDocumentQueryClient client, -diff --git a/sdk/cosmos/cspell.yaml b/sdk/cosmos/cspell.yaml -new file mode 100644 -index 00000000000..94a4002c2c9 ---- /dev/null -+++ b/sdk/cosmos/cspell.yaml -@@ -0,0 +1,6 @@ -+import: -+ - ../../.vscode/cspell.json -+overrides: -+ - filename: "**/sdk/cosmos/*" -+ words: -+ - DCOUNT -diff --git a/sdk/cosmos/docs/readManyByPartitionKey-design.md b/sdk/cosmos/docs/readManyByPartitionKey-design.md -new file mode 100644 -index 00000000000..95d7624f0c8 ---- /dev/null -+++ b/sdk/cosmos/docs/readManyByPartitionKey-design.md -@@ -0,0 +1,169 @@ -+# readManyByPartitionKey ΓÇö Design & Implementation -+ -+## Overview -+ -+New `readManyByPartitionKey` methods on `CosmosAsyncContainer` / `CosmosContainer` that accept a -+`List` (without item-id). The SDK splits the PK values by physical -+partition, generates batched streaming queries per physical partition, and returns results as -+`CosmosPagedFlux` / `CosmosPagedIterable`. -+ -+An optional `SqlQuerySpec` parameter lets callers supply a custom query for projections -+and additional filters. The SDK appends the auto-generated PK WHERE clause to it. -+ -+## Decisions -+ -+| Topic | Decision | -+|---|---| -+| API name | `readManyByPartitionKey` ΓÇö distinct name to avoid ambiguity with existing `readMany(List)` | -+| Return type | `CosmosPagedFlux` (async) / `CosmosPagedIterable` (sync) | -+| Custom query format | `SqlQuerySpec` ΓÇö full query with parameters; SDK ANDs the PK filter | -+| Partial HPK | Supported from the start; prefix PKs fan out via `getOverlappingRanges` | -+| PK deduplication | Done at Spark layer only, not in the SDK | -+| Spark UDF | New `GetCosmosPartitionKeyValue` UDF | -+| Custom query validation | Gateway query plan; reject aggregates/ORDER BY/DISTINCT/GROUP BY/DCount/non-streaming ORDER BY/vector/fulltext | -+| PK list size | No hard upper-bound enforced; SDK batches internally per physical partition (default 1000 PKs per batch, configurable via `COSMOS.READ_MANY_BY_PK_MAX_BATCH_SIZE`) | -+| Eager validation | Null and empty PK list rejected eagerly (not lazily in reactive chain) | -+| Telemetry | Separate span name `readManyByPartitionKeyItems.` (distinct from existing `readManyItems`) | -+| Query construction | Table alias auto-detected from FROM clause; string literals and subqueries handled correctly | -+ -+## Phase 1 ΓÇö SDK Core (`azure-cosmos`) -+ -+### Step 1: New public overloads in CosmosAsyncContainer -+ -+```java -+ CosmosPagedFlux readManyByPartitionKey(List partitionKeys, Class classType) -+ CosmosPagedFlux readManyByPartitionKey(List partitionKeys, -+ CosmosReadManyRequestOptions requestOptions, -+ Class classType) -+ CosmosPagedFlux readManyByPartitionKey(List partitionKeys, -+ SqlQuerySpec customQuery, -+ CosmosReadManyRequestOptions requestOptions, -+ Class classType) -+``` -+ -+All delegate to a private `readManyByPartitionKeyInternalFunc(...)`. -+ -+**Eager validation:** The 4-arg method validates `partitionKeys` is non-null and non-empty before constructing the reactive pipeline, throwing `IllegalArgumentException` synchronously. -+ -+### Step 2: Sync wrappers in CosmosContainer -+ -+Same signatures returning `CosmosPagedIterable`, delegating to the async container. -+ -+### Step 3: Internal orchestration (RxDocumentClientImpl) -+ -+1. Resolve collection metadata + PK definition from cache. -+2. Fetch routing map from `partitionKeyRangeCache` **in parallel with** custom query validation (Step 4). -+3. For each `PartitionKey`: -+ - Compute effective partition key (EPK). -+ - Full PK ΓåÆ `getRangeByEffectivePartitionKey()` (single range). -+ - Partial HPK ΓåÆ compute EPK prefix range ΓåÆ `getOverlappingRanges()` (multiple ranges). -+ **Note:** partial HPK intentionally fans out to multiple physical partitions. -+4. Group PK values by `PartitionKeyRange`. -+5. Per physical partition ΓåÆ split PKs into batches of `maxPksPerPartitionQuery` (configurable, default 1000). -+6. Per batch ΓåÆ build `SqlQuerySpec` with PK WHERE clause (Step 5). -+7. Interleave batches across physical partitions in round-robin order so that bounded concurrency prefers different physical partitions over sequential batches of the same partition. -+8. Execute queries via `queryForReadMany()` with bounded concurrency (`Math.min(batchCount, cpuCount)`). -+9. Return results as `CosmosPagedFlux`. -+ -+### Step 4: Custom query validation -+ -+One-time call per invocation (existing query plan caching applies). Runs **in parallel** with routing map lookup to minimize latency: -+ -+- `QueryPlanRetriever.getQueryPlanThroughGatewayAsync()` for the user query. -+- Reject (`IllegalArgumentException`) if: -+ - `queryInfo.hasGroupBy()` ΓÇö checked first (takes precedence over aggregates since `hasAggregates()` also returns true for GROUP BY queries) -+ - `queryInfo.hasAggregates()` -+ - `queryInfo.hasOrderBy()` -+ - `queryInfo.hasDistinct()` -+ - `queryInfo.hasDCount()` -+ - `queryInfo.hasNonStreamingOrderBy()` -+ - `partitionedQueryExecutionInfo.hasHybridSearchQueryInfo()` -+ -+### Step 5: Query construction -+ -+Query construction is implemented in `ReadManyByPartitionKeyQueryHelper`. The helper: -+- Extracts the table alias from the FROM clause (handles `FROM c`, `FROM root r`, `FROM x WHERE ...`) -+- Handles string literals in queries (parens/keywords inside `'...'` are correctly skipped) -+- Recognizes SQL keywords: WHERE, ORDER, GROUP, JOIN, OFFSET, LIMIT, HAVING -+- Uses parameterized queries (`@__rmPk_` prefix) to prevent SQL injection -+ -+**Single PK (HASH):** -+```sql -+{baseQuery} WHERE {alias}["{pkPath}"] IN (@__rmPk_0, @__rmPk_1, @__rmPk_2) -+``` -+ -+**Full HPK (MULTI_HASH):** -+```sql -+{baseQuery} WHERE ({alias}["{path1}"] = @__rmPk_0 AND {alias}["{path2}"] = @__rmPk_1) -+ OR ({alias}["{path1}"] = @__rmPk_2 AND {alias}["{path2}"] = @__rmPk_3) -+``` -+ -+**Partial HPK (prefix-only):** -+```sql -+{baseQuery} WHERE ({alias}["{path1}"] = @__rmPk_0) -+ OR ({alias}["{path1}"] = @__rmPk_1) -+``` -+ -+If the base query already has a WHERE clause: -+```sql -+{selectAndFrom} WHERE ({existingWhere}) AND ({pkFilter}) -+``` -+ -+### Step 6: Interface wiring -+ -+New method `readManyByPartitionKey` added directly to `AsyncDocumentClient` interface, implemented in `RxDocumentClientImpl`. New `fetchQueryPlanForValidation` static method added to `DocumentQueryExecutionContextFactory` for custom query validation. -+ -+### Step 7: Configuration -+ -+New configurable batch size via system property `COSMOS.READ_MANY_BY_PK_MAX_BATCH_SIZE` or environment variable `COSMOS_READ_MANY_BY_PK_MAX_BATCH_SIZE` (default: 1000, minimum: 1). Follows existing `Configs` patterns. -+ -+## Phase 2 ΓÇö Spark Connector (`azure-cosmos-spark_3`) -+ -+### Step 8: New UDF ΓÇö `GetCosmosPartitionKeyValue` -+ -+- Input: partition key value (single value or Seq for hierarchical PKs). -+- Output: serialized PK string in format `pk([...json...])`. -+- **Null handling:** Throws on null input (Scala convention; callers should filter nulls upstream). -+ -+### Step 9: PK-only serialization helper -+ -+`CosmosPartitionKeyHelper`: -+- `getCosmosPartitionKeyValueString(pkValues: List[Object]): String` ΓÇö serialize to `pk([...])` format. -+- `tryParsePartitionKey(serialized: String): Option[PartitionKey]` ΓÇö deserialize; returns `None` for malformed input including invalid JSON (wrapped in `scala.util.Try`). -+ -+### Step 10: `CosmosItemsDataSource.readManyByPartitionKey` -+ -+Static entry points that accept a DataFrame and Cosmos config. PK extraction supports two modes: -+1. **UDF-produced column**: DataFrame contains `_partitionKeyIdentity` column (from `GetCosmosPartitionKeyValue` UDF). -+2. **Schema-matched columns**: DataFrame columns match the container's PK paths. -+ -+Falls back with `IllegalArgumentException` if neither mode is possible. -+ -+### Step 11: `CosmosReadManyByPartitionKeyReader` -+ -+Orchestrator that resolves schema, initializes and broadcasts client state to executors, then maps each Spark partition to an `ItemsPartitionReaderWithReadManyByPartitionKey`. -+ -+### Step 12: `ItemsPartitionReaderWithReadManyByPartitionKey` -+ -+Spark `PartitionReader[InternalRow]` that: -+- Deduplicates PKs via `LinkedHashMap` (by PK string representation). -+- Passes the pre-built `CosmosReadManyRequestOptions` (with throughput control, diagnostics, custom serializer) to the SDK. -+- Uses `TransientIOErrorsRetryingIterator` for retry handling. -+- Short-circuits empty PK lists to avoid SDK rejection. -+ -+## Phase 3 ΓÇö Testing -+ -+### Unit tests -+- Query construction: single PK, HPK full/partial, custom query composition, table alias detection. -+- Query plan rejection: aggregates, ORDER BY, DISTINCT, GROUP BY (with and without aggregates), DCOUNT. -+- String literal handling: WHERE/parentheses inside string constants. -+- Keyword detection: WHERE, ORDER, GROUP, JOIN, OFFSET, LIMIT, HAVING. -+- PK serialization/deserialization roundtrip (including malformed JSON handling). -+- `findTopLevelWhereIndex` edge cases: subqueries, string literals, case insensitivity. -+ -+### Integration tests -+- End-to-end SDK: single PK basic, projections, filters, empty results, HPK full/partial, request options propagation. -+- Batch size validation: temporarily lowered batch size to exercise batching/interleaving logic. -+- Null/empty PK list rejection (eager validation). -+- Spark connector: `ItemsPartitionReaderWithReadManyByPartitionKey` with known PK values and non-existent PKs. -+- `CosmosPartitionKeyHelper`: single/HPK roundtrip, case insensitivity, malformed input. - diff --git a/sdk/cosmos/.temp/branch-users_fabianm_readManyByPK-vs-origin_main-stat.txt b/sdk/cosmos/.temp/branch-users_fabianm_readManyByPK-vs-origin_main-stat.txt deleted file mode 100644 index 360eb15accb5..000000000000 --- a/sdk/cosmos/.temp/branch-users_fabianm_readManyByPK-vs-origin_main-stat.txt +++ /dev/null @@ -1,76 +0,0 @@ -===== PR #LOCAL-users_fabianm_readManyByPK ===== -Title: Branch comparison users/fabianm/readManyByPK vs origin/main -Author: Fabian Meiswinkel -Status: DIVERGED (ahead 30, behind 4) -Branch: users/fabianm/readManyByPK -> origin/main -Head SHA: 93957f3a8442d730fe67fbc379ef5399f46f5665 -Merge Base: 20313f79ba8dd0dfa97862d0c31dd4b2e44ee671 -URL: N/A (local branch comparison) - ---- Description --- -Adds readManyByPartitionKey API (sync+async) and Spark connector support for PK-only reads, with query-plan-based validation ---- End Description --- - -===== Commits in PR ===== -9770833eb59 Adding readManyByPartitionKey API -ac287bcdf00 Merge branch 'main' of https://github.com/Azure/azure-sdk-for-java into users/fabianm/readManyByPK -9a5b3e96e7e Update sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/ItemsPartitionReaderWithReadManyByPartitionKey.scala -a8720c3c9f2 Update sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/ItemsPartitionReaderWithReadManyByPartitionKey.scala -d499da76fb4 Update sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java -c3c542a33a7 Update sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelper.java -4416354e03e ┬┤Fixing code review comments -3ab3f0d64f5 Merge branch 'main' of https://github.com/Azure/azure-sdk-for-java into users/fabianm/readManyByPK -588a7550c54 Update CosmosAsyncContainer.java -8c5cdb47b31 Merge branch 'main' into users/fabianm/readManyByPK -f5485527a9c Update ReadManyByPartitionKeyTest.java -f68cf02ff71 Fixing test issues -8b6c4b168ea Update CosmosAsyncContainer.java -8ba7f4db2da Merge branch 'main' into users/fabianm/readManyByPK -56b067a9339 Reacted to code review feedback -fa430e918fa Merge branch 'main' into users/fabianm/readManyByPK -d9504c91f34 Fix build issues -73151f09e5f Merge branch 'main' into users/fabianm/readManyByPK -681830e2d4a Fixing changelog -7f745e60641 Merge branch 'main' into users/fabianm/readManyByPK -0b8905dbb01 Addressing code review comments -22abc780ed8 Addressing code review feedback -662b1a4b90e Update CosmosItemsDataSource.scala -c764de9de02 Update CosmosItemsDataSource.scala -e1e6f5a6f73 Merge branch 'main' into users/fabianm/readManyByPK -080ce4a2293 Update RxDocumentClientImpl.java -516bbf3a95a Merge branch 'users/fabianm/readManyByPK' of https://github.com/Azure/azure-sdk-for-java into users/fabianm/readManyByPK -b01f8758eea Fix readManyByPartitionKey retries -7130d4aa35a Fix PK.None -93957f3a844 Update ReadManyByPartitionKeyQueryHelper.java - -===== Files Changed ===== - sdk/cosmos/azure-cosmos-spark_3-3_2-12/CHANGELOG.md (+1 -0) - sdk/cosmos/azure-cosmos-spark_3-4_2-12/CHANGELOG.md (+1 -0) - sdk/cosmos/azure-cosmos-spark_3-5_2-12/CHANGELOG.md (+1 -0) - sdk/cosmos/azure-cosmos-spark_3-5_2-13/CHANGELOG.md (+1 -0) - sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosConfig.scala (+20 -2) - sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosConstants.scala (+1 -0) - sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosItemsDataSource.scala (+125 -1) - sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosPartitionKeyHelper.scala (+45 -0) - sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosReadManyByPartitionKeyReader.scala (+150 -0) - sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/ItemsPartitionReaderWithReadManyByPartitionKey.scala (+249 -0) - sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/TransientIOErrorsRetryingReadManyByPartitionKeyIterator.scala (+259 -0) - sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/udf/GetCosmosPartitionKeyValue.scala (+25 -0) - sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/CosmosConfigSpec.scala (+42 -0) - sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/CosmosPartitionKeyHelperSpec.scala (+104 -0) - sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/ItemsPartitionReaderWithReadManyByPartitionKeyITest.scala (+158 -0) - sdk/cosmos/azure-cosmos-spark_4-0_2-13/CHANGELOG.md (+1 -0) - sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/ReadManyByPartitionKeyTest.java (+462 -0) - sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelperTest.java (+426 -0) - sdk/cosmos/azure-cosmos/CHANGELOG.md (+1 -0) - sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java (+126 -0) - sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosContainer.java (+67 -0) - sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/AsyncDocumentClient.java (+21 -0) - sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/Configs.java (+19 -0) - sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelper.java (+263 -0) - sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java (+292 -0) - sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/DocumentQueryExecutionContextFactory.java (+11 -0) - sdk/cosmos/cspell.yaml (+6 -0) - sdk/cosmos/docs/readManyByPartitionKey-design.md (+169 -0) - - diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosItemsDataSource.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosItemsDataSource.scala index 82e3158d765d..202a31c97bd2 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosItemsDataSource.scala +++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosItemsDataSource.scala @@ -127,13 +127,21 @@ object CosmosItemsDataSource { userProvidedSchema, userConfig.asScala.toMap) + // Resolve the null-handling config up front so both the UDF path and the PK-column path honor it. + val sharedEffectiveConfig = CosmosConfig.getEffectiveConfig( + databaseName = None, + containerName = None, + userConfig.asScala.toMap) + val sharedReadConfig = CosmosReadConfig.parseCosmosReadConfig(sharedEffectiveConfig) + val sharedTreatNullAsNone = sharedReadConfig.readManyByPkTreatNullAsNone + // Option 1: Look for the _partitionKeyIdentity column (produced by GetCosmosPartitionKeyValue UDF) val pkIdentityFieldExtraction = df .schema .find(field => field.name.equals(CosmosConstants.Properties.PartitionKeyIdentity) && field.dataType.equals(StringType)) .map(field => (row: Row) => { val rawValue = row.getString(row.fieldIndex(field.name)) - CosmosPartitionKeyHelper.tryParsePartitionKey(rawValue) + CosmosPartitionKeyHelper.tryParsePartitionKey(rawValue, sharedTreatNullAsNone) .getOrElse(throw new IllegalArgumentException( s"Invalid _partitionKeyIdentity value in row: '$rawValue'. " + "Expected format: pk([...json...])")) @@ -143,11 +151,8 @@ object CosmosItemsDataSource { val pkColumnExtraction: Option[Row => PartitionKey] = if (pkIdentityFieldExtraction.isDefined) { None // no need to resolve PK paths - _partitionKeyIdentity column takes precedence } else { - val effectiveConfig = CosmosConfig.getEffectiveConfig( - databaseName = None, - containerName = None, - userConfig.asScala.toMap) - val readConfig = CosmosReadConfig.parseCosmosReadConfig(effectiveConfig) + val effectiveConfig = sharedEffectiveConfig + val readConfig = sharedReadConfig val containerConfig = CosmosContainerConfig.parseCosmosContainerConfig(effectiveConfig) val sparkEnvironmentInfo = CosmosClientConfiguration.getSparkEnvironmentInfo(None) val calledFrom = s"CosmosItemsDataSource.readManyByPartitionKey" diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosPartitionKeyHelper.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosPartitionKeyHelper.scala index 27776f5c3de6..afea6c26c89d 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosPartitionKeyHelper.scala +++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosPartitionKeyHelper.scala @@ -5,7 +5,7 @@ package com.azure.cosmos.spark import com.azure.cosmos.implementation.routing.PartitionKeyInternal import com.azure.cosmos.implementation.{ImplementationBridgeHelpers, Utils} -import com.azure.cosmos.models.PartitionKey +import com.azure.cosmos.models.{PartitionKey, PartitionKeyBuilder} import com.azure.cosmos.spark.diagnostics.BasicLoggingTrait import java.util @@ -27,16 +27,45 @@ private[spark] object CosmosPartitionKeyHelper extends BasicLoggingTrait { s"pk(${objectMapper.writeValueAsString(partitionKeyValue.asJava)})" } - def tryParsePartitionKey(cosmosPartitionKeyString: String): Option[PartitionKey] = { + def tryParsePartitionKey(cosmosPartitionKeyString: String): Option[PartitionKey] = + tryParsePartitionKey(cosmosPartitionKeyString, treatNullAsNone = false) + + /** + * Parses a pk(...) string into a [[PartitionKey]]. + * + * When treatNullAsNone is true, any JSON null components in the serialized array are mapped to + * [[PartitionKeyBuilder.addNoneValue()]] (meaning the document field is absent/undefined). + * When false, they are mapped to [[PartitionKeyBuilder.addNullValue()]] (JSON null value). + * This matches the spark.cosmos.read.readManyByPk.nullHandling config for the non-UDF column path. + */ + def tryParsePartitionKey( + cosmosPartitionKeyString: String, + treatNullAsNone: Boolean): Option[PartitionKey] = { cosmosPartitionKeyString match { case cosmosPartitionKeyStringRegx(pkValue) => scala.util.Try(Utils.parse(pkValue, classOf[Object])).toOption.flatMap { case arrayList: util.ArrayList[Object @unchecked] => - Some( - ImplementationBridgeHelpers - .PartitionKeyHelper - .getPartitionKeyAccessor - .toPartitionKey(PartitionKeyInternal.fromObjectArray(arrayList.toArray, false))) + val components = arrayList.toArray + if (components.exists(_ == null)) { + // Build via PartitionKeyBuilder so nulls can be disambiguated between + // JSON-null (addNullValue) and undefined (addNoneValue) based on config. + val builder = new PartitionKeyBuilder() + components.foreach { + case null => + if (treatNullAsNone) builder.addNoneValue() else builder.addNullValue() + case s: String => builder.add(s) + case n: java.lang.Number => builder.add(n.doubleValue()) + case b: java.lang.Boolean => builder.add(b.booleanValue()) + case other => builder.add(other.toString) + } + Some(builder.build()) + } else { + Some( + ImplementationBridgeHelpers + .PartitionKeyHelper + .getPartitionKeyAccessor + .toPartitionKey(PartitionKeyInternal.fromObjectArray(components, false))) + } case other => Some(new PartitionKey(other)) } case _ => None diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosReadManyByPartitionKeyReader.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosReadManyByPartitionKeyReader.scala index 207f77eb4766..82225a9039ee 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosReadManyByPartitionKeyReader.scala +++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosReadManyByPartitionKeyReader.scala @@ -64,6 +64,14 @@ private[spark] class CosmosReadManyByPartitionKeyReader( cosmosContainerConfig, clientCacheItems(0).get, clientCacheItems(1)) + // Warm-up readItem: intentionally issues a lookup for a random id/partition-key pair + // on the driver so that the collection/routing-map caches are populated before we serialize + // the client state and broadcast it to executors. This costs ~1 RU + 1 RTT per broadcast build + // (expected 404) but avoids every executor doing the same lookup in parallel on first use. + // Warm-up readItem: intentionally issues a lookup for a random id/partition-key pair + // on the driver so that the collection/routing-map caches are populated before we serialize + // the client state and broadcast it to executors. This costs ~1 RU + 1 RTT per broadcast build + // (expected 404) but avoids every executor doing the same lookup in parallel on first use. try { container.readItem( UUIDs.nonBlockingRandomUUID().toString, diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/ReadManyByPartitionKeyTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/ReadManyByPartitionKeyTest.java index 2c26d564ed24..1eff5f7f18e7 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/ReadManyByPartitionKeyTest.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/ReadManyByPartitionKeyTest.java @@ -227,6 +227,43 @@ public void hpk_readManyByPartitionKey_partialPk_twoLevels() { cleanupContainer(multiHashContainer); } + @Test(groups = {"emulator"}, timeOut = TIMEOUT) + @SuppressWarnings("deprecation") + public void hpk_readManyByPartitionKey_withNoneComponent() { + // Regression test for hierarchical partition key routing with PartitionKey.NONE / addNoneValue() + // at a trailing position. Some documents omit the last PK path (areaCode); they must be + // routed via the NOT IS_DEFINED(c["areaCode"]) predicate and returned only when the caller + // requests that slice via addNoneValue(). + createHpkItems(); + // Insert 3 documents where areaCode is undefined (NONE) under Redmond/98053 + for (int i = 0; i < 3; i++) { + ObjectNode item = com.azure.cosmos.implementation.Utils.getSimpleObjectMapper().createObjectNode(); + item.put("id", UUID.randomUUID().toString()); + item.put("city", "Redmond"); + item.put("zipcode", "98053"); + // deliberately omit areaCode + multiHashContainer.createItem(item); + } + + // Request the NONE slice: Redmond/98053/ + List pkValues = Collections.singletonList( + new PartitionKeyBuilder().add("Redmond").add("98053").addNoneValue().build()); + + CosmosPagedIterable results = multiHashContainer.readManyByPartitionKey(pkValues, ObjectNode.class); + List resultList = results.stream().collect(Collectors.toList()); + + // Only the 3 documents without areaCode should come back — the pre-existing items in + // createHpkItems() all have areaCode defined and live in a different physical partition slice. + assertThat(resultList).hasSize(3); + resultList.forEach(item -> { + assertThat(item.get("city").asText()).isEqualTo("Redmond"); + assertThat(item.get("zipcode").asText()).isEqualTo("98053"); + assertThat(item.has("areaCode")).isFalse(); + }); + + cleanupContainer(multiHashContainer); + } + @Test(groups = {"emulator"}, timeOut = TIMEOUT) public void hpk_readManyByPartitionKey_withProjection() { createHpkItems(); diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java index 4e234667c1c0..8260ed95b19d 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java @@ -165,7 +165,7 @@ private static ImplementationBridgeHelpers.CosmosBatchRequestOptionsHelper.Cosmo private final String createItemSpanName; private final String readAllItemsSpanName; private final String readManyItemsSpanName; - private final String readManyByPartitionKeyItemsSpanName; + private final String readManyByPartitionKeySpanName; private final String readAllItemsOfLogicalPartitionSpanName; private final String queryItemsSpanName; private final String queryChangeFeedSpanName; @@ -199,7 +199,7 @@ protected CosmosAsyncContainer(CosmosAsyncContainer toBeWrappedContainer) { this.createItemSpanName = "createItem." + this.id; this.readAllItemsSpanName = "readAllItems." + this.id; this.readManyItemsSpanName = "readManyItems." + this.id; - this.readManyByPartitionKeyItemsSpanName = "readManyByPartitionKeyItems." + this.id; + this.readManyByPartitionKeySpanName = "readManyByPartitionKey." + this.id; this.readAllItemsOfLogicalPartitionSpanName = "readAllItemsOfLogicalPartition." + this.id; this.queryItemsSpanName = "queryItems." + this.id; this.queryChangeFeedSpanName = "queryChangeFeed." + this.id; @@ -1647,7 +1647,34 @@ public CosmosPagedFlux readManyByPartitionKey( * and/or additional filters (e.g. {@code SELECT * FROM c WHERE c.status = 'active'}). * The SDK will automatically append partition key filtering to the custom query. *

- * The custom query must be a simple streamable query — aggregates, ORDER BY, DISTINCT, + * The custom query must be a simple streamable query - aggregates, ORDER BY, DISTINCT, + * GROUP BY, DCOUNT, vector search, and full-text search are not supported and will be + * rejected. + *

+ * Partial hierarchical partition keys are supported and will fan out to multiple + * physical partitions. + * + * @param the type parameter + * @param partitionKeys list of partition key values to read documents for + * @param customQuery optional custom query for projections/additional filters (null means SELECT * FROM c) + * @param classType class type + * @return a {@link CosmosPagedFlux} containing one or several feed response pages + */ + public CosmosPagedFlux readManyByPartitionKey( + List partitionKeys, + SqlQuerySpec customQuery, + Class classType) { + + return this.readManyByPartitionKey(partitionKeys, customQuery, null, classType); + } + + /** + * Reads many documents matching the provided partition key values with a custom query. + * The custom query can be used to apply projections (e.g. {@code SELECT c.name, c.age FROM c}) + * and/or additional filters (e.g. {@code SELECT * FROM c WHERE c.status = 'active'}). + * The SDK will automatically append partition key filtering to the custom query. + *

+ * The custom query must be a simple streamable query - aggregates, ORDER BY, DISTINCT, * GROUP BY, DCOUNT, vector search, and full-text search are not supported and will be * rejected. *

@@ -1696,19 +1723,24 @@ private Function>> readManyByPa CosmosQueryRequestOptions queryRequestOptions = requestOptions == null ? new CosmosQueryRequestOptions() : queryOptionsAccessor().clone(readManyOptionsAccessor().getImpl(requestOptions)); - queryRequestOptions.setMaxDegreeOfParallelism(-1); + // Honor any caller-provided MaxDegreeOfParallelism; only default to the "unbounded" sentinel + // (-1) when the value is still at the default (0). CosmosReadManyRequestOptions currently does not + // expose MDOP, so this branch is defensive in case it is plumbed through in the future. + if (queryRequestOptions.getMaxDegreeOfParallelism() == 0) { + queryRequestOptions.setMaxDegreeOfParallelism(-1); + } queryRequestOptions.setQueryName("readManyByPartitionKey"); CosmosQueryRequestOptionsBase cosmosQueryRequestOptionsImpl = queryOptionsAccessor().getImpl(queryRequestOptions); - applyPolicies(OperationType.Query, ResourceType.Document, cosmosQueryRequestOptionsImpl, this.readManyByPartitionKeyItemsSpanName); + applyPolicies(OperationType.Query, ResourceType.Document, cosmosQueryRequestOptionsImpl, this.readManyByPartitionKeySpanName); QueryFeedOperationState state = new QueryFeedOperationState( client, - this.readManyByPartitionKeyItemsSpanName, + this.readManyByPartitionKeySpanName, database.getId(), this.getId(), ResourceType.Document, OperationType.Query, - queryOptionsAccessor().getQueryNameOrDefault(queryRequestOptions, this.readManyByPartitionKeyItemsSpanName), + queryOptionsAccessor().getQueryNameOrDefault(queryRequestOptions, this.readManyByPartitionKeySpanName), queryRequestOptions, pagedFluxOptions ); diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosContainer.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosContainer.java index 0bd8be5850c0..33149a498390 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosContainer.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosContainer.java @@ -584,7 +584,34 @@ public CosmosPagedIterable readManyByPartitionKey( * and/or additional filters (e.g. {@code SELECT * FROM c WHERE c.status = 'active'}). * The SDK will automatically append partition key filtering to the custom query. *

- * The custom query must be a simple streamable query — aggregates, ORDER BY, DISTINCT, + * The custom query must be a simple streamable query - aggregates, ORDER BY, DISTINCT, + * GROUP BY, DCOUNT, vector search, and full-text search are not supported and will be + * rejected. + *

+ * Partial hierarchical partition keys are supported and will fan out to multiple + * physical partitions. + * + * @param the type parameter + * @param partitionKeys list of partition key values to read documents for + * @param customQuery optional custom query for projections/additional filters (null means SELECT * FROM c) + * @param classType class type + * @return a {@link CosmosPagedIterable} containing the results + */ + public CosmosPagedIterable readManyByPartitionKey( + List partitionKeys, + SqlQuerySpec customQuery, + Class classType) { + + return getCosmosPagedIterable(this.asyncContainer.readManyByPartitionKey(partitionKeys, customQuery, classType)); + } + + /** + * Reads many documents matching the provided partition key values with a custom query. + * The custom query can be used to apply projections (e.g. {@code SELECT c.name, c.age FROM c}) + * and/or additional filters (e.g. {@code SELECT * FROM c WHERE c.status = 'active'}). + * The SDK will automatically append partition key filtering to the custom query. + *

+ * The custom query must be a simple streamable query - aggregates, ORDER BY, DISTINCT, * GROUP BY, DCOUNT, vector search, and full-text search are not supported and will be * rejected. *

diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelper.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelper.java index 24ac8ea19666..965df746768b 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelper.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelper.java @@ -196,17 +196,19 @@ static String extractTableAlias(String queryText) { // Check if there's an alias after the container name (before WHERE or end) if (afterFrom < queryText.length()) { - char nextChar = Character.toUpperCase(queryText.charAt(afterFrom)); - // If the next token is a keyword (WHERE, ORDER, GROUP, JOIN, OFFSET, LIMIT, HAVING) or end, containerName IS the alias - if (nextChar == 'W' || nextChar == 'O' || nextChar == 'G' || nextChar == 'J' - || nextChar == 'L' || nextChar == 'H') { - // Check if it's actually a keyword - String remaining = upper.substring(afterFrom); - if (remaining.startsWith("WHERE") || remaining.startsWith("ORDER") - || remaining.startsWith("GROUP") || remaining.startsWith("JOIN") - || remaining.startsWith("OFFSET") || remaining.startsWith("LIMIT") - || remaining.startsWith("HAVING")) { - return containerName; + String remaining = upper.substring(afterFrom); + // Reserved keywords that terminate the FROM clause - when the next token is one of these, + // containerName itself IS the alias used throughout the rest of the query. + if (isFollowedByReservedKeyword(remaining)) { + return containerName; + } + // Handle optional AS: "FROM root AS r" -> alias is "r" + if (remaining.startsWith("AS") + && (remaining.length() == 2 || !Character.isLetterOrDigit(remaining.charAt(2)))) { + afterFrom += 2; // skip AS + while (afterFrom < queryText.length() + && Character.isWhitespace(queryText.charAt(afterFrom))) { + afterFrom++; } } // Otherwise the next token is the alias ("FROM root r" -> alias is "r") @@ -225,6 +227,18 @@ static String extractTableAlias(String queryText) { return containerName; } + private static boolean isFollowedByReservedKeyword(String remainingUpper) { + String[] keywords = { "WHERE", "ORDER", "GROUP", "JOIN", "OFFSET", "LIMIT", "HAVING" }; + for (String kw : keywords) { + if (remainingUpper.startsWith(kw) + && (remainingUpper.length() == kw.length() + || !Character.isLetterOrDigit(remainingUpper.charAt(kw.length())))) { + return true; + } + } + return false; + } + /** * Finds the index of a top-level SQL keyword in the query text (case-insensitive), * ignoring occurrences inside parentheses or string literals. @@ -234,14 +248,33 @@ static int findTopLevelKeywordIndex(String queryText, String keyword) { String keywordUpper = keyword.toUpperCase(); int depth = 0; int keyLen = keywordUpper.length(); - for (int i = 0; i <= queryTextUpper.length() - keyLen; i++) { - char ch = queryTextUpper.charAt(i); + int len = queryTextUpper.length(); + for (int i = 0; i <= len - keyLen; i++) { + char ch = queryText.charAt(i); + // Skip single-line comments: -- ... end-of-line + if (ch == '-' && i + 1 < len && queryText.charAt(i + 1) == '-') { + i += 2; + while (i < len && queryText.charAt(i) != '\n' && queryText.charAt(i) != '\r') { + i++; + } + continue; + } + // Skip block comments: /* ... */ + if (ch == '/' && i + 1 < len && queryText.charAt(i + 1) == '*') { + i += 2; + while (i + 1 < len + && !(queryText.charAt(i) == '*' && queryText.charAt(i + 1) == '/')) { + i++; + } + i++; // position on the '/'; loop post-increment moves past it + continue; + } // Skip string literals enclosed in single quotes (handle '' escape) - if (queryText.charAt(i) == '\'') { + if (ch == '\'') { i++; - while (i < queryText.length()) { + while (i < len) { if (queryText.charAt(i) == '\'') { - if (i + 1 < queryText.length() && queryText.charAt(i + 1) == '\'') { + if (i + 1 < len && queryText.charAt(i + 1) == '\'') { i += 2; // escaped quote - skip both continue; } @@ -251,11 +284,12 @@ static int findTopLevelKeywordIndex(String queryText, String keyword) { } continue; } - if (ch == '(') { + char upperCh = queryTextUpper.charAt(i); + if (upperCh == '(') { depth++; - } else if (ch == ')') { + } else if (upperCh == ')') { depth--; - } else if (depth == 0 && ch == keywordUpper.charAt(0) + } else if (depth == 0 && upperCh == keywordUpper.charAt(0) && queryTextUpper.startsWith(keywordUpper, i) && (i == 0 || !Character.isLetterOrDigit(queryTextUpper.charAt(i - 1))) && (i + keyLen >= queryTextUpper.length() || !Character.isLetterOrDigit(queryTextUpper.charAt(i + keyLen)))) { diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java index 3d9082a25bcc..46767c9e9dd5 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java @@ -121,6 +121,7 @@ import java.util.EnumSet; import java.util.HashMap; import java.util.HashSet; +import java.util.LinkedHashMap; import java.util.Iterator; import java.util.List; import java.util.Locale; @@ -4512,9 +4513,6 @@ private Flux> readManyByPartitionKey( for (Map.Entry> entry : partitionRangePkMap.entrySet()) { List allPks = entry.getValue(); - if (allPks.isEmpty()) { - continue; - } List> partitionBatches = new ArrayList<>(); for (int i = 0; i < allPks.size(); i += maxPksPerPartitionQuery) { List batch = allPks.subList( @@ -4621,7 +4619,9 @@ private Map> groupPartitionKeysByPhysicalP PartitionKeyDefinition pkDefinition, CollectionRoutingMap routingMap) { - Map> partitionRangePkMap = new HashMap<>(); + // Use LinkedHashMap so the downstream round-robin interleave is deterministic and the iteration + // order follows insertion order of partition keys (i.e. the order the caller provided). + Map> partitionRangePkMap = new LinkedHashMap<>(); for (PartitionKey pk : partitionKeys) { PartitionKeyInternal pkInternal = BridgeInternal.getPartitionKeyInternal(pk); From e306faef6c378cba267123b87b51f12c45395089 Mon Sep 17 00:00:00 2001 From: Fabian Meiswinkel Date: Fri, 17 Apr 2026 17:27:07 +0000 Subject: [PATCH 24/64] React to code review feedback --- .../cosmos/spark/SparkE2EQueryITest.scala | 122 +++++++++++++++++- .../com/azure/cosmos/spark/CosmosConfig.scala | 5 +- .../cosmos/spark/CosmosItemsDataSource.scala | 13 +- .../spark/CosmosPartitionKeyHelper.scala | 17 ++- .../CosmosReadManyByPartitionKeyReader.scala | 42 +++++- ...tionReaderWithReadManyByPartitionKey.scala | 82 +++++++----- .../udf/GetCosmosPartitionKeyValue.scala | 5 +- .../spark/CosmosPartitionKeyHelperSpec.scala | 27 +++- .../cosmos/ReadManyByPartitionKeyTest.java | 98 +++++++++----- ...ByPartitionKeyQueryPlanValidationTest.java | 81 ++++++++++++ .../ReadManyByPartitionKeyQueryHelper.java | 11 ++ .../implementation/RxDocumentClientImpl.java | 104 ++++++++------- .../DocumentQueryExecutionContextFactory.java | 73 +++++++---- .../docs/readManyByPartitionKey-design.md | 23 ++-- 14 files changed, 536 insertions(+), 167 deletions(-) create mode 100644 sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryPlanValidationTest.java diff --git a/sdk/cosmos/azure-cosmos-spark_3-5/src/test/scala/com/azure/cosmos/spark/SparkE2EQueryITest.scala b/sdk/cosmos/azure-cosmos-spark_3-5/src/test/scala/com/azure/cosmos/spark/SparkE2EQueryITest.scala index 5f9cb1dbdbc8..8476e4aa2026 100644 --- a/sdk/cosmos/azure-cosmos-spark_3-5/src/test/scala/com/azure/cosmos/spark/SparkE2EQueryITest.scala +++ b/sdk/cosmos/azure-cosmos-spark_3-5/src/test/scala/com/azure/cosmos/spark/SparkE2EQueryITest.scala @@ -4,13 +4,20 @@ package com.azure.cosmos.spark import com.azure.cosmos.implementation.TestConfigurations +import com.azure.cosmos.models.{CosmosContainerProperties, CosmosItemRequestOptions, PartitionKey, PartitionKeyDefinition, PartitionKeyDefinitionVersion, PartitionKind, ThroughputProperties} +import com.azure.cosmos.spark.udf.GetCosmosPartitionKeyValue import com.fasterxml.jackson.databind.node.ObjectNode +import org.apache.spark.sql.functions.expr +import org.apache.spark.sql.types.StringType -import java.util.UUID +import java.util.{ArrayList, UUID} + +import scala.collection.JavaConverters._ class SparkE2EQueryITest extends SparkE2EQueryITestBase { + // scalastyle:off multiple.string.literals "spark query" can "return proper Cosmos specific query plan on explain with nullable properties" in { val cosmosEndpoint = TestConfigurations.HOST val cosmosMasterKey = TestConfigurations.MASTER_KEY @@ -67,4 +74,115 @@ class SparkE2EQueryITest val item = rowsArray(0) item.getAs[String]("id") shouldEqual id } -} + + "spark readManyByPartitionKey" can "use a matching top-level partition key column without the UDF" in { + val cosmosEndpoint = TestConfigurations.HOST + val cosmosMasterKey = TestConfigurations.MASTER_KEY + val container = cosmosClient.getDatabase(cosmosDatabase).getContainer(cosmosContainersWithPkAsPartitionKey) + val requestOptions = new CosmosItemRequestOptions() + + Seq("pkA", "pkB").foreach { pkValue => + val item = objectMapper.createObjectNode() + item.put("id", s"item-$pkValue") + item.put("pk", pkValue) + item.put("payload", s"value-$pkValue") + + container.createItem(item, new PartitionKey(pkValue), requestOptions).block() + } + + val cfg = Map( + "spark.cosmos.accountEndpoint" -> cosmosEndpoint, + "spark.cosmos.accountKey" -> cosmosMasterKey, + "spark.cosmos.database" -> cosmosDatabase, + "spark.cosmos.container" -> cosmosContainersWithPkAsPartitionKey, + "spark.cosmos.read.inferSchema.enabled" -> "true" + ) + + val sparkSession = spark + import sparkSession.implicits._ + + val rows = CosmosItemsDataSource + .readManyByPartitionKey(Seq("pkA", "pkB").toDF("pk"), cfg.asJava) + .selectExpr("id", "pk", "payload") + .collect() + + rows should have size 2 + rows.map(_.getAs[String]("id")).toSet shouldEqual Set("item-pkA", "item-pkB") + rows.map(_.getAs[String]("pk")).toSet shouldEqual Set("pkA", "pkB") + rows.map(_.getAs[String]("payload")).toSet shouldEqual Set("value-pkA", "value-pkB") + } + "spark readManyByPartitionKey" can "require the UDF for nested partition key paths and succeed with it" in { + val cosmosEndpoint = TestConfigurations.HOST + val cosmosMasterKey = TestConfigurations.MASTER_KEY + val containerName = s"nested-pk-${UUID.randomUUID()}" + + val pkPaths = new ArrayList[String]() + pkPaths.add("/tenant/id") + + val pkDefinition = new PartitionKeyDefinition() + pkDefinition.setPaths(pkPaths) + pkDefinition.setKind(PartitionKind.HASH) + pkDefinition.setVersion(PartitionKeyDefinitionVersion.V2) + + val containerProperties = new CosmosContainerProperties(containerName, pkDefinition) + cosmosClient + .getDatabase(cosmosDatabase) + .createContainerIfNotExists(containerProperties, ThroughputProperties.createManualThroughput(400)) + .block() + + try { + val container = cosmosClient.getDatabase(cosmosDatabase).getContainer(containerName) + val requestOptions = new CosmosItemRequestOptions() + + Seq("tenantA", "tenantB").foreach { tenantId => + val item = objectMapper.createObjectNode() + item.put("id", s"item-$tenantId") + item.put("payload", s"value-$tenantId") + item.putObject("tenant").put("id", tenantId) + + container.createItem(item, new PartitionKey(tenantId), requestOptions).block() + } + + val cfg = Map( + "spark.cosmos.accountEndpoint" -> cosmosEndpoint, + "spark.cosmos.accountKey" -> cosmosMasterKey, + "spark.cosmos.database" -> cosmosDatabase, + "spark.cosmos.container" -> containerName, + "spark.cosmos.read.inferSchema.enabled" -> "true" + ) + + val sparkSession = spark + import sparkSession.implicits._ + + val missingUdfError = the[IllegalArgumentException] thrownBy { + CosmosItemsDataSource.readManyByPartitionKey(Seq("tenantA").toDF("tenantId"), cfg.asJava) + } + + missingUdfError.getMessage should include("Nested paths cannot be resolved from DataFrame columns automatically") + missingUdfError.getMessage should include("_partitionKeyIdentity") + + spark.udf.register("GetCosmosPartitionKeyValue", new GetCosmosPartitionKeyValue(), StringType) + + val inputDf = Seq("tenantA", "tenantB") + .toDF("tenantId") + .withColumn("_partitionKeyIdentity", expr("GetCosmosPartitionKeyValue(tenantId)")) + + val rows = CosmosItemsDataSource + .readManyByPartitionKey(inputDf, cfg.asJava) + .selectExpr("id", "tenant.id as tenantId") + .collect() + + rows should have size 2 + rows.map(_.getAs[String]("id")).toSet shouldEqual Set("item-tenantA", "item-tenantB") + rows.map(_.getAs[String]("tenantId")).toSet shouldEqual Set("tenantA", "tenantB") + } finally { + cosmosClient + .getDatabase(cosmosDatabase) + .getContainer(containerName) + .delete() + .block() + } + } + + // scalastyle:on multiple.string.literals +} \ No newline at end of file diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosConfig.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosConfig.scala index 4a483ef38a6b..8312e34bf0c1 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosConfig.scala +++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosConfig.scala @@ -1147,8 +1147,9 @@ private object CosmosReadConfig { helpMessage = "Determines how null values in partition key columns are treated for " + "readManyByPartitionKey. 'Null' (default) maps null to a JSON null via addNullValue(), which " + "is appropriate when the document field exists with an explicit null value. 'None' maps null " + - "to PartitionKey.NONE via addNoneValue(), which should only be used when the partition key " + - "path does not exist at all in the document. These two semantics hash to DIFFERENT physical " + + "to PartitionKey.NONE via addNoneValue(), which is only supported for single-path partition keys " + + "and should only be used when the partition key path does not exist at all in the document. " + + "Hierarchical partition keys reject this mode. These two semantics hash to DIFFERENT physical " + "partitions - picking the wrong mode for your data will silently return zero rows." ) diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosItemsDataSource.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosItemsDataSource.scala index 202a31c97bd2..3433e352870b 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosItemsDataSource.scala +++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosItemsDataSource.scala @@ -212,7 +212,7 @@ object CosmosItemsDataSource { // Hierarchical partition key - build level by level val builder = new PartitionKeyBuilder() for (path <- pkPaths) { - addPartitionKeyComponent(builder, row.getAs[Any](path), treatNullAsNone) + addPartitionKeyComponent(builder, row.getAs[Any](path), treatNullAsNone, pkPaths.size) } builder.build() } @@ -233,12 +233,19 @@ object CosmosItemsDataSource { readManyReader.readManyByPartitionKey(df.rdd, pkExtraction) } - private def addPartitionKeyComponent(builder: PartitionKeyBuilder, value: Any, treatNullAsNone: Boolean): Unit = { + private def addPartitionKeyComponent( + builder: PartitionKeyBuilder, + value: Any, + treatNullAsNone: Boolean, + partitionKeyComponentCount: Int): Unit = { value match { case s: String => builder.add(s) case n: Number => builder.add(n.doubleValue()) case b: Boolean => builder.add(b) case null => + CosmosPartitionKeyHelper.validateNoneHandlingForPartitionKeyComponentCount( + partitionKeyComponentCount, + treatNullAsNone) if (treatNullAsNone) builder.addNoneValue() else builder.addNullValue() case other => @@ -255,7 +262,7 @@ object CosmosItemsDataSource { private def buildPartitionKey(value: Any, treatNullAsNone: Boolean): PartitionKey = { val builder = new PartitionKeyBuilder() - addPartitionKeyComponent(builder, value, treatNullAsNone) + addPartitionKeyComponent(builder, value, treatNullAsNone, 1) builder.build() } } diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosPartitionKeyHelper.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosPartitionKeyHelper.scala index afea6c26c89d..29831e9a6fed 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosPartitionKeyHelper.scala +++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosPartitionKeyHelper.scala @@ -15,6 +15,20 @@ import scala.collection.JavaConverters._ // scalastyle:on underscore.import private[spark] object CosmosPartitionKeyHelper extends BasicLoggingTrait { + private[spark] val HierarchicalPartitionKeyNoneHandlingErrorMessage = + s"The configuration '${CosmosConfigNames.ReadManyByPkNullHandling}=None' is not supported for " + + "hierarchical partition keys because PartitionKey.NONE can't be used with multiple paths. " + + "Use 'Null' for explicit JSON null values, filter out rows with missing partition key " + + "components, or provide fully-defined hierarchical partition keys." + + private[spark] def validateNoneHandlingForPartitionKeyComponentCount( + componentCount: Int, + treatNullAsNone: Boolean): Unit = { + if (treatNullAsNone && componentCount > 1) { + throw new IllegalArgumentException(HierarchicalPartitionKeyNoneHandlingErrorMessage) + } + } + // pattern will be recognized // pk(partitionKeyValue) // @@ -47,6 +61,7 @@ private[spark] object CosmosPartitionKeyHelper extends BasicLoggingTrait { case arrayList: util.ArrayList[Object @unchecked] => val components = arrayList.toArray if (components.exists(_ == null)) { + validateNoneHandlingForPartitionKeyComponentCount(components.length, treatNullAsNone) // Build via PartitionKeyBuilder so nulls can be disambiguated between // JSON-null (addNullValue) and undefined (addNoneValue) based on config. val builder = new PartitionKeyBuilder() @@ -71,4 +86,4 @@ private[spark] object CosmosPartitionKeyHelper extends BasicLoggingTrait { case _ => None } } -} +} \ No newline at end of file diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosReadManyByPartitionKeyReader.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosReadManyByPartitionKeyReader.scala index 82225a9039ee..bd6fa5be8bdf 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosReadManyByPartitionKeyReader.scala +++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosReadManyByPartitionKeyReader.scala @@ -138,6 +138,7 @@ private[spark] class CosmosReadManyByPartitionKeyReader( logInfo(s"Creating an ItemsPartitionReaderWithReadManyByPartitionKey for Activity $correlationActivityId to read for " + s"input partition [$partitionIndex] ${tableName}") + val taskContext = TaskContext.get val reader = new ItemsPartitionReaderWithReadManyByPartitionKey( effectiveUserConfig, CosmosReadManyHelper.FullRangeFeedRange, @@ -146,13 +147,46 @@ private[spark] class CosmosReadManyByPartitionKeyReader( clientStates, DiagnosticsConfig.parseDiagnosticsConfig(effectiveUserConfig), sparkEnvironmentInfo, - TaskContext.get, + taskContext, pkIterator) new Iterator[Row] { - override def hasNext: Boolean = reader.next() - - override def next(): Row = reader.getCurrentRow() + private var isClosed = false + + private def closeReader(): Unit = { + if (!isClosed) { + isClosed = true + reader.close() + } + } + + if (taskContext != null) { + taskContext.addTaskCompletionListener[Unit](_ => closeReader()) + } + + override def hasNext: Boolean = { + try { + val hasMore = reader.next() + if (!hasMore) { + closeReader() + } + hasMore + } catch { + case error: Throwable => + closeReader() + throw error + } + } + + override def next(): Row = { + try { + reader.getCurrentRow() + } catch { + case error: Throwable => + closeReader() + throw error + } + } } }, preservesPartitioning = true diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/ItemsPartitionReaderWithReadManyByPartitionKey.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/ItemsPartitionReaderWithReadManyByPartitionKey.scala index 7fb6e0eb0aba..5c994cc2f7f8 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/ItemsPartitionReaderWithReadManyByPartitionKey.scala +++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/ItemsPartitionReaderWithReadManyByPartitionKey.scala @@ -21,6 +21,7 @@ import org.apache.spark.sql.connector.read.PartitionReader import org.apache.spark.sql.types.StructType import java.util +import java.util.concurrent.atomic.AtomicBoolean // scalastyle:off underscore.import import scala.collection.JavaConverters._ @@ -195,31 +196,41 @@ private[spark] case class ItemsPartitionReaderWithReadManyByPartitionKey // Pass the full PK list to the SDK (which batches per physical partition internally). // On transient I/O failures the retry iterator tracks pages already emitted upstream - // and skips them on replay; if a failure occurs mid-page (after items from that page - // have been emitted) the task fails rather than risking row duplication. - private lazy val iterator: CloseableSparkRowItemIterator = - if (pkList.isEmpty) { - EmptySparkRowItemIterator - } else { - new CloseableSparkRowItemIterator { - private val delegate = new TransientIOErrorsRetryingReadManyByPartitionKeyIterator[SparkRowItem]( - cosmosAsyncContainer, - pkList, - readConfig.customQuery.map(_.toSqlQuerySpec), - readManyOptions, - readConfig.maxItemCount, - readConfig.prefetchBufferSize, - operationContextAndListenerTuple, - classOf[SparkRowItem] - ) - - override def hasNext: Boolean = delegate.hasNext - - override def next(): SparkRowItem = delegate.next() - - override def close(): Unit = delegate.close() - } - } + // and skips them on replay; if a failure occurs mid-page (after items from that page have been + // emitted) the task fails rather than risking row duplication. + private val isClosed = new AtomicBoolean(false) + private var iteratorOpt: Option[CloseableSparkRowItemIterator] = None + + private def getOrCreateIterator: CloseableSparkRowItemIterator = iteratorOpt match { + case Some(existing) => existing + case None => + val created = + if (pkList.isEmpty) { + EmptySparkRowItemIterator + } else { + new CloseableSparkRowItemIterator { + private val delegate = new TransientIOErrorsRetryingReadManyByPartitionKeyIterator[SparkRowItem]( + cosmosAsyncContainer, + pkList, + readConfig.customQuery.map(_.toSqlQuerySpec), + readManyOptions, + readConfig.maxItemCount, + readConfig.prefetchBufferSize, + operationContextAndListenerTuple, + classOf[SparkRowItem] + ) + + override def hasNext: Boolean = delegate.hasNext + + override def next(): SparkRowItem = delegate.next() + + override def close(): Unit = delegate.close() + } + } + + iteratorOpt = Some(created) + created + } private val rowSerializer: ExpressionEncoder.Serializer[Row] = RowSerializerPool.getOrCreateSerializer(readSchema) @@ -236,20 +247,23 @@ private[spark] case class ItemsPartitionReaderWithReadManyByPartitionKey } } - override def next(): Boolean = iterator.hasNext + override def next(): Boolean = getOrCreateIterator.hasNext override def get(): InternalRow = { - cosmosRowConverter.fromRowToInternalRow(iterator.next().row, rowSerializer) + cosmosRowConverter.fromRowToInternalRow(getOrCreateIterator.next().row, rowSerializer) } - def getCurrentRow(): Row = iterator.next().row + def getCurrentRow(): Row = getOrCreateIterator.next().row override def close(): Unit = { - this.iterator.close() - RowSerializerPool.returnSerializerToPool(readSchema, rowSerializer) - clientCacheItem.close() - if (throughputControlClientCacheItemOpt.isDefined) { - throughputControlClientCacheItemOpt.get.close() + if (isClosed.compareAndSet(false, true)) { + iteratorOpt.foreach(_.close()) + iteratorOpt = None + RowSerializerPool.returnSerializerToPool(readSchema, rowSerializer) + clientCacheItem.close() + if (throughputControlClientCacheItemOpt.isDefined) { + throughputControlClientCacheItemOpt.get.close() + } } } -} +} \ No newline at end of file diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/udf/GetCosmosPartitionKeyValue.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/udf/GetCosmosPartitionKeyValue.scala index fc038861a19a..4dcc812f3122 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/udf/GetCosmosPartitionKeyValue.scala +++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/udf/GetCosmosPartitionKeyValue.scala @@ -12,8 +12,9 @@ class GetCosmosPartitionKeyValue extends UDF1[Object, String] { // single-level partition key with a JSON null component; parsing that string back via // CosmosPartitionKeyHelper.tryParsePartitionKey yields a PartitionKey built with // addNullValue(). If the caller instead wants PartitionKey.NONE semantics (absent PK - // field) they should filter the null row before calling this UDF and use the - // schema-matched readManyByPartitionKey path with readManyByPk.nullHandling=None. + // field) they should filter the null row before calling this UDF and use the schema-matched + // readManyByPartitionKey path with readManyByPk.nullHandling=None. That None mode is only + // supported for single-path partition keys; hierarchical partition keys reject it. override def call(partitionKeyValue: Object): String = { partitionKeyValue match { case null => diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/CosmosPartitionKeyHelperSpec.scala b/sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/CosmosPartitionKeyHelperSpec.scala index 1ac40e395847..c81113d3a190 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/CosmosPartitionKeyHelperSpec.scala +++ b/sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/CosmosPartitionKeyHelperSpec.scala @@ -88,15 +88,28 @@ class CosmosPartitionKeyHelperSpec extends UnitSpec { pk.isDefined shouldBe false } - it should "produce different partition keys for addNullValue vs addNoneValue in HPK" in { - // addNullValue represents an explicit JSON null for a field that exists with value null - val pkWithNull = new PartitionKeyBuilder().add("Redmond").addNullValue().build() + it should "parse single-path null as PartitionKey.NONE when treatNullAsNone is true" in { + val pk = CosmosPartitionKeyHelper.tryParsePartitionKey("pk([null])", treatNullAsNone = true) - // addNoneValue represents PartitionKey.NONE, meaning the field is absent/undefined - val pkWithNone = new PartitionKeyBuilder().add("Redmond").addNoneValue().build() + pk.isDefined shouldBe true + pk.get shouldEqual PartitionKey.NONE + } + + it should "throw a clear error when None nullHandling is used for hierarchical partition keys" in { + val error = the[IllegalArgumentException] thrownBy { + CosmosPartitionKeyHelper.tryParsePartitionKey("pk([\"Redmond\",null])", treatNullAsNone = true) + } + + error.getMessage should include(CosmosConfigNames.ReadManyByPkNullHandling) + error.getMessage should include("hierarchical partition keys") + } + + it should "reject addNoneValue in hierarchical partition keys" in { + val error = the[IllegalStateException] thrownBy { + new PartitionKeyBuilder().add("Redmond").addNoneValue().build() + } - // These MUST produce different partition key hashes and route to different physical partitions - pkWithNull should not equal pkWithNone + error.getMessage should include("PartitionKey.None can't be used with multiple paths") } //scalastyle:on multiple.string.literals diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/ReadManyByPartitionKeyTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/ReadManyByPartitionKeyTest.java index 1eff5f7f18e7..ca684f1abf9e 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/ReadManyByPartitionKeyTest.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/ReadManyByPartitionKeyTest.java @@ -27,7 +27,9 @@ import java.util.Arrays; import java.util.Collections; import java.util.List; +import java.util.Map; import java.util.UUID; +import java.util.concurrent.atomic.AtomicInteger; import java.util.stream.Collectors; import static org.assertj.core.api.Assertions.assertThat; @@ -230,38 +232,30 @@ public void hpk_readManyByPartitionKey_partialPk_twoLevels() { @Test(groups = {"emulator"}, timeOut = TIMEOUT) @SuppressWarnings("deprecation") public void hpk_readManyByPartitionKey_withNoneComponent() { - // Regression test for hierarchical partition key routing with PartitionKey.NONE / addNoneValue() - // at a trailing position. Some documents omit the last PK path (areaCode); they must be - // routed via the NOT IS_DEFINED(c["areaCode"]) predicate and returned only when the caller - // requests that slice via addNoneValue(). - createHpkItems(); - // Insert 3 documents where areaCode is undefined (NONE) under Redmond/98053 - for (int i = 0; i < 3; i++) { + try { + createHpkItems(); + ObjectNode item = com.azure.cosmos.implementation.Utils.getSimpleObjectMapper().createObjectNode(); item.put("id", UUID.randomUUID().toString()); item.put("city", "Redmond"); item.put("zipcode", "98053"); - // deliberately omit areaCode - multiHashContainer.createItem(item); - } - - // Request the NONE slice: Redmond/98053/ - List pkValues = Collections.singletonList( - new PartitionKeyBuilder().add("Redmond").add("98053").addNoneValue().build()); - CosmosPagedIterable results = multiHashContainer.readManyByPartitionKey(pkValues, ObjectNode.class); - List resultList = results.stream().collect(Collectors.toList()); - - // Only the 3 documents without areaCode should come back — the pre-existing items in - // createHpkItems() all have areaCode defined and live in a different physical partition slice. - assertThat(resultList).hasSize(3); - resultList.forEach(item -> { - assertThat(item.get("city").asText()).isEqualTo("Redmond"); - assertThat(item.get("zipcode").asText()).isEqualTo("98053"); - assertThat(item.has("areaCode")).isFalse(); - }); + try { + multiHashContainer.createItem(item); + fail("Should have thrown CosmosException for HPK item with missing trailing partition key component"); + } catch (CosmosException e) { + assertThat(e.getMessage()).contains("wrong-pk-value"); + } - cleanupContainer(multiHashContainer); + try { + new PartitionKeyBuilder().add("Redmond").add("98053").addNoneValue().build(); + fail("Should have thrown IllegalStateException for HPK addNoneValue"); + } catch (IllegalStateException e) { + assertThat(e.getMessage()).contains("PartitionKey.None can't be used with multiple paths"); + } + } finally { + cleanupContainer(multiHashContainer); + } } @Test(groups = {"emulator"}, timeOut = TIMEOUT) @@ -367,6 +361,22 @@ public void rejectsEmptyPartitionKeyList() { .stream().collect(Collectors.toList()); } + @Test(groups = {"emulator"}, timeOut = TIMEOUT) + public void rejectsOffsetQuery() { + List pkValues = Collections.singletonList(new PartitionKey("pk1")); + SqlQuerySpec offsetQuery = new SqlQuerySpec("SELECT * FROM c OFFSET 0 LIMIT 10"); + + try { + singlePkContainer.readManyByPartitionKey(pkValues, offsetQuery, null, ObjectNode.class) + .stream().collect(Collectors.toList()); + fail("Should have thrown IllegalArgumentException for OFFSET query"); + } catch (IllegalArgumentException e) { + assertThat(e.getMessage()).contains("OFFSET"); + } + } + + + //endregion @@ -393,9 +403,14 @@ public void singlePk_readManyByPartitionKey_withSmallBatchSize() { new PartitionKey("batchPk4")); CosmosPagedIterable results = singlePkContainer.readManyByPartitionKey(pkValues, ObjectNode.class); - List resultList = results.stream().collect(Collectors.toList()); + List> pages = new ArrayList<>(); + results.iterableByPage().forEach(pages::add); + List resultList = pages.stream() + .flatMap(page -> page.getResults().stream()) + .collect(Collectors.toList()); assertThat(resultList).hasSize(8); // 2 items per PK * 4 PKs + assertThat(pages.size()).isGreaterThan(1); resultList.forEach(item -> { String pk = item.get("mypk").asText(); assertThat(pk).isIn("batchPk1", "batchPk2", "batchPk3", "batchPk4"); @@ -417,19 +432,31 @@ public void singlePk_readManyByPartitionKey_withSmallBatchSize() { @Test(groups = {"emulator"}, timeOut = TIMEOUT) public void singlePk_readManyByPartitionKey_withRequestOptions() { - // This test ensures that request options (like throughput control settings) - // are properly propagated through the readManyByPartitionKey path. - // It acts as a regression test for the redundant options construction bug. List items = createSinglePkItems("pkOpts", 3); List pkValues = Collections.singletonList(new PartitionKey("pkOpts")); com.azure.cosmos.models.CosmosReadManyRequestOptions options = new com.azure.cosmos.models.CosmosReadManyRequestOptions(); + AtomicInteger deserializeCount = new AtomicInteger(); + options.setCustomItemSerializer(new CosmosItemSerializerNoExceptionWrapping() { + @Override + public Map serialize(T item) { + return CosmosItemSerializer.DEFAULT_SERIALIZER.serialize(item); + } - CosmosPagedIterable results = singlePkContainer.readManyByPartitionKey( - pkValues, options, ObjectNode.class); - List resultList = results.stream().collect(Collectors.toList()); + @Override + public T deserialize(Map jsonNodeMap, Class classType) { + deserializeCount.incrementAndGet(); + return CosmosItemSerializer.DEFAULT_SERIALIZER.deserialize(jsonNodeMap, classType); + } + }); + + CosmosPagedIterable results = singlePkContainer.readManyByPartitionKey( + pkValues, options, ReadManyByPartitionKeyPojo.class); + List resultList = results.stream().collect(Collectors.toList()); assertThat(resultList).hasSize(3); + assertThat(deserializeCount.get()).isEqualTo(3); + assertThat(resultList.stream().map(item -> item.mypk)).containsOnly("pkOpts"); cleanupContainer(singlePkContainer); } @@ -495,5 +522,10 @@ private void cleanupContainer(CosmosContainer container) { }); } + private static class ReadManyByPartitionKeyPojo { + public String id; + public String mypk; + } + //endregion } diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryPlanValidationTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryPlanValidationTest.java new file mode 100644 index 000000000000..429d1fd8f89c --- /dev/null +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryPlanValidationTest.java @@ -0,0 +1,81 @@ +/* + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. + */ + +package com.azure.cosmos.implementation; + +import com.azure.cosmos.implementation.query.DCountInfo; +import com.azure.cosmos.implementation.query.PartitionedQueryExecutionInfo; +import com.azure.cosmos.implementation.query.QueryInfo; +import com.azure.cosmos.implementation.query.hybridsearch.HybridSearchQueryInfo; +import com.fasterxml.jackson.databind.node.ObjectNode; +import org.testng.annotations.Test; + +import static org.assertj.core.api.Assertions.assertThatCode; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +public class ReadManyByPartitionKeyQueryPlanValidationTest { + + @Test(groups = { "unit" }) + public void rejectsDCountQueryPlan() { + QueryInfo queryInfo = new QueryInfo(); + DCountInfo dCountInfo = new DCountInfo(); + dCountInfo.setDCountAlias("countAlias"); + queryInfo.set("dCountInfo", dCountInfo); + + assertThatThrownBy(() -> RxDocumentClientImpl.validateQueryPlanForReadManyByPartitionKey(createQueryPlan(queryInfo, null))) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("DCOUNT"); + } + + @Test(groups = { "unit" }) + public void rejectsOffsetQueryPlan() { + QueryInfo queryInfo = new QueryInfo(); + queryInfo.set("offset", 10); + + assertThatThrownBy(() -> RxDocumentClientImpl.validateQueryPlanForReadManyByPartitionKey(createQueryPlan(queryInfo, null))) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("OFFSET"); + } + + @Test(groups = { "unit" }) + public void rejectsLimitQueryPlan() { + QueryInfo queryInfo = new QueryInfo(); + queryInfo.set("limit", 10); + + assertThatThrownBy(() -> RxDocumentClientImpl.validateQueryPlanForReadManyByPartitionKey(createQueryPlan(queryInfo, null))) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("LIMIT"); + } + + @Test(groups = { "unit" }) + public void rejectsHybridSearchQueryPlanWithoutDereferencingNullQueryInfo() { + assertThatThrownBy(() -> RxDocumentClientImpl.validateQueryPlanForReadManyByPartitionKey( + createQueryPlan(null, new HybridSearchQueryInfo()))) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("hybrid/vector/full-text"); + } + + @Test(groups = { "unit" }) + public void acceptsSimpleQueryPlan() { + QueryInfo queryInfo = new QueryInfo(); + + assertThatCode(() -> RxDocumentClientImpl.validateQueryPlanForReadManyByPartitionKey(createQueryPlan(queryInfo, null))) + .doesNotThrowAnyException(); + } + + private PartitionedQueryExecutionInfo createQueryPlan(QueryInfo queryInfo, HybridSearchQueryInfo hybridSearchQueryInfo) { + ObjectNode content = Utils.getSimpleObjectMapper().createObjectNode(); + content.put("partitionedQueryExecutionInfoVersion", Constants.PartitionedQueryExecutionInfo.VERSION_1); + + if (queryInfo != null) { + content.set("queryInfo", Utils.getSimpleObjectMapper().valueToTree(queryInfo.getMap())); + } + if (hybridSearchQueryInfo != null) { + content.set("hybridSearchQueryInfo", Utils.getSimpleObjectMapper().createObjectNode()); + } + + return new PartitionedQueryExecutionInfo(content, null); + } +} \ No newline at end of file diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelper.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelper.java index 965df746768b..0bdb867dc3ee 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelper.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelper.java @@ -12,6 +12,7 @@ import java.util.ArrayList; import java.util.List; +import java.util.stream.Collectors; /** * Helper for constructing SqlQuerySpec instances for readManyByPartitionKey operations. @@ -160,6 +161,16 @@ public static SqlQuerySpec createReadManyByPkQuerySpec( return new SqlQuerySpec(finalQuery, parameters); } + static List createPkSelectors(PartitionKeyDefinition partitionKeyDefinition) { + return partitionKeyDefinition.getPaths() + .stream() + .map(PathParser::getPathParts) + .map(pathParts -> pathParts.stream() + .map(pathPart -> "[\"" + pathPart.replace("\"", "\\") + "\"]") + .collect(Collectors.joining())) + .collect(Collectors.toList()); + } + /** * Extracts the table/collection alias from a SQL query's FROM clause. * Handles: "SELECT * FROM c", "SELECT x.id FROM x WHERE ...", "SELECT * FROM root r", etc. diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java index 46767c9e9dd5..1126008946a3 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java @@ -4485,7 +4485,7 @@ private Flux> readManyByPartitionKey( Map> partitionRangePkMap = groupPartitionKeysByPhysicalPartition(partitionKeys, pkDefinition, routingMap); - List partitionKeySelectors = createPkSelectors(pkDefinition); + List partitionKeySelectors = ReadManyByPartitionKeyQueryHelper.createPkSelectors(pkDefinition); String baseQueryText; List baseParameters; @@ -4502,7 +4502,7 @@ private Flux> readManyByPartitionKey( // Build per-physical-partition batched queries. // Each physical partition may have many PKs - split into batches // to avoid oversized SQL queries. Batch size is configurable via - // system property COSMOS.READ_MANY_BY_PK_MAX_BATCH_SIZE (default 1000). + // system property COSMOS.READ_MANY_BY_PK_MAX_BATCH_SIZE (default 100). int maxPksPerPartitionQuery = Configs.getReadManyByPkMaxBatchSize(); // Build batches per partition as a list of lists (one inner list per partition). @@ -4577,41 +4577,62 @@ private Mono validateCustomQueryForReadManyByPartitionKey( RxDocumentClientImpl.this, getOperationContextAndListenerTuple(queryRequestOptions)); return DocumentQueryExecutionContextFactory - .fetchQueryPlanForValidation(this, queryClient, customQuery, resourceLink, queryRequestOptions) - .flatMap(queryPlan -> { - QueryInfo queryInfo = queryPlan.getQueryInfo(); + .fetchQueryPlanForValidation( + this, + queryClient, + customQuery, + resourceLink, + queryRequestOptions, + Configs.isQueryPlanCachingEnabled(), + this.getQueryPlanCache()) + .doOnNext(RxDocumentClientImpl::validateQueryPlanForReadManyByPartitionKey) + .then(); + } - if (queryInfo.hasGroupBy()) { - return Mono.error(new IllegalArgumentException( - "Custom query for readMany by partition key must not contain GROUP BY.")); - } - if (queryInfo.hasAggregates()) { - return Mono.error(new IllegalArgumentException( - "Custom query for readMany by partition key must not contain aggregates.")); - } - if (queryInfo.hasOrderBy()) { - return Mono.error(new IllegalArgumentException( - "Custom query for readMany by partition key must not contain ORDER BY.")); - } - if (queryInfo.hasDistinct()) { - return Mono.error(new IllegalArgumentException( - "Custom query for readMany by partition key must not contain DISTINCT.")); - } - if (queryInfo.hasDCount()) { - return Mono.error(new IllegalArgumentException( - "Custom query for readMany by partition key must not contain DCOUNT.")); - } - if (queryInfo.hasNonStreamingOrderBy()) { - return Mono.error(new IllegalArgumentException( - "Custom query for readMany by partition key must not contain non-streaming ORDER BY.")); - } - if (queryPlan.hasHybridSearchQueryInfo()) { - return Mono.error(new IllegalArgumentException( - "Custom query for readMany by partition key must not contain hybrid/vector/full-text search.")); - } + static void validateQueryPlanForReadManyByPartitionKey(PartitionedQueryExecutionInfo queryPlan) { + if (queryPlan.hasHybridSearchQueryInfo()) { + throw new IllegalArgumentException( + "Custom query for readMany by partition key must not contain hybrid/vector/full-text search."); + } - return Mono.empty(); - }); + QueryInfo queryInfo = queryPlan.getQueryInfo(); + if (queryInfo == null) { + throw new IllegalArgumentException( + "Custom query for readMany by partition key is not supported because query plan details are unavailable."); + } + + if (queryInfo.hasGroupBy()) { + throw new IllegalArgumentException( + "Custom query for readMany by partition key must not contain GROUP BY."); + } + if (queryInfo.hasAggregates()) { + throw new IllegalArgumentException( + "Custom query for readMany by partition key must not contain aggregates."); + } + if (queryInfo.hasOrderBy()) { + throw new IllegalArgumentException( + "Custom query for readMany by partition key must not contain ORDER BY."); + } + if (queryInfo.hasDistinct()) { + throw new IllegalArgumentException( + "Custom query for readMany by partition key must not contain DISTINCT."); + } + if (queryInfo.hasDCount()) { + throw new IllegalArgumentException( + "Custom query for readMany by partition key must not contain DCOUNT."); + } + if (queryInfo.hasOffset()) { + throw new IllegalArgumentException( + "Custom query for readMany by partition key must not contain OFFSET."); + } + if (queryInfo.hasLimit()) { + throw new IllegalArgumentException( + "Custom query for readMany by partition key must not contain LIMIT."); + } + if (queryInfo.hasNonStreamingOrderBy()) { + throw new IllegalArgumentException( + "Custom query for readMany by partition key must not contain non-streaming ORDER BY."); + } } private Map> groupPartitionKeysByPhysicalPartition( @@ -4665,7 +4686,7 @@ private Map getRangeQueryMap( //TODO: Optimise this to include all types of partitionkeydefinitions. ex: c["prop1./ab"]["key1"] Map rangeQueryMap = new HashMap<>(); - List partitionKeySelectors = createPkSelectors(partitionKeyDefinition); + List partitionKeySelectors = ReadManyByPartitionKeyQueryHelper.createPkSelectors(partitionKeyDefinition); for(Map.Entry> entry: partitionRangeItemKeyMap.entrySet()) { SqlQuerySpec sqlQuerySpec; @@ -4759,15 +4780,6 @@ private SqlQuerySpec createReadManyQuerySpec( return new SqlQuerySpec(queryStringBuilder.toString(), parameters); } - private List createPkSelectors(PartitionKeyDefinition partitionKeyDefinition) { - return partitionKeyDefinition.getPaths() - .stream() - .map(pathPart -> StringUtils.substring(pathPart, 1)) // skip starting / - .map(pathPart -> StringUtils.replace(pathPart, "\"", "\\")) // escape quote - .map(part -> "[\"" + part + "\"]") - .collect(Collectors.toList()); - } - private Flux> queryForReadMany( ScopedDiagnosticsFactory diagnosticsFactory, String parentResourceLink, @@ -5289,7 +5301,7 @@ public Flux> readAllDocuments( } PartitionKeyDefinition pkDefinition = collection.getPartitionKey(); - List partitionKeySelectors = createPkSelectors(pkDefinition); + List partitionKeySelectors = ReadManyByPartitionKeyQueryHelper.createPkSelectors(pkDefinition); SqlQuerySpec querySpec = createLogicalPartitionScanQuerySpec(partitionKey, partitionKeySelectors); String resourceLink = parentResourceLinkToQueryLink(collectionLink, ResourceType.Document); diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/DocumentQueryExecutionContextFactory.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/DocumentQueryExecutionContextFactory.java index d8f9614343cb..82506ae361e1 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/DocumentQueryExecutionContextFactory.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/DocumentQueryExecutionContextFactory.java @@ -107,7 +107,6 @@ private static Mono getPartitionKeyRangesAn } Instant startTime = Instant.now(); - Mono queryExecutionInfoMono; if (queryRequestOptionsAccessor() .isQueryPlanRetrievalDisallowed(cosmosQueryRequestOptions)) { @@ -122,37 +121,53 @@ private static Mono getPartitionKeyRangesAn endTime); } + return fetchQueryPlan( + diagnosticsClientContext, + client, + query, + resourceLink, + cosmosQueryRequestOptions, + queryPlanCachingEnabled, + queryPlanCache) + .flatMap( + partitionedQueryExecutionInfo -> { + + Instant endTime = Instant.now(); + + return getTargetRangesFromQueryPlan(cosmosQueryRequestOptions, collection, queryExecutionContext, + partitionedQueryExecutionInfo, startTime, endTime); + }); + } + + private static Mono fetchQueryPlan( + DiagnosticsClientContext diagnosticsClientContext, + IDocumentQueryClient client, + SqlQuerySpec query, + String resourceLink, + CosmosQueryRequestOptions cosmosQueryRequestOptions, + boolean queryPlanCachingEnabled, + Map queryPlanCache) { + if (queryPlanCachingEnabled && - isScopedToSinglePartition(cosmosQueryRequestOptions) && - queryPlanCache.containsKey(query.getQueryText())) { - Instant endTime = Instant.now(); // endTime for query plan diagnostics + isScopedToSinglePartition(cosmosQueryRequestOptions) && + queryPlanCache.containsKey(query.getQueryText())) { PartitionedQueryExecutionInfo partitionedQueryExecutionInfo = queryPlanCache.get(query.getQueryText()); if (partitionedQueryExecutionInfo != null) { logger.debug("Skipping query plan round trip by using the cached plan"); - return getTargetRangesFromQueryPlan(cosmosQueryRequestOptions, collection, queryExecutionContext, - partitionedQueryExecutionInfo, startTime, endTime); + return Mono.just(partitionedQueryExecutionInfo); } } - queryExecutionInfoMono = - QueryPlanRetriever.getQueryPlanThroughGatewayAsync( - diagnosticsClientContext, - client, - query, - resourceLink, - cosmosQueryRequestOptions); - - return queryExecutionInfoMono.flatMap( - partitionedQueryExecutionInfo -> { - - Instant endTime = Instant.now(); - + return QueryPlanRetriever.getQueryPlanThroughGatewayAsync( + diagnosticsClientContext, + client, + query, + resourceLink, + cosmosQueryRequestOptions) + .doOnNext(partitionedQueryExecutionInfo -> { if (queryPlanCachingEnabled && isScopedToSinglePartition(cosmosQueryRequestOptions)) { tryCacheQueryPlan(query, partitionedQueryExecutionInfo, queryPlanCache); } - - return getTargetRangesFromQueryPlan(cosmosQueryRequestOptions, collection, queryExecutionContext, - partitionedQueryExecutionInfo, startTime, endTime); }); } @@ -323,10 +338,18 @@ public static Mono fetchQueryPlanForValidation( IDocumentQueryClient queryClient, SqlQuerySpec sqlQuerySpec, String resourceLink, - CosmosQueryRequestOptions queryRequestOptions) { + CosmosQueryRequestOptions queryRequestOptions, + boolean queryPlanCachingEnabled, + Map queryPlanCache) { - return QueryPlanRetriever.getQueryPlanThroughGatewayAsync( - diagnosticsClientContext, queryClient, sqlQuerySpec, resourceLink, queryRequestOptions); + return fetchQueryPlan( + diagnosticsClientContext, + queryClient, + sqlQuerySpec, + resourceLink, + queryRequestOptions, + queryPlanCachingEnabled, + queryPlanCache); } public static Flux> createDocumentQueryExecutionContextAsync( diff --git a/sdk/cosmos/docs/readManyByPartitionKey-design.md b/sdk/cosmos/docs/readManyByPartitionKey-design.md index 95d7624f0c8b..c193630d4064 100644 --- a/sdk/cosmos/docs/readManyByPartitionKey-design.md +++ b/sdk/cosmos/docs/readManyByPartitionKey-design.md @@ -20,8 +20,8 @@ and additional filters. The SDK appends the auto-generated PK WHERE clause to it | Partial HPK | Supported from the start; prefix PKs fan out via `getOverlappingRanges` | | PK deduplication | Done at Spark layer only, not in the SDK | | Spark UDF | New `GetCosmosPartitionKeyValue` UDF | -| Custom query validation | Gateway query plan; reject aggregates/ORDER BY/DISTINCT/GROUP BY/DCount/non-streaming ORDER BY/vector/fulltext | -| PK list size | No hard upper-bound enforced; SDK batches internally per physical partition (default 1000 PKs per batch, configurable via `COSMOS.READ_MANY_BY_PK_MAX_BATCH_SIZE`) | +| Custom query validation | Gateway query plan via the standard SDK query-plan retrieval path; reject aggregates/ORDER BY/DISTINCT/GROUP BY/DCount/OFFSET/LIMIT/non-streaming ORDER BY/vector/fulltext | +| PK list size | No hard upper-bound enforced; SDK batches internally per physical partition (default 100 PKs per batch, configurable via `COSMOS.READ_MANY_BY_PK_MAX_BATCH_SIZE`) | | Eager validation | Null and empty PK list rejected eagerly (not lazily in reactive chain) | | Telemetry | Separate span name `readManyByPartitionKeyItems.` (distinct from existing `readManyItems`) | | Query construction | Table alias auto-detected from FROM clause; string literals and subqueries handled correctly | @@ -52,14 +52,14 @@ Same signatures returning `CosmosPagedIterable`, delegating to the async cont ### Step 3: Internal orchestration (RxDocumentClientImpl) 1. Resolve collection metadata + PK definition from cache. -2. Fetch routing map from `partitionKeyRangeCache` **in parallel with** custom query validation (Step 4). +2. Fetch routing map from `partitionKeyRangeCache`. 3. For each `PartitionKey`: - Compute effective partition key (EPK). - Full PK → `getRangeByEffectivePartitionKey()` (single range). - Partial HPK → compute EPK prefix range → `getOverlappingRanges()` (multiple ranges). **Note:** partial HPK intentionally fans out to multiple physical partitions. 4. Group PK values by `PartitionKeyRange`. -5. Per physical partition → split PKs into batches of `maxPksPerPartitionQuery` (configurable, default 1000). +5. Per physical partition → split PKs into batches of `maxPksPerPartitionQuery` (configurable, default 100). 6. Per batch → build `SqlQuerySpec` with PK WHERE clause (Step 5). 7. Interleave batches across physical partitions in round-robin order so that bounded concurrency prefers different physical partitions over sequential batches of the same partition. 8. Execute queries via `queryForReadMany()` with bounded concurrency (`Math.min(batchCount, cpuCount)`). @@ -67,7 +67,7 @@ Same signatures returning `CosmosPagedIterable`, delegating to the async cont ### Step 4: Custom query validation -One-time call per invocation (existing query plan caching applies). Runs **in parallel** with routing map lookup to minimize latency: +One-time call per invocation using the same query-plan retrieval path and cacheability rules as regular SDK queries. - `QueryPlanRetriever.getQueryPlanThroughGatewayAsync()` for the user query. - Reject (`IllegalArgumentException`) if: @@ -76,8 +76,11 @@ One-time call per invocation (existing query plan caching applies). Runs **in pa - `queryInfo.hasOrderBy()` - `queryInfo.hasDistinct()` - `queryInfo.hasDCount()` + - `queryInfo.hasOffset()` + - `queryInfo.hasLimit()` - `queryInfo.hasNonStreamingOrderBy()` - `partitionedQueryExecutionInfo.hasHybridSearchQueryInfo()` + - query plan details are unavailable (`queryInfo == null`) ### Step 5: Query construction @@ -115,7 +118,7 @@ New method `readManyByPartitionKey` added directly to `AsyncDocumentClient` inte ### Step 7: Configuration -New configurable batch size via system property `COSMOS.READ_MANY_BY_PK_MAX_BATCH_SIZE` or environment variable `COSMOS_READ_MANY_BY_PK_MAX_BATCH_SIZE` (default: 1000, minimum: 1). Follows existing `Configs` patterns. +New configurable batch size via system property `COSMOS.READ_MANY_BY_PK_MAX_BATCH_SIZE` or environment variable `COSMOS_READ_MANY_BY_PK_MAX_BATCH_SIZE` (default: 100, minimum: 1). Follows existing `Configs` patterns. ## Phase 2 — Spark Connector (`azure-cosmos-spark_3`) @@ -123,13 +126,14 @@ New configurable batch size via system property `COSMOS.READ_MANY_BY_PK_MAX_BATC - Input: partition key value (single value or Seq for hierarchical PKs). - Output: serialized PK string in format `pk([...json...])`. -- **Null handling:** Throws on null input (Scala convention; callers should filter nulls upstream). +- **Null handling:** Null input is serialized as a JSON-null partition key component. If callers need `PartitionKey.NONE` semantics they must use the schema-matched path with `spark.cosmos.read.readManyByPk.nullHandling=None`, which is only supported for single-path partition keys. ### Step 9: PK-only serialization helper `CosmosPartitionKeyHelper`: - `getCosmosPartitionKeyValueString(pkValues: List[Object]): String` — serialize to `pk([...])` format. - `tryParsePartitionKey(serialized: String): Option[PartitionKey]` — deserialize; returns `None` for malformed input including invalid JSON (wrapped in `scala.util.Try`). +- When `spark.cosmos.read.readManyByPk.nullHandling=None` is used, hierarchical partition keys with null components are rejected with a clear error because `PartitionKey.NONE` cannot be used with multiple paths. ### Step 10: `CosmosItemsDataSource.readManyByPartitionKey` @@ -137,11 +141,13 @@ Static entry points that accept a DataFrame and Cosmos config. PK extraction sup 1. **UDF-produced column**: DataFrame contains `_partitionKeyIdentity` column (from `GetCosmosPartitionKeyValue` UDF). 2. **Schema-matched columns**: DataFrame columns match the container's PK paths. +Nested partition key paths are not resolved automatically from DataFrame columns and must use the UDF-produced `_partitionKeyIdentity` column. + Falls back with `IllegalArgumentException` if neither mode is possible. ### Step 11: `CosmosReadManyByPartitionKeyReader` -Orchestrator that resolves schema, initializes and broadcasts client state to executors, then maps each Spark partition to an `ItemsPartitionReaderWithReadManyByPartitionKey`. +Orchestrator that resolves schema, initializes and broadcasts client state to executors, then maps each Spark partition to an `ItemsPartitionReaderWithReadManyByPartitionKey`. The wrapper iterator closes the reader deterministically on exhaustion, on failures, and via Spark task-completion callbacks. ### Step 12: `ItemsPartitionReaderWithReadManyByPartitionKey` @@ -166,4 +172,5 @@ Spark `PartitionReader[InternalRow]` that: - Batch size validation: temporarily lowered batch size to exercise batching/interleaving logic. - Null/empty PK list rejection (eager validation). - Spark connector: `ItemsPartitionReaderWithReadManyByPartitionKey` with known PK values and non-existent PKs. +- Spark public API: nested partition key containers require `_partitionKeyIdentity` and succeed when populated via `GetCosmosPartitionKeyValue`. - `CosmosPartitionKeyHelper`: single/HPK roundtrip, case insensitivity, malformed input. From f34270dc97237a8b423806e2d43fc4d6e4319de9 Mon Sep 17 00:00:00 2001 From: Fabian Meiswinkel Date: Fri, 17 Apr 2026 18:24:40 +0000 Subject: [PATCH 25/64] Addressing code review comments --- .../spark/CosmosPartitionKeyHelper.scala | 5 +++- .../CosmosReadManyByPartitionKeyReader.scala | 5 +--- .../spark/CosmosPartitionKeyHelperSpec.scala | 8 +++++++ .../com/azure/cosmos/CosmosTracerTest.java | 23 +++++++++++++++++++ 4 files changed, 36 insertions(+), 5 deletions(-) diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosPartitionKeyHelper.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosPartitionKeyHelper.scala index 29831e9a6fed..84c3f2fadaf2 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosPartitionKeyHelper.scala +++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosPartitionKeyHelper.scala @@ -71,7 +71,10 @@ private[spark] object CosmosPartitionKeyHelper extends BasicLoggingTrait { case s: String => builder.add(s) case n: java.lang.Number => builder.add(n.doubleValue()) case b: java.lang.Boolean => builder.add(b.booleanValue()) - case other => builder.add(other.toString) + case other => + throw new IllegalArgumentException( + s"Unsupported partition key component type '${other.getClass.getName}' with value '$other'. " + + "Supported types are String, Number (integral or floating-point), Boolean, and null.") } Some(builder.build()) } else { diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosReadManyByPartitionKeyReader.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosReadManyByPartitionKeyReader.scala index bd6fa5be8bdf..796aeea52b91 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosReadManyByPartitionKeyReader.scala +++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosReadManyByPartitionKeyReader.scala @@ -64,10 +64,7 @@ private[spark] class CosmosReadManyByPartitionKeyReader( cosmosContainerConfig, clientCacheItems(0).get, clientCacheItems(1)) - // Warm-up readItem: intentionally issues a lookup for a random id/partition-key pair - // on the driver so that the collection/routing-map caches are populated before we serialize - // the client state and broadcast it to executors. This costs ~1 RU + 1 RTT per broadcast build - // (expected 404) but avoids every executor doing the same lookup in parallel on first use. + // Warm-up readItem: intentionally issues a lookup for a random id/partition-key pair // on the driver so that the collection/routing-map caches are populated before we serialize // the client state and broadcast it to executors. This costs ~1 RU + 1 RTT per broadcast build diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/CosmosPartitionKeyHelperSpec.scala b/sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/CosmosPartitionKeyHelperSpec.scala index c81113d3a190..ba7b37a27065 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/CosmosPartitionKeyHelperSpec.scala +++ b/sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/CosmosPartitionKeyHelperSpec.scala @@ -95,6 +95,14 @@ class CosmosPartitionKeyHelperSpec extends UnitSpec { pk.get shouldEqual PartitionKey.NONE } + it should "throw for unsupported component types in the null-handling builder path" in { + val error = the[IllegalArgumentException] thrownBy { + CosmosPartitionKeyHelper.tryParsePartitionKey("pk([null,{\"nested\":\"value\"}])", treatNullAsNone = false) + } + + error.getMessage should include("Unsupported partition key component type") + error.getMessage should include("java.util.LinkedHashMap") + } it should "throw a clear error when None nullHandling is used for hierarchical partition keys" in { val error = the[IllegalArgumentException] thrownBy { CosmosPartitionKeyHelper.tryParsePartitionKey("pk([\"Redmond\",null])", treatNullAsNone = true) diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/CosmosTracerTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/CosmosTracerTest.java index 06de8524bcbd..9a59e7af8d7c 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/CosmosTracerTest.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/CosmosTracerTest.java @@ -920,6 +920,29 @@ public void cosmosAsyncContainer( "readMany", samplingRate); mockTracer.reset(); + List partitionKeys = createdDocs + .stream() + .map(CosmosItemIdentity::getPartitionKey) + .collect(Collectors.toList()); + feedItemResponse = cosmosAsyncContainer + .readManyByPartitionKey(partitionKeys, ObjectNode.class) + .byPage(1) + .blockFirst(); + assertThat(feedItemResponse).isNotNull(); + assertThat(feedItemResponse.getResults()).isNotEmpty(); + verifyTracerAttributes( + mockTracer, + "readManyByPartitionKey." + cosmosAsyncContainer.getId(), + cosmosAsyncDatabase.getId(), + cosmosAsyncContainer.getId(), + feedItemResponse.getCosmosDiagnostics(), + null, + useLegacyTracing, + enableRequestLevelTracing, + forceThresholdViolations, + "readManyByPartitionKey", + samplingRate); + mockTracer.reset(); } @Test(groups = { "fast", "simple" }, timeOut = 10 * TIMEOUT) From f13de237d75323c8e1e73a6b548e40189e1125cc Mon Sep 17 00:00:00 2001 From: Fabian Meiswinkel Date: Mon, 20 Apr 2026 09:35:23 +0000 Subject: [PATCH 26/64] Fixing nested PK for readMany / readAllItems --- .../com/azure/cosmos/CosmosMultiHashTest.java | 64 +++++++++++++++++++ ...ReadManyByPartitionKeyQueryHelperTest.java | 18 +++--- .../PartitionKeyQueryHelper.java | 27 ++++++++ .../ReadManyByPartitionKeyQueryHelper.java | 10 --- .../implementation/RxDocumentClientImpl.java | 6 +- 5 files changed, 103 insertions(+), 22 deletions(-) create mode 100644 sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/PartitionKeyQueryHelper.java diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/CosmosMultiHashTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/CosmosMultiHashTest.java index 5275af383f2a..5ece83be3156 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/CosmosMultiHashTest.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/CosmosMultiHashTest.java @@ -33,6 +33,7 @@ import java.time.Duration; import java.util.ArrayList; +import java.util.Collections; import java.util.List; import java.util.Map; import java.util.Objects; @@ -51,6 +52,7 @@ public class CosmosMultiHashTest extends TestSuiteBase { private CosmosClient client; private CosmosDatabase createdDatabase; private CosmosContainer createdMultiHashContainer; + private CosmosContainer createdNestedPathContainer; @Factory(dataProvider = "clientBuilders") public CosmosMultiHashTest(CosmosClientBuilder clientBuilder) { @@ -79,6 +81,19 @@ public void before_CosmosMultiHashTest() { //MultiHash collection read createdMultiHashContainer = createdDatabase.getContainer(collectionName); + + String nestedPathCollectionName = UUID.randomUUID().toString(); + PartitionKeyDefinition nestedPartitionKeyDefinition = new PartitionKeyDefinition(); + nestedPartitionKeyDefinition.setKind(PartitionKind.HASH); + nestedPartitionKeyDefinition.setVersion(PartitionKeyDefinitionVersion.V2); + nestedPartitionKeyDefinition.setPaths(Collections.singletonList("/address/city")); + + CosmosContainerProperties nestedPathContainerProperties = getCollectionDefinition( + nestedPathCollectionName, + nestedPartitionKeyDefinition); + + createdDatabase.createContainer(nestedPathContainerProperties); + createdNestedPathContainer = createdDatabase.getContainer(nestedPathCollectionName); } @AfterClass(groups = {"emulator"}, timeOut = SHUTDOWN_TIMEOUT, alwaysRun = true) @@ -155,6 +170,42 @@ private void validateResponse(FeedResponse response, .collect(Collectors.toList()) ); } + @Test(groups = { "emulator" }, timeOut = TIMEOUT) + public void readManySupportsNestedPartitionKeyPaths() { + String city = "nested-readmany-" + UUID.randomUUID(); + + ObjectNode firstItem = createNestedPartitionKeyItem(UUID.randomUUID().toString(), city, "downtown"); + ObjectNode secondItem = createNestedPartitionKeyItem(UUID.randomUUID().toString(), city, "overlake"); + + createdNestedPathContainer.createItem(firstItem); + createdNestedPathContainer.createItem(secondItem); + + List itemList = new ArrayList<>(); + itemList.add(new CosmosItemIdentity(new PartitionKey(city), firstItem.get("id").asText())); + itemList.add(new CosmosItemIdentity(new PartitionKey(city), secondItem.get("id").asText())); + + FeedResponse documentFeedResponse = createdNestedPathContainer.readMany(itemList, ObjectNode.class); + validateResponse(documentFeedResponse, itemList); + } + + @Test(groups = { "emulator" }, timeOut = TIMEOUT) + public void readAllItemsSupportsNestedPartitionKeyPaths() { + String city = "nested-readall-" + UUID.randomUUID(); + + ObjectNode firstItem = createNestedPartitionKeyItem(UUID.randomUUID().toString(), city, "north"); + ObjectNode secondItem = createNestedPartitionKeyItem(UUID.randomUUID().toString(), city, "south"); + ObjectNode otherItem = createNestedPartitionKeyItem(UUID.randomUUID().toString(), city + "-other", "east"); + + createdNestedPathContainer.createItem(firstItem); + createdNestedPathContainer.createItem(secondItem); + createdNestedPathContainer.createItem(otherItem); + + CosmosPagedIterable readAllResults = + createdNestedPathContainer.readAllItems(new PartitionKey(city), ObjectNode.class); + + assertThat(readAllResults.stream().map(item -> item.get("id").asText()).collect(Collectors.toList())) + .containsExactlyInAnyOrder(firstItem.get("id").asText(), secondItem.get("id").asText()); + } @Test(groups = { "emulator" }, timeOut = TIMEOUT) private void validateDocCRUDAndQuery() throws InterruptedException { @@ -550,6 +601,19 @@ private void extractPartitionKeyFromDocumentTests() { } } + private ObjectNode createNestedPartitionKeyItem(String id, String city, String neighborhood) { + ObjectNode item = new ObjectNode(JSON_NODE_FACTORY_INSTANCE); + item.put("id", id); + item.put("type", "nestedPartitionKeyRegression"); + + ObjectNode address = new ObjectNode(JSON_NODE_FACTORY_INSTANCE); + address.put("city", city); + address.put("neighborhood", neighborhood); + item.set("address", address); + + return item; + } + private ArrayList createItems() { ArrayList docs = new ArrayList<>(); diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelperTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelperTest.java index 95c109ba025f..11b85904e6b4 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelperTest.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelperTest.java @@ -18,7 +18,6 @@ import java.util.Arrays; import java.util.Collections; import java.util.List; -import java.util.stream.Collectors; import static org.assertj.core.api.Assertions.assertThat; @@ -393,10 +392,17 @@ public void extractAlias_containerWithHaving() { "SELECT c.status, COUNT(1) FROM c GROUP BY c.status HAVING COUNT(1) > 1")).isEqualTo("c"); } + @Test(groups = { "unit" }) + public void createSelectors_nestedPath() { + PartitionKeyDefinition pkDef = createSinglePkDefinition("/address/city"); + + assertThat(PartitionKeyQueryHelper.createPkSelectors(pkDef)) + .containsExactly("[\"address\"][\"city\"]"); + } + //endregion //region helpers - private PartitionKeyDefinition createSinglePkDefinition(String path) { PartitionKeyDefinition pkDef = new PartitionKeyDefinition(); pkDef.setKind(PartitionKind.HASH); @@ -412,14 +418,8 @@ private PartitionKeyDefinition createMultiHashPkDefinition(String... paths) { pkDef.setPaths(Arrays.asList(paths)); return pkDef; } - private List createSelectors(PartitionKeyDefinition pkDef) { - return pkDef.getPaths() - .stream() - .map(pathPart -> pathPart.substring(1)) // skip starting / - .map(pathPart -> pathPart.replace("\"", "\\")) - .map(part -> "[\"" + part + "\"]") - .collect(Collectors.toList()); + return PartitionKeyQueryHelper.createPkSelectors(pkDef); } //endregion diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/PartitionKeyQueryHelper.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/PartitionKeyQueryHelper.java new file mode 100644 index 000000000000..c346e931bb8f --- /dev/null +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/PartitionKeyQueryHelper.java @@ -0,0 +1,27 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. +package com.azure.cosmos.implementation; + +import com.azure.cosmos.models.PartitionKeyDefinition; + +import java.util.List; +import java.util.stream.Collectors; + +/** + * Helper for constructing SQL partition key selector fragments from partition key definitions. + */ +final class PartitionKeyQueryHelper { + + private PartitionKeyQueryHelper() { + } + + static List createPkSelectors(PartitionKeyDefinition partitionKeyDefinition) { + return partitionKeyDefinition.getPaths() + .stream() + .map(PathParser::getPathParts) + .map(pathParts -> pathParts.stream() + .map(pathPart -> "[\"" + pathPart.replace("\"", "\\") + "\"]") + .collect(Collectors.joining())) + .collect(Collectors.toList()); + } +} \ No newline at end of file diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelper.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelper.java index 0bdb867dc3ee..c8d3094fedfd 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelper.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelper.java @@ -161,16 +161,6 @@ public static SqlQuerySpec createReadManyByPkQuerySpec( return new SqlQuerySpec(finalQuery, parameters); } - static List createPkSelectors(PartitionKeyDefinition partitionKeyDefinition) { - return partitionKeyDefinition.getPaths() - .stream() - .map(PathParser::getPathParts) - .map(pathParts -> pathParts.stream() - .map(pathPart -> "[\"" + pathPart.replace("\"", "\\") + "\"]") - .collect(Collectors.joining())) - .collect(Collectors.toList()); - } - /** * Extracts the table/collection alias from a SQL query's FROM clause. * Handles: "SELECT * FROM c", "SELECT x.id FROM x WHERE ...", "SELECT * FROM root r", etc. diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java index 1126008946a3..57c8d7b0bd22 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java @@ -4485,7 +4485,7 @@ private Flux> readManyByPartitionKey( Map> partitionRangePkMap = groupPartitionKeysByPhysicalPartition(partitionKeys, pkDefinition, routingMap); - List partitionKeySelectors = ReadManyByPartitionKeyQueryHelper.createPkSelectors(pkDefinition); + List partitionKeySelectors = PartitionKeyQueryHelper.createPkSelectors(pkDefinition); String baseQueryText; List baseParameters; @@ -4686,7 +4686,7 @@ private Map getRangeQueryMap( //TODO: Optimise this to include all types of partitionkeydefinitions. ex: c["prop1./ab"]["key1"] Map rangeQueryMap = new HashMap<>(); - List partitionKeySelectors = ReadManyByPartitionKeyQueryHelper.createPkSelectors(partitionKeyDefinition); + List partitionKeySelectors = PartitionKeyQueryHelper.createPkSelectors(partitionKeyDefinition); for(Map.Entry> entry: partitionRangeItemKeyMap.entrySet()) { SqlQuerySpec sqlQuerySpec; @@ -5301,7 +5301,7 @@ public Flux> readAllDocuments( } PartitionKeyDefinition pkDefinition = collection.getPartitionKey(); - List partitionKeySelectors = ReadManyByPartitionKeyQueryHelper.createPkSelectors(pkDefinition); + List partitionKeySelectors = PartitionKeyQueryHelper.createPkSelectors(pkDefinition); SqlQuerySpec querySpec = createLogicalPartitionScanQuerySpec(partitionKey, partitionKeySelectors); String resourceLink = parentResourceLinkToQueryLink(collectionLink, ResourceType.Document); From ce8fca37a59d20781c1fc9d917b346e10ca46172 Mon Sep 17 00:00:00 2001 From: Fabian Meiswinkel Date: Mon, 20 Apr 2026 10:38:00 +0000 Subject: [PATCH 27/64] Fixed changelogs --- sdk/cosmos/azure-cosmos-spark_3-3_2-12/CHANGELOG.md | 1 + sdk/cosmos/azure-cosmos-spark_3-4_2-12/CHANGELOG.md | 1 + sdk/cosmos/azure-cosmos-spark_3-5_2-12/CHANGELOG.md | 1 + sdk/cosmos/azure-cosmos-spark_3-5_2-13/CHANGELOG.md | 1 + sdk/cosmos/azure-cosmos-spark_4-0_2-13/CHANGELOG.md | 1 + sdk/cosmos/azure-cosmos/CHANGELOG.md | 1 + 6 files changed, 6 insertions(+) diff --git a/sdk/cosmos/azure-cosmos-spark_3-3_2-12/CHANGELOG.md b/sdk/cosmos/azure-cosmos-spark_3-3_2-12/CHANGELOG.md index 46cb9d880f87..4987a866c8ab 100644 --- a/sdk/cosmos/azure-cosmos-spark_3-3_2-12/CHANGELOG.md +++ b/sdk/cosmos/azure-cosmos-spark_3-3_2-12/CHANGELOG.md @@ -3,6 +3,7 @@ ### 4.48.0-beta.1 (Unreleased) #### Features Added +* Added new `CosmosItemsDataSource.readManyByPartitionKey` Spark function to execute bulk queries by a list of pk-values with better efficiency. Configure null handling via `spark.cosmos.read.readManyByPk.nullHandling` - default `Null` treats a null PK column as JSON null (`addNullValue`), `None` treats it as `PartitionKey.NONE` (`addNoneValue` / `NOT IS_DEFINED`). These route to different physical partitions - picking the wrong mode silently returns zero rows. Also fixes nested partition keys for `readMany` and `readAllItems` See [PR 48801](https://github.com/Azure/azure-sdk-for-java/pull/48801) #### Breaking Changes diff --git a/sdk/cosmos/azure-cosmos-spark_3-4_2-12/CHANGELOG.md b/sdk/cosmos/azure-cosmos-spark_3-4_2-12/CHANGELOG.md index d1f06f57031a..b1f863457ee4 100644 --- a/sdk/cosmos/azure-cosmos-spark_3-4_2-12/CHANGELOG.md +++ b/sdk/cosmos/azure-cosmos-spark_3-4_2-12/CHANGELOG.md @@ -3,6 +3,7 @@ ### 4.48.0-beta.1 (Unreleased) #### Features Added +* Added new `CosmosItemsDataSource.readManyByPartitionKey` Spark function to execute bulk queries by a list of pk-values with better efficiency. Configure null handling via `spark.cosmos.read.readManyByPk.nullHandling` - default `Null` treats a null PK column as JSON null (`addNullValue`), `None` treats it as `PartitionKey.NONE` (`addNoneValue` / `NOT IS_DEFINED`). These route to different physical partitions - picking the wrong mode silently returns zero rows. Also fixes nested partition keys for `readMany` and `readAllItems` See [PR 48801](https://github.com/Azure/azure-sdk-for-java/pull/48801) #### Breaking Changes diff --git a/sdk/cosmos/azure-cosmos-spark_3-5_2-12/CHANGELOG.md b/sdk/cosmos/azure-cosmos-spark_3-5_2-12/CHANGELOG.md index 93bb3fa96a75..276505324d77 100644 --- a/sdk/cosmos/azure-cosmos-spark_3-5_2-12/CHANGELOG.md +++ b/sdk/cosmos/azure-cosmos-spark_3-5_2-12/CHANGELOG.md @@ -3,6 +3,7 @@ ### 4.48.0-beta.1 (Unreleased) #### Features Added +* Added new `CosmosItemsDataSource.readManyByPartitionKey` Spark function to execute bulk queries by a list of pk-values with better efficiency. Configure null handling via `spark.cosmos.read.readManyByPk.nullHandling` - default `Null` treats a null PK column as JSON null (`addNullValue`), `None` treats it as `PartitionKey.NONE` (`addNoneValue` / `NOT IS_DEFINED`). These route to different physical partitions - picking the wrong mode silently returns zero rows. Also fixes nested partition keys for `readMany` and `readAllItems` See [PR 48801](https://github.com/Azure/azure-sdk-for-java/pull/48801) #### Breaking Changes diff --git a/sdk/cosmos/azure-cosmos-spark_3-5_2-13/CHANGELOG.md b/sdk/cosmos/azure-cosmos-spark_3-5_2-13/CHANGELOG.md index e073483e335c..0b420558a657 100644 --- a/sdk/cosmos/azure-cosmos-spark_3-5_2-13/CHANGELOG.md +++ b/sdk/cosmos/azure-cosmos-spark_3-5_2-13/CHANGELOG.md @@ -3,6 +3,7 @@ ### 4.48.0-beta.1 (Unreleased) #### Features Added +* Added new `CosmosItemsDataSource.readManyByPartitionKey` Spark function to execute bulk queries by a list of pk-values with better efficiency. Configure null handling via `spark.cosmos.read.readManyByPk.nullHandling` - default `Null` treats a null PK column as JSON null (`addNullValue`), `None` treats it as `PartitionKey.NONE` (`addNoneValue` / `NOT IS_DEFINED`). These route to different physical partitions - picking the wrong mode silently returns zero rows. Also fixes nested partition keys for `readMany` and `readAllItems` See [PR 48801](https://github.com/Azure/azure-sdk-for-java/pull/48801) #### Breaking Changes diff --git a/sdk/cosmos/azure-cosmos-spark_4-0_2-13/CHANGELOG.md b/sdk/cosmos/azure-cosmos-spark_4-0_2-13/CHANGELOG.md index 33fa94dda6c0..2527bec5d136 100644 --- a/sdk/cosmos/azure-cosmos-spark_4-0_2-13/CHANGELOG.md +++ b/sdk/cosmos/azure-cosmos-spark_4-0_2-13/CHANGELOG.md @@ -3,6 +3,7 @@ ### 4.48.0-beta.1 (Unreleased) #### Features Added +* Added new `CosmosItemsDataSource.readManyByPartitionKey` Spark function to execute bulk queries by a list of pk-values with better efficiency. Configure null handling via `spark.cosmos.read.readManyByPk.nullHandling` - default `Null` treats a null PK column as JSON null (`addNullValue`), `None` treats it as `PartitionKey.NONE` (`addNoneValue` / `NOT IS_DEFINED`). These route to different physical partitions - picking the wrong mode silently returns zero rows. Also fixes nested partition keys for `readMany` and `readAllItems` See [PR 48801](https://github.com/Azure/azure-sdk-for-java/pull/48801) #### Breaking Changes diff --git a/sdk/cosmos/azure-cosmos/CHANGELOG.md b/sdk/cosmos/azure-cosmos/CHANGELOG.md index 8cbf52e3f916..f4dfae54223a 100644 --- a/sdk/cosmos/azure-cosmos/CHANGELOG.md +++ b/sdk/cosmos/azure-cosmos/CHANGELOG.md @@ -4,6 +4,7 @@ #### Features Added * Added support for change feed with `startFrom` point-in-time on merged partitions by enabling the `CHANGE_FEED_WITH_START_TIME_POST_MERGE` SDK capability. - See [PR 48752](https://github.com/Azure/azure-sdk-for-java/pull/48752) +* Added new `CosmosItemsDataSource.readManyByPartitionKey` Spark function to execute bulk queries by a list of pk-values with better efficiency. Configure null handling via `spark.cosmos.read.readManyByPk.nullHandling` - default `Null` treats a null PK column as JSON null (`addNullValue`), `None` treats it as `PartitionKey.NONE` (`addNoneValue` / `NOT IS_DEFINED`). These route to different physical partitions - picking the wrong mode silently returns zero rows. Also fixes nested partition keys for `readMany` and `readAllItems` See [PR 48801](https://github.com/Azure/azure-sdk-for-java/pull/48801) #### Breaking Changes From 992ef1aafc8495c776f4c658195c14efb00ebc6c Mon Sep 17 00:00:00 2001 From: Fabian Meiswinkel Date: Mon, 20 Apr 2026 10:52:46 +0000 Subject: [PATCH 28/64] Disallow TOP for streamable queries in readManyByPartitionKey-API --- .../ReadManyByPartitionKeyQueryPlanValidationTest.java | 10 ++++++++++ .../cosmos/implementation/RxDocumentClientImpl.java | 4 ++++ 2 files changed, 14 insertions(+) diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryPlanValidationTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryPlanValidationTest.java index 429d1fd8f89c..b9a864cddf65 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryPlanValidationTest.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryPlanValidationTest.java @@ -49,6 +49,16 @@ public void rejectsLimitQueryPlan() { .hasMessageContaining("LIMIT"); } + @Test(groups = { "unit" }) + public void rejectsTopQueryPlan() { + QueryInfo queryInfo = new QueryInfo(); + queryInfo.set("top", 5); + + assertThatThrownBy(() -> RxDocumentClientImpl.validateQueryPlanForReadManyByPartitionKey(createQueryPlan(queryInfo, null))) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("TOP"); + } + @Test(groups = { "unit" }) public void rejectsHybridSearchQueryPlanWithoutDereferencingNullQueryInfo() { assertThatThrownBy(() -> RxDocumentClientImpl.validateQueryPlanForReadManyByPartitionKey( diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java index 57c8d7b0bd22..c1ca60c3f889 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java @@ -4629,6 +4629,10 @@ static void validateQueryPlanForReadManyByPartitionKey(PartitionedQueryExecution throw new IllegalArgumentException( "Custom query for readMany by partition key must not contain LIMIT."); } + if (queryInfo.hasTop()) { + throw new IllegalArgumentException( + "Custom query for readMany by partition key must not contain TOP."); + } if (queryInfo.hasNonStreamingOrderBy()) { throw new IllegalArgumentException( "Custom query for readMany by partition key must not contain non-streaming ORDER BY."); From 84af9c34c82579c0bb2c837f5d3642680ac86c24 Mon Sep 17 00:00:00 2001 From: Fabian Meiswinkel Date: Mon, 20 Apr 2026 11:24:59 +0000 Subject: [PATCH 29/64] Fixing quote recognition in pk for readManyByPartitionKey API --- .../ReadManyByPartitionKeyQueryHelperTest.java | 18 ++++++++++++++++++ ...yByPartitionKeyQueryPlanValidationTest.java | 2 +- .../PartitionKeyQueryHelper.java | 4 ++-- 3 files changed, 21 insertions(+), 3 deletions(-) diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelperTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelperTest.java index 11b85904e6b4..85712945a72f 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelperTest.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelperTest.java @@ -400,6 +400,24 @@ public void createSelectors_nestedPath() { .containsExactly("[\"address\"][\"city\"]"); } + @Test(groups = { "unit" }) + public void createSelectors_escapesQuotesInPathParts() { + // Verify the escaping logic directly: a path part containing a double quote + // must produce \" in the selector, not just a bare backslash. + // Use an unquoted simple path /my so PathParser returns "my" cleanly, + // then verify a path whose PathParser output contains a quote character. + // PathParser for /"my\"field" yields token: my\"field (literal \, ") + PartitionKeyDefinition pkDef = createSinglePkDefinition("/\"my\\\"field\""); + + List selectors = PartitionKeyQueryHelper.createPkSelectors(pkDef); + assertThat(selectors).hasSize(1); + String selector = selectors.get(0); + // Must contain the escaped quote sequence \" + assertThat(selector).contains("\\\""); + // Must NOT be just ["my\field"] (old bug: quote replaced with bare backslash) + assertThat(selector).isNotEqualTo("[\"my\\field\"]"); + } + //endregion //region helpers diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryPlanValidationTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryPlanValidationTest.java index b9a864cddf65..11b69b955893 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryPlanValidationTest.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryPlanValidationTest.java @@ -88,4 +88,4 @@ private PartitionedQueryExecutionInfo createQueryPlan(QueryInfo queryInfo, Hybri return new PartitionedQueryExecutionInfo(content, null); } -} \ No newline at end of file +} diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/PartitionKeyQueryHelper.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/PartitionKeyQueryHelper.java index c346e931bb8f..459e03c8fce3 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/PartitionKeyQueryHelper.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/PartitionKeyQueryHelper.java @@ -20,8 +20,8 @@ static List createPkSelectors(PartitionKeyDefinition partitionKeyDefinit .stream() .map(PathParser::getPathParts) .map(pathParts -> pathParts.stream() - .map(pathPart -> "[\"" + pathPart.replace("\"", "\\") + "\"]") + .map(pathPart -> "[\"" + pathPart.replace("\"", "\\\"") + "\"]") .collect(Collectors.joining())) .collect(Collectors.toList()); } -} \ No newline at end of file +} From b0b73497e78b3551991450df854a4b6dd9cfc2f1 Mon Sep 17 00:00:00 2001 From: Fabian Meiswinkel Date: Mon, 20 Apr 2026 11:45:25 +0000 Subject: [PATCH 30/64] Fixing confusing side effects when calling GetCurrentRow --- ...tionReaderWithReadManyByPartitionKey.scala | 22 +++++++++++++++---- 1 file changed, 18 insertions(+), 4 deletions(-) diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/ItemsPartitionReaderWithReadManyByPartitionKey.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/ItemsPartitionReaderWithReadManyByPartitionKey.scala index 5c994cc2f7f8..822b52c88341 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/ItemsPartitionReaderWithReadManyByPartitionKey.scala +++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/ItemsPartitionReaderWithReadManyByPartitionKey.scala @@ -247,13 +247,27 @@ private[spark] case class ItemsPartitionReaderWithReadManyByPartitionKey } } - override def next(): Boolean = getOrCreateIterator.hasNext + private var currentRow: Option[Row] = None + + override def next(): Boolean = { + val hasMore = getOrCreateIterator.hasNext + if (hasMore) { + currentRow = Some(getOrCreateIterator.next().row) + } else { + currentRow = None + } + hasMore + } override def get(): InternalRow = { - cosmosRowConverter.fromRowToInternalRow(getOrCreateIterator.next().row, rowSerializer) + cosmosRowConverter.fromRowToInternalRow( + currentRow.getOrElse(throw new NoSuchElementException("No current row - next() must be called first")), + rowSerializer) } - def getCurrentRow(): Row = getOrCreateIterator.next().row + def getCurrentRow(): Row = { + currentRow.getOrElse(throw new NoSuchElementException("No current row - next() must be called first")) + } override def close(): Unit = { if (isClosed.compareAndSet(false, true)) { @@ -266,4 +280,4 @@ private[spark] case class ItemsPartitionReaderWithReadManyByPartitionKey } } } -} \ No newline at end of file +} From 42a2fb820cadf012e752e7e5a84fd0dc22f79b3e Mon Sep 17 00:00:00 2001 From: Fabian Meiswinkel Date: Mon, 20 Apr 2026 16:34:18 +0000 Subject: [PATCH 31/64] Fixed code review comments --- .../azure-cosmos-spark_3-3_2-12/CHANGELOG.md | 2 +- .../azure-cosmos-spark_3-4_2-12/CHANGELOG.md | 2 +- .../cosmos/spark/SparkE2EQueryITest.scala | 3 ++- .../azure-cosmos-spark_3-5_2-12/CHANGELOG.md | 2 +- .../azure-cosmos-spark_3-5_2-13/CHANGELOG.md | 2 +- .../azure-cosmos-spark_4-0_2-13/CHANGELOG.md | 2 +- .../ReadManyByPartitionKeyQueryHelperTest.java | 18 ++++++++++++++++++ sdk/cosmos/azure-cosmos/CHANGELOG.md | 2 +- .../com/azure/cosmos/CosmosAsyncContainer.java | 2 ++ .../java/com/azure/cosmos/CosmosContainer.java | 2 ++ .../ReadManyByPartitionKeyQueryHelper.java | 8 ++++++++ 11 files changed, 38 insertions(+), 7 deletions(-) diff --git a/sdk/cosmos/azure-cosmos-spark_3-3_2-12/CHANGELOG.md b/sdk/cosmos/azure-cosmos-spark_3-3_2-12/CHANGELOG.md index 4987a866c8ab..bf15bf284bbd 100644 --- a/sdk/cosmos/azure-cosmos-spark_3-3_2-12/CHANGELOG.md +++ b/sdk/cosmos/azure-cosmos-spark_3-3_2-12/CHANGELOG.md @@ -3,7 +3,7 @@ ### 4.48.0-beta.1 (Unreleased) #### Features Added -* Added new `CosmosItemsDataSource.readManyByPartitionKey` Spark function to execute bulk queries by a list of pk-values with better efficiency. Configure null handling via `spark.cosmos.read.readManyByPk.nullHandling` - default `Null` treats a null PK column as JSON null (`addNullValue`), `None` treats it as `PartitionKey.NONE` (`addNoneValue` / `NOT IS_DEFINED`). These route to different physical partitions - picking the wrong mode silently returns zero rows. Also fixes nested partition keys for `readMany` and `readAllItems` See [PR 48801](https://github.com/Azure/azure-sdk-for-java/pull/48801) +* Added new `CosmosItemsDataSource.readManyByPartitionKey` Spark function to execute bulk queries by a list of pk-values with better efficiency. Configure null handling via `spark.cosmos.read.readManyByPk.nullHandling` - default `Null` treats a null PK column as JSON null (`addNullValue`), `None` treats it as `PartitionKey.NONE` (`addNoneValue` / `NOT IS_DEFINED`). These route to different physical partitions - picking the wrong mode silently returns zero rows. See [PR 48801](https://github.com/Azure/azure-sdk-for-java/pull/48801) #### Breaking Changes diff --git a/sdk/cosmos/azure-cosmos-spark_3-4_2-12/CHANGELOG.md b/sdk/cosmos/azure-cosmos-spark_3-4_2-12/CHANGELOG.md index b1f863457ee4..72e8fdcf99ed 100644 --- a/sdk/cosmos/azure-cosmos-spark_3-4_2-12/CHANGELOG.md +++ b/sdk/cosmos/azure-cosmos-spark_3-4_2-12/CHANGELOG.md @@ -3,7 +3,7 @@ ### 4.48.0-beta.1 (Unreleased) #### Features Added -* Added new `CosmosItemsDataSource.readManyByPartitionKey` Spark function to execute bulk queries by a list of pk-values with better efficiency. Configure null handling via `spark.cosmos.read.readManyByPk.nullHandling` - default `Null` treats a null PK column as JSON null (`addNullValue`), `None` treats it as `PartitionKey.NONE` (`addNoneValue` / `NOT IS_DEFINED`). These route to different physical partitions - picking the wrong mode silently returns zero rows. Also fixes nested partition keys for `readMany` and `readAllItems` See [PR 48801](https://github.com/Azure/azure-sdk-for-java/pull/48801) +* Added new `CosmosItemsDataSource.readManyByPartitionKey` Spark function to execute bulk queries by a list of pk-values with better efficiency. Configure null handling via `spark.cosmos.read.readManyByPk.nullHandling` - default `Null` treats a null PK column as JSON null (`addNullValue`), `None` treats it as `PartitionKey.NONE` (`addNoneValue` / `NOT IS_DEFINED`). These route to different physical partitions - picking the wrong mode silently returns zero rows. See [PR 48801](https://github.com/Azure/azure-sdk-for-java/pull/48801) #### Breaking Changes diff --git a/sdk/cosmos/azure-cosmos-spark_3-5/src/test/scala/com/azure/cosmos/spark/SparkE2EQueryITest.scala b/sdk/cosmos/azure-cosmos-spark_3-5/src/test/scala/com/azure/cosmos/spark/SparkE2EQueryITest.scala index 8476e4aa2026..bac76468dec5 100644 --- a/sdk/cosmos/azure-cosmos-spark_3-5/src/test/scala/com/azure/cosmos/spark/SparkE2EQueryITest.scala +++ b/sdk/cosmos/azure-cosmos-spark_3-5/src/test/scala/com/azure/cosmos/spark/SparkE2EQueryITest.scala @@ -111,6 +111,7 @@ class SparkE2EQueryITest rows.map(_.getAs[String]("pk")).toSet shouldEqual Set("pkA", "pkB") rows.map(_.getAs[String]("payload")).toSet shouldEqual Set("value-pkA", "value-pkB") } + "spark readManyByPartitionKey" can "require the UDF for nested partition key paths and succeed with it" in { val cosmosEndpoint = TestConfigurations.HOST val cosmosMasterKey = TestConfigurations.MASTER_KEY @@ -185,4 +186,4 @@ class SparkE2EQueryITest } // scalastyle:on multiple.string.literals -} \ No newline at end of file +} diff --git a/sdk/cosmos/azure-cosmos-spark_3-5_2-12/CHANGELOG.md b/sdk/cosmos/azure-cosmos-spark_3-5_2-12/CHANGELOG.md index 276505324d77..9e995b3b228f 100644 --- a/sdk/cosmos/azure-cosmos-spark_3-5_2-12/CHANGELOG.md +++ b/sdk/cosmos/azure-cosmos-spark_3-5_2-12/CHANGELOG.md @@ -3,7 +3,7 @@ ### 4.48.0-beta.1 (Unreleased) #### Features Added -* Added new `CosmosItemsDataSource.readManyByPartitionKey` Spark function to execute bulk queries by a list of pk-values with better efficiency. Configure null handling via `spark.cosmos.read.readManyByPk.nullHandling` - default `Null` treats a null PK column as JSON null (`addNullValue`), `None` treats it as `PartitionKey.NONE` (`addNoneValue` / `NOT IS_DEFINED`). These route to different physical partitions - picking the wrong mode silently returns zero rows. Also fixes nested partition keys for `readMany` and `readAllItems` See [PR 48801](https://github.com/Azure/azure-sdk-for-java/pull/48801) +* Added new `CosmosItemsDataSource.readManyByPartitionKey` Spark function to execute bulk queries by a list of pk-values with better efficiency. Configure null handling via `spark.cosmos.read.readManyByPk.nullHandling` - default `Null` treats a null PK column as JSON null (`addNullValue`), `None` treats it as `PartitionKey.NONE` (`addNoneValue` / `NOT IS_DEFINED`). These route to different physical partitions - picking the wrong mode silently returns zero rows. See [PR 48801](https://github.com/Azure/azure-sdk-for-java/pull/48801) #### Breaking Changes diff --git a/sdk/cosmos/azure-cosmos-spark_3-5_2-13/CHANGELOG.md b/sdk/cosmos/azure-cosmos-spark_3-5_2-13/CHANGELOG.md index 0b420558a657..51e4e5760c00 100644 --- a/sdk/cosmos/azure-cosmos-spark_3-5_2-13/CHANGELOG.md +++ b/sdk/cosmos/azure-cosmos-spark_3-5_2-13/CHANGELOG.md @@ -3,7 +3,7 @@ ### 4.48.0-beta.1 (Unreleased) #### Features Added -* Added new `CosmosItemsDataSource.readManyByPartitionKey` Spark function to execute bulk queries by a list of pk-values with better efficiency. Configure null handling via `spark.cosmos.read.readManyByPk.nullHandling` - default `Null` treats a null PK column as JSON null (`addNullValue`), `None` treats it as `PartitionKey.NONE` (`addNoneValue` / `NOT IS_DEFINED`). These route to different physical partitions - picking the wrong mode silently returns zero rows. Also fixes nested partition keys for `readMany` and `readAllItems` See [PR 48801](https://github.com/Azure/azure-sdk-for-java/pull/48801) +* Added new `CosmosItemsDataSource.readManyByPartitionKey` Spark function to execute bulk queries by a list of pk-values with better efficiency. Configure null handling via `spark.cosmos.read.readManyByPk.nullHandling` - default `Null` treats a null PK column as JSON null (`addNullValue`), `None` treats it as `PartitionKey.NONE` (`addNoneValue` / `NOT IS_DEFINED`). These route to different physical partitions - picking the wrong mode silently returns zero rows. See [PR 48801](https://github.com/Azure/azure-sdk-for-java/pull/48801) #### Breaking Changes diff --git a/sdk/cosmos/azure-cosmos-spark_4-0_2-13/CHANGELOG.md b/sdk/cosmos/azure-cosmos-spark_4-0_2-13/CHANGELOG.md index 2527bec5d136..9a58d9ec65d8 100644 --- a/sdk/cosmos/azure-cosmos-spark_4-0_2-13/CHANGELOG.md +++ b/sdk/cosmos/azure-cosmos-spark_4-0_2-13/CHANGELOG.md @@ -3,7 +3,7 @@ ### 4.48.0-beta.1 (Unreleased) #### Features Added -* Added new `CosmosItemsDataSource.readManyByPartitionKey` Spark function to execute bulk queries by a list of pk-values with better efficiency. Configure null handling via `spark.cosmos.read.readManyByPk.nullHandling` - default `Null` treats a null PK column as JSON null (`addNullValue`), `None` treats it as `PartitionKey.NONE` (`addNoneValue` / `NOT IS_DEFINED`). These route to different physical partitions - picking the wrong mode silently returns zero rows. Also fixes nested partition keys for `readMany` and `readAllItems` See [PR 48801](https://github.com/Azure/azure-sdk-for-java/pull/48801) +* Added new `CosmosItemsDataSource.readManyByPartitionKey` Spark function to execute bulk queries by a list of pk-values with better efficiency. Configure null handling via `spark.cosmos.read.readManyByPk.nullHandling` - default `Null` treats a null PK column as JSON null (`addNullValue`), `None` treats it as `PartitionKey.NONE` (`addNoneValue` / `NOT IS_DEFINED`). These route to different physical partitions - picking the wrong mode silently returns zero rows. See [PR 48801](https://github.com/Azure/azure-sdk-for-java/pull/48801) #### Breaking Changes diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelperTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelperTest.java index 85712945a72f..68ab3b0900fb 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelperTest.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelperTest.java @@ -370,6 +370,24 @@ public void singlePk_customQuery_withStringLiteralContainingParens() { assertThat(result.getQueryText()).contains("WHERE (c.msg = 'test(value)WHERE') AND ("); } + @Test(groups = { "unit" }) + public void findWhere_ignoresKeywordInsideDoubleQuotedBracketNotation() { + // c["WHERE"] uses double-quoted bracket notation — the WHERE inside quotes + // must not be matched as the real WHERE keyword. + int idx = ReadManyByPartitionKeyQueryHelper.findTopLevelWhereIndex( + "SELECT c[\"WHERE\"] FROM c WHERE c.status = 'active'"); + // The real WHERE is after "FROM c ", not inside the brackets + assertThat(idx).isEqualTo(25); + } + + @Test(groups = { "unit" }) + public void findWhere_ignoresFromInsideDoubleQuotedBracketNotation() { + // Property named "FROM" in bracket notation should not confuse the FROM-clause parser + String query = "SELECT c[\"FROM\"] FROM c WHERE c.x = 1"; + assertThat(ReadManyByPartitionKeyQueryHelper.extractTableAlias(query)).isEqualTo("c"); + assertThat(ReadManyByPartitionKeyQueryHelper.findTopLevelWhereIndex(query)).isEqualTo(24); + } + //endregion //region OFFSET/LIMIT/HAVING alias detection tests (#9) diff --git a/sdk/cosmos/azure-cosmos/CHANGELOG.md b/sdk/cosmos/azure-cosmos/CHANGELOG.md index f4dfae54223a..96f3ddeb28c0 100644 --- a/sdk/cosmos/azure-cosmos/CHANGELOG.md +++ b/sdk/cosmos/azure-cosmos/CHANGELOG.md @@ -4,7 +4,7 @@ #### Features Added * Added support for change feed with `startFrom` point-in-time on merged partitions by enabling the `CHANGE_FEED_WITH_START_TIME_POST_MERGE` SDK capability. - See [PR 48752](https://github.com/Azure/azure-sdk-for-java/pull/48752) -* Added new `CosmosItemsDataSource.readManyByPartitionKey` Spark function to execute bulk queries by a list of pk-values with better efficiency. Configure null handling via `spark.cosmos.read.readManyByPk.nullHandling` - default `Null` treats a null PK column as JSON null (`addNullValue`), `None` treats it as `PartitionKey.NONE` (`addNoneValue` / `NOT IS_DEFINED`). These route to different physical partitions - picking the wrong mode silently returns zero rows. Also fixes nested partition keys for `readMany` and `readAllItems` See [PR 48801](https://github.com/Azure/azure-sdk-for-java/pull/48801) +* Added new `readManyByPartitionKey` API on `CosmosAsyncContainer` / `CosmosContainer` to bulk-query all documents matching a list of partition key values with better efficiency than issuing individual queries. Also fixes nested partition key selector generation for `readMany` and `readAllItems`. See [PR 48801](https://github.com/Azure/azure-sdk-for-java/pull/48801) #### Breaking Changes diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java index 8260ed95b19d..6e70ecb3d940 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java @@ -1657,6 +1657,7 @@ public CosmosPagedFlux readManyByPartitionKey( * @param the type parameter * @param partitionKeys list of partition key values to read documents for * @param customQuery optional custom query for projections/additional filters (null means SELECT * FROM c) + * - should not contain WHERE clause filters for PK * @param classType class type * @return a {@link CosmosPagedFlux} containing one or several feed response pages */ @@ -1684,6 +1685,7 @@ public CosmosPagedFlux readManyByPartitionKey( * @param the type parameter * @param partitionKeys list of partition key values to read documents for * @param customQuery optional custom query for projections/additional filters (null means SELECT * FROM c) + * - should not contain WHERE clause filters for PK * @param requestOptions the optional request options * @param classType class type * @return a {@link CosmosPagedFlux} containing one or several feed response pages diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosContainer.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosContainer.java index 33149a498390..621173b22209 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosContainer.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosContainer.java @@ -594,6 +594,7 @@ public CosmosPagedIterable readManyByPartitionKey( * @param the type parameter * @param partitionKeys list of partition key values to read documents for * @param customQuery optional custom query for projections/additional filters (null means SELECT * FROM c) + * - should not contain WHERE clause filters for PK * @param classType class type * @return a {@link CosmosPagedIterable} containing the results */ @@ -621,6 +622,7 @@ public CosmosPagedIterable readManyByPartitionKey( * @param the type parameter * @param partitionKeys list of partition key values to read documents for * @param customQuery optional custom query for projections/additional filters (null means SELECT * FROM c) + * - should not contain WHERE clause filters for PK * @param requestOptions the optional request options * @param classType class type * @return a {@link CosmosPagedIterable} containing the results diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelper.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelper.java index c8d3094fedfd..a68866d76fff 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelper.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelper.java @@ -285,6 +285,14 @@ static int findTopLevelKeywordIndex(String queryText, String keyword) { } continue; } + // Skip double-quoted identifiers (e.g. c["WHERE"]) + if (ch == '"') { + i++; + while (i < len && queryText.charAt(i) != '"') { + i++; + } + continue; + } char upperCh = queryTextUpper.charAt(i); if (upperCh == '(') { depth++; From 89ac19e1cc27369fa073511d23eb291b90adb3cc Mon Sep 17 00:00:00 2001 From: Fabian Meiswinkel Date: Mon, 20 Apr 2026 18:05:50 +0000 Subject: [PATCH 32/64] Fixed code review comment --- ...ReadManyByPartitionKeyQueryHelperTest.java | 69 +++++++++++++++++++ .../ReadManyByPartitionKeyQueryHelper.java | 18 +++-- 2 files changed, 83 insertions(+), 4 deletions(-) diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelperTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelperTest.java index 68ab3b0900fb..8e2fa373f09e 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelperTest.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelperTest.java @@ -390,6 +390,75 @@ public void findWhere_ignoresFromInsideDoubleQuotedBracketNotation() { //endregion + //region Dotted/underscore/bracket property access boundary tests + + @Test(groups = { "unit" }) + public void findWhere_ignoresWhereInDottedPropertyAccess() { + // c.where is a property named "where" — not the WHERE keyword + int idx = ReadManyByPartitionKeyQueryHelper.findTopLevelWhereIndex( + "SELECT c.where, c.id FROM c WHERE c.status = 'active'"); + // "SELECT c.where, c.id FROM c " = 28 chars, so WHERE at 28 + assertThat(idx).isEqualTo(28); + } + + @Test(groups = { "unit" }) + public void findWhere_ignoresFromInDottedPropertyAccess() { + // c.from is a property — the real FROM is later + String query = "SELECT c.from FROM c WHERE c.x = 1"; + assertThat(ReadManyByPartitionKeyQueryHelper.extractTableAlias(query)).isEqualTo("c"); + } + + @Test(groups = { "unit" }) + public void findWhere_ignoresOrderInDottedPropertyAccess() { + // c.order shouldn't match ORDER keyword + int idx = ReadManyByPartitionKeyQueryHelper.findTopLevelKeywordIndex( + "SELECT c.order FROM c WHERE c.x = 1", "ORDER"); + assertThat(idx).isEqualTo(-1); // no ORDER BY in this query + } + + @Test(groups = { "unit" }) + public void findWhere_ignoresKeywordInUnderscoredIdentifier() { + // where_clause is an identifier, not WHERE + int idx = ReadManyByPartitionKeyQueryHelper.findTopLevelWhereIndex( + "SELECT c.where_clause FROM c WHERE c.x = 1"); + // "SELECT c.where_clause FROM c " = 29 chars, WHERE at 29 + assertThat(idx).isEqualTo(29); + } + + @Test(groups = { "unit" }) + public void findWhere_ignoresKeywordPrecededByUnderscore() { + // _where is an identifier + int idx = ReadManyByPartitionKeyQueryHelper.findTopLevelWhereIndex( + "SELECT c._where FROM c WHERE c.x = 1"); + assertThat(idx).isEqualTo(23); + } + + @Test(groups = { "unit" }) + public void extractAlias_ignoresKeywordPrefixInPropertyName() { + // c.offset is a property — should not confuse isFollowedByReservedKeyword + assertThat(ReadManyByPartitionKeyQueryHelper.extractTableAlias( + "SELECT c.offset_val FROM c WHERE c.x = 1")).isEqualTo("c"); + } + + @Test(groups = { "unit" }) + public void findWhere_ignoresKeywordFollowedByBracket() { + // WHERE[ should not match as a standalone WHERE keyword + int idx = ReadManyByPartitionKeyQueryHelper.findTopLevelWhereIndex( + "SELECT c.WHERE[0] FROM c WHERE c.x = 1"); + // The real WHERE is at position 25 + assertThat(idx).isEqualTo(25); + } + + @Test(groups = { "unit" }) + public void findWhere_ignoresKeywordFollowedByDollar() { + // WHERE$1 is an identifier, not WHERE + int idx = ReadManyByPartitionKeyQueryHelper.findTopLevelWhereIndex( + "SELECT WHERE$1 FROM c WHERE c.x = 1"); + assertThat(idx).isEqualTo(22); + } + + //endregion + //region OFFSET/LIMIT/HAVING alias detection tests (#9) @Test(groups = { "unit" }) diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelper.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelper.java index a68866d76fff..f2b7f22fee6f 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelper.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelper.java @@ -205,7 +205,7 @@ static String extractTableAlias(String queryText) { } // Handle optional AS: "FROM root AS r" -> alias is "r" if (remaining.startsWith("AS") - && (remaining.length() == 2 || !Character.isLetterOrDigit(remaining.charAt(2)))) { + && (remaining.length() == 2 || !isIdentifierChar(remaining.charAt(2)))) { afterFrom += 2; // skip AS while (afterFrom < queryText.length() && Character.isWhitespace(queryText.charAt(afterFrom))) { @@ -233,7 +233,7 @@ private static boolean isFollowedByReservedKeyword(String remainingUpper) { for (String kw : keywords) { if (remainingUpper.startsWith(kw) && (remainingUpper.length() == kw.length() - || !Character.isLetterOrDigit(remainingUpper.charAt(kw.length())))) { + || !isIdentifierChar(remainingUpper.charAt(kw.length())))) { return true; } } @@ -300,14 +300,24 @@ static int findTopLevelKeywordIndex(String queryText, String keyword) { depth--; } else if (depth == 0 && upperCh == keywordUpper.charAt(0) && queryTextUpper.startsWith(keywordUpper, i) - && (i == 0 || !Character.isLetterOrDigit(queryTextUpper.charAt(i - 1))) - && (i + keyLen >= queryTextUpper.length() || !Character.isLetterOrDigit(queryTextUpper.charAt(i + keyLen)))) { + && (i == 0 || !isIdentifierChar(queryTextUpper.charAt(i - 1))) + && (i + keyLen >= queryTextUpper.length() || !isIdentifierChar(queryTextUpper.charAt(i + keyLen)))) { return i; } } return -1; } + /** + * Returns true if the character can appear in a SQL identifier or property access, + * meaning it should NOT be treated as a word boundary for keyword matching. + * Covers letters, digits, underscore, dot (property access), bracket (bracket notation), + * and dollar sign (system properties). + */ + private static boolean isIdentifierChar(char ch) { + return Character.isLetterOrDigit(ch) || ch == '_' || ch == '.' || ch == '[' || ch == '$'; + } + /** * Finds the index of the top-level WHERE keyword in the query text, * ignoring WHERE that appears inside parentheses (subqueries). From ea28c5d461b15b9123473b616000eee16e2d45fd Mon Sep 17 00:00:00 2001 From: Fabian Meiswinkel Date: Mon, 20 Apr 2026 18:30:31 +0000 Subject: [PATCH 33/64] Avoid getting pkPaths (pk-def) per row --- .../cosmos/spark/CosmosItemsDataSource.scala | 49 ++---------- .../CosmosReadManyByPartitionKeyReader.scala | 76 ++++++++++--------- 2 files changed, 46 insertions(+), 79 deletions(-) diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosItemsDataSource.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosItemsDataSource.scala index 3433e352870b..4dd0f37b164c 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosItemsDataSource.scala +++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosItemsDataSource.scala @@ -135,6 +135,11 @@ object CosmosItemsDataSource { val sharedReadConfig = CosmosReadConfig.parseCosmosReadConfig(sharedEffectiveConfig) val sharedTreatNullAsNone = sharedReadConfig.readManyByPkTreatNullAsNone + // Initialize reader state once: resolves PK paths, infers schema, and broadcasts client caches + // in a single Loan block instead of three separate ones. + val readerState = readManyReader.initializeReaderState() + val (pkPaths, _, _) = readerState + // Option 1: Look for the _partitionKeyIdentity column (produced by GetCosmosPartitionKeyValue UDF) val pkIdentityFieldExtraction = df .schema @@ -151,46 +156,9 @@ object CosmosItemsDataSource { val pkColumnExtraction: Option[Row => PartitionKey] = if (pkIdentityFieldExtraction.isDefined) { None // no need to resolve PK paths - _partitionKeyIdentity column takes precedence } else { - val effectiveConfig = sharedEffectiveConfig - val readConfig = sharedReadConfig - val containerConfig = CosmosContainerConfig.parseCosmosContainerConfig(effectiveConfig) - val sparkEnvironmentInfo = CosmosClientConfiguration.getSparkEnvironmentInfo(None) - val calledFrom = s"CosmosItemsDataSource.readManyByPartitionKey" - val treatNullAsNone = readConfig.readManyByPkTreatNullAsNone - - val pkPaths = Loan( - List[Option[CosmosClientCacheItem]]( - Some( - CosmosClientCache( - CosmosClientConfiguration( - effectiveConfig, - readConsistencyStrategy = readConfig.readConsistencyStrategy, - sparkEnvironmentInfo), - None, - calledFrom)), - ThroughputControlHelper.getThroughputControlClientCacheItem( - effectiveConfig, - calledFrom, - None, - sparkEnvironmentInfo) - )) - .to(clientCacheItems => { - val container = - ThroughputControlHelper.getContainer( - effectiveConfig, - containerConfig, - clientCacheItems(0).get, - clientCacheItems(1)) - - val pkDefinition = SparkBridgeInternal - .getContainerPropertiesFromCollectionCache(container) - .getPartitionKeyDefinition - - pkDefinition.getPaths.asScala.map(_.stripPrefix("/")).toList - }) + val treatNullAsNone = sharedReadConfig.readManyByPkTreatNullAsNone // Nested PK paths (containing /) cannot be resolved from top-level DataFrame columns. - // Surface an explicit error so users know to use the UDF-produced _partitionKeyIdentity column. if (pkPaths.exists(_.contains("/"))) { throw new IllegalArgumentException( "Container has nested partition key path(s) " + pkPaths.mkString("[", ",", "]") + ". " + @@ -203,13 +171,10 @@ object CosmosItemsDataSource { val allPkColumnsPresent = pkPaths.forall(path => dfFieldNames.contains(path)) if (allPkColumnsPresent && pkPaths.nonEmpty) { - // pkPaths already defined above Some((row: Row) => { if (pkPaths.size == 1) { - // Single partition key buildPartitionKey(row.getAs[Any](pkPaths.head), treatNullAsNone) } else { - // Hierarchical partition key - build level by level val builder = new PartitionKeyBuilder() for (path <- pkPaths) { addPartitionKeyComponent(builder, row.getAs[Any](path), treatNullAsNone, pkPaths.size) @@ -230,7 +195,7 @@ object CosmosItemsDataSource { "Either add a '_partitionKeyIdentity' column (using the GetCosmosPartitionKeyValue UDF) " + "or ensure the DataFrame contains columns matching the container's partition key paths.")) - readManyReader.readManyByPartitionKey(df.rdd, pkExtraction) + readManyReader.readManyByPartitionKey(df.rdd, pkExtraction, readerState) } private def addPartitionKeyComponent( diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosReadManyByPartitionKeyReader.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosReadManyByPartitionKeyReader.scala index 796aeea52b91..a49bbe2867ad 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosReadManyByPartitionKeyReader.scala +++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosReadManyByPartitionKeyReader.scala @@ -2,7 +2,7 @@ // Licensed under the MIT License. package com.azure.cosmos.spark -import com.azure.cosmos.{CosmosException, ReadConsistencyStrategy} +import com.azure.cosmos.{CosmosException, ReadConsistencyStrategy, SparkBridgeInternal} import com.azure.cosmos.implementation.{CosmosClientMetadataCachesSnapshot, UUIDs} import com.azure.cosmos.models.PartitionKey import com.azure.cosmos.spark.CosmosPredicates.assertOnSparkDriver @@ -16,6 +16,10 @@ import org.apache.spark.sql.types.StructType import java.util.UUID +// scalastyle:off underscore.import +import scala.collection.JavaConverters._ +// scalastyle:on underscore.import + private[spark] class CosmosReadManyByPartitionKeyReader( val userProvidedSchema: StructType, val userConfig: Map[String, String] @@ -39,8 +43,15 @@ private[spark] class CosmosReadManyByPartitionKeyReader( val sparkEnvironmentInfo: String = CosmosClientConfiguration.getSparkEnvironmentInfo(Some(sparkSession)) logTrace(s"Instantiated ${this.getClass.getSimpleName} for $tableName") - private[spark] def initializeAndBroadcastCosmosClientStatesForContainer(): Broadcast[CosmosClientMetadataCachesSnapshots] = { - val calledFrom = s"CosmosReadManyByPartitionKeyReader($tableName).initializeAndBroadcastCosmosClientStateForContainer" + /** + * Resolves the partition key paths for the target container. + * Uses a single Loan block that also infers the schema (if needed) and warms + * the client metadata caches for broadcast — avoiding three separate Loan round-trips. + * + * @return (pkPaths, schema, broadcastClientStates) + */ + private[spark] def initializeReaderState(): (List[String], StructType, Broadcast[CosmosClientMetadataCachesSnapshots]) = { + val calledFrom = s"CosmosReadManyByPartitionKeyReader($tableName).initializeReaderState" Loan( List[Option[CosmosClientCacheItem]]( Some( @@ -65,10 +76,21 @@ private[spark] class CosmosReadManyByPartitionKeyReader( clientCacheItems(0).get, clientCacheItems(1)) - // Warm-up readItem: intentionally issues a lookup for a random id/partition-key pair - // on the driver so that the collection/routing-map caches are populated before we serialize - // the client state and broadcast it to executors. This costs ~1 RU + 1 RTT per broadcast build - // (expected 404) but avoids every executor doing the same lookup in parallel on first use. + // 1. Resolve PK paths from the collection cache + val pkPaths = SparkBridgeInternal + .getContainerPropertiesFromCollectionCache(container) + .getPartitionKeyDefinition + .getPaths.asScala.map(_.stripPrefix("/")).toList + + // 2. Infer schema if not user-provided + val schema = Option.apply(userProvidedSchema).getOrElse( + CosmosTableSchemaInferrer.inferSchema( + clientCacheItems(0).get, + clientCacheItems(1), + effectiveUserConfig, + ItemsTable.defaultSchemaForInferenceDisabled)) + + // 3. Warm-up readItem so collection/routing-map caches are populated before broadcast try { container.readItem( UUIDs.nonBlockingRandomUUID().toString, @@ -76,14 +98,12 @@ private[spark] class CosmosReadManyByPartitionKeyReader( classOf[ObjectNode]) .block() } catch { - // The warm-up readItem is only used to hydrate the collection/routing-map caches. - // A 404 (item not found) is expected, but we log other CosmosExceptions at debug to - // aid diagnosis (auth failures, throttling, etc.) while not failing reader setup. case ex: CosmosException => logDebug(s"Warm-up readItem for metadata caches completed with exception: ${ex.getMessage}", ex) None } + // 4. Serialize and broadcast client state val state = new CosmosClientMetadataCachesSnapshot() state.serialize(clientCacheItems(0).get.cosmosClient) @@ -94,37 +114,19 @@ private[spark] class CosmosReadManyByPartitionKeyReader( } val metadataSnapshots = CosmosClientMetadataCachesSnapshots(state, throughputControlState) - sparkSession.sparkContext.broadcast(metadataSnapshots) + val broadcastStates = sparkSession.sparkContext.broadcast(metadataSnapshots) + + (pkPaths, schema, broadcastStates) }) } - def readManyByPartitionKey(inputRdd: RDD[Row], pkExtraction: Row => PartitionKey): DataFrame = { - val correlationActivityId = UUIDs.nonBlockingRandomUUID() - val calledFrom = s"CosmosReadManyByPartitionKeyReader.readManyByPartitionKey($correlationActivityId)" - val schema = Loan( - List[Option[CosmosClientCacheItem]]( - Some(CosmosClientCache( - CosmosClientConfiguration( - effectiveUserConfig, - readConsistencyStrategy = readConfig.readConsistencyStrategy, - sparkEnvironmentInfo), - None, - calledFrom - )), - ThroughputControlHelper.getThroughputControlClientCacheItem( - effectiveUserConfig, - calledFrom, - None, - sparkEnvironmentInfo) - )) - .to(clientCacheItems => Option.apply(userProvidedSchema).getOrElse( - CosmosTableSchemaInferrer.inferSchema( - clientCacheItems(0).get, - clientCacheItems(1), - effectiveUserConfig, - ItemsTable.defaultSchemaForInferenceDisabled))) + def readManyByPartitionKey( + inputRdd: RDD[Row], + pkExtraction: Row => PartitionKey, + readerState: (List[String], StructType, Broadcast[CosmosClientMetadataCachesSnapshots])): DataFrame = { - val clientStates = initializeAndBroadcastCosmosClientStatesForContainer + val correlationActivityId = UUIDs.nonBlockingRandomUUID() + val (_, schema, clientStates) = readerState sparkSession.sqlContext.createDataFrame( inputRdd.mapPartitionsWithIndex( From 35e9d73883059a7f65dbf3e215aecaab40c0a3f0 Mon Sep 17 00:00:00 2001 From: Fabian Meiswinkel Date: Mon, 20 Apr 2026 18:37:02 +0000 Subject: [PATCH 34/64] Update CosmosReadManyByPartitionKeyReader.scala --- .../spark/CosmosReadManyByPartitionKeyReader.scala | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosReadManyByPartitionKeyReader.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosReadManyByPartitionKeyReader.scala index a49bbe2867ad..d7dc00617295 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosReadManyByPartitionKeyReader.scala +++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosReadManyByPartitionKeyReader.scala @@ -76,11 +76,13 @@ private[spark] class CosmosReadManyByPartitionKeyReader( clientCacheItems(0).get, clientCacheItems(1)) - // 1. Resolve PK paths from the collection cache - val pkPaths = SparkBridgeInternal - .getContainerPropertiesFromCollectionCache(container) - .getPartitionKeyDefinition - .getPaths.asScala.map(_.stripPrefix("/")).toList + // 1. Resolve PK paths from the collection cache (with transient-error retry) + val pkPaths = TransientErrorsRetryPolicy.executeWithRetry(() => { + SparkBridgeInternal + .getContainerPropertiesFromCollectionCache(container) + .getPartitionKeyDefinition + .getPaths.asScala.map(_.stripPrefix("/")).toList + }) // 2. Infer schema if not user-provided val schema = Option.apply(userProvidedSchema).getOrElse( From 78bbd1598d20629626be6a186fd63d4d62e39348 Mon Sep 17 00:00:00 2001 From: Fabian Meiswinkel Date: Mon, 20 Apr 2026 19:53:15 +0000 Subject: [PATCH 35/64] Renamed readManyByPartitionKey to readManyByPartitionKeys --- .../azure-cosmos-spark_3-3_2-12/CHANGELOG.md | 2 +- .../azure-cosmos-spark_3-4_2-12/CHANGELOG.md | 2 +- .../cosmos/spark/SparkE2EQueryITest.scala | 10 +++--- .../azure-cosmos-spark_3-5_2-12/CHANGELOG.md | 2 +- .../azure-cosmos-spark_3-5_2-13/CHANGELOG.md | 2 +- .../com/azure/cosmos/spark/CosmosConfig.scala | 2 +- .../cosmos/spark/CosmosItemsDataSource.scala | 10 +++--- .../CosmosReadManyByPartitionKeyReader.scala | 2 +- ...tionReaderWithReadManyByPartitionKey.scala | 2 +- ...tryingReadManyByPartitionKeyIterator.scala | 8 ++--- .../udf/GetCosmosPartitionKeyValue.scala | 2 +- .../azure-cosmos-spark_4-0_2-13/CHANGELOG.md | 2 +- .../com/azure/cosmos/CosmosTracerTest.java | 6 ++-- .../cosmos/ReadManyByPartitionKeyTest.java | 36 +++++++++---------- sdk/cosmos/azure-cosmos/CHANGELOG.md | 2 +- .../azure/cosmos/CosmosAsyncContainer.java | 20 +++++------ .../com/azure/cosmos/CosmosContainer.java | 16 ++++----- .../implementation/AsyncDocumentClient.java | 2 +- .../azure/cosmos/implementation/Configs.java | 2 +- .../ReadManyByPartitionKeyQueryHelper.java | 4 +-- .../implementation/RxDocumentClientImpl.java | 8 ++--- .../docs/readManyByPartitionKey-design.md | 16 ++++----- 22 files changed, 79 insertions(+), 79 deletions(-) diff --git a/sdk/cosmos/azure-cosmos-spark_3-3_2-12/CHANGELOG.md b/sdk/cosmos/azure-cosmos-spark_3-3_2-12/CHANGELOG.md index bf15bf284bbd..a7c445911584 100644 --- a/sdk/cosmos/azure-cosmos-spark_3-3_2-12/CHANGELOG.md +++ b/sdk/cosmos/azure-cosmos-spark_3-3_2-12/CHANGELOG.md @@ -3,7 +3,7 @@ ### 4.48.0-beta.1 (Unreleased) #### Features Added -* Added new `CosmosItemsDataSource.readManyByPartitionKey` Spark function to execute bulk queries by a list of pk-values with better efficiency. Configure null handling via `spark.cosmos.read.readManyByPk.nullHandling` - default `Null` treats a null PK column as JSON null (`addNullValue`), `None` treats it as `PartitionKey.NONE` (`addNoneValue` / `NOT IS_DEFINED`). These route to different physical partitions - picking the wrong mode silently returns zero rows. See [PR 48801](https://github.com/Azure/azure-sdk-for-java/pull/48801) +* Added new `CosmosItemsDataSource.readManyByPartitionKeys` Spark function to execute bulk queries by a list of pk-values with better efficiency. Configure null handling via `spark.cosmos.read.readManyByPk.nullHandling` - default `Null` treats a null PK column as JSON null (`addNullValue`), `None` treats it as `PartitionKey.NONE` (`addNoneValue` / `NOT IS_DEFINED`). These route to different physical partitions - picking the wrong mode silently returns zero rows. See [PR 48801](https://github.com/Azure/azure-sdk-for-java/pull/48801) #### Breaking Changes diff --git a/sdk/cosmos/azure-cosmos-spark_3-4_2-12/CHANGELOG.md b/sdk/cosmos/azure-cosmos-spark_3-4_2-12/CHANGELOG.md index 72e8fdcf99ed..9ab39ee9e943 100644 --- a/sdk/cosmos/azure-cosmos-spark_3-4_2-12/CHANGELOG.md +++ b/sdk/cosmos/azure-cosmos-spark_3-4_2-12/CHANGELOG.md @@ -3,7 +3,7 @@ ### 4.48.0-beta.1 (Unreleased) #### Features Added -* Added new `CosmosItemsDataSource.readManyByPartitionKey` Spark function to execute bulk queries by a list of pk-values with better efficiency. Configure null handling via `spark.cosmos.read.readManyByPk.nullHandling` - default `Null` treats a null PK column as JSON null (`addNullValue`), `None` treats it as `PartitionKey.NONE` (`addNoneValue` / `NOT IS_DEFINED`). These route to different physical partitions - picking the wrong mode silently returns zero rows. See [PR 48801](https://github.com/Azure/azure-sdk-for-java/pull/48801) +* Added new `CosmosItemsDataSource.readManyByPartitionKeys` Spark function to execute bulk queries by a list of pk-values with better efficiency. Configure null handling via `spark.cosmos.read.readManyByPk.nullHandling` - default `Null` treats a null PK column as JSON null (`addNullValue`), `None` treats it as `PartitionKey.NONE` (`addNoneValue` / `NOT IS_DEFINED`). These route to different physical partitions - picking the wrong mode silently returns zero rows. See [PR 48801](https://github.com/Azure/azure-sdk-for-java/pull/48801) #### Breaking Changes diff --git a/sdk/cosmos/azure-cosmos-spark_3-5/src/test/scala/com/azure/cosmos/spark/SparkE2EQueryITest.scala b/sdk/cosmos/azure-cosmos-spark_3-5/src/test/scala/com/azure/cosmos/spark/SparkE2EQueryITest.scala index bac76468dec5..79e52075c2cf 100644 --- a/sdk/cosmos/azure-cosmos-spark_3-5/src/test/scala/com/azure/cosmos/spark/SparkE2EQueryITest.scala +++ b/sdk/cosmos/azure-cosmos-spark_3-5/src/test/scala/com/azure/cosmos/spark/SparkE2EQueryITest.scala @@ -75,7 +75,7 @@ class SparkE2EQueryITest item.getAs[String]("id") shouldEqual id } - "spark readManyByPartitionKey" can "use a matching top-level partition key column without the UDF" in { + "spark readManyByPartitionKeys" can "use a matching top-level partition key column without the UDF" in { val cosmosEndpoint = TestConfigurations.HOST val cosmosMasterKey = TestConfigurations.MASTER_KEY val container = cosmosClient.getDatabase(cosmosDatabase).getContainer(cosmosContainersWithPkAsPartitionKey) @@ -102,7 +102,7 @@ class SparkE2EQueryITest import sparkSession.implicits._ val rows = CosmosItemsDataSource - .readManyByPartitionKey(Seq("pkA", "pkB").toDF("pk"), cfg.asJava) + .readManyByPartitionKeys(Seq("pkA", "pkB").toDF("pk"), cfg.asJava) .selectExpr("id", "pk", "payload") .collect() @@ -112,7 +112,7 @@ class SparkE2EQueryITest rows.map(_.getAs[String]("payload")).toSet shouldEqual Set("value-pkA", "value-pkB") } - "spark readManyByPartitionKey" can "require the UDF for nested partition key paths and succeed with it" in { + "spark readManyByPartitionKeys" can "require the UDF for nested partition key paths and succeed with it" in { val cosmosEndpoint = TestConfigurations.HOST val cosmosMasterKey = TestConfigurations.MASTER_KEY val containerName = s"nested-pk-${UUID.randomUUID()}" @@ -156,7 +156,7 @@ class SparkE2EQueryITest import sparkSession.implicits._ val missingUdfError = the[IllegalArgumentException] thrownBy { - CosmosItemsDataSource.readManyByPartitionKey(Seq("tenantA").toDF("tenantId"), cfg.asJava) + CosmosItemsDataSource.readManyByPartitionKeys(Seq("tenantA").toDF("tenantId"), cfg.asJava) } missingUdfError.getMessage should include("Nested paths cannot be resolved from DataFrame columns automatically") @@ -169,7 +169,7 @@ class SparkE2EQueryITest .withColumn("_partitionKeyIdentity", expr("GetCosmosPartitionKeyValue(tenantId)")) val rows = CosmosItemsDataSource - .readManyByPartitionKey(inputDf, cfg.asJava) + .readManyByPartitionKeys(inputDf, cfg.asJava) .selectExpr("id", "tenant.id as tenantId") .collect() diff --git a/sdk/cosmos/azure-cosmos-spark_3-5_2-12/CHANGELOG.md b/sdk/cosmos/azure-cosmos-spark_3-5_2-12/CHANGELOG.md index 9e995b3b228f..1e8bc11ffe69 100644 --- a/sdk/cosmos/azure-cosmos-spark_3-5_2-12/CHANGELOG.md +++ b/sdk/cosmos/azure-cosmos-spark_3-5_2-12/CHANGELOG.md @@ -3,7 +3,7 @@ ### 4.48.0-beta.1 (Unreleased) #### Features Added -* Added new `CosmosItemsDataSource.readManyByPartitionKey` Spark function to execute bulk queries by a list of pk-values with better efficiency. Configure null handling via `spark.cosmos.read.readManyByPk.nullHandling` - default `Null` treats a null PK column as JSON null (`addNullValue`), `None` treats it as `PartitionKey.NONE` (`addNoneValue` / `NOT IS_DEFINED`). These route to different physical partitions - picking the wrong mode silently returns zero rows. See [PR 48801](https://github.com/Azure/azure-sdk-for-java/pull/48801) +* Added new `CosmosItemsDataSource.readManyByPartitionKeys` Spark function to execute bulk queries by a list of pk-values with better efficiency. Configure null handling via `spark.cosmos.read.readManyByPk.nullHandling` - default `Null` treats a null PK column as JSON null (`addNullValue`), `None` treats it as `PartitionKey.NONE` (`addNoneValue` / `NOT IS_DEFINED`). These route to different physical partitions - picking the wrong mode silently returns zero rows. See [PR 48801](https://github.com/Azure/azure-sdk-for-java/pull/48801) #### Breaking Changes diff --git a/sdk/cosmos/azure-cosmos-spark_3-5_2-13/CHANGELOG.md b/sdk/cosmos/azure-cosmos-spark_3-5_2-13/CHANGELOG.md index 51e4e5760c00..b5dc43bab19d 100644 --- a/sdk/cosmos/azure-cosmos-spark_3-5_2-13/CHANGELOG.md +++ b/sdk/cosmos/azure-cosmos-spark_3-5_2-13/CHANGELOG.md @@ -3,7 +3,7 @@ ### 4.48.0-beta.1 (Unreleased) #### Features Added -* Added new `CosmosItemsDataSource.readManyByPartitionKey` Spark function to execute bulk queries by a list of pk-values with better efficiency. Configure null handling via `spark.cosmos.read.readManyByPk.nullHandling` - default `Null` treats a null PK column as JSON null (`addNullValue`), `None` treats it as `PartitionKey.NONE` (`addNoneValue` / `NOT IS_DEFINED`). These route to different physical partitions - picking the wrong mode silently returns zero rows. See [PR 48801](https://github.com/Azure/azure-sdk-for-java/pull/48801) +* Added new `CosmosItemsDataSource.readManyByPartitionKeys` Spark function to execute bulk queries by a list of pk-values with better efficiency. Configure null handling via `spark.cosmos.read.readManyByPk.nullHandling` - default `Null` treats a null PK column as JSON null (`addNullValue`), `None` treats it as `PartitionKey.NONE` (`addNoneValue` / `NOT IS_DEFINED`). These route to different physical partitions - picking the wrong mode silently returns zero rows. See [PR 48801](https://github.com/Azure/azure-sdk-for-java/pull/48801) #### Breaking Changes diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosConfig.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosConfig.scala index 8312e34bf0c1..1a1025cd7ad2 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosConfig.scala +++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosConfig.scala @@ -1145,7 +1145,7 @@ private object CosmosReadConfig { defaultValue = Some("Null"), parseFromStringFunction = value => value, helpMessage = "Determines how null values in partition key columns are treated for " + - "readManyByPartitionKey. 'Null' (default) maps null to a JSON null via addNullValue(), which " + + "readManyByPartitionKeys. 'Null' (default) maps null to a JSON null via addNullValue(), which " + "is appropriate when the document field exists with an explicit null value. 'None' maps null " + "to PartitionKey.NONE via addNoneValue(), which is only supported for single-path partition keys " + "and should only be used when the partition key path does not exist at all in the document. " + diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosItemsDataSource.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosItemsDataSource.scala index 4dd0f37b164c..ac5ed255c008 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosItemsDataSource.scala +++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosItemsDataSource.scala @@ -114,11 +114,11 @@ object CosmosItemsDataSource { readManyReader.readMany(df.rdd, readManyFilterExtraction) } - def readManyByPartitionKey(df: DataFrame, userConfig: java.util.Map[String, String]): DataFrame = { - readManyByPartitionKey(df, userConfig, null) + def readManyByPartitionKeys(df: DataFrame, userConfig: java.util.Map[String, String]): DataFrame = { + readManyByPartitionKeys(df, userConfig, null) } - def readManyByPartitionKey( + def readManyByPartitionKeys( df: DataFrame, userConfig: java.util.Map[String, String], userProvidedSchema: StructType): DataFrame = { @@ -195,7 +195,7 @@ object CosmosItemsDataSource { "Either add a '_partitionKeyIdentity' column (using the GetCosmosPartitionKeyValue UDF) " + "or ensure the DataFrame contains columns matching the container's partition key paths.")) - readManyReader.readManyByPartitionKey(df.rdd, pkExtraction, readerState) + readManyReader.readManyByPartitionKeys(df.rdd, pkExtraction, readerState) } private def addPartitionKeyComponent( @@ -220,7 +220,7 @@ object CosmosItemsDataSource { throw new IllegalArgumentException( s"Unsupported partition key column type '${other.getClass.getName}' with value '$other'. " + "Supported types are String, Number (integral or floating-point), Boolean, and null. " + - "For other source types, convert the column before calling readManyByPartitionKey or use " + + "For other source types, convert the column before calling readManyByPartitionKeys or use " + "the GetCosmosPartitionKeyValue UDF to produce a '_partitionKeyIdentity' column.") } } diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosReadManyByPartitionKeyReader.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosReadManyByPartitionKeyReader.scala index d7dc00617295..5cde39941c0d 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosReadManyByPartitionKeyReader.scala +++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosReadManyByPartitionKeyReader.scala @@ -122,7 +122,7 @@ private[spark] class CosmosReadManyByPartitionKeyReader( }) } - def readManyByPartitionKey( + def readManyByPartitionKeys( inputRdd: RDD[Row], pkExtraction: Row => PartitionKey, readerState: (List[String], StructType, Broadcast[CosmosClientMetadataCachesSnapshots])): DataFrame = { diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/ItemsPartitionReaderWithReadManyByPartitionKey.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/ItemsPartitionReaderWithReadManyByPartitionKey.scala index 822b52c88341..a92b4c833cac 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/ItemsPartitionReaderWithReadManyByPartitionKey.scala +++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/ItemsPartitionReaderWithReadManyByPartitionKey.scala @@ -156,7 +156,7 @@ private[spark] case class ItemsPartitionReaderWithReadManyByPartitionKey } ) - // Collect all PK values upfront - readManyByPartitionKey needs the full list to + // Collect all PK values upfront - readManyByPartitionKeys needs the full list to // group by physical partition (the SDK batches internally per physical partition). // Deduplicate using the canonical PartitionKeyInternal JSON representation so that // equivalent PKs built from different runtime types (Int vs Long vs Double) are diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/TransientIOErrorsRetryingReadManyByPartitionKeyIterator.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/TransientIOErrorsRetryingReadManyByPartitionKeyIterator.scala index f82855218cce..8c7e5999fa57 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/TransientIOErrorsRetryingReadManyByPartitionKeyIterator.scala +++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/TransientIOErrorsRetryingReadManyByPartitionKeyIterator.scala @@ -17,7 +17,7 @@ import scala.util.Random import scala.util.control.Breaks /** - * Retry-safe iterator for readManyByPartitionKey. The full partition-key list is passed to the + * Retry-safe iterator for readManyByPartitionKeys. The full partition-key list is passed to the * SDK in a single call - the SDK is responsible for fan-out and per-physical-partition batching * (see Configs.getReadManyByPkMaxBatchSize()). This iterator therefore wraps a single * CosmosPagedIterable and, on transient I/O failures, re-creates the underlying flux and @@ -90,9 +90,9 @@ private[spark] class TransientIOErrorsRetryingReadManyByPartitionKeyIterator[TSp case None => val pagedFlux = customQuery match { case Some(query) => - container.readManyByPartitionKey(partitionKeys, query, queryOptions, classType) + container.readManyByPartitionKeys(partitionKeys, query, queryOptions, classType) case None => - container.readManyByPartitionKey(partitionKeys, queryOptions, classType) + container.readManyByPartitionKeys(partitionKeys, queryOptions, classType) } val rawIterator = new CosmosPagedIterable[TSparkRow]( @@ -113,7 +113,7 @@ private[spark] class TransientIOErrorsRetryingReadManyByPartitionKeyIterator[TSp // The server returned fewer pages than before - cannot safely replay. // Surface a clean error rather than silently emitting a truncated result. throw new IllegalStateException( - s"readManyByPartitionKey retry replay failed: expected to skip $pagesCommitted " + + s"readManyByPartitionKeys retry replay failed: expected to skip $pagesCommitted " + s"already-emitted pages but only $skipped were available. Context: $operationContextString") } diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/udf/GetCosmosPartitionKeyValue.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/udf/GetCosmosPartitionKeyValue.scala index 4dcc812f3122..cb4b2c5d5fbc 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/udf/GetCosmosPartitionKeyValue.scala +++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/udf/GetCosmosPartitionKeyValue.scala @@ -13,7 +13,7 @@ class GetCosmosPartitionKeyValue extends UDF1[Object, String] { // CosmosPartitionKeyHelper.tryParsePartitionKey yields a PartitionKey built with // addNullValue(). If the caller instead wants PartitionKey.NONE semantics (absent PK // field) they should filter the null row before calling this UDF and use the schema-matched - // readManyByPartitionKey path with readManyByPk.nullHandling=None. That None mode is only + // readManyByPartitionKeys path with readManyByPk.nullHandling=None. That None mode is only // supported for single-path partition keys; hierarchical partition keys reject it. override def call(partitionKeyValue: Object): String = { partitionKeyValue match { diff --git a/sdk/cosmos/azure-cosmos-spark_4-0_2-13/CHANGELOG.md b/sdk/cosmos/azure-cosmos-spark_4-0_2-13/CHANGELOG.md index 9a58d9ec65d8..f686d3ee37da 100644 --- a/sdk/cosmos/azure-cosmos-spark_4-0_2-13/CHANGELOG.md +++ b/sdk/cosmos/azure-cosmos-spark_4-0_2-13/CHANGELOG.md @@ -3,7 +3,7 @@ ### 4.48.0-beta.1 (Unreleased) #### Features Added -* Added new `CosmosItemsDataSource.readManyByPartitionKey` Spark function to execute bulk queries by a list of pk-values with better efficiency. Configure null handling via `spark.cosmos.read.readManyByPk.nullHandling` - default `Null` treats a null PK column as JSON null (`addNullValue`), `None` treats it as `PartitionKey.NONE` (`addNoneValue` / `NOT IS_DEFINED`). These route to different physical partitions - picking the wrong mode silently returns zero rows. See [PR 48801](https://github.com/Azure/azure-sdk-for-java/pull/48801) +* Added new `CosmosItemsDataSource.readManyByPartitionKeys` Spark function to execute bulk queries by a list of pk-values with better efficiency. Configure null handling via `spark.cosmos.read.readManyByPk.nullHandling` - default `Null` treats a null PK column as JSON null (`addNullValue`), `None` treats it as `PartitionKey.NONE` (`addNoneValue` / `NOT IS_DEFINED`). These route to different physical partitions - picking the wrong mode silently returns zero rows. See [PR 48801](https://github.com/Azure/azure-sdk-for-java/pull/48801) #### Breaking Changes diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/CosmosTracerTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/CosmosTracerTest.java index 9a59e7af8d7c..f8e2979f544d 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/CosmosTracerTest.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/CosmosTracerTest.java @@ -925,14 +925,14 @@ public void cosmosAsyncContainer( .map(CosmosItemIdentity::getPartitionKey) .collect(Collectors.toList()); feedItemResponse = cosmosAsyncContainer - .readManyByPartitionKey(partitionKeys, ObjectNode.class) + .readManyByPartitionKeys(partitionKeys, ObjectNode.class) .byPage(1) .blockFirst(); assertThat(feedItemResponse).isNotNull(); assertThat(feedItemResponse.getResults()).isNotEmpty(); verifyTracerAttributes( mockTracer, - "readManyByPartitionKey." + cosmosAsyncContainer.getId(), + "readManyByPartitionKeys." + cosmosAsyncContainer.getId(), cosmosAsyncDatabase.getId(), cosmosAsyncContainer.getId(), feedItemResponse.getCosmosDiagnostics(), @@ -940,7 +940,7 @@ public void cosmosAsyncContainer( useLegacyTracing, enableRequestLevelTracing, forceThresholdViolations, - "readManyByPartitionKey", + "readManyByPartitionKeys", samplingRate); mockTracer.reset(); } diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/ReadManyByPartitionKeyTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/ReadManyByPartitionKeyTest.java index ca684f1abf9e..e031cd6abacc 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/ReadManyByPartitionKeyTest.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/ReadManyByPartitionKeyTest.java @@ -99,7 +99,7 @@ public void singlePk_readManyByPartitionKey_basic() { new PartitionKey("pk1"), new PartitionKey("pk2")); - CosmosPagedIterable results = singlePkContainer.readManyByPartitionKey(pkValues, ObjectNode.class); + CosmosPagedIterable results = singlePkContainer.readManyByPartitionKeys(pkValues, ObjectNode.class); List resultList = results.stream().collect(Collectors.toList()); assertThat(resultList).hasSize(5); // 3 + 2 @@ -119,7 +119,7 @@ public void singlePk_readManyByPartitionKey_withProjection() { List pkValues = Collections.singletonList(new PartitionKey("pkProj")); SqlQuerySpec customQuery = new SqlQuerySpec("SELECT c.id, c.mypk FROM c"); - CosmosPagedIterable results = singlePkContainer.readManyByPartitionKey( + CosmosPagedIterable results = singlePkContainer.readManyByPartitionKeys( pkValues, customQuery, null, ObjectNode.class); List resultList = results.stream().collect(Collectors.toList()); @@ -144,7 +144,7 @@ public void singlePk_readManyByPartitionKey_withAdditionalFilter() { "SELECT * FROM c WHERE c.status = @status", Arrays.asList(new SqlParameter("@status", "active"))); - CosmosPagedIterable results = singlePkContainer.readManyByPartitionKey( + CosmosPagedIterable results = singlePkContainer.readManyByPartitionKeys( pkValues, customQuery, null, ObjectNode.class); List resultList = results.stream().collect(Collectors.toList()); @@ -160,7 +160,7 @@ public void singlePk_readManyByPartitionKey_withAdditionalFilter() { public void singlePk_readManyByPartitionKey_emptyResults() { List pkValues = Collections.singletonList(new PartitionKey("nonExistent")); - CosmosPagedIterable results = singlePkContainer.readManyByPartitionKey(pkValues, ObjectNode.class); + CosmosPagedIterable results = singlePkContainer.readManyByPartitionKeys(pkValues, ObjectNode.class); List resultList = results.stream().collect(Collectors.toList()); assertThat(resultList).isEmpty(); @@ -179,7 +179,7 @@ public void hpk_readManyByPartitionKey_fullPk() { new PartitionKeyBuilder().add("Redmond").add("98053").add(1).build(), new PartitionKeyBuilder().add("Pittsburgh").add("15232").add(2).build()); - CosmosPagedIterable results = multiHashContainer.readManyByPartitionKey(pkValues, ObjectNode.class); + CosmosPagedIterable results = multiHashContainer.readManyByPartitionKeys(pkValues, ObjectNode.class); List resultList = results.stream().collect(Collectors.toList()); // Redmond/98053/1 has 2 items, Pittsburgh/15232/2 has 1 item @@ -196,7 +196,7 @@ public void hpk_readManyByPartitionKey_partialPk_singleLevel() { List pkValues = Collections.singletonList( new PartitionKeyBuilder().add("Redmond").build()); - CosmosPagedIterable results = multiHashContainer.readManyByPartitionKey(pkValues, ObjectNode.class); + CosmosPagedIterable results = multiHashContainer.readManyByPartitionKeys(pkValues, ObjectNode.class); List resultList = results.stream().collect(Collectors.toList()); // Redmond has 3 items total (2 with 98053/1 and 1 with 12345/1) @@ -216,7 +216,7 @@ public void hpk_readManyByPartitionKey_partialPk_twoLevels() { List pkValues = Collections.singletonList( new PartitionKeyBuilder().add("Redmond").add("98053").build()); - CosmosPagedIterable results = multiHashContainer.readManyByPartitionKey(pkValues, ObjectNode.class); + CosmosPagedIterable results = multiHashContainer.readManyByPartitionKeys(pkValues, ObjectNode.class); List resultList = results.stream().collect(Collectors.toList()); // Redmond/98053 has 2 items @@ -267,7 +267,7 @@ public void hpk_readManyByPartitionKey_withProjection() { SqlQuerySpec customQuery = new SqlQuerySpec("SELECT c.id, c.city FROM c"); - CosmosPagedIterable results = multiHashContainer.readManyByPartitionKey( + CosmosPagedIterable results = multiHashContainer.readManyByPartitionKeys( pkValues, customQuery, null, ObjectNode.class); List resultList = results.stream().collect(Collectors.toList()); @@ -286,7 +286,7 @@ public void rejectsAggregateQuery() { SqlQuerySpec aggregateQuery = new SqlQuerySpec("SELECT COUNT(1) FROM c"); try { - singlePkContainer.readManyByPartitionKey(pkValues, aggregateQuery, null, ObjectNode.class) + singlePkContainer.readManyByPartitionKeys(pkValues, aggregateQuery, null, ObjectNode.class) .stream().collect(Collectors.toList()); fail("Should have thrown IllegalArgumentException for aggregate query"); } catch (IllegalArgumentException e) { @@ -300,7 +300,7 @@ public void rejectsOrderByQuery() { SqlQuerySpec orderByQuery = new SqlQuerySpec("SELECT * FROM c ORDER BY c.id"); try { - singlePkContainer.readManyByPartitionKey(pkValues, orderByQuery, null, ObjectNode.class) + singlePkContainer.readManyByPartitionKeys(pkValues, orderByQuery, null, ObjectNode.class) .stream().collect(Collectors.toList()); fail("Should have thrown IllegalArgumentException for ORDER BY query"); } catch (IllegalArgumentException e) { @@ -314,7 +314,7 @@ public void rejectsDistinctQuery() { SqlQuerySpec distinctQuery = new SqlQuerySpec("SELECT DISTINCT c.mypk FROM c"); try { - singlePkContainer.readManyByPartitionKey(pkValues, distinctQuery, null, ObjectNode.class) + singlePkContainer.readManyByPartitionKeys(pkValues, distinctQuery, null, ObjectNode.class) .stream().collect(Collectors.toList()); fail("Should have thrown IllegalArgumentException for DISTINCT query"); } catch (IllegalArgumentException e) { @@ -328,7 +328,7 @@ public void rejectsGroupByQuery() { SqlQuerySpec groupByQuery = new SqlQuerySpec("SELECT c.mypk FROM c GROUP BY c.mypk"); try { - singlePkContainer.readManyByPartitionKey(pkValues, groupByQuery, null, ObjectNode.class) + singlePkContainer.readManyByPartitionKeys(pkValues, groupByQuery, null, ObjectNode.class) .stream().collect(Collectors.toList()); fail("Should have thrown IllegalArgumentException for GROUP BY query"); } catch (IllegalArgumentException e) { @@ -342,7 +342,7 @@ public void rejectsGroupByWithAggregateQuery() { SqlQuerySpec groupByWithAggregateQuery = new SqlQuerySpec("SELECT c.mypk, COUNT(1) as cnt FROM c GROUP BY c.mypk"); try { - singlePkContainer.readManyByPartitionKey(pkValues, groupByWithAggregateQuery, null, ObjectNode.class) + singlePkContainer.readManyByPartitionKeys(pkValues, groupByWithAggregateQuery, null, ObjectNode.class) .stream().collect(Collectors.toList()); fail("Should have thrown IllegalArgumentException for GROUP BY with aggregate query"); } catch (IllegalArgumentException e) { @@ -352,12 +352,12 @@ public void rejectsGroupByWithAggregateQuery() { @Test(groups = {"emulator"}, timeOut = TIMEOUT, expectedExceptions = IllegalArgumentException.class) public void rejectsNullPartitionKeyList() { - singlePkContainer.readManyByPartitionKey((List) null, ObjectNode.class); + singlePkContainer.readManyByPartitionKeys((List) null, ObjectNode.class); } @Test(groups = {"emulator"}, timeOut = TIMEOUT, expectedExceptions = IllegalArgumentException.class) public void rejectsEmptyPartitionKeyList() { - singlePkContainer.readManyByPartitionKey(new ArrayList<>(), ObjectNode.class) + singlePkContainer.readManyByPartitionKeys(new ArrayList<>(), ObjectNode.class) .stream().collect(Collectors.toList()); } @@ -367,7 +367,7 @@ public void rejectsOffsetQuery() { SqlQuerySpec offsetQuery = new SqlQuerySpec("SELECT * FROM c OFFSET 0 LIMIT 10"); try { - singlePkContainer.readManyByPartitionKey(pkValues, offsetQuery, null, ObjectNode.class) + singlePkContainer.readManyByPartitionKeys(pkValues, offsetQuery, null, ObjectNode.class) .stream().collect(Collectors.toList()); fail("Should have thrown IllegalArgumentException for OFFSET query"); } catch (IllegalArgumentException e) { @@ -402,7 +402,7 @@ public void singlePk_readManyByPartitionKey_withSmallBatchSize() { new PartitionKey("batchPk3"), new PartitionKey("batchPk4")); - CosmosPagedIterable results = singlePkContainer.readManyByPartitionKey(pkValues, ObjectNode.class); + CosmosPagedIterable results = singlePkContainer.readManyByPartitionKeys(pkValues, ObjectNode.class); List> pages = new ArrayList<>(); results.iterableByPage().forEach(pages::add); List resultList = pages.stream() @@ -450,7 +450,7 @@ public T deserialize(Map jsonNodeMap, Class classType) { } }); - CosmosPagedIterable results = singlePkContainer.readManyByPartitionKey( + CosmosPagedIterable results = singlePkContainer.readManyByPartitionKeys( pkValues, options, ReadManyByPartitionKeyPojo.class); List resultList = results.stream().collect(Collectors.toList()); diff --git a/sdk/cosmos/azure-cosmos/CHANGELOG.md b/sdk/cosmos/azure-cosmos/CHANGELOG.md index 96f3ddeb28c0..042527892a09 100644 --- a/sdk/cosmos/azure-cosmos/CHANGELOG.md +++ b/sdk/cosmos/azure-cosmos/CHANGELOG.md @@ -4,7 +4,7 @@ #### Features Added * Added support for change feed with `startFrom` point-in-time on merged partitions by enabling the `CHANGE_FEED_WITH_START_TIME_POST_MERGE` SDK capability. - See [PR 48752](https://github.com/Azure/azure-sdk-for-java/pull/48752) -* Added new `readManyByPartitionKey` API on `CosmosAsyncContainer` / `CosmosContainer` to bulk-query all documents matching a list of partition key values with better efficiency than issuing individual queries. Also fixes nested partition key selector generation for `readMany` and `readAllItems`. See [PR 48801](https://github.com/Azure/azure-sdk-for-java/pull/48801) +* Added new `readManyByPartitionKeys` API on `CosmosAsyncContainer` / `CosmosContainer` to bulk-query all documents matching a list of partition key values with better efficiency than issuing individual queries. Also fixes nested partition key selector generation for `readMany` and `readAllItems`. See [PR 48801](https://github.com/Azure/azure-sdk-for-java/pull/48801) #### Breaking Changes diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java index 6e70ecb3d940..b24ab9991333 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java @@ -199,7 +199,7 @@ protected CosmosAsyncContainer(CosmosAsyncContainer toBeWrappedContainer) { this.createItemSpanName = "createItem." + this.id; this.readAllItemsSpanName = "readAllItems." + this.id; this.readManyItemsSpanName = "readManyItems." + this.id; - this.readManyByPartitionKeySpanName = "readManyByPartitionKey." + this.id; + this.readManyByPartitionKeySpanName = "readManyByPartitionKeys." + this.id; this.readAllItemsOfLogicalPartitionSpanName = "readAllItemsOfLogicalPartition." + this.id; this.queryItemsSpanName = "queryItems." + this.id; this.queryChangeFeedSpanName = "queryChangeFeed." + this.id; @@ -1614,11 +1614,11 @@ private Mono> readManyInternal( * @param classType class type * @return a {@link CosmosPagedFlux} containing one or several feed response pages */ - public CosmosPagedFlux readManyByPartitionKey( + public CosmosPagedFlux readManyByPartitionKeys( List partitionKeys, Class classType) { - return this.readManyByPartitionKey(partitionKeys, null, null, classType); + return this.readManyByPartitionKeys(partitionKeys, null, null, classType); } /** @@ -1633,12 +1633,12 @@ public CosmosPagedFlux readManyByPartitionKey( * @param classType class type * @return a {@link CosmosPagedFlux} containing one or several feed response pages */ - public CosmosPagedFlux readManyByPartitionKey( + public CosmosPagedFlux readManyByPartitionKeys( List partitionKeys, CosmosReadManyRequestOptions requestOptions, Class classType) { - return this.readManyByPartitionKey(partitionKeys, null, requestOptions, classType); + return this.readManyByPartitionKeys(partitionKeys, null, requestOptions, classType); } /** @@ -1661,12 +1661,12 @@ public CosmosPagedFlux readManyByPartitionKey( * @param classType class type * @return a {@link CosmosPagedFlux} containing one or several feed response pages */ - public CosmosPagedFlux readManyByPartitionKey( + public CosmosPagedFlux readManyByPartitionKeys( List partitionKeys, SqlQuerySpec customQuery, Class classType) { - return this.readManyByPartitionKey(partitionKeys, customQuery, null, classType); + return this.readManyByPartitionKeys(partitionKeys, customQuery, null, classType); } /** @@ -1690,7 +1690,7 @@ public CosmosPagedFlux readManyByPartitionKey( * @param classType class type * @return a {@link CosmosPagedFlux} containing one or several feed response pages */ - public CosmosPagedFlux readManyByPartitionKey( + public CosmosPagedFlux readManyByPartitionKeys( List partitionKeys, SqlQuerySpec customQuery, CosmosReadManyRequestOptions requestOptions, @@ -1731,7 +1731,7 @@ private Function>> readManyByPa if (queryRequestOptions.getMaxDegreeOfParallelism() == 0) { queryRequestOptions.setMaxDegreeOfParallelism(-1); } - queryRequestOptions.setQueryName("readManyByPartitionKey"); + queryRequestOptions.setQueryName("readManyByPartitionKeys"); CosmosQueryRequestOptionsBase cosmosQueryRequestOptionsImpl = queryOptionsAccessor().getImpl(queryRequestOptions); applyPolicies(OperationType.Query, ResourceType.Document, cosmosQueryRequestOptionsImpl, this.readManyByPartitionKeySpanName); @@ -1751,7 +1751,7 @@ private Function>> readManyByPa return CosmosBridgeInternal .getAsyncDocumentClient(this.getDatabase()) - .readManyByPartitionKey( + .readManyByPartitionKeys( partitionKeys, customQuery, BridgeInternal.getLink(this), diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosContainer.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosContainer.java index 621173b22209..6acf391f3244 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosContainer.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosContainer.java @@ -551,11 +551,11 @@ public FeedResponse readMany( * @param classType class type * @return a {@link CosmosPagedIterable} containing the results */ - public CosmosPagedIterable readManyByPartitionKey( + public CosmosPagedIterable readManyByPartitionKeys( List partitionKeys, Class classType) { - return getCosmosPagedIterable(this.asyncContainer.readManyByPartitionKey(partitionKeys, classType)); + return getCosmosPagedIterable(this.asyncContainer.readManyByPartitionKeys(partitionKeys, classType)); } /** @@ -570,12 +570,12 @@ public CosmosPagedIterable readManyByPartitionKey( * @param classType class type * @return a {@link CosmosPagedIterable} containing the results */ - public CosmosPagedIterable readManyByPartitionKey( + public CosmosPagedIterable readManyByPartitionKeys( List partitionKeys, CosmosReadManyRequestOptions requestOptions, Class classType) { - return getCosmosPagedIterable(this.asyncContainer.readManyByPartitionKey(partitionKeys, requestOptions, classType)); + return getCosmosPagedIterable(this.asyncContainer.readManyByPartitionKeys(partitionKeys, requestOptions, classType)); } /** @@ -598,12 +598,12 @@ public CosmosPagedIterable readManyByPartitionKey( * @param classType class type * @return a {@link CosmosPagedIterable} containing the results */ - public CosmosPagedIterable readManyByPartitionKey( + public CosmosPagedIterable readManyByPartitionKeys( List partitionKeys, SqlQuerySpec customQuery, Class classType) { - return getCosmosPagedIterable(this.asyncContainer.readManyByPartitionKey(partitionKeys, customQuery, classType)); + return getCosmosPagedIterable(this.asyncContainer.readManyByPartitionKeys(partitionKeys, customQuery, classType)); } /** @@ -627,13 +627,13 @@ public CosmosPagedIterable readManyByPartitionKey( * @param classType class type * @return a {@link CosmosPagedIterable} containing the results */ - public CosmosPagedIterable readManyByPartitionKey( + public CosmosPagedIterable readManyByPartitionKeys( List partitionKeys, SqlQuerySpec customQuery, CosmosReadManyRequestOptions requestOptions, Class classType) { - return getCosmosPagedIterable(this.asyncContainer.readManyByPartitionKey(partitionKeys, customQuery, requestOptions, classType)); + return getCosmosPagedIterable(this.asyncContainer.readManyByPartitionKeys(partitionKeys, customQuery, requestOptions, classType)); } /** diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/AsyncDocumentClient.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/AsyncDocumentClient.java index 8e2499c9039f..d2d6413489ad 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/AsyncDocumentClient.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/AsyncDocumentClient.java @@ -1598,7 +1598,7 @@ Mono> readMany( * @param the type parameter * @return a Flux with feed response pages of documents */ - Flux> readManyByPartitionKey( + Flux> readManyByPartitionKeys( List partitionKeys, SqlQuerySpec customQuery, String collectionLink, diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/Configs.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/Configs.java index 1f5952302525..b896430e85cb 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/Configs.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/Configs.java @@ -250,7 +250,7 @@ public class Configs { public static final String MIN_TARGET_BULK_MICRO_BATCH_SIZE_VARIABLE = "COSMOS_MIN_TARGET_BULK_MICRO_BATCH_SIZE"; public static final int DEFAULT_MIN_TARGET_BULK_MICRO_BATCH_SIZE = 1; - // readManyByPartitionKey: max number of PK values per query per physical partition + // readManyByPartitionKeys: max number of PK values per query per physical partition public static final String READ_MANY_BY_PK_MAX_BATCH_SIZE = "COSMOS.READ_MANY_BY_PK_MAX_BATCH_SIZE"; public static final String READ_MANY_BY_PK_MAX_BATCH_SIZE_VARIABLE = "COSMOS_READ_MANY_BY_PK_MAX_BATCH_SIZE"; public static final int DEFAULT_READ_MANY_BY_PK_MAX_BATCH_SIZE = 100; diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelper.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelper.java index f2b7f22fee6f..39a33a1415a1 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelper.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelper.java @@ -15,7 +15,7 @@ import java.util.stream.Collectors; /** - * Helper for constructing SqlQuerySpec instances for readManyByPartitionKey operations. + * Helper for constructing SqlQuerySpec instances for readManyByPartitionKeys operations. * This class is not intended to be used directly by end-users. */ public class ReadManyByPartitionKeyQueryHelper { @@ -39,7 +39,7 @@ public static SqlQuerySpec createReadManyByPkQuerySpec( if (name != null && name.startsWith(PK_PARAM_PREFIX)) { throw new IllegalArgumentException( "Custom query parameter name '" + name + "' collides with the reserved " + - "readManyByPartitionKey internal prefix '" + PK_PARAM_PREFIX + "'. Rename the parameter."); + "readManyByPartitionKeys internal prefix '" + PK_PARAM_PREFIX + "'. Rename the parameter."); } } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java index c1ca60c3f889..1e5d7f9e8f50 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java @@ -4367,7 +4367,7 @@ private Mono>> readMany( } @Override - public Flux> readManyByPartitionKey( + public Flux> readManyByPartitionKeys( List partitionKeys, SqlQuerySpec customQuery, String collectionLink, @@ -4379,7 +4379,7 @@ public Flux> readManyByPartitionKey( final ScopedDiagnosticsFactory diagnosticsFactory = new ScopedDiagnosticsFactory(this, true); state.registerDiagnosticsFactory( - () -> {}, // we never want to reset in readManyByPartitionKey + () -> {}, // we never want to reset in readManyByPartitionKeys (ctx) -> diagnosticsFactory.merge(ctx) ); @@ -4396,7 +4396,7 @@ public Flux> readManyByPartitionKey( return ObservableHelper .fluxInlineIfPossibleAsObs( - () -> readManyByPartitionKey( + () -> readManyByPartitionKeys( partitionKeys, customQuery, collectionLink, state, diagnosticsFactory, klass), staleResourceRetryPolicy ) @@ -4431,7 +4431,7 @@ public Flux> readManyByPartitionKey( }); } - private Flux> readManyByPartitionKey( + private Flux> readManyByPartitionKeys( List partitionKeys, SqlQuerySpec customQuery, String collectionLink, diff --git a/sdk/cosmos/docs/readManyByPartitionKey-design.md b/sdk/cosmos/docs/readManyByPartitionKey-design.md index c193630d4064..7f23d63011dd 100644 --- a/sdk/cosmos/docs/readManyByPartitionKey-design.md +++ b/sdk/cosmos/docs/readManyByPartitionKey-design.md @@ -1,8 +1,8 @@ -# readManyByPartitionKey — Design & Implementation +# readManyByPartitionKeys — Design & Implementation ## Overview -New `readManyByPartitionKey` methods on `CosmosAsyncContainer` / `CosmosContainer` that accept a +New `readManyByPartitionKeys` methods on `CosmosAsyncContainer` / `CosmosContainer` that accept a `List` (without item-id). The SDK splits the PK values by physical partition, generates batched streaming queries per physical partition, and returns results as `CosmosPagedFlux` / `CosmosPagedIterable`. @@ -14,7 +14,7 @@ and additional filters. The SDK appends the auto-generated PK WHERE clause to it | Topic | Decision | |---|---| -| API name | `readManyByPartitionKey` — distinct name to avoid ambiguity with existing `readMany(List)` | +| API name | `readManyByPartitionKeys` — distinct name to avoid ambiguity with existing `readMany(List)` | | Return type | `CosmosPagedFlux` (async) / `CosmosPagedIterable` (sync) | | Custom query format | `SqlQuerySpec` — full query with parameters; SDK ANDs the PK filter | | Partial HPK | Supported from the start; prefix PKs fan out via `getOverlappingRanges` | @@ -31,11 +31,11 @@ and additional filters. The SDK appends the auto-generated PK WHERE clause to it ### Step 1: New public overloads in CosmosAsyncContainer ```java - CosmosPagedFlux readManyByPartitionKey(List partitionKeys, Class classType) - CosmosPagedFlux readManyByPartitionKey(List partitionKeys, + CosmosPagedFlux readManyByPartitionKeys(List partitionKeys, Class classType) + CosmosPagedFlux readManyByPartitionKeys(List partitionKeys, CosmosReadManyRequestOptions requestOptions, Class classType) - CosmosPagedFlux readManyByPartitionKey(List partitionKeys, + CosmosPagedFlux readManyByPartitionKeys(List partitionKeys, SqlQuerySpec customQuery, CosmosReadManyRequestOptions requestOptions, Class classType) @@ -114,7 +114,7 @@ If the base query already has a WHERE clause: ### Step 6: Interface wiring -New method `readManyByPartitionKey` added directly to `AsyncDocumentClient` interface, implemented in `RxDocumentClientImpl`. New `fetchQueryPlanForValidation` static method added to `DocumentQueryExecutionContextFactory` for custom query validation. +New method `readManyByPartitionKeys` added directly to `AsyncDocumentClient` interface, implemented in `RxDocumentClientImpl`. New `fetchQueryPlanForValidation` static method added to `DocumentQueryExecutionContextFactory` for custom query validation. ### Step 7: Configuration @@ -135,7 +135,7 @@ New configurable batch size via system property `COSMOS.READ_MANY_BY_PK_MAX_BATC - `tryParsePartitionKey(serialized: String): Option[PartitionKey]` — deserialize; returns `None` for malformed input including invalid JSON (wrapped in `scala.util.Try`). - When `spark.cosmos.read.readManyByPk.nullHandling=None` is used, hierarchical partition keys with null components are rejected with a clear error because `PartitionKey.NONE` cannot be used with multiple paths. -### Step 10: `CosmosItemsDataSource.readManyByPartitionKey` +### Step 10: `CosmosItemsDataSource.readManyByPartitionKeys` Static entry points that accept a DataFrame and Cosmos config. PK extraction supports two modes: 1. **UDF-produced column**: DataFrame contains `_partitionKeyIdentity` column (from `GetCosmosPartitionKeyValue` UDF). From dd831253a57b3b5f6900fe21bac98c43acede1c8 Mon Sep 17 00:00:00 2001 From: Fabian Meiswinkel Date: Mon, 20 Apr 2026 21:30:20 +0000 Subject: [PATCH 36/64] Fixing code review comments --- ...tionReaderWithReadManyByPartitionKey.scala | 26 +++- .../TransientIOErrorsRetryingIterator.scala | 2 +- ...tryingReadManyByPartitionKeyIterator.scala | 146 +++++------------- ...ientIOErrorsRetryingReadManyIterator.scala | 22 +-- 4 files changed, 64 insertions(+), 132 deletions(-) diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/ItemsPartitionReaderWithReadManyByPartitionKey.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/ItemsPartitionReaderWithReadManyByPartitionKey.scala index a92b4c833cac..bd25d0b601eb 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/ItemsPartitionReaderWithReadManyByPartitionKey.scala +++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/ItemsPartitionReaderWithReadManyByPartitionKey.scala @@ -11,6 +11,7 @@ import com.azure.cosmos.models.{CosmosReadManyRequestOptions, ModelBridgeInterna import com.azure.cosmos.spark.BulkWriter.getThreadInfo import com.azure.cosmos.spark.CosmosTableSchemaInferrer.IdAttributeName import com.azure.cosmos.spark.diagnostics.{DetailedFeedDiagnosticsProvider, DiagnosticsContext, DiagnosticsLoader, LoggerHelper, SparkTaskContext} +import com.azure.cosmos.util.CosmosPagedFlux import com.fasterxml.jackson.databind.node.ObjectNode import org.apache.spark.TaskContext import org.apache.spark.broadcast.Broadcast @@ -195,9 +196,9 @@ private[spark] case class ItemsPartitionReaderWithReadManyByPartitionKey } // Pass the full PK list to the SDK (which batches per physical partition internally). - // On transient I/O failures the retry iterator tracks pages already emitted upstream - // and skips them on replay; if a failure occurs mid-page (after items from that page have been - // emitted) the task fails rather than risking row duplication. + // On transient I/O failures the retry iterator re-creates the underlying flux from the + // continuation token of the last fully-committed page, matching the pattern used by + // TransientIOErrorsRetryingIterator for queries and change feed. private val isClosed = new AtomicBoolean(false) private var iteratorOpt: Option[CloseableSparkRowItemIterator] = None @@ -209,11 +210,22 @@ private[spark] case class ItemsPartitionReaderWithReadManyByPartitionKey EmptySparkRowItemIterator } else { new CloseableSparkRowItemIterator { + private val customQueryOpt = readConfig.customQuery.map(_.toSqlQuerySpec) + + // Factory that creates a CosmosPagedFlux from an optional continuation token. + // On the first call continuationToken is null (start from scratch); on retry + // it is the continuation token from the last fully-drained page. + private val fluxFactory: String => CosmosPagedFlux[SparkRowItem] = { (_: String) => + customQueryOpt match { + case Some(query) => + cosmosAsyncContainer.readManyByPartitionKeys(pkList, query, readManyOptions, classOf[SparkRowItem]) + case None => + cosmosAsyncContainer.readManyByPartitionKeys(pkList, readManyOptions, classOf[SparkRowItem]) + } + } + private val delegate = new TransientIOErrorsRetryingReadManyByPartitionKeyIterator[SparkRowItem]( - cosmosAsyncContainer, - pkList, - readConfig.customQuery.map(_.toSqlQuerySpec), - readManyOptions, + fluxFactory, readConfig.maxItemCount, readConfig.prefetchBufferSize, operationContextAndListenerTuple, diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/TransientIOErrorsRetryingIterator.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/TransientIOErrorsRetryingIterator.scala index b227d1a5ffb1..946d085878b1 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/TransientIOErrorsRetryingIterator.scala +++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/TransientIOErrorsRetryingIterator.scala @@ -275,7 +275,7 @@ private class TransientIOErrorsRetryingIterator[TSparkRow] } } -private object TransientIOErrorsRetryingIterator extends BasicLoggingTrait { +private[spark] object TransientIOErrorsRetryingIterator extends BasicLoggingTrait { private val maxConcurrency = SparkUtils.getNumberOfHostCPUCores val executorService: ExecutorService = new ThreadPoolExecutor( diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/TransientIOErrorsRetryingReadManyByPartitionKeyIterator.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/TransientIOErrorsRetryingReadManyByPartitionKeyIterator.scala index 8c7e5999fa57..eb2ff8142227 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/TransientIOErrorsRetryingReadManyByPartitionKeyIterator.scala +++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/TransientIOErrorsRetryingReadManyByPartitionKeyIterator.scala @@ -3,32 +3,37 @@ package com.azure.cosmos.spark -import com.azure.cosmos.{CosmosAsyncContainer, CosmosException} +import com.azure.cosmos.CosmosException import com.azure.cosmos.implementation.OperationCancelledException import com.azure.cosmos.implementation.spark.OperationContextAndListenerTuple -import com.azure.cosmos.models.{CosmosReadManyRequestOptions, FeedResponse, PartitionKey, SqlQuerySpec} +import com.azure.cosmos.models.FeedResponse import com.azure.cosmos.spark.diagnostics.BasicLoggingTrait -import com.azure.cosmos.util.CosmosPagedIterable +import com.azure.cosmos.util.{CosmosPagedFlux, CosmosPagedIterable} -import java.util.concurrent.{ExecutorService, SynchronousQueue, ThreadPoolExecutor, TimeUnit, TimeoutException} -import java.util.concurrent.atomic.AtomicLong +import java.util.concurrent.TimeoutException +import java.util.concurrent.atomic.{AtomicLong, AtomicReference} import scala.concurrent.{Await, ExecutionContext, Future} import scala.util.Random import scala.util.control.Breaks +// scalastyle:off underscore.import +import scala.collection.JavaConverters._ +// scalastyle:on underscore.import + /** - * Retry-safe iterator for readManyByPartitionKeys. The full partition-key list is passed to the - * SDK in a single call - the SDK is responsible for fan-out and per-physical-partition batching - * (see Configs.getReadManyByPkMaxBatchSize()). This iterator therefore wraps a single - * CosmosPagedIterable and, on transient I/O failures, re-creates the underlying flux and - * skips the pages that were already emitted upstream so no row is delivered twice. + * Retry-safe iterator for readManyByPartitionKeys that uses continuation-token-based + * replay (matching the pattern of TransientIOErrorsRetryingIterator). + * + * On transient I/O failures the iterator re-creates the underlying CosmosPagedFlux from the + * continuation token of the last fully-committed page. A page is "committed" only after all + * its items have been drained by the caller; the continuation token is captured from each + * FeedResponse and used as the resume point on retry. This avoids the correctness issues of + * page-count-based skipping (where the server is not guaranteed to return the same page + * boundaries across requests). */ private[spark] class TransientIOErrorsRetryingReadManyByPartitionKeyIterator[TSparkRow] ( - val container: CosmosAsyncContainer, - val partitionKeys: java.util.List[PartitionKey], - val customQuery: Option[SqlQuerySpec], - val queryOptions: CosmosReadManyRequestOptions, + val cosmosPagedFluxFactory: String => CosmosPagedFlux[TSparkRow], val pageSize: Int, val pagePrefetchBufferSize: Int, val operationContextAndListener: Option[OperationContextAndListenerTuple], @@ -43,6 +48,9 @@ private[spark] class TransientIOErrorsRetryingReadManyByPartitionKeyIterator[TSp scala.concurrent.duration.SECONDS) private val rnd = Random + // scalastyle:off null + private val lastContinuationToken = new AtomicReference[String](null) + // scalastyle:on null private val retryCount = new AtomicLong(0) private lazy val operationContextString = operationContextAndListener match { case Some(o) => if (o.getOperationContext != null) { @@ -53,17 +61,6 @@ private[spark] class TransientIOErrorsRetryingReadManyByPartitionKeyIterator[TSp case None => "n/a" } - // Number of pages that have been fully emitted upstream. On retry, we recreate the flux - // and skip this many pages before emitting any item, so already-delivered rows are not - // re-emitted. A page is "committed" as soon as we surface its first item to the caller - - // subsequent failures while still inside that page cannot be recovered from without - // risking duplication, so we fail fast in that case. - private var pagesCommitted: Long = 0 - // Whether the currently-buffered page has emitted at least one item. If true, we have - // passed the point of no return for this page: any transient failure here must surface, - // because we cannot partially-skip within a page on retry. - private var currentPagePartiallyConsumed: Boolean = false - private[spark] var currentFeedResponseIterator: Option[BufferedIterator[FeedResponse[TSparkRow]]] = None private[spark] var currentItemIterator: Option[BufferedIterator[TSparkRow]] = None @@ -88,39 +85,18 @@ private[spark] class TransientIOErrorsRetryingReadManyByPartitionKeyIterator[TSp val feedResponseIterator = currentFeedResponseIterator match { case Some(existing) => existing case None => - val pagedFlux = customQuery match { - case Some(query) => - container.readManyByPartitionKeys(partitionKeys, query, queryOptions, classType) - case None => - container.readManyByPartitionKeys(partitionKeys, queryOptions, classType) - } - - val rawIterator = new CosmosPagedIterable[TSparkRow]( - pagedFlux, - pageSize, - pagePrefetchBufferSize + val newPagedFlux = cosmosPagedFluxFactory.apply(lastContinuationToken.get) + currentFeedResponseIterator = Some( + new CosmosPagedIterable[TSparkRow]( + newPagedFlux, + pageSize, + pagePrefetchBufferSize + ) + .iterableByPage() + .iterator + .asScala + .buffered ) - .iterableByPage() - .iterator - - // Skip pages already emitted upstream (replay-safe retry). - var skipped: Long = 0 - while (skipped < pagesCommitted && rawIterator.hasNext) { - rawIterator.next() - skipped += 1 - } - if (skipped < pagesCommitted) { - // The server returned fewer pages than before - cannot safely replay. - // Surface a clean error rather than silently emitting a truncated result. - throw new IllegalStateException( - s"readManyByPartitionKeys retry replay failed: expected to skip $pagesCommitted " + - s"already-emitted pages but only $skipped were available. Context: $operationContextString") - } - - // scalastyle:off underscore.import - import scala.collection.JavaConverters._ - // scalastyle:on underscore.import - currentFeedResponseIterator = Some(rawIterator.asScala.buffered) currentFeedResponseIterator.get } @@ -134,13 +110,13 @@ private[spark] class TransientIOErrorsRetryingReadManyByPartitionKeyIterator[TSp } catch { case endToEndTimeoutException: OperationCancelledException => val message = s"End-to-end timeout hit when trying to retrieve the next page. " + - s"Context: $operationContextString" + s"ContinuationToken: $lastContinuationToken, Context: $operationContextString" logError(message, throwable = endToEndTimeoutException) throw endToEndTimeoutException case timeoutException: TimeoutException => val message = s"Attempting to retrieve the next page timed out. " + - s"Context: $operationContextString" + s"ContinuationToken: $lastContinuationToken, Context: $operationContextString" logError(message, timeoutException) val exception = new OperationCancelledException(message, null) exception.setStackTrace(timeoutException.getStackTrace) @@ -156,23 +132,18 @@ private[spark] class TransientIOErrorsRetryingReadManyByPartitionKeyIterator[TSp operationContextAndListener.get.getOperationContext, feedResponse) } - // scalastyle:off underscore.import - import scala.collection.JavaConverters._ - // scalastyle:on underscore.import val iteratorCandidate = feedResponse.getResults.iterator().asScala.buffered + lastContinuationToken.set(feedResponse.getContinuationToken) if (iteratorCandidate.hasNext) { currentItemIterator = Some(iteratorCandidate) - currentPagePartiallyConsumed = false Some(true) } else { - // empty page - count it as committed (no items to replay) and try again - pagesCommitted += 1 + // empty page interleaved - attempt to get next FeedResponse None } } else { // Flux exhausted - currentFeedResponseIterator = None Some(false) } } @@ -183,9 +154,6 @@ private[spark] class TransientIOErrorsRetryingReadManyByPartitionKeyIterator[TSp case Some(iterator) => if (iterator.hasNext) { true } else { - // Entire page drained -> it is now committed for replay-skipping purposes. - pagesCommitted += 1 - currentPagePartiallyConsumed = false currentItemIterator = None false } @@ -194,15 +162,11 @@ private[spark] class TransientIOErrorsRetryingReadManyByPartitionKeyIterator[TSp } override def next(): TSparkRow = { - executeWithRetry("next", () => { - val value = currentItemIterator.get.next() - currentPagePartiallyConsumed = true - value - }) + currentItemIterator.get.next() } override def head: TSparkRow = { - executeWithRetry("head", () => currentItemIterator.get.head) + currentItemIterator.get.head } private[spark] def executeWithRetry[T](methodName: String, func: () => T): T = { @@ -221,17 +185,6 @@ private[spark] class TransientIOErrorsRetryingReadManyByPartitionKeyIterator[TSp catch { case cosmosException: CosmosException => if (Exceptions.canBeTransientFailure(cosmosException.getStatusCode, cosmosException.getSubStatusCode)) { - if (currentPagePartiallyConsumed) { - // We have already emitted items from the current page upstream. Replaying - // the flux would re-skip only completed pages, not items within a page - - // which would cause silent duplication. Fail the task instead. - logError( - s"Transient failure in TransientIOErrorsRetryingReadManyByPartitionKeyIterator." + - s"$methodName after items from the current page were already emitted - " + - s"cannot safely retry without duplicating rows.", - cosmosException) - throw cosmosException - } val retryCountSnapshot = retryCount.incrementAndGet() if (retryCountSnapshot > maxRetryCount) { logError( @@ -244,7 +197,7 @@ private[spark] class TransientIOErrorsRetryingReadManyByPartitionKeyIterator[TSp s"Transient failure handled in " + s"TransientIOErrorsRetryingReadManyByPartitionKeyIterator.$methodName -" + s" will be retried (attempt#$retryCountSnapshot) in ${retryIntervalInMs}ms " + - s"(pagesCommitted=$pagesCommitted)", + s"(continuationToken=$lastContinuationToken)", cosmosException) } } else { @@ -253,7 +206,6 @@ private[spark] class TransientIOErrorsRetryingReadManyByPartitionKeyIterator[TSp case other: Throwable => throw other } - // Reset iterators; pagesCommitted is intentionally preserved so replay can skip them. currentItemIterator = None currentFeedResponseIterator = None Thread.sleep(retryIntervalInMs) @@ -263,10 +215,6 @@ private[spark] class TransientIOErrorsRetryingReadManyByPartitionKeyIterator[TSp returnValue.get } - // Clean up iterator references - the underlying Reactor subscription from - // CosmosPagedIterable.iterator will be cleaned up when the iterator is GC'd. - // This matches the behavior of TransientIOErrorsRetryingIterator; any still-prefetched - // pages are discarded with the iterator. override def close(): Unit = { currentItemIterator = None currentFeedResponseIterator = None @@ -274,17 +222,5 @@ private[spark] class TransientIOErrorsRetryingReadManyByPartitionKeyIterator[TSp } private object TransientIOErrorsRetryingReadManyByPartitionKeyIterator extends BasicLoggingTrait { - private val maxConcurrency = SparkUtils.getNumberOfHostCPUCores - - val executorService: ExecutorService = new ThreadPoolExecutor( - maxConcurrency, - maxConcurrency, - 0L, - TimeUnit.MILLISECONDS, - new SynchronousQueue(), - SparkUtils.daemonThreadFactory(), - new ThreadPoolExecutor.CallerRunsPolicy() - ) - - val executionContext: ExecutionContext = ExecutionContext.fromExecutorService(executorService) + val executionContext: ExecutionContext = TransientIOErrorsRetryingIterator.executionContext } diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/TransientIOErrorsRetryingReadManyIterator.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/TransientIOErrorsRetryingReadManyIterator.scala index c51c5c1226e1..920bc0ac5523 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/TransientIOErrorsRetryingReadManyIterator.scala +++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/TransientIOErrorsRetryingReadManyIterator.scala @@ -10,7 +10,7 @@ import com.azure.cosmos.models.{CosmosItemIdentity, CosmosReadManyRequestOptions import com.azure.cosmos.spark.diagnostics.BasicLoggingTrait import java.time.Duration -import java.util.concurrent.{ExecutorService, SynchronousQueue, ThreadPoolExecutor, TimeUnit, TimeoutException} +import java.util.concurrent.TimeoutException import scala.concurrent.{Await, ExecutionContext, Future} // scalastyle:off underscore.import @@ -174,21 +174,5 @@ private[spark] class TransientIOErrorsRetryingReadManyIterator[TSparkRow] } private object TransientIOErrorsRetryingReadManyIterator extends BasicLoggingTrait { - private val maxConcurrency = SparkUtils.getNumberOfHostCPUCores - - val executorService: ExecutorService = new ThreadPoolExecutor( - maxConcurrency, - maxConcurrency, - 0L, - TimeUnit.MILLISECONDS, - // A synchronous queue does not have any internal capacity, not even a capacity of one. - new SynchronousQueue(), - SparkUtils.daemonThreadFactory(), - // if all worker threads are busy, - // this policy makes the caller thread execute the task. - // This provides a simple feedback control mechanism that will slow down the rate that new tasks are submitted. - new ThreadPoolExecutor.CallerRunsPolicy() - ) - - val executionContext = ExecutionContext.fromExecutorService(executorService) -} \ No newline at end of file + val executionContext: ExecutionContext = TransientIOErrorsRetryingIterator.executionContext +} From 2246e3fac7ec62d84b789cef81cc1d0f916cb995 Mon Sep 17 00:00:00 2001 From: Fabian Meiswinkel Date: Mon, 20 Apr 2026 22:36:21 +0000 Subject: [PATCH 37/64] added test --- .../cosmos/ReadManyByPartitionKeyTest.java | 21 +++++++++ ...ReadManyByPartitionKeyQueryHelperTest.java | 46 +++++++++++++++++++ 2 files changed, 67 insertions(+) diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/ReadManyByPartitionKeyTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/ReadManyByPartitionKeyTest.java index e031cd6abacc..afbd7ba35216 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/ReadManyByPartitionKeyTest.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/ReadManyByPartitionKeyTest.java @@ -375,6 +375,27 @@ public void rejectsOffsetQuery() { } } + @Test(groups = {"emulator"}, timeOut = TIMEOUT) + public void rejectsTopQuery() { + List pkValues = Collections.singletonList(new PartitionKey("pk1")); + SqlQuerySpec topQuery = new SqlQuerySpec("SELECT TOP 5 * FROM c"); + + try { + singlePkContainer.readManyByPartitionKeys(pkValues, topQuery, null, ObjectNode.class) + .stream().collect(Collectors.toList()); + fail("Should have thrown IllegalArgumentException for TOP query"); + } catch (IllegalArgumentException e) { + assertThat(e.getMessage()).contains("TOP"); + } + } + + // DCOUNT, standalone LIMIT, and hybrid/vector/full-text search cannot be tested against the + // emulator: DCOUNT is not recognized as a built-in function, standalone LIMIT is not valid + // Cosmos SQL syntax (only valid with OFFSET, already covered by rejectsOffsetQuery), and + // hybrid search requires vector indexes. All three are covered by unit tests in + // ReadManyByPartitionKeyQueryPlanValidationTest (rejectsDCountQueryPlan, rejectsLimitQueryPlan, + // rejectsHybridSearchQueryPlanWithoutDereferencingNullQueryInfo). + //endregion diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelperTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelperTest.java index 8e2fa373f09e..492bf4dd91f3 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelperTest.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelperTest.java @@ -507,6 +507,52 @@ public void createSelectors_escapesQuotesInPathParts() { //endregion + //region PartitionKey.NONE query generation tests + + @Test(groups = { "unit" }) + public void singlePk_nonePartitionKey_generatesNotIsDefined() { + PartitionKeyDefinition pkDef = createSinglePkDefinition("/mypk"); + List selectors = createSelectors(pkDef); + List pkValues = Collections.singletonList(PartitionKey.NONE); + + SqlQuerySpec result = ReadManyByPartitionKeyQueryHelper.createReadManyByPkQuerySpec( + "SELECT * FROM c", new ArrayList<>(), pkValues, selectors, pkDef); + + assertThat(result.getQueryText()).contains("NOT IS_DEFINED(c[\"mypk\"])"); + assertThat(result.getParameters()).isEmpty(); + } + + @Test(groups = { "unit" }) + public void singlePk_mixedNoneAndNormal_generatesCombinedFilter() { + PartitionKeyDefinition pkDef = createSinglePkDefinition("/mypk"); + List selectors = createSelectors(pkDef); + List pkValues = Arrays.asList(new PartitionKey("pk1"), PartitionKey.NONE); + + SqlQuerySpec result = ReadManyByPartitionKeyQueryHelper.createReadManyByPkQuerySpec( + "SELECT * FROM c", new ArrayList<>(), pkValues, selectors, pkDef); + + assertThat(result.getQueryText()).contains("IN ("); + assertThat(result.getQueryText()).contains("NOT IS_DEFINED(c[\"mypk\"])"); + assertThat(result.getQueryText()).contains("OR"); + } + + @Test(groups = { "unit" }) + public void hpk_nonePartitionKey_generatesNotIsDefinedForAllPaths() { + PartitionKeyDefinition pkDef = createMultiHashPkDefinition("/city", "/zipcode"); + List selectors = createSelectors(pkDef); + List pkValues = Collections.singletonList(PartitionKey.NONE); + + SqlQuerySpec result = ReadManyByPartitionKeyQueryHelper.createReadManyByPkQuerySpec( + "SELECT * FROM c", new ArrayList<>(), pkValues, selectors, pkDef); + + assertThat(result.getQueryText()).contains("NOT IS_DEFINED(c[\"city\"])"); + assertThat(result.getQueryText()).contains("NOT IS_DEFINED(c[\"zipcode\"])"); + assertThat(result.getQueryText()).contains("AND"); + assertThat(result.getParameters()).isEmpty(); + } + + //endregion + //region helpers private PartitionKeyDefinition createSinglePkDefinition(String path) { PartitionKeyDefinition pkDef = new PartitionKeyDefinition(); From fdfa15fca9c082f820b422f60f12825f4f576757 Mon Sep 17 00:00:00 2001 From: Fabian Meiswinkel Date: Tue, 21 Apr 2026 22:39:12 +0000 Subject: [PATCH 38/64] Fixing continuation Token --- ...tionReaderWithReadManyByPartitionKey.scala | 17 +- ...tryingReadManyByPartitionKeyIterator.scala | 6 +- ...ientIOErrorsRetryingReadManyIterator.scala | 6 +- .../cosmos/ReadManyByPartitionKeyTest.java | 151 ++++++ ...nyByPartitionKeyContinuationTokenTest.java | 272 +++++++++++ .../azure/cosmos/CosmosAsyncContainer.java | 15 + .../implementation/AsyncDocumentClient.java | 4 +- .../CosmosReadManyRequestOptionsImpl.java | 24 + .../ImplementationBridgeHelpers.java | 1 + ...adManyByPartitionKeyContinuationToken.java | 305 ++++++++++++ .../implementation/RxDocumentClientImpl.java | 439 ++++++++++++++---- .../models/CosmosReadManyRequestOptions.java | 5 + .../docs/readManyByPartitionKey-design.md | 134 +++++- 13 files changed, 1270 insertions(+), 109 deletions(-) create mode 100644 sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyContinuationTokenTest.java create mode 100644 sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyContinuationToken.java diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/ItemsPartitionReaderWithReadManyByPartitionKey.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/ItemsPartitionReaderWithReadManyByPartitionKey.scala index bd25d0b601eb..e1158dddffaa 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/ItemsPartitionReaderWithReadManyByPartitionKey.scala +++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/ItemsPartitionReaderWithReadManyByPartitionKey.scala @@ -215,12 +215,23 @@ private[spark] case class ItemsPartitionReaderWithReadManyByPartitionKey // Factory that creates a CosmosPagedFlux from an optional continuation token. // On the first call continuationToken is null (start from scratch); on retry // it is the continuation token from the last fully-drained page. - private val fluxFactory: String => CosmosPagedFlux[SparkRowItem] = { (_: String) => + private val fluxFactory: String => CosmosPagedFlux[SparkRowItem] = { (continuationToken: String) => + // Clone options and set continuation token if provided so the SDK + // resumes from where the previous attempt left off. + val effectiveOptions = new CosmosReadManyRequestOptions(readManyOptions) + if (continuationToken != null) { + val effectiveImpl = ImplementationBridgeHelpers + .CosmosReadManyRequestOptionsHelper + .getCosmosReadManyRequestOptionsAccessor + .getImpl(effectiveOptions) + .asInstanceOf[com.azure.cosmos.implementation.CosmosReadManyRequestOptionsImpl] + effectiveImpl.setRequestContinuation(continuationToken) + } customQueryOpt match { case Some(query) => - cosmosAsyncContainer.readManyByPartitionKeys(pkList, query, readManyOptions, classOf[SparkRowItem]) + cosmosAsyncContainer.readManyByPartitionKeys(pkList, query, effectiveOptions, classOf[SparkRowItem]) case None => - cosmosAsyncContainer.readManyByPartitionKeys(pkList, readManyOptions, classOf[SparkRowItem]) + cosmosAsyncContainer.readManyByPartitionKeys(pkList, effectiveOptions, classOf[SparkRowItem]) } } diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/TransientIOErrorsRetryingReadManyByPartitionKeyIterator.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/TransientIOErrorsRetryingReadManyByPartitionKeyIterator.scala index eb2ff8142227..be2e0e2fe299 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/TransientIOErrorsRetryingReadManyByPartitionKeyIterator.scala +++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/TransientIOErrorsRetryingReadManyByPartitionKeyIterator.scala @@ -105,7 +105,7 @@ private[spark] class TransientIOErrorsRetryingReadManyByPartitionKeyIterator[TSp Await.result( Future { feedResponseIterator.hasNext - }(TransientIOErrorsRetryingReadManyByPartitionKeyIterator.executionContext), + }(TransientIOErrorsRetryingIterator.executionContext), maxPageRetrievalTimeout) } catch { case endToEndTimeoutException: OperationCancelledException => @@ -220,7 +220,3 @@ private[spark] class TransientIOErrorsRetryingReadManyByPartitionKeyIterator[TSp currentFeedResponseIterator = None } } - -private object TransientIOErrorsRetryingReadManyByPartitionKeyIterator extends BasicLoggingTrait { - val executionContext: ExecutionContext = TransientIOErrorsRetryingIterator.executionContext -} diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/TransientIOErrorsRetryingReadManyIterator.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/TransientIOErrorsRetryingReadManyIterator.scala index 920bc0ac5523..02feecd6901c 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/TransientIOErrorsRetryingReadManyIterator.scala +++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/TransientIOErrorsRetryingReadManyIterator.scala @@ -87,7 +87,7 @@ private[spark] class TransientIOErrorsRetryingReadManyIterator[TSparkRow] .getCosmosAsyncContainerAccessor .readMany(container, readManyFilterList.asJava, queryOptionsWithEnd2EndTimeout, classType) .block() - }(TransientIOErrorsRetryingReadManyIterator.executionContext), + }(TransientIOErrorsRetryingIterator.executionContext), maxPageRetrievalTimeout) } catch { case endToEndTimeoutException: OperationCancelledException => @@ -172,7 +172,3 @@ private[spark] class TransientIOErrorsRetryingReadManyIterator[TSparkRow] override def close(): Unit = {} } - -private object TransientIOErrorsRetryingReadManyIterator extends BasicLoggingTrait { - val executionContext: ExecutionContext = TransientIOErrorsRetryingIterator.executionContext -} diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/ReadManyByPartitionKeyTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/ReadManyByPartitionKeyTest.java index afbd7ba35216..508725a23112 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/ReadManyByPartitionKeyTest.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/ReadManyByPartitionKeyTest.java @@ -5,6 +5,8 @@ package com.azure.cosmos; +import com.azure.cosmos.implementation.CosmosReadManyRequestOptionsImpl; +import com.azure.cosmos.implementation.ImplementationBridgeHelpers; import com.azure.cosmos.models.CosmosContainerProperties; import com.azure.cosmos.models.CosmosItemRequestOptions; import com.azure.cosmos.models.FeedResponse; @@ -484,6 +486,155 @@ public T deserialize(Map jsonNodeMap, Class classType) { //endregion + //region Continuation token tests + + @Test(groups = {"emulator"}, timeOut = TIMEOUT) + public void singlePk_sequentialExecution_emitsContinuationTokens() { + // Use small batch size so we get multiple batches (and thus multiple continuation tokens) + String originalValue = System.getProperty("COSMOS.READ_MANY_BY_PK_MAX_BATCH_SIZE"); + try { + System.setProperty("COSMOS.READ_MANY_BY_PK_MAX_BATCH_SIZE", "2"); + + // Create items across 3 PKs + createSinglePkItems("seqPk1", 3); + createSinglePkItems("seqPk2", 3); + createSinglePkItems("seqPk3", 3); + + List pkValues = Arrays.asList( + new PartitionKey("seqPk1"), + new PartitionKey("seqPk2"), + new PartitionKey("seqPk3")); + + // Use the async container to collect FeedResponse pages with continuation tokens + CosmosAsyncContainer asyncContainer = client.asyncClient() + .getDatabase(preExistingDatabaseId) + .getContainer(singlePkContainer.getId()); + + List> pages = asyncContainer + .readManyByPartitionKeys(pkValues, ObjectNode.class) + .byPage() + .collectList() + .block(); + + assertThat(pages).isNotNull(); + assertThat(pages).isNotEmpty(); + + // All non-final pages should have a non-null continuation token + for (int i = 0; i < pages.size() - 1; i++) { + assertThat(pages.get(i).getContinuationToken()) + .as("Page %d should have a continuation token", i) + .isNotNull() + .isNotEmpty(); + } + + // The final page should have null continuation + assertThat(pages.get(pages.size() - 1).getContinuationToken()) + .as("Last page should have null continuation token") + .isNull(); + + // Total items should be 9 + long totalItems = pages.stream() + .mapToLong(p -> p.getResults().size()) + .sum(); + assertThat(totalItems).isEqualTo(9); + + cleanupContainer(singlePkContainer); + } finally { + if (originalValue != null) { + System.setProperty("COSMOS.READ_MANY_BY_PK_MAX_BATCH_SIZE", originalValue); + } else { + System.clearProperty("COSMOS.READ_MANY_BY_PK_MAX_BATCH_SIZE"); + } + } + } + + @Test(groups = {"emulator"}, timeOut = TIMEOUT) + public void singlePk_continuationToken_resumesCorrectly() { + // Use small batch size to force multiple batches + String originalValue = System.getProperty("COSMOS.READ_MANY_BY_PK_MAX_BATCH_SIZE"); + try { + System.setProperty("COSMOS.READ_MANY_BY_PK_MAX_BATCH_SIZE", "1"); + + // Create items across 3 PKs + createSinglePkItems("resumePk1", 2); + createSinglePkItems("resumePk2", 2); + createSinglePkItems("resumePk3", 2); + + List pkValues = Arrays.asList( + new PartitionKey("resumePk1"), + new PartitionKey("resumePk2"), + new PartitionKey("resumePk3")); + + CosmosAsyncContainer asyncContainer = client.asyncClient() + .getDatabase(preExistingDatabaseId) + .getContainer(singlePkContainer.getId()); + + // First pass: collect all items + List> allPages = asyncContainer + .readManyByPartitionKeys(pkValues, ObjectNode.class) + .byPage() + .collectList() + .block(); + + assertThat(allPages).isNotNull(); + assertThat(allPages.size()).isGreaterThan(1); + + // Pick a continuation token from the first page + String continuationAfterFirstPage = allPages.get(0).getContinuationToken(); + assertThat(continuationAfterFirstPage).isNotNull(); + List itemsFromFirstPage = allPages.get(0).getResults(); + + // Second pass: resume from the continuation token + com.azure.cosmos.models.CosmosReadManyRequestOptions options2 = + new com.azure.cosmos.models.CosmosReadManyRequestOptions(); + ((com.azure.cosmos.implementation.CosmosReadManyRequestOptionsImpl) + ImplementationBridgeHelpers.CosmosReadManyRequestOptionsHelper + .getCosmosReadManyRequestOptionsAccessor() + .getImpl(options2)) + .setRequestContinuation(continuationAfterFirstPage); + + List> remainingPages = asyncContainer + .readManyByPartitionKeys(pkValues, options2, ObjectNode.class) + .byPage() + .collectList() + .block(); + + assertThat(remainingPages).isNotNull(); + + // Collect all item ids + List firstPageIds = itemsFromFirstPage.stream() + .map(n -> n.get("id").asText()) + .collect(Collectors.toList()); + List remainingIds = remainingPages.stream() + .flatMap(p -> p.getResults().stream()) + .map(n -> n.get("id").asText()) + .collect(Collectors.toList()); + List allIds = allPages.stream() + .flatMap(p -> p.getResults().stream()) + .map(n -> n.get("id").asText()) + .collect(Collectors.toList()); + + // Items from first page + remaining should equal all items (no overlap, no gap) + List combined = new ArrayList<>(firstPageIds); + combined.addAll(remainingIds); + + assertThat(combined).hasSameElementsAs(allIds); + + // No duplicates + assertThat(combined).doesNotHaveDuplicates(); + + cleanupContainer(singlePkContainer); + } finally { + if (originalValue != null) { + System.setProperty("COSMOS.READ_MANY_BY_PK_MAX_BATCH_SIZE", originalValue); + } else { + System.clearProperty("COSMOS.READ_MANY_BY_PK_MAX_BATCH_SIZE"); + } + } + } + + //endregion + //region helper methods private List createSinglePkItems(String pkValue, int count) { diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyContinuationTokenTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyContinuationTokenTest.java new file mode 100644 index 000000000000..f92c2c902e06 --- /dev/null +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyContinuationTokenTest.java @@ -0,0 +1,272 @@ +/* + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. + */ + +package com.azure.cosmos.implementation; + +import com.azure.cosmos.implementation.routing.Range; +import com.azure.cosmos.models.SqlParameter; +import com.azure.cosmos.models.SqlQuerySpec; +import org.testng.annotations.Test; + +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.List; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +public class ReadManyByPartitionKeyContinuationTokenTest { + + private static final String TEST_COLLECTION_RID = "dbs/testDb/colls/testColl"; + private static final int TEST_QUERY_HASH = 12345; + + @Test(groups = { "unit" }) + public void roundtrip_withBackendContinuation() { + List> remaining = Arrays.asList( + new Range<>("05C1E0", "0BF333", true, false), + new Range<>("0BF333", "FF", true, false)); + Range current = new Range<>("", "05C1E0", true, false); + String backendCont = "eyJDb21wb3NpdGVUb2tlbg=="; + + ReadManyByPartitionKeyContinuationToken token = + new ReadManyByPartitionKeyContinuationToken( + remaining, current, backendCont, TEST_COLLECTION_RID, TEST_QUERY_HASH); + + String serialized = token.serialize(); + assertThat(serialized).isNotNull().isNotEmpty(); + + ReadManyByPartitionKeyContinuationToken deserialized = + ReadManyByPartitionKeyContinuationToken.deserialize(serialized); + + assertThat(deserialized.getBackendContinuation()).isEqualTo(backendCont); + assertThat(deserialized.getCollectionRid()).isEqualTo(TEST_COLLECTION_RID); + assertThat(deserialized.getQueryHash()).isEqualTo(TEST_QUERY_HASH); + + ReadManyByPartitionKeyContinuationToken.BatchDefinition currentBatch = deserialized.getCurrentBatch(); + assertThat(currentBatch.getPartitionScope().getMin()).isEqualTo(""); + assertThat(currentBatch.getPartitionScope().getMax()).isEqualTo("05C1E0"); + + List remainingBatches = + deserialized.getRemainingBatches(); + assertThat(remainingBatches).hasSize(2); + assertThat(remainingBatches.get(0).getPartitionScope().getMin()).isEqualTo("05C1E0"); + assertThat(remainingBatches.get(0).getPartitionScope().getMax()).isEqualTo("0BF333"); + assertThat(remainingBatches.get(1).getPartitionScope().getMin()).isEqualTo("0BF333"); + assertThat(remainingBatches.get(1).getPartitionScope().getMax()).isEqualTo("FF"); + } + + @Test(groups = { "unit" }) + public void roundtrip_withNullBackendContinuation() { + List> remaining = Collections.singletonList( + new Range<>("0BF333", "FF", true, false)); + Range current = new Range<>("05C1E0", "0BF333", true, false); + + ReadManyByPartitionKeyContinuationToken token = + new ReadManyByPartitionKeyContinuationToken( + remaining, current, null, TEST_COLLECTION_RID, TEST_QUERY_HASH); + + String serialized = token.serialize(); + ReadManyByPartitionKeyContinuationToken deserialized = + ReadManyByPartitionKeyContinuationToken.deserialize(serialized); + + assertThat(deserialized.getBackendContinuation()).isNull(); + assertThat(deserialized.getCurrentBatch().getPartitionScope().getMin()).isEqualTo("05C1E0"); + assertThat(deserialized.getCurrentBatch().getPartitionScope().getMax()).isEqualTo("0BF333"); + assertThat(deserialized.getRemainingBatches()).hasSize(1); + } + + @Test(groups = { "unit" }) + public void roundtrip_emptyRemainingBatches() { + Range current = new Range<>("0BF333", "FF", true, false); + + ReadManyByPartitionKeyContinuationToken token = + new ReadManyByPartitionKeyContinuationToken( + Collections.emptyList(), current, "someCont", TEST_COLLECTION_RID, TEST_QUERY_HASH); + + String serialized = token.serialize(); + ReadManyByPartitionKeyContinuationToken deserialized = + ReadManyByPartitionKeyContinuationToken.deserialize(serialized); + + assertThat(deserialized.getRemainingBatches()).isEmpty(); + assertThat(deserialized.getCurrentBatch().getPartitionScope().getMin()).isEqualTo("0BF333"); + assertThat(deserialized.getBackendContinuation()).isEqualTo("someCont"); + } + + @Test(groups = { "unit" }) + public void roundtrip_lastBatchNoContinuation() { + Range current = new Range<>("0BF333", "FF", true, false); + + ReadManyByPartitionKeyContinuationToken token = + new ReadManyByPartitionKeyContinuationToken( + Collections.emptyList(), current, null, TEST_COLLECTION_RID, TEST_QUERY_HASH); + + String serialized = token.serialize(); + ReadManyByPartitionKeyContinuationToken deserialized = + ReadManyByPartitionKeyContinuationToken.deserialize(serialized); + + assertThat(deserialized.getRemainingBatches()).isEmpty(); + assertThat(deserialized.getBackendContinuation()).isNull(); + } + + @Test(groups = { "unit" }) + public void deserialize_malformedInput_throws() { + assertThatThrownBy(() -> + ReadManyByPartitionKeyContinuationToken.deserialize("not-valid-base64!!!") + ).isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("Failed to deserialize"); + } + + @Test(groups = { "unit" }) + public void deserialize_emptyString_throws() { + assertThatThrownBy(() -> + ReadManyByPartitionKeyContinuationToken.deserialize("") + ).isInstanceOf(IllegalArgumentException.class); + } + + @Test(groups = { "unit" }) + public void deserialize_null_throws() { + assertThatThrownBy(() -> + ReadManyByPartitionKeyContinuationToken.deserialize(null) + ).isInstanceOf(NullPointerException.class); + } + + @Test(groups = { "unit" }) + public void serialized_isBase64() { + Range current = new Range<>("", "FF", true, false); + ReadManyByPartitionKeyContinuationToken token = + new ReadManyByPartitionKeyContinuationToken( + Collections.emptyList(), current, null, TEST_COLLECTION_RID, 0); + + String serialized = token.serialize(); + + // Should be valid Base64 (no whitespace, no special chars except +/=) + assertThat(serialized).matches("[A-Za-z0-9+/=]+"); + + // Decoding should produce valid JSON + String json = new String( + java.util.Base64.getDecoder().decode(serialized), + java.nio.charset.StandardCharsets.UTF_8); + assertThat(json).startsWith("{"); + assertThat(json).endsWith("}"); + } + + @Test(groups = { "unit" }) + public void rangesPreserveMinMaxInclusive() { + Range current = new Range<>("AB", "CD", true, false); + ReadManyByPartitionKeyContinuationToken token = + new ReadManyByPartitionKeyContinuationToken( + Collections.singletonList(new Range<>("CD", "EF", true, false)), + current, null, TEST_COLLECTION_RID, TEST_QUERY_HASH); + + String serialized = token.serialize(); + ReadManyByPartitionKeyContinuationToken deserialized = + ReadManyByPartitionKeyContinuationToken.deserialize(serialized); + + // Verify that deserialized ranges have the canonical min-inclusive, max-exclusive form + Range currentScope = deserialized.getCurrentBatch().getPartitionScope(); + assertThat(currentScope.isMinInclusive()).isTrue(); + assertThat(currentScope.isMaxInclusive()).isFalse(); + Range remainingScope = deserialized.getRemainingBatches().get(0).getPartitionScope(); + assertThat(remainingScope.isMinInclusive()).isTrue(); + assertThat(remainingScope.isMaxInclusive()).isFalse(); + } + + @Test(groups = { "unit" }) + public void collectionRidAndQueryHash_roundtrip() { + Range current = new Range<>("", "FF", true, false); + String rid = "dbs/myDb/colls/myColl"; + int hash = 98765; + + ReadManyByPartitionKeyContinuationToken token = + new ReadManyByPartitionKeyContinuationToken( + Collections.emptyList(), current, null, rid, hash); + + String serialized = token.serialize(); + ReadManyByPartitionKeyContinuationToken deserialized = + ReadManyByPartitionKeyContinuationToken.deserialize(serialized); + + assertThat(deserialized.getCollectionRid()).isEqualTo(rid); + assertThat(deserialized.getQueryHash()).isEqualTo(hash); + } + + @Test(groups = { "unit" }) + public void computeQueryHash_nullSpec_returnsZero() { + assertThat(ReadManyByPartitionKeyContinuationToken.computeQueryHash(null)).isEqualTo(0); + } + + @Test(groups = { "unit" }) + public void computeQueryHash_sameQueryText_sameHash() { + SqlQuerySpec spec1 = new SqlQuerySpec("SELECT * FROM c WHERE c.pk = @pk", + Collections.singletonList(new SqlParameter("@pk", "value1"))); + SqlQuerySpec spec2 = new SqlQuerySpec("SELECT * FROM c WHERE c.pk = @pk", + Collections.singletonList(new SqlParameter("@pk", "value1"))); + + assertThat(ReadManyByPartitionKeyContinuationToken.computeQueryHash(spec1)) + .isEqualTo(ReadManyByPartitionKeyContinuationToken.computeQueryHash(spec2)); + } + + @Test(groups = { "unit" }) + public void computeQueryHash_differentParams_differentHash() { + SqlQuerySpec spec1 = new SqlQuerySpec("SELECT * FROM c WHERE c.pk = @pk", + Collections.singletonList(new SqlParameter("@pk", "value1"))); + SqlQuerySpec spec2 = new SqlQuerySpec("SELECT * FROM c WHERE c.pk = @pk", + Collections.singletonList(new SqlParameter("@pk", "value2"))); + + assertThat(ReadManyByPartitionKeyContinuationToken.computeQueryHash(spec1)) + .isNotEqualTo(ReadManyByPartitionKeyContinuationToken.computeQueryHash(spec2)); + } + + @Test(groups = { "unit" }) + public void computeQueryHash_differentQueryText_differentHash() { + SqlQuerySpec spec1 = new SqlQuerySpec("SELECT * FROM c"); + SqlQuerySpec spec2 = new SqlQuerySpec("SELECT c.id FROM c"); + + assertThat(ReadManyByPartitionKeyContinuationToken.computeQueryHash(spec1)) + .isNotEqualTo(ReadManyByPartitionKeyContinuationToken.computeQueryHash(spec2)); + } + + @Test(groups = { "unit" }) + public void computeQueryHash_noParams_stableHash() { + SqlQuerySpec spec = new SqlQuerySpec("SELECT * FROM c"); + int hash1 = ReadManyByPartitionKeyContinuationToken.computeQueryHash(spec); + int hash2 = ReadManyByPartitionKeyContinuationToken.computeQueryHash(spec); + assertThat(hash1).isEqualTo(hash2); + } + + @Test(groups = { "unit" }) + public void batchDefinition_roundtrip() { + Range partitionScope = new Range<>("", "05C1E0", true, false); + Range batchFilter = new Range<>("01", "03", true, false); + + ReadManyByPartitionKeyContinuationToken.BatchDefinition bd = + new ReadManyByPartitionKeyContinuationToken.BatchDefinition(partitionScope, batchFilter); + + List remaining = Collections.singletonList( + new ReadManyByPartitionKeyContinuationToken.BatchDefinition( + new Range<>("05C1E0", "FF", true, false), + new Range<>("05C1E0", "0A", true, false))); + + ReadManyByPartitionKeyContinuationToken token = + new ReadManyByPartitionKeyContinuationToken( + remaining, bd, "cont", TEST_COLLECTION_RID, TEST_QUERY_HASH); + + String serialized = token.serialize(); + ReadManyByPartitionKeyContinuationToken deserialized = + ReadManyByPartitionKeyContinuationToken.deserialize(serialized); + + ReadManyByPartitionKeyContinuationToken.BatchDefinition currentBatch = deserialized.getCurrentBatch(); + assertThat(currentBatch.getPartitionScope().getMin()).isEqualTo(""); + assertThat(currentBatch.getPartitionScope().getMax()).isEqualTo("05C1E0"); + assertThat(currentBatch.getBatchFilter().getMin()).isEqualTo("01"); + assertThat(currentBatch.getBatchFilter().getMax()).isEqualTo("03"); + + ReadManyByPartitionKeyContinuationToken.BatchDefinition remainingBatch = + deserialized.getRemainingBatches().get(0); + assertThat(remainingBatch.getPartitionScope().getMin()).isEqualTo("05C1E0"); + assertThat(remainingBatch.getBatchFilter().getMin()).isEqualTo("05C1E0"); + assertThat(remainingBatch.getBatchFilter().getMax()).isEqualTo("0A"); + } +} diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java index b24ab9991333..fbc3508d5be1 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java @@ -1721,6 +1721,13 @@ private Function>> readManyByPa CosmosAsyncClient client = this.getDatabase().getClient(); + // Extract continuation token before entering the reactive chain. + // It will be set on the cloned CosmosQueryRequestOptions so that + // QueryFeedOperationState picks it up automatically. + String requestContinuation = requestOptions != null + ? readManyOptionsAccessor().getRequestContinuation(requestOptions) + : null; + return (pagedFluxOptions -> { CosmosQueryRequestOptions queryRequestOptions = requestOptions == null ? new CosmosQueryRequestOptions() @@ -1732,6 +1739,14 @@ private Function>> readManyByPa queryRequestOptions.setMaxDegreeOfParallelism(-1); } queryRequestOptions.setQueryName("readManyByPartitionKeys"); + + // Set the composite continuation token on the cloned query options so that + // QueryFeedOperationState carries it through to RxDocumentClientImpl. + if (requestContinuation != null) { + ModelBridgeInternal.setQueryRequestOptionsContinuationToken( + queryRequestOptions, requestContinuation); + } + CosmosQueryRequestOptionsBase cosmosQueryRequestOptionsImpl = queryOptionsAccessor().getImpl(queryRequestOptions); applyPolicies(OperationType.Query, ResourceType.Document, cosmosQueryRequestOptionsImpl, this.readManyByPartitionKeySpanName); diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/AsyncDocumentClient.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/AsyncDocumentClient.java index d2d6413489ad..0996af4010b3 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/AsyncDocumentClient.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/AsyncDocumentClient.java @@ -1593,7 +1593,7 @@ Mono> readMany( * @param partitionKeys list of partition key values to read documents for * @param customQuery optional custom query (for projections/additional filters) - null means SELECT * FROM c * @param collectionLink link for the documentcollection/container to be queried - * @param state the query operation state + * @param state the query operation state (may carry a composite continuation token via requestContinuation) * @param klass class type * @param the type parameter * @return a Flux with feed response pages of documents @@ -1672,7 +1672,7 @@ Flux> readAllDocuments( */ void enableSDKThroughputControlGroup(SDKThroughputControlGroupInternal group, Mono throughputQueryMono); - + /*** * Enable server throughput control group. * diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/CosmosReadManyRequestOptionsImpl.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/CosmosReadManyRequestOptionsImpl.java index 05ee42c66753..9828a4cb7ad6 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/CosmosReadManyRequestOptionsImpl.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/CosmosReadManyRequestOptionsImpl.java @@ -3,6 +3,9 @@ package com.azure.cosmos.implementation; public class CosmosReadManyRequestOptionsImpl extends CosmosQueryRequestOptionsBase { + + private String requestContinuation; + /** * Instantiates a new read many request options. */ @@ -16,6 +19,27 @@ public CosmosReadManyRequestOptionsImpl() { */ public CosmosReadManyRequestOptionsImpl(CosmosReadManyRequestOptionsImpl options) { super(options); + this.requestContinuation = options.requestContinuation; + } + + /** + * Gets the composite continuation token for readManyByPartitionKeys. + * + * @return the continuation token, or null if not set. + */ + public String getRequestContinuation() { + return this.requestContinuation; + } + + /** + * Sets the composite continuation token for readManyByPartitionKeys. + * + * @param requestContinuation the continuation token from a previous invocation. + * @return this instance. + */ + public CosmosReadManyRequestOptionsImpl setRequestContinuation(String requestContinuation) { + this.requestContinuation = requestContinuation; + return this; } @Override diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ImplementationBridgeHelpers.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ImplementationBridgeHelpers.java index ee98360d2af4..246e9f232317 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ImplementationBridgeHelpers.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ImplementationBridgeHelpers.java @@ -351,6 +351,7 @@ public static CosmosReadManyRequestOptionsAccessor getCosmosReadManyRequestOptio public interface CosmosReadManyRequestOptionsAccessor { CosmosQueryRequestOptionsBase getImpl(CosmosReadManyRequestOptions options); + String getRequestContinuation(CosmosReadManyRequestOptions options); } } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyContinuationToken.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyContinuationToken.java new file mode 100644 index 000000000000..5c40a22c63ad --- /dev/null +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyContinuationToken.java @@ -0,0 +1,305 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. +package com.azure.cosmos.implementation; + +import com.azure.cosmos.implementation.routing.Range; +import com.azure.cosmos.models.SqlParameter; +import com.azure.cosmos.models.SqlQuerySpec; +import com.fasterxml.jackson.annotation.JsonCreator; +import com.fasterxml.jackson.annotation.JsonIgnore; +import com.fasterxml.jackson.annotation.JsonProperty; + +import java.nio.charset.StandardCharsets; +import java.util.ArrayList; +import java.util.Base64; +import java.util.List; + +import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkArgument; +import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkNotNull; + +/** + * Composite continuation token for {@code readManyByPartitionKeys} operations. + *

+ * Captures the state needed to resume a readManyByPartitionKeys operation: + *

    + *
  • {@code remainingBatches} — batch definitions of batches not yet started
    • + *
  • {@code currentBatch} — batch definition of the batch currently being processed
    • + *
  • {@code backendContinuation} — backend query continuation within the current batch (nullable)
    • + *
+ * Each batch definition has two EPK ranges: + *
    + *
  • {@code partitionScope} — the physical partition's EPK range at the time the operation + * started. Used for routing queries (via FeedRange). Only causes fan-out overhead if a + * split has actually occurred.
    • + *
  • {@code batchFilter} — the EPK sub-range that identifies which PKs belong to this batch. + * Within a physical partition, PKs are sorted by EPK and split into batches; the filter + * range spans the EPKs of the PKs in this batch. Used to reconstruct the exact same set + * of PKs per batch when resuming from a continuation token.
    • + *
+ * EPK ranges are used instead of PkRangeIds so the token survives partition splits. + *

+ * Serialized as JSON → Base64 to keep the token opaque. + */ +public final class ReadManyByPartitionKeyContinuationToken { + + private static final String REMAINING_BATCHES_PROPERTY = "rb"; + private static final String CURRENT_BATCH_PROPERTY = "cb"; + private static final String BACKEND_CONTINUATION_PROPERTY = "bc"; + private static final String COLLECTION_RID_PROPERTY = "cr"; + private static final String QUERY_HASH_PROPERTY = "qh"; + + @JsonProperty(REMAINING_BATCHES_PROPERTY) + private final List remainingBatches; + + @JsonProperty(CURRENT_BATCH_PROPERTY) + private final BatchDefinitionDto currentBatch; + + @JsonProperty(BACKEND_CONTINUATION_PROPERTY) + private final String backendContinuation; + + @JsonProperty(COLLECTION_RID_PROPERTY) + private final String collectionRid; + + @JsonProperty(QUERY_HASH_PROPERTY) + private final int queryHash; + + @JsonCreator + ReadManyByPartitionKeyContinuationToken( + @JsonProperty(REMAINING_BATCHES_PROPERTY) List remainingBatches, + @JsonProperty(CURRENT_BATCH_PROPERTY) BatchDefinitionDto currentBatch, + @JsonProperty(BACKEND_CONTINUATION_PROPERTY) String backendContinuation, + @JsonProperty(COLLECTION_RID_PROPERTY) String collectionRid, + @JsonProperty(QUERY_HASH_PROPERTY) int queryHash) { + + this.remainingBatches = remainingBatches; + this.currentBatch = currentBatch; + this.backendContinuation = backendContinuation; + this.collectionRid = collectionRid; + this.queryHash = queryHash; + } + + public ReadManyByPartitionKeyContinuationToken( + List remainingBatches, + BatchDefinition currentBatch, + String backendContinuation, + String collectionRid, + int queryHash) { + + checkNotNull(currentBatch, "Argument 'currentBatch' must not be null."); + checkNotNull(remainingBatches, "Argument 'remainingBatches' must not be null."); + + this.remainingBatches = new ArrayList<>(remainingBatches.size()); + for (BatchDefinition bd : remainingBatches) { + this.remainingBatches.add(BatchDefinitionDto.fromBatchDefinition(bd)); + } + this.currentBatch = BatchDefinitionDto.fromBatchDefinition(currentBatch); + this.backendContinuation = backendContinuation; + this.collectionRid = collectionRid; + this.queryHash = queryHash; + } + + /** + * Convenience constructor that accepts EPK ranges directly (without separate partitionScope/batchFilter). + * Each range is used as both the partitionScope and batchFilter in the BatchDefinition. + * This is an interim API until the full batch redesign with distinct partitionScope/batchFilter is wired up. + */ + public ReadManyByPartitionKeyContinuationToken( + List> remainingBatchRanges, + Range currentBatchRange, + String backendContinuation, + String collectionRid, + int queryHash) { + + checkNotNull(currentBatchRange, "Argument 'currentBatchRange' must not be null."); + checkNotNull(remainingBatchRanges, "Argument 'remainingBatchRanges' must not be null."); + + this.remainingBatches = new ArrayList<>(remainingBatchRanges.size()); + for (Range range : remainingBatchRanges) { + BatchDefinition bd = new BatchDefinition(range, range); + this.remainingBatches.add(BatchDefinitionDto.fromBatchDefinition(bd)); + } + this.currentBatch = BatchDefinitionDto.fromBatchDefinition( + new BatchDefinition(currentBatchRange, currentBatchRange)); + this.backendContinuation = backendContinuation; + this.collectionRid = collectionRid; + this.queryHash = queryHash; + } + + @JsonIgnore + public List getRemainingBatches() { + List result = new ArrayList<>(remainingBatches.size()); + for (BatchDefinitionDto dto : remainingBatches) { + result.add(dto.toBatchDefinition()); + } + return result; + } + + @JsonIgnore + public BatchDefinition getCurrentBatch() { + return currentBatch.toBatchDefinition(); + } + + public String getBackendContinuation() { + return backendContinuation; + } + + public String getCollectionRid() { + return collectionRid; + } + + public int getQueryHash() { + return queryHash; + } + + /** + * Computes a stable hash for a SqlQuerySpec (or null for the default SELECT * FROM c query). + * Hashes over both the query text and all parameter names/values to detect when a continuation + * token is reused with a different query or different parameter values. + */ + public static int computeQueryHash(SqlQuerySpec querySpec) { + if (querySpec == null) { + return 0; + } + int hash = 17; + String queryText = querySpec.getQueryText(); + hash = 31 * hash + (queryText != null ? queryText.hashCode() : 0); + List params = querySpec.getParameters(); + if (params != null) { + for (SqlParameter param : params) { + hash = 31 * hash + (param.getName() != null ? param.getName().hashCode() : 0); + Object value = param.getValue(Object.class); + hash = 31 * hash + (value != null ? value.hashCode() : 0); + } + } + return hash; + } + + /** + * Serializes this token to a Base64-encoded JSON string. + */ + public String serialize() { + try { + String json = Utils.getSimpleObjectMapper().writeValueAsString(this); + return Base64.getEncoder().encodeToString(json.getBytes(StandardCharsets.UTF_8)); + } catch (Exception e) { + throw new IllegalStateException("Failed to serialize ReadManyByPartitionKeyContinuationToken.", e); + } + } + + /** + * Deserializes a Base64-encoded JSON string into a continuation token. + * + * @param serialized the serialized token (Base64 of JSON) + * @return the deserialized token + * @throws IllegalArgumentException if the token is malformed + */ + public static ReadManyByPartitionKeyContinuationToken deserialize(String serialized) { + checkNotNull(serialized, "Argument 'serialized' must not be null."); + checkArgument(!serialized.isEmpty(), "Argument 'serialized' must not be empty."); + + try { + byte[] decoded = Base64.getDecoder().decode(serialized); + String json = new String(decoded, StandardCharsets.UTF_8); + return Utils.getSimpleObjectMapper().readValue(json, ReadManyByPartitionKeyContinuationToken.class); + } catch (Exception e) { + throw new IllegalArgumentException( + "Failed to deserialize ReadManyByPartitionKeyContinuationToken. " + + "The continuation token may be malformed or from an incompatible version.", e); + } + } + + /** + * Identifies a single batch in a readManyByPartitionKeys operation. + *

+ * Each batch has two EPK ranges: + *

    + *
  • {@code partitionScope} — the physical partition's full EPK range. Used for routing + * the query to the correct physical partition(s) via FeedRange.
    • + *
  • {@code batchFilter} — the EPK sub-range within the partition that identifies which + * PKs belong to this batch. PKs whose EPK falls within [filterMin, filterMax) are + * part of this batch.
    • + *
+ */ + public static final class BatchDefinition { + private final Range partitionScope; + private final Range batchFilter; + + public BatchDefinition(Range partitionScope, Range batchFilter) { + this.partitionScope = checkNotNull(partitionScope, "Argument 'partitionScope' must not be null."); + this.batchFilter = checkNotNull(batchFilter, "Argument 'batchFilter' must not be null."); + } + + public Range getPartitionScope() { + return partitionScope; + } + + public Range getBatchFilter() { + return batchFilter; + } + } + + /** + * Compact DTO for JSON serialization of a batch definition. + */ + static final class BatchDefinitionDto { + private final EpkRangeDto ps; + private final EpkRangeDto bf; + + @JsonCreator + BatchDefinitionDto( + @JsonProperty("ps") EpkRangeDto ps, + @JsonProperty("bf") EpkRangeDto bf) { + this.ps = ps; + this.bf = bf; + } + + @JsonProperty("ps") + EpkRangeDto getPs() { return ps; } + + @JsonProperty("bf") + EpkRangeDto getBf() { return bf; } + + static BatchDefinitionDto fromBatchDefinition(BatchDefinition bd) { + return new BatchDefinitionDto( + EpkRangeDto.fromRange(bd.partitionScope), + EpkRangeDto.fromRange(bd.batchFilter)); + } + + BatchDefinition toBatchDefinition() { + return new BatchDefinition(ps.toRange(), bf.toRange()); + } + } + + /** + * Minimal DTO for EPK range serialization. Uses short field names to keep the + * serialized token compact. + */ + static final class EpkRangeDto { + @JsonProperty("min") + private final String min; + @JsonProperty("max") + private final String max; + + @JsonCreator + EpkRangeDto( + @JsonProperty("min") String min, + @JsonProperty("max") String max) { + this.min = min; + this.max = max; + } + + static EpkRangeDto fromRange(Range range) { + if (!range.isMinInclusive()) { + throw new IllegalArgumentException("EPK ranges must be minInclusive."); + } + if (range.isMaxInclusive()) { + throw new IllegalArgumentException("EPK ranges must be maxExclusive."); + } + return new EpkRangeDto(range.getMin(), range.getMax()); + } + + Range toRange() { + return new Range<>(min, max, true, false); + } + } +} diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java index 1e5d7f9e8f50..e1a4518c9805 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java @@ -118,6 +118,7 @@ import java.util.Arrays; import java.util.Collection; import java.util.Collections; +import java.util.Comparator; import java.util.EnumSet; import java.util.HashMap; import java.util.HashSet; @@ -4439,6 +4440,8 @@ private Flux> readManyByPartitionKeys( ScopedDiagnosticsFactory diagnosticsFactory, Class klass) { + String requestContinuation = state.getRequestContinuation(); + String resourceLink = parentResourceLinkToQueryLink(collectionLink, ResourceType.Document); RxDocumentServiceRequest request = RxDocumentServiceRequest.create(diagnosticsFactory, OperationType.Query, @@ -4458,14 +4461,34 @@ private Flux> readManyByPartitionKeys( final PartitionKeyDefinition pkDefinition = collection.getPartitionKey(); - Mono> valueHolderMono = partitionKeyRangeCache - .tryLookupAsync( - BridgeInternal.getMetaDataDiagnosticContext(request.requestContext.cosmosDiagnostics), - collection.getResourceId(), - null, - null); + final String collectionRid = collection.getResourceId(); + final int queryHash = ReadManyByPartitionKeyContinuationToken.computeQueryHash(customQuery); + + // When a continuation token is present, skip the routing map lookup and batch + // construction - the token already contains the batch definitions. Only resolve + // the routing map and build batches on the very first call (no continuation). + if (requestContinuation != null) { + // Continuation path: validate collection/query, then resume from token + ReadManyByPartitionKeyContinuationToken parsedContinuation = + ReadManyByPartitionKeyContinuationToken.deserialize(requestContinuation); + if (!collectionRid.equals(parsedContinuation.getCollectionRid())) { + return Flux.error(new IllegalArgumentException( + "Continuation token was created for a different collection (rid mismatch). " + + "Expected: " + collectionRid + ", token has: " + parsedContinuation.getCollectionRid())); + } + if (queryHash != parsedContinuation.getQueryHash()) { + return Flux.error(new IllegalArgumentException( + "Continuation token was created with a different query (hash mismatch). " + + "The same query must be used when resuming from a continuation token.")); + } - // Validate custom query if provided + return buildSequentialFluxFromContinuation( + parsedContinuation, partitionKeys, customQuery, pkDefinition, + resourceLink, state, diagnosticsFactory, klass, + collectionRid, queryHash); + } + + // First-call path: validate custom query, resolve routing map, build batches Mono queryValidationMono; if (customQuery != null) { queryValidationMono = validateCustomQueryForReadManyByPartitionKey( @@ -4474,6 +4497,13 @@ private Flux> readManyByPartitionKeys( queryValidationMono = Mono.empty(); } + Mono> valueHolderMono = partitionKeyRangeCache + .tryLookupAsync( + BridgeInternal.getMetaDataDiagnosticContext(request.requestContext.cosmosDiagnostics), + collection.getResourceId(), + null, + null); + return valueHolderMono .delayUntil(ignored -> queryValidationMono) .flatMapMany(routingMapHolder -> { @@ -4482,90 +4512,335 @@ private Flux> readManyByPartitionKeys( return Flux.error(new IllegalStateException("Failed to get routing map.")); } - Map> partitionRangePkMap = - groupPartitionKeysByPhysicalPartition(partitionKeys, pkDefinition, routingMap); + return buildSequentialFluxFromScratch( + partitionKeys, customQuery, pkDefinition, routingMap, + resourceLink, state, diagnosticsFactory, klass, + collectionRid, queryHash); + }); + }); + } - List partitionKeySelectors = PartitionKeyQueryHelper.createPkSelectors(pkDefinition); + /** + * Builds the sequential Flux of FeedResponse pages for readManyByPartitionKeys when starting + * from scratch (no continuation token). Groups PKs by physical partition, creates batches, + * sorts by EPK, and executes sequentially. + */ + private Flux> buildSequentialFluxFromScratch( + List partitionKeys, + SqlQuerySpec customQuery, + PartitionKeyDefinition pkDefinition, + CollectionRoutingMap routingMap, + String resourceLink, + QueryFeedOperationState state, + ScopedDiagnosticsFactory diagnosticsFactory, + Class klass, + String collectionRid, + int queryHash) { - String baseQueryText; - List baseParameters; - if (customQuery != null) { - baseQueryText = customQuery.getQueryText(); - baseParameters = customQuery.getParameters() != null - ? new ArrayList<>(customQuery.getParameters()) - : new ArrayList<>(); - } else { - baseQueryText = "SELECT * FROM c"; - baseParameters = new ArrayList<>(); - } + Map> partitionRangePkMap = + groupPartitionKeysByPhysicalPartition(partitionKeys, pkDefinition, routingMap); - // Build per-physical-partition batched queries. - // Each physical partition may have many PKs - split into batches - // to avoid oversized SQL queries. Batch size is configurable via - // system property COSMOS.READ_MANY_BY_PK_MAX_BATCH_SIZE (default 100). - int maxPksPerPartitionQuery = Configs.getReadManyByPkMaxBatchSize(); - - // Build batches per partition as a list of lists (one inner list per partition). - // Then interleave in round-robin order so that concurrent execution - // prefers different physical partitions over multiple batches of the same partition. - List>> batchesPerPartition = new ArrayList<>(); - int maxBatchesPerPartition = 0; - - for (Map.Entry> entry : partitionRangePkMap.entrySet()) { - List allPks = entry.getValue(); - List> partitionBatches = new ArrayList<>(); - for (int i = 0; i < allPks.size(); i += maxPksPerPartitionQuery) { - List batch = allPks.subList( - i, Math.min(i + maxPksPerPartitionQuery, allPks.size())); - SqlQuerySpec querySpec = ReadManyByPartitionKeyQueryHelper - .createReadManyByPkQuerySpec( - baseQueryText, baseParameters, batch, - partitionKeySelectors, pkDefinition); - partitionBatches.add( - Collections.singletonMap(entry.getKey(), querySpec)); - } - batchesPerPartition.add(partitionBatches); - maxBatchesPerPartition = Math.max(maxBatchesPerPartition, partitionBatches.size()); - } + List partitionKeySelectors = PartitionKeyQueryHelper.createPkSelectors(pkDefinition); - if (batchesPerPartition.isEmpty()) { - return Flux.empty(); - } + String baseQueryText; + List baseParameters; + if (customQuery != null) { + baseQueryText = customQuery.getQueryText(); + baseParameters = customQuery.getParameters() != null + ? new ArrayList<>(customQuery.getParameters()) + : new ArrayList<>(); + } else { + baseQueryText = "SELECT * FROM c"; + baseParameters = new ArrayList<>(); + } - // Interleave batches across physical partitions so that the first batch for - // each partition is kicked off before the second batch of any partition. With bounded - // concurrency this spreads the initial wave of work across the cluster; note that - // skewed distributions (one partition with N batches, another with 1) will eventually - // fall back to sequential execution once the short partitions are drained. - List> interleavedBatches = new ArrayList<>(); - for (int batchIdx = 0; batchIdx < maxBatchesPerPartition; batchIdx++) { - for (List> partitionBatches : batchesPerPartition) { - if (batchIdx < partitionBatches.size()) { - interleavedBatches.add(partitionBatches.get(batchIdx)); - } - } - } + int maxPksPerPartitionQuery = Configs.getReadManyByPkMaxBatchSize(); - // Execute all batches with bounded concurrency. - List>> queryFluxes = interleavedBatches - .stream() - .map(batchMap -> queryForReadMany( - diagnosticsFactory, - resourceLink, - new SqlQuerySpec(DUMMY_SQL_QUERY), - state.getQueryOptions(), - klass, - ResourceType.Document, - collection, - Collections.unmodifiableMap(batchMap))) - .collect(Collectors.toList()); + List allBatches = new ArrayList<>(); - int fluxConcurrency = Math.min(queryFluxes.size(), - Math.max(Configs.getCPUCnt(), 1)); + for (Map.Entry> entry : partitionRangePkMap.entrySet()) { + PartitionKeyRange pkRange = entry.getKey(); + Range partitionScope = pkRange.toRange(); + List allPks = entry.getValue(); - return Flux.merge(Flux.fromIterable(queryFluxes), fluxConcurrency, 1); - }); + // Sort PKs by EPK within each partition so that sub-batching is deterministic + // and the batchFilter range can be computed from first/last EPK in the batch. + allPks.sort(Comparator.comparing(pk -> getEffectivePartitionKeyString(pk, pkDefinition))); + + for (int i = 0; i < allPks.size(); i += maxPksPerPartitionQuery) { + int batchEnd = Math.min(i + maxPksPerPartitionQuery, allPks.size()); + List batch = allPks.subList(i, batchEnd); + + // batchFilter is [epk(first), maxExclusive) where maxExclusive is: + // - epk of the next PK after this batch (first PK of the next batch), or + // - partitionScope.getMax() if this is the last batch in the partition + String batchMinInclusive = getEffectivePartitionKeyString(batch.get(0), pkDefinition); + String batchMaxExclusive = batchEnd < allPks.size() + ? getEffectivePartitionKeyString(allPks.get(batchEnd), pkDefinition) + : partitionScope.getMax(); + Range batchFilter = new Range<>(batchMinInclusive, batchMaxExclusive, true, false); + + SqlQuerySpec querySpec = ReadManyByPartitionKeyQueryHelper + .createReadManyByPkQuerySpec( + baseQueryText, baseParameters, batch, + partitionKeySelectors, pkDefinition); + allBatches.add(new BatchDescriptor(partitionScope, batchFilter, querySpec)); + } + } + + if (allBatches.isEmpty()) { + return Flux.empty(); + } + + // Sort batches by batchFilter minInclusive EPK for deterministic sequential processing + allBatches.sort(Comparator.comparing(bd -> bd.batchFilter.getMin())); + + return buildSequentialBatchFlux( + allBatches, 0, null, + resourceLink, state, diagnosticsFactory, klass, + collectionRid, queryHash); + } + + /** + * Builds the sequential Flux of FeedResponse pages for readManyByPartitionKeys when resuming + * from a continuation token. Reconstructs batches from the token's batch definitions + * without needing the routing map. + */ + private Flux> buildSequentialFluxFromContinuation( + ReadManyByPartitionKeyContinuationToken parsedContinuation, + List partitionKeys, + SqlQuerySpec customQuery, + PartitionKeyDefinition pkDefinition, + String resourceLink, + QueryFeedOperationState state, + ScopedDiagnosticsFactory diagnosticsFactory, + Class klass, + String collectionRid, + int queryHash) { + + List partitionKeySelectors = PartitionKeyQueryHelper.createPkSelectors(pkDefinition); + + String baseQueryText; + List baseParameters; + if (customQuery != null) { + baseQueryText = customQuery.getQueryText(); + baseParameters = customQuery.getParameters() != null + ? new ArrayList<>(customQuery.getParameters()) + : new ArrayList<>(); + } else { + baseQueryText = "SELECT * FROM c"; + baseParameters = new ArrayList<>(); + } + + // Reconstruct batches from the continuation token's batch definitions. + // The current batch + remaining batches define the work left to do. + ReadManyByPartitionKeyContinuationToken.BatchDefinition currentBatchDef = + parsedContinuation.getCurrentBatch(); + List remainingBatchDefs = + parsedContinuation.getRemainingBatches(); + String initialBackendContinuation = parsedContinuation.getBackendContinuation(); + + // Build ordered list: current batch first, then remaining + List allBatchDefs = new ArrayList<>(); + allBatchDefs.add(currentBatchDef); + allBatchDefs.addAll(remainingBatchDefs); + + // For each batch definition, filter PKs that belong to it by batchFilter EPK range + // (not partitionScope — which covers the entire physical partition). + // Use partitionScope for FeedRange routing; batchFilter for PK filtering. + List allBatches = new ArrayList<>(); + + for (ReadManyByPartitionKeyContinuationToken.BatchDefinition batchDef : allBatchDefs) { + Range partitionScope = batchDef.getPartitionScope(); + Range batchFilter = batchDef.getBatchFilter(); + + // Filter PKs whose EPK falls within this batch's batchFilter range + List batchPks = filterPartitionKeysByEpkRange( + partitionKeys, pkDefinition, batchFilter); + + if (batchPks.isEmpty()) { + continue; + } + + SqlQuerySpec querySpec = ReadManyByPartitionKeyQueryHelper + .createReadManyByPkQuerySpec( + baseQueryText, baseParameters, batchPks, + partitionKeySelectors, pkDefinition); + + allBatches.add(new BatchDescriptor(partitionScope, batchFilter, querySpec)); + } + + if (allBatches.isEmpty()) { + return Flux.empty(); + } + + return buildSequentialBatchFlux( + allBatches, 0, initialBackendContinuation, + resourceLink, state, diagnosticsFactory, klass, + collectionRid, queryHash); + } + + /** + * Filters partition keys to those whose EPK falls within the given range. + */ + private List filterPartitionKeysByEpkRange( + List partitionKeys, + PartitionKeyDefinition pkDefinition, + Range epkRange) { + + List result = new ArrayList<>(); + for (PartitionKey pk : partitionKeys) { + String epk = getEffectivePartitionKeyString(pk, pkDefinition); + + // Check if EPK falls within [min, max) range + if (epk.compareTo(epkRange.getMin()) >= 0 && epk.compareTo(epkRange.getMax()) < 0) { + result.add(pk); + } + } + return result; + } + + /** + * Builds the sequential Flux that executes batches one at a time, stamping each + * FeedResponse with the composite continuation token. + */ + private Flux> buildSequentialBatchFlux( + List allBatches, + int startBatchIndex, + String initialBackendContinuation, + String resourceLink, + QueryFeedOperationState state, + ScopedDiagnosticsFactory diagnosticsFactory, + Class klass, + String collectionRid, + int queryHash) { + + List>> sequentialFluxes = new ArrayList<>(); + for (int i = startBatchIndex; i < allBatches.size(); i++) { + final int batchIndex = i; + final BatchDescriptor bd = allBatches.get(i); + final String backendContinuation = (i == startBatchIndex) ? initialBackendContinuation : null; + + // Remaining batch definitions for the continuation token (batches after the current one) + final List remainingAfterThis = new ArrayList<>(); + for (int j = batchIndex + 1; j < allBatches.size(); j++) { + BatchDescriptor remaining = allBatches.get(j); + remainingAfterThis.add( + new ReadManyByPartitionKeyContinuationToken.BatchDefinition( + remaining.partitionScope, remaining.batchFilter)); + } + + // Clone query options for this batch + CosmosQueryRequestOptions batchQueryOptions = queryOptionsAccessor() + .clone(state.getQueryOptions()); + queryOptionsAccessor().disallowQueryPlanRetrieval(batchQueryOptions); + + // Set FeedRange for routing via partitionScope EPK — handles splits transparently + batchQueryOptions.setFeedRange(new FeedRangeEpkImpl(bd.partitionScope)); + + if (backendContinuation != null) { + ModelBridgeInternal.setQueryRequestOptionsContinuationToken( + batchQueryOptions, backendContinuation); + } + + // Execute the batch query using the standard query API with FeedRange for routing. + // FeedRangeEpkImpl(partitionScope) handles routing to the correct physical partition(s), + // including transparently handling partition splits. + Flux> batchFlux = createQueryInternal( + diagnosticsFactory, + resourceLink, + bd.querySpec, + batchQueryOptions, + klass, + ResourceType.Document, + documentQueryClientImpl(RxDocumentClientImpl.this, getOperationContextAndListenerTuple(batchQueryOptions)), + UUIDs.nonBlockingRandomUUID(), + new AtomicBoolean(false)); + + // Current batch as a BatchDefinition for the continuation token + final ReadManyByPartitionKeyContinuationToken.BatchDefinition currentBatchDef = + new ReadManyByPartitionKeyContinuationToken.BatchDefinition( + bd.partitionScope, bd.batchFilter); + + // Stamp each FeedResponse with the composite continuation token + Flux> stampedFlux = batchFlux.map(feedResponse -> { + String backendCont = feedResponse.getContinuationToken(); + boolean isLastBatch = remainingAfterThis.isEmpty(); + boolean batchExhausted = (backendCont == null); + + if (isLastBatch && batchExhausted) { + ModelBridgeInternal.setFeedResponseContinuationToken(null, feedResponse); + } else if (batchExhausted) { + ReadManyByPartitionKeyContinuationToken.BatchDefinition nextBatchDef = + remainingAfterThis.get(0); + List remaining = + remainingAfterThis.size() > 1 + ? remainingAfterThis.subList(1, remainingAfterThis.size()) + : Collections.emptyList(); + ReadManyByPartitionKeyContinuationToken compositeContinuation = + new ReadManyByPartitionKeyContinuationToken( + remaining, nextBatchDef, null, collectionRid, queryHash); + ModelBridgeInternal.setFeedResponseContinuationToken( + compositeContinuation.serialize(), feedResponse); + } else { + ReadManyByPartitionKeyContinuationToken compositeContinuation = + new ReadManyByPartitionKeyContinuationToken( + remainingAfterThis, currentBatchDef, backendCont, + collectionRid, queryHash); + ModelBridgeInternal.setFeedResponseContinuationToken( + compositeContinuation.serialize(), feedResponse); + } + return feedResponse; }); + + sequentialFluxes.add(stampedFlux); + } + + return Flux.concat(sequentialFluxes); + } + + /** + * Returns the effective partition key string for a PartitionKey, handling PartitionKey.NONE + * by mapping to UndefinedPartitionKey. + */ + private static String getEffectivePartitionKeyString( + PartitionKey pk, + PartitionKeyDefinition pkDefinition) { + + PartitionKeyInternal pkInternal = BridgeInternal.getPartitionKeyInternal(pk); + PartitionKeyInternal effectivePkInternal = pkInternal.getComponents() == null + ? PartitionKeyInternal.UndefinedPartitionKey + : pkInternal; + + return PartitionKeyInternalHelper + .getEffectivePartitionKeyString(effectivePkInternal, pkDefinition); + } + + /** + * Descriptor for a single batch in readManyByPartitionKeys. + *

+ * Each batch carries two EPK ranges: + *

    + *
  • {@code partitionScope} — the physical partition's full EPK range, used for + * FeedRange-based routing (set on CosmosQueryRequestOptions). Survives splits + * because FeedRangeEpkImpl transparently fans out.
    • + *
  • {@code batchFilter} — the EPK sub-range covering only the PKs in this batch. + * When a physical partition has more PKs than maxBatchSize, multiple batches share + * the same partitionScope but have distinct batchFilter ranges. Used to + * reconstruct the correct PK set per batch when resuming from a continuation + * token.
    • + *
+ */ + private static final class BatchDescriptor { + final Range partitionScope; + final Range batchFilter; + final SqlQuerySpec querySpec; + + BatchDescriptor(Range partitionScope, Range batchFilter, SqlQuerySpec querySpec) { + this.partitionScope = partitionScope; + this.batchFilter = batchFilter; + this.querySpec = querySpec; + } } private Mono validateCustomQueryForReadManyByPartitionKey( diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/CosmosReadManyRequestOptions.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/CosmosReadManyRequestOptions.java index f6e570258042..e428a676e151 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/CosmosReadManyRequestOptions.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/CosmosReadManyRequestOptions.java @@ -380,6 +380,11 @@ static void initialize() { public CosmosQueryRequestOptionsBase getImpl(CosmosReadManyRequestOptions options) { return options.actualRequestOptions; } + + @Override + public String getRequestContinuation(CosmosReadManyRequestOptions options) { + return options.actualRequestOptions.getRequestContinuation(); + } }); } diff --git a/sdk/cosmos/docs/readManyByPartitionKey-design.md b/sdk/cosmos/docs/readManyByPartitionKey-design.md index 7f23d63011dd..2e169a1b1ed5 100644 --- a/sdk/cosmos/docs/readManyByPartitionKey-design.md +++ b/sdk/cosmos/docs/readManyByPartitionKey-design.md @@ -12,19 +12,19 @@ and additional filters. The SDK appends the auto-generated PK WHERE clause to it ## Decisions -| Topic | Decision | -|---|---| -| API name | `readManyByPartitionKeys` — distinct name to avoid ambiguity with existing `readMany(List)` | -| Return type | `CosmosPagedFlux` (async) / `CosmosPagedIterable` (sync) | -| Custom query format | `SqlQuerySpec` — full query with parameters; SDK ANDs the PK filter | -| Partial HPK | Supported from the start; prefix PKs fan out via `getOverlappingRanges` | -| PK deduplication | Done at Spark layer only, not in the SDK | -| Spark UDF | New `GetCosmosPartitionKeyValue` UDF | +| Topic | Decision | +|-------------------------|----------------------------------------------------------------------------------------------------------------------------------------------------------------------------| +| API name | `readManyByPartitionKeys` — distinct name to avoid ambiguity with existing `readMany(List)` | +| Return type | `CosmosPagedFlux` (async) / `CosmosPagedIterable` (sync) | +| Custom query format | `SqlQuerySpec` — full query with parameters; SDK ANDs the PK filter | +| Partial HPK | Supported from the start; prefix PKs fan out via `getOverlappingRanges` | +| PK deduplication | Done at Spark layer only, not in the SDK | +| Spark UDF | New `GetCosmosPartitionKeyValue` UDF | | Custom query validation | Gateway query plan via the standard SDK query-plan retrieval path; reject aggregates/ORDER BY/DISTINCT/GROUP BY/DCount/OFFSET/LIMIT/non-streaming ORDER BY/vector/fulltext | -| PK list size | No hard upper-bound enforced; SDK batches internally per physical partition (default 100 PKs per batch, configurable via `COSMOS.READ_MANY_BY_PK_MAX_BATCH_SIZE`) | -| Eager validation | Null and empty PK list rejected eagerly (not lazily in reactive chain) | -| Telemetry | Separate span name `readManyByPartitionKeyItems.` (distinct from existing `readManyItems`) | -| Query construction | Table alias auto-detected from FROM clause; string literals and subqueries handled correctly | +| PK list size | No hard upper-bound enforced; SDK batches internally per physical partition (default 100 PKs per batch, configurable via `COSMOS.READ_MANY_BY_PK_MAX_BATCH_SIZE`) | +| Eager validation | Null and empty PK list rejected eagerly (not lazily in reactive chain) | +| Telemetry | Separate span name `readManyByPartitionKeyItems.` (distinct from existing `readManyItems`) | +| Query construction | Table alias auto-detected from FROM clause; string literals and subqueries handled correctly | ## Phase 1 — SDK Core (`azure-cosmos`) @@ -85,29 +85,34 @@ One-time call per invocation using the same query-plan retrieval path and cachea ### Step 5: Query construction Query construction is implemented in `ReadManyByPartitionKeyQueryHelper`. The helper: + - Extracts the table alias from the FROM clause (handles `FROM c`, `FROM root r`, `FROM x WHERE ...`) - Handles string literals in queries (parens/keywords inside `'...'` are correctly skipped) - Recognizes SQL keywords: WHERE, ORDER, GROUP, JOIN, OFFSET, LIMIT, HAVING - Uses parameterized queries (`@__rmPk_` prefix) to prevent SQL injection **Single PK (HASH):** + ```sql {baseQuery} WHERE {alias}["{pkPath}"] IN (@__rmPk_0, @__rmPk_1, @__rmPk_2) ``` **Full HPK (MULTI_HASH):** + ```sql {baseQuery} WHERE ({alias}["{path1}"] = @__rmPk_0 AND {alias}["{path2}"] = @__rmPk_1) OR ({alias}["{path1}"] = @__rmPk_2 AND {alias}["{path2}"] = @__rmPk_3) ``` **Partial HPK (prefix-only):** + ```sql {baseQuery} WHERE ({alias}["{path1}"] = @__rmPk_0) OR ({alias}["{path1}"] = @__rmPk_1) ``` If the base query already has a WHERE clause: + ```sql {selectAndFrom} WHERE ({existingWhere}) AND ({pkFilter}) ``` @@ -131,6 +136,7 @@ New configurable batch size via system property `COSMOS.READ_MANY_BY_PK_MAX_BATC ### Step 9: PK-only serialization helper `CosmosPartitionKeyHelper`: + - `getCosmosPartitionKeyValueString(pkValues: List[Object]): String` — serialize to `pk([...])` format. - `tryParsePartitionKey(serialized: String): Option[PartitionKey]` — deserialize; returns `None` for malformed input including invalid JSON (wrapped in `scala.util.Try`). - When `spark.cosmos.read.readManyByPk.nullHandling=None` is used, hierarchical partition keys with null components are rejected with a clear error because `PartitionKey.NONE` cannot be used with multiple paths. @@ -138,6 +144,7 @@ New configurable batch size via system property `COSMOS.READ_MANY_BY_PK_MAX_BATC ### Step 10: `CosmosItemsDataSource.readManyByPartitionKeys` Static entry points that accept a DataFrame and Cosmos config. PK extraction supports two modes: + 1. **UDF-produced column**: DataFrame contains `_partitionKeyIdentity` column (from `GetCosmosPartitionKeyValue` UDF). 2. **Schema-matched columns**: DataFrame columns match the container's PK paths. @@ -152,6 +159,7 @@ Orchestrator that resolves schema, initializes and broadcasts client state to ex ### Step 12: `ItemsPartitionReaderWithReadManyByPartitionKey` Spark `PartitionReader[InternalRow]` that: + - Deduplicates PKs via `LinkedHashMap` (by PK string representation). - Passes the pre-built `CosmosReadManyRequestOptions` (with throughput control, diagnostics, custom serializer) to the SDK. - Uses `TransientIOErrorsRetryingIterator` for retry handling. @@ -160,6 +168,7 @@ Spark `PartitionReader[InternalRow]` that: ## Phase 3 — Testing ### Unit tests + - Query construction: single PK, HPK full/partial, custom query composition, table alias detection. - Query plan rejection: aggregates, ORDER BY, DISTINCT, GROUP BY (with and without aggregates), DCOUNT. - String literal handling: WHERE/parentheses inside string constants. @@ -168,9 +177,110 @@ Spark `PartitionReader[InternalRow]` that: - `findTopLevelWhereIndex` edge cases: subqueries, string literals, case insensitivity. ### Integration tests + - End-to-end SDK: single PK basic, projections, filters, empty results, HPK full/partial, request options propagation. - Batch size validation: temporarily lowered batch size to exercise batching/interleaving logic. - Null/empty PK list rejection (eager validation). - Spark connector: `ItemsPartitionReaderWithReadManyByPartitionKey` with known PK values and non-existent PKs. - Spark public API: nested partition key containers require `_partitionKeyIdentity` and succeed when populated via `GetCosmosPartitionKeyValue`. - `CosmosPartitionKeyHelper`: single/HPK roundtrip, case insensitivity, malformed input. + +## Phase 4 — Continuation Token Support for Retry Safety + +### Motivation + +The Spark `TransientIOErrorsRetryingReadManyByPartitionKeyIterator` replays from a +continuation token on transient I/O failures. For this to work, `readManyByPartitionKeys` +must (a) emit a meaningful continuation token on each `FeedResponse`, and (b) accept a +continuation token that resumes where the previous attempt left off. + +Because the SDK batches PK values across multiple physical partitions, the continuation +token must capture not only the backend query continuation within a single batch but also +which batches remain. Additionally, batches must be processed sequentially so that a +continuation token unambiguously identifies a position in the result stream. + +### Composite continuation token + +A new internal class `ReadManyByPartitionKeyContinuationToken` captures the resume state: + +| Field | Type | Description | +|-----------------------|-----------------------|----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------| +| `remainingBatches` | `List>` | EPK ranges (minInclusive, maxExclusive) of batches not yet started. Uses EPK ranges — **not** PkRangeIds — so the token survives partition splits. Reuses the existing `Range` type. | +| `currentBatch` | `Range` | EPK range of the batch currently being processed. | +| `backendContinuation` | `String` (nullable) | Backend query continuation within the current batch. `null` means "start from the beginning of this batch". | + +EPK ranges are represented as `Range` (the existing SDK type used throughout the +routing layer) with `isMinInclusive = true` and `isMaxExclusive = false`. + +**Serialization:** JSON → Base64, keeping the token opaque to callers. Example decoded: + +```json +{ + "remainingBatches": [ + {"min": "05C1E0", "max": "0BF333"}, + {"min": "0BF333", "max": "FF"} + ], + "currentBatch": {"min": "", "max": "05C1E0"}, + "backendContinuation": "eyJDb21wb3NpdGVUb2tlbi..." +} +``` + +### Sequential batch execution + +Batches are sorted by `minInclusive` EPK (lexicographic) and processed one at a time +via `Flux.concat` instead of `Flux.merge`. Each `FeedResponse` carries the composite +continuation token reflecting the current position: + +```text +FeedResponse 1: remaining=[B2, B3], current=B1, backend=null +FeedResponse 2: remaining=[B2, B3], current=B1, backend= + (B1 exhausted) +FeedResponse 3: remaining=[B3], current=B2, backend=null +FeedResponse 4: remaining=[B3], current=B2, backend= + (B2 exhausted) +FeedResponse 5: remaining=[], current=B3, backend=null +FeedResponse 6: remaining=[], current=B3, backend= + (done — final FeedResponse has null continuation token) +``` + +Sequential execution is always used — batches are never interleaved concurrently. +This ensures every `FeedResponse` carries a usable continuation token regardless of +whether the caller plans to resume or not. + +### API changes + +**`CosmosReadManyRequestOptionsImpl`:** Add `requestContinuation` field with +getter/setter. The impl already extends `CosmosQueryRequestOptionsBase`. + +**`RxDocumentClientImpl.readManyByPartitionKeys`:** + +- **With continuation token:** Deserialize the composite token. Restore batch list from + `currentBatch` + `remainingBatches` (no PK→routing-map computation needed). Pass + `backendContinuation` to the `CosmosQueryRequestOptions` for the first batch. Process + remaining batches in EPK order via `Flux.concat`. +- **Without continuation:** Compute batches as today, map each `PartitionKeyRange` to its + EPK range, sort by `minInclusive`, execute via `Flux.concat`, stamp each `FeedResponse` + with the composite token. + +**`CosmosAsyncContainer.readManyByPartitionKeyInternalFunc`:** Pass continuation token +from request options through to `RxDocumentClientImpl`. + +### Spark reader integration + +**`ItemsPartitionReaderWithReadManyByPartitionKey`:** The `fluxFactory` lambda applies +the continuation token: when non-null, set it on a cloned `CosmosReadManyRequestOptions` +before calling `readManyByPartitionKeys`. + +**`TransientIOErrorsRetryingReadManyByPartitionKeyIterator`:** No changes needed — the +existing retry loop already captures `lastContinuationToken` from each `FeedResponse` +and passes it to `fluxFactory` on retry. Once the factory applies the token, retry +works correctly. + +### EPK range stability across partition splits + +The token uses EPK ranges, not PkRangeIds. If a split occurs between retries: + +- The query engine handles EPK ranges spanning multiple new physical partitions transparently. +- A stale backend continuation is rejected by the backend; the SDK's + `StaleResourceRetryPolicy` retries with a fresh routing map while the EPK range ensures + correct targeting. From 5da1906642c976439523c7de7e621da701491031 Mon Sep 17 00:00:00 2001 From: Fabian Meiswinkel Date: Tue, 21 Apr 2026 23:11:42 +0000 Subject: [PATCH 39/64] Update RxDocumentClientImpl.java --- .../implementation/RxDocumentClientImpl.java | 23 ++++++++++--------- 1 file changed, 12 insertions(+), 11 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java index e1a4518c9805..278fb583e3d9 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java @@ -4596,7 +4596,7 @@ private Flux> buildSequentialFluxFromScratch( allBatches.sort(Comparator.comparing(bd -> bd.batchFilter.getMin())); return buildSequentialBatchFlux( - allBatches, 0, null, + allBatches, null, resourceLink, state, diagnosticsFactory, klass, collectionRid, queryHash); } @@ -4675,7 +4675,7 @@ private Flux> buildSequentialFluxFromContinuation( } return buildSequentialBatchFlux( - allBatches, 0, initialBackendContinuation, + allBatches, initialBackendContinuation, resourceLink, state, diagnosticsFactory, klass, collectionRid, queryHash); } @@ -4692,8 +4692,7 @@ private List filterPartitionKeysByEpkRange( for (PartitionKey pk : partitionKeys) { String epk = getEffectivePartitionKeyString(pk, pkDefinition); - // Check if EPK falls within [min, max) range - if (epk.compareTo(epkRange.getMin()) >= 0 && epk.compareTo(epkRange.getMax()) < 0) { + if (epkRange.contains(epk)) { result.add(pk); } } @@ -4706,7 +4705,6 @@ private List filterPartitionKeysByEpkRange( */ private Flux> buildSequentialBatchFlux( List allBatches, - int startBatchIndex, String initialBackendContinuation, String resourceLink, QueryFeedOperationState state, @@ -4716,10 +4714,10 @@ private Flux> buildSequentialBatchFlux( int queryHash) { List>> sequentialFluxes = new ArrayList<>(); - for (int i = startBatchIndex; i < allBatches.size(); i++) { + for (int i = 0; i < allBatches.size(); i++) { final int batchIndex = i; final BatchDescriptor bd = allBatches.get(i); - final String backendContinuation = (i == startBatchIndex) ? initialBackendContinuation : null; + final String backendContinuation = (i == 0) ? initialBackendContinuation : null; // Remaining batch definitions for the continuation token (batches after the current one) final List remainingAfterThis = new ArrayList<>(); @@ -4738,10 +4736,13 @@ private Flux> buildSequentialBatchFlux( // Set FeedRange for routing via partitionScope EPK — handles splits transparently batchQueryOptions.setFeedRange(new FeedRangeEpkImpl(bd.partitionScope)); - if (backendContinuation != null) { - ModelBridgeInternal.setQueryRequestOptionsContinuationToken( - batchQueryOptions, backendContinuation); - } + // Always set the continuation token on the cloned options — even when null. + // The cloned options may still carry the external caller's continuation token, + // which the backend wouldn't know how to interpret. For the first batch we pass + // the backend continuation from the composite token; for subsequent batches null + // (meaning "start from the beginning of this batch"). + ModelBridgeInternal.setQueryRequestOptionsContinuationToken( + batchQueryOptions, backendContinuation); // Execute the batch query using the standard query API with FeedRange for routing. // FeedRangeEpkImpl(partitionScope) handles routing to the correct physical partition(s), From 79d248733d25a1596b455ea821c040722c8b17bd Mon Sep 17 00:00:00 2001 From: Fabian Meiswinkel Date: Tue, 21 Apr 2026 23:49:58 +0000 Subject: [PATCH 40/64] Update readManyByPartitionKey-design.md --- sdk/cosmos/docs/readManyByPartitionKey-design.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/sdk/cosmos/docs/readManyByPartitionKey-design.md b/sdk/cosmos/docs/readManyByPartitionKey-design.md index 2e169a1b1ed5..9fa747c50ec3 100644 --- a/sdk/cosmos/docs/readManyByPartitionKey-design.md +++ b/sdk/cosmos/docs/readManyByPartitionKey-design.md @@ -221,7 +221,7 @@ routing layer) with `isMinInclusive = true` and `isMaxExclusive = false`. {"min": "0BF333", "max": "FF"} ], "currentBatch": {"min": "", "max": "05C1E0"}, - "backendContinuation": "eyJDb21wb3NpdGVUb2tlbi..." + "backendContinuation": "eyJDb21wb3NpdGVUb2" } ``` From f12582f6f49cfc89728d11301fdf6be20ff0e157 Mon Sep 17 00:00:00 2001 From: Fabian Meiswinkel Date: Wed, 22 Apr 2026 00:50:13 +0000 Subject: [PATCH 41/64] Fixed code review comments --- .../cosmos/spark/SparkE2EQueryITest.scala | 78 +++++++ .../cosmos/spark/CosmosItemsDataSource.scala | 26 ++- ...tionReaderWithReadManyByPartitionKey.scala | 36 +-- .../cosmos/ReadManyByPartitionKeyTest.java | 80 +++++++ ...nyByPartitionKeyContinuationTokenTest.java | 44 +++- .../azure/cosmos/CosmosAsyncContainer.java | 16 +- .../com/azure/cosmos/CosmosContainer.java | 12 +- ...adManyByPartitionKeyContinuationToken.java | 218 +++++++++++++++++- .../implementation/RxDocumentClientImpl.java | 198 ++++++++-------- .../docs/readManyByPartitionKey-design.md | 8 +- 10 files changed, 570 insertions(+), 146 deletions(-) diff --git a/sdk/cosmos/azure-cosmos-spark_3-5/src/test/scala/com/azure/cosmos/spark/SparkE2EQueryITest.scala b/sdk/cosmos/azure-cosmos-spark_3-5/src/test/scala/com/azure/cosmos/spark/SparkE2EQueryITest.scala index 79e52075c2cf..aa4fe2e3250d 100644 --- a/sdk/cosmos/azure-cosmos-spark_3-5/src/test/scala/com/azure/cosmos/spark/SparkE2EQueryITest.scala +++ b/sdk/cosmos/azure-cosmos-spark_3-5/src/test/scala/com/azure/cosmos/spark/SparkE2EQueryITest.scala @@ -185,5 +185,83 @@ class SparkE2EQueryITest } } + + "spark readManyByPartitionKeys" can "support partial top-level hierarchical partition keys from DataFrame columns without the UDF" in { + val cosmosEndpoint = TestConfigurations.HOST + val cosmosMasterKey = TestConfigurations.MASTER_KEY + val containerName = s"top-level-hpk-${UUID.randomUUID()}" + + val pkPaths = new ArrayList[String]() + pkPaths.add("/tenant") + pkPaths.add("/region") + pkPaths.add("/team") + + val pkDefinition = new PartitionKeyDefinition() + pkDefinition.setPaths(pkPaths) + pkDefinition.setKind(PartitionKind.MULTI_HASH) + pkDefinition.setVersion(PartitionKeyDefinitionVersion.V2) + + val containerProperties = new CosmosContainerProperties(containerName, pkDefinition) + cosmosClient + .getDatabase(cosmosDatabase) + .createContainerIfNotExists(containerProperties, ThroughputProperties.createManualThroughput(400)) + .block() + + try { + val container = cosmosClient.getDatabase(cosmosDatabase).getContainer(containerName) + val requestOptions = new CosmosItemRequestOptions() + + Seq( + ("tenantA", "east", "sales", "item-a1"), + ("tenantA", "west", "hr", "item-a2"), + ("tenantB", "east", "sales", "item-b1") + ).foreach { case (tenant, region, team, id) => + val item = objectMapper.createObjectNode() + item.put("id", id) + item.put("tenant", tenant) + item.put("region", region) + item.put("team", team) + item.put("payload", s"$tenant-$region-$team") + + val pk = new PartitionKeyBuilder().add(tenant).add(region).add(team).build() + container.createItem(item, pk, requestOptions).block() + } + + val cfg = Map( + "spark.cosmos.accountEndpoint" -> cosmosEndpoint, + "spark.cosmos.accountKey" -> cosmosMasterKey, + "spark.cosmos.database" -> cosmosDatabase, + "spark.cosmos.container" -> containerName, + "spark.cosmos.read.inferSchema.enabled" -> "true" + ) + + val sparkSession = spark + import sparkSession.implicits._ + + val tenantRows = CosmosItemsDataSource + .readManyByPartitionKeys(Seq("tenantA").toDF("tenant"), cfg.asJava) + .selectExpr("id", "tenant", "region", "team") + .collect() + + tenantRows should have size 2 + tenantRows.map(_.getAs[String]("id")).toSet shouldEqual Set("item-a1", "item-a2") + tenantRows.map(_.getAs[String]("tenant")).toSet shouldEqual Set("tenantA") + + val tenantRegionRows = CosmosItemsDataSource + .readManyByPartitionKeys(Seq(("tenantA", "east")).toDF("tenant", "region"), cfg.asJava) + .selectExpr("id", "tenant", "region", "team") + .collect() + + tenantRegionRows should have size 1 + tenantRegionRows.head.getAs[String]("id") shouldEqual "item-a1" + } finally { + cosmosClient + .getDatabase(cosmosDatabase) + .getContainer(containerName) + .delete() + .block() + } + } + // scalastyle:on multiple.string.literals } diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosItemsDataSource.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosItemsDataSource.scala index ac5ed255c008..b8b4a184c78e 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosItemsDataSource.scala +++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosItemsDataSource.scala @@ -166,18 +166,28 @@ object CosmosItemsDataSource { "'_partitionKeyIdentity' column produced by the GetCosmosPartitionKeyValue UDF.") } - // Check if ALL PK path columns exist in the DataFrame schema + // Allow DataFrames to provide a contiguous top-level prefix of the container's + // hierarchical partition key paths. For example: tenant, or tenant + region. val dfFieldNames = df.schema.fieldNames.toSet - val allPkColumnsPresent = pkPaths.forall(path => dfFieldNames.contains(path)) + val matchedPrefix = pkPaths.takeWhile(path => dfFieldNames.contains(path)) + val hasNonPrefixMatch = pkPaths.drop(matchedPrefix.size).exists(path => dfFieldNames.contains(path)) - if (allPkColumnsPresent && pkPaths.nonEmpty) { + if (hasNonPrefixMatch) { + throw new IllegalArgumentException( + "DataFrame columns matching the container's partition key paths must form a contiguous top-level prefix " + + "(for example: tenant, or tenant + region). " + + "For nested or non-prefix partition key extraction, add a '_partitionKeyIdentity' column produced " + + "by the GetCosmosPartitionKeyValue UDF.") + } + + if (matchedPrefix.nonEmpty) { Some((row: Row) => { - if (pkPaths.size == 1) { - buildPartitionKey(row.getAs[Any](pkPaths.head), treatNullAsNone) + if (matchedPrefix.size == 1) { + buildPartitionKey(row.getAs[Any](matchedPrefix.head), treatNullAsNone) } else { val builder = new PartitionKeyBuilder() - for (path <- pkPaths) { - addPartitionKeyComponent(builder, row.getAs[Any](path), treatNullAsNone, pkPaths.size) + for (path <- matchedPrefix) { + addPartitionKeyComponent(builder, row.getAs[Any](path), treatNullAsNone, matchedPrefix.size) } builder.build() } @@ -193,7 +203,7 @@ object CosmosItemsDataSource { throw new IllegalArgumentException( "Cannot determine partition key extraction from the input DataFrame. " + "Either add a '_partitionKeyIdentity' column (using the GetCosmosPartitionKeyValue UDF) " + - "or ensure the DataFrame contains columns matching the container's partition key paths.")) + "or ensure the DataFrame contains columns matching a top-level prefix of the container's partition key paths.")) readManyReader.readManyByPartitionKeys(df.rdd, pkExtraction, readerState) } diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/ItemsPartitionReaderWithReadManyByPartitionKey.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/ItemsPartitionReaderWithReadManyByPartitionKey.scala index e1158dddffaa..cfe5757716fd 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/ItemsPartitionReaderWithReadManyByPartitionKey.scala +++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/ItemsPartitionReaderWithReadManyByPartitionKey.scala @@ -6,7 +6,6 @@ package com.azure.cosmos.spark import com.azure.cosmos.{CosmosAsyncContainer, CosmosEndToEndOperationLatencyPolicyConfigBuilder, CosmosItemSerializerNoExceptionWrapping, SparkBridgeInternal} import com.azure.cosmos.implementation.spark.OperationContextAndListenerTuple import com.azure.cosmos.implementation.{ImplementationBridgeHelpers, ObjectNodeMap, SparkRowItem, Utils} -import com.azure.cosmos.BridgeInternal import com.azure.cosmos.models.{CosmosReadManyRequestOptions, ModelBridgeInternal, PartitionKey, PartitionKeyDefinition, SqlQuerySpec} import com.azure.cosmos.spark.BulkWriter.getThreadInfo import com.azure.cosmos.spark.CosmosTableSchemaInferrer.IdAttributeName @@ -157,18 +156,13 @@ private[spark] case class ItemsPartitionReaderWithReadManyByPartitionKey } ) - // Collect all PK values upfront - readManyByPartitionKeys needs the full list to - // group by physical partition (the SDK batches internally per physical partition). - // Deduplicate using the canonical PartitionKeyInternal JSON representation so that - // equivalent PKs built from different runtime types (Int vs Long vs Double) are - // collapsed, and distinct PKs that happen to toString() identically are not. + // Collect all PK values upfront - the SDK now owns normalization and deduplication, + // so Spark preserves the caller's input as-is and relies on the SDK's set-based semantics. + // Callers should still dedupe their DataFrame input when practical to avoid extra work. private lazy val pkList = { - val seen = new java.util.LinkedHashMap[String, PartitionKey]() - readManyPartitionKeys.foreach(pk => { - val key = BridgeInternal.getPartitionKeyInternal(pk).toJson - seen.putIfAbsent(key, pk) - }) - new java.util.ArrayList[PartitionKey](seen.values()) + val values = new java.util.ArrayList[PartitionKey]() + readManyPartitionKeys.foreach(values.add) + values } private val endToEndTimeoutPolicy = @@ -216,22 +210,14 @@ private[spark] case class ItemsPartitionReaderWithReadManyByPartitionKey // On the first call continuationToken is null (start from scratch); on retry // it is the continuation token from the last fully-drained page. private val fluxFactory: String => CosmosPagedFlux[SparkRowItem] = { (continuationToken: String) => - // Clone options and set continuation token if provided so the SDK - // resumes from where the previous attempt left off. - val effectiveOptions = new CosmosReadManyRequestOptions(readManyOptions) - if (continuationToken != null) { - val effectiveImpl = ImplementationBridgeHelpers - .CosmosReadManyRequestOptionsHelper - .getCosmosReadManyRequestOptionsAccessor - .getImpl(effectiveOptions) - .asInstanceOf[com.azure.cosmos.implementation.CosmosReadManyRequestOptionsImpl] - effectiveImpl.setRequestContinuation(continuationToken) - } + // Reuse the preconfigured request options and only update the continuation + // token so the SDK resumes from the last fully committed page. + readManyOptionsImpl.setRequestContinuation(continuationToken) customQueryOpt match { case Some(query) => - cosmosAsyncContainer.readManyByPartitionKeys(pkList, query, effectiveOptions, classOf[SparkRowItem]) + cosmosAsyncContainer.readManyByPartitionKeys(pkList, query, readManyOptions, classOf[SparkRowItem]) case None => - cosmosAsyncContainer.readManyByPartitionKeys(pkList, effectiveOptions, classOf[SparkRowItem]) + cosmosAsyncContainer.readManyByPartitionKeys(pkList, readManyOptions, classOf[SparkRowItem]) } } diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/ReadManyByPartitionKeyTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/ReadManyByPartitionKeyTest.java index 508725a23112..95b2c0d87607 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/ReadManyByPartitionKeyTest.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/ReadManyByPartitionKeyTest.java @@ -633,6 +633,86 @@ public void singlePk_continuationToken_resumesCorrectly() { } } + + @Test(groups = {"emulator"}, timeOut = TIMEOUT) + public void singlePk_continuationToken_resumesCorrectly_whenInputContainsDuplicatePartitionKeys() { + String originalValue = System.getProperty("COSMOS.READ_MANY_BY_PK_MAX_BATCH_SIZE"); + try { + System.setProperty("COSMOS.READ_MANY_BY_PK_MAX_BATCH_SIZE", "1"); + + createSinglePkItems("dupResumePk1", 2); + createSinglePkItems("dupResumePk2", 2); + createSinglePkItems("dupResumePk3", 2); + + List pkValues = Arrays.asList( + new PartitionKey("dupResumePk2"), + new PartitionKey("dupResumePk1"), + new PartitionKey("dupResumePk2"), + new PartitionKey("dupResumePk3"), + new PartitionKey("dupResumePk1")); + + CosmosAsyncContainer asyncContainer = client.asyncClient() + .getDatabase(preExistingDatabaseId) + .getContainer(singlePkContainer.getId()); + + List> allPages = asyncContainer + .readManyByPartitionKeys(pkValues, ObjectNode.class) + .byPage() + .collectList() + .block(); + + assertThat(allPages).isNotNull(); + assertThat(allPages.size()).isGreaterThan(1); + + String continuationAfterFirstPage = allPages.get(0).getContinuationToken(); + assertThat(continuationAfterFirstPage).isNotNull(); + List itemsFromFirstPage = allPages.get(0).getResults(); + + com.azure.cosmos.models.CosmosReadManyRequestOptions options2 = + new com.azure.cosmos.models.CosmosReadManyRequestOptions(); + ((CosmosReadManyRequestOptionsImpl) + ImplementationBridgeHelpers.CosmosReadManyRequestOptionsHelper + .getCosmosReadManyRequestOptionsAccessor() + .getImpl(options2)) + .setRequestContinuation(continuationAfterFirstPage); + + List> remainingPages = asyncContainer + .readManyByPartitionKeys(pkValues, options2, ObjectNode.class) + .byPage() + .collectList() + .block(); + + assertThat(remainingPages).isNotNull(); + + List firstPageIds = itemsFromFirstPage.stream() + .map(n -> n.get("id").asText()) + .collect(Collectors.toList()); + List remainingIds = remainingPages.stream() + .flatMap(p -> p.getResults().stream()) + .map(n -> n.get("id").asText()) + .collect(Collectors.toList()); + List allIds = allPages.stream() + .flatMap(p -> p.getResults().stream()) + .map(n -> n.get("id").asText()) + .collect(Collectors.toList()); + + List combined = new ArrayList<>(firstPageIds); + combined.addAll(remainingIds); + + assertThat(combined).hasSameElementsAs(allIds); + assertThat(combined).doesNotHaveDuplicates(); + assertThat(combined).hasSize(6); + + cleanupContainer(singlePkContainer); + } finally { + if (originalValue != null) { + System.setProperty("COSMOS.READ_MANY_BY_PK_MAX_BATCH_SIZE", originalValue); + } else { + System.clearProperty("COSMOS.READ_MANY_BY_PK_MAX_BATCH_SIZE"); + } + } + } + //endregion //region helper methods diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyContinuationTokenTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyContinuationTokenTest.java index f92c2c902e06..b7f2c7e34b31 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyContinuationTokenTest.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyContinuationTokenTest.java @@ -21,7 +21,8 @@ public class ReadManyByPartitionKeyContinuationTokenTest { private static final String TEST_COLLECTION_RID = "dbs/testDb/colls/testColl"; - private static final int TEST_QUERY_HASH = 12345; + private static final String TEST_QUERY_HASH = "12345"; + private static final String TEST_PARTITION_KEY_SET_HASH = "67890"; @Test(groups = { "unit" }) public void roundtrip_withBackendContinuation() { @@ -178,7 +179,7 @@ public void rangesPreserveMinMaxInclusive() { public void collectionRidAndQueryHash_roundtrip() { Range current = new Range<>("", "FF", true, false); String rid = "dbs/myDb/colls/myColl"; - int hash = 98765; + String hash = "98765"; ReadManyByPartitionKeyContinuationToken token = new ReadManyByPartitionKeyContinuationToken( @@ -192,9 +193,42 @@ public void collectionRidAndQueryHash_roundtrip() { assertThat(deserialized.getQueryHash()).isEqualTo(hash); } + @Test(groups = { "unit" }) + public void partitionKeySetHash_roundtrip() { + Range current = new Range<>("", "FF", true, false); + + ReadManyByPartitionKeyContinuationToken token = + new ReadManyByPartitionKeyContinuationToken( + Collections.emptyList(), current, null, TEST_COLLECTION_RID, TEST_QUERY_HASH, TEST_PARTITION_KEY_SET_HASH); + + String serialized = token.serialize(); + ReadManyByPartitionKeyContinuationToken deserialized = + ReadManyByPartitionKeyContinuationToken.deserialize(serialized); + + assertThat(deserialized.getPartitionKeySetHash()).isEqualTo(TEST_PARTITION_KEY_SET_HASH); + } + + @Test(groups = { "unit" }) + public void computePartitionKeySetHash_isStableAcrossDuplicateAndReorderedEpks() { + List epks1 = Arrays.asList("BB", "AA", "BB", "CC"); + List epks2 = Arrays.asList("CC", "AA", "BB"); + + assertThat(ReadManyByPartitionKeyContinuationToken.computePartitionKeySetHash(epks1)) + .isEqualTo(ReadManyByPartitionKeyContinuationToken.computePartitionKeySetHash(epks2)); + } + + + @Test(groups = { "unit" }) + public void computePartitionKeySetHash_returnsStableHexDigest() { + String hash = ReadManyByPartitionKeyContinuationToken.computePartitionKeySetHash( + Arrays.asList("BB", "AA", "BB", "CC")); + + assertThat(hash).matches("[0-9a-f]{32}"); + } + @Test(groups = { "unit" }) public void computeQueryHash_nullSpec_returnsZero() { - assertThat(ReadManyByPartitionKeyContinuationToken.computeQueryHash(null)).isEqualTo(0); + assertThat(ReadManyByPartitionKeyContinuationToken.computeQueryHash(null)).isEqualTo("0"); } @Test(groups = { "unit" }) @@ -231,8 +265,8 @@ public void computeQueryHash_differentQueryText_differentHash() { @Test(groups = { "unit" }) public void computeQueryHash_noParams_stableHash() { SqlQuerySpec spec = new SqlQuerySpec("SELECT * FROM c"); - int hash1 = ReadManyByPartitionKeyContinuationToken.computeQueryHash(spec); - int hash2 = ReadManyByPartitionKeyContinuationToken.computeQueryHash(spec); + String hash1 = ReadManyByPartitionKeyContinuationToken.computeQueryHash(spec); + String hash2 = ReadManyByPartitionKeyContinuationToken.computeQueryHash(spec); assertThat(hash1).isEqualTo(hash2); } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java index fbc3508d5be1..12458b458338 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java @@ -1607,7 +1607,8 @@ private Mono> readManyInternal( * Reads many documents matching the provided partition key values. * Unlike {@link #readMany(List, Class)} this method does not require item ids - it queries * all documents matching the provided partition key values. Uses {@code SELECT * FROM c} - * as the base query. + * as the base query. Duplicate partition key inputs are normalized with set-based semantics + * before batching, so repeated keys do not duplicate the results. * * @param the type parameter * @param partitionKeys list of partition key values to read documents for @@ -1625,7 +1626,8 @@ public CosmosPagedFlux readManyByPartitionKeys( * Reads many documents matching the provided partition key values. * Unlike {@link #readMany(List, Class)} this method does not require item ids - it queries * all documents matching the provided partition key values. Uses {@code SELECT * FROM c} - * as the base query. + * as the base query. Duplicate partition key inputs are normalized with set-based semantics + * before batching, so repeated keys do not duplicate the results. * * @param the type parameter * @param partitionKeys list of partition key values to read documents for @@ -1652,7 +1654,8 @@ public CosmosPagedFlux readManyByPartitionKeys( * rejected. *

* Partial hierarchical partition keys are supported and will fan out to multiple - * physical partitions. + * physical partitions. Duplicate partition key inputs are normalized with set-based semantics + * before batching. * * @param the type parameter * @param partitionKeys list of partition key values to read documents for @@ -1680,7 +1683,8 @@ public CosmosPagedFlux readManyByPartitionKeys( * rejected. *

* Partial hierarchical partition keys are supported and will fan out to multiple - * physical partitions. + * physical partitions. Duplicate partition key inputs are normalized with set-based semantics + * before batching. * * @param the type parameter * @param partitionKeys list of partition key values to read documents for @@ -1709,8 +1713,10 @@ public CosmosPagedFlux readManyByPartitionKeys( } } + List partitionKeysSnapshot = new ArrayList<>(partitionKeys); + return UtilBridgeInternal.createCosmosPagedFlux( - readManyByPartitionKeyInternalFunc(partitionKeys, customQuery, requestOptions, classType)); + readManyByPartitionKeyInternalFunc(partitionKeysSnapshot, customQuery, requestOptions, classType)); } private Function>> readManyByPartitionKeyInternalFunc( diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosContainer.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosContainer.java index 6acf391f3244..e61e4ee1f422 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosContainer.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosContainer.java @@ -544,7 +544,8 @@ public FeedResponse readMany( * Reads many documents matching the provided partition key values. * Unlike {@link #readMany(List, Class)} this method does not require item ids - it queries * all documents matching the provided partition key values. Uses {@code SELECT * FROM c} - * as the base query. + * as the base query. Duplicate partition key inputs are normalized with set-based semantics + * before batching, so repeated keys do not duplicate the results. * * @param the type parameter * @param partitionKeys list of partition key values to read documents for @@ -562,7 +563,8 @@ public CosmosPagedIterable readManyByPartitionKeys( * Reads many documents matching the provided partition key values. * Unlike {@link #readMany(List, Class)} this method does not require item ids - it queries * all documents matching the provided partition key values. Uses {@code SELECT * FROM c} - * as the base query. + * as the base query. Duplicate partition key inputs are normalized with set-based semantics + * before batching, so repeated keys do not duplicate the results. * * @param the type parameter * @param partitionKeys list of partition key values to read documents for @@ -589,7 +591,8 @@ public CosmosPagedIterable readManyByPartitionKeys( * rejected. *

* Partial hierarchical partition keys are supported and will fan out to multiple - * physical partitions. + * physical partitions. Duplicate partition key inputs are normalized with set-based semantics + * before batching. * * @param the type parameter * @param partitionKeys list of partition key values to read documents for @@ -617,7 +620,8 @@ public CosmosPagedIterable readManyByPartitionKeys( * rejected. *

* Partial hierarchical partition keys are supported and will fan out to multiple - * physical partitions. + * physical partitions. Duplicate partition key inputs are normalized with set-based semantics + * before batching. * * @param the type parameter * @param partitionKeys list of partition key values to read documents for diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyContinuationToken.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyContinuationToken.java index 5c40a22c63ad..b25faff9e934 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyContinuationToken.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyContinuationToken.java @@ -2,16 +2,20 @@ // Licensed under the MIT License. package com.azure.cosmos.implementation; +import com.azure.cosmos.implementation.routing.MurmurHash3_128; import com.azure.cosmos.implementation.routing.Range; +import com.azure.cosmos.implementation.routing.UInt128; import com.azure.cosmos.models.SqlParameter; import com.azure.cosmos.models.SqlQuerySpec; import com.fasterxml.jackson.annotation.JsonCreator; import com.fasterxml.jackson.annotation.JsonIgnore; import com.fasterxml.jackson.annotation.JsonProperty; +import java.io.ByteArrayOutputStream; import java.nio.charset.StandardCharsets; import java.util.ArrayList; import java.util.Base64; +import java.util.Collections; import java.util.List; import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkArgument; @@ -47,6 +51,7 @@ public final class ReadManyByPartitionKeyContinuationToken { private static final String BACKEND_CONTINUATION_PROPERTY = "bc"; private static final String COLLECTION_RID_PROPERTY = "cr"; private static final String QUERY_HASH_PROPERTY = "qh"; + private static final String PARTITION_KEY_SET_HASH_PROPERTY = "ph"; @JsonProperty(REMAINING_BATCHES_PROPERTY) private final List remainingBatches; @@ -61,7 +66,10 @@ public final class ReadManyByPartitionKeyContinuationToken { private final String collectionRid; @JsonProperty(QUERY_HASH_PROPERTY) - private final int queryHash; + private final String queryHash; + + @JsonProperty(PARTITION_KEY_SET_HASH_PROPERTY) + private final String partitionKeySetHash; @JsonCreator ReadManyByPartitionKeyContinuationToken( @@ -69,13 +77,15 @@ public final class ReadManyByPartitionKeyContinuationToken { @JsonProperty(CURRENT_BATCH_PROPERTY) BatchDefinitionDto currentBatch, @JsonProperty(BACKEND_CONTINUATION_PROPERTY) String backendContinuation, @JsonProperty(COLLECTION_RID_PROPERTY) String collectionRid, - @JsonProperty(QUERY_HASH_PROPERTY) int queryHash) { + @JsonProperty(QUERY_HASH_PROPERTY) String queryHash, + @JsonProperty(PARTITION_KEY_SET_HASH_PROPERTY) String partitionKeySetHash) { this.remainingBatches = remainingBatches; this.currentBatch = currentBatch; this.backendContinuation = backendContinuation; this.collectionRid = collectionRid; this.queryHash = queryHash; + this.partitionKeySetHash = partitionKeySetHash; } public ReadManyByPartitionKeyContinuationToken( @@ -85,6 +95,39 @@ public ReadManyByPartitionKeyContinuationToken( String collectionRid, int queryHash) { + this(remainingBatches, currentBatch, backendContinuation, collectionRid, String.valueOf(queryHash), "0"); + } + + public ReadManyByPartitionKeyContinuationToken( + List remainingBatches, + BatchDefinition currentBatch, + String backendContinuation, + String collectionRid, + int queryHash, + int partitionKeySetHash) { + + this(remainingBatches, currentBatch, backendContinuation, collectionRid, + String.valueOf(queryHash), String.valueOf(partitionKeySetHash)); + } + + public ReadManyByPartitionKeyContinuationToken( + List remainingBatches, + BatchDefinition currentBatch, + String backendContinuation, + String collectionRid, + String queryHash) { + + this(remainingBatches, currentBatch, backendContinuation, collectionRid, queryHash, "0"); + } + + public ReadManyByPartitionKeyContinuationToken( + List remainingBatches, + BatchDefinition currentBatch, + String backendContinuation, + String collectionRid, + String queryHash, + String partitionKeySetHash) { + checkNotNull(currentBatch, "Argument 'currentBatch' must not be null."); checkNotNull(remainingBatches, "Argument 'remainingBatches' must not be null."); @@ -96,6 +139,7 @@ public ReadManyByPartitionKeyContinuationToken( this.backendContinuation = backendContinuation; this.collectionRid = collectionRid; this.queryHash = queryHash; + this.partitionKeySetHash = partitionKeySetHash; } /** @@ -110,6 +154,39 @@ public ReadManyByPartitionKeyContinuationToken( String collectionRid, int queryHash) { + this(remainingBatchRanges, currentBatchRange, backendContinuation, collectionRid, String.valueOf(queryHash), "0"); + } + + public ReadManyByPartitionKeyContinuationToken( + List> remainingBatchRanges, + Range currentBatchRange, + String backendContinuation, + String collectionRid, + int queryHash, + int partitionKeySetHash) { + + this(remainingBatchRanges, currentBatchRange, backendContinuation, collectionRid, + String.valueOf(queryHash), String.valueOf(partitionKeySetHash)); + } + + public ReadManyByPartitionKeyContinuationToken( + List> remainingBatchRanges, + Range currentBatchRange, + String backendContinuation, + String collectionRid, + String queryHash) { + + this(remainingBatchRanges, currentBatchRange, backendContinuation, collectionRid, queryHash, "0"); + } + + public ReadManyByPartitionKeyContinuationToken( + List> remainingBatchRanges, + Range currentBatchRange, + String backendContinuation, + String collectionRid, + String queryHash, + String partitionKeySetHash) { + checkNotNull(currentBatchRange, "Argument 'currentBatchRange' must not be null."); checkNotNull(remainingBatchRanges, "Argument 'remainingBatchRanges' must not be null."); @@ -123,6 +200,7 @@ public ReadManyByPartitionKeyContinuationToken( this.backendContinuation = backendContinuation; this.collectionRid = collectionRid; this.queryHash = queryHash; + this.partitionKeySetHash = partitionKeySetHash; } @JsonIgnore @@ -147,16 +225,50 @@ public String getCollectionRid() { return collectionRid; } - public int getQueryHash() { + public String getQueryHash() { return queryHash; } + public String getPartitionKeySetHash() { + return partitionKeySetHash; + } + /** * Computes a stable hash for a SqlQuerySpec (or null for the default SELECT * FROM c query). * Hashes over both the query text and all parameter names/values to detect when a continuation * token is reused with a different query or different parameter values. */ - public static int computeQueryHash(SqlQuerySpec querySpec) { + public static String computeQueryHash(SqlQuerySpec querySpec) { + if (querySpec == null) { + return "0"; + } + + try { + ByteArrayOutputStream output = new ByteArrayOutputStream(); + updateHashInput(output, querySpec.getQueryText()); + + List params = querySpec.getParameters(); + if (params != null) { + for (SqlParameter param : params) { + updateHashInput(output, param.getName()); + + Object value = param.getValue(Object.class); + if (value == null) { + updateHashInput(output, null); + } else { + output.write(Utils.getSimpleObjectMapper().writeValueAsBytes(value)); + output.write(0); + } + } + } + + return murmurHash128Hex(output.toByteArray()); + } catch (Exception e) { + throw new IllegalStateException("Failed to compute stable query hash for continuation token.", e); + } + } + + static int computeLegacyQueryHash(SqlQuerySpec querySpec) { if (querySpec == null) { return 0; } @@ -174,6 +286,104 @@ public static int computeQueryHash(SqlQuerySpec querySpec) { return hash; } + /** + * Computes a stable hash for the normalized set of partition key EPK values. + * Duplicate and reordered inputs intentionally produce the same digest. + */ + public static String computePartitionKeySetHash(List partitionKeyEpks) { + if (partitionKeyEpks == null || partitionKeyEpks.isEmpty()) { + return "0"; + } + + List normalizedEpks = new ArrayList<>(partitionKeyEpks.size()); + for (String epk : partitionKeyEpks) { + if (epk != null) { + normalizedEpks.add(epk); + } + } + + if (normalizedEpks.isEmpty()) { + return "0"; + } + + Collections.sort(normalizedEpks); + + ByteArrayOutputStream output = new ByteArrayOutputStream(); + String previous = null; + for (String epk : normalizedEpks) { + if (epk.equals(previous)) { + continue; + } + + updateHashInput(output, epk); + previous = epk; + } + + return murmurHash128Hex(output.toByteArray()); + } + + static int computeLegacyPartitionKeySetHash(List partitionKeyEpks) { + if (partitionKeyEpks == null || partitionKeyEpks.isEmpty()) { + return 0; + } + + List normalizedEpks = new ArrayList<>(partitionKeyEpks.size()); + for (String epk : partitionKeyEpks) { + if (epk != null) { + normalizedEpks.add(epk); + } + } + + if (normalizedEpks.isEmpty()) { + return 0; + } + + Collections.sort(normalizedEpks); + + int hash = 17; + String previous = null; + for (String epk : normalizedEpks) { + if (epk.equals(previous)) { + continue; + } + + hash = 31 * hash + epk.hashCode(); + previous = epk; + } + + return hash; + } + + private static void updateHashInput(ByteArrayOutputStream output, String value) { + if (value != null) { + output.writeBytes(value.getBytes(StandardCharsets.UTF_8)); + } + output.write(0); + } + + private static String murmurHash128Hex(byte[] bytes) { + if (bytes == null || bytes.length == 0) { + return "0"; + } + + UInt128 hash = MurmurHash3_128.hash128(bytes, bytes.length); + return toFixedHex(hash.getHigh()) + toFixedHex(hash.getLow()); + } + + private static String toFixedHex(long value) { + String hex = Long.toHexString(value); + if (hex.length() >= 16) { + return hex; + } + + StringBuilder builder = new StringBuilder(16); + for (int i = hex.length(); i < 16; i++) { + builder.append('0'); + } + builder.append(hex); + return builder.toString(); + } + /** * Serializes this token to a Base64-encoded JSON string. */ diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java index 278fb583e3d9..3604ed885f9c 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java @@ -4462,7 +4462,13 @@ private Flux> readManyByPartitionKeys( final PartitionKeyDefinition pkDefinition = collection.getPartitionKey(); final String collectionRid = collection.getResourceId(); - final int queryHash = ReadManyByPartitionKeyContinuationToken.computeQueryHash(customQuery); + final String queryHash = ReadManyByPartitionKeyContinuationToken.computeQueryHash(customQuery); + final List normalizedPartitionKeys = + normalizePartitionKeys(partitionKeys, pkDefinition); + final String partitionKeySetHash = ReadManyByPartitionKeyContinuationToken.computePartitionKeySetHash( + normalizedPartitionKeys.stream() + .map(normalizedPk -> normalizedPk.effectivePartitionKeyString) + .collect(Collectors.toList())); // When a continuation token is present, skip the routing map lookup and batch // construction - the token already contains the batch definitions. Only resolve @@ -4476,16 +4482,30 @@ private Flux> readManyByPartitionKeys( "Continuation token was created for a different collection (rid mismatch). " + "Expected: " + collectionRid + ", token has: " + parsedContinuation.getCollectionRid())); } - if (queryHash != parsedContinuation.getQueryHash()) { + String legacyQueryHash = String.valueOf( + ReadManyByPartitionKeyContinuationToken.computeLegacyQueryHash(customQuery)); + if (!queryHash.equals(parsedContinuation.getQueryHash()) + && !legacyQueryHash.equals(parsedContinuation.getQueryHash())) { return Flux.error(new IllegalArgumentException( "Continuation token was created with a different query (hash mismatch). " + "The same query must be used when resuming from a continuation token.")); } + String legacyPartitionKeySetHash = String.valueOf( + ReadManyByPartitionKeyContinuationToken.computeLegacyPartitionKeySetHash( + normalizedPartitionKeys.stream() + .map(normalizedPk -> normalizedPk.effectivePartitionKeyString) + .collect(Collectors.toList()))); + if (!partitionKeySetHash.equals(parsedContinuation.getPartitionKeySetHash()) + && !legacyPartitionKeySetHash.equals(parsedContinuation.getPartitionKeySetHash())) { + return Flux.error(new IllegalArgumentException( + "Continuation token was created with a different partition-key set (hash mismatch). " + + "The same normalized set of partition key values must be used when resuming.")); + } return buildSequentialFluxFromContinuation( - parsedContinuation, partitionKeys, customQuery, pkDefinition, + parsedContinuation, normalizedPartitionKeys, customQuery, pkDefinition, resourceLink, state, diagnosticsFactory, klass, - collectionRid, queryHash); + collectionRid, queryHash, partitionKeySetHash); } // First-call path: validate custom query, resolve routing map, build batches @@ -4513,20 +4533,49 @@ private Flux> readManyByPartitionKeys( } return buildSequentialFluxFromScratch( - partitionKeys, customQuery, pkDefinition, routingMap, + normalizedPartitionKeys, customQuery, pkDefinition, routingMap, resourceLink, state, diagnosticsFactory, klass, - collectionRid, queryHash); + collectionRid, queryHash, partitionKeySetHash); }); }); } + /** + * Normalizes the input partition key list into a deterministic, set-based representation. + * Duplicates are removed by effective partition key string and the remaining keys are sorted + * lexicographically by EPK so batching and continuation-token hashes stay stable. + */ + private List normalizePartitionKeys( + List partitionKeys, + PartitionKeyDefinition pkDefinition) { + + Map normalizedByEpk = new LinkedHashMap<>(); + + for (PartitionKey pk : partitionKeys) { + PartitionKeyInternal pkInternal = BridgeInternal.getPartitionKeyInternal(pk); + PartitionKeyInternal effectivePkInternal = pkInternal.getComponents() == null + ? PartitionKeyInternal.UndefinedPartitionKey + : pkInternal; + String effectivePartitionKeyString = PartitionKeyInternalHelper + .getEffectivePartitionKeyString(effectivePkInternal, pkDefinition); + + normalizedByEpk.putIfAbsent( + effectivePartitionKeyString, + new NormalizedPartitionKey(pk, effectivePkInternal, effectivePartitionKeyString)); + } + + List normalizedPartitionKeys = new ArrayList<>(normalizedByEpk.values()); + normalizedPartitionKeys.sort(Comparator.comparing(normalizedPk -> normalizedPk.effectivePartitionKeyString)); + return normalizedPartitionKeys; + } + /** * Builds the sequential Flux of FeedResponse pages for readManyByPartitionKeys when starting * from scratch (no continuation token). Groups PKs by physical partition, creates batches, * sorts by EPK, and executes sequentially. */ private Flux> buildSequentialFluxFromScratch( - List partitionKeys, + List normalizedPartitionKeys, SqlQuerySpec customQuery, PartitionKeyDefinition pkDefinition, CollectionRoutingMap routingMap, @@ -4535,10 +4584,11 @@ private Flux> buildSequentialFluxFromScratch( ScopedDiagnosticsFactory diagnosticsFactory, Class klass, String collectionRid, - int queryHash) { + String queryHash, + String partitionKeySetHash) { - Map> partitionRangePkMap = - groupPartitionKeysByPhysicalPartition(partitionKeys, pkDefinition, routingMap); + Map> partitionRangePkMap = + groupPartitionKeysByPhysicalPartition(normalizedPartitionKeys, pkDefinition, routingMap); List partitionKeySelectors = PartitionKeyQueryHelper.createPkSelectors(pkDefinition); @@ -4558,25 +4608,26 @@ private Flux> buildSequentialFluxFromScratch( List allBatches = new ArrayList<>(); - for (Map.Entry> entry : partitionRangePkMap.entrySet()) { + for (Map.Entry> entry : partitionRangePkMap.entrySet()) { PartitionKeyRange pkRange = entry.getKey(); Range partitionScope = pkRange.toRange(); - List allPks = entry.getValue(); - - // Sort PKs by EPK within each partition so that sub-batching is deterministic - // and the batchFilter range can be computed from first/last EPK in the batch. - allPks.sort(Comparator.comparing(pk -> getEffectivePartitionKeyString(pk, pkDefinition))); + List allPks = entry.getValue(); + // The per-range list already preserves the globally normalized EPK order, + // so no additional in-partition sort is required here. for (int i = 0; i < allPks.size(); i += maxPksPerPartitionQuery) { int batchEnd = Math.min(i + maxPksPerPartitionQuery, allPks.size()); - List batch = allPks.subList(i, batchEnd); + List batch = allPks.subList(i, batchEnd) + .stream() + .map(normalizedPk -> normalizedPk.partitionKey) + .collect(Collectors.toList()); // batchFilter is [epk(first), maxExclusive) where maxExclusive is: // - epk of the next PK after this batch (first PK of the next batch), or // - partitionScope.getMax() if this is the last batch in the partition - String batchMinInclusive = getEffectivePartitionKeyString(batch.get(0), pkDefinition); + String batchMinInclusive = allPks.get(i).effectivePartitionKeyString; String batchMaxExclusive = batchEnd < allPks.size() - ? getEffectivePartitionKeyString(allPks.get(batchEnd), pkDefinition) + ? allPks.get(batchEnd).effectivePartitionKeyString : partitionScope.getMax(); Range batchFilter = new Range<>(batchMinInclusive, batchMaxExclusive, true, false); @@ -4598,7 +4649,7 @@ private Flux> buildSequentialFluxFromScratch( return buildSequentialBatchFlux( allBatches, null, resourceLink, state, diagnosticsFactory, klass, - collectionRid, queryHash); + collectionRid, queryHash, partitionKeySetHash); } /** @@ -4608,7 +4659,7 @@ private Flux> buildSequentialFluxFromScratch( */ private Flux> buildSequentialFluxFromContinuation( ReadManyByPartitionKeyContinuationToken parsedContinuation, - List partitionKeys, + List normalizedPartitionKeys, SqlQuerySpec customQuery, PartitionKeyDefinition pkDefinition, String resourceLink, @@ -4616,7 +4667,8 @@ private Flux> buildSequentialFluxFromContinuation( ScopedDiagnosticsFactory diagnosticsFactory, Class klass, String collectionRid, - int queryHash) { + String queryHash, + String partitionKeySetHash) { List partitionKeySelectors = PartitionKeyQueryHelper.createPkSelectors(pkDefinition); @@ -4632,31 +4684,24 @@ private Flux> buildSequentialFluxFromContinuation( baseParameters = new ArrayList<>(); } - // Reconstruct batches from the continuation token's batch definitions. - // The current batch + remaining batches define the work left to do. ReadManyByPartitionKeyContinuationToken.BatchDefinition currentBatchDef = parsedContinuation.getCurrentBatch(); List remainingBatchDefs = parsedContinuation.getRemainingBatches(); String initialBackendContinuation = parsedContinuation.getBackendContinuation(); - // Build ordered list: current batch first, then remaining List allBatchDefs = new ArrayList<>(); allBatchDefs.add(currentBatchDef); allBatchDefs.addAll(remainingBatchDefs); - // For each batch definition, filter PKs that belong to it by batchFilter EPK range - // (not partitionScope — which covers the entire physical partition). - // Use partitionScope for FeedRange routing; batchFilter for PK filtering. List allBatches = new ArrayList<>(); for (ReadManyByPartitionKeyContinuationToken.BatchDefinition batchDef : allBatchDefs) { Range partitionScope = batchDef.getPartitionScope(); Range batchFilter = batchDef.getBatchFilter(); - // Filter PKs whose EPK falls within this batch's batchFilter range List batchPks = filterPartitionKeysByEpkRange( - partitionKeys, pkDefinition, batchFilter); + normalizedPartitionKeys, batchFilter); if (batchPks.isEmpty()) { continue; @@ -4677,23 +4722,20 @@ private Flux> buildSequentialFluxFromContinuation( return buildSequentialBatchFlux( allBatches, initialBackendContinuation, resourceLink, state, diagnosticsFactory, klass, - collectionRid, queryHash); + collectionRid, queryHash, partitionKeySetHash); } /** * Filters partition keys to those whose EPK falls within the given range. */ private List filterPartitionKeysByEpkRange( - List partitionKeys, - PartitionKeyDefinition pkDefinition, + List normalizedPartitionKeys, Range epkRange) { List result = new ArrayList<>(); - for (PartitionKey pk : partitionKeys) { - String epk = getEffectivePartitionKeyString(pk, pkDefinition); - - if (epkRange.contains(epk)) { - result.add(pk); + for (NormalizedPartitionKey normalizedPk : normalizedPartitionKeys) { + if (epkRange.contains(normalizedPk.effectivePartitionKeyString)) { + result.add(normalizedPk.partitionKey); } } return result; @@ -4711,7 +4753,8 @@ private Flux> buildSequentialBatchFlux( ScopedDiagnosticsFactory diagnosticsFactory, Class klass, String collectionRid, - int queryHash) { + String queryHash, + String partitionKeySetHash) { List>> sequentialFluxes = new ArrayList<>(); for (int i = 0; i < allBatches.size(); i++) { @@ -4719,7 +4762,6 @@ private Flux> buildSequentialBatchFlux( final BatchDescriptor bd = allBatches.get(i); final String backendContinuation = (i == 0) ? initialBackendContinuation : null; - // Remaining batch definitions for the continuation token (batches after the current one) final List remainingAfterThis = new ArrayList<>(); for (int j = batchIndex + 1; j < allBatches.size(); j++) { BatchDescriptor remaining = allBatches.get(j); @@ -4728,25 +4770,15 @@ private Flux> buildSequentialBatchFlux( remaining.partitionScope, remaining.batchFilter)); } - // Clone query options for this batch CosmosQueryRequestOptions batchQueryOptions = queryOptionsAccessor() .clone(state.getQueryOptions()); queryOptionsAccessor().disallowQueryPlanRetrieval(batchQueryOptions); - // Set FeedRange for routing via partitionScope EPK — handles splits transparently batchQueryOptions.setFeedRange(new FeedRangeEpkImpl(bd.partitionScope)); - // Always set the continuation token on the cloned options — even when null. - // The cloned options may still carry the external caller's continuation token, - // which the backend wouldn't know how to interpret. For the first batch we pass - // the backend continuation from the composite token; for subsequent batches null - // (meaning "start from the beginning of this batch"). ModelBridgeInternal.setQueryRequestOptionsContinuationToken( batchQueryOptions, backendContinuation); - // Execute the batch query using the standard query API with FeedRange for routing. - // FeedRangeEpkImpl(partitionScope) handles routing to the correct physical partition(s), - // including transparently handling partition splits. Flux> batchFlux = createQueryInternal( diagnosticsFactory, resourceLink, @@ -4758,12 +4790,10 @@ private Flux> buildSequentialBatchFlux( UUIDs.nonBlockingRandomUUID(), new AtomicBoolean(false)); - // Current batch as a BatchDefinition for the continuation token final ReadManyByPartitionKeyContinuationToken.BatchDefinition currentBatchDef = new ReadManyByPartitionKeyContinuationToken.BatchDefinition( bd.partitionScope, bd.batchFilter); - // Stamp each FeedResponse with the composite continuation token Flux> stampedFlux = batchFlux.map(feedResponse -> { String backendCont = feedResponse.getContinuationToken(); boolean isLastBatch = remainingAfterThis.isEmpty(); @@ -4780,14 +4810,14 @@ private Flux> buildSequentialBatchFlux( : Collections.emptyList(); ReadManyByPartitionKeyContinuationToken compositeContinuation = new ReadManyByPartitionKeyContinuationToken( - remaining, nextBatchDef, null, collectionRid, queryHash); + remaining, nextBatchDef, null, collectionRid, queryHash, partitionKeySetHash); ModelBridgeInternal.setFeedResponseContinuationToken( compositeContinuation.serialize(), feedResponse); } else { ReadManyByPartitionKeyContinuationToken compositeContinuation = new ReadManyByPartitionKeyContinuationToken( remainingAfterThis, currentBatchDef, backendCont, - collectionRid, queryHash); + collectionRid, queryHash, partitionKeySetHash); ModelBridgeInternal.setFeedResponseContinuationToken( compositeContinuation.serialize(), feedResponse); } @@ -4797,24 +4827,24 @@ private Flux> buildSequentialBatchFlux( sequentialFluxes.add(stampedFlux); } - return Flux.concat(sequentialFluxes); + int fluxConcurrency = Math.max(1, Math.min(Configs.getCPUCnt(), sequentialFluxes.size())); + return Flux.mergeSequential(sequentialFluxes, fluxConcurrency, 1); } - /** - * Returns the effective partition key string for a PartitionKey, handling PartitionKey.NONE - * by mapping to UndefinedPartitionKey. - */ - private static String getEffectivePartitionKeyString( - PartitionKey pk, - PartitionKeyDefinition pkDefinition) { + private static final class NormalizedPartitionKey { + final PartitionKey partitionKey; + final PartitionKeyInternal effectivePkInternal; + final String effectivePartitionKeyString; - PartitionKeyInternal pkInternal = BridgeInternal.getPartitionKeyInternal(pk); - PartitionKeyInternal effectivePkInternal = pkInternal.getComponents() == null - ? PartitionKeyInternal.UndefinedPartitionKey - : pkInternal; + private NormalizedPartitionKey( + PartitionKey partitionKey, + PartitionKeyInternal effectivePkInternal, + String effectivePartitionKeyString) { - return PartitionKeyInternalHelper - .getEffectivePartitionKeyString(effectivePkInternal, pkDefinition); + this.partitionKey = partitionKey; + this.effectivePkInternal = effectivePkInternal; + this.effectivePartitionKeyString = effectivePartitionKeyString; + } } /** @@ -4915,45 +4945,31 @@ static void validateQueryPlanForReadManyByPartitionKey(PartitionedQueryExecution } } - private Map> groupPartitionKeysByPhysicalPartition( - List partitionKeys, + private Map> groupPartitionKeysByPhysicalPartition( + List normalizedPartitionKeys, PartitionKeyDefinition pkDefinition, CollectionRoutingMap routingMap) { - // Use LinkedHashMap so the downstream round-robin interleave is deterministic and the iteration - // order follows insertion order of partition keys (i.e. the order the caller provided). - Map> partitionRangePkMap = new LinkedHashMap<>(); - - for (PartitionKey pk : partitionKeys) { - PartitionKeyInternal pkInternal = BridgeInternal.getPartitionKeyInternal(pk); - - // PartitionKey.NONE wraps NonePartitionKey which has components = null. - // For routing purposes, treat NONE as UndefinedPartitionKey - documents ingested - // without a partition key path are stored with the undefined EPK. - PartitionKeyInternal effectivePkInternal = pkInternal.getComponents() == null - ? PartitionKeyInternal.UndefinedPartitionKey - : pkInternal; + Map> partitionRangePkMap = new LinkedHashMap<>(); - int componentCount = effectivePkInternal.getComponents().size(); + for (NormalizedPartitionKey normalizedPk : normalizedPartitionKeys) { + int componentCount = normalizedPk.effectivePkInternal.getComponents().size(); int definedPathCount = pkDefinition.getPaths().size(); List targetRanges; if (pkDefinition.getKind() == PartitionKind.MULTI_HASH && componentCount < definedPathCount) { - // Partial HPK - compute EPK prefix range and find all overlapping physical partitions Range epkRange = PartitionKeyInternalHelper.getEPKRangeForPrefixPartitionKey( - effectivePkInternal, pkDefinition); + normalizedPk.effectivePkInternal, pkDefinition); targetRanges = routingMap.getOverlappingRanges(epkRange); } else { - // Full PK - maps to exactly one physical partition - String effectivePartitionKeyString = PartitionKeyInternalHelper - .getEffectivePartitionKeyString(effectivePkInternal, pkDefinition); - PartitionKeyRange range = routingMap.getRangeByEffectivePartitionKey(effectivePartitionKeyString); + PartitionKeyRange range = routingMap.getRangeByEffectivePartitionKey( + normalizedPk.effectivePartitionKeyString); targetRanges = Collections.singletonList(range); } for (PartitionKeyRange range : targetRanges) { - partitionRangePkMap.computeIfAbsent(range, k -> new ArrayList<>()).add(pk); + partitionRangePkMap.computeIfAbsent(range, k -> new ArrayList<>()).add(normalizedPk); } } diff --git a/sdk/cosmos/docs/readManyByPartitionKey-design.md b/sdk/cosmos/docs/readManyByPartitionKey-design.md index 9fa747c50ec3..f0df3bce99bd 100644 --- a/sdk/cosmos/docs/readManyByPartitionKey-design.md +++ b/sdk/cosmos/docs/readManyByPartitionKey-design.md @@ -18,7 +18,7 @@ and additional filters. The SDK appends the auto-generated PK WHERE clause to it | Return type | `CosmosPagedFlux` (async) / `CosmosPagedIterable` (sync) | | Custom query format | `SqlQuerySpec` — full query with parameters; SDK ANDs the PK filter | | Partial HPK | Supported from the start; prefix PKs fan out via `getOverlappingRanges` | -| PK deduplication | Done at Spark layer only, not in the SDK | +| PK deduplication | SDK normalizes and deduplicates the partition-key set by EPK before batching; Spark callers should still dedupe the input DataFrame when practical for efficiency | | Spark UDF | New `GetCosmosPartitionKeyValue` UDF | | Custom query validation | Gateway query plan via the standard SDK query-plan retrieval path; reject aggregates/ORDER BY/DISTINCT/GROUP BY/DCount/OFFSET/LIMIT/non-streaming ORDER BY/vector/fulltext | | PK list size | No hard upper-bound enforced; SDK batches internally per physical partition (default 100 PKs per batch, configurable via `COSMOS.READ_MANY_BY_PK_MAX_BATCH_SIZE`) | @@ -148,7 +148,7 @@ Static entry points that accept a DataFrame and Cosmos config. PK extraction sup 1. **UDF-produced column**: DataFrame contains `_partitionKeyIdentity` column (from `GetCosmosPartitionKeyValue` UDF). 2. **Schema-matched columns**: DataFrame columns match the container's PK paths. -Nested partition key paths are not resolved automatically from DataFrame columns and must use the UDF-produced `_partitionKeyIdentity` column. +Top-level DataFrame columns may supply a full or prefix hierarchical partition key directly. Nested partition key paths are not resolved automatically and must use the UDF-produced `_partitionKeyIdentity` column. Falls back with `IllegalArgumentException` if neither mode is possible. @@ -160,7 +160,7 @@ Orchestrator that resolves schema, initializes and broadcasts client state to ex Spark `PartitionReader[InternalRow]` that: -- Deduplicates PKs via `LinkedHashMap` (by PK string representation). +- Preserves the caller's PK list and lets the SDK normalize/dedupe by effective partition key; callers should still dedupe the DataFrame upstream when practical. - Passes the pre-built `CosmosReadManyRequestOptions` (with throughput control, diagnostics, custom serializer) to the SDK. - Uses `TransientIOErrorsRetryingIterator` for retry handling. - Short-circuits empty PK lists to avoid SDK rejection. @@ -228,7 +228,7 @@ routing layer) with `isMinInclusive = true` and `isMaxExclusive = false`. ### Sequential batch execution Batches are sorted by `minInclusive` EPK (lexicographic) and processed one at a time -via `Flux.concat` instead of `Flux.merge`. Each `FeedResponse` carries the composite +via ordered sequential merging with a small amount of Reactor prefetch. Each `FeedResponse` carries the composite continuation token reflecting the current position: ```text From 5b44e3407bd8d27eddf4341dcfe22d1d5dd2a9db Mon Sep 17 00:00:00 2001 From: Fabian Meiswinkel Date: Wed, 22 Apr 2026 10:01:25 +0000 Subject: [PATCH 42/64] Changed the continuation token to not carry partitionScopes anymore --- .../azure-cosmos-spark_3-3_2-12/CHANGELOG.md | 1 + .../azure-cosmos-spark_3-4_2-12/CHANGELOG.md | 1 + .../cosmos/spark/SparkE2EQueryITest.scala | 4 +- .../azure-cosmos-spark_3-5_2-12/CHANGELOG.md | 1 + .../azure-cosmos-spark_3-5_2-13/CHANGELOG.md | 1 + .../com/azure/cosmos/spark/CosmosConfig.scala | 20 +- .../cosmos/spark/CosmosItemsDataSource.scala | 1 - .../spark/CosmosPartitionKeyHelper.scala | 2 +- .../CosmosReadManyByPartitionKeyReader.scala | 7 +- ...tionReaderWithReadManyByPartitionKey.scala | 44 +- ...tryingReadManyByPartitionKeyIterator.scala | 9 + .../azure-cosmos-spark_4-0_2-13/CHANGELOG.md | 1 + .../cosmos/ReadManyByPartitionKeyTest.java | 24 +- ...nyByPartitionKeyContinuationTokenTest.java | 170 +++++--- sdk/cosmos/azure-cosmos/CHANGELOG.md | 4 +- .../azure/cosmos/CosmosAsyncContainer.java | 44 +- .../com/azure/cosmos/CosmosContainer.java | 5 +- .../implementation/AsyncDocumentClient.java | 3 + .../azure/cosmos/implementation/Configs.java | 38 +- ...dManyByPartitionKeyRequestOptionsImpl.java | 113 +++++ .../CosmosReadManyRequestOptionsImpl.java | 24 -- .../ImplementationBridgeHelpers.java | 39 +- ...adManyByPartitionKeyContinuationToken.java | 272 ++++-------- .../ReadManyByPartitionKeyQueryHelper.java | 19 +- .../implementation/RxDocumentClientImpl.java | 143 +++++-- ...sReadManyByPartitionKeyRequestOptions.java | 402 ++++++++++++++++++ .../models/CosmosReadManyRequestOptions.java | 5 - .../main/java/com/azure/cosmos/util/Beta.java | 4 +- sdk/cosmos/cspell.yaml | 1 + .../docs/readManyByPartitionKey-design.md | 385 ++++++++++++----- 30 files changed, 1290 insertions(+), 497 deletions(-) create mode 100644 sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/CosmosReadManyByPartitionKeyRequestOptionsImpl.java create mode 100644 sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/CosmosReadManyByPartitionKeyRequestOptions.java diff --git a/sdk/cosmos/azure-cosmos-spark_3-3_2-12/CHANGELOG.md b/sdk/cosmos/azure-cosmos-spark_3-3_2-12/CHANGELOG.md index a7c445911584..87600e4128c1 100644 --- a/sdk/cosmos/azure-cosmos-spark_3-3_2-12/CHANGELOG.md +++ b/sdk/cosmos/azure-cosmos-spark_3-3_2-12/CHANGELOG.md @@ -4,6 +4,7 @@ #### Features Added * Added new `CosmosItemsDataSource.readManyByPartitionKeys` Spark function to execute bulk queries by a list of pk-values with better efficiency. Configure null handling via `spark.cosmos.read.readManyByPk.nullHandling` - default `Null` treats a null PK column as JSON null (`addNullValue`), `None` treats it as `PartitionKey.NONE` (`addNoneValue` / `NOT IS_DEFINED`). These route to different physical partitions - picking the wrong mode silently returns zero rows. See [PR 48801](https://github.com/Azure/azure-sdk-for-java/pull/48801) +* Added Spark config `spark.cosmos.read.readManyByPk.maxConcurrentBatchPrefetch` (default `1`) to bound the per-task prefetch parallelism the SDK uses inside `readManyByPartitionKeys`. See [PR 48801](https://github.com/Azure/azure-sdk-for-java/pull/48801) #### Breaking Changes diff --git a/sdk/cosmos/azure-cosmos-spark_3-4_2-12/CHANGELOG.md b/sdk/cosmos/azure-cosmos-spark_3-4_2-12/CHANGELOG.md index 9ab39ee9e943..56b546f71d87 100644 --- a/sdk/cosmos/azure-cosmos-spark_3-4_2-12/CHANGELOG.md +++ b/sdk/cosmos/azure-cosmos-spark_3-4_2-12/CHANGELOG.md @@ -4,6 +4,7 @@ #### Features Added * Added new `CosmosItemsDataSource.readManyByPartitionKeys` Spark function to execute bulk queries by a list of pk-values with better efficiency. Configure null handling via `spark.cosmos.read.readManyByPk.nullHandling` - default `Null` treats a null PK column as JSON null (`addNullValue`), `None` treats it as `PartitionKey.NONE` (`addNoneValue` / `NOT IS_DEFINED`). These route to different physical partitions - picking the wrong mode silently returns zero rows. See [PR 48801](https://github.com/Azure/azure-sdk-for-java/pull/48801) +* Added Spark config `spark.cosmos.read.readManyByPk.maxConcurrentBatchPrefetch` (default `1`) to bound the per-task prefetch parallelism the SDK uses inside `readManyByPartitionKeys`. See [PR 48801](https://github.com/Azure/azure-sdk-for-java/pull/48801) #### Breaking Changes diff --git a/sdk/cosmos/azure-cosmos-spark_3-5/src/test/scala/com/azure/cosmos/spark/SparkE2EQueryITest.scala b/sdk/cosmos/azure-cosmos-spark_3-5/src/test/scala/com/azure/cosmos/spark/SparkE2EQueryITest.scala index aa4fe2e3250d..c5bce27bcda9 100644 --- a/sdk/cosmos/azure-cosmos-spark_3-5/src/test/scala/com/azure/cosmos/spark/SparkE2EQueryITest.scala +++ b/sdk/cosmos/azure-cosmos-spark_3-5/src/test/scala/com/azure/cosmos/spark/SparkE2EQueryITest.scala @@ -4,7 +4,7 @@ package com.azure.cosmos.spark import com.azure.cosmos.implementation.TestConfigurations -import com.azure.cosmos.models.{CosmosContainerProperties, CosmosItemRequestOptions, PartitionKey, PartitionKeyDefinition, PartitionKeyDefinitionVersion, PartitionKind, ThroughputProperties} +import com.azure.cosmos.models.{CosmosContainerProperties, CosmosItemRequestOptions, PartitionKey, PartitionKeyBuilder, PartitionKeyDefinition, PartitionKeyDefinitionVersion, PartitionKind, ThroughputProperties} import com.azure.cosmos.spark.udf.GetCosmosPartitionKeyValue import com.fasterxml.jackson.databind.node.ObjectNode import org.apache.spark.sql.functions.expr @@ -111,7 +111,7 @@ class SparkE2EQueryITest rows.map(_.getAs[String]("pk")).toSet shouldEqual Set("pkA", "pkB") rows.map(_.getAs[String]("payload")).toSet shouldEqual Set("value-pkA", "value-pkB") } - + "spark readManyByPartitionKeys" can "require the UDF for nested partition key paths and succeed with it" in { val cosmosEndpoint = TestConfigurations.HOST val cosmosMasterKey = TestConfigurations.MASTER_KEY diff --git a/sdk/cosmos/azure-cosmos-spark_3-5_2-12/CHANGELOG.md b/sdk/cosmos/azure-cosmos-spark_3-5_2-12/CHANGELOG.md index 1e8bc11ffe69..acc6d23dbe2b 100644 --- a/sdk/cosmos/azure-cosmos-spark_3-5_2-12/CHANGELOG.md +++ b/sdk/cosmos/azure-cosmos-spark_3-5_2-12/CHANGELOG.md @@ -4,6 +4,7 @@ #### Features Added * Added new `CosmosItemsDataSource.readManyByPartitionKeys` Spark function to execute bulk queries by a list of pk-values with better efficiency. Configure null handling via `spark.cosmos.read.readManyByPk.nullHandling` - default `Null` treats a null PK column as JSON null (`addNullValue`), `None` treats it as `PartitionKey.NONE` (`addNoneValue` / `NOT IS_DEFINED`). These route to different physical partitions - picking the wrong mode silently returns zero rows. See [PR 48801](https://github.com/Azure/azure-sdk-for-java/pull/48801) +* Added Spark config `spark.cosmos.read.readManyByPk.maxConcurrentBatchPrefetch` (default `1`) to bound the per-task prefetch parallelism the SDK uses inside `readManyByPartitionKeys`. See [PR 48801](https://github.com/Azure/azure-sdk-for-java/pull/48801) #### Breaking Changes diff --git a/sdk/cosmos/azure-cosmos-spark_3-5_2-13/CHANGELOG.md b/sdk/cosmos/azure-cosmos-spark_3-5_2-13/CHANGELOG.md index b5dc43bab19d..ea6312673520 100644 --- a/sdk/cosmos/azure-cosmos-spark_3-5_2-13/CHANGELOG.md +++ b/sdk/cosmos/azure-cosmos-spark_3-5_2-13/CHANGELOG.md @@ -4,6 +4,7 @@ #### Features Added * Added new `CosmosItemsDataSource.readManyByPartitionKeys` Spark function to execute bulk queries by a list of pk-values with better efficiency. Configure null handling via `spark.cosmos.read.readManyByPk.nullHandling` - default `Null` treats a null PK column as JSON null (`addNullValue`), `None` treats it as `PartitionKey.NONE` (`addNoneValue` / `NOT IS_DEFINED`). These route to different physical partitions - picking the wrong mode silently returns zero rows. See [PR 48801](https://github.com/Azure/azure-sdk-for-java/pull/48801) +* Added Spark config `spark.cosmos.read.readManyByPk.maxConcurrentBatchPrefetch` (default `1`) to bound the per-task prefetch parallelism the SDK uses inside `readManyByPartitionKeys`. See [PR 48801](https://github.com/Azure/azure-sdk-for-java/pull/48801) #### Breaking Changes diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosConfig.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosConfig.scala index 1a1025cd7ad2..75b47274cd2b 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosConfig.scala +++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosConfig.scala @@ -93,6 +93,7 @@ private[spark] object CosmosConfigNames { val ReadRuntimeFilteringEnabled = "spark.cosmos.read.runtimeFiltering.enabled" val ReadManyFilteringEnabled = "spark.cosmos.read.readManyFiltering.enabled" val ReadManyByPkNullHandling = "spark.cosmos.read.readManyByPk.nullHandling" + val ReadManyByPkMaxConcurrentBatchPrefetch = "spark.cosmos.read.readManyByPk.maxConcurrentBatchPrefetch" val ViewsRepositoryPath = "spark.cosmos.views.repositoryPath" val DiagnosticsMode = "spark.cosmos.diagnostics" val DiagnosticsSamplingMaxCount = "spark.cosmos.diagnostics.sampling.maxCount" @@ -228,6 +229,7 @@ private[spark] object CosmosConfigNames { ReadRuntimeFilteringEnabled, ReadManyFilteringEnabled, ReadManyByPkNullHandling, + ReadManyByPkMaxConcurrentBatchPrefetch, ViewsRepositoryPath, DiagnosticsMode, DiagnosticsSamplingIntervalInSeconds, @@ -1045,7 +1047,8 @@ private case class CosmosReadConfig(readConsistencyStrategy: ReadConsistencyStra runtimeFilteringEnabled: Boolean, readManyFilteringConfig: CosmosReadManyFilteringConfig, responseContinuationTokenLimitInKb: Option[Int] = None, - readManyByPkTreatNullAsNone: Boolean = false) + readManyByPkTreatNullAsNone: Boolean = false, + readManyByPkMaxConcurrentBatchPrefetch: Int = 1) private object SchemaConversionModes extends Enumeration { type SchemaConversionMode = Value @@ -1153,6 +1156,17 @@ private object CosmosReadConfig { "partitions - picking the wrong mode for your data will silently return zero rows." ) + private val ReadManyByPkMaxConcurrentBatchPrefetch = CosmosConfigEntry[Int]( + key = CosmosConfigNames.ReadManyByPkMaxConcurrentBatchPrefetch, + mandatory = false, + defaultValue = Some(1), + parseFromStringFunction = value => Math.max(1, value.toInt), + helpMessage = "The maximum number of per-physical-partition batches whose first page is prefetched " + + "concurrently inside a single Spark task by the SDK's readManyByPartitionKeys execution. The " + + "default is `1`, because Spark already parallelises across tasks - increase this when individual " + + "tasks span many physical partitions and additional intra-task prefetch is desired." + ) + def parseCosmosReadConfig(cfg: Map[String, String]): CosmosReadConfig = { val forceEventualConsistency = CosmosConfigEntry.parse(cfg, ForceEventualConsistency) val readConsistencyStrategyOverride = CosmosConfigEntry.parse(cfg, ReadConsistencyStrategyOverride) @@ -1177,6 +1191,7 @@ private object CosmosReadConfig { val readManyFilteringConfig = CosmosReadManyFilteringConfig.parseCosmosReadManyFilterConfig(cfg) val readManyByPkNullHandling = CosmosConfigEntry.parse(cfg, ReadManyByPkNullHandling) val readManyByPkTreatNullAsNone = readManyByPkNullHandling.getOrElse("Null").equalsIgnoreCase("None") + val readManyByPkMaxConcurrentBatchPrefetch = CosmosConfigEntry.parse(cfg, ReadManyByPkMaxConcurrentBatchPrefetch).getOrElse(1) val effectiveReadConsistencyStrategy = if (readConsistencyStrategyOverride.getOrElse(ReadConsistencyStrategy.DEFAULT) != ReadConsistencyStrategy.DEFAULT) { readConsistencyStrategyOverride.get @@ -1209,7 +1224,8 @@ private object CosmosReadConfig { runtimeFilteringEnabled.get, readManyFilteringConfig, responseContinuationTokenLimitInKb, - readManyByPkTreatNullAsNone) + readManyByPkTreatNullAsNone, + readManyByPkMaxConcurrentBatchPrefetch) } } diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosItemsDataSource.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosItemsDataSource.scala index b8b4a184c78e..ef310298f29a 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosItemsDataSource.scala +++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosItemsDataSource.scala @@ -5,7 +5,6 @@ package com.azure.cosmos.spark import com.azure.cosmos.models.{CosmosItemIdentity, PartitionKey, PartitionKeyBuilder} import com.azure.cosmos.spark.CosmosPredicates.assertOnSparkDriver import com.azure.cosmos.spark.diagnostics.BasicLoggingTrait -import com.azure.cosmos.SparkBridgeInternal import org.apache.spark.sql.{DataFrame, Row, SparkSession} import java.util diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosPartitionKeyHelper.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosPartitionKeyHelper.scala index 84c3f2fadaf2..a96572c5f632 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosPartitionKeyHelper.scala +++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosPartitionKeyHelper.scala @@ -34,7 +34,7 @@ private[spark] object CosmosPartitionKeyHelper extends BasicLoggingTrait { // // (?i) : The whole matching is case-insensitive // pk[(](.*)[)]: partitionKey Value - private val cosmosPartitionKeyStringRegx = """(?i)pk[(](.*)[)]""".r + private val cosmosPartitionKeyStringRegx = """(?i)^pk\((.*)\)$""".r private val objectMapper = Utils.getSimpleObjectMapper def getCosmosPartitionKeyValueString(partitionKeyValue: List[Object]): String = { diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosReadManyByPartitionKeyReader.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosReadManyByPartitionKeyReader.scala index 5cde39941c0d..eda6e755faf2 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosReadManyByPartitionKeyReader.scala +++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosReadManyByPartitionKeyReader.scala @@ -100,9 +100,10 @@ private[spark] class CosmosReadManyByPartitionKeyReader( classOf[ObjectNode]) .block() } catch { - case ex: CosmosException => - logDebug(s"Warm-up readItem for metadata caches completed with exception: ${ex.getMessage}", ex) - None + case _: CosmosException => + // Expected when the random read targets a non-existent item; we only need the + // routing map / collection cache populated as a side-effect of the call. + () } // 4. Serialize and broadcast client state diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/ItemsPartitionReaderWithReadManyByPartitionKey.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/ItemsPartitionReaderWithReadManyByPartitionKey.scala index cfe5757716fd..b254ad51464d 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/ItemsPartitionReaderWithReadManyByPartitionKey.scala +++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/ItemsPartitionReaderWithReadManyByPartitionKey.scala @@ -3,10 +3,10 @@ package com.azure.cosmos.spark -import com.azure.cosmos.{CosmosAsyncContainer, CosmosEndToEndOperationLatencyPolicyConfigBuilder, CosmosItemSerializerNoExceptionWrapping, SparkBridgeInternal} +import com.azure.cosmos.{CosmosAsyncContainer, CosmosEndToEndOperationLatencyPolicyConfigBuilder, CosmosItemSerializer, CosmosItemSerializerNoExceptionWrapping, SparkBridgeInternal} import com.azure.cosmos.implementation.spark.OperationContextAndListenerTuple import com.azure.cosmos.implementation.{ImplementationBridgeHelpers, ObjectNodeMap, SparkRowItem, Utils} -import com.azure.cosmos.models.{CosmosReadManyRequestOptions, ModelBridgeInternal, PartitionKey, PartitionKeyDefinition, SqlQuerySpec} +import com.azure.cosmos.models.{CosmosReadManyByPartitionKeyRequestOptions, ModelBridgeInternal, PartitionKey, PartitionKeyDefinition, SqlQuerySpec} import com.azure.cosmos.spark.BulkWriter.getThreadInfo import com.azure.cosmos.spark.CosmosTableSchemaInferrer.IdAttributeName import com.azure.cosmos.spark.diagnostics.{DetailedFeedDiagnosticsProvider, DiagnosticsContext, DiagnosticsLoader, LoggerHelper, SparkTaskContext} @@ -43,14 +43,15 @@ private[spark] case class ItemsPartitionReaderWithReadManyByPartitionKey private lazy val log = LoggerHelper.getLogger(diagnosticsConfig, this.getClass) - private val readManyOptions = new CosmosReadManyRequestOptions() + private val readManyOptions = new CosmosReadManyByPartitionKeyRequestOptions() private val readManyOptionsImpl = ImplementationBridgeHelpers - .CosmosReadManyRequestOptionsHelper - .getCosmosReadManyRequestOptionsAccessor + .CosmosReadManyByPartitionKeyRequestOptionsHelper + .getCosmosReadManyByPartitionKeyRequestOptionsAccessor .getImpl(readManyOptions) private val readConfig = CosmosReadConfig.parseCosmosReadConfig(config) ThroughputControlHelper.populateThroughputControlGroupName(readManyOptionsImpl, readConfig.throughputControlConfig) + readManyOptions.setMaxConcurrentBatchPrefetch(readConfig.readManyByPkMaxConcurrentBatchPrefetch) private val operationContext = { assert(taskContext != null) @@ -120,13 +121,16 @@ private[spark] case class ItemsPartitionReaderWithReadManyByPartitionKey readManyOptionsImpl .setCustomItemSerializer( new CosmosItemSerializerNoExceptionWrapping { + // The reader-side path of readManyByPartitionKeys never produces items that the + // SDK needs to serialize back out (no item bodies are sent on the wire other than + // SqlParameter values, which do not flow through this serializer). However, some + // generic SDK code paths (e.g. SqlParameter#getValueAsBytes via the configured + // serializer) will defensively call serialize() during diagnostics or during query + // plan handling. Delegating to the default Jackson-based serializer keeps those + // code paths working and avoids spurious UnsupportedOperationExceptions while we + // still own deserialization (the hot path) below. override def serialize[T](item: T): util.Map[String, AnyRef] = { - throw new UnsupportedOperationException( - s"Serialization is not supported by the custom item serializer in " + - s"ItemsPartitionReaderWithReadManyByPartitionKey; this serializer is intended " + - s"for deserializing read-many responses into SparkRowItem only. " + - s"Unexpected item type: ${if (item == null) "null" else item.getClass.getName}" - ) + CosmosItemSerializer.DEFAULT_SERIALIZER.serialize(item) } override def deserialize[T](jsonNodeMap: util.Map[String, AnyRef], classType: Class[T]): T = { @@ -159,9 +163,25 @@ private[spark] case class ItemsPartitionReaderWithReadManyByPartitionKey // Collect all PK values upfront - the SDK now owns normalization and deduplication, // so Spark preserves the caller's input as-is and relies on the SDK's set-based semantics. // Callers should still dedupe their DataFrame input when practical to avoid extra work. + // + // NOTE on memory footprint: every PartitionKey from this iterator is materialized into a + // single ArrayList here, and the SDK in turn keeps the normalized set, the EPK->PK map, + // and the per-batch BatchDescriptor lists alive for the lifetime of the reader. For a + // single Spark partition with O(N) input rows this is O(N) memory on the executor; if + // upstream Spark partitioning sends millions of distinct PKs to one task this can become + // a noticeable allocation. Repartition upstream when N is very large. + private val PK_COUNT_LARGE_INPUT_WARN_THRESHOLD = 200000 private lazy val pkList = { val values = new java.util.ArrayList[PartitionKey]() readManyPartitionKeys.foreach(values.add) + if (values.size() > PK_COUNT_LARGE_INPUT_WARN_THRESHOLD) { + log.logWarning( + s"ItemsPartitionReaderWithReadManyByPartitionKey received ${values.size()} partition " + + s"keys for a single Spark partition (feedRange=$feedRange). Large PK lists materialize " + + s"the full set in memory plus the SDK's normalized batch metadata; consider increasing " + + s"upstream Spark parallelism so each task processes <= " + + s"$PK_COUNT_LARGE_INPUT_WARN_THRESHOLD distinct partition keys.") + } values } @@ -212,7 +232,7 @@ private[spark] case class ItemsPartitionReaderWithReadManyByPartitionKey private val fluxFactory: String => CosmosPagedFlux[SparkRowItem] = { (continuationToken: String) => // Reuse the preconfigured request options and only update the continuation // token so the SDK resumes from the last fully committed page. - readManyOptionsImpl.setRequestContinuation(continuationToken) + readManyOptions.setContinuationToken(continuationToken) customQueryOpt match { case Some(query) => cosmosAsyncContainer.readManyByPartitionKeys(pkList, query, readManyOptions, classOf[SparkRowItem]) diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/TransientIOErrorsRetryingReadManyByPartitionKeyIterator.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/TransientIOErrorsRetryingReadManyByPartitionKeyIterator.scala index be2e0e2fe299..28044d00a069 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/TransientIOErrorsRetryingReadManyByPartitionKeyIterator.scala +++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/TransientIOErrorsRetryingReadManyByPartitionKeyIterator.scala @@ -133,6 +133,15 @@ private[spark] class TransientIOErrorsRetryingReadManyByPartitionKeyIterator[TSp feedResponse) } val iteratorCandidate = feedResponse.getResults.iterator().asScala.buffered + // INVARIANT: it is safe to record the continuation token BEFORE the items in this + // FeedResponse have been drained because executeWithRetry only wraps `hasNext`, + // and the buffered iterator is always fully drained before the next page is fetched. + // If a transient failure occurs while draining items from the *current* page, the + // page is replayed from this continuation token, which is the start of *this* page + // (the server's continuation token points at the next page), so on retry the SDK + // re-issues the request that produced the page we were processing. Items already + // emitted to the caller may be re-emitted; readManyByPartitionKeys is idempotent + // and returns documents (not deltas), so duplicates from replay are acceptable. lastContinuationToken.set(feedResponse.getContinuationToken) if (iteratorCandidate.hasNext) { diff --git a/sdk/cosmos/azure-cosmos-spark_4-0_2-13/CHANGELOG.md b/sdk/cosmos/azure-cosmos-spark_4-0_2-13/CHANGELOG.md index f686d3ee37da..28c503b4dfec 100644 --- a/sdk/cosmos/azure-cosmos-spark_4-0_2-13/CHANGELOG.md +++ b/sdk/cosmos/azure-cosmos-spark_4-0_2-13/CHANGELOG.md @@ -4,6 +4,7 @@ #### Features Added * Added new `CosmosItemsDataSource.readManyByPartitionKeys` Spark function to execute bulk queries by a list of pk-values with better efficiency. Configure null handling via `spark.cosmos.read.readManyByPk.nullHandling` - default `Null` treats a null PK column as JSON null (`addNullValue`), `None` treats it as `PartitionKey.NONE` (`addNoneValue` / `NOT IS_DEFINED`). These route to different physical partitions - picking the wrong mode silently returns zero rows. See [PR 48801](https://github.com/Azure/azure-sdk-for-java/pull/48801) +* Added Spark config `spark.cosmos.read.readManyByPk.maxConcurrentBatchPrefetch` (default `1`) to bound the per-task prefetch parallelism the SDK uses inside `readManyByPartitionKeys`. See [PR 48801](https://github.com/Azure/azure-sdk-for-java/pull/48801) #### Breaking Changes diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/ReadManyByPartitionKeyTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/ReadManyByPartitionKeyTest.java index 95b2c0d87607..c5f3e46978ed 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/ReadManyByPartitionKeyTest.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/ReadManyByPartitionKeyTest.java @@ -5,8 +5,6 @@ package com.azure.cosmos; -import com.azure.cosmos.implementation.CosmosReadManyRequestOptionsImpl; -import com.azure.cosmos.implementation.ImplementationBridgeHelpers; import com.azure.cosmos.models.CosmosContainerProperties; import com.azure.cosmos.models.CosmosItemRequestOptions; import com.azure.cosmos.models.FeedResponse; @@ -458,7 +456,7 @@ public void singlePk_readManyByPartitionKey_withRequestOptions() { List items = createSinglePkItems("pkOpts", 3); List pkValues = Collections.singletonList(new PartitionKey("pkOpts")); - com.azure.cosmos.models.CosmosReadManyRequestOptions options = new com.azure.cosmos.models.CosmosReadManyRequestOptions(); + com.azure.cosmos.models.CosmosReadManyByPartitionKeyRequestOptions options = new com.azure.cosmos.models.CosmosReadManyByPartitionKeyRequestOptions(); AtomicInteger deserializeCount = new AtomicInteger(); options.setCustomItemSerializer(new CosmosItemSerializerNoExceptionWrapping() { @Override @@ -585,13 +583,9 @@ public void singlePk_continuationToken_resumesCorrectly() { List itemsFromFirstPage = allPages.get(0).getResults(); // Second pass: resume from the continuation token - com.azure.cosmos.models.CosmosReadManyRequestOptions options2 = - new com.azure.cosmos.models.CosmosReadManyRequestOptions(); - ((com.azure.cosmos.implementation.CosmosReadManyRequestOptionsImpl) - ImplementationBridgeHelpers.CosmosReadManyRequestOptionsHelper - .getCosmosReadManyRequestOptionsAccessor() - .getImpl(options2)) - .setRequestContinuation(continuationAfterFirstPage); + com.azure.cosmos.models.CosmosReadManyByPartitionKeyRequestOptions options2 = + new com.azure.cosmos.models.CosmosReadManyByPartitionKeyRequestOptions(); + options2.setContinuationToken(continuationAfterFirstPage); List> remainingPages = asyncContainer .readManyByPartitionKeys(pkValues, options2, ObjectNode.class) @@ -668,13 +662,9 @@ public void singlePk_continuationToken_resumesCorrectly_whenInputContainsDuplica assertThat(continuationAfterFirstPage).isNotNull(); List itemsFromFirstPage = allPages.get(0).getResults(); - com.azure.cosmos.models.CosmosReadManyRequestOptions options2 = - new com.azure.cosmos.models.CosmosReadManyRequestOptions(); - ((CosmosReadManyRequestOptionsImpl) - ImplementationBridgeHelpers.CosmosReadManyRequestOptionsHelper - .getCosmosReadManyRequestOptionsAccessor() - .getImpl(options2)) - .setRequestContinuation(continuationAfterFirstPage); + com.azure.cosmos.models.CosmosReadManyByPartitionKeyRequestOptions options2 = + new com.azure.cosmos.models.CosmosReadManyByPartitionKeyRequestOptions(); + options2.setContinuationToken(continuationAfterFirstPage); List> remainingPages = asyncContainer .readManyByPartitionKeys(pkValues, options2, ObjectNode.class) diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyContinuationTokenTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyContinuationTokenTest.java index b7f2c7e34b31..35309dcda135 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyContinuationTokenTest.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyContinuationTokenTest.java @@ -24,17 +24,30 @@ public class ReadManyByPartitionKeyContinuationTokenTest { private static final String TEST_QUERY_HASH = "12345"; private static final String TEST_PARTITION_KEY_SET_HASH = "67890"; + /** + * Builds a BatchDefinition whose batchFilter is the half-open EPK range [min, max). + */ + private static ReadManyByPartitionKeyContinuationToken.BatchDefinition bd(String min, String max) { + return new ReadManyByPartitionKeyContinuationToken.BatchDefinition( + new Range<>(min, max, true, false)); + } + + private static List bds( + ReadManyByPartitionKeyContinuationToken.BatchDefinition... defs) { + return new ArrayList<>(Arrays.asList(defs)); + } + @Test(groups = { "unit" }) public void roundtrip_withBackendContinuation() { - List> remaining = Arrays.asList( - new Range<>("05C1E0", "0BF333", true, false), - new Range<>("0BF333", "FF", true, false)); - Range current = new Range<>("", "05C1E0", true, false); + List remaining = bds( + bd("05C1E0", "0BF333"), + bd("0BF333", "FF")); + ReadManyByPartitionKeyContinuationToken.BatchDefinition current = bd("", "05C1E0"); String backendCont = "eyJDb21wb3NpdGVUb2tlbg=="; ReadManyByPartitionKeyContinuationToken token = new ReadManyByPartitionKeyContinuationToken( - remaining, current, backendCont, TEST_COLLECTION_RID, TEST_QUERY_HASH); + remaining, current, backendCont, TEST_COLLECTION_RID, TEST_QUERY_HASH, TEST_PARTITION_KEY_SET_HASH); String serialized = token.serialize(); assertThat(serialized).isNotNull().isNotEmpty(); @@ -45,64 +58,60 @@ public void roundtrip_withBackendContinuation() { assertThat(deserialized.getBackendContinuation()).isEqualTo(backendCont); assertThat(deserialized.getCollectionRid()).isEqualTo(TEST_COLLECTION_RID); assertThat(deserialized.getQueryHash()).isEqualTo(TEST_QUERY_HASH); + assertThat(deserialized.getPartitionKeySetHash()).isEqualTo(TEST_PARTITION_KEY_SET_HASH); ReadManyByPartitionKeyContinuationToken.BatchDefinition currentBatch = deserialized.getCurrentBatch(); - assertThat(currentBatch.getPartitionScope().getMin()).isEqualTo(""); - assertThat(currentBatch.getPartitionScope().getMax()).isEqualTo("05C1E0"); + assertThat(currentBatch.getBatchFilter().getMin()).isEqualTo(""); + assertThat(currentBatch.getBatchFilter().getMax()).isEqualTo("05C1E0"); List remainingBatches = deserialized.getRemainingBatches(); assertThat(remainingBatches).hasSize(2); - assertThat(remainingBatches.get(0).getPartitionScope().getMin()).isEqualTo("05C1E0"); - assertThat(remainingBatches.get(0).getPartitionScope().getMax()).isEqualTo("0BF333"); - assertThat(remainingBatches.get(1).getPartitionScope().getMin()).isEqualTo("0BF333"); - assertThat(remainingBatches.get(1).getPartitionScope().getMax()).isEqualTo("FF"); + assertThat(remainingBatches.get(0).getBatchFilter().getMin()).isEqualTo("05C1E0"); + assertThat(remainingBatches.get(0).getBatchFilter().getMax()).isEqualTo("0BF333"); + assertThat(remainingBatches.get(1).getBatchFilter().getMin()).isEqualTo("0BF333"); + assertThat(remainingBatches.get(1).getBatchFilter().getMax()).isEqualTo("FF"); } @Test(groups = { "unit" }) public void roundtrip_withNullBackendContinuation() { - List> remaining = Collections.singletonList( - new Range<>("0BF333", "FF", true, false)); - Range current = new Range<>("05C1E0", "0BF333", true, false); - ReadManyByPartitionKeyContinuationToken token = new ReadManyByPartitionKeyContinuationToken( - remaining, current, null, TEST_COLLECTION_RID, TEST_QUERY_HASH); + bds(bd("0BF333", "FF")), bd("05C1E0", "0BF333"), + null, TEST_COLLECTION_RID, TEST_QUERY_HASH, TEST_PARTITION_KEY_SET_HASH); String serialized = token.serialize(); ReadManyByPartitionKeyContinuationToken deserialized = ReadManyByPartitionKeyContinuationToken.deserialize(serialized); assertThat(deserialized.getBackendContinuation()).isNull(); - assertThat(deserialized.getCurrentBatch().getPartitionScope().getMin()).isEqualTo("05C1E0"); - assertThat(deserialized.getCurrentBatch().getPartitionScope().getMax()).isEqualTo("0BF333"); + assertThat(deserialized.getCurrentBatch().getBatchFilter().getMin()).isEqualTo("05C1E0"); + assertThat(deserialized.getCurrentBatch().getBatchFilter().getMax()).isEqualTo("0BF333"); assertThat(deserialized.getRemainingBatches()).hasSize(1); } @Test(groups = { "unit" }) public void roundtrip_emptyRemainingBatches() { - Range current = new Range<>("0BF333", "FF", true, false); - ReadManyByPartitionKeyContinuationToken token = new ReadManyByPartitionKeyContinuationToken( - Collections.emptyList(), current, "someCont", TEST_COLLECTION_RID, TEST_QUERY_HASH); + Collections.emptyList(), bd("0BF333", "FF"), + "someCont", TEST_COLLECTION_RID, TEST_QUERY_HASH, TEST_PARTITION_KEY_SET_HASH); String serialized = token.serialize(); ReadManyByPartitionKeyContinuationToken deserialized = ReadManyByPartitionKeyContinuationToken.deserialize(serialized); assertThat(deserialized.getRemainingBatches()).isEmpty(); - assertThat(deserialized.getCurrentBatch().getPartitionScope().getMin()).isEqualTo("0BF333"); + assertThat(deserialized.getCurrentBatch().getBatchFilter().getMin()).isEqualTo("0BF333"); assertThat(deserialized.getBackendContinuation()).isEqualTo("someCont"); } @Test(groups = { "unit" }) public void roundtrip_lastBatchNoContinuation() { - Range current = new Range<>("0BF333", "FF", true, false); - ReadManyByPartitionKeyContinuationToken token = new ReadManyByPartitionKeyContinuationToken( - Collections.emptyList(), current, null, TEST_COLLECTION_RID, TEST_QUERY_HASH); + Collections.emptyList(), bd("0BF333", "FF"), + null, TEST_COLLECTION_RID, TEST_QUERY_HASH, TEST_PARTITION_KEY_SET_HASH); String serialized = token.serialize(); ReadManyByPartitionKeyContinuationToken deserialized = @@ -114,10 +123,11 @@ public void roundtrip_lastBatchNoContinuation() { @Test(groups = { "unit" }) public void deserialize_malformedInput_throws() { + // Either the base64 decoder or the JSON parsing layer rejects garbage; both raise + // IllegalArgumentException, which is the contract callers depend on. assertThatThrownBy(() -> ReadManyByPartitionKeyContinuationToken.deserialize("not-valid-base64!!!") - ).isInstanceOf(IllegalArgumentException.class) - .hasMessageContaining("Failed to deserialize"); + ).isInstanceOf(IllegalArgumentException.class); } @Test(groups = { "unit" }) @@ -135,55 +145,99 @@ public void deserialize_null_throws() { } @Test(groups = { "unit" }) - public void serialized_isBase64() { - Range current = new Range<>("", "FF", true, false); + public void deserialize_unsupportedVersion_throws() { + // Hand-craft a token JSON with version=999 (a future format) and ensure it is rejected. + String json = "{\"v\":999,\"rb\":[],\"cb\":{\"bf\":{\"min\":\"\",\"max\":\"FF\"}}," + + "\"bc\":null,\"cr\":\"" + TEST_COLLECTION_RID + "\",\"qh\":\"" + TEST_QUERY_HASH + "\",\"ph\":\"" + TEST_PARTITION_KEY_SET_HASH + "\"}"; + String serialized = java.util.Base64.getEncoder().encodeToString(json.getBytes(java.nio.charset.StandardCharsets.UTF_8)); + + // The root cause carries the precise "Unsupported version" message; the outer wrapper + // gives a generic "Failed to deserialize" hint. Match against the full chain. + assertThatThrownBy(() -> ReadManyByPartitionKeyContinuationToken.deserialize(serialized)) + .isInstanceOf(IllegalArgumentException.class) + .hasStackTraceContaining("Unsupported readManyByPartitionKeys continuation token version"); + } + + @Test(groups = { "unit" }) + public void serialized_includesVersionField_andIsBase64() { ReadManyByPartitionKeyContinuationToken token = new ReadManyByPartitionKeyContinuationToken( - Collections.emptyList(), current, null, TEST_COLLECTION_RID, 0); + Collections.emptyList(), bd("", "FF"), null, + TEST_COLLECTION_RID, TEST_QUERY_HASH, TEST_PARTITION_KEY_SET_HASH); String serialized = token.serialize(); - - // Should be valid Base64 (no whitespace, no special chars except +/=) assertThat(serialized).matches("[A-Za-z0-9+/=]+"); - // Decoding should produce valid JSON String json = new String( java.util.Base64.getDecoder().decode(serialized), java.nio.charset.StandardCharsets.UTF_8); assertThat(json).startsWith("{"); assertThat(json).endsWith("}"); + // The wire format must include the version field so future SDKs can detect/reject + // tokens from incompatible versions. + assertThat(json).contains("\"v\":1"); + } + + @Test(groups = { "unit" }) + public void serialized_doesNotIncludePartitionScope() { + // The partition routing scope is intentionally NOT persisted in the continuation token. + // It is rederived at execution time from the live PartitionKeyRange cache so partition + // splits never cause stale routing information to be embedded in a token. + ReadManyByPartitionKeyContinuationToken token = + new ReadManyByPartitionKeyContinuationToken( + bds(bd("AA", "BB"), bd("BB", "CC")), + bd("", "AA"), null, + TEST_COLLECTION_RID, TEST_QUERY_HASH, TEST_PARTITION_KEY_SET_HASH); + + String json = new String( + java.util.Base64.getDecoder().decode(token.serialize()), + java.nio.charset.StandardCharsets.UTF_8); + + assertThat(json).doesNotContain("\"ps\""); + assertThat(json).contains("\"bf\""); + } + + @Test(groups = { "unit" }) + public void version_roundtripsAsCurrent() { + ReadManyByPartitionKeyContinuationToken token = + new ReadManyByPartitionKeyContinuationToken( + Collections.emptyList(), bd("", "FF"), null, + TEST_COLLECTION_RID, TEST_QUERY_HASH, TEST_PARTITION_KEY_SET_HASH); + + ReadManyByPartitionKeyContinuationToken deserialized = + ReadManyByPartitionKeyContinuationToken.deserialize(token.serialize()); + + assertThat(deserialized.getVersion()).isEqualTo(ReadManyByPartitionKeyContinuationToken.CURRENT_VERSION); + assertThat(ReadManyByPartitionKeyContinuationToken.CURRENT_VERSION).isEqualTo(1); } @Test(groups = { "unit" }) public void rangesPreserveMinMaxInclusive() { - Range current = new Range<>("AB", "CD", true, false); ReadManyByPartitionKeyContinuationToken token = new ReadManyByPartitionKeyContinuationToken( - Collections.singletonList(new Range<>("CD", "EF", true, false)), - current, null, TEST_COLLECTION_RID, TEST_QUERY_HASH); + bds(bd("CD", "EF")), bd("AB", "CD"), null, + TEST_COLLECTION_RID, TEST_QUERY_HASH, TEST_PARTITION_KEY_SET_HASH); String serialized = token.serialize(); ReadManyByPartitionKeyContinuationToken deserialized = ReadManyByPartitionKeyContinuationToken.deserialize(serialized); - // Verify that deserialized ranges have the canonical min-inclusive, max-exclusive form - Range currentScope = deserialized.getCurrentBatch().getPartitionScope(); - assertThat(currentScope.isMinInclusive()).isTrue(); - assertThat(currentScope.isMaxInclusive()).isFalse(); - Range remainingScope = deserialized.getRemainingBatches().get(0).getPartitionScope(); - assertThat(remainingScope.isMinInclusive()).isTrue(); - assertThat(remainingScope.isMaxInclusive()).isFalse(); + Range currentFilter = deserialized.getCurrentBatch().getBatchFilter(); + assertThat(currentFilter.isMinInclusive()).isTrue(); + assertThat(currentFilter.isMaxInclusive()).isFalse(); + Range remainingFilter = deserialized.getRemainingBatches().get(0).getBatchFilter(); + assertThat(remainingFilter.isMinInclusive()).isTrue(); + assertThat(remainingFilter.isMaxInclusive()).isFalse(); } @Test(groups = { "unit" }) public void collectionRidAndQueryHash_roundtrip() { - Range current = new Range<>("", "FF", true, false); String rid = "dbs/myDb/colls/myColl"; String hash = "98765"; ReadManyByPartitionKeyContinuationToken token = new ReadManyByPartitionKeyContinuationToken( - Collections.emptyList(), current, null, rid, hash); + Collections.emptyList(), bd("", "FF"), null, rid, hash, TEST_PARTITION_KEY_SET_HASH); String serialized = token.serialize(); ReadManyByPartitionKeyContinuationToken deserialized = @@ -195,11 +249,10 @@ public void collectionRidAndQueryHash_roundtrip() { @Test(groups = { "unit" }) public void partitionKeySetHash_roundtrip() { - Range current = new Range<>("", "FF", true, false); - ReadManyByPartitionKeyContinuationToken token = new ReadManyByPartitionKeyContinuationToken( - Collections.emptyList(), current, null, TEST_COLLECTION_RID, TEST_QUERY_HASH, TEST_PARTITION_KEY_SET_HASH); + Collections.emptyList(), bd("", "FF"), null, + TEST_COLLECTION_RID, TEST_QUERY_HASH, TEST_PARTITION_KEY_SET_HASH); String serialized = token.serialize(); ReadManyByPartitionKeyContinuationToken deserialized = @@ -217,7 +270,6 @@ public void computePartitionKeySetHash_isStableAcrossDuplicateAndReorderedEpks() .isEqualTo(ReadManyByPartitionKeyContinuationToken.computePartitionKeySetHash(epks2)); } - @Test(groups = { "unit" }) public void computePartitionKeySetHash_returnsStableHexDigest() { String hash = ReadManyByPartitionKeyContinuationToken.computePartitionKeySetHash( @@ -272,35 +324,25 @@ public void computeQueryHash_noParams_stableHash() { @Test(groups = { "unit" }) public void batchDefinition_roundtrip() { - Range partitionScope = new Range<>("", "05C1E0", true, false); - Range batchFilter = new Range<>("01", "03", true, false); - - ReadManyByPartitionKeyContinuationToken.BatchDefinition bd = - new ReadManyByPartitionKeyContinuationToken.BatchDefinition(partitionScope, batchFilter); - - List remaining = Collections.singletonList( - new ReadManyByPartitionKeyContinuationToken.BatchDefinition( - new Range<>("05C1E0", "FF", true, false), - new Range<>("05C1E0", "0A", true, false))); + ReadManyByPartitionKeyContinuationToken.BatchDefinition currentBd = bd("01", "03"); + List remaining = + Collections.singletonList(bd("05C1E0", "0A")); ReadManyByPartitionKeyContinuationToken token = new ReadManyByPartitionKeyContinuationToken( - remaining, bd, "cont", TEST_COLLECTION_RID, TEST_QUERY_HASH); + remaining, currentBd, "cont", TEST_COLLECTION_RID, TEST_QUERY_HASH, TEST_PARTITION_KEY_SET_HASH); String serialized = token.serialize(); ReadManyByPartitionKeyContinuationToken deserialized = ReadManyByPartitionKeyContinuationToken.deserialize(serialized); ReadManyByPartitionKeyContinuationToken.BatchDefinition currentBatch = deserialized.getCurrentBatch(); - assertThat(currentBatch.getPartitionScope().getMin()).isEqualTo(""); - assertThat(currentBatch.getPartitionScope().getMax()).isEqualTo("05C1E0"); assertThat(currentBatch.getBatchFilter().getMin()).isEqualTo("01"); assertThat(currentBatch.getBatchFilter().getMax()).isEqualTo("03"); ReadManyByPartitionKeyContinuationToken.BatchDefinition remainingBatch = deserialized.getRemainingBatches().get(0); - assertThat(remainingBatch.getPartitionScope().getMin()).isEqualTo("05C1E0"); assertThat(remainingBatch.getBatchFilter().getMin()).isEqualTo("05C1E0"); assertThat(remainingBatch.getBatchFilter().getMax()).isEqualTo("0A"); } -} +} \ No newline at end of file diff --git a/sdk/cosmos/azure-cosmos/CHANGELOG.md b/sdk/cosmos/azure-cosmos/CHANGELOG.md index ca5e400ace55..36ba92cdee2f 100644 --- a/sdk/cosmos/azure-cosmos/CHANGELOG.md +++ b/sdk/cosmos/azure-cosmos/CHANGELOG.md @@ -4,11 +4,13 @@ #### Features Added * Added support for change feed with `startFrom` point-in-time on merged partitions by enabling the `CHANGE_FEED_WITH_START_TIME_POST_MERGE` SDK capability. - See [PR 48752](https://github.com/Azure/azure-sdk-for-java/pull/48752) -* Added new `readManyByPartitionKeys` API on `CosmosAsyncContainer` / `CosmosContainer` to bulk-query all documents matching a list of partition key values with better efficiency than issuing individual queries. Also fixes nested partition key selector generation for `readMany` and `readAllItems`. See [PR 48801](https://github.com/Azure/azure-sdk-for-java/pull/48801) +* Added new `readManyByPartitionKeys` API on `CosmosAsyncContainer` / `CosmosContainer` to bulk-query all documents matching a list of partition key values with better efficiency than issuing individual queries. See [PR 48801](https://github.com/Azure/azure-sdk-for-java/pull/48801) +* Added `CosmosReadManyByPartitionKeyRequestOptions` - a dedicated request-options type for `readManyByPartitionKeys` that exposes `setContinuationToken(String)` for resuming previous invocations and `setMaxConcurrentBatchPrefetch(int)` to bound per-call prefetch parallelism. See [PR 48801](https://github.com/Azure/azure-sdk-for-java/pull/48801) #### Breaking Changes #### Bugs Fixed +* Fixed `readMany` and `readAllItems` returning incorrect results on containers whose partition key path is nested (e.g. `/address/city`) due to malformed selector generation. - See [PR 48801](https://github.com/Azure/azure-sdk-for-java/pull/48801) * Fixed an issue where the throughput control `throughputQueryMono` was always subscribed even when `targetThroughput` is used (not `targetThroughputThreshold`), causing unnecessary `throughputSettings/read` permission requirement for AAD principals. - See [PR 48800](https://github.com/Azure/azure-sdk-for-java/pull/48800) * Fixed JVM `` deadlock when multiple threads concurrently trigger Cosmos SDK class loading for the first time. - See [PR 48689](https://github.com/Azure/azure-sdk-for-java/pull/48689) * Fixed an issue where `CustomItemSerializer` was incorrectly applied to internal SDK query pipeline structures (e.g., `OrderByRowResult`, `Document`), causing deserialization failures in ORDER BY, GROUP BY, aggregate, DISTINCT, and hybrid search queries. - See [PR 48811](https://github.com/Azure/azure-sdk-for-java/pull/48811) diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java index 12458b458338..5d2009d24d4f 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java @@ -60,6 +60,7 @@ import com.azure.cosmos.models.CosmosPatchItemRequestOptions; import com.azure.cosmos.models.CosmosPatchOperations; import com.azure.cosmos.models.CosmosQueryRequestOptions; +import com.azure.cosmos.models.CosmosReadManyByPartitionKeyRequestOptions; import com.azure.cosmos.models.CosmosReadManyRequestOptions; import com.azure.cosmos.models.FeedRange; import com.azure.cosmos.models.FeedResponse; @@ -126,6 +127,10 @@ private static ImplementationBridgeHelpers.CosmosReadManyRequestOptionsHelper.Co return ImplementationBridgeHelpers.CosmosReadManyRequestOptionsHelper.getCosmosReadManyRequestOptionsAccessor(); } + private static ImplementationBridgeHelpers.CosmosReadManyByPartitionKeyRequestOptionsHelper.CosmosReadManyByPartitionKeyRequestOptionsAccessor readManyByPkOptionsAccessor() { + return ImplementationBridgeHelpers.CosmosReadManyByPartitionKeyRequestOptionsHelper.getCosmosReadManyByPartitionKeyRequestOptionsAccessor(); + } + private static ImplementationBridgeHelpers.CosmosDiagnosticsContextHelper.CosmosDiagnosticsContextAccessor ctxAccessor() { return ImplementationBridgeHelpers.CosmosDiagnosticsContextHelper.getCosmosDiagnosticsContextAccessor(); } @@ -148,6 +153,12 @@ private static ImplementationBridgeHelpers.CosmosBatchRequestOptionsHelper.Cosmo private final static Logger logger = LoggerFactory.getLogger(CosmosAsyncContainer.class); + // Sentinel values for CosmosQueryRequestOptions.maxDegreeOfParallelism used when no + // caller-supplied value is present: 0 == "uninitialized / SDK-chooses", + // -1 == "unbounded" (the value RxDocumentClientImpl interprets as no concurrency cap). + private static final int DEFAULT_MAX_DEGREE_OF_PARALLELISM = 0; + private static final int UNBOUNDED_MAX_DEGREE_OF_PARALLELISM = -1; + private final CosmosAsyncDatabase database; private final String id; private final String link; @@ -1637,7 +1648,7 @@ public CosmosPagedFlux readManyByPartitionKeys( */ public CosmosPagedFlux readManyByPartitionKeys( List partitionKeys, - CosmosReadManyRequestOptions requestOptions, + CosmosReadManyByPartitionKeyRequestOptions requestOptions, Class classType) { return this.readManyByPartitionKeys(partitionKeys, null, requestOptions, classType); @@ -1697,7 +1708,7 @@ public CosmosPagedFlux readManyByPartitionKeys( public CosmosPagedFlux readManyByPartitionKeys( List partitionKeys, SqlQuerySpec customQuery, - CosmosReadManyRequestOptions requestOptions, + CosmosReadManyByPartitionKeyRequestOptions requestOptions, Class classType) { if (partitionKeys == null) { @@ -1722,7 +1733,7 @@ public CosmosPagedFlux readManyByPartitionKeys( private Function>> readManyByPartitionKeyInternalFunc( List partitionKeys, SqlQuerySpec customQuery, - CosmosReadManyRequestOptions requestOptions, + CosmosReadManyByPartitionKeyRequestOptions requestOptions, Class classType) { CosmosAsyncClient client = this.getDatabase().getClient(); @@ -1731,18 +1742,30 @@ private Function>> readManyByPa // It will be set on the cloned CosmosQueryRequestOptions so that // QueryFeedOperationState picks it up automatically. String requestContinuation = requestOptions != null - ? readManyOptionsAccessor().getRequestContinuation(requestOptions) + ? readManyByPkOptionsAccessor().getContinuationToken(requestOptions) + : null; + + // Resolve the max-concurrent-batch-prefetch knob; default to availableProcessors() so + // the SDK fully utilises CPU when the caller does not override. + Integer prefetchOverride = requestOptions != null + ? readManyByPkOptionsAccessor().getMaxConcurrentBatchPrefetch(requestOptions) : null; + int maxConcurrentBatchPrefetch = prefetchOverride != null + ? prefetchOverride + : Configs.getCPUCnt(); return (pagedFluxOptions -> { CosmosQueryRequestOptions queryRequestOptions = requestOptions == null ? new CosmosQueryRequestOptions() - : queryOptionsAccessor().clone(readManyOptionsAccessor().getImpl(requestOptions)); - // Honor any caller-provided MaxDegreeOfParallelism; only default to the "unbounded" sentinel - // (-1) when the value is still at the default (0). CosmosReadManyRequestOptions currently does not - // expose MDOP, so this branch is defensive in case it is plumbed through in the future. - if (queryRequestOptions.getMaxDegreeOfParallelism() == 0) { - queryRequestOptions.setMaxDegreeOfParallelism(-1); + : queryOptionsAccessor().clone(readManyByPkOptionsAccessor().getImpl(requestOptions)); + // CosmosQueryRequestOptionsBase initializes MaxDegreeOfParallelism to 0 (the + // "uninitialized / SDK-chooses" sentinel); -1 is the "unbounded" sentinel that + // RxDocumentClientImpl recognizes for query parallelism. Honor any caller-provided + // value but default the uninitialized case to unbounded for readManyByPartitionKeys. + // CosmosReadManyRequestOptions does not currently expose MDOP, so this only matters + // if it is plumbed through in the future. + if (queryRequestOptions.getMaxDegreeOfParallelism() == DEFAULT_MAX_DEGREE_OF_PARALLELISM) { + queryRequestOptions.setMaxDegreeOfParallelism(UNBOUNDED_MAX_DEGREE_OF_PARALLELISM); } queryRequestOptions.setQueryName("readManyByPartitionKeys"); @@ -1777,6 +1800,7 @@ private Function>> readManyByPa customQuery, BridgeInternal.getLink(this), state, + maxConcurrentBatchPrefetch, classType) .map(response -> prepareFeedResponse(response, false)); }); diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosContainer.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosContainer.java index e61e4ee1f422..a733dc272079 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosContainer.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosContainer.java @@ -21,6 +21,7 @@ import com.azure.cosmos.models.CosmosPatchItemRequestOptions; import com.azure.cosmos.models.CosmosPatchOperations; import com.azure.cosmos.models.CosmosQueryRequestOptions; +import com.azure.cosmos.models.CosmosReadManyByPartitionKeyRequestOptions; import com.azure.cosmos.models.CosmosReadManyRequestOptions; import com.azure.cosmos.models.FeedRange; import com.azure.cosmos.models.FeedResponse; @@ -574,7 +575,7 @@ public CosmosPagedIterable readManyByPartitionKeys( */ public CosmosPagedIterable readManyByPartitionKeys( List partitionKeys, - CosmosReadManyRequestOptions requestOptions, + CosmosReadManyByPartitionKeyRequestOptions requestOptions, Class classType) { return getCosmosPagedIterable(this.asyncContainer.readManyByPartitionKeys(partitionKeys, requestOptions, classType)); @@ -634,7 +635,7 @@ public CosmosPagedIterable readManyByPartitionKeys( public CosmosPagedIterable readManyByPartitionKeys( List partitionKeys, SqlQuerySpec customQuery, - CosmosReadManyRequestOptions requestOptions, + CosmosReadManyByPartitionKeyRequestOptions requestOptions, Class classType) { return getCosmosPagedIterable(this.asyncContainer.readManyByPartitionKeys(partitionKeys, customQuery, requestOptions, classType)); diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/AsyncDocumentClient.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/AsyncDocumentClient.java index 0996af4010b3..3018188f274c 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/AsyncDocumentClient.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/AsyncDocumentClient.java @@ -1594,6 +1594,8 @@ Mono> readMany( * @param customQuery optional custom query (for projections/additional filters) - null means SELECT * FROM c * @param collectionLink link for the documentcollection/container to be queried * @param state the query operation state (may carry a composite continuation token via requestContinuation) + * @param maxConcurrentBatchPrefetch the maximum number of per-physical-partition batches whose first + * page is prefetched concurrently. Must be >= 1. * @param klass class type * @param the type parameter * @return a Flux with feed response pages of documents @@ -1603,6 +1605,7 @@ Flux> readManyByPartitionKeys( SqlQuerySpec customQuery, String collectionLink, QueryFeedOperationState state, + int maxConcurrentBatchPrefetch, Class klass); /** diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/Configs.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/Configs.java index b896430e85cb..18eef0544e18 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/Configs.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/Configs.java @@ -828,19 +828,45 @@ public static int getMinTargetBulkMicroBatchSize() { } public static int getReadManyByPkMaxBatchSize() { - String valueFromSystemProperty = System.getProperty(READ_MANY_BY_PK_MAX_BATCH_SIZE); - if (valueFromSystemProperty != null && !valueFromSystemProperty.isEmpty()) { - return Math.max(1, Integer.parseInt(valueFromSystemProperty)); + Integer parsed = parsePositiveInt(System.getProperty(READ_MANY_BY_PK_MAX_BATCH_SIZE), READ_MANY_BY_PK_MAX_BATCH_SIZE); + if (parsed != null) { + return parsed; } - String valueFromEnvVariable = System.getenv(READ_MANY_BY_PK_MAX_BATCH_SIZE_VARIABLE); - if (valueFromEnvVariable != null && !valueFromEnvVariable.isEmpty()) { - return Math.max(1, Integer.parseInt(valueFromEnvVariable)); + parsed = parsePositiveInt(System.getenv(READ_MANY_BY_PK_MAX_BATCH_SIZE_VARIABLE), READ_MANY_BY_PK_MAX_BATCH_SIZE_VARIABLE); + if (parsed != null) { + return parsed; } return DEFAULT_READ_MANY_BY_PK_MAX_BATCH_SIZE; } + /** + * Parses a non-empty string as a positive integer (>= 1). On parse failure or + * non-positive result, logs a WARN and returns null so the caller can fall back + * to its default. A null/empty input is also treated as "no value". + */ + private static Integer parsePositiveInt(String value, String configName) { + if (value == null || value.isEmpty()) { + return null; + } + try { + int parsed = Integer.parseInt(value); + if (parsed < 1) { + logger.warn( + "Ignoring invalid value '{}' for config '{}'. Value must be >= 1. Falling back to default.", + value, configName); + return null; + } + return parsed; + } catch (NumberFormatException e) { + logger.warn( + "Ignoring non-numeric value '{}' for config '{}'. Falling back to default.", + value, configName); + return null; + } + } + public static int getMaxBulkMicroBatchConcurrency() { String valueFromSystemProperty = System.getProperty(MAX_BULK_MICRO_BATCH_CONCURRENCY); if (valueFromSystemProperty != null && !valueFromSystemProperty.isEmpty()) { diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/CosmosReadManyByPartitionKeyRequestOptionsImpl.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/CosmosReadManyByPartitionKeyRequestOptionsImpl.java new file mode 100644 index 000000000000..0480beaa3e4d --- /dev/null +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/CosmosReadManyByPartitionKeyRequestOptionsImpl.java @@ -0,0 +1,113 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. +package com.azure.cosmos.implementation; + +/** + * Internal implementation backing the public {@code CosmosReadManyByPartitionKeyRequestOptions} + * facade. Holds state specific to the {@code readManyByPartitionKeys} operation. + */ +public class CosmosReadManyByPartitionKeyRequestOptionsImpl + extends CosmosQueryRequestOptionsBase { + + private String continuationToken; + private Integer maxConcurrentBatchPrefetch; + private Integer maxItemCount; + + public CosmosReadManyByPartitionKeyRequestOptionsImpl() { + super(); + } + + public CosmosReadManyByPartitionKeyRequestOptionsImpl(CosmosReadManyByPartitionKeyRequestOptionsImpl options) { + super(options); + this.continuationToken = options.continuationToken; + this.maxConcurrentBatchPrefetch = options.maxConcurrentBatchPrefetch; + this.maxItemCount = options.maxItemCount; + } + + /** + * Gets the composite continuation token for readManyByPartitionKeys. + * + * @return the continuation token, or null if not set. + */ + public String getContinuationToken() { + return this.continuationToken; + } + + /** + * Sets the composite continuation token for readManyByPartitionKeys. + * + * @param continuationToken the continuation token from a previous invocation. + * @return this instance. + */ + public CosmosReadManyByPartitionKeyRequestOptionsImpl setContinuationToken(String continuationToken) { + this.continuationToken = continuationToken; + return this; + } + + /** + * Gets the maximum number of per-physical-partition batches whose first page is + * prefetched concurrently. {@code null} means the SDK default applies. + * + * @return the max concurrent batch prefetch, or null if not set. + */ + public Integer getMaxConcurrentBatchPrefetch() { + return this.maxConcurrentBatchPrefetch; + } + + /** + * Sets the maximum number of per-physical-partition batches whose first page is + * prefetched concurrently. + * + * @param maxConcurrentBatchPrefetch the max concurrent batch prefetch (must be >= 1). + * @return this instance. + */ + public CosmosReadManyByPartitionKeyRequestOptionsImpl setMaxConcurrentBatchPrefetch(int maxConcurrentBatchPrefetch) { + this.maxConcurrentBatchPrefetch = maxConcurrentBatchPrefetch; + return this; + } + + @Override + public Integer getMaxItemCount() { + return this.maxItemCount; + } + + public CosmosReadManyByPartitionKeyRequestOptionsImpl setMaxItemCount(Integer maxItemCount) { + this.maxItemCount = maxItemCount; + return this; + } + + @Override + public Boolean isContentResponseOnWriteEnabled() { + return null; + } + + @Override + public Boolean getNonIdempotentWriteRetriesEnabled() { + return null; + } + + @Override + public Boolean isScanInQueryEnabled() { + return null; + } + + @Override + public Integer getMaxDegreeOfParallelism() { + return null; + } + + @Override + public Integer getMaxBufferedItemCount() { + return null; + } + + @Override + public Integer getMaxPrefetchPageCount() { + return null; + } + + @Override + public String getQueryNameOrDefault(String defaultQueryName) { + return null; + } +} diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/CosmosReadManyRequestOptionsImpl.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/CosmosReadManyRequestOptionsImpl.java index 9828a4cb7ad6..05ee42c66753 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/CosmosReadManyRequestOptionsImpl.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/CosmosReadManyRequestOptionsImpl.java @@ -3,9 +3,6 @@ package com.azure.cosmos.implementation; public class CosmosReadManyRequestOptionsImpl extends CosmosQueryRequestOptionsBase { - - private String requestContinuation; - /** * Instantiates a new read many request options. */ @@ -19,27 +16,6 @@ public CosmosReadManyRequestOptionsImpl() { */ public CosmosReadManyRequestOptionsImpl(CosmosReadManyRequestOptionsImpl options) { super(options); - this.requestContinuation = options.requestContinuation; - } - - /** - * Gets the composite continuation token for readManyByPartitionKeys. - * - * @return the continuation token, or null if not set. - */ - public String getRequestContinuation() { - return this.requestContinuation; - } - - /** - * Sets the composite continuation token for readManyByPartitionKeys. - * - * @param requestContinuation the continuation token from a previous invocation. - * @return this instance. - */ - public CosmosReadManyRequestOptionsImpl setRequestContinuation(String requestContinuation) { - this.requestContinuation = requestContinuation; - return this; } @Override diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ImplementationBridgeHelpers.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ImplementationBridgeHelpers.java index 246e9f232317..0be7b33d967b 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ImplementationBridgeHelpers.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ImplementationBridgeHelpers.java @@ -69,6 +69,7 @@ import com.azure.cosmos.models.CosmosPatchItemRequestOptions; import com.azure.cosmos.models.CosmosPatchOperations; import com.azure.cosmos.models.CosmosQueryRequestOptions; +import com.azure.cosmos.models.CosmosReadManyByPartitionKeyRequestOptions; import com.azure.cosmos.models.CosmosReadManyRequestOptions; import com.azure.cosmos.models.FeedRange; import com.azure.cosmos.models.FeedResponse; @@ -351,7 +352,43 @@ public static CosmosReadManyRequestOptionsAccessor getCosmosReadManyRequestOptio public interface CosmosReadManyRequestOptionsAccessor { CosmosQueryRequestOptionsBase getImpl(CosmosReadManyRequestOptions options); - String getRequestContinuation(CosmosReadManyRequestOptions options); + } + } + + public static final class CosmosReadManyByPartitionKeyRequestOptionsHelper { + private final static AtomicBoolean cosmosReadManyByPkRequestOptionsClassLoaded = new AtomicBoolean(false); + private final static AtomicReference accessor = new AtomicReference<>(); + + private CosmosReadManyByPartitionKeyRequestOptionsHelper() {} + + public static void setCosmosReadManyByPartitionKeyRequestOptionsAccessor( + final CosmosReadManyByPartitionKeyRequestOptionsAccessor newAccessor) { + if (!accessor.compareAndSet(null, newAccessor)) { + logger.debug("CosmosReadManyByPartitionKeyRequestOptionsAccessor already initialized!"); + } else { + logger.debug("Setting CosmosReadManyByPartitionKeyRequestOptionsAccessor..."); + cosmosReadManyByPkRequestOptionsClassLoaded.set(true); + } + } + + public static CosmosReadManyByPartitionKeyRequestOptionsAccessor getCosmosReadManyByPartitionKeyRequestOptionsAccessor() { + if (!cosmosReadManyByPkRequestOptionsClassLoaded.get()) { + logger.debug("Initializing CosmosReadManyByPartitionKeyRequestOptionsAccessor..."); + initializeAllAccessors(); + } + + CosmosReadManyByPartitionKeyRequestOptionsAccessor snapshot = accessor.get(); + if (snapshot == null) { + logger.error("CosmosReadManyByPartitionKeyRequestOptionsAccessor is not initialized yet!"); + } + + return snapshot; + } + + public interface CosmosReadManyByPartitionKeyRequestOptionsAccessor { + CosmosQueryRequestOptionsBase getImpl(CosmosReadManyByPartitionKeyRequestOptions options); + String getContinuationToken(CosmosReadManyByPartitionKeyRequestOptions options); + Integer getMaxConcurrentBatchPrefetch(CosmosReadManyByPartitionKeyRequestOptions options); } } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyContinuationToken.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyContinuationToken.java index b25faff9e934..77ad2dd16f9b 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyContinuationToken.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyContinuationToken.java @@ -26,26 +26,33 @@ *

* Captures the state needed to resume a readManyByPartitionKeys operation: *

    - *
  • {@code remainingBatches} — batch definitions of batches not yet started
    • - *
  • {@code currentBatch} — batch definition of the batch currently being processed
    • - *
  • {@code backendContinuation} — backend query continuation within the current batch (nullable)
    • + *
  • {@code remainingBatches} - batch definitions of batches not yet started
    • + *
  • {@code currentBatch} - batch definition of the batch currently being processed
    • + *
  • {@code backendContinuation} - backend query continuation within the current batch (nullable)
    • *
- * Each batch definition has two EPK ranges: - *
    - *
  • {@code partitionScope} — the physical partition's EPK range at the time the operation - * started. Used for routing queries (via FeedRange). Only causes fan-out overhead if a - * split has actually occurred.
    • - *
  • {@code batchFilter} — the EPK sub-range that identifies which PKs belong to this batch. - * Within a physical partition, PKs are sorted by EPK and split into batches; the filter - * range spans the EPKs of the PKs in this batch. Used to reconstruct the exact same set - * of PKs per batch when resuming from a continuation token.
    • - *
- * EPK ranges are used instead of PkRangeIds so the token survives partition splits. + * Each batch is identified solely by its {@link BatchDefinition#getBatchFilter() batchFilter} + * EPK range - the half-open range {@code [minInclusive, maxExclusive)} that contains the EPKs + * of all PKs assigned to the batch. The physical-partition routing range used at execution + * time (the FeedRange set on {@code CosmosQueryRequestOptions}) is not persisted; + * it is rederived at execution time from the current PartitionKeyRange cache by taking the + * union of all partition-key-range EPK ranges that overlap the batch filter + * ({@code [min(minEpk), max(maxEpk))}). *

- * Serialized as JSON → Base64 to keep the token opaque. + * This means the token survives partition splits without ever encoding stale partition + * boundaries: after a split, the rederived routing range exactly matches the new physical + * partition boundaries (one or more of them), keeping query-RU cost minimal. It also keeps + * the serialized token small (one EPK range per batch instead of two). + *

+ * Serialized as JSON -> Base64 to keep the token opaque. The serialized form embeds a + * {@code "v"} version field so future format evolutions can be detected and rejected + * (or migrated) cleanly without silently misinterpreting an older token. */ public final class ReadManyByPartitionKeyContinuationToken { + /** Wire format version. Bump on any breaking change to the JSON shape. */ + public static final int CURRENT_VERSION = 1; + + private static final String VERSION_PROPERTY = "v"; private static final String REMAINING_BATCHES_PROPERTY = "rb"; private static final String CURRENT_BATCH_PROPERTY = "cb"; private static final String BACKEND_CONTINUATION_PROPERTY = "bc"; @@ -53,6 +60,9 @@ public final class ReadManyByPartitionKeyContinuationToken { private static final String QUERY_HASH_PROPERTY = "qh"; private static final String PARTITION_KEY_SET_HASH_PROPERTY = "ph"; + @JsonProperty(VERSION_PROPERTY) + private final int version; + @JsonProperty(REMAINING_BATCHES_PROPERTY) private final List remainingBatches; @@ -71,55 +81,12 @@ public final class ReadManyByPartitionKeyContinuationToken { @JsonProperty(PARTITION_KEY_SET_HASH_PROPERTY) private final String partitionKeySetHash; - @JsonCreator - ReadManyByPartitionKeyContinuationToken( - @JsonProperty(REMAINING_BATCHES_PROPERTY) List remainingBatches, - @JsonProperty(CURRENT_BATCH_PROPERTY) BatchDefinitionDto currentBatch, - @JsonProperty(BACKEND_CONTINUATION_PROPERTY) String backendContinuation, - @JsonProperty(COLLECTION_RID_PROPERTY) String collectionRid, - @JsonProperty(QUERY_HASH_PROPERTY) String queryHash, - @JsonProperty(PARTITION_KEY_SET_HASH_PROPERTY) String partitionKeySetHash) { - - this.remainingBatches = remainingBatches; - this.currentBatch = currentBatch; - this.backendContinuation = backendContinuation; - this.collectionRid = collectionRid; - this.queryHash = queryHash; - this.partitionKeySetHash = partitionKeySetHash; - } - - public ReadManyByPartitionKeyContinuationToken( - List remainingBatches, - BatchDefinition currentBatch, - String backendContinuation, - String collectionRid, - int queryHash) { - - this(remainingBatches, currentBatch, backendContinuation, collectionRid, String.valueOf(queryHash), "0"); - } - - public ReadManyByPartitionKeyContinuationToken( - List remainingBatches, - BatchDefinition currentBatch, - String backendContinuation, - String collectionRid, - int queryHash, - int partitionKeySetHash) { - - this(remainingBatches, currentBatch, backendContinuation, collectionRid, - String.valueOf(queryHash), String.valueOf(partitionKeySetHash)); - } - - public ReadManyByPartitionKeyContinuationToken( - List remainingBatches, - BatchDefinition currentBatch, - String backendContinuation, - String collectionRid, - String queryHash) { - - this(remainingBatches, currentBatch, backendContinuation, collectionRid, queryHash, "0"); - } - + /** + * Production constructor used by RxDocumentClientImpl when stamping each FeedResponse + * with a continuation token. Callers supply just the batch filter for both the current + * batch and every remaining batch, plus all three identity fingerprints + * (collectionRid + queryHash + partitionKeySetHash). Routing scopes are not persisted. + */ public ReadManyByPartitionKeyContinuationToken( List remainingBatches, BatchDefinition currentBatch, @@ -131,6 +98,7 @@ public ReadManyByPartitionKeyContinuationToken( checkNotNull(currentBatch, "Argument 'currentBatch' must not be null."); checkNotNull(remainingBatches, "Argument 'remainingBatches' must not be null."); + this.version = CURRENT_VERSION; this.remainingBatches = new ArrayList<>(remainingBatches.size()); for (BatchDefinition bd : remainingBatches) { this.remainingBatches.add(BatchDefinitionDto.fromBatchDefinition(bd)); @@ -143,66 +111,44 @@ public ReadManyByPartitionKeyContinuationToken( } /** - * Convenience constructor that accepts EPK ranges directly (without separate partitionScope/batchFilter). - * Each range is used as both the partitionScope and batchFilter in the BatchDefinition. - * This is an interim API until the full batch redesign with distinct partitionScope/batchFilter is wired up. + * Deserialization constructor invoked by Jackson. Validates the version field so a + * token from an incompatible future SDK is rejected with a clear error rather than + * being silently misinterpreted. */ - public ReadManyByPartitionKeyContinuationToken( - List> remainingBatchRanges, - Range currentBatchRange, - String backendContinuation, - String collectionRid, - int queryHash) { - - this(remainingBatchRanges, currentBatchRange, backendContinuation, collectionRid, String.valueOf(queryHash), "0"); - } - - public ReadManyByPartitionKeyContinuationToken( - List> remainingBatchRanges, - Range currentBatchRange, - String backendContinuation, - String collectionRid, - int queryHash, - int partitionKeySetHash) { - - this(remainingBatchRanges, currentBatchRange, backendContinuation, collectionRid, - String.valueOf(queryHash), String.valueOf(partitionKeySetHash)); - } - - public ReadManyByPartitionKeyContinuationToken( - List> remainingBatchRanges, - Range currentBatchRange, - String backendContinuation, - String collectionRid, - String queryHash) { - - this(remainingBatchRanges, currentBatchRange, backendContinuation, collectionRid, queryHash, "0"); - } - - public ReadManyByPartitionKeyContinuationToken( - List> remainingBatchRanges, - Range currentBatchRange, - String backendContinuation, - String collectionRid, - String queryHash, - String partitionKeySetHash) { - - checkNotNull(currentBatchRange, "Argument 'currentBatchRange' must not be null."); - checkNotNull(remainingBatchRanges, "Argument 'remainingBatchRanges' must not be null."); + @JsonCreator + ReadManyByPartitionKeyContinuationToken( + @JsonProperty(VERSION_PROPERTY) Integer version, + @JsonProperty(REMAINING_BATCHES_PROPERTY) List remainingBatches, + @JsonProperty(CURRENT_BATCH_PROPERTY) BatchDefinitionDto currentBatch, + @JsonProperty(BACKEND_CONTINUATION_PROPERTY) String backendContinuation, + @JsonProperty(COLLECTION_RID_PROPERTY) String collectionRid, + @JsonProperty(QUERY_HASH_PROPERTY) String queryHash, + @JsonProperty(PARTITION_KEY_SET_HASH_PROPERTY) String partitionKeySetHash) { - this.remainingBatches = new ArrayList<>(remainingBatchRanges.size()); - for (Range range : remainingBatchRanges) { - BatchDefinition bd = new BatchDefinition(range, range); - this.remainingBatches.add(BatchDefinitionDto.fromBatchDefinition(bd)); + // Tokens written before the version field existed will deserialize with version == null. + // Treat null as version 1 (the format that existed when this field was introduced) to + // remain forward-compatible with any tokens emitted by an in-flight pre-versioned beta. + int effectiveVersion = (version == null) ? CURRENT_VERSION : version; + if (effectiveVersion != CURRENT_VERSION) { + throw new IllegalArgumentException( + "Unsupported readManyByPartitionKeys continuation token version: " + effectiveVersion + + ". This SDK supports version " + CURRENT_VERSION + "."); } - this.currentBatch = BatchDefinitionDto.fromBatchDefinition( - new BatchDefinition(currentBatchRange, currentBatchRange)); + + this.version = effectiveVersion; + this.remainingBatches = remainingBatches; + this.currentBatch = currentBatch; this.backendContinuation = backendContinuation; this.collectionRid = collectionRid; this.queryHash = queryHash; this.partitionKeySetHash = partitionKeySetHash; } + @JsonIgnore + public int getVersion() { + return version; + } + @JsonIgnore public List getRemainingBatches() { List result = new ArrayList<>(remainingBatches.size()); @@ -268,24 +214,6 @@ public static String computeQueryHash(SqlQuerySpec querySpec) { } } - static int computeLegacyQueryHash(SqlQuerySpec querySpec) { - if (querySpec == null) { - return 0; - } - int hash = 17; - String queryText = querySpec.getQueryText(); - hash = 31 * hash + (queryText != null ? queryText.hashCode() : 0); - List params = querySpec.getParameters(); - if (params != null) { - for (SqlParameter param : params) { - hash = 31 * hash + (param.getName() != null ? param.getName().hashCode() : 0); - Object value = param.getValue(Object.class); - hash = 31 * hash + (value != null ? value.hashCode() : 0); - } - } - return hash; - } - /** * Computes a stable hash for the normalized set of partition key EPK values. * Duplicate and reordered inputs intentionally produce the same digest. @@ -322,41 +250,10 @@ public static String computePartitionKeySetHash(List partitionKeyEpks) { return murmurHash128Hex(output.toByteArray()); } - static int computeLegacyPartitionKeySetHash(List partitionKeyEpks) { - if (partitionKeyEpks == null || partitionKeyEpks.isEmpty()) { - return 0; - } - - List normalizedEpks = new ArrayList<>(partitionKeyEpks.size()); - for (String epk : partitionKeyEpks) { - if (epk != null) { - normalizedEpks.add(epk); - } - } - - if (normalizedEpks.isEmpty()) { - return 0; - } - - Collections.sort(normalizedEpks); - - int hash = 17; - String previous = null; - for (String epk : normalizedEpks) { - if (epk.equals(previous)) { - continue; - } - - hash = 31 * hash + epk.hashCode(); - previous = epk; - } - - return hash; - } - private static void updateHashInput(ByteArrayOutputStream output, String value) { if (value != null) { - output.writeBytes(value.getBytes(StandardCharsets.UTF_8)); + byte[] bytes = value.getBytes(StandardCharsets.UTF_8); + output.write(bytes, 0, bytes.length); } output.write(0); } @@ -411,6 +308,9 @@ public static ReadManyByPartitionKeyContinuationToken deserialize(String seriali byte[] decoded = Base64.getDecoder().decode(serialized); String json = new String(decoded, StandardCharsets.UTF_8); return Utils.getSimpleObjectMapper().readValue(json, ReadManyByPartitionKeyContinuationToken.class); + } catch (IllegalArgumentException e) { + // Preserve our own version-mismatch IllegalArgumentException without wrapping. + throw e; } catch (Exception e) { throw new IllegalArgumentException( "Failed to deserialize ReadManyByPartitionKeyContinuationToken. " + @@ -419,30 +319,22 @@ public static ReadManyByPartitionKeyContinuationToken deserialize(String seriali } /** - * Identifies a single batch in a readManyByPartitionKeys operation. + * Identifies a single batch in a readManyByPartitionKeys operation by its EPK filter range. *

- * Each batch has two EPK ranges: - *

    - *
  • {@code partitionScope} — the physical partition's full EPK range. Used for routing - * the query to the correct physical partition(s) via FeedRange.
    • - *
  • {@code batchFilter} — the EPK sub-range within the partition that identifies which - * PKs belong to this batch. PKs whose EPK falls within [filterMin, filterMax) are - * part of this batch.
    • - *
+ * The {@code batchFilter} is the half-open EPK range {@code [minInclusive, maxExclusive)} + * containing the EPKs of all PKs assigned to the batch. It is the only piece of routing + * data persisted in the continuation token; the physical-partition scope used as the + * query FeedRange is rederived at execution time from the current PartitionKeyRange + * cache (union of overlapping partition-key-range EPK ranges) so partition splits do not + * cause stale routing information to be embedded in the token. */ public static final class BatchDefinition { - private final Range partitionScope; private final Range batchFilter; - public BatchDefinition(Range partitionScope, Range batchFilter) { - this.partitionScope = checkNotNull(partitionScope, "Argument 'partitionScope' must not be null."); + public BatchDefinition(Range batchFilter) { this.batchFilter = checkNotNull(batchFilter, "Argument 'batchFilter' must not be null."); } - public Range getPartitionScope() { - return partitionScope; - } - public Range getBatchFilter() { return batchFilter; } @@ -450,33 +342,25 @@ public Range getBatchFilter() { /** * Compact DTO for JSON serialization of a batch definition. + * Persists only the batch filter; routing scope is rederived at execution time. */ static final class BatchDefinitionDto { - private final EpkRangeDto ps; private final EpkRangeDto bf; @JsonCreator - BatchDefinitionDto( - @JsonProperty("ps") EpkRangeDto ps, - @JsonProperty("bf") EpkRangeDto bf) { - this.ps = ps; + BatchDefinitionDto(@JsonProperty("bf") EpkRangeDto bf) { this.bf = bf; } - @JsonProperty("ps") - EpkRangeDto getPs() { return ps; } - @JsonProperty("bf") EpkRangeDto getBf() { return bf; } static BatchDefinitionDto fromBatchDefinition(BatchDefinition bd) { - return new BatchDefinitionDto( - EpkRangeDto.fromRange(bd.partitionScope), - EpkRangeDto.fromRange(bd.batchFilter)); + return new BatchDefinitionDto(EpkRangeDto.fromRange(bd.batchFilter)); } BatchDefinition toBatchDefinition() { - return new BatchDefinition(ps.toRange(), bf.toRange()); + return new BatchDefinition(bf.toRange()); } } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelper.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelper.java index 39a33a1415a1..7384788920c4 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelper.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelper.java @@ -108,20 +108,11 @@ public static SqlQuerySpec createReadManyByPkQuerySpec( PartitionKeyInternal pkInternal = BridgeInternal.getPartitionKeyInternal(pkValues.get(i)); Object[] pkComponents = pkInternal.toObjectArray(); - // PartitionKey.NONE - generate NOT IS_DEFINED for all PK paths - if (pkComponents == null) { - pkFilter.append("("); - for (int j = 0; j < partitionKeySelectors.size(); j++) { - if (j > 0) { - pkFilter.append(" AND "); - } - pkFilter.append("NOT IS_DEFINED("); - pkFilter.append(tableAlias); - pkFilter.append(partitionKeySelectors.get(j)); - pkFilter.append(")"); - } - pkFilter.append(")"); - } else { + // PartitionKey.NONE for hierarchical keys is impossible at this point: the builder + // (CosmosPartitionKeyHelper.validateNoneHandlingForPartitionKeyComponentCount and + // PartitionKeyBuilder.addNoneValue) reject NONE on multi-path PK definitions before + // reaching this code path. So pkComponents is guaranteed non-null here. + { pkFilter.append("("); for (int j = 0; j < pkComponents.length; j++) { String pkParamName = PK_PARAM_PREFIX + paramCount; diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java index 3604ed885f9c..5abb617808b2 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java @@ -122,8 +122,8 @@ import java.util.EnumSet; import java.util.HashMap; import java.util.HashSet; -import java.util.LinkedHashMap; import java.util.Iterator; +import java.util.LinkedHashMap; import java.util.List; import java.util.Locale; import java.util.Map; @@ -4373,10 +4373,13 @@ public Flux> readManyByPartitionKeys( SqlQuerySpec customQuery, String collectionLink, QueryFeedOperationState state, + int maxConcurrentBatchPrefetch, Class klass) { checkNotNull(partitionKeys, "Argument 'partitionKeys' must not be null."); checkArgument(!partitionKeys.isEmpty(), "Argument 'partitionKeys' must not be empty."); + checkArgument(maxConcurrentBatchPrefetch >= 1, + "Argument 'maxConcurrentBatchPrefetch' must be greater than or equal to 1."); final ScopedDiagnosticsFactory diagnosticsFactory = new ScopedDiagnosticsFactory(this, true); state.registerDiagnosticsFactory( @@ -4398,7 +4401,7 @@ public Flux> readManyByPartitionKeys( return ObservableHelper .fluxInlineIfPossibleAsObs( () -> readManyByPartitionKeys( - partitionKeys, customQuery, collectionLink, state, diagnosticsFactory, klass), + partitionKeys, customQuery, collectionLink, state, diagnosticsFactory, maxConcurrentBatchPrefetch, klass), staleResourceRetryPolicy ) .onErrorMap(throwable -> { @@ -4438,6 +4441,7 @@ private Flux> readManyByPartitionKeys( String collectionLink, QueryFeedOperationState state, ScopedDiagnosticsFactory diagnosticsFactory, + int maxConcurrentBatchPrefetch, Class klass) { String requestContinuation = state.getRequestContinuation(); @@ -4474,7 +4478,11 @@ private Flux> readManyByPartitionKeys( // construction - the token already contains the batch definitions. Only resolve // the routing map and build batches on the very first call (no continuation). if (requestContinuation != null) { - // Continuation path: validate collection/query, then resume from token + // Continuation path: validate collection/query, then resume from token. + // The routing map is still required because the persisted token only carries + // batch FILTER ranges; the FeedRange used at execution time (for low-RU + // routing to a single physical partition) is rederived from the current + // routing map per batch via resolvePartitionScopeFromBatchFilter(). ReadManyByPartitionKeyContinuationToken parsedContinuation = ReadManyByPartitionKeyContinuationToken.deserialize(requestContinuation); if (!collectionRid.equals(parsedContinuation.getCollectionRid())) { @@ -4482,30 +4490,35 @@ private Flux> readManyByPartitionKeys( "Continuation token was created for a different collection (rid mismatch). " + "Expected: " + collectionRid + ", token has: " + parsedContinuation.getCollectionRid())); } - String legacyQueryHash = String.valueOf( - ReadManyByPartitionKeyContinuationToken.computeLegacyQueryHash(customQuery)); - if (!queryHash.equals(parsedContinuation.getQueryHash()) - && !legacyQueryHash.equals(parsedContinuation.getQueryHash())) { + if (!queryHash.equals(parsedContinuation.getQueryHash())) { return Flux.error(new IllegalArgumentException( "Continuation token was created with a different query (hash mismatch). " + "The same query must be used when resuming from a continuation token.")); } - String legacyPartitionKeySetHash = String.valueOf( - ReadManyByPartitionKeyContinuationToken.computeLegacyPartitionKeySetHash( - normalizedPartitionKeys.stream() - .map(normalizedPk -> normalizedPk.effectivePartitionKeyString) - .collect(Collectors.toList()))); - if (!partitionKeySetHash.equals(parsedContinuation.getPartitionKeySetHash()) - && !legacyPartitionKeySetHash.equals(parsedContinuation.getPartitionKeySetHash())) { + if (!partitionKeySetHash.equals(parsedContinuation.getPartitionKeySetHash())) { return Flux.error(new IllegalArgumentException( "Continuation token was created with a different partition-key set (hash mismatch). " + "The same normalized set of partition key values must be used when resuming.")); } - return buildSequentialFluxFromContinuation( - parsedContinuation, normalizedPartitionKeys, customQuery, pkDefinition, - resourceLink, state, diagnosticsFactory, klass, - collectionRid, queryHash, partitionKeySetHash); + Mono> resumeRoutingMapMono = partitionKeyRangeCache + .tryLookupAsync( + BridgeInternal.getMetaDataDiagnosticContext(request.requestContext.cosmosDiagnostics), + collection.getResourceId(), + null, + null); + + return resumeRoutingMapMono.flatMapMany(resumeRoutingMapHolder -> { + CollectionRoutingMap resumeRoutingMap = resumeRoutingMapHolder.v; + if (resumeRoutingMap == null) { + return Flux.error(new IllegalStateException( + "Failed to get routing map for readManyByPartitionKeys continuation.")); + } + return buildSequentialFluxFromContinuation( + parsedContinuation, normalizedPartitionKeys, customQuery, pkDefinition, + resumeRoutingMap, resourceLink, state, diagnosticsFactory, klass, + collectionRid, queryHash, partitionKeySetHash, maxConcurrentBatchPrefetch); + }); } // First-call path: validate custom query, resolve routing map, build batches @@ -4535,7 +4548,7 @@ private Flux> readManyByPartitionKeys( return buildSequentialFluxFromScratch( normalizedPartitionKeys, customQuery, pkDefinition, routingMap, resourceLink, state, diagnosticsFactory, klass, - collectionRid, queryHash, partitionKeySetHash); + collectionRid, queryHash, partitionKeySetHash, maxConcurrentBatchPrefetch); }); }); } @@ -4585,7 +4598,8 @@ private Flux> buildSequentialFluxFromScratch( Class klass, String collectionRid, String queryHash, - String partitionKeySetHash) { + String partitionKeySetHash, + int maxConcurrentBatchPrefetch) { Map> partitionRangePkMap = groupPartitionKeysByPhysicalPartition(normalizedPartitionKeys, pkDefinition, routingMap); @@ -4649,7 +4663,7 @@ private Flux> buildSequentialFluxFromScratch( return buildSequentialBatchFlux( allBatches, null, resourceLink, state, diagnosticsFactory, klass, - collectionRid, queryHash, partitionKeySetHash); + collectionRid, queryHash, partitionKeySetHash, maxConcurrentBatchPrefetch); } /** @@ -4662,13 +4676,15 @@ private Flux> buildSequentialFluxFromContinuation( List normalizedPartitionKeys, SqlQuerySpec customQuery, PartitionKeyDefinition pkDefinition, + CollectionRoutingMap routingMap, String resourceLink, QueryFeedOperationState state, ScopedDiagnosticsFactory diagnosticsFactory, Class klass, String collectionRid, String queryHash, - String partitionKeySetHash) { + String partitionKeySetHash, + int maxConcurrentBatchPrefetch) { List partitionKeySelectors = PartitionKeyQueryHelper.createPkSelectors(pkDefinition); @@ -4697,7 +4713,6 @@ private Flux> buildSequentialFluxFromContinuation( List allBatches = new ArrayList<>(); for (ReadManyByPartitionKeyContinuationToken.BatchDefinition batchDef : allBatchDefs) { - Range partitionScope = batchDef.getPartitionScope(); Range batchFilter = batchDef.getBatchFilter(); List batchPks = filterPartitionKeysByEpkRange( @@ -4707,6 +4722,13 @@ private Flux> buildSequentialFluxFromContinuation( continue; } + // Rederive the routing scope from the CURRENT routing map. After a partition + // split this naturally yields the new (potentially smaller) physical-partition + // boundaries that exactly cover this batch filter, keeping query RUs minimal. + // If the cache is briefly stale right after a split, the SDK's stale-resource + // retry will refresh it and rerun, so any RU-cost elevation is bounded. + Range partitionScope = resolvePartitionScopeFromBatchFilter(batchFilter, routingMap); + SqlQuerySpec querySpec = ReadManyByPartitionKeyQueryHelper .createReadManyByPkQuerySpec( baseQueryText, baseParameters, batchPks, @@ -4722,7 +4744,7 @@ private Flux> buildSequentialFluxFromContinuation( return buildSequentialBatchFlux( allBatches, initialBackendContinuation, resourceLink, state, diagnosticsFactory, klass, - collectionRid, queryHash, partitionKeySetHash); + collectionRid, queryHash, partitionKeySetHash, maxConcurrentBatchPrefetch); } /** @@ -4754,7 +4776,8 @@ private Flux> buildSequentialBatchFlux( Class klass, String collectionRid, String queryHash, - String partitionKeySetHash) { + String partitionKeySetHash, + int maxConcurrentBatchPrefetch) { List>> sequentialFluxes = new ArrayList<>(); for (int i = 0; i < allBatches.size(); i++) { @@ -4765,9 +4788,10 @@ private Flux> buildSequentialBatchFlux( final List remainingAfterThis = new ArrayList<>(); for (int j = batchIndex + 1; j < allBatches.size(); j++) { BatchDescriptor remaining = allBatches.get(j); + // Only the batch FILTER is persisted; the routing scope is rederived per batch + // at resume time from the live routing-map cache. remainingAfterThis.add( - new ReadManyByPartitionKeyContinuationToken.BatchDefinition( - remaining.partitionScope, remaining.batchFilter)); + new ReadManyByPartitionKeyContinuationToken.BatchDefinition(remaining.batchFilter)); } CosmosQueryRequestOptions batchQueryOptions = queryOptionsAccessor() @@ -4791,8 +4815,7 @@ private Flux> buildSequentialBatchFlux( new AtomicBoolean(false)); final ReadManyByPartitionKeyContinuationToken.BatchDefinition currentBatchDef = - new ReadManyByPartitionKeyContinuationToken.BatchDefinition( - bd.partitionScope, bd.batchFilter); + new ReadManyByPartitionKeyContinuationToken.BatchDefinition(bd.batchFilter); Flux> stampedFlux = batchFlux.map(feedResponse -> { String backendCont = feedResponse.getContinuationToken(); @@ -4827,10 +4850,50 @@ private Flux> buildSequentialBatchFlux( sequentialFluxes.add(stampedFlux); } - int fluxConcurrency = Math.max(1, Math.min(Configs.getCPUCnt(), sequentialFluxes.size())); + int fluxConcurrency = Math.max(1, Math.min(maxConcurrentBatchPrefetch, sequentialFluxes.size())); return Flux.mergeSequential(sequentialFluxes, fluxConcurrency, 1); } + /** + * Computes the EPK routing scope for a batch from its filter range and the current + * collection routing map. Looks up all partition-key ranges that overlap the batch + * filter and returns {@code [min(minInclusive), max(maxExclusive))} across them so + * the FeedRange set on the per-batch query options exactly aligns with one or more + * physical-partition boundaries (yielding minimal query RU cost). + *

+ * If a partition split happened between the time the continuation token was emitted + * and resume time, the rederived scope reflects the post-split boundaries; the SDK's + * stale-resource retry refreshes the cache promptly if the local cache snapshot was + * still pre-split. Either way no stale boundary is encoded in the token itself. + * + * @throws IllegalStateException if the routing map has no overlap for the batch filter. + */ + private static Range resolvePartitionScopeFromBatchFilter( + Range batchFilter, + CollectionRoutingMap routingMap) { + + List overlapping = routingMap.getOverlappingRanges(batchFilter); + if (overlapping == null || overlapping.isEmpty()) { + throw new IllegalStateException( + "Routing map returned no overlapping partition key ranges for batch filter " + + batchFilter + "."); + } + + String minInclusive = overlapping.get(0).getMinInclusive(); + String maxExclusive = overlapping.get(0).getMaxExclusive(); + for (int i = 1; i < overlapping.size(); i++) { + PartitionKeyRange r = overlapping.get(i); + if (r.getMinInclusive().compareTo(minInclusive) < 0) { + minInclusive = r.getMinInclusive(); + } + if (r.getMaxExclusive().compareTo(maxExclusive) > 0) { + maxExclusive = r.getMaxExclusive(); + } + } + + return new Range<>(minInclusive, maxExclusive, true, false); + } + private static final class NormalizedPartitionKey { final PartitionKey partitionKey; final PartitionKeyInternal effectivePkInternal; @@ -4848,18 +4911,22 @@ private NormalizedPartitionKey( } /** - * Descriptor for a single batch in readManyByPartitionKeys. + * Descriptor for a single batch during execution of readManyByPartitionKeys. *

* Each batch carries two EPK ranges: *

    - *
  • {@code partitionScope} — the physical partition's full EPK range, used for - * FeedRange-based routing (set on CosmosQueryRequestOptions). Survives splits - * because FeedRangeEpkImpl transparently fans out.
    • + *
  • {@code partitionScope} — the physical partition's EPK range, used as the + * FeedRange on CosmosQueryRequestOptions so the backend query request is sent + * to the matching physical partition with minimal RU overhead. This value is + * not persisted in the continuation token; it is derived per + * batch from the current routing map (via + * {@link #resolvePartitionScopeFromBatchFilter}) so post-split refreshes are + * reflected automatically.
    • *
  • {@code batchFilter} — the EPK sub-range covering only the PKs in this batch. - * When a physical partition has more PKs than maxBatchSize, multiple batches share - * the same partitionScope but have distinct batchFilter ranges. Used to - * reconstruct the correct PK set per batch when resuming from a continuation - * token.
    • + * This IS persisted in the continuation token. When a physical partition has + * more PKs than maxBatchSize, multiple batches share the same partitionScope + * at execution time but have distinct batchFilter ranges. Used to reconstruct + * the correct PK set per batch when resuming. *
*/ private static final class BatchDescriptor { diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/CosmosReadManyByPartitionKeyRequestOptions.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/CosmosReadManyByPartitionKeyRequestOptions.java new file mode 100644 index 000000000000..9a02eeb1259d --- /dev/null +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/CosmosReadManyByPartitionKeyRequestOptions.java @@ -0,0 +1,402 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.azure.cosmos.models; + +import com.azure.cosmos.CosmosDiagnosticsThresholds; +import com.azure.cosmos.CosmosEndToEndOperationLatencyPolicyConfig; +import com.azure.cosmos.CosmosItemSerializer; +import com.azure.cosmos.ReadConsistencyStrategy; +import com.azure.cosmos.implementation.CosmosQueryRequestOptionsBase; +import com.azure.cosmos.implementation.CosmosReadManyByPartitionKeyRequestOptionsImpl; +import com.azure.cosmos.implementation.ImplementationBridgeHelpers; +import com.azure.cosmos.util.Beta; + +import java.time.Duration; +import java.util.List; +import java.util.Set; + +/** + * Specifies the options associated with the {@code readManyByPartitionKeys} operation + * in the Azure Cosmos DB database service. + *

+ * This is distinct from {@link CosmosReadManyRequestOptions} (used by the + * {@code readMany(List<CosmosItemIdentity>)} API). It exposes only the knobs that are + * applicable to {@code readManyByPartitionKeys} — for example, properties that influence + * query parallelism inside a single physical partition or a feed range filter are + * intentionally not exposed because the operation is fully managed by the SDK. + */ +public final class CosmosReadManyByPartitionKeyRequestOptions { + private final CosmosReadManyByPartitionKeyRequestOptionsImpl actualRequestOptions; + + /** + * Instantiates a new readManyByPartitionKeys request options. + */ + public CosmosReadManyByPartitionKeyRequestOptions() { + this.actualRequestOptions = new CosmosReadManyByPartitionKeyRequestOptionsImpl(); + } + + /** + * Copy constructor. + * + * @param options the options to copy. + */ + CosmosReadManyByPartitionKeyRequestOptions(CosmosReadManyByPartitionKeyRequestOptions options) { + this.actualRequestOptions = new CosmosReadManyByPartitionKeyRequestOptionsImpl(options.actualRequestOptions); + } + + /** + * Gets the composite continuation token used to resume a previous + * {@code readManyByPartitionKeys} invocation. + * + * @return the continuation token, or null if not set. + */ + public String getContinuationToken() { + return this.actualRequestOptions.getContinuationToken(); + } + + /** + * Sets the composite continuation token used to resume a previous + * {@code readManyByPartitionKeys} invocation. The token must have been returned by a prior + * invocation of {@code readManyByPartitionKeys} on the same container, with the same + * partition-key set and the same custom query. + * + * @param continuationToken the composite continuation token from a previous invocation. + * @return the {@link CosmosReadManyByPartitionKeyRequestOptions} for fluent chaining. + */ + public CosmosReadManyByPartitionKeyRequestOptions setContinuationToken(String continuationToken) { + this.actualRequestOptions.setContinuationToken(continuationToken); + return this; + } + + /** + * Gets the maximum number of per-physical-partition batches whose first page is prefetched + * concurrently. This bounds the prefetch parallelism the SDK uses while sequentially + * draining batches. + * + * @return the max concurrent batch prefetch, or null if the SDK default is in effect. + */ + public Integer getMaxConcurrentBatchPrefetch() { + return this.actualRequestOptions.getMaxConcurrentBatchPrefetch(); + } + + /** + * Sets the maximum number of per-physical-partition batches whose first page is prefetched + * concurrently. The default is {@code Runtime.getRuntime().availableProcessors()}. + *

+ * Increase this to trade memory for lower end-to-end latency on wide containers; decrease it + * (e.g. to {@code 1}) when running in environments where a single task already saturates the + * network/CPU and additional prefetch only adds memory pressure. + * + * @param maxConcurrentBatchPrefetch the max concurrent batch prefetch (must be >= 1). + * @return the {@link CosmosReadManyByPartitionKeyRequestOptions} for fluent chaining. + * @throws IllegalArgumentException if {@code maxConcurrentBatchPrefetch} is < 1. + */ + public CosmosReadManyByPartitionKeyRequestOptions setMaxConcurrentBatchPrefetch(int maxConcurrentBatchPrefetch) { + if (maxConcurrentBatchPrefetch < 1) { + throw new IllegalArgumentException( + "Argument 'maxConcurrentBatchPrefetch' must be greater than or equal to 1."); + } + this.actualRequestOptions.setMaxConcurrentBatchPrefetch(maxConcurrentBatchPrefetch); + return this; + } + + /** + * Gets the read consistency strategy for the request. + * + * @return the read consistency strategy. + */ + @Beta(value = Beta.SinceVersion.V4_69_0, warningText = Beta.PREVIEW_SUBJECT_TO_CHANGE_WARNING) + public ReadConsistencyStrategy getReadConsistencyStrategy() { + return this.actualRequestOptions.getReadConsistencyStrategy(); + } + + /** + * Sets the read consistency strategy required for the request. + * + * @param readConsistencyStrategy the read consistency strategy. + * @return the {@link CosmosReadManyByPartitionKeyRequestOptions} for fluent chaining. + */ + @Beta(value = Beta.SinceVersion.V4_69_0, warningText = Beta.PREVIEW_SUBJECT_TO_CHANGE_WARNING) + public CosmosReadManyByPartitionKeyRequestOptions setReadConsistencyStrategy( + ReadConsistencyStrategy readConsistencyStrategy) { + this.actualRequestOptions.setReadConsistencyStrategy(readConsistencyStrategy); + return this; + } + + /** + * Gets the session token for use with session consistency. + * + * @return the session token. + */ + public String getSessionToken() { + return this.actualRequestOptions.getSessionToken(); + } + + /** + * Sets the session token for use with session consistency. + * + * @param sessionToken the session token. + * @return the {@link CosmosReadManyByPartitionKeyRequestOptions} for fluent chaining. + */ + public CosmosReadManyByPartitionKeyRequestOptions setSessionToken(String sessionToken) { + this.actualRequestOptions.setSessionToken(sessionToken); + return this; + } + + /** + * Sets the maximum size (in kilobytes) of the backend continuation token embedded inside the + * composite {@code readManyByPartitionKeys} continuation token. + *

+ * Note: this only constrains the per-batch backend continuation that the SDK wraps inside + * the public composite token; the public composite token itself is always larger because it + * also carries the remaining batch definitions, query hash, and partition-key-set hash. + * + * @param limitInKb backend continuation token size limit (must be >= 1). + * @return the {@link CosmosReadManyByPartitionKeyRequestOptions} for fluent chaining. + */ + public CosmosReadManyByPartitionKeyRequestOptions setMaxBackendContinuationTokenSizeInKb(int limitInKb) { + this.actualRequestOptions.setResponseContinuationTokenLimitInKb(limitInKb); + return this; + } + + /** + * Gets the maximum size (in kilobytes) of the backend continuation token embedded inside the + * composite {@code readManyByPartitionKeys} continuation token. Returns 0 if not set. + * + * @return the configured backend continuation token size limit, or 0 if not set. + */ + public int getMaxBackendContinuationTokenSizeInKb() { + return this.actualRequestOptions.getResponseContinuationTokenLimitInKb(); + } + + /** + * Gets the maximum number of items returned in a single page. + * + * @return the max item count, or null if not set (the SDK default applies). + */ + public Integer getMaxItemCount() { + return this.actualRequestOptions.getMaxItemCount(); + } + + /** + * Sets the maximum number of items returned in a single page. + * + * @param maxItemCount the maximum number of items per page. + * @return the {@link CosmosReadManyByPartitionKeyRequestOptions} for fluent chaining. + */ + public CosmosReadManyByPartitionKeyRequestOptions setMaxItemCount(int maxItemCount) { + this.actualRequestOptions.setMaxItemCount(maxItemCount); + return this; + } + + /** + * Sets the {@link CosmosEndToEndOperationLatencyPolicyConfig} to be used for the request. + * + * @param cosmosEndToEndOperationLatencyPolicyConfig the {@link CosmosEndToEndOperationLatencyPolicyConfig} + * @return the {@link CosmosReadManyByPartitionKeyRequestOptions} for fluent chaining. + */ + public CosmosReadManyByPartitionKeyRequestOptions setCosmosEndToEndOperationLatencyPolicyConfig( + CosmosEndToEndOperationLatencyPolicyConfig cosmosEndToEndOperationLatencyPolicyConfig) { + this.actualRequestOptions + .setCosmosEndToEndOperationLatencyPolicyConfig(cosmosEndToEndOperationLatencyPolicyConfig); + return this; + } + + /** + * List of regions to be excluded for the request/retries. + * + * @param excludeRegions the regions to exclude + * @return the {@link CosmosReadManyByPartitionKeyRequestOptions} + */ + public CosmosReadManyByPartitionKeyRequestOptions setExcludedRegions(List excludeRegions) { + this.actualRequestOptions.setExcludedRegions(excludeRegions); + return this; + } + + /** + * Gets the list of regions to exclude for the request/retries. + * + * @return the list of excluded regions + */ + public List getExcludedRegions() { + return this.actualRequestOptions.getExcludedRegions(); + } + + /** + * Gets the option to enable populate query metrics. By default query metrics are enabled. + * + * @return whether query metrics are enabled + */ + public boolean isQueryMetricsEnabled() { + return this.actualRequestOptions.isQueryMetricsEnabled(); + } + + /** + * Sets the option to enable/disable query metrics. + * + * @param queryMetricsEnabled whether to enable or disable query metrics + * @return the {@link CosmosReadManyByPartitionKeyRequestOptions} for fluent chaining. + */ + public CosmosReadManyByPartitionKeyRequestOptions setQueryMetricsEnabled(boolean queryMetricsEnabled) { + this.actualRequestOptions.setQueryMetricsEnabled(queryMetricsEnabled); + return this; + } + + /** + * Gets the throughput control group name. + * + * @return the throughput control group name. + */ + public String getThroughputControlGroupName() { + return this.actualRequestOptions.getThroughputControlGroupName(); + } + + /** + * Sets the throughput control group name. + * + * @param throughputControlGroupName the throughput control group name. + * @return the {@link CosmosReadManyByPartitionKeyRequestOptions} for fluent chaining. + */ + public CosmosReadManyByPartitionKeyRequestOptions setThroughputControlGroupName(String throughputControlGroupName) { + this.actualRequestOptions.setThroughputControlGroupName(throughputControlGroupName); + return this; + } + + /** + * Gets the dedicated gateway request options. + * + * @return the dedicated gateway request options. + */ + public DedicatedGatewayRequestOptions getDedicatedGatewayRequestOptions() { + return this.actualRequestOptions.getDedicatedGatewayRequestOptions(); + } + + /** + * Sets the dedicated gateway request options. + * + * @param dedicatedGatewayRequestOptions the dedicated gateway request options. + * @return the {@link CosmosReadManyByPartitionKeyRequestOptions} for fluent chaining. + */ + public CosmosReadManyByPartitionKeyRequestOptions setDedicatedGatewayRequestOptions( + DedicatedGatewayRequestOptions dedicatedGatewayRequestOptions) { + this.actualRequestOptions.setDedicatedGatewayRequestOptions(dedicatedGatewayRequestOptions); + return this; + } + + /** + * Gets the latency threshold for diagnostics on tracer. + * + * @return the latency threshold for diagnostics on tracer. + */ + public Duration getThresholdForDiagnosticsOnTracer() { + return this.actualRequestOptions.getThresholdForDiagnosticsOnTracer(); + } + + /** + * Sets the latency threshold for diagnostics on tracer. + * + * @param thresholdForDiagnosticsOnTracer the latency threshold for diagnostics on tracer. + * @return the {@link CosmosReadManyByPartitionKeyRequestOptions} for fluent chaining. + */ + public CosmosReadManyByPartitionKeyRequestOptions setThresholdForDiagnosticsOnTracer( + Duration thresholdForDiagnosticsOnTracer) { + this.actualRequestOptions.setThresholdForDiagnosticsOnTracer(thresholdForDiagnosticsOnTracer); + return this; + } + + /** + * Allows overriding the diagnostic thresholds for a specific operation. + * + * @param operationSpecificThresholds the diagnostic threshold override for this operation + * @return the {@link CosmosReadManyByPartitionKeyRequestOptions} for fluent chaining. + */ + public CosmosReadManyByPartitionKeyRequestOptions setDiagnosticsThresholds( + CosmosDiagnosticsThresholds operationSpecificThresholds) { + this.actualRequestOptions.setDiagnosticsThresholds(operationSpecificThresholds); + return this; + } + + /** + * Gets the diagnostic thresholds used as an override for a specific operation. + * + * @return the diagnostic thresholds for this operation. + */ + public CosmosDiagnosticsThresholds getDiagnosticsThresholds() { + return this.actualRequestOptions.getDiagnosticsThresholds(); + } + + /** + * Gets the custom item serializer defined for this instance of request options. + * + * @return the custom item serializer. + */ + public CosmosItemSerializer getCustomItemSerializer() { + return this.actualRequestOptions.getCustomItemSerializer(); + } + + /** + * Sets a custom item serializer to be used for this operation. + * + * @param customItemSerializer the custom item serializer for this operation. + * @return the {@link CosmosReadManyByPartitionKeyRequestOptions} for fluent chaining. + */ + public CosmosReadManyByPartitionKeyRequestOptions setCustomItemSerializer( + CosmosItemSerializer customItemSerializer) { + this.actualRequestOptions.setCustomItemSerializer(customItemSerializer); + return this; + } + + /** + * Sets the custom keyword identifiers. + * + * @param keywordIdentifiers the custom keyword identifiers. + * @return the {@link CosmosReadManyByPartitionKeyRequestOptions} for fluent chaining. + */ + public CosmosReadManyByPartitionKeyRequestOptions setKeywordIdentifiers(Set keywordIdentifiers) { + this.actualRequestOptions.setKeywordIdentifiers(keywordIdentifiers); + return this; + } + + /** + * Gets the custom keyword identifiers. + * + * @return the custom keyword identifiers. + */ + public Set getKeywordIdentifiers() { + return this.actualRequestOptions.getKeywordIdentifiers(); + } + + CosmosQueryRequestOptionsBase getImpl() { + return this.actualRequestOptions; + } + + /////////////////////////////////////////////////////////////////////////////////////////// + // the following helper/accessor only helps to access this class outside of this package.// + /////////////////////////////////////////////////////////////////////////////////////////// + static void initialize() { + ImplementationBridgeHelpers.CosmosReadManyByPartitionKeyRequestOptionsHelper + .setCosmosReadManyByPartitionKeyRequestOptionsAccessor( + new ImplementationBridgeHelpers.CosmosReadManyByPartitionKeyRequestOptionsHelper + .CosmosReadManyByPartitionKeyRequestOptionsAccessor() { + @Override + public CosmosQueryRequestOptionsBase getImpl( + CosmosReadManyByPartitionKeyRequestOptions options) { + return options.actualRequestOptions; + } + + @Override + public String getContinuationToken(CosmosReadManyByPartitionKeyRequestOptions options) { + return options.actualRequestOptions.getContinuationToken(); + } + + @Override + public Integer getMaxConcurrentBatchPrefetch( + CosmosReadManyByPartitionKeyRequestOptions options) { + return options.actualRequestOptions.getMaxConcurrentBatchPrefetch(); + } + }); + } + + static { initialize(); } +} diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/CosmosReadManyRequestOptions.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/CosmosReadManyRequestOptions.java index e428a676e151..f6e570258042 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/CosmosReadManyRequestOptions.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/CosmosReadManyRequestOptions.java @@ -380,11 +380,6 @@ static void initialize() { public CosmosQueryRequestOptionsBase getImpl(CosmosReadManyRequestOptions options) { return options.actualRequestOptions; } - - @Override - public String getRequestContinuation(CosmosReadManyRequestOptions options) { - return options.actualRequestOptions.getRequestContinuation(); - } }); } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/util/Beta.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/util/Beta.java index 5ccd99db2fd0..53072b09326d 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/util/Beta.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/util/Beta.java @@ -107,6 +107,8 @@ public enum SinceVersion { /** v4.71.0 */ V4_71_0, /** v4.74.0 */ - V4_74_0 + V4_74_0, + /** v4.75.0 */ + V4_75_0 } } diff --git a/sdk/cosmos/cspell.yaml b/sdk/cosmos/cspell.yaml index 94a4002c2c9c..5cbd5ca804e7 100644 --- a/sdk/cosmos/cspell.yaml +++ b/sdk/cosmos/cspell.yaml @@ -4,3 +4,4 @@ overrides: - filename: "**/sdk/cosmos/*" words: - DCOUNT + - dedupe diff --git a/sdk/cosmos/docs/readManyByPartitionKey-design.md b/sdk/cosmos/docs/readManyByPartitionKey-design.md index f0df3bce99bd..17fe086e4018 100644 --- a/sdk/cosmos/docs/readManyByPartitionKey-design.md +++ b/sdk/cosmos/docs/readManyByPartitionKey-design.md @@ -1,32 +1,56 @@ -# readManyByPartitionKeys — Design & Implementation +# readManyByPartitionKeys - Design & Implementation ## Overview New `readManyByPartitionKeys` methods on `CosmosAsyncContainer` / `CosmosContainer` that accept a -`List` (without item-id). The SDK splits the PK values by physical -partition, generates batched streaming queries per physical partition, and returns results as -`CosmosPagedFlux` / `CosmosPagedIterable`. +`List` (without item ids). The SDK splits the PK values by physical partition, +generates batched streaming queries per physical partition, and returns results as +`CosmosPagedFlux` / `CosmosPagedIterable`. The result stream is strictly batch-by-batch +ordered (see Phase 4) so each `FeedResponse` carries a usable composite continuation token. An optional `SqlQuerySpec` parameter lets callers supply a custom query for projections -and additional filters. The SDK appends the auto-generated PK WHERE clause to it. +and additional filters. The SDK appends the auto-generated PK WHERE clause to it and rejects +non-streamable shapes (aggregates / ORDER BY / DISTINCT / etc.) up front via a one-time +gateway query-plan validation. + +## High-level flow + +```mermaid +flowchart TD + A["Caller
readManyByPartitionKeys(pks, ?customQuery, ?options)"] --> B{"Continuation
token in options?"} + B -- no --> C[Resolve collection + PK definition] + C --> D[Normalize PKs
dedupe by EPK + sort] + D --> E[Optional: validate custom query
via gateway query plan] + E --> F[Look up routing map
PartitionKeyRangeCache] + F --> G[Group PKs per physical partition,
split into batches of N PKs,
compute batchFilter EPK range] + G --> H[Sort all batches by
batchFilter.minInclusive] + H --> I[mergeSequential
execution] + B -- yes --> J[Deserialize composite token] + J --> K[Validate version + collectionRid
+ queryHash + partitionKeySetHash] + K --> L[Look up routing map
PartitionKeyRangeCache] + L --> M[For each persisted batchFilter:
filter PKs into batch +
resolve partitionScope from cache] + M --> I + I --> N["FeedResponse stream
(each stamped with composite token)"] +``` ## Decisions | Topic | Decision | |-------------------------|----------------------------------------------------------------------------------------------------------------------------------------------------------------------------| -| API name | `readManyByPartitionKeys` — distinct name to avoid ambiguity with existing `readMany(List)` | +| API name | `readManyByPartitionKeys` - distinct name to avoid ambiguity with existing `readMany(List)` | | Return type | `CosmosPagedFlux` (async) / `CosmosPagedIterable` (sync) | -| Custom query format | `SqlQuerySpec` — full query with parameters; SDK ANDs the PK filter | +| Custom query format | `SqlQuerySpec` - full query with parameters; SDK ANDs the PK filter | | Partial HPK | Supported from the start; prefix PKs fan out via `getOverlappingRanges` | | PK deduplication | SDK normalizes and deduplicates the partition-key set by EPK before batching; Spark callers should still dedupe the input DataFrame when practical for efficiency | | Spark UDF | New `GetCosmosPartitionKeyValue` UDF | | Custom query validation | Gateway query plan via the standard SDK query-plan retrieval path; reject aggregates/ORDER BY/DISTINCT/GROUP BY/DCount/OFFSET/LIMIT/non-streaming ORDER BY/vector/fulltext | | PK list size | No hard upper-bound enforced; SDK batches internally per physical partition (default 100 PKs per batch, configurable via `COSMOS.READ_MANY_BY_PK_MAX_BATCH_SIZE`) | | Eager validation | Null and empty PK list rejected eagerly (not lazily in reactive chain) | -| Telemetry | Separate span name `readManyByPartitionKeyItems.` (distinct from existing `readManyItems`) | +| Telemetry | Separate span name `readManyByPartitionKeys.` (distinct from existing `readManyItems`) | | Query construction | Table alias auto-detected from FROM clause; string literals and subqueries handled correctly | +| Continuation tokens | Composite token persists only the per-batch EPK FILTER; the routing scope (`FeedRange`) is rederived from the live `PartitionKeyRangeCache` per batch on resume so post-split refreshes are reflected automatically | -## Phase 1 — SDK Core (`azure-cosmos`) +## Phase 1 - SDK Core (`azure-cosmos`) ### Step 1: New public overloads in CosmosAsyncContainer @@ -37,41 +61,86 @@ and additional filters. The SDK appends the auto-generated PK WHERE clause to it Class classType) CosmosPagedFlux readManyByPartitionKeys(List partitionKeys, SqlQuerySpec customQuery, - CosmosReadManyRequestOptions requestOptions, + Class classType) + CosmosPagedFlux readManyByPartitionKeys(List partitionKeys, + SqlQuerySpec customQuery, + CosmosReadManyByPartitionKeyRequestOptions requestOptions, Class classType) ``` -All delegate to a private `readManyByPartitionKeyInternalFunc(...)`. +All four overloads delegate to a private `readManyByPartitionKeyInternalFunc(...)`. + +> **Note:** the request-options type for `readManyByPartitionKeys` is the dedicated +> `CosmosReadManyByPartitionKeyRequestOptions` class (not the shared +> `CosmosReadManyRequestOptions` used by `readMany(List)`). It exposes +> `setContinuationToken(String)` / `getContinuationToken()` directly on the public surface +> and a `setMaxConcurrentBatchPrefetch(int)` knob that defaults to +> `Runtime.getRuntime().availableProcessors()` (the Spark connector overrides this default to +> `1` and exposes a config key `spark.cosmos.read.readManyByPk.maxConcurrentBatchPrefetch` for +> per-task tuning). -**Eager validation:** The 4-arg method validates `partitionKeys` is non-null and non-empty before constructing the reactive pipeline, throwing `IllegalArgumentException` synchronously. +**Eager validation:** The 4-arg method validates `partitionKeys` is non-null and non-empty before +constructing the reactive pipeline, throwing `IllegalArgumentException` synchronously. ### Step 2: Sync wrappers in CosmosContainer -Same signatures returning `CosmosPagedIterable`, delegating to the async container. +Same four signatures returning `CosmosPagedIterable`, each delegating to the async container. ### Step 3: Internal orchestration (RxDocumentClientImpl) +#### Step 3a: First-call path (no continuation token) + 1. Resolve collection metadata + PK definition from cache. -2. Fetch routing map from `partitionKeyRangeCache`. -3. For each `PartitionKey`: - - Compute effective partition key (EPK). - - Full PK → `getRangeByEffectivePartitionKey()` (single range). - - Partial HPK → compute EPK prefix range → `getOverlappingRanges()` (multiple ranges). - **Note:** partial HPK intentionally fans out to multiple physical partitions. -4. Group PK values by `PartitionKeyRange`. -5. Per physical partition → split PKs into batches of `maxPksPerPartitionQuery` (configurable, default 100). -6. Per batch → build `SqlQuerySpec` with PK WHERE clause (Step 5). -7. Interleave batches across physical partitions in round-robin order so that bounded concurrency prefers different physical partitions over sequential batches of the same partition. -8. Execute queries via `queryForReadMany()` with bounded concurrency (`Math.min(batchCount, cpuCount)`). -9. Return results as `CosmosPagedFlux`. +2. Normalize the input PK list: compute the effective partition key (EPK) for each input, + deduplicate by EPK string, sort the surviving entries by EPK. +3. Optionally validate the custom query (Step 4) - one gateway query-plan call, executed + in parallel with the routing-map lookup. +4. Fetch the routing map from `partitionKeyRangeCache`. +5. For each `PartitionKeyRange` (physical partition): + - Bucket the normalized PKs whose EPK falls in the partition's range. + - Split into batches of `Configs.getReadManyByPkMaxBatchSize()` PKs (default 100). + - For each batch, compute the **batchFilter** EPK range: + `[epkOf(firstPkInBatch), epkOf(firstPkInNextBatch))` - or, for the last batch in a + partition, `[epkOf(firstPkInBatch), partitionScope.maxExclusive)`. + - The batch's **partitionScope** is the partition's EPK range itself. + - Build the per-batch `SqlQuerySpec` (Step 5). +6. Sort all batches across all partitions by `batchFilter.minInclusive` for deterministic + global EPK ordering. +7. Execute sequentially with `Flux.mergeSequential(sources, maxConcurrency, prefetch=1)` + so emission order is strictly batch-by-batch (items from batch *N+1* are never interleaved + with items from batch *N*) while still allowing eager subscription to the next batch + (small Reactor prefetch) for throughput. `maxConcurrency = Math.min(batchCount, cpuCount)`. +8. As each `FeedResponse` flows through, stamp it with a composite continuation token (Phase 4). + +#### Step 3b: Continuation path (token present) + +1. Resolve collection metadata + PK definition from cache (same as 3a step 1). +2. Normalize the input PK list (same as 3a step 2). +3. Deserialize the composite token; validate `version`, `collectionRid`, `queryHash`, + `partitionKeySetHash`. Any mismatch throws `IllegalArgumentException` eagerly inside + the reactive chain. +4. Fetch the routing map from `partitionKeyRangeCache` - required to rederive each batch's + `partitionScope` for FeedRange routing. +5. For each persisted `batchFilter` (current + remaining): + - Filter the normalized PK list down to PKs whose EPK falls inside `batchFilter`. + Skip the batch if the filter happens to match no PKs (e.g. caller dropped some inputs). + - Compute the routing scope by calling `resolvePartitionScopeFromBatchFilter(batchFilter, routingMap)`: + `[min(minEpk), max(maxEpk))` across all `PartitionKeyRange`s overlapping the batch + filter. This naturally returns the post-split boundaries when a split has happened. + - Build the per-batch `SqlQuerySpec`. +6. Stitch the sequence: the first batch starts from the persisted `backendContinuation`; + subsequent batches start fresh. +7. Execute the same sequential merge as the first-call path. ### Step 4: Custom query validation -One-time call per invocation using the same query-plan retrieval path and cacheability rules as regular SDK queries. +One-time call per invocation using the same query-plan retrieval path and cacheability rules as +regular SDK queries. - `QueryPlanRetriever.getQueryPlanThroughGatewayAsync()` for the user query. - Reject (`IllegalArgumentException`) if: - - `queryInfo.hasGroupBy()` — checked first (takes precedence over aggregates since `hasAggregates()` also returns true for GROUP BY queries) + - `queryInfo.hasGroupBy()` - checked first (takes precedence over aggregates since + `hasAggregates()` also returns true for GROUP BY queries) - `queryInfo.hasAggregates()` - `queryInfo.hasOrderBy()` - `queryInfo.hasDistinct()` @@ -82,6 +151,9 @@ One-time call per invocation using the same query-plan retrieval path and cachea - `partitionedQueryExecutionInfo.hasHybridSearchQueryInfo()` - query plan details are unavailable (`queryInfo == null`) +The validator is called in parallel with the routing-map lookup so it does not add a +serial round-trip. + ### Step 5: Query construction Query construction is implemented in `ReadManyByPartitionKeyQueryHelper`. The helper: @@ -89,7 +161,9 @@ Query construction is implemented in `ReadManyByPartitionKeyQueryHelper`. The he - Extracts the table alias from the FROM clause (handles `FROM c`, `FROM root r`, `FROM x WHERE ...`) - Handles string literals in queries (parens/keywords inside `'...'` are correctly skipped) - Recognizes SQL keywords: WHERE, ORDER, GROUP, JOIN, OFFSET, LIMIT, HAVING -- Uses parameterized queries (`@__rmPk_` prefix) to prevent SQL injection +- Uses parameterized queries (`@__rmPk_` prefix) to prevent SQL injection. The prefix is + reserved: any caller-supplied `SqlParameter` whose name starts with `@__rmPk_` is rejected + to prevent collisions with the auto-generated PK parameters. **Single PK (HASH):** @@ -119,117 +193,198 @@ If the base query already has a WHERE clause: ### Step 6: Interface wiring -New method `readManyByPartitionKeys` added directly to `AsyncDocumentClient` interface, implemented in `RxDocumentClientImpl`. New `fetchQueryPlanForValidation` static method added to `DocumentQueryExecutionContextFactory` for custom query validation. +New method `readManyByPartitionKeys` added directly to `AsyncDocumentClient` interface, +implemented in `RxDocumentClientImpl`. The query executor is `createQueryInternal` (the same +internal entry point used by the streaming query API), invoked once per batch. + +A new `fetchQueryPlanForValidation` static method on `DocumentQueryExecutionContextFactory` +exposes the query-plan retrieval path for the custom-query validator. ### Step 7: Configuration -New configurable batch size via system property `COSMOS.READ_MANY_BY_PK_MAX_BATCH_SIZE` or environment variable `COSMOS_READ_MANY_BY_PK_MAX_BATCH_SIZE` (default: 100, minimum: 1). Follows existing `Configs` patterns. +New configurable batch size via system property `COSMOS.READ_MANY_BY_PK_MAX_BATCH_SIZE` or +environment variable `COSMOS_READ_MANY_BY_PK_MAX_BATCH_SIZE` (default: 100, minimum: 1). +Both inputs are parsed via a shared safe-parse helper that logs a WARN and falls back to +the default on `NumberFormatException` or non-positive values. -## Phase 2 — Spark Connector (`azure-cosmos-spark_3`) +## Phase 2 - Spark Connector (`azure-cosmos-spark_3`) -### Step 8: New UDF — `GetCosmosPartitionKeyValue` +### Step 8: New UDF - `GetCosmosPartitionKeyValue` - Input: partition key value (single value or Seq for hierarchical PKs). - Output: serialized PK string in format `pk([...json...])`. -- **Null handling:** Null input is serialized as a JSON-null partition key component. If callers need `PartitionKey.NONE` semantics they must use the schema-matched path with `spark.cosmos.read.readManyByPk.nullHandling=None`, which is only supported for single-path partition keys. +- **Null handling:** Null input is serialized as a JSON-null partition key component. If callers + need `PartitionKey.NONE` semantics they must use the schema-matched path with + `spark.cosmos.read.readManyByPk.nullHandling=None`, which is only supported for single-path + partition keys. ### Step 9: PK-only serialization helper `CosmosPartitionKeyHelper`: -- `getCosmosPartitionKeyValueString(pkValues: List[Object]): String` — serialize to `pk([...])` format. -- `tryParsePartitionKey(serialized: String): Option[PartitionKey]` — deserialize; returns `None` for malformed input including invalid JSON (wrapped in `scala.util.Try`). -- When `spark.cosmos.read.readManyByPk.nullHandling=None` is used, hierarchical partition keys with null components are rejected with a clear error because `PartitionKey.NONE` cannot be used with multiple paths. +- `getCosmosPartitionKeyValueString(pkValues: List[Object]): String` - serialize to `pk([...])` format. +- `tryParsePartitionKey(serialized: String): Option[PartitionKey]` - deserialize; returns `None` + for malformed input including invalid JSON. The parser regex is anchored + (`^pk\((.*)\)$`) so substrings that merely contain a `pk(...)` literal are not accepted. +- When `spark.cosmos.read.readManyByPk.nullHandling=None` is used, hierarchical partition keys + with null components are rejected with a clear error because `PartitionKey.NONE` cannot be + used with multiple paths. ### Step 10: `CosmosItemsDataSource.readManyByPartitionKeys` Static entry points that accept a DataFrame and Cosmos config. PK extraction supports two modes: -1. **UDF-produced column**: DataFrame contains `_partitionKeyIdentity` column (from `GetCosmosPartitionKeyValue` UDF). +1. **UDF-produced column**: DataFrame contains `_partitionKeyIdentity` column (from + `GetCosmosPartitionKeyValue` UDF). 2. **Schema-matched columns**: DataFrame columns match the container's PK paths. -Top-level DataFrame columns may supply a full or prefix hierarchical partition key directly. Nested partition key paths are not resolved automatically and must use the UDF-produced `_partitionKeyIdentity` column. +Top-level DataFrame columns may supply a full or prefix hierarchical partition key directly. +Nested partition key paths are not resolved automatically and must use the UDF-produced +`_partitionKeyIdentity` column. Falls back with `IllegalArgumentException` if neither mode is possible. ### Step 11: `CosmosReadManyByPartitionKeyReader` -Orchestrator that resolves schema, initializes and broadcasts client state to executors, then maps each Spark partition to an `ItemsPartitionReaderWithReadManyByPartitionKey`. The wrapper iterator closes the reader deterministically on exhaustion, on failures, and via Spark task-completion callbacks. +Orchestrator that resolves schema, initializes and broadcasts client state to executors, then maps +each Spark partition to an `ItemsPartitionReaderWithReadManyByPartitionKey`. The wrapper iterator +closes the reader deterministically on exhaustion, on failures, and via Spark task-completion +callbacks. ### Step 12: `ItemsPartitionReaderWithReadManyByPartitionKey` Spark `PartitionReader[InternalRow]` that: -- Preserves the caller's PK list and lets the SDK normalize/dedupe by effective partition key; callers should still dedupe the DataFrame upstream when practical. -- Passes the pre-built `CosmosReadManyRequestOptions` (with throughput control, diagnostics, custom serializer) to the SDK. -- Uses `TransientIOErrorsRetryingIterator` for retry handling. +- Preserves the caller's PK list and lets the SDK normalize/dedupe by effective partition key; + callers should still dedupe the DataFrame upstream when practical. Logs a WARN if a single + partition reader receives more than 200000 PKs (large input materializes the full set in + memory plus the SDK's normalized batch metadata). +- Passes the pre-built request options (with throughput control, diagnostics, custom serializer) + to the SDK. +- Wraps each batch in a `TransientIOErrorsRetryingReadManyByPartitionKeyIterator` for retry + handling - on transient I/O failures the iterator re-creates the underlying `CosmosPagedFlux` + using the continuation token of the last fully-drained page (Phase 4). - Short-circuits empty PK lists to avoid SDK rejection. -## Phase 3 — Testing +## Phase 3 - Testing ### Unit tests - Query construction: single PK, HPK full/partial, custom query composition, table alias detection. -- Query plan rejection: aggregates, ORDER BY, DISTINCT, GROUP BY (with and without aggregates), DCOUNT. +- Query plan rejection: aggregates, ORDER BY, DISTINCT, GROUP BY (with and without aggregates), + DCOUNT. - String literal handling: WHERE/parentheses inside string constants. - Keyword detection: WHERE, ORDER, GROUP, JOIN, OFFSET, LIMIT, HAVING. - PK serialization/deserialization roundtrip (including malformed JSON handling). - `findTopLevelWhereIndex` edge cases: subqueries, string literals, case insensitivity. +- Continuation token: roundtrip, version-field presence, version-mismatch rejection, + partition-key-set hash stability across reordered/duplicate inputs, partition scope is NOT + persisted. +- Reserved parameter prefix: caller-supplied `@__rmPk_*` parameter is rejected. ### Integration tests -- End-to-end SDK: single PK basic, projections, filters, empty results, HPK full/partial, request options propagation. -- Batch size validation: temporarily lowered batch size to exercise batching/interleaving logic. +- End-to-end SDK: single PK basic, projections, filters, empty results, HPK full/partial, + request options propagation. +- Batch size validation: temporarily lowered batch size to exercise batching/sequential-merge + logic. +- Continuation tokens: round-trip resume, resume with reordered PK list (must produce same + results because the partitionKeySetHash is order-invariant), resume rejection on changed + query / changed PK set / different collection. - Null/empty PK list rejection (eager validation). -- Spark connector: `ItemsPartitionReaderWithReadManyByPartitionKey` with known PK values and non-existent PKs. -- Spark public API: nested partition key containers require `_partitionKeyIdentity` and succeed when populated via `GetCosmosPartitionKeyValue`. +- Spark connector: `ItemsPartitionReaderWithReadManyByPartitionKey` with known PK values and + non-existent PKs. +- Spark public API: nested partition key containers require `_partitionKeyIdentity` and succeed + when populated via `GetCosmosPartitionKeyValue`. - `CosmosPartitionKeyHelper`: single/HPK roundtrip, case insensitivity, malformed input. -## Phase 4 — Continuation Token Support for Retry Safety +## Phase 4 - Continuation Token Support for Retry Safety ### Motivation -The Spark `TransientIOErrorsRetryingReadManyByPartitionKeyIterator` replays from a -continuation token on transient I/O failures. For this to work, `readManyByPartitionKeys` -must (a) emit a meaningful continuation token on each `FeedResponse`, and (b) accept a -continuation token that resumes where the previous attempt left off. +The Spark `TransientIOErrorsRetryingReadManyByPartitionKeyIterator` replays from a continuation +token on transient I/O failures. For this to work, `readManyByPartitionKeys` must (a) emit a +meaningful continuation token on each `FeedResponse`, and (b) accept a continuation token that +resumes where the previous attempt left off. -Because the SDK batches PK values across multiple physical partitions, the continuation -token must capture not only the backend query continuation within a single batch but also -which batches remain. Additionally, batches must be processed sequentially so that a -continuation token unambiguously identifies a position in the result stream. +Because the SDK batches PK values across multiple physical partitions, the continuation token +must capture not only the backend query continuation within a single batch but also which +batches remain. Additionally, batches must be processed sequentially so that a continuation +token unambiguously identifies a position in the result stream. ### Composite continuation token -A new internal class `ReadManyByPartitionKeyContinuationToken` captures the resume state: - -| Field | Type | Description | -|-----------------------|-----------------------|----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------| -| `remainingBatches` | `List>` | EPK ranges (minInclusive, maxExclusive) of batches not yet started. Uses EPK ranges — **not** PkRangeIds — so the token survives partition splits. Reuses the existing `Range` type. | -| `currentBatch` | `Range` | EPK range of the batch currently being processed. | -| `backendContinuation` | `String` (nullable) | Backend query continuation within the current batch. `null` means "start from the beginning of this batch". | +A new internal class `ReadManyByPartitionKeyContinuationToken` captures the resume state. Each +remaining/current batch is identified by a single `BatchDefinition` carrying just the +**batchFilter** EPK range (`[minInclusive, maxExclusive)` containing the EPKs of all PKs in that +batch). The **partitionScope** used as the FeedRange at execution time is intentionally NOT +persisted - it is rederived per batch on resume. + +| Field | Type | Description | +|-----------------------|----------------------------|----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------| +| `version` | `int` (wire field `v`) | Wire format version. Current = 1. Tokens with an unknown version are rejected with `IllegalArgumentException`. | +| `remainingBatches` | `List` | EPK FILTER range of each batch not yet started. | +| `currentBatch` | `BatchDefinition` | EPK FILTER range of the batch currently being processed. | +| `backendContinuation` | `String` (nullable) | Backend query continuation within the current batch. `null` means "start from the beginning of this batch". | +| `collectionRid` | `String` | Resource id of the collection. Token is rejected on resume if the rid does not match. | +| `queryHash` | `String` (Murmur3-128 hex) | Stable hash of the custom `SqlQuerySpec` (text + parameter names + parameter values). Token is rejected on resume if the hash does not match. | +| `partitionKeySetHash` | `String` (Murmur3-128 hex) | Stable hash of the deduplicated, sorted set of input EPKs. Order- and duplicate-invariant: a caller may shuffle the input PK list across resume attempts and still get the same data back. | EPK ranges are represented as `Range` (the existing SDK type used throughout the routing layer) with `isMinInclusive = true` and `isMaxExclusive = false`. -**Serialization:** JSON → Base64, keeping the token opaque to callers. Example decoded: +### Why batchFilter only, not partitionScope + +Persisting the `partitionScope` (the physical-partition EPK boundary) in the token would have +two drawbacks: + +- **Token bloat**: doubles the number of EPK ranges in the serialized form. +- **Stale-after-split routing**: if the partition splits between the moment the token was + emitted and the moment a caller resumes, the persisted `partitionScope` would no longer match + any current physical partition exactly, so the backend would silently fan the query out + across multiple partitions for the entire range - inflating RU charge until the token is + retired. + +Instead, the SDK rederives `partitionScope` per batch at execution time: + +```text +partitionScope = union(getOverlappingRanges(batchFilter)) + = [ min(minInclusive of overlapping ranges), + max(maxExclusive of overlapping ranges) ) +``` + +When the routing-map cache is fresh, this yields a range that exactly matches one or more +physical partitions, so the backend hits only the intended partition(s) at minimum RU cost. +When the cache is briefly stale immediately after a split, the SDK's `StaleResourceRetryPolicy` +refreshes it on the first miss, so any RU-cost elevation is bounded to a single retried +request. + +### Serialization + +JSON -> Base64, keeping the token opaque to callers. The serialized form uses short field +names (`v`, `rb`, `cb`, `bc`, `cr`, `qh`, `ph`) and a compact `BatchDefinition` shape with one +EPK range per batch (`bf` = batch filter). Example decoded: ```json { - "remainingBatches": [ - {"min": "05C1E0", "max": "0BF333"}, - {"min": "0BF333", "max": "FF"} + "v": 1, + "rb": [ + {"bf": {"min": "05C1E0", "max": "0BF333"}}, + {"bf": {"min": "0BF333", "max": "FF"}} ], - "currentBatch": {"min": "", "max": "05C1E0"}, - "backendContinuation": "eyJDb21wb3NpdGVUb2" + "cb": {"bf": {"min": "", "max": "05C1E0"}}, + "bc": "eyJDb21wb3NpdGVUb2", + "cr": "dbs/myDb/colls/myColl", + "qh": "", + "ph": "" } ``` -### Sequential batch execution +### Sequential batch execution and stamping -Batches are sorted by `minInclusive` EPK (lexicographic) and processed one at a time -via ordered sequential merging with a small amount of Reactor prefetch. Each `FeedResponse` carries the composite -continuation token reflecting the current position: +Batches are sorted by `batchFilter.minInclusive` (lexicographic) and processed via +`Flux.mergeSequential` so emission order is strictly batch-by-batch. Each `FeedResponse` +is stamped with a composite token reflecting its position in the stream: ```text FeedResponse 1: remaining=[B2, B3], current=B1, backend=null @@ -240,47 +395,79 @@ FeedResponse 4: remaining=[B3], current=B2, backend= (B2 exhausted) FeedResponse 5: remaining=[], current=B3, backend=null FeedResponse 6: remaining=[], current=B3, backend= - (done — final FeedResponse has null continuation token) + (done - final FeedResponse has null continuation token) ``` -Sequential execution is always used — batches are never interleaved concurrently. -This ensures every `FeedResponse` carries a usable continuation token regardless of -whether the caller plans to resume or not. +Reactor's `mergeSequential` is allowed to subscribe-eagerly to the next batch (small prefetch) +while the current batch is still draining, so backend-side prefetch of the next batch is +permitted and improves throughput, but it never reorders the emitted items. This guarantees +every `FeedResponse` carries a usable continuation token regardless of whether the caller plans +to resume or not. + +### Resume sequence + +```mermaid +sequenceDiagram + autonumber + participant Caller + participant SDK as RxDocumentClientImpl + participant Cache as PartitionKeyRangeCache + participant Backend + Caller->>SDK: readManyByPartitionKeys(pks, ?customQuery, options(token)) + SDK->>SDK: deserialize token
verify v + collectionRid
+ queryHash + partitionKeySetHash + SDK->>Cache: tryLookupAsync(collectionRid) + Cache-->>SDK: routingMap (post-split if cache refreshed) + loop per persisted batchFilter + SDK->>SDK: filterPksByEpkRange(batchFilter) + SDK->>SDK: resolvePartitionScopeFromBatchFilter(batchFilter, routingMap) + SDK->>Backend: createQueryInternal(SELECT WHERE pkFilter, FeedRange=partitionScope, [backendCont if 1st]) + Backend-->>SDK: FeedResponse pages + SDK->>SDK: stamp each FeedResponse with composite token (current=this, remaining=tail) + SDK-->>Caller: FeedResponse + end +``` ### API changes -**`CosmosReadManyRequestOptionsImpl`:** Add `requestContinuation` field with -getter/setter. The impl already extends `CosmosQueryRequestOptionsBase`. - -**`RxDocumentClientImpl.readManyByPartitionKeys`:** +**`CosmosReadManyByPartitionKeyRequestOptionsImpl`:** Holds `continuationToken`, +`maxConcurrentBatchPrefetch`, and `maxItemCount` fields with corresponding getters and setters; +all other knobs (consistency, session token, regions, throughput control, dedicated gateway, +diagnostics, custom item serializer, keyword identifiers, e2e latency policy) are inherited +from `CosmosQueryRequestOptionsBase`. The public `CosmosReadManyByPartitionKeyRequestOptions` +facade exposes only the surface that is meaningful for this operation - +`MaxBackendContinuationTokenSizeInKb` (renamed from the inherited +`ResponseContinuationTokenLimitInKb` to make clear it limits the *backend* token wrapped inside +the composite token), `MaxItemCount`, `MaxConcurrentBatchPrefetch`, `ContinuationToken`, plus +the standard query-options forwarders. -- **With continuation token:** Deserialize the composite token. Restore batch list from - `currentBatch` + `remainingBatches` (no PK→routing-map computation needed). Pass - `backendContinuation` to the `CosmosQueryRequestOptions` for the first batch. Process - remaining batches in EPK order via `Flux.concat`. -- **Without continuation:** Compute batches as today, map each `PartitionKeyRange` to its - EPK range, sort by `minInclusive`, execute via `Flux.concat`, stamp each `FeedResponse` - with the composite token. +**`RxDocumentClientImpl.readManyByPartitionKeys`:** see Step 3 above. -**`CosmosAsyncContainer.readManyByPartitionKeyInternalFunc`:** Pass continuation token -from request options through to `RxDocumentClientImpl`. +**`CosmosAsyncContainer.readManyByPartitionKeyInternalFunc`:** Pulls the continuation token +out of the request options (before entering the reactive chain) and threads it into the +cloned `CosmosQueryRequestOptions` so `QueryFeedOperationState` carries it through to +`RxDocumentClientImpl`. ### Spark reader integration -**`ItemsPartitionReaderWithReadManyByPartitionKey`:** The `fluxFactory` lambda applies -the continuation token: when non-null, set it on a cloned `CosmosReadManyRequestOptions` -before calling `readManyByPartitionKeys`. +**`ItemsPartitionReaderWithReadManyByPartitionKey`:** The `fluxFactory` lambda applies the +continuation token: when non-null, set it on the preconfigured request options before calling +`readManyByPartitionKeys`. -**`TransientIOErrorsRetryingReadManyByPartitionKeyIterator`:** No changes needed — the -existing retry loop already captures `lastContinuationToken` from each `FeedResponse` -and passes it to `fluxFactory` on retry. Once the factory applies the token, retry -works correctly. +**`TransientIOErrorsRetryingReadManyByPartitionKeyIterator`:** A lightweight clone of the +existing `TransientIOErrorsRetryingIterator` adapted for `readManyByPartitionKeys`. Captures +`lastContinuationToken` from each fully-drained `FeedResponse` and passes it back to the +factory on retry. The iterator must NEVER call `.subscribe()` on the `CosmosPagedFlux` it +holds (a bug previously fixed for the regular query iterator); it only consumes the +`iterableByPage().iterator()` view. ### EPK range stability across partition splits The token uses EPK ranges, not PkRangeIds. If a split occurs between retries: -- The query engine handles EPK ranges spanning multiple new physical partitions transparently. +- Routing scope is rederived from the current routing map (Step 3b/5), so the FeedRange used + for the resumed query exactly matches the post-split physical partition boundaries. - A stale backend continuation is rejected by the backend; the SDK's `StaleResourceRetryPolicy` retries with a fresh routing map while the EPK range ensures correct targeting. +- The `partitionKeySetHash` and `queryHash` are unaffected by split (they hash the input PK + set and the user query, not the routing layout). From da578f6e1421e221416561a192ed7cbec296c0dd Mon Sep 17 00:00:00 2001 From: Fabian Meiswinkel Date: Wed, 22 Apr 2026 11:02:29 +0000 Subject: [PATCH 43/64] Update cspell.yaml --- sdk/cosmos/cspell.yaml | 1 + 1 file changed, 1 insertion(+) diff --git a/sdk/cosmos/cspell.yaml b/sdk/cosmos/cspell.yaml index 5cbd5ca804e7..2237242789e8 100644 --- a/sdk/cosmos/cspell.yaml +++ b/sdk/cosmos/cspell.yaml @@ -5,3 +5,4 @@ overrides: words: - DCOUNT - dedupe + - colls From f393c3f7b331a0038f8a64b7fb50d8afb3b8b734 Mon Sep 17 00:00:00 2001 From: Fabian Meiswinkel Date: Wed, 22 Apr 2026 13:34:48 +0000 Subject: [PATCH 44/64] Fixing test failure --- ...nyByPartitionKeyContinuationTokenTest.java | 29 +++++++++++++++++ ...ReadManyByPartitionKeyQueryHelperTest.java | 29 +++++++++++++++++ .../ReadManyByPartitionKeyQueryHelper.java | 10 +++--- .../implementation/RxDocumentClientImpl.java | 31 +++++++++++++++++-- .../com/azure/cosmos/models/FeedResponse.java | 17 +++++++--- .../docs/readManyByPartitionKey-design.md | 2 +- 6 files changed, 105 insertions(+), 13 deletions(-) diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyContinuationTokenTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyContinuationTokenTest.java index 35309dcda135..403683c69452 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyContinuationTokenTest.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyContinuationTokenTest.java @@ -6,6 +6,8 @@ package com.azure.cosmos.implementation; import com.azure.cosmos.implementation.routing.Range; +import com.azure.cosmos.models.FeedResponse; +import com.azure.cosmos.models.ModelBridgeInternal; import com.azure.cosmos.models.SqlParameter; import com.azure.cosmos.models.SqlQuerySpec; import org.testng.annotations.Test; @@ -13,7 +15,9 @@ import java.util.ArrayList; import java.util.Arrays; import java.util.Collections; +import java.util.HashMap; import java.util.List; +import java.util.Map; import static org.assertj.core.api.Assertions.assertThat; import static org.assertj.core.api.Assertions.assertThatThrownBy; @@ -121,6 +125,31 @@ public void roundtrip_lastBatchNoContinuation() { assertThat(deserialized.getBackendContinuation()).isNull(); } + @Test(groups = { "unit" }) + public void setFeedResponseContinuationToken_handlesEmptyHeadersWithoutCopyingNormalCase() { + Map immutableEmptyHeaders = Collections.emptyMap(); + FeedResponse emptyResponse = ModelBridgeInternal.createFeedResponse( + Collections.emptyList(), + immutableEmptyHeaders); + + ModelBridgeInternal.setFeedResponseContinuationToken(null, emptyResponse); + + assertThat(emptyResponse.getContinuationToken()).isNull(); + assertThat(emptyResponse.getResponseHeaders()).isNotSameAs(immutableEmptyHeaders); + assertThat(emptyResponse.getResponseHeaders()).isEmpty(); + + Map normalHeaders = new HashMap<>(); + normalHeaders.put(HttpConstants.HttpHeaders.ACTIVITY_ID, "test-activity-id"); + FeedResponse normalResponse = ModelBridgeInternal.createFeedResponse( + Collections.emptyList(), + normalHeaders); + + ModelBridgeInternal.setFeedResponseContinuationToken("token", normalResponse); + + assertThat(normalResponse.getContinuationToken()).isEqualTo("token"); + assertThat(normalResponse.getResponseHeaders()).isSameAs(normalHeaders); + } + @Test(groups = { "unit" }) public void deserialize_malformedInput_throws() { // Either the base64 decoder or the JSON parsing layer rejects garbage; both raise diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelperTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelperTest.java index 492bf4dd91f3..8a4b63069fc3 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelperTest.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelperTest.java @@ -20,6 +20,7 @@ import java.util.List; import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; public class ReadManyByPartitionKeyQueryHelperTest { @@ -194,6 +195,34 @@ public void hpk_customQuery_withWhere() { assertThat(result.getParameters()).hasSize(3); // @status + 2 pk params } + @Test(groups = { "unit" }) + public void normalizePartitionKeys_removesSubsumedFullHpkValues() { + PartitionKeyDefinition pkDef = createMultiHashPkDefinition("/city", "/zipcode", "/areaCode"); + + List pkValues = Arrays.asList( + new PartitionKeyBuilder().add("Redmond").build(), + new PartitionKeyBuilder().add("Redmond").add("98052").add(1).build(), + new PartitionKeyBuilder().add("Seattle").build()); + + List normalizedPartitionKeys = + RxDocumentClientImpl.normalizePartitionKeys(pkValues, pkDef); + + assertThat(normalizedPartitionKeys) + .extracting(normalizedPk -> normalizedPk.effectivePkInternal.toJson()) + .containsExactly("[\"Redmond\"]", "[\"Seattle\"]"); + } + + @Test(groups = { "unit" }) + public void normalizePartitionKeys_rejectsNoneForMultiHashPartitionKeys() { + PartitionKeyDefinition pkDef = createMultiHashPkDefinition("/city", "/zipcode"); + + assertThatThrownBy(() -> RxDocumentClientImpl.normalizePartitionKeys( + Collections.singletonList(PartitionKey.NONE), + pkDef)) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("PartitionKey.NONE is not supported for multi-path partition keys"); + } + //endregion //region findTopLevelWhereIndex tests diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelper.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelper.java index 7384788920c4..3b0d0f490165 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelper.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelper.java @@ -12,7 +12,6 @@ import java.util.ArrayList; import java.util.List; -import java.util.stream.Collectors; /** * Helper for constructing SqlQuerySpec instances for readManyByPartitionKeys operations. @@ -108,10 +107,11 @@ public static SqlQuerySpec createReadManyByPkQuerySpec( PartitionKeyInternal pkInternal = BridgeInternal.getPartitionKeyInternal(pkValues.get(i)); Object[] pkComponents = pkInternal.toObjectArray(); - // PartitionKey.NONE for hierarchical keys is impossible at this point: the builder - // (CosmosPartitionKeyHelper.validateNoneHandlingForPartitionKeyComponentCount and - // PartitionKeyBuilder.addNoneValue) reject NONE on multi-path PK definitions before - // reaching this code path. So pkComponents is guaranteed non-null here. + if (pkComponents == null) { + throw new IllegalArgumentException( + "PartitionKey.NONE is not supported for multi-path partition keys in readManyByPartitionKeys."); + } + { pkFilter.append("("); for (int j = 0; j < pkComponents.length; j++) { diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java index 5abb617808b2..5fbf82eb9830 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java @@ -4558,7 +4558,7 @@ private Flux> readManyByPartitionKeys( * Duplicates are removed by effective partition key string and the remaining keys are sorted * lexicographically by EPK so batching and continuation-token hashes stay stable. */ - private List normalizePartitionKeys( + static List normalizePartitionKeys( List partitionKeys, PartitionKeyDefinition pkDefinition) { @@ -4566,6 +4566,11 @@ private List normalizePartitionKeys( for (PartitionKey pk : partitionKeys) { PartitionKeyInternal pkInternal = BridgeInternal.getPartitionKeyInternal(pk); + if (pkDefinition.getKind() == PartitionKind.MULTI_HASH && pkInternal.getComponents() == null) { + throw new IllegalArgumentException( + "PartitionKey.NONE is not supported for multi-path partition keys in readManyByPartitionKeys."); + } + PartitionKeyInternal effectivePkInternal = pkInternal.getComponents() == null ? PartitionKeyInternal.UndefinedPartitionKey : pkInternal; @@ -4579,7 +4584,27 @@ private List normalizePartitionKeys( List normalizedPartitionKeys = new ArrayList<>(normalizedByEpk.values()); normalizedPartitionKeys.sort(Comparator.comparing(normalizedPk -> normalizedPk.effectivePartitionKeyString)); - return normalizedPartitionKeys; + + if (pkDefinition.getKind() != PartitionKind.MULTI_HASH || normalizedPartitionKeys.size() < 2) { + return normalizedPartitionKeys; + } + + List collapsedPartitionKeys = new ArrayList<>(normalizedPartitionKeys.size()); + for (NormalizedPartitionKey candidate : normalizedPartitionKeys) { + if (!collapsedPartitionKeys.isEmpty()) { + NormalizedPartitionKey previous = collapsedPartitionKeys.get(collapsedPartitionKeys.size() - 1); + if (previous.effectivePkInternal.contains(candidate.effectivePkInternal)) { + // The previous PK is a prefix of the current PK, so the current PK is fully + // subsumed by the previous read scope and would only cause duplicate results + // if it ended up in a separate batch. + continue; + } + } + + collapsedPartitionKeys.add(candidate); + } + + return collapsedPartitionKeys; } /** @@ -4894,7 +4919,7 @@ private static Range resolvePartitionScopeFromBatchFilter( return new Range<>(minInclusive, maxExclusive, true, false); } - private static final class NormalizedPartitionKey { + static final class NormalizedPartitionKey { final PartitionKey partitionKey; final PartitionKeyInternal effectivePkInternal; final String effectivePartitionKeyString; diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/FeedResponse.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/FeedResponse.java index fed27eb48945..5369143aee94 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/FeedResponse.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/FeedResponse.java @@ -44,7 +44,7 @@ private static ImplementationBridgeHelpers.CosmosDiagnosticsHelper.CosmosDiagnos private static final Pattern DELIMITER_CHARS_PATTERN = Pattern.compile(Constants.Quota.DELIMITER_CHARS); private final List results; - private final Map header; + private Map header; private final HashMap usageHeaders; private final HashMap quotaHeaders; private final boolean useEtagAsContinuation; @@ -55,6 +55,10 @@ private static ImplementationBridgeHelpers.CosmosDiagnosticsHelper.CosmosDiagnos private QueryInfo queryInfo; private QueryInfo.QueryPlanDiagnosticsContext queryPlanDiagnosticsContext; + private static Map ensureMutableHeadersMap(Map headers) { + return headers == null ? new HashMap<>() : headers; + } + FeedResponse(List results, Map headers) { this(results, headers, false, false, new ConcurrentHashMap<>()); } @@ -112,7 +116,7 @@ private FeedResponse( boolean nochanges, ConcurrentMap queryMetricsMap) { this.results = results; - this.header = header; + this.header = ensureMutableHeadersMap(header); this.usageHeaders = new HashMap<>(); this.quotaHeaders = new HashMap<>(); this.useEtagAsContinuation = useEtagAsContinuation; @@ -129,7 +133,7 @@ private FeedResponse( ConcurrentMap queryMetricsMap, CosmosDiagnostics diagnostics) { this.results = results; - this.header = header; + this.header = ensureMutableHeadersMap(header); this.usageHeaders = new HashMap<>(); this.quotaHeaders = new HashMap<>(); this.useEtagAsContinuation = useEtagAsContinuation; @@ -430,9 +434,14 @@ void setContinuationToken(String continuationToken) { ? HttpConstants.HttpHeaders.E_TAG : HttpConstants.HttpHeaders.CONTINUATION; + setContinuationTokenInternal(headerName, continuationToken); + } + + private void setContinuationTokenInternal(String headerName, String continuationToken) { if (!Strings.isNullOrWhiteSpace(continuationToken)) { + this.header = ensureMutableHeadersMap(this.header); this.header.put(headerName, continuationToken); - } else { + } else if (this.header != null && !this.header.isEmpty()) { this.header.remove(headerName); } } diff --git a/sdk/cosmos/docs/readManyByPartitionKey-design.md b/sdk/cosmos/docs/readManyByPartitionKey-design.md index 17fe086e4018..1075c29cd8ce 100644 --- a/sdk/cosmos/docs/readManyByPartitionKey-design.md +++ b/sdk/cosmos/docs/readManyByPartitionKey-design.md @@ -327,7 +327,7 @@ persisted - it is rederived per batch on resume. | `currentBatch` | `BatchDefinition` | EPK FILTER range of the batch currently being processed. | | `backendContinuation` | `String` (nullable) | Backend query continuation within the current batch. `null` means "start from the beginning of this batch". | | `collectionRid` | `String` | Resource id of the collection. Token is rejected on resume if the rid does not match. | -| `queryHash` | `String` (Murmur3-128 hex) | Stable hash of the custom `SqlQuerySpec` (text + parameter names + parameter values). Token is rejected on resume if the hash does not match. | +| `queryHash` | `String` (Murmur3-128 hex) | Stable hash of the custom `SqlQuerySpec` (text + parameter names + parameter values). Token is rejected on resume if the hash does not match. Resume intentionally requires the exact same query shape, not just a semantically equivalent query. | | `partitionKeySetHash` | `String` (Murmur3-128 hex) | Stable hash of the deduplicated, sorted set of input EPKs. Order- and duplicate-invariant: a caller may shuffle the input PK list across resume attempts and still get the same data back. | EPK ranges are represented as `Range` (the existing SDK type used throughout the From d4722c2cfaeb23b83fc75849b69a4ff73e39a5d5 Mon Sep 17 00:00:00 2001 From: Fabian Meiswinkel Date: Wed, 22 Apr 2026 14:02:02 +0000 Subject: [PATCH 45/64] Fixed test issues --- .../spark/CosmosPartitionKeyHelper.scala | 2 +- ...tryingReadManyByPartitionKeyIterator.scala | 26 ++++++---- .../udf/GetCosmosPartitionKeyValue.scala | 2 +- ...nyByPartitionKeyContinuationTokenTest.java | 2 +- ...ReadManyByPartitionKeyQueryHelperTest.java | 52 +++++++++++++++++++ ...ByPartitionKeyQueryPlanValidationTest.java | 12 ++--- .../azure/cosmos/CosmosAsyncContainer.java | 14 ++--- .../implementation/RxDocumentClientImpl.java | 8 +-- .../com/azure/cosmos/models/FeedResponse.java | 5 +- 9 files changed, 90 insertions(+), 33 deletions(-) diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosPartitionKeyHelper.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosPartitionKeyHelper.scala index a96572c5f632..68c51dde4c3d 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosPartitionKeyHelper.scala +++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosPartitionKeyHelper.scala @@ -89,4 +89,4 @@ private[spark] object CosmosPartitionKeyHelper extends BasicLoggingTrait { case _ => None } } -} \ No newline at end of file +} diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/TransientIOErrorsRetryingReadManyByPartitionKeyIterator.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/TransientIOErrorsRetryingReadManyByPartitionKeyIterator.scala index 28044d00a069..bd57bf48cb8c 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/TransientIOErrorsRetryingReadManyByPartitionKeyIterator.scala +++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/TransientIOErrorsRetryingReadManyByPartitionKeyIterator.scala @@ -50,6 +50,7 @@ private[spark] class TransientIOErrorsRetryingReadManyByPartitionKeyIterator[TSp private val rnd = Random // scalastyle:off null private val lastContinuationToken = new AtomicReference[String](null) + private val pendingContinuationToken = new AtomicReference[String](null) // scalastyle:on null private val retryCount = new AtomicLong(0) private lazy val operationContextString = operationContextAndListener match { @@ -82,6 +83,14 @@ private[spark] class TransientIOErrorsRetryingReadManyByPartitionKeyIterator[TSp if (hasBufferedNext) { Some(true) } else { + // All items from the previous page have been drained — promote the pending + // continuation token so that any retry resumes from the NEXT page rather than + // replaying items the caller has already consumed. + val pending = pendingContinuationToken.getAndSet(null) + if (pending != null) { + lastContinuationToken.set(pending) + } + val feedResponseIterator = currentFeedResponseIterator match { case Some(existing) => existing case None => @@ -133,16 +142,12 @@ private[spark] class TransientIOErrorsRetryingReadManyByPartitionKeyIterator[TSp feedResponse) } val iteratorCandidate = feedResponse.getResults.iterator().asScala.buffered - // INVARIANT: it is safe to record the continuation token BEFORE the items in this - // FeedResponse have been drained because executeWithRetry only wraps `hasNext`, - // and the buffered iterator is always fully drained before the next page is fetched. - // If a transient failure occurs while draining items from the *current* page, the - // page is replayed from this continuation token, which is the start of *this* page - // (the server's continuation token points at the next page), so on retry the SDK - // re-issues the request that produced the page we were processing. Items already - // emitted to the caller may be re-emitted; readManyByPartitionKeys is idempotent - // and returns documents (not deltas), so duplicates from replay are acceptable. - lastContinuationToken.set(feedResponse.getContinuationToken) + // Store the continuation token from this page as "pending". It will only be promoted + // to lastContinuationToken (the retry resume point) after all items in this page have + // been drained by the caller. This ensures that on a transient failure mid-page the + // retry resumes from the PREVIOUS page's continuation (i.e. re-fetches the current + // page) and never skips items. + pendingContinuationToken.set(feedResponse.getContinuationToken) if (iteratorCandidate.hasNext) { currentItemIterator = Some(iteratorCandidate) @@ -217,6 +222,7 @@ private[spark] class TransientIOErrorsRetryingReadManyByPartitionKeyIterator[TSp currentItemIterator = None currentFeedResponseIterator = None + pendingContinuationToken.set(null) Thread.sleep(retryIntervalInMs) } } diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/udf/GetCosmosPartitionKeyValue.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/udf/GetCosmosPartitionKeyValue.scala index cb4b2c5d5fbc..3e536269f2fb 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/udf/GetCosmosPartitionKeyValue.scala +++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/udf/GetCosmosPartitionKeyValue.scala @@ -26,4 +26,4 @@ class GetCosmosPartitionKeyValue extends UDF1[Object, String] { CosmosPartitionKeyHelper.getCosmosPartitionKeyValueString(List(partitionKeyValue)) } } -} \ No newline at end of file +} diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyContinuationTokenTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyContinuationTokenTest.java index 403683c69452..032218a08e61 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyContinuationTokenTest.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyContinuationTokenTest.java @@ -374,4 +374,4 @@ public void batchDefinition_roundtrip() { assertThat(remainingBatch.getBatchFilter().getMin()).isEqualTo("05C1E0"); assertThat(remainingBatch.getBatchFilter().getMax()).isEqualTo("0A"); } -} \ No newline at end of file +} diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelperTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelperTest.java index 8a4b63069fc3..b1f240182deb 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelperTest.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelperTest.java @@ -12,6 +12,7 @@ import com.azure.cosmos.models.PartitionKind; import com.azure.cosmos.models.SqlParameter; import com.azure.cosmos.models.SqlQuerySpec; +import com.azure.cosmos.implementation.routing.Range; import org.testng.annotations.Test; import java.util.ArrayList; @@ -223,6 +224,57 @@ public void normalizePartitionKeys_rejectsNoneForMultiHashPartitionKeys() { .hasMessageContaining("PartitionKey.NONE is not supported for multi-path partition keys"); } + @Test(groups = { "unit" }) + public void normalizePartitionKeys_epkAtBatchBoundary_isCorrectlyContained() { + // Verifies that a PK whose EPK sits exactly at a batchFilter boundary is correctly + // included in its owning range and excluded from the next range. The batchFilter is + // [minInclusive, maxExclusive) — a PK at the minInclusive boundary must be included, + // and a PK at the maxExclusive boundary must be excluded (it belongs to the next batch). + PartitionKeyDefinition pkDef = createSinglePkDefinition("/mypk"); + List pkValues = Arrays.asList( + new PartitionKey("alpha"), + new PartitionKey("bravo"), + new PartitionKey("charlie")); + + List normalized = + RxDocumentClientImpl.normalizePartitionKeys(pkValues, pkDef); + + // After normalization the EPKs are sorted; take the middle one's EPK as a boundary + assertThat(normalized).hasSize(3); + String middleEpk = normalized.get(1).effectivePartitionKeyString; + + // Range [middleEpk, ...) must include the middle PK (minInclusive) + Range rangeIncluding = new Range<>(middleEpk, "FF", true, false); + assertThat(rangeIncluding.contains(middleEpk)).isTrue(); + + // Range [..., middleEpk) must exclude the middle PK (maxExclusive) + Range rangeExcluding = new Range<>("", middleEpk, true, false); + assertThat(rangeExcluding.contains(middleEpk)).isFalse(); + } + + @Test(groups = { "unit" }) + public void normalizePartitionKeys_noneInHpk_alwaysRejected() { + // PartitionKey.NONE must be rejected for MULTI_HASH regardless of whether this is the + // first call or a continuation-resume call (normalizePartitionKeys is invoked on both + // code paths). This test verifies the invariant holds for mixed inputs too. + PartitionKeyDefinition pkDef = createMultiHashPkDefinition("/city", "/zipcode"); + + // NONE alone + assertThatThrownBy(() -> RxDocumentClientImpl.normalizePartitionKeys( + Collections.singletonList(PartitionKey.NONE), pkDef)) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("PartitionKey.NONE is not supported for multi-path partition keys"); + + // NONE mixed with valid HPK values + assertThatThrownBy(() -> RxDocumentClientImpl.normalizePartitionKeys( + Arrays.asList( + new PartitionKeyBuilder().add("Redmond").add("98052").build(), + PartitionKey.NONE), + pkDef)) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("PartitionKey.NONE is not supported for multi-path partition keys"); + } + //endregion //region findTopLevelWhereIndex tests diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryPlanValidationTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryPlanValidationTest.java index 11b69b955893..2bcaa8810d4f 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryPlanValidationTest.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryPlanValidationTest.java @@ -24,7 +24,7 @@ public void rejectsDCountQueryPlan() { dCountInfo.setDCountAlias("countAlias"); queryInfo.set("dCountInfo", dCountInfo); - assertThatThrownBy(() -> RxDocumentClientImpl.validateQueryPlanForReadManyByPartitionKey(createQueryPlan(queryInfo, null))) + assertThatThrownBy(() -> RxDocumentClientImpl.validateQueryPlanForReadManyByPartitionKeys(createQueryPlan(queryInfo, null))) .isInstanceOf(IllegalArgumentException.class) .hasMessageContaining("DCOUNT"); } @@ -34,7 +34,7 @@ public void rejectsOffsetQueryPlan() { QueryInfo queryInfo = new QueryInfo(); queryInfo.set("offset", 10); - assertThatThrownBy(() -> RxDocumentClientImpl.validateQueryPlanForReadManyByPartitionKey(createQueryPlan(queryInfo, null))) + assertThatThrownBy(() -> RxDocumentClientImpl.validateQueryPlanForReadManyByPartitionKeys(createQueryPlan(queryInfo, null))) .isInstanceOf(IllegalArgumentException.class) .hasMessageContaining("OFFSET"); } @@ -44,7 +44,7 @@ public void rejectsLimitQueryPlan() { QueryInfo queryInfo = new QueryInfo(); queryInfo.set("limit", 10); - assertThatThrownBy(() -> RxDocumentClientImpl.validateQueryPlanForReadManyByPartitionKey(createQueryPlan(queryInfo, null))) + assertThatThrownBy(() -> RxDocumentClientImpl.validateQueryPlanForReadManyByPartitionKeys(createQueryPlan(queryInfo, null))) .isInstanceOf(IllegalArgumentException.class) .hasMessageContaining("LIMIT"); } @@ -54,14 +54,14 @@ public void rejectsTopQueryPlan() { QueryInfo queryInfo = new QueryInfo(); queryInfo.set("top", 5); - assertThatThrownBy(() -> RxDocumentClientImpl.validateQueryPlanForReadManyByPartitionKey(createQueryPlan(queryInfo, null))) + assertThatThrownBy(() -> RxDocumentClientImpl.validateQueryPlanForReadManyByPartitionKeys(createQueryPlan(queryInfo, null))) .isInstanceOf(IllegalArgumentException.class) .hasMessageContaining("TOP"); } @Test(groups = { "unit" }) public void rejectsHybridSearchQueryPlanWithoutDereferencingNullQueryInfo() { - assertThatThrownBy(() -> RxDocumentClientImpl.validateQueryPlanForReadManyByPartitionKey( + assertThatThrownBy(() -> RxDocumentClientImpl.validateQueryPlanForReadManyByPartitionKeys( createQueryPlan(null, new HybridSearchQueryInfo()))) .isInstanceOf(IllegalArgumentException.class) .hasMessageContaining("hybrid/vector/full-text"); @@ -71,7 +71,7 @@ public void rejectsHybridSearchQueryPlanWithoutDereferencingNullQueryInfo() { public void acceptsSimpleQueryPlan() { QueryInfo queryInfo = new QueryInfo(); - assertThatCode(() -> RxDocumentClientImpl.validateQueryPlanForReadManyByPartitionKey(createQueryPlan(queryInfo, null))) + assertThatCode(() -> RxDocumentClientImpl.validateQueryPlanForReadManyByPartitionKeys(createQueryPlan(queryInfo, null))) .doesNotThrowAnyException(); } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java index 5d2009d24d4f..1842c5044947 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java @@ -176,7 +176,7 @@ private static ImplementationBridgeHelpers.CosmosBatchRequestOptionsHelper.Cosmo private final String createItemSpanName; private final String readAllItemsSpanName; private final String readManyItemsSpanName; - private final String readManyByPartitionKeySpanName; + private final String readManyByPartitionKeysSpanName; private final String readAllItemsOfLogicalPartitionSpanName; private final String queryItemsSpanName; private final String queryChangeFeedSpanName; @@ -210,7 +210,7 @@ protected CosmosAsyncContainer(CosmosAsyncContainer toBeWrappedContainer) { this.createItemSpanName = "createItem." + this.id; this.readAllItemsSpanName = "readAllItems." + this.id; this.readManyItemsSpanName = "readManyItems." + this.id; - this.readManyByPartitionKeySpanName = "readManyByPartitionKeys." + this.id; + this.readManyByPartitionKeysSpanName = "readManyByPartitionKeys." + this.id; this.readAllItemsOfLogicalPartitionSpanName = "readAllItemsOfLogicalPartition." + this.id; this.queryItemsSpanName = "queryItems." + this.id; this.queryChangeFeedSpanName = "queryChangeFeed." + this.id; @@ -1727,10 +1727,10 @@ public CosmosPagedFlux readManyByPartitionKeys( List partitionKeysSnapshot = new ArrayList<>(partitionKeys); return UtilBridgeInternal.createCosmosPagedFlux( - readManyByPartitionKeyInternalFunc(partitionKeysSnapshot, customQuery, requestOptions, classType)); + readManyByPartitionKeysInternalFunc(partitionKeysSnapshot, customQuery, requestOptions, classType)); } - private Function>> readManyByPartitionKeyInternalFunc( + private Function>> readManyByPartitionKeysInternalFunc( List partitionKeys, SqlQuerySpec customQuery, CosmosReadManyByPartitionKeyRequestOptions requestOptions, @@ -1777,16 +1777,16 @@ private Function>> readManyByPa } CosmosQueryRequestOptionsBase cosmosQueryRequestOptionsImpl = queryOptionsAccessor().getImpl(queryRequestOptions); - applyPolicies(OperationType.Query, ResourceType.Document, cosmosQueryRequestOptionsImpl, this.readManyByPartitionKeySpanName); + applyPolicies(OperationType.Query, ResourceType.Document, cosmosQueryRequestOptionsImpl, this.readManyByPartitionKeysSpanName); QueryFeedOperationState state = new QueryFeedOperationState( client, - this.readManyByPartitionKeySpanName, + this.readManyByPartitionKeysSpanName, database.getId(), this.getId(), ResourceType.Document, OperationType.Query, - queryOptionsAccessor().getQueryNameOrDefault(queryRequestOptions, this.readManyByPartitionKeySpanName), + queryOptionsAccessor().getQueryNameOrDefault(queryRequestOptions, this.readManyByPartitionKeysSpanName), queryRequestOptions, pagedFluxOptions ); diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java index 5fbf82eb9830..079e67e4102a 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java @@ -4524,7 +4524,7 @@ private Flux> readManyByPartitionKeys( // First-call path: validate custom query, resolve routing map, build batches Mono queryValidationMono; if (customQuery != null) { - queryValidationMono = validateCustomQueryForReadManyByPartitionKey( + queryValidationMono = validateCustomQueryForReadManyByPartitionKeys( customQuery, resourceLink, state.getQueryOptions()); } else { queryValidationMono = Mono.empty(); @@ -4966,7 +4966,7 @@ private static final class BatchDescriptor { } } - private Mono validateCustomQueryForReadManyByPartitionKey( + private Mono validateCustomQueryForReadManyByPartitionKeys( SqlQuerySpec customQuery, String resourceLink, CosmosQueryRequestOptions queryRequestOptions) { @@ -4983,11 +4983,11 @@ private Mono validateCustomQueryForReadManyByPartitionKey( queryRequestOptions, Configs.isQueryPlanCachingEnabled(), this.getQueryPlanCache()) - .doOnNext(RxDocumentClientImpl::validateQueryPlanForReadManyByPartitionKey) + .doOnNext(RxDocumentClientImpl::validateQueryPlanForReadManyByPartitionKeys) .then(); } - static void validateQueryPlanForReadManyByPartitionKey(PartitionedQueryExecutionInfo queryPlan) { + static void validateQueryPlanForReadManyByPartitionKeys(PartitionedQueryExecutionInfo queryPlan) { if (queryPlan.hasHybridSearchQueryInfo()) { throw new IllegalArgumentException( "Custom query for readMany by partition key must not contain hybrid/vector/full-text search."); diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/FeedResponse.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/FeedResponse.java index 5369143aee94..73e2b96339b2 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/FeedResponse.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/FeedResponse.java @@ -44,7 +44,7 @@ private static ImplementationBridgeHelpers.CosmosDiagnosticsHelper.CosmosDiagnos private static final Pattern DELIMITER_CHARS_PATTERN = Pattern.compile(Constants.Quota.DELIMITER_CHARS); private final List results; - private Map header; + private final Map header; private final HashMap usageHeaders; private final HashMap quotaHeaders; private final boolean useEtagAsContinuation; @@ -439,9 +439,8 @@ void setContinuationToken(String continuationToken) { private void setContinuationTokenInternal(String headerName, String continuationToken) { if (!Strings.isNullOrWhiteSpace(continuationToken)) { - this.header = ensureMutableHeadersMap(this.header); this.header.put(headerName, continuationToken); - } else if (this.header != null && !this.header.isEmpty()) { + } else if (!this.header.isEmpty()) { this.header.remove(headerName); } } From 26df2ab407a75165cd4cb25ec3def23b0d738bfb Mon Sep 17 00:00:00 2001 From: Fabian Meiswinkel Date: Wed, 22 Apr 2026 15:19:01 +0000 Subject: [PATCH 46/64] Update FeedResponse.java --- .../src/main/java/com/azure/cosmos/models/FeedResponse.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/FeedResponse.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/FeedResponse.java index 73e2b96339b2..5f397e90aa22 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/FeedResponse.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/FeedResponse.java @@ -440,7 +440,7 @@ void setContinuationToken(String continuationToken) { private void setContinuationTokenInternal(String headerName, String continuationToken) { if (!Strings.isNullOrWhiteSpace(continuationToken)) { this.header.put(headerName, continuationToken); - } else if (!this.header.isEmpty()) { + } else if (!this.header.isEmpty() && this.header.containsKey(headerName)) { this.header.remove(headerName); } } From a45aabbf28db195a0e05a84152c7cff8321915a8 Mon Sep 17 00:00:00 2001 From: Fabian Meiswinkel Date: Wed, 22 Apr 2026 15:32:21 +0000 Subject: [PATCH 47/64] Update FeedResponse.java --- .../java/com/azure/cosmos/models/FeedResponse.java | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/FeedResponse.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/FeedResponse.java index 5f397e90aa22..f023dc7a1146 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/FeedResponse.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/FeedResponse.java @@ -441,6 +441,16 @@ private void setContinuationTokenInternal(String headerName, String continuation if (!Strings.isNullOrWhiteSpace(continuationToken)) { this.header.put(headerName, continuationToken); } else if (!this.header.isEmpty() && this.header.containsKey(headerName)) { + // The query API returns unmodifiable header collections for empty + // responses (no documents returned - when only header set is request charge) + // the protection here to check for existence of the header before attempting + // to remove it would not be robust enough against unknown headers + // but since we only ever call our own query pipeline + // avoiding cloning in all cases and gating on continuation header + // existence is a reaosnable trade-off - test coverage exists that uncovered + // the problem - so, this acts as regression test as well + // --> the etst coverage is in ItemsPartitionReaderWithReadManyByPartitionKeyITest + // it should "return empty results for non-existent partition keys" this.header.remove(headerName); } } From 7299a2ca2c356a39e40e21b68c8c59047c3e37f1 Mon Sep 17 00:00:00 2001 From: Fabian Meiswinkel Date: Wed, 22 Apr 2026 20:09:42 +0000 Subject: [PATCH 48/64] Fixes test issues --- .../CosmosReadManyByPartitionKeyReader.scala | 6 +- ...tionReaderWithReadManyByPartitionKey.scala | 3 + .../TransientIOErrorsRetryingIterator.scala | 104 +++++++++++------- ...tryingReadManyByPartitionKeyIterator.scala | 55 ++------- .../cosmos/ReadManyByPartitionKeyTest.java | 23 ++++ ...nyByPartitionKeyContinuationTokenTest.java | 2 +- ...ReadManyByPartitionKeyQueryHelperTest.java | 16 +-- .../azure/cosmos/CosmosAsyncContainer.java | 6 +- .../com/azure/cosmos/models/FeedResponse.java | 15 ++- 9 files changed, 130 insertions(+), 100 deletions(-) diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosReadManyByPartitionKeyReader.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosReadManyByPartitionKeyReader.scala index eda6e755faf2..0bfd8339b4ec 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosReadManyByPartitionKeyReader.scala +++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosReadManyByPartitionKeyReader.scala @@ -15,6 +15,7 @@ import org.apache.spark.sql.{DataFrame, Row, SparkSession} import org.apache.spark.sql.types.StructType import java.util.UUID +import java.util.concurrent.atomic.AtomicBoolean // scalastyle:off underscore.import import scala.collection.JavaConverters._ @@ -153,11 +154,10 @@ private[spark] class CosmosReadManyByPartitionKeyReader( pkIterator) new Iterator[Row] { - private var isClosed = false + private val isClosed = new AtomicBoolean(false) private def closeReader(): Unit = { - if (!isClosed) { - isClosed = true + if (isClosed.compareAndSet(false, true)) { reader.close() } } diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/ItemsPartitionReaderWithReadManyByPartitionKey.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/ItemsPartitionReaderWithReadManyByPartitionKey.scala index b254ad51464d..9b1eeed06287 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/ItemsPartitionReaderWithReadManyByPartitionKey.scala +++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/ItemsPartitionReaderWithReadManyByPartitionKey.scala @@ -229,6 +229,9 @@ private[spark] case class ItemsPartitionReaderWithReadManyByPartitionKey // Factory that creates a CosmosPagedFlux from an optional continuation token. // On the first call continuationToken is null (start from scratch); on retry // it is the continuation token from the last fully-drained page. + // NOTE: mutating the shared readManyOptions before each call is safe because + // the SDK clones the options internally at the start of each call, and the + // iterator is single-threaded (mutation in step N completes before step N+1). private val fluxFactory: String => CosmosPagedFlux[SparkRowItem] = { (continuationToken: String) => // Reuse the preconfigured request options and only update the continuation // token so the SDK resumes from the last fully committed page. diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/TransientIOErrorsRetryingIterator.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/TransientIOErrorsRetryingIterator.scala index 946d085878b1..c949cd846f6c 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/TransientIOErrorsRetryingIterator.scala +++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/TransientIOErrorsRetryingIterator.scala @@ -50,7 +50,6 @@ private class TransientIOErrorsRetryingIterator[TSparkRow] 5 + CosmosConstants.readOperationEndToEndTimeoutInSeconds, scala.concurrent.duration.SECONDS) - private val rnd = Random // scalastyle:off null private val lastContinuationToken = new AtomicReference[String](null) // scalastyle:on null @@ -204,46 +203,17 @@ private class TransientIOErrorsRetryingIterator[TSparkRow] } private[spark] def executeWithRetry[T](methodName: String, func: () => T): T = { - val loop = new Breaks() - var returnValue: Option[T] = None - - loop.breakable { - while (true) { - val retryIntervalInMs = rnd.nextInt(maxRetryIntervalInMs) - - try { - returnValue = Some(func()) - retryCount.set(0) - loop.break - } - catch { - case cosmosException: CosmosException => - if (Exceptions.canBeTransientFailure(cosmosException.getStatusCode, cosmosException.getSubStatusCode)) { - val retryCountSnapshot = retryCount.incrementAndGet() - if (retryCountSnapshot > maxRetryCount) { - logError( - s"Too many transient failure retry attempts in TransientIOErrorsRetryingIterator.$methodName", - cosmosException) - throw cosmosException - } else { - logWarning( - s"Transient failure handled in TransientIOErrorsRetryingIterator.$methodName -" + - s" will be retried (attempt#$retryCountSnapshot) in ${retryIntervalInMs}ms", - cosmosException) - } - } else { - throw cosmosException - } - case other: Throwable => throw other - } - + TransientIOErrorsRetryingIterator.executeWithRetry( + "TransientIOErrorsRetryingIterator", + methodName, + func, + maxRetryCount, + maxRetryIntervalInMs, + retryCount, + () => { currentItemIterator = None currentFeedResponseIterator = None - Thread.sleep(retryIntervalInMs) - } - } - - returnValue.get + }) } private[this] def validateNextLsn(itemIterator: BufferedIterator[TSparkRow]): Boolean = { @@ -293,4 +263,60 @@ private[spark] object TransientIOErrorsRetryingIterator extends BasicLoggingTrai ) val executionContext = ExecutionContext.fromExecutorService(executorService) + + /** + * Shared retry logic for transient I/O failures. Both TransientIOErrorsRetryingIterator + * and TransientIOErrorsRetryingReadManyByPartitionKeyIterator delegate to this method + * to avoid duplicating the retry loop, backoff, and transient-failure classification. + */ + def executeWithRetry[T]( + callerName: String, + methodName: String, + func: () => T, + maxRetryCount: Long, + maxRetryIntervalInMs: Int, + retryCount: AtomicLong, + onRetry: () => Unit): T = { + + val rnd = Random + val loop = new Breaks() + var returnValue: Option[T] = None + + loop.breakable { + while (true) { + val retryIntervalInMs = rnd.nextInt(maxRetryIntervalInMs) + + try { + returnValue = Some(func()) + retryCount.set(0) + loop.break + } + catch { + case cosmosException: CosmosException => + if (Exceptions.canBeTransientFailure(cosmosException.getStatusCode, cosmosException.getSubStatusCode)) { + val retryCountSnapshot = retryCount.incrementAndGet() + if (retryCountSnapshot > maxRetryCount) { + logError( + s"Too many transient failure retry attempts in $callerName.$methodName", + cosmosException) + throw cosmosException + } else { + logWarning( + s"Transient failure handled in $callerName.$methodName -" + + s" will be retried (attempt#$retryCountSnapshot) in ${retryIntervalInMs}ms", + cosmosException) + } + } else { + throw cosmosException + } + case other: Throwable => throw other + } + + onRetry() + Thread.sleep(retryIntervalInMs) + } + } + + returnValue.get + } } diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/TransientIOErrorsRetryingReadManyByPartitionKeyIterator.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/TransientIOErrorsRetryingReadManyByPartitionKeyIterator.scala index bd57bf48cb8c..3b7e96e8cd45 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/TransientIOErrorsRetryingReadManyByPartitionKeyIterator.scala +++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/TransientIOErrorsRetryingReadManyByPartitionKeyIterator.scala @@ -3,7 +3,6 @@ package com.azure.cosmos.spark -import com.azure.cosmos.CosmosException import com.azure.cosmos.implementation.OperationCancelledException import com.azure.cosmos.implementation.spark.OperationContextAndListenerTuple import com.azure.cosmos.models.FeedResponse @@ -13,8 +12,6 @@ import com.azure.cosmos.util.{CosmosPagedFlux, CosmosPagedIterable} import java.util.concurrent.TimeoutException import java.util.concurrent.atomic.{AtomicLong, AtomicReference} import scala.concurrent.{Await, ExecutionContext, Future} -import scala.util.Random -import scala.util.control.Breaks // scalastyle:off underscore.import import scala.collection.JavaConverters._ @@ -47,7 +44,6 @@ private[spark] class TransientIOErrorsRetryingReadManyByPartitionKeyIterator[TSp 5 + CosmosConstants.readOperationEndToEndTimeoutInSeconds, scala.concurrent.duration.SECONDS) - private val rnd = Random // scalastyle:off null private val lastContinuationToken = new AtomicReference[String](null) private val pendingContinuationToken = new AtomicReference[String](null) @@ -184,52 +180,19 @@ private[spark] class TransientIOErrorsRetryingReadManyByPartitionKeyIterator[TSp } private[spark] def executeWithRetry[T](methodName: String, func: () => T): T = { - val loop = new Breaks() - var returnValue: Option[T] = None - - loop.breakable { - while (true) { - val retryIntervalInMs = rnd.nextInt(maxRetryIntervalInMs) - - try { - returnValue = Some(func()) - retryCount.set(0) - loop.break - } - catch { - case cosmosException: CosmosException => - if (Exceptions.canBeTransientFailure(cosmosException.getStatusCode, cosmosException.getSubStatusCode)) { - val retryCountSnapshot = retryCount.incrementAndGet() - if (retryCountSnapshot > maxRetryCount) { - logError( - s"Too many transient failure retry attempts in " + - s"TransientIOErrorsRetryingReadManyByPartitionKeyIterator.$methodName", - cosmosException) - throw cosmosException - } else { - logWarning( - s"Transient failure handled in " + - s"TransientIOErrorsRetryingReadManyByPartitionKeyIterator.$methodName -" + - s" will be retried (attempt#$retryCountSnapshot) in ${retryIntervalInMs}ms " + - s"(continuationToken=$lastContinuationToken)", - cosmosException) - } - } else { - throw cosmosException - } - case other: Throwable => throw other - } - + TransientIOErrorsRetryingIterator.executeWithRetry( + "TransientIOErrorsRetryingReadManyByPartitionKeyIterator", + methodName, + func, + maxRetryCount, + maxRetryIntervalInMs, + retryCount, + () => { currentItemIterator = None currentFeedResponseIterator = None pendingContinuationToken.set(null) - Thread.sleep(retryIntervalInMs) - } - } - - returnValue.get + }) } - override def close(): Unit = { currentItemIterator = None currentFeedResponseIterator = None diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/ReadManyByPartitionKeyTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/ReadManyByPartitionKeyTest.java index c5f3e46978ed..247c868ec6da 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/ReadManyByPartitionKeyTest.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/ReadManyByPartitionKeyTest.java @@ -482,6 +482,29 @@ public T deserialize(Map jsonNodeMap, Class classType) { cleanupContainer(singlePkContainer); } + @Test(groups = {"emulator"}, timeOut = TIMEOUT) + public void singlePk_readManyByPartitionKey_withRequestOptionsAndMaxConcurrentBatchPrefetch() { + // Regression test: passing non-null CosmosReadManyByPartitionKeyRequestOptions + // with maxConcurrentBatchPrefetch set should not throw NullPointerException + // from auto-unboxing a null MaxDegreeOfParallelism during options cloning. + List items = createSinglePkItems("pkMdop", 3); + + List pkValues = Collections.singletonList(new PartitionKey("pkMdop")); + com.azure.cosmos.models.CosmosReadManyByPartitionKeyRequestOptions options = + new com.azure.cosmos.models.CosmosReadManyByPartitionKeyRequestOptions(); + options.setMaxConcurrentBatchPrefetch(2); + + CosmosPagedIterable results = singlePkContainer.readManyByPartitionKeys( + pkValues, options, ObjectNode.class); + List resultList = results.stream().collect(Collectors.toList()); + + assertThat(resultList).hasSize(3); + resultList.forEach(item -> { + assertThat(item.get("mypk").asText()).isEqualTo("pkMdop"); + }); + + cleanupContainer(singlePkContainer); + } //endregion //region Continuation token tests diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyContinuationTokenTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyContinuationTokenTest.java index 032218a08e61..7e3d0c5c4fbf 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyContinuationTokenTest.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyContinuationTokenTest.java @@ -135,7 +135,7 @@ public void setFeedResponseContinuationToken_handlesEmptyHeadersWithoutCopyingNo ModelBridgeInternal.setFeedResponseContinuationToken(null, emptyResponse); assertThat(emptyResponse.getContinuationToken()).isNull(); - assertThat(emptyResponse.getResponseHeaders()).isNotSameAs(immutableEmptyHeaders); + assertThat(emptyResponse.getResponseHeaders()).isSameAs(immutableEmptyHeaders); assertThat(emptyResponse.getResponseHeaders()).isEmpty(); Map normalHeaders = new HashMap<>(); diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelperTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelperTest.java index b1f240182deb..9987a184e420 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelperTest.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelperTest.java @@ -618,18 +618,18 @@ public void singlePk_mixedNoneAndNormal_generatesCombinedFilter() { } @Test(groups = { "unit" }) - public void hpk_nonePartitionKey_generatesNotIsDefinedForAllPaths() { + public void hpk_nonePartitionKey_throwsForMultiHash() { PartitionKeyDefinition pkDef = createMultiHashPkDefinition("/city", "/zipcode"); List selectors = createSelectors(pkDef); List pkValues = Collections.singletonList(PartitionKey.NONE); - SqlQuerySpec result = ReadManyByPartitionKeyQueryHelper.createReadManyByPkQuerySpec( - "SELECT * FROM c", new ArrayList<>(), pkValues, selectors, pkDef); - - assertThat(result.getQueryText()).contains("NOT IS_DEFINED(c[\"city\"])"); - assertThat(result.getQueryText()).contains("NOT IS_DEFINED(c[\"zipcode\"])"); - assertThat(result.getQueryText()).contains("AND"); - assertThat(result.getParameters()).isEmpty(); + // PartitionKey.NONE is not supported for multi-path partition keys - + // the SDK rejects it in normalizePartitionKeys before reaching the query helper. + // The query helper itself also rejects it defensively. + assertThatThrownBy(() -> ReadManyByPartitionKeyQueryHelper.createReadManyByPkQuerySpec( + "SELECT * FROM c", new ArrayList<>(), pkValues, selectors, pkDef)) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("PartitionKey.NONE is not supported for multi-path partition keys"); } //endregion diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java index 1842c5044947..2b6a49700ac0 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java @@ -1764,7 +1764,11 @@ private Function>> readManyByPa // value but default the uninitialized case to unbounded for readManyByPartitionKeys. // CosmosReadManyRequestOptions does not currently expose MDOP, so this only matters // if it is plumbed through in the future. - if (queryRequestOptions.getMaxDegreeOfParallelism() == DEFAULT_MAX_DEGREE_OF_PARALLELISM) { + // When cloning from CosmosReadManyByPartitionKeyRequestOptionsImpl (which does not + // define maxDegreeOfParallelism), the cloned CosmosQueryRequestOptions may have a + // null backing Integer. Guard against NPE from auto-unboxing by checking for null. + Integer mdop = queryOptionsAccessor().getImpl(queryRequestOptions).getMaxDegreeOfParallelism(); + if (mdop == null || mdop == DEFAULT_MAX_DEGREE_OF_PARALLELISM) { queryRequestOptions.setMaxDegreeOfParallelism(UNBOUNDED_MAX_DEGREE_OF_PARALLELISM); } queryRequestOptions.setQueryName("readManyByPartitionKeys"); diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/FeedResponse.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/FeedResponse.java index f023dc7a1146..959cca07b2b4 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/FeedResponse.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/FeedResponse.java @@ -55,6 +55,17 @@ private static ImplementationBridgeHelpers.CosmosDiagnosticsHelper.CosmosDiagnos private QueryInfo queryInfo; private QueryInfo.QueryPlanDiagnosticsContext queryPlanDiagnosticsContext; + // All header maps are produced by the SDK's own query pipeline. Non-null maps + // are always mutable (HashMap or ConcurrentHashMap) - the SDK intentionally + // allows callers to add/modify headers on FeedResponse. The only known + // exception is empty-page responses where the query pipeline may pass null. + // We do NOT clone non-null maps here to avoid unnecessary allocations on every + // FeedResponse construction - the wider blast radius of cloning (every query, + // change feed, readMany response) is not justified by the narrow null case. + // If a future code path introduces an immutable non-null header map, the + // setContinuationTokenInternal method will fail fast with + // UnsupportedOperationException, and the fix should be to make the upstream + // pipeline emit a mutable map rather than adding defensive cloning here. private static Map ensureMutableHeadersMap(Map headers) { return headers == null ? new HashMap<>() : headers; } @@ -447,9 +458,9 @@ private void setContinuationTokenInternal(String headerName, String continuation // to remove it would not be robust enough against unknown headers // but since we only ever call our own query pipeline // avoiding cloning in all cases and gating on continuation header - // existence is a reaosnable trade-off - test coverage exists that uncovered + // existence is a reasonable trade-off - test coverage exists that uncovered // the problem - so, this acts as regression test as well - // --> the etst coverage is in ItemsPartitionReaderWithReadManyByPartitionKeyITest + // --> the test coverage is in ItemsPartitionReaderWithReadManyByPartitionKeyITest // it should "return empty results for non-existent partition keys" this.header.remove(headerName); } From dc00653622213db76e3c544c1348173d60e14e66 Mon Sep 17 00:00:00 2001 From: Fabian Meiswinkel Date: Wed, 22 Apr 2026 20:27:05 +0000 Subject: [PATCH 49/64] Fixed null handinlgin in ReadManyByParittionKeyToken JsonCreator --- ...nyByPartitionKeyContinuationTokenTest.java | 24 +++++++++++++++++++ ...adManyByPartitionKeyContinuationToken.java | 21 +++++++++++++--- 2 files changed, 42 insertions(+), 3 deletions(-) diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyContinuationTokenTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyContinuationTokenTest.java index 7e3d0c5c4fbf..86057a353021 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyContinuationTokenTest.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyContinuationTokenTest.java @@ -173,6 +173,30 @@ public void deserialize_null_throws() { ).isInstanceOf(NullPointerException.class); } + @Test(groups = { "unit" }) + public void deserialize_nullRemainingBatches_throws() { + // Hand-craft a token with "rb":null - a malformed or tampered token. + String json = "{\"v\":1,\"rb\":null,\"cb\":{\"bf\":{\"min\":\"\",\"max\":\"FF\"}}," + + "\"bc\":null,\"cr\":\"" + TEST_COLLECTION_RID + "\",\"qh\":\"" + TEST_QUERY_HASH + "\",\"ph\":\"" + TEST_PARTITION_KEY_SET_HASH + "\"}"; + String serialized = java.util.Base64.getEncoder().encodeToString(json.getBytes(java.nio.charset.StandardCharsets.UTF_8)); + + assertThatThrownBy(() -> ReadManyByPartitionKeyContinuationToken.deserialize(serialized)) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("remainingBatches"); + } + + @Test(groups = { "unit" }) + public void deserialize_nullCurrentBatch_throws() { + // Hand-craft a token with "cb":null - a malformed or tampered token. + String json = "{\"v\":1,\"rb\":[],\"cb\":null," + + "\"bc\":null,\"cr\":\"" + TEST_COLLECTION_RID + "\",\"qh\":\"" + TEST_QUERY_HASH + "\",\"ph\":\"" + TEST_PARTITION_KEY_SET_HASH + "\"}"; + String serialized = java.util.Base64.getEncoder().encodeToString(json.getBytes(java.nio.charset.StandardCharsets.UTF_8)); + + assertThatThrownBy(() -> ReadManyByPartitionKeyContinuationToken.deserialize(serialized)) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("currentBatch"); + } + @Test(groups = { "unit" }) public void deserialize_unsupportedVersion_throws() { // Hand-craft a token JSON with version=999 (a future format) and ensure it is rejected. diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyContinuationToken.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyContinuationToken.java index 77ad2dd16f9b..813bab25453d 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyContinuationToken.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyContinuationToken.java @@ -136,6 +136,15 @@ public ReadManyByPartitionKeyContinuationToken( } this.version = effectiveVersion; + + if (remainingBatches == null) { + throw new IllegalArgumentException( + "Malformed readManyByPartitionKeys continuation token: 'remainingBatches' is required."); + } + if (currentBatch == null) { + throw new IllegalArgumentException( + "Malformed readManyByPartitionKeys continuation token: 'currentBatch' is required."); + } this.remainingBatches = remainingBatches; this.currentBatch = currentBatch; this.backendContinuation = backendContinuation; @@ -309,12 +318,18 @@ public static ReadManyByPartitionKeyContinuationToken deserialize(String seriali String json = new String(decoded, StandardCharsets.UTF_8); return Utils.getSimpleObjectMapper().readValue(json, ReadManyByPartitionKeyContinuationToken.class); } catch (IllegalArgumentException e) { - // Preserve our own version-mismatch IllegalArgumentException without wrapping. + // Preserve our own IllegalArgumentException (version mismatch, null fields) without wrapping. throw e; } catch (Exception e) { + // Jackson wraps constructor-thrown IllegalArgumentException inside JsonMappingException. + // Unwrap to preserve the actionable error message from our null checks. + Throwable cause = e.getCause(); + if (cause instanceof IllegalArgumentException) { + throw (IllegalArgumentException) cause; + } throw new IllegalArgumentException( - "Failed to deserialize ReadManyByPartitionKeyContinuationToken. " + - "The continuation token may be malformed or from an incompatible version.", e); + "Failed to deserialize ReadManyByPartitionKeyContinuationToken. " + + "The continuation token may be malformed or from an incompatible version.", e); } } From c47f9a674152a360c8b3e46b4575b746f552c939 Mon Sep 17 00:00:00 2001 From: Fabian Meiswinkel Date: Wed, 22 Apr 2026 20:57:28 +0000 Subject: [PATCH 50/64] Update CosmosAsyncContainer.java --- .../main/java/com/azure/cosmos/CosmosAsyncContainer.java | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java index 2b6a49700ac0..4546c2db4632 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java @@ -1696,6 +1696,12 @@ public CosmosPagedFlux readManyByPartitionKeys( * Partial hierarchical partition keys are supported and will fan out to multiple * physical partitions. Duplicate partition key inputs are normalized with set-based semantics * before batching. + *

+ * {@link PartitionKey#NONE} is allowed for single-path partition key containers — it + * queries documents where the partition key path is absent (generates + * {@code NOT IS_DEFINED(...)}). For hierarchical (multi-path) partition key containers, + * {@link PartitionKey#NONE} is rejected because {@code addNoneValue()} is not supported + * with multiple paths. * * @param the type parameter * @param partitionKeys list of partition key values to read documents for From 1f2f14b99afcc713226822a364b72835bf2ff14a Mon Sep 17 00:00:00 2001 From: Fabian Meiswinkel Date: Wed, 22 Apr 2026 21:08:16 +0000 Subject: [PATCH 51/64] Renamed CosmosReadManyPartitionKeyRequestOptions to CosmosReadManyPartitionKeysRequestOptions --- ...tionReaderWithReadManyByPartitionKey.scala | 8 +- .../cosmos/ReadManyByPartitionKeyTest.java | 16 ++-- sdk/cosmos/azure-cosmos/CHANGELOG.md | 2 +- .../azure/cosmos/CosmosAsyncContainer.java | 14 +-- .../com/azure/cosmos/CosmosContainer.java | 6 +- ...anyByPartitionKeysRequestOptionsImpl.java} | 16 ++-- .../ImplementationBridgeHelpers.java | 32 +++---- ...eadManyByPartitionKeysRequestOptions.java} | 88 +++++++++---------- .../docs/readManyByPartitionKey-design.md | 10 +-- 9 files changed, 96 insertions(+), 96 deletions(-) rename sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/{CosmosReadManyByPartitionKeyRequestOptionsImpl.java => CosmosReadManyByPartitionKeysRequestOptionsImpl.java} (81%) rename sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/{CosmosReadManyByPartitionKeyRequestOptions.java => CosmosReadManyByPartitionKeysRequestOptions.java} (78%) diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/ItemsPartitionReaderWithReadManyByPartitionKey.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/ItemsPartitionReaderWithReadManyByPartitionKey.scala index 9b1eeed06287..a8cbadba5f70 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/ItemsPartitionReaderWithReadManyByPartitionKey.scala +++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/ItemsPartitionReaderWithReadManyByPartitionKey.scala @@ -6,7 +6,7 @@ package com.azure.cosmos.spark import com.azure.cosmos.{CosmosAsyncContainer, CosmosEndToEndOperationLatencyPolicyConfigBuilder, CosmosItemSerializer, CosmosItemSerializerNoExceptionWrapping, SparkBridgeInternal} import com.azure.cosmos.implementation.spark.OperationContextAndListenerTuple import com.azure.cosmos.implementation.{ImplementationBridgeHelpers, ObjectNodeMap, SparkRowItem, Utils} -import com.azure.cosmos.models.{CosmosReadManyByPartitionKeyRequestOptions, ModelBridgeInternal, PartitionKey, PartitionKeyDefinition, SqlQuerySpec} +import com.azure.cosmos.models.{CosmosReadManyByPartitionKeysRequestOptions, ModelBridgeInternal, PartitionKey, PartitionKeyDefinition, SqlQuerySpec} import com.azure.cosmos.spark.BulkWriter.getThreadInfo import com.azure.cosmos.spark.CosmosTableSchemaInferrer.IdAttributeName import com.azure.cosmos.spark.diagnostics.{DetailedFeedDiagnosticsProvider, DiagnosticsContext, DiagnosticsLoader, LoggerHelper, SparkTaskContext} @@ -43,10 +43,10 @@ private[spark] case class ItemsPartitionReaderWithReadManyByPartitionKey private lazy val log = LoggerHelper.getLogger(diagnosticsConfig, this.getClass) - private val readManyOptions = new CosmosReadManyByPartitionKeyRequestOptions() + private val readManyOptions = new CosmosReadManyByPartitionKeysRequestOptions() private val readManyOptionsImpl = ImplementationBridgeHelpers - .CosmosReadManyByPartitionKeyRequestOptionsHelper - .getCosmosReadManyByPartitionKeyRequestOptionsAccessor + .CosmosReadManyByPartitionKeysRequestOptionsHelper + .getCosmosReadManyByPartitionKeysRequestOptionsAccessor .getImpl(readManyOptions) private val readConfig = CosmosReadConfig.parseCosmosReadConfig(config) diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/ReadManyByPartitionKeyTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/ReadManyByPartitionKeyTest.java index 247c868ec6da..965e2f3fc31e 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/ReadManyByPartitionKeyTest.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/ReadManyByPartitionKeyTest.java @@ -456,7 +456,7 @@ public void singlePk_readManyByPartitionKey_withRequestOptions() { List items = createSinglePkItems("pkOpts", 3); List pkValues = Collections.singletonList(new PartitionKey("pkOpts")); - com.azure.cosmos.models.CosmosReadManyByPartitionKeyRequestOptions options = new com.azure.cosmos.models.CosmosReadManyByPartitionKeyRequestOptions(); + com.azure.cosmos.models.CosmosReadManyByPartitionKeysRequestOptions options = new com.azure.cosmos.models.CosmosReadManyByPartitionKeysRequestOptions(); AtomicInteger deserializeCount = new AtomicInteger(); options.setCustomItemSerializer(new CosmosItemSerializerNoExceptionWrapping() { @Override @@ -484,14 +484,14 @@ public T deserialize(Map jsonNodeMap, Class classType) { @Test(groups = {"emulator"}, timeOut = TIMEOUT) public void singlePk_readManyByPartitionKey_withRequestOptionsAndMaxConcurrentBatchPrefetch() { - // Regression test: passing non-null CosmosReadManyByPartitionKeyRequestOptions + // Regression test: passing non-null CosmosReadManyByPartitionKeysRequestOptions // with maxConcurrentBatchPrefetch set should not throw NullPointerException // from auto-unboxing a null MaxDegreeOfParallelism during options cloning. List items = createSinglePkItems("pkMdop", 3); List pkValues = Collections.singletonList(new PartitionKey("pkMdop")); - com.azure.cosmos.models.CosmosReadManyByPartitionKeyRequestOptions options = - new com.azure.cosmos.models.CosmosReadManyByPartitionKeyRequestOptions(); + com.azure.cosmos.models.CosmosReadManyByPartitionKeysRequestOptions options = + new com.azure.cosmos.models.CosmosReadManyByPartitionKeysRequestOptions(); options.setMaxConcurrentBatchPrefetch(2); CosmosPagedIterable results = singlePkContainer.readManyByPartitionKeys( @@ -606,8 +606,8 @@ public void singlePk_continuationToken_resumesCorrectly() { List itemsFromFirstPage = allPages.get(0).getResults(); // Second pass: resume from the continuation token - com.azure.cosmos.models.CosmosReadManyByPartitionKeyRequestOptions options2 = - new com.azure.cosmos.models.CosmosReadManyByPartitionKeyRequestOptions(); + com.azure.cosmos.models.CosmosReadManyByPartitionKeysRequestOptions options2 = + new com.azure.cosmos.models.CosmosReadManyByPartitionKeysRequestOptions(); options2.setContinuationToken(continuationAfterFirstPage); List> remainingPages = asyncContainer @@ -685,8 +685,8 @@ public void singlePk_continuationToken_resumesCorrectly_whenInputContainsDuplica assertThat(continuationAfterFirstPage).isNotNull(); List itemsFromFirstPage = allPages.get(0).getResults(); - com.azure.cosmos.models.CosmosReadManyByPartitionKeyRequestOptions options2 = - new com.azure.cosmos.models.CosmosReadManyByPartitionKeyRequestOptions(); + com.azure.cosmos.models.CosmosReadManyByPartitionKeysRequestOptions options2 = + new com.azure.cosmos.models.CosmosReadManyByPartitionKeysRequestOptions(); options2.setContinuationToken(continuationAfterFirstPage); List> remainingPages = asyncContainer diff --git a/sdk/cosmos/azure-cosmos/CHANGELOG.md b/sdk/cosmos/azure-cosmos/CHANGELOG.md index 36ba92cdee2f..05949f653fbe 100644 --- a/sdk/cosmos/azure-cosmos/CHANGELOG.md +++ b/sdk/cosmos/azure-cosmos/CHANGELOG.md @@ -5,7 +5,7 @@ #### Features Added * Added support for change feed with `startFrom` point-in-time on merged partitions by enabling the `CHANGE_FEED_WITH_START_TIME_POST_MERGE` SDK capability. - See [PR 48752](https://github.com/Azure/azure-sdk-for-java/pull/48752) * Added new `readManyByPartitionKeys` API on `CosmosAsyncContainer` / `CosmosContainer` to bulk-query all documents matching a list of partition key values with better efficiency than issuing individual queries. See [PR 48801](https://github.com/Azure/azure-sdk-for-java/pull/48801) -* Added `CosmosReadManyByPartitionKeyRequestOptions` - a dedicated request-options type for `readManyByPartitionKeys` that exposes `setContinuationToken(String)` for resuming previous invocations and `setMaxConcurrentBatchPrefetch(int)` to bound per-call prefetch parallelism. See [PR 48801](https://github.com/Azure/azure-sdk-for-java/pull/48801) +* Added `CosmosReadManyByPartitionKeysRequestOptions` - a dedicated request-options type for `readManyByPartitionKeys` that exposes `setContinuationToken(String)` for resuming previous invocations and `setMaxConcurrentBatchPrefetch(int)` to bound per-call prefetch parallelism. See [PR 48801](https://github.com/Azure/azure-sdk-for-java/pull/48801) #### Breaking Changes diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java index 4546c2db4632..041e2b68aeb0 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java @@ -60,7 +60,7 @@ import com.azure.cosmos.models.CosmosPatchItemRequestOptions; import com.azure.cosmos.models.CosmosPatchOperations; import com.azure.cosmos.models.CosmosQueryRequestOptions; -import com.azure.cosmos.models.CosmosReadManyByPartitionKeyRequestOptions; +import com.azure.cosmos.models.CosmosReadManyByPartitionKeysRequestOptions; import com.azure.cosmos.models.CosmosReadManyRequestOptions; import com.azure.cosmos.models.FeedRange; import com.azure.cosmos.models.FeedResponse; @@ -127,8 +127,8 @@ private static ImplementationBridgeHelpers.CosmosReadManyRequestOptionsHelper.Co return ImplementationBridgeHelpers.CosmosReadManyRequestOptionsHelper.getCosmosReadManyRequestOptionsAccessor(); } - private static ImplementationBridgeHelpers.CosmosReadManyByPartitionKeyRequestOptionsHelper.CosmosReadManyByPartitionKeyRequestOptionsAccessor readManyByPkOptionsAccessor() { - return ImplementationBridgeHelpers.CosmosReadManyByPartitionKeyRequestOptionsHelper.getCosmosReadManyByPartitionKeyRequestOptionsAccessor(); + private static ImplementationBridgeHelpers.CosmosReadManyByPartitionKeysRequestOptionsHelper.CosmosReadManyByPartitionKeysRequestOptionsAccessor readManyByPkOptionsAccessor() { + return ImplementationBridgeHelpers.CosmosReadManyByPartitionKeysRequestOptionsHelper.getCosmosReadManyByPartitionKeysRequestOptionsAccessor(); } private static ImplementationBridgeHelpers.CosmosDiagnosticsContextHelper.CosmosDiagnosticsContextAccessor ctxAccessor() { @@ -1648,7 +1648,7 @@ public CosmosPagedFlux readManyByPartitionKeys( */ public CosmosPagedFlux readManyByPartitionKeys( List partitionKeys, - CosmosReadManyByPartitionKeyRequestOptions requestOptions, + CosmosReadManyByPartitionKeysRequestOptions requestOptions, Class classType) { return this.readManyByPartitionKeys(partitionKeys, null, requestOptions, classType); @@ -1714,7 +1714,7 @@ public CosmosPagedFlux readManyByPartitionKeys( public CosmosPagedFlux readManyByPartitionKeys( List partitionKeys, SqlQuerySpec customQuery, - CosmosReadManyByPartitionKeyRequestOptions requestOptions, + CosmosReadManyByPartitionKeysRequestOptions requestOptions, Class classType) { if (partitionKeys == null) { @@ -1739,7 +1739,7 @@ public CosmosPagedFlux readManyByPartitionKeys( private Function>> readManyByPartitionKeysInternalFunc( List partitionKeys, SqlQuerySpec customQuery, - CosmosReadManyByPartitionKeyRequestOptions requestOptions, + CosmosReadManyByPartitionKeysRequestOptions requestOptions, Class classType) { CosmosAsyncClient client = this.getDatabase().getClient(); @@ -1770,7 +1770,7 @@ private Function>> readManyByPa // value but default the uninitialized case to unbounded for readManyByPartitionKeys. // CosmosReadManyRequestOptions does not currently expose MDOP, so this only matters // if it is plumbed through in the future. - // When cloning from CosmosReadManyByPartitionKeyRequestOptionsImpl (which does not + // When cloning from CosmosReadManyByPartitionKeysRequestOptionsImpl (which does not // define maxDegreeOfParallelism), the cloned CosmosQueryRequestOptions may have a // null backing Integer. Guard against NPE from auto-unboxing by checking for null. Integer mdop = queryOptionsAccessor().getImpl(queryRequestOptions).getMaxDegreeOfParallelism(); diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosContainer.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosContainer.java index a733dc272079..d0146a75157b 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosContainer.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosContainer.java @@ -21,7 +21,7 @@ import com.azure.cosmos.models.CosmosPatchItemRequestOptions; import com.azure.cosmos.models.CosmosPatchOperations; import com.azure.cosmos.models.CosmosQueryRequestOptions; -import com.azure.cosmos.models.CosmosReadManyByPartitionKeyRequestOptions; +import com.azure.cosmos.models.CosmosReadManyByPartitionKeysRequestOptions; import com.azure.cosmos.models.CosmosReadManyRequestOptions; import com.azure.cosmos.models.FeedRange; import com.azure.cosmos.models.FeedResponse; @@ -575,7 +575,7 @@ public CosmosPagedIterable readManyByPartitionKeys( */ public CosmosPagedIterable readManyByPartitionKeys( List partitionKeys, - CosmosReadManyByPartitionKeyRequestOptions requestOptions, + CosmosReadManyByPartitionKeysRequestOptions requestOptions, Class classType) { return getCosmosPagedIterable(this.asyncContainer.readManyByPartitionKeys(partitionKeys, requestOptions, classType)); @@ -635,7 +635,7 @@ public CosmosPagedIterable readManyByPartitionKeys( public CosmosPagedIterable readManyByPartitionKeys( List partitionKeys, SqlQuerySpec customQuery, - CosmosReadManyByPartitionKeyRequestOptions requestOptions, + CosmosReadManyByPartitionKeysRequestOptions requestOptions, Class classType) { return getCosmosPagedIterable(this.asyncContainer.readManyByPartitionKeys(partitionKeys, customQuery, requestOptions, classType)); diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/CosmosReadManyByPartitionKeyRequestOptionsImpl.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/CosmosReadManyByPartitionKeysRequestOptionsImpl.java similarity index 81% rename from sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/CosmosReadManyByPartitionKeyRequestOptionsImpl.java rename to sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/CosmosReadManyByPartitionKeysRequestOptionsImpl.java index 0480beaa3e4d..115e390a47ae 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/CosmosReadManyByPartitionKeyRequestOptionsImpl.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/CosmosReadManyByPartitionKeysRequestOptionsImpl.java @@ -3,21 +3,21 @@ package com.azure.cosmos.implementation; /** - * Internal implementation backing the public {@code CosmosReadManyByPartitionKeyRequestOptions} + * Internal implementation backing the public {@code CosmosReadManyByPartitionKeysRequestOptions} * facade. Holds state specific to the {@code readManyByPartitionKeys} operation. */ -public class CosmosReadManyByPartitionKeyRequestOptionsImpl - extends CosmosQueryRequestOptionsBase { +public class CosmosReadManyByPartitionKeysRequestOptionsImpl + extends CosmosQueryRequestOptionsBase { private String continuationToken; private Integer maxConcurrentBatchPrefetch; private Integer maxItemCount; - public CosmosReadManyByPartitionKeyRequestOptionsImpl() { + public CosmosReadManyByPartitionKeysRequestOptionsImpl() { super(); } - public CosmosReadManyByPartitionKeyRequestOptionsImpl(CosmosReadManyByPartitionKeyRequestOptionsImpl options) { + public CosmosReadManyByPartitionKeysRequestOptionsImpl(CosmosReadManyByPartitionKeysRequestOptionsImpl options) { super(options); this.continuationToken = options.continuationToken; this.maxConcurrentBatchPrefetch = options.maxConcurrentBatchPrefetch; @@ -39,7 +39,7 @@ public String getContinuationToken() { * @param continuationToken the continuation token from a previous invocation. * @return this instance. */ - public CosmosReadManyByPartitionKeyRequestOptionsImpl setContinuationToken(String continuationToken) { + public CosmosReadManyByPartitionKeysRequestOptionsImpl setContinuationToken(String continuationToken) { this.continuationToken = continuationToken; return this; } @@ -61,7 +61,7 @@ public Integer getMaxConcurrentBatchPrefetch() { * @param maxConcurrentBatchPrefetch the max concurrent batch prefetch (must be >= 1). * @return this instance. */ - public CosmosReadManyByPartitionKeyRequestOptionsImpl setMaxConcurrentBatchPrefetch(int maxConcurrentBatchPrefetch) { + public CosmosReadManyByPartitionKeysRequestOptionsImpl setMaxConcurrentBatchPrefetch(int maxConcurrentBatchPrefetch) { this.maxConcurrentBatchPrefetch = maxConcurrentBatchPrefetch; return this; } @@ -71,7 +71,7 @@ public Integer getMaxItemCount() { return this.maxItemCount; } - public CosmosReadManyByPartitionKeyRequestOptionsImpl setMaxItemCount(Integer maxItemCount) { + public CosmosReadManyByPartitionKeysRequestOptionsImpl setMaxItemCount(Integer maxItemCount) { this.maxItemCount = maxItemCount; return this; } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ImplementationBridgeHelpers.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ImplementationBridgeHelpers.java index 0be7b33d967b..ab882d89cf64 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ImplementationBridgeHelpers.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ImplementationBridgeHelpers.java @@ -69,7 +69,7 @@ import com.azure.cosmos.models.CosmosPatchItemRequestOptions; import com.azure.cosmos.models.CosmosPatchOperations; import com.azure.cosmos.models.CosmosQueryRequestOptions; -import com.azure.cosmos.models.CosmosReadManyByPartitionKeyRequestOptions; +import com.azure.cosmos.models.CosmosReadManyByPartitionKeysRequestOptions; import com.azure.cosmos.models.CosmosReadManyRequestOptions; import com.azure.cosmos.models.FeedRange; import com.azure.cosmos.models.FeedResponse; @@ -355,40 +355,40 @@ public interface CosmosReadManyRequestOptionsAccessor { } } - public static final class CosmosReadManyByPartitionKeyRequestOptionsHelper { + public static final class CosmosReadManyByPartitionKeysRequestOptionsHelper { private final static AtomicBoolean cosmosReadManyByPkRequestOptionsClassLoaded = new AtomicBoolean(false); - private final static AtomicReference accessor = new AtomicReference<>(); + private final static AtomicReference accessor = new AtomicReference<>(); - private CosmosReadManyByPartitionKeyRequestOptionsHelper() {} + private CosmosReadManyByPartitionKeysRequestOptionsHelper() {} - public static void setCosmosReadManyByPartitionKeyRequestOptionsAccessor( - final CosmosReadManyByPartitionKeyRequestOptionsAccessor newAccessor) { + public static void setCosmosReadManyByPartitionKeysRequestOptionsAccessor( + final CosmosReadManyByPartitionKeysRequestOptionsAccessor newAccessor) { if (!accessor.compareAndSet(null, newAccessor)) { - logger.debug("CosmosReadManyByPartitionKeyRequestOptionsAccessor already initialized!"); + logger.debug("CosmosReadManyByPartitionKeysRequestOptionsAccessor already initialized!"); } else { - logger.debug("Setting CosmosReadManyByPartitionKeyRequestOptionsAccessor..."); + logger.debug("Setting CosmosReadManyByPartitionKeysRequestOptionsAccessor..."); cosmosReadManyByPkRequestOptionsClassLoaded.set(true); } } - public static CosmosReadManyByPartitionKeyRequestOptionsAccessor getCosmosReadManyByPartitionKeyRequestOptionsAccessor() { + public static CosmosReadManyByPartitionKeysRequestOptionsAccessor getCosmosReadManyByPartitionKeysRequestOptionsAccessor() { if (!cosmosReadManyByPkRequestOptionsClassLoaded.get()) { - logger.debug("Initializing CosmosReadManyByPartitionKeyRequestOptionsAccessor..."); + logger.debug("Initializing CosmosReadManyByPartitionKeysRequestOptionsAccessor..."); initializeAllAccessors(); } - CosmosReadManyByPartitionKeyRequestOptionsAccessor snapshot = accessor.get(); + CosmosReadManyByPartitionKeysRequestOptionsAccessor snapshot = accessor.get(); if (snapshot == null) { - logger.error("CosmosReadManyByPartitionKeyRequestOptionsAccessor is not initialized yet!"); + logger.error("CosmosReadManyByPartitionKeysRequestOptionsAccessor is not initialized yet!"); } return snapshot; } - public interface CosmosReadManyByPartitionKeyRequestOptionsAccessor { - CosmosQueryRequestOptionsBase getImpl(CosmosReadManyByPartitionKeyRequestOptions options); - String getContinuationToken(CosmosReadManyByPartitionKeyRequestOptions options); - Integer getMaxConcurrentBatchPrefetch(CosmosReadManyByPartitionKeyRequestOptions options); + public interface CosmosReadManyByPartitionKeysRequestOptionsAccessor { + CosmosQueryRequestOptionsBase getImpl(CosmosReadManyByPartitionKeysRequestOptions options); + String getContinuationToken(CosmosReadManyByPartitionKeysRequestOptions options); + Integer getMaxConcurrentBatchPrefetch(CosmosReadManyByPartitionKeysRequestOptions options); } } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/CosmosReadManyByPartitionKeyRequestOptions.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/CosmosReadManyByPartitionKeysRequestOptions.java similarity index 78% rename from sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/CosmosReadManyByPartitionKeyRequestOptions.java rename to sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/CosmosReadManyByPartitionKeysRequestOptions.java index 9a02eeb1259d..30d1cf1f8a07 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/CosmosReadManyByPartitionKeyRequestOptions.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/CosmosReadManyByPartitionKeysRequestOptions.java @@ -8,7 +8,7 @@ import com.azure.cosmos.CosmosItemSerializer; import com.azure.cosmos.ReadConsistencyStrategy; import com.azure.cosmos.implementation.CosmosQueryRequestOptionsBase; -import com.azure.cosmos.implementation.CosmosReadManyByPartitionKeyRequestOptionsImpl; +import com.azure.cosmos.implementation.CosmosReadManyByPartitionKeysRequestOptionsImpl; import com.azure.cosmos.implementation.ImplementationBridgeHelpers; import com.azure.cosmos.util.Beta; @@ -26,14 +26,14 @@ * query parallelism inside a single physical partition or a feed range filter are * intentionally not exposed because the operation is fully managed by the SDK. */ -public final class CosmosReadManyByPartitionKeyRequestOptions { - private final CosmosReadManyByPartitionKeyRequestOptionsImpl actualRequestOptions; +public final class CosmosReadManyByPartitionKeysRequestOptions { + private final CosmosReadManyByPartitionKeysRequestOptionsImpl actualRequestOptions; /** * Instantiates a new readManyByPartitionKeys request options. */ - public CosmosReadManyByPartitionKeyRequestOptions() { - this.actualRequestOptions = new CosmosReadManyByPartitionKeyRequestOptionsImpl(); + public CosmosReadManyByPartitionKeysRequestOptions() { + this.actualRequestOptions = new CosmosReadManyByPartitionKeysRequestOptionsImpl(); } /** @@ -41,8 +41,8 @@ public CosmosReadManyByPartitionKeyRequestOptions() { * * @param options the options to copy. */ - CosmosReadManyByPartitionKeyRequestOptions(CosmosReadManyByPartitionKeyRequestOptions options) { - this.actualRequestOptions = new CosmosReadManyByPartitionKeyRequestOptionsImpl(options.actualRequestOptions); + CosmosReadManyByPartitionKeysRequestOptions(CosmosReadManyByPartitionKeysRequestOptions options) { + this.actualRequestOptions = new CosmosReadManyByPartitionKeysRequestOptionsImpl(options.actualRequestOptions); } /** @@ -62,9 +62,9 @@ public String getContinuationToken() { * partition-key set and the same custom query. * * @param continuationToken the composite continuation token from a previous invocation. - * @return the {@link CosmosReadManyByPartitionKeyRequestOptions} for fluent chaining. + * @return the {@link CosmosReadManyByPartitionKeysRequestOptions} for fluent chaining. */ - public CosmosReadManyByPartitionKeyRequestOptions setContinuationToken(String continuationToken) { + public CosmosReadManyByPartitionKeysRequestOptions setContinuationToken(String continuationToken) { this.actualRequestOptions.setContinuationToken(continuationToken); return this; } @@ -89,10 +89,10 @@ public Integer getMaxConcurrentBatchPrefetch() { * network/CPU and additional prefetch only adds memory pressure. * * @param maxConcurrentBatchPrefetch the max concurrent batch prefetch (must be >= 1). - * @return the {@link CosmosReadManyByPartitionKeyRequestOptions} for fluent chaining. + * @return the {@link CosmosReadManyByPartitionKeysRequestOptions} for fluent chaining. * @throws IllegalArgumentException if {@code maxConcurrentBatchPrefetch} is < 1. */ - public CosmosReadManyByPartitionKeyRequestOptions setMaxConcurrentBatchPrefetch(int maxConcurrentBatchPrefetch) { + public CosmosReadManyByPartitionKeysRequestOptions setMaxConcurrentBatchPrefetch(int maxConcurrentBatchPrefetch) { if (maxConcurrentBatchPrefetch < 1) { throw new IllegalArgumentException( "Argument 'maxConcurrentBatchPrefetch' must be greater than or equal to 1."); @@ -115,10 +115,10 @@ public ReadConsistencyStrategy getReadConsistencyStrategy() { * Sets the read consistency strategy required for the request. * * @param readConsistencyStrategy the read consistency strategy. - * @return the {@link CosmosReadManyByPartitionKeyRequestOptions} for fluent chaining. + * @return the {@link CosmosReadManyByPartitionKeysRequestOptions} for fluent chaining. */ @Beta(value = Beta.SinceVersion.V4_69_0, warningText = Beta.PREVIEW_SUBJECT_TO_CHANGE_WARNING) - public CosmosReadManyByPartitionKeyRequestOptions setReadConsistencyStrategy( + public CosmosReadManyByPartitionKeysRequestOptions setReadConsistencyStrategy( ReadConsistencyStrategy readConsistencyStrategy) { this.actualRequestOptions.setReadConsistencyStrategy(readConsistencyStrategy); return this; @@ -137,9 +137,9 @@ public String getSessionToken() { * Sets the session token for use with session consistency. * * @param sessionToken the session token. - * @return the {@link CosmosReadManyByPartitionKeyRequestOptions} for fluent chaining. + * @return the {@link CosmosReadManyByPartitionKeysRequestOptions} for fluent chaining. */ - public CosmosReadManyByPartitionKeyRequestOptions setSessionToken(String sessionToken) { + public CosmosReadManyByPartitionKeysRequestOptions setSessionToken(String sessionToken) { this.actualRequestOptions.setSessionToken(sessionToken); return this; } @@ -153,9 +153,9 @@ public CosmosReadManyByPartitionKeyRequestOptions setSessionToken(String session * also carries the remaining batch definitions, query hash, and partition-key-set hash. * * @param limitInKb backend continuation token size limit (must be >= 1). - * @return the {@link CosmosReadManyByPartitionKeyRequestOptions} for fluent chaining. + * @return the {@link CosmosReadManyByPartitionKeysRequestOptions} for fluent chaining. */ - public CosmosReadManyByPartitionKeyRequestOptions setMaxBackendContinuationTokenSizeInKb(int limitInKb) { + public CosmosReadManyByPartitionKeysRequestOptions setMaxBackendContinuationTokenSizeInKb(int limitInKb) { this.actualRequestOptions.setResponseContinuationTokenLimitInKb(limitInKb); return this; } @@ -183,9 +183,9 @@ public Integer getMaxItemCount() { * Sets the maximum number of items returned in a single page. * * @param maxItemCount the maximum number of items per page. - * @return the {@link CosmosReadManyByPartitionKeyRequestOptions} for fluent chaining. + * @return the {@link CosmosReadManyByPartitionKeysRequestOptions} for fluent chaining. */ - public CosmosReadManyByPartitionKeyRequestOptions setMaxItemCount(int maxItemCount) { + public CosmosReadManyByPartitionKeysRequestOptions setMaxItemCount(int maxItemCount) { this.actualRequestOptions.setMaxItemCount(maxItemCount); return this; } @@ -194,9 +194,9 @@ public CosmosReadManyByPartitionKeyRequestOptions setMaxItemCount(int maxItemCou * Sets the {@link CosmosEndToEndOperationLatencyPolicyConfig} to be used for the request. * * @param cosmosEndToEndOperationLatencyPolicyConfig the {@link CosmosEndToEndOperationLatencyPolicyConfig} - * @return the {@link CosmosReadManyByPartitionKeyRequestOptions} for fluent chaining. + * @return the {@link CosmosReadManyByPartitionKeysRequestOptions} for fluent chaining. */ - public CosmosReadManyByPartitionKeyRequestOptions setCosmosEndToEndOperationLatencyPolicyConfig( + public CosmosReadManyByPartitionKeysRequestOptions setCosmosEndToEndOperationLatencyPolicyConfig( CosmosEndToEndOperationLatencyPolicyConfig cosmosEndToEndOperationLatencyPolicyConfig) { this.actualRequestOptions .setCosmosEndToEndOperationLatencyPolicyConfig(cosmosEndToEndOperationLatencyPolicyConfig); @@ -207,9 +207,9 @@ public CosmosReadManyByPartitionKeyRequestOptions setCosmosEndToEndOperationLate * List of regions to be excluded for the request/retries. * * @param excludeRegions the regions to exclude - * @return the {@link CosmosReadManyByPartitionKeyRequestOptions} + * @return the {@link CosmosReadManyByPartitionKeysRequestOptions} */ - public CosmosReadManyByPartitionKeyRequestOptions setExcludedRegions(List excludeRegions) { + public CosmosReadManyByPartitionKeysRequestOptions setExcludedRegions(List excludeRegions) { this.actualRequestOptions.setExcludedRegions(excludeRegions); return this; } @@ -236,9 +236,9 @@ public boolean isQueryMetricsEnabled() { * Sets the option to enable/disable query metrics. * * @param queryMetricsEnabled whether to enable or disable query metrics - * @return the {@link CosmosReadManyByPartitionKeyRequestOptions} for fluent chaining. + * @return the {@link CosmosReadManyByPartitionKeysRequestOptions} for fluent chaining. */ - public CosmosReadManyByPartitionKeyRequestOptions setQueryMetricsEnabled(boolean queryMetricsEnabled) { + public CosmosReadManyByPartitionKeysRequestOptions setQueryMetricsEnabled(boolean queryMetricsEnabled) { this.actualRequestOptions.setQueryMetricsEnabled(queryMetricsEnabled); return this; } @@ -256,9 +256,9 @@ public String getThroughputControlGroupName() { * Sets the throughput control group name. * * @param throughputControlGroupName the throughput control group name. - * @return the {@link CosmosReadManyByPartitionKeyRequestOptions} for fluent chaining. + * @return the {@link CosmosReadManyByPartitionKeysRequestOptions} for fluent chaining. */ - public CosmosReadManyByPartitionKeyRequestOptions setThroughputControlGroupName(String throughputControlGroupName) { + public CosmosReadManyByPartitionKeysRequestOptions setThroughputControlGroupName(String throughputControlGroupName) { this.actualRequestOptions.setThroughputControlGroupName(throughputControlGroupName); return this; } @@ -276,9 +276,9 @@ public DedicatedGatewayRequestOptions getDedicatedGatewayRequestOptions() { * Sets the dedicated gateway request options. * * @param dedicatedGatewayRequestOptions the dedicated gateway request options. - * @return the {@link CosmosReadManyByPartitionKeyRequestOptions} for fluent chaining. + * @return the {@link CosmosReadManyByPartitionKeysRequestOptions} for fluent chaining. */ - public CosmosReadManyByPartitionKeyRequestOptions setDedicatedGatewayRequestOptions( + public CosmosReadManyByPartitionKeysRequestOptions setDedicatedGatewayRequestOptions( DedicatedGatewayRequestOptions dedicatedGatewayRequestOptions) { this.actualRequestOptions.setDedicatedGatewayRequestOptions(dedicatedGatewayRequestOptions); return this; @@ -297,9 +297,9 @@ public Duration getThresholdForDiagnosticsOnTracer() { * Sets the latency threshold for diagnostics on tracer. * * @param thresholdForDiagnosticsOnTracer the latency threshold for diagnostics on tracer. - * @return the {@link CosmosReadManyByPartitionKeyRequestOptions} for fluent chaining. + * @return the {@link CosmosReadManyByPartitionKeysRequestOptions} for fluent chaining. */ - public CosmosReadManyByPartitionKeyRequestOptions setThresholdForDiagnosticsOnTracer( + public CosmosReadManyByPartitionKeysRequestOptions setThresholdForDiagnosticsOnTracer( Duration thresholdForDiagnosticsOnTracer) { this.actualRequestOptions.setThresholdForDiagnosticsOnTracer(thresholdForDiagnosticsOnTracer); return this; @@ -309,9 +309,9 @@ public CosmosReadManyByPartitionKeyRequestOptions setThresholdForDiagnosticsOnTr * Allows overriding the diagnostic thresholds for a specific operation. * * @param operationSpecificThresholds the diagnostic threshold override for this operation - * @return the {@link CosmosReadManyByPartitionKeyRequestOptions} for fluent chaining. + * @return the {@link CosmosReadManyByPartitionKeysRequestOptions} for fluent chaining. */ - public CosmosReadManyByPartitionKeyRequestOptions setDiagnosticsThresholds( + public CosmosReadManyByPartitionKeysRequestOptions setDiagnosticsThresholds( CosmosDiagnosticsThresholds operationSpecificThresholds) { this.actualRequestOptions.setDiagnosticsThresholds(operationSpecificThresholds); return this; @@ -339,9 +339,9 @@ public CosmosItemSerializer getCustomItemSerializer() { * Sets a custom item serializer to be used for this operation. * * @param customItemSerializer the custom item serializer for this operation. - * @return the {@link CosmosReadManyByPartitionKeyRequestOptions} for fluent chaining. + * @return the {@link CosmosReadManyByPartitionKeysRequestOptions} for fluent chaining. */ - public CosmosReadManyByPartitionKeyRequestOptions setCustomItemSerializer( + public CosmosReadManyByPartitionKeysRequestOptions setCustomItemSerializer( CosmosItemSerializer customItemSerializer) { this.actualRequestOptions.setCustomItemSerializer(customItemSerializer); return this; @@ -351,9 +351,9 @@ public CosmosReadManyByPartitionKeyRequestOptions setCustomItemSerializer( * Sets the custom keyword identifiers. * * @param keywordIdentifiers the custom keyword identifiers. - * @return the {@link CosmosReadManyByPartitionKeyRequestOptions} for fluent chaining. + * @return the {@link CosmosReadManyByPartitionKeysRequestOptions} for fluent chaining. */ - public CosmosReadManyByPartitionKeyRequestOptions setKeywordIdentifiers(Set keywordIdentifiers) { + public CosmosReadManyByPartitionKeysRequestOptions setKeywordIdentifiers(Set keywordIdentifiers) { this.actualRequestOptions.setKeywordIdentifiers(keywordIdentifiers); return this; } @@ -375,24 +375,24 @@ CosmosQueryRequestOptionsBase getImpl() { // the following helper/accessor only helps to access this class outside of this package.// /////////////////////////////////////////////////////////////////////////////////////////// static void initialize() { - ImplementationBridgeHelpers.CosmosReadManyByPartitionKeyRequestOptionsHelper - .setCosmosReadManyByPartitionKeyRequestOptionsAccessor( - new ImplementationBridgeHelpers.CosmosReadManyByPartitionKeyRequestOptionsHelper - .CosmosReadManyByPartitionKeyRequestOptionsAccessor() { + ImplementationBridgeHelpers.CosmosReadManyByPartitionKeysRequestOptionsHelper + .setCosmosReadManyByPartitionKeysRequestOptionsAccessor( + new ImplementationBridgeHelpers.CosmosReadManyByPartitionKeysRequestOptionsHelper + .CosmosReadManyByPartitionKeysRequestOptionsAccessor() { @Override public CosmosQueryRequestOptionsBase getImpl( - CosmosReadManyByPartitionKeyRequestOptions options) { + CosmosReadManyByPartitionKeysRequestOptions options) { return options.actualRequestOptions; } @Override - public String getContinuationToken(CosmosReadManyByPartitionKeyRequestOptions options) { + public String getContinuationToken(CosmosReadManyByPartitionKeysRequestOptions options) { return options.actualRequestOptions.getContinuationToken(); } @Override public Integer getMaxConcurrentBatchPrefetch( - CosmosReadManyByPartitionKeyRequestOptions options) { + CosmosReadManyByPartitionKeysRequestOptions options) { return options.actualRequestOptions.getMaxConcurrentBatchPrefetch(); } }); diff --git a/sdk/cosmos/docs/readManyByPartitionKey-design.md b/sdk/cosmos/docs/readManyByPartitionKey-design.md index 1075c29cd8ce..210a98657942 100644 --- a/sdk/cosmos/docs/readManyByPartitionKey-design.md +++ b/sdk/cosmos/docs/readManyByPartitionKey-design.md @@ -57,21 +57,21 @@ flowchart TD ```java CosmosPagedFlux readManyByPartitionKeys(List partitionKeys, Class classType) CosmosPagedFlux readManyByPartitionKeys(List partitionKeys, - CosmosReadManyRequestOptions requestOptions, + CosmosReadManyByPartitionKeysRequestOptions requestOptions, Class classType) CosmosPagedFlux readManyByPartitionKeys(List partitionKeys, SqlQuerySpec customQuery, Class classType) CosmosPagedFlux readManyByPartitionKeys(List partitionKeys, SqlQuerySpec customQuery, - CosmosReadManyByPartitionKeyRequestOptions requestOptions, + CosmosReadManyByPartitionKeysRequestOptions requestOptions, Class classType) ``` All four overloads delegate to a private `readManyByPartitionKeyInternalFunc(...)`. > **Note:** the request-options type for `readManyByPartitionKeys` is the dedicated -> `CosmosReadManyByPartitionKeyRequestOptions` class (not the shared +> `CosmosReadManyByPartitionKeysRequestOptions` class (not the shared > `CosmosReadManyRequestOptions` used by `readMany(List)`). It exposes > `setContinuationToken(String)` / `getContinuationToken()` directly on the public surface > and a `setMaxConcurrentBatchPrefetch(int)` knob that defaults to @@ -429,11 +429,11 @@ sequenceDiagram ### API changes -**`CosmosReadManyByPartitionKeyRequestOptionsImpl`:** Holds `continuationToken`, +**`CosmosReadManyByPartitionKeysRequestOptionsImpl`:** Holds `continuationToken`, `maxConcurrentBatchPrefetch`, and `maxItemCount` fields with corresponding getters and setters; all other knobs (consistency, session token, regions, throughput control, dedicated gateway, diagnostics, custom item serializer, keyword identifiers, e2e latency policy) are inherited -from `CosmosQueryRequestOptionsBase`. The public `CosmosReadManyByPartitionKeyRequestOptions` +from `CosmosQueryRequestOptionsBase`. The public `CosmosReadManyByPartitionKeysRequestOptions` facade exposes only the surface that is meaningful for this operation - `MaxBackendContinuationTokenSizeInKb` (renamed from the inherited `ResponseContinuationTokenLimitInKb` to make clear it limits the *backend* token wrapped inside From 9c75a32a1ce8ca8941014b1e64b1a212cf9cf188 Mon Sep 17 00:00:00 2001 From: Fabian Meiswinkel Date: Wed, 22 Apr 2026 23:06:02 +0000 Subject: [PATCH 52/64] Code review fixes --- .../com/azure/cosmos/spark/CosmosConfig.scala | 4 +- .../cosmos/spark/CosmosItemsDataSource.scala | 17 ++---- .../CosmosReadManyByPartitionKeyReader.scala | 6 +-- ...tionReaderWithReadManyByPartitionKey.scala | 27 +++++++--- ...tryingReadManyByPartitionKeyIterator.scala | 6 +++ ...ReadManyByPartitionKeyQueryHelperTest.java | 52 +++++++++++++++++++ ...adManyByPartitionKeyContinuationToken.java | 34 ++++++++---- .../implementation/RxDocumentClientImpl.java | 2 +- .../com/azure/cosmos/models/FeedResponse.java | 2 +- .../docs/readManyByPartitionKey-design.md | 1 + 10 files changed, 115 insertions(+), 36 deletions(-) diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosConfig.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosConfig.scala index 75b47274cd2b..990d18d0b487 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosConfig.scala +++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosConfig.scala @@ -1160,10 +1160,10 @@ private object CosmosReadConfig { key = CosmosConfigNames.ReadManyByPkMaxConcurrentBatchPrefetch, mandatory = false, defaultValue = Some(1), - parseFromStringFunction = value => Math.max(1, value.toInt), + parseFromStringFunction = value => Math.min(64, Math.max(1, value.toInt)), helpMessage = "The maximum number of per-physical-partition batches whose first page is prefetched " + "concurrently inside a single Spark task by the SDK's readManyByPartitionKeys execution. The " + - "default is `1`, because Spark already parallelises across tasks - increase this when individual " + + "default is `1` - max is `64`, because Spark already parallelises across tasks - increase this when individual " + "tasks span many physical partitions and additional intra-task prefetch is desired." ) diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosItemsDataSource.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosItemsDataSource.scala index ef310298f29a..7723b566db12 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosItemsDataSource.scala +++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosItemsDataSource.scala @@ -126,18 +126,11 @@ object CosmosItemsDataSource { userProvidedSchema, userConfig.asScala.toMap) - // Resolve the null-handling config up front so both the UDF path and the PK-column path honor it. - val sharedEffectiveConfig = CosmosConfig.getEffectiveConfig( - databaseName = None, - containerName = None, - userConfig.asScala.toMap) - val sharedReadConfig = CosmosReadConfig.parseCosmosReadConfig(sharedEffectiveConfig) - val sharedTreatNullAsNone = sharedReadConfig.readManyByPkTreatNullAsNone - - // Initialize reader state once: resolves PK paths, infers schema, and broadcasts client caches - // in a single Loan block instead of three separate ones. + // Initialize reader state once: resolves PK paths, infers schema, broadcasts client caches, + // and returns the resolved treatNullAsNone flag from the reader's parsed config - avoiding + // duplicate config parsing between the data source and the reader. val readerState = readManyReader.initializeReaderState() - val (pkPaths, _, _) = readerState + val (pkPaths, _, _, sharedTreatNullAsNone) = readerState // Option 1: Look for the _partitionKeyIdentity column (produced by GetCosmosPartitionKeyValue UDF) val pkIdentityFieldExtraction = df @@ -155,7 +148,7 @@ object CosmosItemsDataSource { val pkColumnExtraction: Option[Row => PartitionKey] = if (pkIdentityFieldExtraction.isDefined) { None // no need to resolve PK paths - _partitionKeyIdentity column takes precedence } else { - val treatNullAsNone = sharedReadConfig.readManyByPkTreatNullAsNone + val treatNullAsNone = sharedTreatNullAsNone // Nested PK paths (containing /) cannot be resolved from top-level DataFrame columns. if (pkPaths.exists(_.contains("/"))) { diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosReadManyByPartitionKeyReader.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosReadManyByPartitionKeyReader.scala index 0bfd8339b4ec..9e3f2287742f 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosReadManyByPartitionKeyReader.scala +++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosReadManyByPartitionKeyReader.scala @@ -51,7 +51,7 @@ private[spark] class CosmosReadManyByPartitionKeyReader( * * @return (pkPaths, schema, broadcastClientStates) */ - private[spark] def initializeReaderState(): (List[String], StructType, Broadcast[CosmosClientMetadataCachesSnapshots]) = { + private[spark] def initializeReaderState(): (List[String], StructType, Broadcast[CosmosClientMetadataCachesSnapshots], Boolean) = { val calledFrom = s"CosmosReadManyByPartitionKeyReader($tableName).initializeReaderState" Loan( List[Option[CosmosClientCacheItem]]( @@ -120,14 +120,14 @@ private[spark] class CosmosReadManyByPartitionKeyReader( val metadataSnapshots = CosmosClientMetadataCachesSnapshots(state, throughputControlState) val broadcastStates = sparkSession.sparkContext.broadcast(metadataSnapshots) - (pkPaths, schema, broadcastStates) + (pkPaths, schema, broadcastStates, readConfig.readManyByPkTreatNullAsNone) }) } def readManyByPartitionKeys( inputRdd: RDD[Row], pkExtraction: Row => PartitionKey, - readerState: (List[String], StructType, Broadcast[CosmosClientMetadataCachesSnapshots])): DataFrame = { + readerState: (List[String], StructType, Broadcast[CosmosClientMetadataCachesSnapshots], Boolean)): DataFrame = { val correlationActivityId = UUIDs.nonBlockingRandomUUID() val (_, schema, clientStates) = readerState diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/ItemsPartitionReaderWithReadManyByPartitionKey.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/ItemsPartitionReaderWithReadManyByPartitionKey.scala index a8cbadba5f70..7a1a915443a8 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/ItemsPartitionReaderWithReadManyByPartitionKey.scala +++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/ItemsPartitionReaderWithReadManyByPartitionKey.scala @@ -229,18 +229,29 @@ private[spark] case class ItemsPartitionReaderWithReadManyByPartitionKey // Factory that creates a CosmosPagedFlux from an optional continuation token. // On the first call continuationToken is null (start from scratch); on retry // it is the continuation token from the last fully-drained page. - // NOTE: mutating the shared readManyOptions before each call is safe because - // the SDK clones the options internally at the start of each call, and the - // iterator is single-threaded (mutation in step N completes before step N+1). + // A fresh CosmosReadManyByPartitionKeysRequestOptions instance is created per + // call to avoid mutating the shared readManyOptions object, which would be + // fragile if the SDK ever stopped cloning options internally. private val fluxFactory: String => CosmosPagedFlux[SparkRowItem] = { (continuationToken: String) => - // Reuse the preconfigured request options and only update the continuation - // token so the SDK resumes from the last fully committed page. - readManyOptions.setContinuationToken(continuationToken) + val perCallOptions = new CosmosReadManyByPartitionKeysRequestOptions() + perCallOptions.setMaxConcurrentBatchPrefetch(readConfig.readManyByPkMaxConcurrentBatchPrefetch) + perCallOptions.setContinuationToken(continuationToken) + val perCallOptionsImpl = ImplementationBridgeHelpers + .CosmosReadManyByPartitionKeysRequestOptionsHelper + .getCosmosReadManyByPartitionKeysRequestOptionsAccessor + .getImpl(perCallOptions) + ThroughputControlHelper.populateThroughputControlGroupName(perCallOptionsImpl, readConfig.throughputControlConfig) + perCallOptionsImpl.setCosmosEndToEndOperationLatencyPolicyConfig(endToEndTimeoutPolicy) + if (operationContextAndListenerTuple.isDefined) { + perCallOptionsImpl.setOperationContextAndListenerTuple(operationContextAndListenerTuple.get) + } + perCallOptionsImpl + .setCustomItemSerializer(readManyOptionsImpl.getCustomItemSerializer) customQueryOpt match { case Some(query) => - cosmosAsyncContainer.readManyByPartitionKeys(pkList, query, readManyOptions, classOf[SparkRowItem]) + cosmosAsyncContainer.readManyByPartitionKeys(pkList, query, perCallOptions, classOf[SparkRowItem]) case None => - cosmosAsyncContainer.readManyByPartitionKeys(pkList, readManyOptions, classOf[SparkRowItem]) + cosmosAsyncContainer.readManyByPartitionKeys(pkList, perCallOptions, classOf[SparkRowItem]) } } diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/TransientIOErrorsRetryingReadManyByPartitionKeyIterator.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/TransientIOErrorsRetryingReadManyByPartitionKeyIterator.scala index 3b7e96e8cd45..6779a4117dd8 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/TransientIOErrorsRetryingReadManyByPartitionKeyIterator.scala +++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/TransientIOErrorsRetryingReadManyByPartitionKeyIterator.scala @@ -27,6 +27,12 @@ import scala.collection.JavaConverters._ * FeedResponse and used as the resume point on retry. This avoids the correctness issues of * page-count-based skipping (where the server is not guaranteed to return the same page * boundaries across requests). + * + * Note: transient I/O failures can only occur during {@code hasNext} (which fetches the next + * FeedResponse page from the network). Once a page has been fetched, iterating over its + * in-memory items ({@code next()}) performs no I/O and cannot trigger a retry. Therefore, + * partially-consumed pages are never replayed and the iterator provides exactly-once + * delivery in practice. */ private[spark] class TransientIOErrorsRetryingReadManyByPartitionKeyIterator[TSparkRow] ( diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelperTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelperTest.java index 9987a184e420..4a9bbc0c2da9 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelperTest.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelperTest.java @@ -275,6 +275,58 @@ public void normalizePartitionKeys_noneInHpk_alwaysRejected() { .hasMessageContaining("PartitionKey.NONE is not supported for multi-path partition keys"); } + + @Test(groups = { "unit" }) + public void normalizePartitionKeys_twoPartialPrefixesAtDifferentDepths() { + // Two partial PKs at different prefix depths: (Redmond) and (Redmond, 98052). + // (Redmond) is a prefix of (Redmond, 98052), so (Redmond, 98052) should be collapsed. + PartitionKeyDefinition pkDef = createMultiHashPkDefinition("/city", "/zipcode", "/areaCode"); + + List pkValues = Arrays.asList( + new PartitionKeyBuilder().add("Redmond").build(), + new PartitionKeyBuilder().add("Redmond").add("98052").build()); + + List normalized = + RxDocumentClientImpl.normalizePartitionKeys(pkValues, pkDef); + + assertThat(normalized) + .extracting(normalizedPk -> normalizedPk.effectivePkInternal.toJson()) + .containsExactly("[\"Redmond\"]"); + } + + @Test(groups = { "unit" }) + public void normalizePartitionKeys_partialPkDoesNotSubsumeUnrelatedPk() { + // (Redmond) and (Seattle, 98052) are unrelated - neither subsumes the other. + PartitionKeyDefinition pkDef = createMultiHashPkDefinition("/city", "/zipcode", "/areaCode"); + + List pkValues = Arrays.asList( + new PartitionKeyBuilder().add("Redmond").build(), + new PartitionKeyBuilder().add("Seattle").add("98052").build()); + + List normalized = + RxDocumentClientImpl.normalizePartitionKeys(pkValues, pkDef); + + assertThat(normalized).hasSize(2); + assertThat(normalized) + .extracting(normalizedPk -> normalizedPk.effectivePkInternal.toJson()) + .containsExactly("[\"Redmond\"]", "[\"Seattle\",\"98052\"]"); + } + + @Test(groups = { "unit" }) + public void normalizePartitionKeys_allPartialPksPreserved() { + // All PKs are partial (single-level) and unrelated - all should survive. + PartitionKeyDefinition pkDef = createMultiHashPkDefinition("/city", "/zipcode", "/areaCode"); + + List pkValues = Arrays.asList( + new PartitionKeyBuilder().add("Redmond").build(), + new PartitionKeyBuilder().add("Seattle").build(), + new PartitionKeyBuilder().add("Pittsburgh").build()); + + List normalized = + RxDocumentClientImpl.normalizePartitionKeys(pkValues, pkDef); + + assertThat(normalized).hasSize(3); + } //endregion //region findTopLevelWhereIndex tests diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyContinuationToken.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyContinuationToken.java index 813bab25453d..75449d7af4be 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyContinuationToken.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyContinuationToken.java @@ -228,22 +228,38 @@ public static String computeQueryHash(SqlQuerySpec querySpec) { * Duplicate and reordered inputs intentionally produce the same digest. */ public static String computePartitionKeySetHash(List partitionKeyEpks) { + return computePartitionKeySetHash(partitionKeyEpks, false); + } + + /** + * Computes a stable hash for the normalized set of partition key EPK values. + * When {@code alreadySorted} is true, the input is assumed to be + * already sorted (e.g. from {@code normalizePartitionKeys}), + * skipping the O(n log n) sort pass. + */ + public static String computePartitionKeySetHash(List partitionKeyEpks, boolean alreadySorted) { if (partitionKeyEpks == null || partitionKeyEpks.isEmpty()) { return "0"; } - List normalizedEpks = new ArrayList<>(partitionKeyEpks.size()); - for (String epk : partitionKeyEpks) { - if (epk != null) { - normalizedEpks.add(epk); + List normalizedEpks = null; + if (alreadySorted) { + normalizedEpks = partitionKeyEpks; + } else { + + normalizedEpks = new ArrayList<>(partitionKeyEpks.size()); + for (String epk : partitionKeyEpks) { + if (epk != null) { + normalizedEpks.add(epk); + } } - } - if (normalizedEpks.isEmpty()) { - return "0"; - } + if (normalizedEpks.isEmpty()) { + return "0"; + } - Collections.sort(normalizedEpks); + Collections.sort(normalizedEpks); + } ByteArrayOutputStream output = new ByteArrayOutputStream(); String previous = null; diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java index 079e67e4102a..6a4338f1ccd2 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java @@ -4562,7 +4562,7 @@ static List normalizePartitionKeys( List partitionKeys, PartitionKeyDefinition pkDefinition) { - Map normalizedByEpk = new LinkedHashMap<>(); + Map normalizedByEpk = new HashMap<>(); for (PartitionKey pk : partitionKeys) { PartitionKeyInternal pkInternal = BridgeInternal.getPartitionKeyInternal(pk); diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/FeedResponse.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/FeedResponse.java index 959cca07b2b4..da8764c95f85 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/FeedResponse.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/FeedResponse.java @@ -160,7 +160,7 @@ private FeedResponse( // NOTE - it is important to use HashMap over ConcurrentHashMap here because some keys/values might be null // and this is not allowed in ConcurrentHashMap - while it is ok in HashMap - this.header = toBeCloned.header != null ? new HashMap<>(toBeCloned.header) : null; + this.header = toBeCloned.header != null ? new HashMap<>(toBeCloned.header) : new HashMap<>(); this.usageHeaders = toBeCloned.usageHeaders != null ? new HashMap<>(toBeCloned.usageHeaders) : null; this.quotaHeaders = toBeCloned.quotaHeaders != null ? new HashMap<>(toBeCloned.quotaHeaders) : null; diff --git a/sdk/cosmos/docs/readManyByPartitionKey-design.md b/sdk/cosmos/docs/readManyByPartitionKey-design.md index 210a98657942..7edc808ba854 100644 --- a/sdk/cosmos/docs/readManyByPartitionKey-design.md +++ b/sdk/cosmos/docs/readManyByPartitionKey-design.md @@ -147,6 +147,7 @@ regular SDK queries. - `queryInfo.hasDCount()` - `queryInfo.hasOffset()` - `queryInfo.hasLimit()` + - `queryInfo.hasTop()` - `queryInfo.hasNonStreamingOrderBy()` - `partitionedQueryExecutionInfo.hasHybridSearchQueryInfo()` - query plan details are unavailable (`queryInfo == null`) From e8ea36817c9d25af920f3fc62393269caab5e675 Mon Sep 17 00:00:00 2001 From: Fabian Meiswinkel Date: Wed, 22 Apr 2026 23:35:21 +0000 Subject: [PATCH 53/64] Fixed test issues --- .../CosmosReadManyByPartitionKeyReader.scala | 2 +- ...ReadManyByPartitionKeyQueryHelperTest.java | 83 ++++++++++++++++++- .../azure/cosmos/CosmosAsyncContainer.java | 11 +++ .../ReadManyByPartitionKeyQueryHelper.java | 5 +- .../cosmos/models/ModelBridgeInternal.java | 1 + 5 files changed, 95 insertions(+), 7 deletions(-) diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosReadManyByPartitionKeyReader.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosReadManyByPartitionKeyReader.scala index 9e3f2287742f..e3ea4298b3f3 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosReadManyByPartitionKeyReader.scala +++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosReadManyByPartitionKeyReader.scala @@ -130,7 +130,7 @@ private[spark] class CosmosReadManyByPartitionKeyReader( readerState: (List[String], StructType, Broadcast[CosmosClientMetadataCachesSnapshots], Boolean)): DataFrame = { val correlationActivityId = UUIDs.nonBlockingRandomUUID() - val (_, schema, clientStates) = readerState + val (_, schema, clientStates, _) = readerState sparkSession.sqlContext.createDataFrame( inputRdd.mapPartitionsWithIndex( diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelperTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelperTest.java index 4a9bbc0c2da9..c40a25108e31 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelperTest.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelperTest.java @@ -36,7 +36,7 @@ public void singlePk_defaultQuery_singleValue() { SqlQuerySpec result = ReadManyByPartitionKeyQueryHelper.createReadManyByPkQuerySpec( "SELECT * FROM c", new ArrayList<>(), pkValues, selectors, pkDef); - assertThat(result.getQueryText()).contains("SELECT * FROM c WHERE"); + assertThat(result.getQueryText()).contains("SELECT * FROM c\n WHERE"); assertThat(result.getQueryText()).contains("IN ("); assertThat(result.getQueryText()).contains("@__rmPk_0"); assertThat(result.getParameters()).hasSize(1); @@ -71,7 +71,7 @@ public void singlePk_customQuery_noWhere() { SqlQuerySpec result = ReadManyByPartitionKeyQueryHelper.createReadManyByPkQuerySpec( "SELECT c.name, c.age FROM c", new ArrayList<>(), pkValues, selectors, pkDef); - assertThat(result.getQueryText()).startsWith("SELECT c.name, c.age FROM c WHERE"); + assertThat(result.getQueryText()).startsWith("SELECT c.name, c.age FROM c\n WHERE"); assertThat(result.getQueryText()).contains("IN ("); } @@ -109,7 +109,7 @@ public void hpk_fullPk_defaultQuery() { SqlQuerySpec result = ReadManyByPartitionKeyQueryHelper.createReadManyByPkQuerySpec( "SELECT * FROM c", new ArrayList<>(), pkValues, selectors, pkDef); - assertThat(result.getQueryText()).contains("SELECT * FROM c WHERE"); + assertThat(result.getQueryText()).contains("SELECT * FROM c\n WHERE"); // Should use OR/AND pattern, not IN assertThat(result.getQueryText()).doesNotContain("IN ("); assertThat(result.getQueryText()).contains("c[\"city\"] = @__rmPk_0"); @@ -378,7 +378,7 @@ public void singlePk_customAlias() { SqlQuerySpec result = ReadManyByPartitionKeyQueryHelper.createReadManyByPkQuerySpec( "SELECT x.id, x.mypk FROM x", new ArrayList<>(), pkValues, selectors, pkDef); - assertThat(result.getQueryText()).startsWith("SELECT x.id, x.mypk FROM x WHERE"); + assertThat(result.getQueryText()).startsWith("SELECT x.id, x.mypk FROM x\n WHERE"); assertThat(result.getQueryText()).contains("x[\"mypk\"] IN ("); assertThat(result.getQueryText()).doesNotContain("c[\"mypk\"]"); } @@ -686,6 +686,81 @@ public void hpk_nonePartitionKey_throwsForMultiHash() { //endregion + //region Trailing single-line comment tests + + @Test(groups = { "unit" }) + public void singlePk_customQuery_trailingSingleLineComment_noWhere() { + // A trailing -- comment must not swallow the appended WHERE clause + PartitionKeyDefinition pkDef = createSinglePkDefinition("/mypk"); + List selectors = createSelectors(pkDef); + List pkValues = Collections.singletonList(new PartitionKey("pk1")); + + SqlQuerySpec result = ReadManyByPartitionKeyQueryHelper.createReadManyByPkQuerySpec( + "SELECT * FROM c -- my note", new ArrayList<>(), pkValues, selectors, pkDef); + + // WHERE must appear on a new line so it is not inside the -- comment + assertThat(result.getQueryText()).contains("\n WHERE"); + assertThat(result.getQueryText()).contains("IN ("); + assertThat(result.getQueryText()).contains("@__rmPk_0"); + assertThat(result.getParameters()).hasSize(1); + } + + @Test(groups = { "unit" }) + public void singlePk_customQuery_trailingBlockComment_noWhere() { + // Block comments are not affected, but verify the query still works + PartitionKeyDefinition pkDef = createSinglePkDefinition("/mypk"); + List selectors = createSelectors(pkDef); + List pkValues = Collections.singletonList(new PartitionKey("pk1")); + + SqlQuerySpec result = ReadManyByPartitionKeyQueryHelper.createReadManyByPkQuerySpec( + "SELECT * FROM c /* note */", new ArrayList<>(), pkValues, selectors, pkDef); + + assertThat(result.getQueryText()).contains("WHERE"); + assertThat(result.getQueryText()).contains("IN ("); + assertThat(result.getParameters()).hasSize(1); + } + + @Test(groups = { "unit" }) + public void hpk_customQuery_trailingSingleLineComment_noWhere() { + // HPK path with trailing -- comment. + // Note: extractTableAlias does not skip comments, so "-- comment" after the + // FROM alias is mis-parsed as an alias token. This test verifies the \n WHERE + // fix and parameter correctness; alias-aware comment handling is a separate concern. + PartitionKeyDefinition pkDef = createMultiHashPkDefinition("/city", "/zipcode"); + List selectors = createSelectors(pkDef); + + PartitionKey pk = new PartitionKeyBuilder().add("Redmond").add("98052").build(); + List pkValues = Collections.singletonList(pk); + + SqlQuerySpec result = ReadManyByPartitionKeyQueryHelper.createReadManyByPkQuerySpec( + "SELECT * FROM c -- WHERE comment", new ArrayList<>(), pkValues, selectors, pkDef); + + assertThat(result.getQueryText()).contains("\n WHERE"); + assertThat(result.getQueryText()).contains("@__rmPk_0"); + assertThat(result.getQueryText()).contains("@__rmPk_1"); + assertThat(result.getParameters()).hasSize(2); + } + + @Test(groups = { "unit" }) + public void singlePk_customQuery_trailingComment_withExistingWhere() { + // When the query already has a WHERE clause, trailing comment after the condition + // is handled by the AND-merge path, not the no-WHERE path + PartitionKeyDefinition pkDef = createSinglePkDefinition("/mypk"); + List selectors = createSelectors(pkDef); + List pkValues = Collections.singletonList(new PartitionKey("pk1")); + + List baseParams = new ArrayList<>(); + baseParams.add(new SqlParameter("@minAge", 18)); + + SqlQuerySpec result = ReadManyByPartitionKeyQueryHelper.createReadManyByPkQuerySpec( + "SELECT * FROM c WHERE c.age > @minAge -- filter", baseParams, pkValues, selectors, pkDef); + + assertThat(result.getQueryText()).contains("WHERE (c.age > @minAge -- filter) AND ("); + assertThat(result.getQueryText()).contains("IN ("); + } + + //endregion + //region helpers private PartitionKeyDefinition createSinglePkDefinition(String path) { PartitionKeyDefinition pkDef = new PartitionKeyDefinition(); diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java index 041e2b68aeb0..1225773903d4 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java @@ -1777,6 +1777,17 @@ private Function>> readManyByPa if (mdop == null || mdop == DEFAULT_MAX_DEGREE_OF_PARALLELISM) { queryRequestOptions.setMaxDegreeOfParallelism(UNBOUNDED_MAX_DEGREE_OF_PARALLELISM); } + + // maxItemCount lives in CosmosQueryRequestOptionsImpl but is not copied by the + // clone path through CosmosQueryRequestOptionsBase, so propagate it explicitly. + Integer maxItems = queryOptionsAccessor().getImpl(queryRequestOptions).getMaxItemCount(); + if (maxItems == null && requestOptions != null) { + Integer userMaxItems = readManyByPkOptionsAccessor().getImpl(requestOptions).getMaxItemCount(); + if (userMaxItems != null) { + ModelBridgeInternal.setQueryRequestOptionsMaxItemCount(queryRequestOptions, userMaxItems); + } + } + queryRequestOptions.setQueryName("readManyByPartitionKeys"); // Set the composite continuation token on the cloned query options so that diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelper.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelper.java index 3b0d0f490165..6eadd103cad3 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelper.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelper.java @@ -145,8 +145,9 @@ public static SqlQuerySpec createReadManyByPkQuerySpec( String afterWhere = baseQueryText.substring(whereIndex + 5); // skip "WHERE" finalQuery = beforeWhere + "WHERE (" + afterWhere.trim() + ") AND (" + pkFilter.toString().trim() + ")"; } else { - // No WHERE - add one - finalQuery = baseQueryText + " WHERE" + pkFilter.toString(); + // No WHERE - add one. Use \n before WHERE so that a trailing single-line comment + // (-- ...) in the base query does not swallow the WHERE clause. + finalQuery = baseQueryText + "\n WHERE" + pkFilter.toString(); } return new SqlQuerySpec(finalQuery, parameters); diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/ModelBridgeInternal.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/ModelBridgeInternal.java index c34df7c4e925..7dde81e06e98 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/ModelBridgeInternal.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/ModelBridgeInternal.java @@ -758,6 +758,7 @@ public static void initializeAllAccessors() { CosmosItemRequestOptions.initialize(); CosmosItemResponse.initialize(); CosmosPatchOperations.initialize(); + CosmosReadManyByPartitionKeysRequestOptions.initialize(); CosmosReadManyRequestOptions.initialize(); CosmosQueryRequestOptions.initialize(); CosmosOperationDetails.initialize(); From f3aa1daaf284c8b8e169df3853d58f2e3666babb Mon Sep 17 00:00:00 2001 From: Scott Beddall Date: Thu, 23 Apr 2026 00:57:14 +0000 Subject: [PATCH 54/64] updates to patch sparse checkout to account for spring and cosmos while they are still excluded from the build list --- .../Generate-ServiceDirectories-From-Project-List.ps1 | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/eng/scripts/Generate-ServiceDirectories-From-Project-List.ps1 b/eng/scripts/Generate-ServiceDirectories-From-Project-List.ps1 index 1541fc32845e..7f944a986bbe 100644 --- a/eng/scripts/Generate-ServiceDirectories-From-Project-List.ps1 +++ b/eng/scripts/Generate-ServiceDirectories-From-Project-List.ps1 @@ -77,6 +77,12 @@ foreach($file in Get-ChildItem -Path $SourcesDirectory -Filter pom*.xml -Recurse $sparseCheckoutDirectories = @() $serviceDirectories = @() $sparseCheckoutDirectories += "/sdk/parents" +# in place until we can come up with a better path to honoring ExcludePaths +# in archetype-sdk-client. Given that the validation for these is triggered for outside paths, but `analyze` still needs to +# process those folders, we'll just include them in the sparse checkout and service directory lists for now. +# This is only two additional folders and it allows us to avoid a lot of complexity around trying to special case the ExcludePaths in archetype-sdk-client. +$sparseCheckoutDirectories += "/sdk/cosmos" +$sparseCheckoutDirectories += "/sdk/spring" foreach ($project in $ProjectList) { if ($sparseCheckoutDirHash.ContainsKey($project)) { $sparseCheckoutDirectories += $sparseCheckoutDirHash[$project] From 2c404e664111ee2d95a4b61b3c21e9782386f76e Mon Sep 17 00:00:00 2001 From: Scott Beddall <45376673+scbedd@users.noreply.github.com> Date: Wed, 22 Apr 2026 18:10:50 -0700 Subject: [PATCH 55/64] Update eng/scripts/Generate-ServiceDirectories-From-Project-List.ps1 Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- eng/scripts/Generate-ServiceDirectories-From-Project-List.ps1 | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/eng/scripts/Generate-ServiceDirectories-From-Project-List.ps1 b/eng/scripts/Generate-ServiceDirectories-From-Project-List.ps1 index 7f944a986bbe..9ea7d818e29c 100644 --- a/eng/scripts/Generate-ServiceDirectories-From-Project-List.ps1 +++ b/eng/scripts/Generate-ServiceDirectories-From-Project-List.ps1 @@ -78,7 +78,7 @@ $sparseCheckoutDirectories = @() $serviceDirectories = @() $sparseCheckoutDirectories += "/sdk/parents" # in place until we can come up with a better path to honoring ExcludePaths -# in archetype-sdk-client. Given that the validation for these is triggered for outside paths, but `analyze` still needs to +# in archetype-sdk-client. Given that the validation for these is triggered for outside paths, but `analyze` still needs to # process those folders, we'll just include them in the sparse checkout and service directory lists for now. # This is only two additional folders and it allows us to avoid a lot of complexity around trying to special case the ExcludePaths in archetype-sdk-client. $sparseCheckoutDirectories += "/sdk/cosmos" From 112cf71fe607cf0c18a42f99e4834c9f807194b8 Mon Sep 17 00:00:00 2001 From: Fabian Meiswinkel Date: Thu, 23 Apr 2026 01:28:46 +0000 Subject: [PATCH 56/64] Fixing code review comments --- ...ReadManyByPartitionKeyQueryHelperTest.java | 98 ++++++++++++++++++- .../azure/cosmos/CosmosAsyncContainer.java | 15 +-- ...ManyByPartitionKeysRequestOptionsImpl.java | 2 +- .../ReadManyByPartitionKeyQueryHelper.java | 9 +- ...ReadManyByPartitionKeysRequestOptions.java | 2 +- 5 files changed, 109 insertions(+), 17 deletions(-) diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelperTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelperTest.java index c40a25108e31..d9cab76b16ac 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelperTest.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelperTest.java @@ -88,7 +88,7 @@ public void singlePk_customQuery_withExistingWhere() { "SELECT * FROM c WHERE c.age > @minAge", baseParams, pkValues, selectors, pkDef); // Should AND the PK filter to the existing WHERE clause - assertThat(result.getQueryText()).contains("WHERE (c.age > @minAge) AND ("); + assertThat(result.getQueryText()).contains("WHERE (c.age > @minAge\n) AND ("); assertThat(result.getQueryText()).contains("IN ("); assertThat(result.getParameters()).hasSize(2); // @minAge + @__rmPk_0 assertThat(result.getParameters().get(0).getName()).isEqualTo("@minAge"); @@ -191,7 +191,7 @@ public void hpk_customQuery_withWhere() { SqlQuerySpec result = ReadManyByPartitionKeyQueryHelper.createReadManyByPkQuerySpec( "SELECT c.name FROM c WHERE c.status = @status", baseParams, pkValues, selectors, pkDef); - assertThat(result.getQueryText()).contains("WHERE (c.status = @status) AND ("); + assertThat(result.getQueryText()).contains("WHERE (c.status = @status\n) AND ("); assertThat(result.getQueryText()).contains("c[\"city\"] = @__rmPk_0"); assertThat(result.getParameters()).hasSize(3); // @status + 2 pk params } @@ -395,7 +395,7 @@ public void singlePk_customAlias_withWhere() { SqlQuerySpec result = ReadManyByPartitionKeyQueryHelper.createReadManyByPkQuerySpec( "SELECT x.id, x.mypk FROM x WHERE x.category = @cat", baseParams, pkValues, selectors, pkDef); - assertThat(result.getQueryText()).contains("WHERE (x.category = @cat) AND (x[\"mypk\"] IN ("); + assertThat(result.getQueryText()).contains("WHERE (x.category = @cat\n) AND (x[\"mypk\"] IN ("); } @Test(groups = { "unit" }) @@ -500,7 +500,7 @@ public void singlePk_customQuery_withStringLiteralContainingParens() { "SELECT * FROM c WHERE c.msg = 'test(value)WHERE'", baseParams, pkValues, selectors, pkDef); // Should correctly AND the PK filter to the real WHERE clause - assertThat(result.getQueryText()).contains("WHERE (c.msg = 'test(value)WHERE') AND ("); + assertThat(result.getQueryText()).contains("WHERE (c.msg = 'test(value)WHERE'\n) AND ("); } @Test(groups = { "unit" }) @@ -755,10 +755,98 @@ public void singlePk_customQuery_trailingComment_withExistingWhere() { SqlQuerySpec result = ReadManyByPartitionKeyQueryHelper.createReadManyByPkQuerySpec( "SELECT * FROM c WHERE c.age > @minAge -- filter", baseParams, pkValues, selectors, pkDef); - assertThat(result.getQueryText()).contains("WHERE (c.age > @minAge -- filter) AND ("); + assertThat(result.getQueryText()).contains("WHERE (c.age > @minAge -- filter\n) AND ("); assertThat(result.getQueryText()).contains("IN ("); } + @Test(groups = { "unit" }) + public void singlePk_existingWhere_trailingLineComment_andClauseNotSwallowed() { + // Without the \n before ), the -- comment would swallow ") AND (" and the entire + // PK filter, producing a query that returns unfiltered results. + PartitionKeyDefinition pkDef = createSinglePkDefinition("/mypk"); + List selectors = createSelectors(pkDef); + List pkValues = Arrays.asList(new PartitionKey("pk1"), new PartitionKey("pk2")); + + List baseParams = new ArrayList<>(); + baseParams.add(new SqlParameter("@status", "active")); + + SqlQuerySpec result = ReadManyByPartitionKeyQueryHelper.createReadManyByPkQuerySpec( + "SELECT * FROM c WHERE c.status = @status -- only active", baseParams, pkValues, selectors, pkDef); + + // The \n must break the -- comment so that ) AND ( is on a new line + String queryText = result.getQueryText(); + assertThat(queryText).contains("-- only active\n) AND ("); + assertThat(queryText).contains("IN ("); + assertThat(queryText).contains("@__rmPk_0"); + assertThat(queryText).contains("@__rmPk_1"); + assertThat(result.getParameters()).hasSize(3); // @status + 2 pk params + } + + @Test(groups = { "unit" }) + public void hpk_existingWhere_trailingLineComment_andClauseNotSwallowed() { + // HPK variant: trailing -- in existing WHERE must not swallow the AND clause + PartitionKeyDefinition pkDef = createMultiHashPkDefinition("/city", "/zipcode"); + List selectors = createSelectors(pkDef); + + List baseParams = new ArrayList<>(); + baseParams.add(new SqlParameter("@active", true)); + + PartitionKey pk = new PartitionKeyBuilder().add("Redmond").add("98052").build(); + List pkValues = Collections.singletonList(pk); + + SqlQuerySpec result = ReadManyByPartitionKeyQueryHelper.createReadManyByPkQuerySpec( + "SELECT * FROM c WHERE c.active = @active -- check", baseParams, pkValues, selectors, pkDef); + + String queryText = result.getQueryText(); + assertThat(queryText).contains("-- check\n) AND ("); + assertThat(queryText).contains("c[\"city\"] = @__rmPk_0"); + assertThat(queryText).contains("c[\"zipcode\"] = @__rmPk_1"); + assertThat(result.getParameters()).hasSize(3); + } + + @Test(groups = { "unit" }) + public void singlePk_existingWhere_multipleLineComments_andClauseNotSwallowed() { + // WHERE clause with a -- comment mid-condition (on its own line) and at the end + PartitionKeyDefinition pkDef = createSinglePkDefinition("/mypk"); + List selectors = createSelectors(pkDef); + List pkValues = Collections.singletonList(new PartitionKey("pk1")); + + List baseParams = new ArrayList<>(); + baseParams.add(new SqlParameter("@minAge", 18)); + baseParams.add(new SqlParameter("@maxAge", 65)); + + SqlQuerySpec result = ReadManyByPartitionKeyQueryHelper.createReadManyByPkQuerySpec( + "SELECT * FROM c WHERE c.age > @minAge -- lower bound\n AND c.age < @maxAge -- upper bound", + baseParams, pkValues, selectors, pkDef); + + String queryText = result.getQueryText(); + // The final -- upper bound must be broken by \n before ) + assertThat(queryText).contains("-- upper bound\n) AND ("); + assertThat(queryText).contains("IN ("); + assertThat(queryText).contains("@__rmPk_0"); + assertThat(result.getParameters()).hasSize(3); + } + + @Test(groups = { "unit" }) + public void singlePk_existingWhere_blockComment_noNewlineNeeded() { + // Block comments (/* */) are terminated by */ and don't need \n protection, + // but the \n is harmless. Verify the query is still correct. + PartitionKeyDefinition pkDef = createSinglePkDefinition("/mypk"); + List selectors = createSelectors(pkDef); + List pkValues = Collections.singletonList(new PartitionKey("pk1")); + + List baseParams = new ArrayList<>(); + baseParams.add(new SqlParameter("@val", 42)); + + SqlQuerySpec result = ReadManyByPartitionKeyQueryHelper.createReadManyByPkQuerySpec( + "SELECT * FROM c WHERE c.x = @val /* note */", baseParams, pkValues, selectors, pkDef); + + String queryText = result.getQueryText(); + assertThat(queryText).contains("WHERE (c.x = @val /* note */\n) AND ("); + assertThat(queryText).contains("IN ("); + assertThat(result.getParameters()).hasSize(2); + } + //endregion //region helpers diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java index 1225773903d4..7dc2ca475330 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java @@ -1717,12 +1717,13 @@ public CosmosPagedFlux readManyByPartitionKeys( CosmosReadManyByPartitionKeysRequestOptions requestOptions, Class classType) { - if (partitionKeys == null) { - throw new IllegalArgumentException("Argument 'partitionKeys' must not be null."); - } + checkNotNull(partitionKeys, "Argument 'partitionKeys' must not be null."); + checkNotNull(classType, "Argument 'classType' must not be null."); + if (partitionKeys.isEmpty()) { throw new IllegalArgumentException("Argument 'partitionKeys' must not be empty."); } + for (PartitionKey pk : partitionKeys) { if (pk == null) { throw new IllegalArgumentException( @@ -1751,14 +1752,16 @@ private Function>> readManyByPa ? readManyByPkOptionsAccessor().getContinuationToken(requestOptions) : null; - // Resolve the max-concurrent-batch-prefetch knob; default to availableProcessors() so - // the SDK fully utilises CPU when the caller does not override. + // Resolve the max-concurrent-batch-prefetch knob. Unlike queryItems (which fans out + // many small reads across partitions), readManyByPartitionKeys typically returns large + // result sets per partition, so a conservative default avoids overwhelming individual + // partitions with parallel requests and spiking RU consumption. Integer prefetchOverride = requestOptions != null ? readManyByPkOptionsAccessor().getMaxConcurrentBatchPrefetch(requestOptions) : null; int maxConcurrentBatchPrefetch = prefetchOverride != null ? prefetchOverride - : Configs.getCPUCnt(); + : Math.max(1, Math.min(Configs.getCPUCnt(), 8)); return (pagedFluxOptions -> { CosmosQueryRequestOptions queryRequestOptions = requestOptions == null diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/CosmosReadManyByPartitionKeysRequestOptionsImpl.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/CosmosReadManyByPartitionKeysRequestOptionsImpl.java index 115e390a47ae..d92eb0d2f1ab 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/CosmosReadManyByPartitionKeysRequestOptionsImpl.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/CosmosReadManyByPartitionKeysRequestOptionsImpl.java @@ -108,6 +108,6 @@ public Integer getMaxPrefetchPageCount() { @Override public String getQueryNameOrDefault(String defaultQueryName) { - return null; + return defaultQueryName; } } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelper.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelper.java index 6eadd103cad3..f498a452fad3 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelper.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelper.java @@ -12,6 +12,7 @@ import java.util.ArrayList; import java.util.List; +import java.util.Locale; /** * Helper for constructing SqlQuerySpec instances for readManyByPartitionKeys operations. @@ -143,7 +144,7 @@ public static SqlQuerySpec createReadManyByPkQuerySpec( // Base query has WHERE - AND our PK filter String beforeWhere = baseQueryText.substring(0, whereIndex); String afterWhere = baseQueryText.substring(whereIndex + 5); // skip "WHERE" - finalQuery = beforeWhere + "WHERE (" + afterWhere.trim() + ") AND (" + pkFilter.toString().trim() + ")"; + finalQuery = beforeWhere + "WHERE (" + afterWhere.trim() + "\n) AND (" + pkFilter.toString().trim() + ")"; } else { // No WHERE - add one. Use \n before WHERE so that a trailing single-line comment // (-- ...) in the base query does not swallow the WHERE clause. @@ -159,7 +160,7 @@ public static SqlQuerySpec createReadManyByPkQuerySpec( * Returns the alias used after FROM (last token before WHERE or end of FROM clause). */ static String extractTableAlias(String queryText) { - String upper = queryText.toUpperCase(); + String upper = queryText.toUpperCase(Locale.ROOT); int fromIndex = findTopLevelKeywordIndex(upper, "FROM"); if (fromIndex < 0) { return DEFAULT_TABLE_ALIAS; @@ -237,8 +238,8 @@ private static boolean isFollowedByReservedKeyword(String remainingUpper) { * ignoring occurrences inside parentheses or string literals. */ static int findTopLevelKeywordIndex(String queryText, String keyword) { - String queryTextUpper = queryText.toUpperCase(); - String keywordUpper = keyword.toUpperCase(); + String queryTextUpper = queryText.toUpperCase(Locale.ROOT); + String keywordUpper = keyword.toUpperCase(Locale.ROOT); int depth = 0; int keyLen = keywordUpper.length(); int len = queryTextUpper.length(); diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/CosmosReadManyByPartitionKeysRequestOptions.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/CosmosReadManyByPartitionKeysRequestOptions.java index 30d1cf1f8a07..9cfe713ef69a 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/CosmosReadManyByPartitionKeysRequestOptions.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/CosmosReadManyByPartitionKeysRequestOptions.java @@ -82,7 +82,7 @@ public Integer getMaxConcurrentBatchPrefetch() { /** * Sets the maximum number of per-physical-partition batches whose first page is prefetched - * concurrently. The default is {@code Runtime.getRuntime().availableProcessors()}. + * concurrently. The default is {@code Math.min(1, Math.max(64, Runtime.getRuntime().availableProcessors()))}. *

* Increase this to trade memory for lower end-to-end latency on wide containers; decrease it * (e.g. to {@code 1}) when running in environments where a single task already saturates the From cbd8a9329fbf338c21861f00446c3ae64ccd59fa Mon Sep 17 00:00:00 2001 From: Fabian Meiswinkel Date: Thu, 23 Apr 2026 09:55:03 +0000 Subject: [PATCH 57/64] Fixes code review comments --- .../java/com/azure/cosmos/ReadManyByPartitionKeyTest.java | 2 +- .../main/java/com/azure/cosmos/CosmosAsyncContainer.java | 6 ++++++ .../azure/cosmos/implementation/RxDocumentClientImpl.java | 7 +++++-- .../CosmosReadManyByPartitionKeysRequestOptions.java | 2 +- 4 files changed, 13 insertions(+), 4 deletions(-) diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/ReadManyByPartitionKeyTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/ReadManyByPartitionKeyTest.java index 965e2f3fc31e..79f7800bf426 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/ReadManyByPartitionKeyTest.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/ReadManyByPartitionKeyTest.java @@ -350,7 +350,7 @@ public void rejectsGroupByWithAggregateQuery() { } } - @Test(groups = {"emulator"}, timeOut = TIMEOUT, expectedExceptions = IllegalArgumentException.class) + @Test(groups = {"emulator"}, timeOut = TIMEOUT, expectedExceptions = NullPointerException.class) public void rejectsNullPartitionKeyList() { singlePkContainer.readManyByPartitionKeys((List) null, ObjectNode.class); } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java index 7dc2ca475330..3bd56a58d385 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java @@ -1620,6 +1620,8 @@ private Mono> readManyInternal( * all documents matching the provided partition key values. Uses {@code SELECT * FROM c} * as the base query. Duplicate partition key inputs are normalized with set-based semantics * before batching, so repeated keys do not duplicate the results. + *

+ * {@link PartitionKey#NONE} is supported for single-path partition key containers only. * * @param the type parameter * @param partitionKeys list of partition key values to read documents for @@ -1639,6 +1641,8 @@ public CosmosPagedFlux readManyByPartitionKeys( * all documents matching the provided partition key values. Uses {@code SELECT * FROM c} * as the base query. Duplicate partition key inputs are normalized with set-based semantics * before batching, so repeated keys do not duplicate the results. + *

+ * {@link PartitionKey#NONE} is supported for single-path partition key containers only. * * @param the type parameter * @param partitionKeys list of partition key values to read documents for @@ -1667,6 +1671,8 @@ public CosmosPagedFlux readManyByPartitionKeys( * Partial hierarchical partition keys are supported and will fan out to multiple * physical partitions. Duplicate partition key inputs are normalized with set-based semantics * before batching. + *

+ * {@link PartitionKey#NONE} is supported for single-path partition key containers only. * * @param the type parameter * @param partitionKeys list of partition key values to read documents for diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java index 6a4338f1ccd2..2773ecdf6cf6 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java @@ -4682,8 +4682,11 @@ private Flux> buildSequentialFluxFromScratch( return Flux.empty(); } - // Sort batches by batchFilter minInclusive EPK for deterministic sequential processing - allBatches.sort(Comparator.comparing(bd -> bd.batchFilter.getMin())); + // Sort batches by batchFilter EPK range for deterministic sequential processing. + // Tie-break on maxExclusive so that prefix HPK batches sharing the same minInclusive + // (the prefix EPK) always have a stable order regardless of routing-map traversal order. + allBatches.sort(Comparator.comparing((BatchDescriptor bd) -> bd.batchFilter.getMin()) + .thenComparing(bd -> bd.batchFilter.getMax())); return buildSequentialBatchFlux( allBatches, null, diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/CosmosReadManyByPartitionKeysRequestOptions.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/CosmosReadManyByPartitionKeysRequestOptions.java index 9cfe713ef69a..8b46dcc9c987 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/CosmosReadManyByPartitionKeysRequestOptions.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/CosmosReadManyByPartitionKeysRequestOptions.java @@ -82,7 +82,7 @@ public Integer getMaxConcurrentBatchPrefetch() { /** * Sets the maximum number of per-physical-partition batches whose first page is prefetched - * concurrently. The default is {@code Math.min(1, Math.max(64, Runtime.getRuntime().availableProcessors()))}. + * concurrently. The default is {@code Math.max(1, Math.min(Runtime.getRuntime().availableProcessors(), 8))}. *

* Increase this to trade memory for lower end-to-end latency on wide containers; decrease it * (e.g. to {@code 1}) when running in environments where a single task already saturates the From cf678fad997868fed66948e8db4feabd5b55df63 Mon Sep 17 00:00:00 2001 From: Fabian Meiswinkel Date: Thu, 23 Apr 2026 10:17:53 +0000 Subject: [PATCH 58/64] Fixed code review comments --- .../cosmos/spark/CosmosPartitionKeyHelper.scala | 1 + ...PartitionReaderWithReadManyByPartitionKey.scala | 14 +++----------- .../spark/CosmosPartitionKeyHelperSpec.scala | 14 ++++++++++++++ 3 files changed, 18 insertions(+), 11 deletions(-) diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosPartitionKeyHelper.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosPartitionKeyHelper.scala index 68c51dde4c3d..234974212983 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosPartitionKeyHelper.scala +++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosPartitionKeyHelper.scala @@ -55,6 +55,7 @@ private[spark] object CosmosPartitionKeyHelper extends BasicLoggingTrait { def tryParsePartitionKey( cosmosPartitionKeyString: String, treatNullAsNone: Boolean): Option[PartitionKey] = { + require(cosmosPartitionKeyString != null, "Argument 'cosmosPartitionKeyString' must not be null.") cosmosPartitionKeyString match { case cosmosPartitionKeyStringRegx(pkValue) => scala.util.Try(Utils.parse(pkValue, classOf[Object])).toOption.flatMap { diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/ItemsPartitionReaderWithReadManyByPartitionKey.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/ItemsPartitionReaderWithReadManyByPartitionKey.scala index 7a1a915443a8..cfcb7397ec03 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/ItemsPartitionReaderWithReadManyByPartitionKey.scala +++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/ItemsPartitionReaderWithReadManyByPartitionKey.scala @@ -121,17 +121,9 @@ private[spark] case class ItemsPartitionReaderWithReadManyByPartitionKey readManyOptionsImpl .setCustomItemSerializer( new CosmosItemSerializerNoExceptionWrapping { - // The reader-side path of readManyByPartitionKeys never produces items that the - // SDK needs to serialize back out (no item bodies are sent on the wire other than - // SqlParameter values, which do not flow through this serializer). However, some - // generic SDK code paths (e.g. SqlParameter#getValueAsBytes via the configured - // serializer) will defensively call serialize() during diagnostics or during query - // plan handling. Delegating to the default Jackson-based serializer keeps those - // code paths working and avoids spurious UnsupportedOperationExceptions while we - // still own deserialization (the hot path) below. - override def serialize[T](item: T): util.Map[String, AnyRef] = { - CosmosItemSerializer.DEFAULT_SERIALIZER.serialize(item) - } + // The base class (CosmosItemSerializerNoExceptionWrapping) sets canSerialize = false, + // which prevents the SDK from calling serialize() — so no override is needed here. + // Only deserialization (the hot path) is customized below. override def deserialize[T](jsonNodeMap: util.Map[String, AnyRef], classType: Class[T]): T = { if (jsonNodeMap == null) { diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/CosmosPartitionKeyHelperSpec.scala b/sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/CosmosPartitionKeyHelperSpec.scala index ba7b37a27065..d1d543b8a59f 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/CosmosPartitionKeyHelperSpec.scala +++ b/sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/CosmosPartitionKeyHelperSpec.scala @@ -120,6 +120,20 @@ class CosmosPartitionKeyHelperSpec extends UnitSpec { error.getMessage should include("PartitionKey.None can't be used with multiple paths") } + it should "throw IllegalArgumentException for null input" in { + val error = the[IllegalArgumentException] thrownBy { + CosmosPartitionKeyHelper.tryParsePartitionKey(null) + } + error.getMessage should include("must not be null") + } + + it should "throw IllegalArgumentException for null input with treatNullAsNone" in { + val error = the[IllegalArgumentException] thrownBy { + CosmosPartitionKeyHelper.tryParsePartitionKey(null, treatNullAsNone = true) + } + error.getMessage should include("must not be null") + } + //scalastyle:on multiple.string.literals //scalastyle:on magic.number } From bbfa6981a6572d010ace724d3d4e31bd544cb358 Mon Sep 17 00:00:00 2001 From: Fabian Meiswinkel Date: Thu, 23 Apr 2026 10:51:33 +0000 Subject: [PATCH 59/64] Added e2e test for resuming readManyByPartitionKeys with continuation token --- ...ngReadManyByPartitionKeyIteratorSpec.scala | 356 ++++++++++++++++++ .../cosmos/ReadManyByPartitionKeyTest.java | 183 +++++++++ 2 files changed, 539 insertions(+) create mode 100644 sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/TransientIOErrorsRetryingReadManyByPartitionKeyIteratorSpec.scala diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/TransientIOErrorsRetryingReadManyByPartitionKeyIteratorSpec.scala b/sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/TransientIOErrorsRetryingReadManyByPartitionKeyIteratorSpec.scala new file mode 100644 index 000000000000..64047f41634c --- /dev/null +++ b/sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/TransientIOErrorsRetryingReadManyByPartitionKeyIteratorSpec.scala @@ -0,0 +1,356 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. +package com.azure.cosmos.spark + +import com.azure.cosmos.CosmosException +import com.azure.cosmos.implementation.SparkRowItem +import com.azure.cosmos.models.{FeedResponse, ModelBridgeInternal} +import com.azure.cosmos.spark.diagnostics.BasicLoggingTrait +import com.azure.cosmos.util.UtilBridgeInternal +import com.fasterxml.jackson.core.JsonProcessingException +import com.fasterxml.jackson.databind.ObjectMapper +import com.fasterxml.jackson.databind.node.ObjectNode +import reactor.core.publisher.Flux + +import java.util.concurrent.ConcurrentHashMap +import java.util.concurrent.atomic.AtomicLong + +// scalastyle:off underscore.import +import scala.collection.JavaConverters._ +// scalastyle:on underscore.import + +//scalastyle:off magic.number +//scalastyle:off multiple.string.literals +class TransientIOErrorsRetryingReadManyByPartitionKeyIteratorSpec extends UnitSpec with BasicLoggingTrait { + + private val rnd = scala.util.Random + private val pageSize = 2 + private val cosmosSerializationConfig = CosmosSerializationConfig( + SerializationInclusionModes.Always, + SerializationDateTimeConversionModes.Default + ) + + private val cosmosRowConverter = CosmosRowConverter.get(cosmosSerializationConfig) + private val objectMapper = new ObjectMapper + + "TransientIOErrors" should "be retried without duplicates or missing records" in { + + val pageCount = 30 + val transientErrorCount = new AtomicLong(0) + val iterator = new TransientIOErrorsRetryingReadManyByPartitionKeyIterator[SparkRowItem]( + continuationToken => generateMockedCosmosPagedFlux( + continuationToken, pageCount, 0.2, transientErrorCount, injectEmptyPages = false), + pageSize, + 1, + None, + classOf[SparkRowItem] + ) + iterator.maxRetryIntervalInMs = 5 + + val items = drainAll(iterator) + items.size shouldEqual (pageCount * pageSize) + + transientErrorCount.get > 0 shouldEqual true + assertNoDuplicates(items) + } + + "TransientIOErrors" should "be retried without duplicates when empty pages exist" in { + + val pageCount = 30 + val transientErrorCount = new AtomicLong(0) + val iterator = new TransientIOErrorsRetryingReadManyByPartitionKeyIterator[SparkRowItem]( + continuationToken => generateMockedCosmosPagedFlux( + continuationToken, pageCount, 0.2, transientErrorCount, injectEmptyPages = true), + pageSize, + 1, + None, + classOf[SparkRowItem] + ) + iterator.maxRetryIntervalInMs = 5 + + // Pages 2,4,6,8,10,12,14,16,18,20 are empty (10 empty pages out of first 20) + val items = drainAll(iterator) + items.size shouldEqual ((pageCount - 10) * pageSize) + + transientErrorCount.get > 0 shouldEqual true + assertNoDuplicates(items) + } + + "Continuation token replay" should "not re-yield already consumed items after transient error" in { + + // Inject a transient error at a deterministic point (after page 3) so we can assert + // that pages 1-3's items are not duplicated after the retry. + val pageCount = 10 + val errorInjectedAfterPage = 3 + val errorInjected = new AtomicLong(0) + val factoryCallCount = new AtomicLong(0) + + val iterator = new TransientIOErrorsRetryingReadManyByPartitionKeyIterator[SparkRowItem]( + continuationToken => { + factoryCallCount.incrementAndGet() + generateMockedCosmosPagedFluxWithDeterministicError( + continuationToken, pageCount, errorInjectedAfterPage, errorInjected) + }, + pageSize, + 1, + None, + classOf[SparkRowItem] + ) + iterator.maxRetryIntervalInMs = 5 + + val items = drainAll(iterator) + items.size shouldEqual (pageCount * pageSize) + + errorInjected.get shouldEqual 1 + factoryCallCount.get shouldEqual 2 // initial + 1 retry + assertNoDuplicates(items) + } + + "Continuation token" should "be passed to factory on retry" in { + + val pageCount = 30 + val capturedContinuationTokens = new java.util.concurrent.CopyOnWriteArrayList[String]() + val transientErrorCount = new AtomicLong(0) + + val iterator = new TransientIOErrorsRetryingReadManyByPartitionKeyIterator[SparkRowItem]( + continuationToken => { + capturedContinuationTokens.add(continuationToken) + generateMockedCosmosPagedFlux( + continuationToken, pageCount, 0.15, transientErrorCount, injectEmptyPages = false) + }, + pageSize, + 1, + None, + classOf[SparkRowItem] + ) + iterator.maxRetryIntervalInMs = 5 + + drainAll(iterator) + + transientErrorCount.get > 0 shouldEqual true + + // First call should have null continuation token (start from beginning) + capturedContinuationTokens.get(0) shouldEqual null + + // Subsequent retry calls should have non-null continuation tokens + // (resume from last committed page) + val retryTokens = capturedContinuationTokens.asScala.drop(1) + retryTokens should not be empty + retryTokens.foreach(_ should not be null) + } + + "Non-transient errors" should "not be retried and propagate immediately" in { + + val iterator = new TransientIOErrorsRetryingReadManyByPartitionKeyIterator[SparkRowItem]( + _ => { + val flux = Flux.error[FeedResponse[SparkRowItem]](new DummyNonTransientCosmosException) + UtilBridgeInternal.createCosmosPagedFlux(_ => flux) + }, + pageSize, + 1, + None, + classOf[SparkRowItem] + ) + iterator.maxRetryIntervalInMs = 5 + + val thrown = the[CosmosException] thrownBy { + iterator.hasNext + } + thrown.getStatusCode shouldEqual 404 + } + + "Iterator close" should "clear internal state without extra factory calls" in { + + val factoryCallCount = new AtomicLong(0) + val iterator = new TransientIOErrorsRetryingReadManyByPartitionKeyIterator[SparkRowItem]( + continuationToken => { + factoryCallCount.incrementAndGet() + generateMockedCosmosPagedFlux( + continuationToken, 30, 0.0, new AtomicLong(0), injectEmptyPages = false) + }, + pageSize, + 1, + None, + classOf[SparkRowItem] + ) + + iterator.hasNext shouldEqual true + iterator.next() + + factoryCallCount.get shouldEqual 1 + + iterator.currentFeedResponseIterator.isDefined shouldEqual true + + iterator.close() + + iterator.currentFeedResponseIterator shouldEqual None + iterator.currentItemIterator shouldEqual None + + factoryCallCount.get shouldEqual 1 + } + + "Iterator close" should "be safe to call multiple times" in { + + val iterator = new TransientIOErrorsRetryingReadManyByPartitionKeyIterator[SparkRowItem]( + continuationToken => generateMockedCosmosPagedFlux( + continuationToken, 30, 0.0, new AtomicLong(0), injectEmptyPages = false), + pageSize, + 1, + None, + classOf[SparkRowItem] + ) + + iterator.hasNext shouldEqual true + iterator.next() + + iterator.close() + iterator.close() + iterator.close() + // No exception should be thrown + } + + // --- helpers --- + + private def drainAll(iterator: TransientIOErrorsRetryingReadManyByPartitionKeyIterator[SparkRowItem]): List[SparkRowItem] = { + val buffer = scala.collection.mutable.ListBuffer[SparkRowItem]() + while (iterator.hasNext) { + buffer += iterator.next() + } + buffer.toList + } + + private def assertNoDuplicates(items: List[SparkRowItem]): Unit = { + val ids = items.map(item => item.row.getString(0)) // "id" field + ids.size shouldEqual ids.distinct.size + } + + @throws[JsonProcessingException] + private def getDocumentDefinition(documentId: String) = { + val json = s"""{"id":"$documentId"}""" + val node = objectMapper.readValue(json, classOf[ObjectNode]) + SparkRowItem( + cosmosRowConverter.fromObjectNodeToRow( + ItemsTable.defaultSchemaForInferenceDisabled, + node, + SchemaConversionModes.Strict + ), + None) + } + + private def generateMockedCosmosPagedFlux( + continuationToken: String, + pageCount: Int, + errorRate: Double, + transientErrorCounter: AtomicLong, + injectEmptyPages: Boolean + ) = { + + require(pageCount > 20) + + val flux = generateFeedResponseFlux( + "Batch", pageCount, errorRate, + Option(continuationToken), transientErrorCounter, injectEmptyPages) + + UtilBridgeInternal.createCosmosPagedFlux(_ => flux) + } + + private def generateMockedCosmosPagedFluxWithDeterministicError( + continuationToken: String, + pageCount: Int, + errorAfterPage: Int, + errorInjected: AtomicLong + ) = { + + val flux = generateFeedResponseFluxWithDeterministicError( + "Batch", pageCount, errorAfterPage, + Option(continuationToken), errorInjected) + + UtilBridgeInternal.createCosmosPagedFlux(_ => flux) + } + + private def generateFeedResponseFlux( + prefix: String, + pageCount: Int, + errorRate: Double, + requestContinuationToken: Option[String], + transientErrorCounter: AtomicLong, + injectEmptyPages: Boolean + ): Flux[FeedResponse[SparkRowItem]] = { + + val responses = Array.range(1, pageCount + 1) + .map(i => generateFeedResponse( + prefix, i, + if (injectEmptyPages && i > 1 && i <= 20 && i % 2 == 0) -1 else 1)) + .filter(response => requestContinuationToken.isEmpty || + requestContinuationToken.get == null || + requestContinuationToken.get < response.getContinuationToken) + + Flux + .fromArray(responses) + .map(response => if (errorRate > 0 && rnd.nextDouble() < errorRate) { + transientErrorCounter.incrementAndGet() + throw new DummyTransientCosmosException + } else { + response + }) + } + + private def generateFeedResponseFluxWithDeterministicError( + prefix: String, + pageCount: Int, + errorAfterPage: Int, + requestContinuationToken: Option[String], + errorInjected: AtomicLong + ): Flux[FeedResponse[SparkRowItem]] = { + + val responses = Array.range(1, pageCount + 1) + .map(i => generateFeedResponse(prefix, i, 1)) + .filter(response => requestContinuationToken.isEmpty || + requestContinuationToken.get == null || + requestContinuationToken.get < response.getContinuationToken) + + Flux + .fromArray(responses) + .map(response => { + // Extract page sequence number from continuation token to inject error deterministically + val token = response.getContinuationToken + val pageNum = token.split("_Page")(1).split("_")(0).toInt + if (pageNum == errorAfterPage + 1 && errorInjected.compareAndSet(0, 1)) { + throw new DummyTransientCosmosException + } else { + response + } + }) + } + + private def generateFeedResponse( + prefix: String, + pageSequenceNumber: Int, + documentStartIndex: Int + ): FeedResponse[SparkRowItem] = { + + val continuationToken = f"$prefix%s_Page$pageSequenceNumber%05d_ContinuationToken" + val items = if (documentStartIndex < 0) { + Array.empty[SparkRowItem] + } else { + val id1 = f"$prefix%s_Page$pageSequenceNumber%05d_$documentStartIndex%05d" + val id2 = f"$prefix%s_Page$pageSequenceNumber%05d_${documentStartIndex + 1}%05d" + Array[SparkRowItem](getDocumentDefinition(id1), getDocumentDefinition(id2)) + } + + val r = ModelBridgeInternal.createFeedResponse( + items.toList.asJava, + new ConcurrentHashMap[String, String] + ) + ModelBridgeInternal.setFeedResponseContinuationToken(continuationToken, r) + r + } + + private class DummyTransientCosmosException + extends CosmosException(500, "Dummy Internal Server Error") + + private class DummyNonTransientCosmosException + extends CosmosException(404, "Dummy Not Found Error") +} +//scalastyle:on magic.number +//scalastyle:on multiple.string.literals diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/ReadManyByPartitionKeyTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/ReadManyByPartitionKeyTest.java index 79f7800bf426..e3ddb2ad8f6c 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/ReadManyByPartitionKeyTest.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/ReadManyByPartitionKeyTest.java @@ -16,6 +16,13 @@ import com.azure.cosmos.models.SqlParameter; import com.azure.cosmos.models.SqlQuerySpec; import com.azure.cosmos.rx.TestSuiteBase; +import com.azure.cosmos.test.faultinjection.CosmosFaultInjectionHelper; +import com.azure.cosmos.test.faultinjection.FaultInjectionConditionBuilder; +import com.azure.cosmos.test.faultinjection.FaultInjectionOperationType; +import com.azure.cosmos.test.faultinjection.FaultInjectionResultBuilders; +import com.azure.cosmos.test.faultinjection.FaultInjectionRule; +import com.azure.cosmos.test.faultinjection.FaultInjectionRuleBuilder; +import com.azure.cosmos.test.faultinjection.FaultInjectionServerErrorType; import com.azure.cosmos.util.CosmosPagedIterable; import com.fasterxml.jackson.databind.node.ObjectNode; import org.testng.annotations.AfterClass; @@ -728,6 +735,182 @@ public void singlePk_continuationToken_resumesCorrectly_whenInputContainsDuplica //endregion + //region Fault injection continuation-token replay tests + + @Test(groups = {"emulator"}, timeOut = TIMEOUT * 3) + public void singlePk_faultInjection_transientQueryFailure_resumesWithoutDuplicatesOrLoss() { + // Use small batch size to get multiple batches (continuation tokens between batches) + String originalValue = System.getProperty("COSMOS.READ_MANY_BY_PK_MAX_BATCH_SIZE"); + try { + System.setProperty("COSMOS.READ_MANY_BY_PK_MAX_BATCH_SIZE", "2"); + + // Create enough items across 4 PKs so that each PK produces multiple pages + createSinglePkItems("fiPk1", 5); + createSinglePkItems("fiPk2", 5); + createSinglePkItems("fiPk3", 5); + createSinglePkItems("fiPk4", 5); + + List pkValues = Arrays.asList( + new PartitionKey("fiPk1"), + new PartitionKey("fiPk2"), + new PartitionKey("fiPk3"), + new PartitionKey("fiPk4")); + + // First: collect baseline (no faults) to know exactly which items exist + List baselineResults = singlePkContainer + .readManyByPartitionKeys(pkValues, ObjectNode.class) + .stream().collect(Collectors.toList()); + assertThat(baselineResults).hasSize(20); + List baselineIds = baselineResults.stream() + .map(n -> n.get("id").asText()) + .sorted() + .collect(Collectors.toList()); + + // Now inject a transient 503 (Service Unavailable) on query operations with a + // limited hit count. This simulates network blips mid-iteration — the SDK's + // continuation-token-based retry must resume from the last committed page. + CosmosAsyncContainer asyncContainer = client.asyncClient() + .getDatabase(preExistingDatabaseId) + .getContainer(singlePkContainer.getId()); + + FaultInjectionRule queryTransientFaultRule = new FaultInjectionRuleBuilder( + "readManyByPk-transient-503") + .condition(new FaultInjectionConditionBuilder() + .operationType(FaultInjectionOperationType.QUERY_ITEM) + .build()) + .result(FaultInjectionResultBuilders + .getResultBuilder(FaultInjectionServerErrorType.SERVICE_UNAVAILABLE) + .build()) + .hitLimit(3) // inject exactly 3 transient failures then stop + .build(); + + CosmosFaultInjectionHelper + .configureFaultInjectionRules(asyncContainer, Arrays.asList(queryTransientFaultRule)) + .block(); + + try { + // Read with fault injection active — SDK should retry and deliver all items + List faultInjectedResults = singlePkContainer + .readManyByPartitionKeys(pkValues, ObjectNode.class) + .stream().collect(Collectors.toList()); + + List faultInjectedIds = faultInjectedResults.stream() + .map(n -> n.get("id").asText()) + .sorted() + .collect(Collectors.toList()); + + // Assert: no duplicates + assertThat(faultInjectedIds).doesNotHaveDuplicates(); + + // Assert: all baseline items are present (no data loss) + assertThat(faultInjectedIds).hasSameElementsAs(baselineIds); + + // Assert: exact count + assertThat(faultInjectedIds).hasSize(20); + } finally { + queryTransientFaultRule.disable(); + } + + cleanupContainer(singlePkContainer); + } finally { + if (originalValue != null) { + System.setProperty("COSMOS.READ_MANY_BY_PK_MAX_BATCH_SIZE", originalValue); + } else { + System.clearProperty("COSMOS.READ_MANY_BY_PK_MAX_BATCH_SIZE"); + } + } + } + + @Test(groups = {"emulator"}, timeOut = TIMEOUT * 3) + public void singlePk_faultInjection_transientFailure_continuationResumeNoLoss() { + // Tests the manual continuation token resume path: drain first page, inject fault, + // resume from continuation token — verifying the union of both calls covers all items. + String originalValue = System.getProperty("COSMOS.READ_MANY_BY_PK_MAX_BATCH_SIZE"); + try { + System.setProperty("COSMOS.READ_MANY_BY_PK_MAX_BATCH_SIZE", "1"); + + createSinglePkItems("ctResPk1", 3); + createSinglePkItems("ctResPk2", 3); + createSinglePkItems("ctResPk3", 3); + + List pkValues = Arrays.asList( + new PartitionKey("ctResPk1"), + new PartitionKey("ctResPk2"), + new PartitionKey("ctResPk3")); + + CosmosAsyncContainer asyncContainer = client.asyncClient() + .getDatabase(preExistingDatabaseId) + .getContainer(singlePkContainer.getId()); + + // Drain just the first page to get a continuation token + FeedResponse firstPage = asyncContainer + .readManyByPartitionKeys(pkValues, ObjectNode.class) + .byPage() + .blockFirst(); + + assertThat(firstPage).isNotNull(); + String continuationAfterFirstPage = firstPage.getContinuationToken(); + assertThat(continuationAfterFirstPage).isNotNull(); + List firstPageIds = firstPage.getResults().stream() + .map(n -> n.get("id").asText()) + .collect(Collectors.toList()); + + // Inject a transient fault, then resume from the continuation token + FaultInjectionRule resumeFaultRule = new FaultInjectionRuleBuilder( + "readManyByPk-resume-503") + .condition(new FaultInjectionConditionBuilder() + .operationType(FaultInjectionOperationType.QUERY_ITEM) + .build()) + .result(FaultInjectionResultBuilders + .getResultBuilder(FaultInjectionServerErrorType.SERVICE_UNAVAILABLE) + .build()) + .hitLimit(2) + .build(); + + CosmosFaultInjectionHelper + .configureFaultInjectionRules(asyncContainer, Arrays.asList(resumeFaultRule)) + .block(); + + try { + com.azure.cosmos.models.CosmosReadManyByPartitionKeysRequestOptions resumeOptions = + new com.azure.cosmos.models.CosmosReadManyByPartitionKeysRequestOptions(); + resumeOptions.setContinuationToken(continuationAfterFirstPage); + + List resumedResults = asyncContainer + .readManyByPartitionKeys(pkValues, resumeOptions, ObjectNode.class) + .byPage() + .flatMapIterable(FeedResponse::getResults) + .collectList() + .block(); + + assertThat(resumedResults).isNotNull(); + + List resumedIds = resumedResults.stream() + .map(n -> n.get("id").asText()) + .collect(Collectors.toList()); + + // Union of first page + resumed should cover all 9 items + List allIds = new ArrayList<>(firstPageIds); + allIds.addAll(resumedIds); + + assertThat(allIds).doesNotHaveDuplicates(); + assertThat(allIds).hasSize(9); + } finally { + resumeFaultRule.disable(); + } + + cleanupContainer(singlePkContainer); + } finally { + if (originalValue != null) { + System.setProperty("COSMOS.READ_MANY_BY_PK_MAX_BATCH_SIZE", originalValue); + } else { + System.clearProperty("COSMOS.READ_MANY_BY_PK_MAX_BATCH_SIZE"); + } + } + } + + //endregion + //region helper methods private List createSinglePkItems(String pkValue, int count) { From 587cac6d117d04afc0335d109045cf5828e3d531 Mon Sep 17 00:00:00 2001 From: Fabian Meiswinkel Date: Thu, 23 Apr 2026 13:50:38 +0000 Subject: [PATCH 60/64] Added and fixed tests --- ..._readManyByPartitionKeysAfterCreation.java | 307 +++++++++++++++++ ...tionWithAvailabilityStrategyTestsBase.java | 316 ++++++++++++++++++ .../cosmos/ReadManyByPartitionKeyTest.java | 185 +++------- ...InjectionServerErrorRuleOnDirectTests.java | 98 ++++++ 4 files changed, 770 insertions(+), 136 deletions(-) create mode 100644 sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/FITests_readManyByPartitionKeysAfterCreation.java diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/FITests_readManyByPartitionKeysAfterCreation.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/FITests_readManyByPartitionKeysAfterCreation.java new file mode 100644 index 000000000000..f618951522b2 --- /dev/null +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/FITests_readManyByPartitionKeysAfterCreation.java @@ -0,0 +1,307 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. +package com.azure.cosmos; + +import com.azure.cosmos.models.CosmosReadManyByPartitionKeysRequestOptions; +import com.azure.cosmos.models.FeedRange; +import com.azure.cosmos.models.FeedResponse; +import com.azure.cosmos.models.PartitionKey; +import com.azure.cosmos.test.faultinjection.CosmosFaultInjectionHelper; +import com.azure.cosmos.test.faultinjection.FaultInjectionConditionBuilder; +import com.azure.cosmos.test.faultinjection.FaultInjectionEndpointBuilder; +import com.azure.cosmos.test.faultinjection.FaultInjectionOperationType; +import com.azure.cosmos.test.faultinjection.FaultInjectionResultBuilders; +import com.azure.cosmos.test.faultinjection.FaultInjectionRule; +import com.azure.cosmos.test.faultinjection.FaultInjectionRuleBuilder; +import com.azure.cosmos.test.faultinjection.FaultInjectionServerErrorType; +import com.azure.cosmos.implementation.HttpConstants; +import com.azure.cosmos.implementation.TestConfigurations; +import com.fasterxml.jackson.databind.node.ObjectNode; +import org.apache.commons.lang3.ArrayUtils; +import org.testng.annotations.Test; + +import java.time.Duration; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.List; +import java.util.UUID; +import java.util.function.BiConsumer; +import java.util.function.Consumer; +import java.util.function.Function; +import java.util.stream.Collectors; + +import static org.assertj.core.api.Assertions.assertThat; + +public class FITests_readManyByPartitionKeysAfterCreation + extends FaultInjectionWithAvailabilityStrategyTestsBase { + + @Test(groups = {"fi-multi-master"}, dataProvider = "testConfigs_readManyByPartitionKeysAfterCreation", retryAnalyzer = FlakyTestRetryAnalyzer.class) + public void readManyByPartitionKeysAfterCreation( + String testCaseId, + Duration endToEndTimeout, + ThresholdBasedAvailabilityStrategy availabilityStrategy, + CosmosRegionSwitchHint regionSwitchHint, + ConnectionMode connectionMode, + Function readManyByPkOperation, + BiConsumer faultInjectionCallback, + BiConsumer validateStatusCode, + int expectedDiagnosticsContextCount, + Consumer[] firstDiagnosticsContextValidations, + Consumer[] otherDiagnosticsContextValidations, + Consumer responseValidator, + int numberOfOtherDocumentsWithSameId, + int numberOfOtherDocumentsWithSamePk, + boolean shouldInjectPreferredRegionsInClient) { + + execute( + testCaseId, + endToEndTimeout, + availabilityStrategy, + regionSwitchHint, + null, + notSpecifiedWhetherIdempotentWriteRetriesAreEnabled, + ArrayUtils.toArray(FaultInjectionOperationType.QUERY_ITEM), + readManyByPkOperation, + faultInjectionCallback, + validateStatusCode, + expectedDiagnosticsContextCount, + firstDiagnosticsContextValidations, + otherDiagnosticsContextValidations, + responseValidator, + numberOfOtherDocumentsWithSameId, + numberOfOtherDocumentsWithSamePk, + false, + connectionMode, + shouldInjectPreferredRegionsInClient); + } + + /** + * Validates continuation-token resume after fault injection causes a deterministic error. + * + * Strategy: use FeedRange-scoped fault injection so that queries against one physical + * partition fail while queries against other partitions succeed. With batch size 1, + * readManyByPartitionKeys processes one PK-batch at a time sequentially. The first + * batch(es) targeting the non-faulted partition succeed and emit pages with continuation + * tokens. When iteration reaches the faulted partition, the error surfaces to the caller. + * + * 1. Create documents across multiple PKs (spread across partitions) + * 2. Collect baseline (no faults) — all ids + * 3. Get feed ranges; pick the second one to fault + * 4. Inject sustained SERVICE_UNAVAILABLE scoped to that feed range + * 5. Iterate page-by-page; collect items + continuation tokens from successful pages + * 6. When error occurs: validate it's expected, capture last good continuation + * 7. Disable the fault injection rule + * 8. Resume from the last good continuation token + * 9. Assert: union of items before error + items from resume = all baseline items, no duplicates + */ + @Test(groups = {"fi-multi-master"}, timeOut = 180000, retryAnalyzer = FlakyTestRetryAnalyzer.class) + public void readManyByPartitionKeys_continuationResumeAfterFaultInjection() { + + String originalBatchSize = System.getProperty("COSMOS.READ_MANY_BY_PK_MAX_BATCH_SIZE"); + try { + // batch size 1 = one PK per batch = sequential processing across partitions + System.setProperty("COSMOS.READ_MANY_BY_PK_MAX_BATCH_SIZE", "1"); + + CosmosAsyncClient client = new CosmosClientBuilder() + .endpoint(TestConfigurations.HOST) + .key(TestConfigurations.MASTER_KEY) + .contentResponseOnWriteEnabled(true) + .directMode() + .buildAsyncClient(); + + try { + CosmosAsyncContainer container = client + .getDatabase(this.getTestDatabaseId()) + .getContainer(this.getTestContainerId()); + + String uniqueTag = UUID.randomUUID().toString().substring(0, 8); + + // Create items across 3 PKs, 3 items each = 9 items total + List pkValues = Arrays.asList( + "ctResumePk1_" + uniqueTag, + "ctResumePk2_" + uniqueTag, + "ctResumePk3_" + uniqueTag); + + List allCreatedItems = new ArrayList<>(); + for (String pk : pkValues) { + for (int i = 0; i < 3; i++) { + ObjectNode item = com.azure.cosmos.implementation.Utils + .getSimpleObjectMapper().createObjectNode(); + item.put("id", UUID.randomUUID().toString()); + item.put("mypk", pk); + container.createItem(item).block(); + allCreatedItems.add(item); + } + } + + List partitionKeys = pkValues.stream() + .map(PartitionKey::new) + .collect(Collectors.toList()); + + // Step 1: Baseline — drain all pages without faults to know the complete set of ids + List> baselinePages = container + .readManyByPartitionKeys(partitionKeys, ObjectNode.class) + .byPage() + .collectList() + .block(); + + assertThat(baselinePages).isNotNull(); + assertThat(baselinePages.size()).isGreaterThan(1); // with batch size 1 there must be multiple pages + + List baselineIds = baselinePages.stream() + .flatMap(p -> p.getResults().stream()) + .map(n -> n.get("id").asText()) + .sorted() + .collect(Collectors.toList()); + assertThat(baselineIds).hasSize(9); + assertThat(baselineIds).doesNotHaveDuplicates(); + + // Step 2: Get feed ranges and pick the LAST one to fault. + // With batch size 1, readManyByPartitionKeys processes batches sorted by EPK. + // Faulting the last feed range ensures the first batches succeed (giving us + // pages with continuation tokens) before the faulted partition is reached. + List feedRanges = container.getFeedRanges().block(); + assertThat(feedRanges).isNotNull(); + assertThat(feedRanges.size()).isGreaterThanOrEqualTo(1); + FeedRange faultedFeedRange = feedRanges.get(feedRanges.size() - 1); + + // Step 3: Inject sustained SERVICE_UNAVAILABLE scoped to the last feed range + FaultInjectionRule partitionScopedRule = new FaultInjectionRuleBuilder( + "readManyByPk-ct-resume-partition-scoped") + .condition(new FaultInjectionConditionBuilder() + .operationType(FaultInjectionOperationType.QUERY_ITEM) + .endpoints(new FaultInjectionEndpointBuilder(faultedFeedRange) + .replicaCount(4) + .includePrimary(true) + .build()) + .build()) + .result(FaultInjectionResultBuilders + .getResultBuilder(FaultInjectionServerErrorType.SERVICE_UNAVAILABLE) + .build()) + .duration(Duration.ofSeconds(120)) + .build(); + + CosmosFaultInjectionHelper + .configureFaultInjectionRules(container, Collections.singletonList(partitionScopedRule)) + .block(); + + // Step 4: Drain page-by-page. Pages from non-faulted partitions succeed; + // when the faulted partition is reached, the error surfaces. + List itemsBeforeError = new ArrayList<>(); + String lastGoodContinuation = null; + boolean errorOccurred = false; + + CosmosEndToEndOperationLatencyPolicyConfig e2ePolicy = + new CosmosEndToEndOperationLatencyPolicyConfigBuilder(Duration.ofSeconds(10)) + .enable(true) + .build(); + + CosmosReadManyByPartitionKeysRequestOptions faultOptions = + new CosmosReadManyByPartitionKeysRequestOptions(); + faultOptions.setCosmosEndToEndOperationLatencyPolicyConfig(e2ePolicy); + + try { + // Use Flux iteration (toIterable) so we can capture per-page state + for (FeedResponse page : container + .readManyByPartitionKeys(partitionKeys, faultOptions, ObjectNode.class) + .byPage() + .toIterable()) { + + for (ObjectNode item : page.getResults()) { + itemsBeforeError.add(item.get("id").asText()); + } + if (page.getContinuationToken() != null) { + lastGoodContinuation = page.getContinuationToken(); + } + } + } catch (Exception e) { + errorOccurred = true; + } + + // Step 5: The fault injection MUST have caused an error — pages from the + // faulted partition cannot succeed with SERVICE_UNAVAILABLE on all replicas. + assertThat(errorOccurred) + .as("Fault injection on the last feed range must cause an error") + .isTrue(); + + // We must have captured at least one continuation token from successful pages + assertThat(lastGoodContinuation) + .as("At least one page must have succeeded before the faulted partition") + .isNotNull(); + + // Items collected so far must be a strict subset of the baseline + assertThat(itemsBeforeError).doesNotHaveDuplicates(); + assertThat(itemsBeforeError.size()).isGreaterThan(0); + assertThat(itemsBeforeError.size()).isLessThan(baselineIds.size()); + + // Step 6: Disable fault injection rule + partitionScopedRule.disable(); + + // Step 7: Resume from the last good continuation token + CosmosReadManyByPartitionKeysRequestOptions resumeOptions = + new CosmosReadManyByPartitionKeysRequestOptions(); + resumeOptions.setContinuationToken(lastGoodContinuation); + + List> resumedPages = container + .readManyByPartitionKeys(partitionKeys, resumeOptions, ObjectNode.class) + .byPage() + .collectList() + .block(); + + assertThat(resumedPages).isNotNull(); + + List resumedIds = resumedPages.stream() + .flatMap(p -> p.getResults().stream()) + .map(n -> n.get("id").asText()) + .collect(Collectors.toList()); + + // Step 8: Assert completeness — union of before + after = all baseline items + List combined = new ArrayList<>(itemsBeforeError); + combined.addAll(resumedIds); + + assertThat(combined).doesNotHaveDuplicates(); + assertThat(combined).hasSameElementsAs(baselineIds); + + // Cleanup + for (ObjectNode item : allCreatedItems) { + try { + container.deleteItem( + item.get("id").asText(), + new PartitionKey(item.get("mypk").asText())).block(); + } catch (Exception ignore) { } + } + + } finally { + safeClose(client); + } + } finally { + if (originalBatchSize != null) { + System.setProperty("COSMOS.READ_MANY_BY_PK_MAX_BATCH_SIZE", originalBatchSize); + } else { + System.clearProperty("COSMOS.READ_MANY_BY_PK_MAX_BATCH_SIZE"); + } + } + } + + // Helper to access testDatabaseId from base class + private String getTestDatabaseId() { + try { + java.lang.reflect.Field f = FaultInjectionWithAvailabilityStrategyTestsBase.class.getDeclaredField("testDatabaseId"); + f.setAccessible(true); + return (String) f.get(this); + } catch (Exception e) { + throw new RuntimeException(e); + } + } + + private String getTestContainerId() { + try { + java.lang.reflect.Field f = FaultInjectionWithAvailabilityStrategyTestsBase.class.getDeclaredField("testContainerId"); + f.setAccessible(true); + return (String) f.get(this); + } catch (Exception e) { + throw new RuntimeException(e); + } + } +} diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/FaultInjectionWithAvailabilityStrategyTestsBase.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/FaultInjectionWithAvailabilityStrategyTestsBase.java index dbb158438a9e..5b1243560ecf 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/FaultInjectionWithAvailabilityStrategyTestsBase.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/FaultInjectionWithAvailabilityStrategyTestsBase.java @@ -31,6 +31,7 @@ import com.azure.cosmos.models.FeedResponse; import com.azure.cosmos.models.PartitionKey; import com.azure.cosmos.models.PartitionKeyDefinition; +import com.azure.cosmos.models.CosmosReadManyByPartitionKeysRequestOptions; import com.azure.cosmos.models.ThroughputProperties; import com.azure.cosmos.rx.TestSuiteBase; import com.azure.cosmos.test.faultinjection.CosmosFaultInjectionHelper; @@ -3887,6 +3888,321 @@ public Object[][] testConfigs_readManyAfterCreation() { return addBooleanFlagsToAllTestConfigs(testConfigs_readManyAfterCreation); } + private CosmosResponseWrapper readManyByPartitionKeysCore( + ItemOperationInvocationParameters params, + int numberOfOtherDocumentsWithSamePk + ) { + + List pkValues = new ArrayList<>(); + pkValues.add(new PartitionKey(params.idAndPkValuePair.getRight())); + + CosmosReadManyByPartitionKeysRequestOptions options = new CosmosReadManyByPartitionKeysRequestOptions(); + + CosmosEndToEndOperationLatencyPolicyConfig e2ePolicy = ImplementationBridgeHelpers + .CosmosItemRequestOptionsHelper + .getCosmosItemRequestOptionsAccessor() + .getEndToEndOperationLatencyPolicyConfig(params.options); + options.setCosmosEndToEndOperationLatencyPolicyConfig(e2ePolicy); + + List> returnedPages; + try { + returnedPages = params.container + .readManyByPartitionKeys(pkValues, options, ObjectNode.class) + .byPage() + .collectList() + .block(); + } catch (CosmosException cosmosException) { + return new CosmosResponseWrapper( + cosmosException.getDiagnostics() != null + ? new CosmosDiagnosticsContext[] { cosmosException.getDiagnostics().getDiagnosticsContext() } + : null, + cosmosException.getStatusCode(), + cosmosException.getSubStatusCode(), + null); + } + + ArrayList foundCtxs = new ArrayList<>(); + + if (returnedPages == null || returnedPages.isEmpty()) { + return new CosmosResponseWrapper( + null, + HttpConstants.StatusCodes.OK, + HttpConstants.SubStatusCodes.UNKNOWN, + 0L); + } + + long totalRecordCount = 0L; + for (FeedResponse page : returnedPages) { + if (page.getCosmosDiagnostics() != null) { + foundCtxs.add(page.getCosmosDiagnostics().getDiagnosticsContext()); + } else { + foundCtxs.add(null); + } + + if (page.getResults() != null && page.getResults().size() > 0) { + totalRecordCount += page.getResults().size(); + } + } + + return new CosmosResponseWrapper( + foundCtxs.toArray(new CosmosDiagnosticsContext[0]), + HttpConstants.StatusCodes.OK, + HttpConstants.SubStatusCodes.UNKNOWN, + totalRecordCount); + } + + @DataProvider(name = "testConfigs_readManyByPartitionKeysAfterCreation") + public Object[][] testConfigs_readManyByPartitionKeysAfterCreation() { + + final int ENOUGH_DOCS_SAME_PK_TO_EXCEED_PAGE_SIZE = 10; + final int NO_OTHER_DOCS_WITH_SAME_PK = 0; + final int NO_OTHER_DOCS_WITH_SAME_ID = 0; + final int ENOUGH_DOCS_OTHER_PK_TO_HIT_EVERY_PARTITION = PHYSICAL_PARTITION_COUNT * 10; + final int SINGLE_REGION = 1; + final int TWO_REGIONS = 2; + + BiConsumer validateExpectedRecordCount = (response, expectedRecordCount) -> { + if (expectedRecordCount != null) { + assertThat(response).isNotNull(); + assertThat(response.getTotalRecordCount()).isNotNull(); + assertThat(response.getTotalRecordCount()).isEqualTo(expectedRecordCount); + } + }; + + Consumer validateExactlyOneRecordReturned = + (response) -> validateExpectedRecordCount.accept(response, 1L); + + Consumer validateAllRecordsSamePartitionReturned = + (response) -> validateExpectedRecordCount.accept( + response, + 1L + ENOUGH_DOCS_SAME_PK_TO_EXCEED_PAGE_SIZE); + + BiConsumer validateCtxRegions = + (ctx, expectedNumberOfRegionsContacted) -> { + assertThat(ctx).isNotNull(); + if (ctx != null) { + assertThat(ctx.getContactedRegionNames().size()).isEqualTo(expectedNumberOfRegionsContacted); + } + }; + + Consumer validateCtxSingleRegion = + (ctx) -> validateCtxRegions.accept(ctx, SINGLE_REGION); + + Consumer validateCtxTwoRegions = + (ctx) -> validateCtxRegions.accept(ctx, TWO_REGIONS); + + Consumer validateCtxOnlyFeedResponses = + (ctx) -> { + assertThat(ctx).isNotNull(); + if (ctx != null) { + assertThat(ctx.getDiagnostics()).isNotNull(); + assertThat(ctx.getDiagnostics().size()).isGreaterThanOrEqualTo(1); + } + }; + + Function + readManyByPkSinglePartition = (inputParams) -> + readManyByPartitionKeysCore(inputParams, ENOUGH_DOCS_SAME_PK_TO_EXCEED_PAGE_SIZE); + + Function + readManyByPkSingleDoc = (inputParams) -> + readManyByPartitionKeysCore(inputParams, NO_OTHER_DOCS_WITH_SAME_PK); + + Object[][] testConfigs = new Object[][] { + // CONFIG description + // new Object[] { + // TestId - name identifying the test case + // End-to-end timeout + // Availability Strategy used + // Region switch hint + // ConnectionMode + // readManyByPartitionKeys operation callback + // Failure injection callback + // Status code/sub status code validation callback + // Expected number of DiagnosticsContext instances + // Diagnostics context validation callback applied to the first DiagnosticsContext + // Diagnostics context validation callback applied to all other DiagnosticsContext + // Consumer - callback to validate the response + // numberOfOtherDocumentsWithSameId + // numberOfOtherDocumentsWithSamePk + // }, + + // readManyByPartitionKeys - single partition, no failures, no availability strategy + new Object[] { + "ReadManyByPk_SinglePartition_AllGood_NoAvailabilityStrategy", + ONE_SECOND_DURATION, + noAvailabilityStrategy, + noRegionSwitchHint, + ConnectionMode.DIRECT, + readManyByPkSinglePartition, + noFailureInjection, + validateStatusCodeIs200Ok, + 1, + ArrayUtils.toArray( + validateCtxSingleRegion, + validateCtxOnlyFeedResponses + ), + null, + validateAllRecordsSamePartitionReturned, + NO_OTHER_DOCS_WITH_SAME_ID, + ENOUGH_DOCS_SAME_PK_TO_EXCEED_PAGE_SIZE + }, + + // readManyByPartitionKeys - single doc, no failures, no availability strategy + new Object[] { + "ReadManyByPk_SingleDoc_AllGood_NoAvailabilityStrategy", + ONE_SECOND_DURATION, + noAvailabilityStrategy, + noRegionSwitchHint, + ConnectionMode.DIRECT, + readManyByPkSingleDoc, + noFailureInjection, + validateStatusCodeIs200Ok, + 1, + ArrayUtils.toArray( + validateCtxSingleRegion, + validateCtxOnlyFeedResponses + ), + null, + validateExactlyOneRecordReturned, + NO_OTHER_DOCS_WITH_SAME_ID, + NO_OTHER_DOCS_WITH_SAME_PK + }, + + // readManyByPartitionKeys - 408 timeout in first region, eager availability strategy + // Should succeed via hedging to second region + new Object[] { + "ReadManyByPk_SinglePartition_408_FirstRegionOnly_EagerAvailabilityStrategy", + THREE_SECOND_DURATION, + eagerThresholdAvailabilityStrategy, + noRegionSwitchHint, + ConnectionMode.DIRECT, + readManyByPkSinglePartition, + injectTransitTimeoutIntoFirstRegionOnly, + validateStatusCodeIs200Ok, + 1, + ArrayUtils.toArray( + validateCtxTwoRegions, + validateCtxOnlyFeedResponses + ), + null, + validateAllRecordsSamePartitionReturned, + NO_OTHER_DOCS_WITH_SAME_ID, + ENOUGH_DOCS_SAME_PK_TO_EXCEED_PAGE_SIZE + }, + + // readManyByPartitionKeys - 404/1002 in first region, remote preferred, reluctant strategy + // Client retry policy should failover to second region before hedging kicks in + new Object[] { + "ReadManyByPk_SinglePartition_404-1002_RemotePreferred_FirstRegionOnly_ReluctantAvailabilityStrategy", + Duration.ofSeconds(10), + reluctantThresholdAvailabilityStrategy, + noRegionSwitchHint, + ConnectionMode.DIRECT, + readManyByPkSinglePartition, + injectReadSessionNotAvailableIntoFirstRegionOnly, + validateStatusCodeIs200Ok, + 1, + ArrayUtils.toArray( + validateCtxTwoRegions, + validateCtxOnlyFeedResponses + ), + null, + validateAllRecordsSamePartitionReturned, + NO_OTHER_DOCS_WITH_SAME_ID, + ENOUGH_DOCS_SAME_PK_TO_EXCEED_PAGE_SIZE + }, + + // readManyByPartitionKeys - 429/3200 in first region, default availability strategy + // Should succeed via hedging to second region + new Object[] { + "ReadManyByPk_SinglePartition_429-3200_FirstRegionOnly_DefaultAvailabilityStrategy", + THREE_SECOND_DURATION, + defaultAvailabilityStrategy, + noRegionSwitchHint, + ConnectionMode.DIRECT, + readManyByPkSinglePartition, + injectRequestRateTooLargeIntoFirstRegionOnly, + validateStatusCodeIs200Ok, + 1, + ArrayUtils.toArray( + validateCtxTwoRegions, + validateCtxOnlyFeedResponses + ), + null, + validateAllRecordsSamePartitionReturned, + NO_OTHER_DOCS_WITH_SAME_ID, + ENOUGH_DOCS_SAME_PK_TO_EXCEED_PAGE_SIZE + }, + + // readManyByPartitionKeys - 429/3200 in first region, no availability strategy + // Should time out since no hedging and 429 retries locally until timeout + new Object[] { + "ReadManyByPk_SinglePartition_429-3200_FirstRegionOnly_NoAvailabilityStrategy", + THREE_SECOND_DURATION, + noAvailabilityStrategy, + noRegionSwitchHint, + ConnectionMode.DIRECT, + readManyByPkSinglePartition, + injectRequestRateTooLargeIntoFirstRegionOnly, + validateStatusCodeIsOperationCancelled, + 1, + ArrayUtils.toArray( + validateCtxSingleRegion, + validateCtxOnlyFeedResponses + ), + null, + null, + NO_OTHER_DOCS_WITH_SAME_ID, + ENOUGH_DOCS_SAME_PK_TO_EXCEED_PAGE_SIZE + }, + + // readManyByPartitionKeys - 503 in first region, no availability strategy + // ClientRetryPolicy will failover to second region + new Object[] { + "ReadManyByPk_SinglePartition_503_FirstRegionOnly_NoAvailabilityStrategy", + Duration.ofSeconds(90), + noAvailabilityStrategy, + noRegionSwitchHint, + ConnectionMode.DIRECT, + readManyByPkSinglePartition, + injectServiceUnavailableIntoFirstRegionOnly, + validateStatusCodeIs200Ok, + 1, + ArrayUtils.toArray( + validateCtxTwoRegions, + validateCtxOnlyFeedResponses + ), + null, + validateAllRecordsSamePartitionReturned, + NO_OTHER_DOCS_WITH_SAME_ID, + ENOUGH_DOCS_SAME_PK_TO_EXCEED_PAGE_SIZE + }, + + // readManyByPartitionKeys - 503 in all regions, eager availability strategy + // Both regions fail - should get 503 + new Object[] { + "ReadManyByPk_SinglePartition_503_AllRegions_EagerAvailabilityStrategy", + Duration.ofSeconds(10), + eagerThresholdAvailabilityStrategy, + noRegionSwitchHint, + ConnectionMode.DIRECT, + readManyByPkSinglePartition, + injectServiceUnavailableIntoAllRegions, + validateStatusCodeIsServiceUnavailable, + 1, + ArrayUtils.toArray( + validateCtxTwoRegions + ), + null, + null, + NO_OTHER_DOCS_WITH_SAME_ID, + ENOUGH_DOCS_SAME_PK_TO_EXCEED_PAGE_SIZE + }, + }; + + return addBooleanFlagsToAllTestConfigs(testConfigs); + } private CosmosResponseWrapper readAllReturnsTotalRecordCountCore( diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/ReadManyByPartitionKeyTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/ReadManyByPartitionKeyTest.java index e3ddb2ad8f6c..4b287457ee4b 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/ReadManyByPartitionKeyTest.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/ReadManyByPartitionKeyTest.java @@ -16,13 +16,6 @@ import com.azure.cosmos.models.SqlParameter; import com.azure.cosmos.models.SqlQuerySpec; import com.azure.cosmos.rx.TestSuiteBase; -import com.azure.cosmos.test.faultinjection.CosmosFaultInjectionHelper; -import com.azure.cosmos.test.faultinjection.FaultInjectionConditionBuilder; -import com.azure.cosmos.test.faultinjection.FaultInjectionOperationType; -import com.azure.cosmos.test.faultinjection.FaultInjectionResultBuilders; -import com.azure.cosmos.test.faultinjection.FaultInjectionRule; -import com.azure.cosmos.test.faultinjection.FaultInjectionRuleBuilder; -import com.azure.cosmos.test.faultinjection.FaultInjectionServerErrorType; import com.azure.cosmos.util.CosmosPagedIterable; import com.fasterxml.jackson.databind.node.ObjectNode; import org.testng.annotations.AfterClass; @@ -735,146 +728,66 @@ public void singlePk_continuationToken_resumesCorrectly_whenInputContainsDuplica //endregion - //region Fault injection continuation-token replay tests + //region Continuation-token resume correctness tests @Test(groups = {"emulator"}, timeOut = TIMEOUT * 3) - public void singlePk_faultInjection_transientQueryFailure_resumesWithoutDuplicatesOrLoss() { - // Use small batch size to get multiple batches (continuation tokens between batches) + public void singlePk_continuationToken_resumeAtEveryPageBoundary_noLossNoDuplicates() { + // Validates that resuming from ANY page boundary produces a complete, duplicate-free + // result set when combined with the items from earlier pages. String originalValue = System.getProperty("COSMOS.READ_MANY_BY_PK_MAX_BATCH_SIZE"); try { System.setProperty("COSMOS.READ_MANY_BY_PK_MAX_BATCH_SIZE", "2"); - // Create enough items across 4 PKs so that each PK produces multiple pages - createSinglePkItems("fiPk1", 5); - createSinglePkItems("fiPk2", 5); - createSinglePkItems("fiPk3", 5); - createSinglePkItems("fiPk4", 5); + createSinglePkItems("resumePk1", 4); + createSinglePkItems("resumePk2", 4); + createSinglePkItems("resumePk3", 4); + createSinglePkItems("resumePk4", 4); List pkValues = Arrays.asList( - new PartitionKey("fiPk1"), - new PartitionKey("fiPk2"), - new PartitionKey("fiPk3"), - new PartitionKey("fiPk4")); - - // First: collect baseline (no faults) to know exactly which items exist - List baselineResults = singlePkContainer - .readManyByPartitionKeys(pkValues, ObjectNode.class) - .stream().collect(Collectors.toList()); - assertThat(baselineResults).hasSize(20); - List baselineIds = baselineResults.stream() - .map(n -> n.get("id").asText()) - .sorted() - .collect(Collectors.toList()); - - // Now inject a transient 503 (Service Unavailable) on query operations with a - // limited hit count. This simulates network blips mid-iteration — the SDK's - // continuation-token-based retry must resume from the last committed page. - CosmosAsyncContainer asyncContainer = client.asyncClient() - .getDatabase(preExistingDatabaseId) - .getContainer(singlePkContainer.getId()); - - FaultInjectionRule queryTransientFaultRule = new FaultInjectionRuleBuilder( - "readManyByPk-transient-503") - .condition(new FaultInjectionConditionBuilder() - .operationType(FaultInjectionOperationType.QUERY_ITEM) - .build()) - .result(FaultInjectionResultBuilders - .getResultBuilder(FaultInjectionServerErrorType.SERVICE_UNAVAILABLE) - .build()) - .hitLimit(3) // inject exactly 3 transient failures then stop - .build(); - - CosmosFaultInjectionHelper - .configureFaultInjectionRules(asyncContainer, Arrays.asList(queryTransientFaultRule)) - .block(); - - try { - // Read with fault injection active — SDK should retry and deliver all items - List faultInjectedResults = singlePkContainer - .readManyByPartitionKeys(pkValues, ObjectNode.class) - .stream().collect(Collectors.toList()); - - List faultInjectedIds = faultInjectedResults.stream() - .map(n -> n.get("id").asText()) - .sorted() - .collect(Collectors.toList()); - - // Assert: no duplicates - assertThat(faultInjectedIds).doesNotHaveDuplicates(); - - // Assert: all baseline items are present (no data loss) - assertThat(faultInjectedIds).hasSameElementsAs(baselineIds); - - // Assert: exact count - assertThat(faultInjectedIds).hasSize(20); - } finally { - queryTransientFaultRule.disable(); - } - - cleanupContainer(singlePkContainer); - } finally { - if (originalValue != null) { - System.setProperty("COSMOS.READ_MANY_BY_PK_MAX_BATCH_SIZE", originalValue); - } else { - System.clearProperty("COSMOS.READ_MANY_BY_PK_MAX_BATCH_SIZE"); - } - } - } - - @Test(groups = {"emulator"}, timeOut = TIMEOUT * 3) - public void singlePk_faultInjection_transientFailure_continuationResumeNoLoss() { - // Tests the manual continuation token resume path: drain first page, inject fault, - // resume from continuation token — verifying the union of both calls covers all items. - String originalValue = System.getProperty("COSMOS.READ_MANY_BY_PK_MAX_BATCH_SIZE"); - try { - System.setProperty("COSMOS.READ_MANY_BY_PK_MAX_BATCH_SIZE", "1"); - - createSinglePkItems("ctResPk1", 3); - createSinglePkItems("ctResPk2", 3); - createSinglePkItems("ctResPk3", 3); - - List pkValues = Arrays.asList( - new PartitionKey("ctResPk1"), - new PartitionKey("ctResPk2"), - new PartitionKey("ctResPk3")); + new PartitionKey("resumePk1"), + new PartitionKey("resumePk2"), + new PartitionKey("resumePk3"), + new PartitionKey("resumePk4")); CosmosAsyncContainer asyncContainer = client.asyncClient() .getDatabase(preExistingDatabaseId) .getContainer(singlePkContainer.getId()); - // Drain just the first page to get a continuation token - FeedResponse firstPage = asyncContainer + // Collect all pages in a single pass + List> allPages = asyncContainer .readManyByPartitionKeys(pkValues, ObjectNode.class) .byPage() - .blockFirst(); + .collectList() + .block(); - assertThat(firstPage).isNotNull(); - String continuationAfterFirstPage = firstPage.getContinuationToken(); - assertThat(continuationAfterFirstPage).isNotNull(); - List firstPageIds = firstPage.getResults().stream() + assertThat(allPages).isNotNull(); + assertThat(allPages.size()).isGreaterThan(1); + + List allIds = allPages.stream() + .flatMap(p -> p.getResults().stream()) .map(n -> n.get("id").asText()) .collect(Collectors.toList()); - - // Inject a transient fault, then resume from the continuation token - FaultInjectionRule resumeFaultRule = new FaultInjectionRuleBuilder( - "readManyByPk-resume-503") - .condition(new FaultInjectionConditionBuilder() - .operationType(FaultInjectionOperationType.QUERY_ITEM) - .build()) - .result(FaultInjectionResultBuilders - .getResultBuilder(FaultInjectionServerErrorType.SERVICE_UNAVAILABLE) - .build()) - .hitLimit(2) - .build(); - - CosmosFaultInjectionHelper - .configureFaultInjectionRules(asyncContainer, Arrays.asList(resumeFaultRule)) - .block(); - - try { + assertThat(allIds).hasSize(16); + assertThat(allIds).doesNotHaveDuplicates(); + + // For each page boundary, resume from that boundary's continuation token + // and verify: items before + items after = all items, no duplicates + for (int splitAt = 0; splitAt < allPages.size() - 1; splitAt++) { + String continuation = allPages.get(splitAt).getContinuationToken(); + if (continuation == null) { + continue; // last page has null continuation + } + + // Items from pages 0..splitAt + List beforeIds = new ArrayList<>(); + for (int p = 0; p <= splitAt; p++) { + allPages.get(p).getResults().forEach(n -> beforeIds.add(n.get("id").asText())); + } + + // Resume from continuation com.azure.cosmos.models.CosmosReadManyByPartitionKeysRequestOptions resumeOptions = new com.azure.cosmos.models.CosmosReadManyByPartitionKeysRequestOptions(); - resumeOptions.setContinuationToken(continuationAfterFirstPage); + resumeOptions.setContinuationToken(continuation); List resumedResults = asyncContainer .readManyByPartitionKeys(pkValues, resumeOptions, ObjectNode.class) @@ -884,19 +797,19 @@ public void singlePk_faultInjection_transientFailure_continuationResumeNoLoss() .block(); assertThat(resumedResults).isNotNull(); - - List resumedIds = resumedResults.stream() + List afterIds = resumedResults.stream() .map(n -> n.get("id").asText()) .collect(Collectors.toList()); - // Union of first page + resumed should cover all 9 items - List allIds = new ArrayList<>(firstPageIds); - allIds.addAll(resumedIds); + List combined = new ArrayList<>(beforeIds); + combined.addAll(afterIds); - assertThat(allIds).doesNotHaveDuplicates(); - assertThat(allIds).hasSize(9); - } finally { - resumeFaultRule.disable(); + assertThat(combined) + .as("Resume at page boundary %d should produce all items", splitAt) + .doesNotHaveDuplicates(); + assertThat(combined) + .as("Resume at page boundary %d should cover all 16 items", splitAt) + .hasSameElementsAs(allIds); } cleanupContainer(singlePkContainer); diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/faultinjection/FaultInjectionServerErrorRuleOnDirectTests.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/faultinjection/FaultInjectionServerErrorRuleOnDirectTests.java index 2ed7e7e091f4..2cf84c721107 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/faultinjection/FaultInjectionServerErrorRuleOnDirectTests.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/faultinjection/FaultInjectionServerErrorRuleOnDirectTests.java @@ -24,6 +24,7 @@ import com.azure.cosmos.models.CosmosItemResponse; import com.azure.cosmos.models.CosmosQueryRequestOptions; import com.azure.cosmos.models.FeedRange; +import com.azure.cosmos.models.FeedResponse; import com.azure.cosmos.models.PartitionKey; import com.azure.cosmos.test.faultinjection.CosmosFaultInjectionHelper; import com.azure.cosmos.test.faultinjection.FaultInjectionConditionBuilder; @@ -1078,6 +1079,103 @@ public void faultInjectionServerErrorRuleTests_HitLimit() throws JsonProcessingE } } + @Test(groups = {"multi-region"}, timeOut = TIMEOUT * 3) + public void readManyByPartitionKeys_goneError_retriedInternally_noLossNoDuplicates() { + // Injects GONE (410/0) errors on query operations targeting readManyByPartitionKeys. + // 410/0 triggers a partition-key-range refresh + retry inside the SDK's + // DocumentProducer retry loop (direct mode): + // + // DocumentProducer.executeFeedOperationCore [DocumentProducer.java:144] + // → ObservableHelper.inlineIfPossibleAsObs(func, finalRetryPolicy) [DocumentProducer.java:147] + // → BackoffRetryUtility.executeRetry() [ObservableHelper.java:43] + // → client.executeQueryAsync(request) [IDocumentQueryClient impl in RxDocumentClientImpl.java:5371] + // → RxDocumentClientImpl.query(request) [RxDocumentClientImpl.java:1822] + // → getStoreProxy(request).processMessage(request) + // → [Direct: RntbdTransportClient] ← GONE 410/0 injected by fault injection + // + // The finalRetryPolicy is a ClientRetryPolicy created via: + // resetSessionTokenRetryPolicy.getRequestPolicy() [ParallelDocumentQueryExecutionContextBase.java:115] + // ClientRetryPolicy handles 410 (GONE) by refreshing the partition key range cache + // and retrying in the same region — this works regardless of region count. + + CosmosAsyncContainer container = getSharedMultiPartitionCosmosContainer(clientWithoutPreferredRegions); + + // Create test items across 3 partition keys + String pk1 = "readManyGone_" + UUID.randomUUID(); + String pk2 = "readManyGone_" + UUID.randomUUID(); + String pk3 = "readManyGone_" + UUID.randomUUID(); + + List createdItems = new ArrayList<>(); + for (String pk : Arrays.asList(pk1, pk2, pk3)) { + for (int i = 0; i < 3; i++) { + TestObject item = TestObject.create(pk); + container.createItem(item).block(); + createdItems.add(item); + } + } + + List pkValues = Arrays.asList( + new PartitionKey(pk1), + new PartitionKey(pk2), + new PartitionKey(pk3)); + + // Baseline — no faults + List baselineResults = container + .readManyByPartitionKeys(pkValues, TestObject.class) + .byPage() + .flatMapIterable(FeedResponse::getResults) + .collectList() + .block(); + + assertThat(baselineResults).isNotNull(); + assertThat(baselineResults.size()).isEqualTo(9); + List baselineIds = baselineResults.stream() + .map(TestObject::getId) + .sorted() + .collect(Collectors.toList()); + + // Inject GONE (410/0) on QUERY_ITEM — direct mode, no connectionType needed + FaultInjectionRule goneRule = new FaultInjectionRuleBuilder("readManyByPk-gone-direct") + .condition(new FaultInjectionConditionBuilder() + .operationType(FaultInjectionOperationType.QUERY_ITEM) + .build()) + .result(FaultInjectionResultBuilders + .getResultBuilder(FaultInjectionServerErrorType.GONE) + .build()) + .hitLimit(3) + .build(); + + CosmosFaultInjectionHelper + .configureFaultInjectionRules(container, Arrays.asList(goneRule)) + .block(); + + try { + List faultInjectedResults = container + .readManyByPartitionKeys(pkValues, TestObject.class) + .byPage() + .flatMapIterable(FeedResponse::getResults) + .collectList() + .block(); + + assertThat(faultInjectedResults).isNotNull(); + List faultInjectedIds = faultInjectedResults.stream() + .map(TestObject::getId) + .sorted() + .collect(Collectors.toList()); + + org.assertj.core.api.Assertions.assertThat(faultInjectedIds).doesNotHaveDuplicates(); + org.assertj.core.api.Assertions.assertThat(faultInjectedIds).hasSameElementsAs(baselineIds); + org.assertj.core.api.Assertions.assertThat(faultInjectedIds).hasSize(9); + } finally { + goneRule.disable(); + } + + // Cleanup + for (TestObject item : createdItems) { + container.deleteItem(item.getId(), new PartitionKey(item.getMypk())).block(); + } + } + @AfterClass(groups = {"multi-region", "long", "fast", "fi-multi-master", "multi-region-strong"}, timeOut = SHUTDOWN_TIMEOUT, alwaysRun = true) public void afterClass() { safeClose(clientWithoutPreferredRegions); From 81a99f3e12463398ba2cb4540ae87c20b7be96b0 Mon Sep 17 00:00:00 2001 From: Fabian Meiswinkel Date: Thu, 23 Apr 2026 14:08:13 +0000 Subject: [PATCH 61/64] Update CosmosItemSerializerNoExceptionWrapping.scala --- .../cosmos/CosmosItemSerializerNoExceptionWrapping.scala | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/CosmosItemSerializerNoExceptionWrapping.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/CosmosItemSerializerNoExceptionWrapping.scala index a354cee66efc..9dabaf82c8f6 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/CosmosItemSerializerNoExceptionWrapping.scala +++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/CosmosItemSerializerNoExceptionWrapping.scala @@ -14,4 +14,12 @@ private[cosmos] abstract class CosmosItemSerializerNoExceptionWrapping extends C .CosmosItemSerializerHelper .getCosmosItemSerializerAccessor .setCanSerialize(this, false) + + // canSerialize is set to false above, so the SDK will never call serialize(). + // This default implementation satisfies the abstract method contract and throws if + // called unexpectedly. + override def serialize[T](item: T): java.util.Map[String, AnyRef] = { + throw new UnsupportedOperationException( + "serialize() is not supported on CosmosItemSerializerNoExceptionWrapping (canSerialize = false)") + } } From 5beab22f13d902c67bacf26bfc9c180deb715942 Mon Sep 17 00:00:00 2001 From: Fabian Meiswinkel Date: Thu, 23 Apr 2026 16:44:22 +0000 Subject: [PATCH 62/64] Update ReadManyByPartitionKeyContinuationToken.java --- .../ReadManyByPartitionKeyContinuationToken.java | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyContinuationToken.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyContinuationToken.java index 75449d7af4be..5c972ed3f10f 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyContinuationToken.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyContinuationToken.java @@ -128,7 +128,7 @@ public ReadManyByPartitionKeyContinuationToken( // Tokens written before the version field existed will deserialize with version == null. // Treat null as version 1 (the format that existed when this field was introduced) to // remain forward-compatible with any tokens emitted by an in-flight pre-versioned beta. - int effectiveVersion = (version == null) ? CURRENT_VERSION : version; + int effectiveVersion = (version == null) ? 1 : version; if (effectiveVersion != CURRENT_VERSION) { throw new IllegalArgumentException( "Unsupported readManyByPartitionKeys continuation token version: " + effectiveVersion @@ -307,21 +307,21 @@ private static String toFixedHex(long value) { } /** - * Serializes this token to a Base64-encoded JSON string. + * Serializes this token to a URL-safe Base64-encoded JSON string. */ public String serialize() { try { String json = Utils.getSimpleObjectMapper().writeValueAsString(this); - return Base64.getEncoder().encodeToString(json.getBytes(StandardCharsets.UTF_8)); + return Base64.getUrlEncoder().encodeToString(json.getBytes(StandardCharsets.UTF_8)); } catch (Exception e) { throw new IllegalStateException("Failed to serialize ReadManyByPartitionKeyContinuationToken.", e); } } /** - * Deserializes a Base64-encoded JSON string into a continuation token. + * Deserializes a URL-safe Base64-encoded JSON string into a continuation token. * - * @param serialized the serialized token (Base64 of JSON) + * @param serialized the serialized token (URL-safe Base64 of JSON) * @return the deserialized token * @throws IllegalArgumentException if the token is malformed */ @@ -330,7 +330,7 @@ public static ReadManyByPartitionKeyContinuationToken deserialize(String seriali checkArgument(!serialized.isEmpty(), "Argument 'serialized' must not be empty."); try { - byte[] decoded = Base64.getDecoder().decode(serialized); + byte[] decoded = Base64.getUrlDecoder().decode(serialized); String json = new String(decoded, StandardCharsets.UTF_8); return Utils.getSimpleObjectMapper().readValue(json, ReadManyByPartitionKeyContinuationToken.class); } catch (IllegalArgumentException e) { From 853bd04fc9a8816fe824c099497ca116639093b4 Mon Sep 17 00:00:00 2001 From: Fabian Meiswinkel Date: Thu, 23 Apr 2026 18:20:31 +0000 Subject: [PATCH 63/64] Update DocumentQueryExecutionContextFactory.java --- .../DocumentQueryExecutionContextFactory.java | 19 +++++++++++++------ 1 file changed, 13 insertions(+), 6 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/DocumentQueryExecutionContextFactory.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/DocumentQueryExecutionContextFactory.java index 82506ae361e1..11528ed4f7ea 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/DocumentQueryExecutionContextFactory.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/DocumentQueryExecutionContextFactory.java @@ -247,11 +247,11 @@ private static Mono getTargetRangesFromEmpt "Query plan retrieval must not be suppressed when not using FeedRanges"); } - QueryInfo queryInfo = QueryInfo.EMPTY; - queryInfo.setQueryPlanDiagnosticsContext( - new QueryInfo.QueryPlanDiagnosticsContext( - planFetchStartTime, - planFetchEndTime)); + // Do NOT use QueryInfo.EMPTY here — setQueryPlanDiagnosticsContext would mutate + // the shared static singleton. Instead, capture the diagnostics context and create + // a fresh per-request QueryInfo inside the reactive chain. + QueryInfo.QueryPlanDiagnosticsContext queryPlanDiagnosticsContext = + new QueryInfo.QueryPlanDiagnosticsContext(planFetchStartTime, planFetchEndTime); FeedRange userProvidedFeedRange = cosmosQueryRequestOptions.getFeedRange(); Mono> targetRange = queryExecutionContext @@ -264,7 +264,14 @@ private static Mono getTargetRangesFromEmpt ).map(pkRanges -> pkRanges.stream().map(PartitionKeyRange::toRange).collect(Collectors.toList())); return Mono.zip(targetRange, allRanges) - .map(tuple -> new PartitionKeyRangesAndQueryInfos(queryInfo, null, Collections.singletonList(tuple.getT1()), tuple.getT2())); + .map(tuple -> { + // Set diagnostics on a per-request basis inside the reactive chain, + // not on the shared EMPTY singleton. Create a fresh QueryInfo only here. + QueryInfo perRequestQueryInfo = new QueryInfo(); + perRequestQueryInfo.setQueryPlanDiagnosticsContext(queryPlanDiagnosticsContext); + return new PartitionKeyRangesAndQueryInfos( + perRequestQueryInfo, null, Collections.singletonList(tuple.getT1()), tuple.getT2()); + }); } synchronized private static void tryCacheQueryPlan( From 53ad60da624cbe025c1837aea8944048a2dd072e Mon Sep 17 00:00:00 2001 From: Fabian Meiswinkel Date: Thu, 23 Apr 2026 20:58:50 +0000 Subject: [PATCH 64/64] Update FaultInjectionWithAvailabilityStrategyTestsBase.java --- ...tionWithAvailabilityStrategyTestsBase.java | 24 +++++++------------ 1 file changed, 9 insertions(+), 15 deletions(-) diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/FaultInjectionWithAvailabilityStrategyTestsBase.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/FaultInjectionWithAvailabilityStrategyTestsBase.java index 5b1243560ecf..296a1c06b7ac 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/FaultInjectionWithAvailabilityStrategyTestsBase.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/FaultInjectionWithAvailabilityStrategyTestsBase.java @@ -3905,21 +3905,15 @@ private CosmosResponseWrapper readManyByPartitionKeysCore( options.setCosmosEndToEndOperationLatencyPolicyConfig(e2ePolicy); List> returnedPages; - try { - returnedPages = params.container - .readManyByPartitionKeys(pkValues, options, ObjectNode.class) - .byPage() - .collectList() - .block(); - } catch (CosmosException cosmosException) { - return new CosmosResponseWrapper( - cosmosException.getDiagnostics() != null - ? new CosmosDiagnosticsContext[] { cosmosException.getDiagnostics().getDiagnosticsContext() } - : null, - cosmosException.getStatusCode(), - cosmosException.getSubStatusCode(), - null); - } + // Let CosmosException propagate to execute()'s catch block — it handles + // status code + sub-status code validation and diagnostics context extraction + // correctly for error cases (the response-level validation path passes null + // for sub-status which breaks validators like validateStatusCodeIsServiceUnavailable). + returnedPages = params.container + .readManyByPartitionKeys(pkValues, options, ObjectNode.class) + .byPage() + .collectList() + .block(); ArrayList foundCtxs = new ArrayList<>();