Azure · FabianMeiswinkel · Apr 13, 2026 · Apr 13, 2026 · Apr 14, 2026 · Apr 14, 2026
@@ -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
 

@@ -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
 

@@ -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
 

@@ -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
 

@@ -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)
   }
 }
 

@@ -45,6 +45,7 @@ private[cosmos] object CosmosConstants {
     val Id = "id"
     val ETag = "_etag"
     val ItemIdentity = "_itemIdentity"
+    val PartitionKeyIdentity = "_partitionKeyIdentity"
   }
 
   object StatusCodes {

@@ -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()
+  }
 }
@@ -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
+    }
+  }
+}