fix: CustomItemSerializer not applied correctly in queries and SqlParameter [Cosmos]

sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/CosmosItemSerializerNoExceptionWrapping.scala

@@ -9,4 +9,9 @@ private[cosmos] abstract class CosmosItemSerializerNoExceptionWrapping extends C
     .CosmosItemSerializerHelper
     .getCosmosItemSerializerAccessor
     .setShouldWrapSerializationExceptions(this, false)
+
+  ImplementationBridgeHelpers
+    .CosmosItemSerializerHelper
+    .getCosmosItemSerializerAccessor
+    .setCanSerialize(this, false)
 }

sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/CosmosItemSerializerNoExceptionWrapping.java

@@ -14,5 +14,10 @@ public CosmosItemSerializerNoExceptionWrapping() {
             .CosmosItemSerializerHelper
             .getCosmosItemSerializerAccessor()
             .setShouldWrapSerializationExceptions(this, false);
+
+        ImplementationBridgeHelpers
+            .CosmosItemSerializerHelper
+            .getCosmosItemSerializerAccessor()
+            .setCanSerialize(this, false);
     }
 }

sdk/cosmos/azure-cosmos/CHANGELOG.md

@@ -10,6 +10,8 @@
 * 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 an issue where change feed with `startFrom` point-in-time returned `400` on merged partitions by enabling the `CHANGE_FEED_WITH_START_TIME_POST_MERGE` SDK capability.
 * Fixed JVM `<clinit>` 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)
+* Fixed an issue where `SqlParameter` ignored the configured `CustomItemSerializer`, always using the internal default serializer instead. - See [PR 48811](https://github.com/Azure/azure-sdk-for-java/pull/48811)
 
 #### Other Changes
 

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosItemSerializer.java

@@ -3,27 +3,16 @@
 
 package com.azure.cosmos;
 
-import com.azure.cosmos.implementation.ApiType;
 import com.azure.cosmos.implementation.BadRequestException;
 import com.azure.cosmos.implementation.Configs;
-import com.azure.cosmos.implementation.ConnectionPolicy;
-import com.azure.cosmos.implementation.CosmosClientMetadataCachesSnapshot;
 import com.azure.cosmos.implementation.DefaultCosmosItemSerializer;
 import com.azure.cosmos.implementation.HttpConstants;
 import com.azure.cosmos.implementation.ImplementationBridgeHelpers;
-import com.azure.cosmos.implementation.JsonSerializable;
-import com.azure.cosmos.implementation.ObjectNodeMap;
-import com.azure.cosmos.implementation.PrimitiveJsonNodeMap;
 import com.azure.cosmos.implementation.Utils;
-import com.fasterxml.jackson.databind.JsonNode;
 import com.fasterxml.jackson.databind.ObjectMapper;
-import com.fasterxml.jackson.databind.node.ObjectNode;
 
-import java.io.IOException;
 import java.util.Map;
 
-import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkNotNull;
-
 /**
  * The {@link CosmosItemSerializer} allows customizing the serialization of Cosmos Items - either to transform payload (for
  * example wrap/unwrap in custom envelopes) or use custom serialization settings or json serializer stacks.
@@ -47,6 +36,7 @@ public abstract class CosmosItemSerializer {
         new DefaultCosmosItemSerializer(Utils.getSimpleObjectMapper());
 
     private boolean shouldWrapSerializationExceptions;
+    private boolean canSerialize;
 
     private ObjectMapper mapper = Utils.getSimpleObjectMapper();
 
@@ -55,6 +45,7 @@ public abstract class CosmosItemSerializer {
      */
     protected CosmosItemSerializer() {
         this.shouldWrapSerializationExceptions = true;
+        this.canSerialize = true;
     }
 
     /**
@@ -139,6 +130,14 @@ void setShouldWrapSerializationExceptions(boolean enabled) {
         this.shouldWrapSerializationExceptions = enabled;
     }
 
+    void setCanSerialize(boolean canSerialize) {
+        this.canSerialize = canSerialize;
+    }
+
+    boolean canSerialize() {
+        return this.canSerialize;
+    }
+
 
     ///////////////////////////////////////////////////////////////////////////////////////////
     // the following helper/accessor only helps to access this class outside of this package.//
@@ -172,6 +171,15 @@ public ObjectMapper getItemObjectMapper(CosmosItemSerializer serializer) {
                 }
 
                 @Override
+                public boolean canSerialize(CosmosItemSerializer serializer) {
+                    return serializer.canSerialize();
+                }
+
+                @Override
+                public void setCanSerialize(CosmosItemSerializer serializer, boolean canSerialize) {
+                    serializer.setCanSerialize(canSerialize);
+                }
+
                 public CosmosItemSerializer getInternalDefaultSerializer() {
                     return INTERNAL_DEFAULT_SERIALIZER;
                 }

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ImplementationBridgeHelpers.java

@@ -77,6 +77,7 @@
 import com.azure.cosmos.models.PartitionKeyDefinition;
 import com.azure.cosmos.models.PriorityLevel;
 import com.azure.cosmos.models.ShowQueryMode;
+import com.azure.cosmos.models.SqlParameter;
 import com.azure.cosmos.models.SqlQuerySpec;
 import com.azure.cosmos.util.CosmosPagedFlux;
 import com.azure.cosmos.util.UtilBridgeInternal;
@@ -1893,6 +1894,8 @@ void setShouldWrapSerializationExceptions(
             void setItemObjectMapper(CosmosItemSerializer serializer, ObjectMapper mapper);
             ObjectMapper getItemObjectMapper(CosmosItemSerializer serializer);
             CosmosItemSerializer getInternalDefaultSerializer();
+            void setCanSerialize(CosmosItemSerializer serializer, boolean canSerialize);
+            boolean canSerialize(CosmosItemSerializer serializer);
         }
     }
 
@@ -1934,4 +1937,72 @@ ReadConsistencyStrategy getEffectiveReadConsistencyStrategy(
                 ReadConsistencyStrategy clientLevelReadConsistencyStrategy);
         }
     }
+
+    public static final class SqlQuerySpecHelper {
+        private static final AtomicReference<SqlQuerySpecAccessor> accessor = new AtomicReference<>();
+        private static final AtomicBoolean sqlQuerySpecClassLoaded = new AtomicBoolean(false);
+
+        private SqlQuerySpecHelper() {}
+
+        public static void setSqlQuerySpecAccessor(final SqlQuerySpecAccessor newAccessor) {
+            if (!accessor.compareAndSet(null, newAccessor)) {
+                logger.debug("SqlQuerySpecAccessor already initialized!");
+            } else {
+                logger.debug("Setting SqlQuerySpecAccessor...");
+                sqlQuerySpecClassLoaded.set(true);
+            }
+        }
+
+        public static SqlQuerySpecAccessor getSqlQuerySpecAccessor() {
+            if (!sqlQuerySpecClassLoaded.get()) {
+                logger.debug("Initializing SqlQuerySpecAccessor...");
+                initializeAllAccessors();
+            }
+
+            SqlQuerySpecAccessor snapshot = accessor.get();
+            if (snapshot == null) {
+                logger.error("SqlQuerySpecAccessor is not initialized yet!");
+            }
+
+            return snapshot;
+        }
+
+        public interface SqlQuerySpecAccessor {
+            void applySerializerToParameters(SqlQuerySpec sqlQuerySpec, CosmosItemSerializer serializer);
+        }
+    }
+
+    public static final class SqlParameterHelper {
+        private static final AtomicReference<SqlParameterAccessor> accessor = new AtomicReference<>();
+        private static final AtomicBoolean sqlParameterClassLoaded = new AtomicBoolean(false);
+
+        private SqlParameterHelper() {}
+
+        public static void setSqlParameterAccessor(final SqlParameterAccessor newAccessor) {
+            if (!accessor.compareAndSet(null, newAccessor)) {
+                logger.debug("SqlParameterAccessor already initialized!");
+            } else {
+                logger.debug("Setting SqlParameterAccessor...");
+                sqlParameterClassLoaded.set(true);
+            }
+        }
+
+        public static SqlParameterAccessor getSqlParameterAccessor() {
+            if (!sqlParameterClassLoaded.get()) {
+                logger.debug("Initializing SqlParameterAccessor...");
+                initializeAllAccessors();
+            }
+
+            SqlParameterAccessor snapshot = accessor.get();
+            if (snapshot == null) {
+                logger.error("SqlParameterAccessor is not initialized yet!");
+            }
+
+            return snapshot;
+        }
+
+        public interface SqlParameterAccessor {
+            SqlParameter createCopy(SqlParameter original);
+        }
+    }
 }

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/PipelinedDocumentQueryExecutionContext.java

@@ -63,29 +63,26 @@ private static BiFunction<String, PipelinedDocumentQueryParams<Document>, Flux<I
         if (queryInfo.hasOrderBy()) {
             createBaseComponentFunction = (continuationToken, documentQueryParams) -> {
                 CosmosQueryRequestOptions orderByCosmosQueryRequestOptions =
-                    qryOptAccessor().clone(requestOptions);
+                    cloneOptionsForInternalPipeline(requestOptions);
                 if (queryInfo.hasNonStreamingOrderBy()) {
                     if (continuationToken != null) {
                         throw new NonStreamingOrderByBadRequestException(
                             HttpConstants.StatusCodes.BADREQUEST,
                             "Can not use a continuation token for a vector search query");
                     }
-                    qryOptAccessor().getImpl(orderByCosmosQueryRequestOptions).setCustomItemSerializer(null);
+
                     documentQueryParams.setCosmosQueryRequestOptions(orderByCosmosQueryRequestOptions);
                     return NonStreamingOrderByDocumentQueryExecutionContext.createAsync(diagnosticsClientContext, client, documentQueryParams, collection);
                 } else {
                     ModelBridgeInternal.setQueryRequestOptionsContinuationToken(orderByCosmosQueryRequestOptions, continuationToken);
-                    qryOptAccessor().getImpl(orderByCosmosQueryRequestOptions).setCustomItemSerializer(null);
                     documentQueryParams.setCosmosQueryRequestOptions(orderByCosmosQueryRequestOptions);
                     return OrderByDocumentQueryExecutionContext.createAsync(diagnosticsClientContext, client, documentQueryParams, collection);
                 }
             };
         } else {
 
             createBaseComponentFunction = (continuationToken, documentQueryParams) -> {
-                CosmosQueryRequestOptions parallelCosmosQueryRequestOptions =
-                    qryOptAccessor().clone(requestOptions);
-                qryOptAccessor().getImpl(parallelCosmosQueryRequestOptions).setCustomItemSerializer(null);
+                CosmosQueryRequestOptions parallelCosmosQueryRequestOptions = cloneOptionsForInternalPipeline(requestOptions);
                 ModelBridgeInternal.setQueryRequestOptionsContinuationToken(parallelCosmosQueryRequestOptions, continuationToken);
 
                 documentQueryParams.setCosmosQueryRequestOptions(parallelCosmosQueryRequestOptions);
@@ -122,7 +119,7 @@ private static BiFunction<String, PipelinedDocumentQueryParams<Document>, Flux<I
 
         createBaseComponentFunction = (continuationToken, documentQueryParams) -> {
             CosmosQueryRequestOptions orderByCosmosQueryRequestOptions =
-                qryOptAccessor().clone(requestOptions);
+                cloneOptionsForInternalPipeline(requestOptions);
 
             documentQueryParams.setCosmosQueryRequestOptions(orderByCosmosQueryRequestOptions);
             return HybridSearchDocumentQueryExecutionContext.createAsync(diagnosticsClientContext, client, documentQueryParams, collection);
@@ -273,4 +270,24 @@ public Flux<FeedResponse<T>> executeAsync() {
                 )
             );
     }
+
+    /**
+     * Clones query request options and neutralizes the custom item serializer to
+     * DEFAULT_SERIALIZER. Internal pipeline stages must always use the default
+     * serializer; custom serialization is applied at the top-level pipeline boundary.
+     *
+     * All internal pipeline stages that work with Document/intermediate types MUST
+     * use this method rather than cloning options directly via
+     * {@code qryOptAccessor().clone(...)}. Bypassing this method would leave the
+     * custom serializer active in the inner pipeline, causing incorrect
+     * serialization of intermediate results. See PR #48811 for context.
+     */
+    private static CosmosQueryRequestOptions cloneOptionsForInternalPipeline(
+        CosmosQueryRequestOptions source) {
+
+        CosmosQueryRequestOptions cloned = qryOptAccessor().clone(source);
+        qryOptAccessor().getImpl(cloned)
+            .setCustomItemSerializer(CosmosItemSerializer.DEFAULT_SERIALIZER);
+        return cloned;
+    }
 }

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/PipelinedDocumentQueryParams.java

@@ -21,7 +21,7 @@ public class PipelinedDocumentQueryParams<T> {
     private final String collectionRid;
     private final ResourceType resourceTypeEnum;
     private final Class<T> resourceType;
-    private final SqlQuerySpec query;
+    private SqlQuerySpec query;
     private final String resourceLink;
     private final UUID correlatedActivityId;
     private CosmosQueryRequestOptions cosmosQueryRequestOptions;
@@ -101,6 +101,10 @@ public SqlQuerySpec getQuery() {
         return query;
     }
 
+    public void setQuery(SqlQuerySpec query) {
+        this.query = query;
+    }
+
     public String getResourceLink() {
         return resourceLink;
     }

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/PipelinedQueryExecutionContextBase.java

@@ -5,12 +5,20 @@
 import com.azure.cosmos.CosmosItemSerializer;
 import com.azure.cosmos.implementation.DiagnosticsClientContext;
 import com.azure.cosmos.implementation.DocumentCollection;
+import com.azure.cosmos.implementation.ImplementationBridgeHelpers;
+import com.azure.cosmos.implementation.ImplementationBridgeHelpers.CosmosItemSerializerHelper.CosmosItemSerializerAccessor;
+import com.azure.cosmos.implementation.ImplementationBridgeHelpers.SqlParameterHelper.SqlParameterAccessor;
+import com.azure.cosmos.implementation.ImplementationBridgeHelpers.SqlQuerySpecHelper.SqlQuerySpecAccessor;
 import com.azure.cosmos.implementation.Utils;
 import com.azure.cosmos.implementation.query.hybridsearch.HybridSearchQueryInfo;
 import com.azure.cosmos.models.CosmosQueryRequestOptions;
 import com.azure.cosmos.models.ModelBridgeInternal;
+import com.azure.cosmos.models.SqlParameter;
+import com.azure.cosmos.models.SqlQuerySpec;
 import reactor.core.publisher.Flux;
 
+import java.util.ArrayList;
+import java.util.List;
 import java.util.function.BiFunction;
 
 /**
@@ -60,6 +68,28 @@ public static <T> Flux<PipelinedQueryExecutionContextBase<T>> createAsync(
 
         CosmosItemSerializer candidateSerializer = client.getEffectiveItemSerializer(cosmosQueryRequestOptions);
 
+        // Apply the effective custom serializer to SqlParameter values so that
+        // query parameters are serialized consistently with stored document values.
+        // Clone the SqlQuerySpec first to avoid mutating the caller's original object,
+        // which could race if the same instance is reused across concurrent queries.
+        if (candidateSerializer != CosmosItemSerializer.DEFAULT_SERIALIZER
+            && cosmosItemSerializerAccessor().canSerialize(candidateSerializer)) {
+            SqlQuerySpec original = initParams.getQuery();
+            if (original != null) {
+                List<SqlParameter> originalParams = original.getParameters();
+                if (originalParams != null && !originalParams.isEmpty()) {
+                    List<SqlParameter> clonedParams = new ArrayList<>(originalParams.size());
+                    SqlParameterAccessor paramAccessor = sqlParameterAccessor();
+                    for (SqlParameter p : originalParams) {
+                        clonedParams.add(paramAccessor.createCopy(p));
+                    }
+                    SqlQuerySpec clonedQuery = new SqlQuerySpec(original.getQueryText(), clonedParams);
+                    sqlQuerySpecAccessor().applySerializerToParameters(clonedQuery, candidateSerializer);
+                    initParams.setQuery(clonedQuery);
+                }
+            }
+        }
+
         if (hybridSearchQueryInfo != null) {
             return PipelinedDocumentQueryExecutionContext.createHybridAsyncCore(
                 diagnosticsClientContext,
@@ -194,4 +224,16 @@ public QueryInfo getQueryInfo() {
     public HybridSearchQueryInfo getHybridSearchQueryInfo() {
         return this.hybridSearchQueryInfo;
     }
+
+    private static SqlQuerySpecAccessor sqlQuerySpecAccessor() {
+        return ImplementationBridgeHelpers.SqlQuerySpecHelper.getSqlQuerySpecAccessor();
+    }
+
+    private static SqlParameterAccessor sqlParameterAccessor() {
+        return ImplementationBridgeHelpers.SqlParameterHelper.getSqlParameterAccessor();
+    }
+
+    private static CosmosItemSerializerAccessor cosmosItemSerializerAccessor() {
+        return ImplementationBridgeHelpers.CosmosItemSerializerHelper.getCosmosItemSerializerAccessor();
+    }
 }

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/QueryPlanRetriever.java

@@ -15,6 +15,7 @@
 import com.azure.cosmos.models.CosmosQueryRequestOptions;
 import com.azure.cosmos.models.ModelBridgeInternal;
 import com.azure.cosmos.models.PartitionKey;
+import com.azure.cosmos.models.SqlParameter;
 import com.azure.cosmos.models.SqlQuerySpec;
 import com.azure.cosmos.implementation.BackoffRetryUtility;
 import com.azure.cosmos.implementation.DocumentClientRetryPolicy;
@@ -26,6 +27,7 @@
 import com.fasterxml.jackson.databind.node.ObjectNode;
 import reactor.core.publisher.Mono;
 
+import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
@@ -105,7 +107,13 @@ static Mono<PartitionedQueryExecutionInfo> getQueryPlanThroughGatewayAsync(Diagn
                                                                                  resourceLink,
                                                                                  requestHeaders);
         queryPlanRequest.useGatewayMode = true;
-        queryPlanRequest.setByteBuffer(ModelBridgeInternal.serializeJsonToByteBuffer(sqlQuerySpec));
+
+        // Create a defensive copy to prevent concurrent modification of the shared
+        // SqlQuerySpec's internal ObjectNode when multiple threads retrieve the query
+        // plan simultaneously. Each copy has its own property bag, avoiding the race
+        // condition on the non-thread-safe ObjectNode/LinkedHashMap backing store.
+        SqlQuerySpec querySpecCopy = copyQuerySpec(sqlQuerySpec);
+        queryPlanRequest.setByteBuffer(ModelBridgeInternal.serializeJsonToByteBuffer(querySpecCopy));
 
         CosmosEndToEndOperationLatencyPolicyConfig end2EndConfig = queryOptionsAccessor()
             .getImpl(nonNullRequestOptions)
@@ -162,4 +170,10 @@ static Mono<PartitionedQueryExecutionInfo> getQueryPlanThroughGatewayAsync(Diagn
                 return throwable;
             });
     }
+
+    private static SqlQuerySpec copyQuerySpec(SqlQuerySpec original) {
+        List<SqlParameter> params = original.getParameters();
+        List<SqlParameter> copiedParams = params != null ? new ArrayList<>(params) : null;
+        return new SqlQuerySpec(original.getQueryText(), copiedParams);
+    }
 }