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

sdk/cosmos/azure-cosmos/CHANGELOG.md

@@ -8,6 +8,8 @@
 
 #### Bugs Fixed
 * 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 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 48702](https://github.com/Azure/azure-sdk-for-java/pull/48702)
+* Fixed an issue where `SqlParameter` ignored the configured `CustomItemSerializer`, always using the internal default serializer instead. - See [PR 48702](https://github.com/Azure/azure-sdk-for-java/pull/48702)
 
 #### Other Changes
 

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;
@@ -1933,4 +1934,39 @@ 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);
+            SqlParameter cloneSqlParameter(SqlParameter original);
+        }
+    }
 }

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

@@ -65,12 +65,12 @@ private static BiFunction<String, PipelinedDocumentQueryParams<Document>, Flux<I
                             HttpConstants.StatusCodes.BADREQUEST,
                             "Can not use a continuation token for a vector search query");
                     }
-                    qryOptAccessor.getImpl(orderByCosmosQueryRequestOptions).setCustomItemSerializer(null);
+                    qryOptAccessor.getImpl(orderByCosmosQueryRequestOptions).setCustomItemSerializer(CosmosItemSerializer.DEFAULT_SERIALIZER);
                     documentQueryParams.setCosmosQueryRequestOptions(orderByCosmosQueryRequestOptions);
                     return NonStreamingOrderByDocumentQueryExecutionContext.createAsync(diagnosticsClientContext, client, documentQueryParams, collection);
                 } else {
                     ModelBridgeInternal.setQueryRequestOptionsContinuationToken(orderByCosmosQueryRequestOptions, continuationToken);
-                    qryOptAccessor.getImpl(orderByCosmosQueryRequestOptions).setCustomItemSerializer(null);
+                    qryOptAccessor.getImpl(orderByCosmosQueryRequestOptions).setCustomItemSerializer(CosmosItemSerializer.DEFAULT_SERIALIZER);
                     documentQueryParams.setCosmosQueryRequestOptions(orderByCosmosQueryRequestOptions);
                     return OrderByDocumentQueryExecutionContext.createAsync(diagnosticsClientContext, client, documentQueryParams, collection);
                 }
@@ -80,7 +80,7 @@ private static BiFunction<String, PipelinedDocumentQueryParams<Document>, Flux<I
             createBaseComponentFunction = (continuationToken, documentQueryParams) -> {
                 CosmosQueryRequestOptions parallelCosmosQueryRequestOptions =
                     qryOptAccessor.clone(requestOptions);
-                qryOptAccessor.getImpl(parallelCosmosQueryRequestOptions).setCustomItemSerializer(null);
+                qryOptAccessor.getImpl(parallelCosmosQueryRequestOptions).setCustomItemSerializer(CosmosItemSerializer.DEFAULT_SERIALIZER);
                 ModelBridgeInternal.setQueryRequestOptionsContinuationToken(parallelCosmosQueryRequestOptions, continuationToken);
 
                 documentQueryParams.setCosmosQueryRequestOptions(parallelCosmosQueryRequestOptions);
@@ -119,6 +119,7 @@ private static BiFunction<String, PipelinedDocumentQueryParams<Document>, Flux<I
             CosmosQueryRequestOptions orderByCosmosQueryRequestOptions =
                 qryOptAccessor.clone(requestOptions);
 
+            qryOptAccessor.getImpl(orderByCosmosQueryRequestOptions).setCustomItemSerializer(CosmosItemSerializer.DEFAULT_SERIALIZER);
             documentQueryParams.setCosmosQueryRequestOptions(orderByCosmosQueryRequestOptions);
             return HybridSearchDocumentQueryExecutionContext.createAsync(diagnosticsClientContext, client, documentQueryParams, collection);
         };

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,17 @@
 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.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 +65,26 @@ 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) {
+            SqlQuerySpec original = initParams.getQuery();
+            List<SqlParameter> clonedParams = new ArrayList<>();
+            List<SqlParameter> originalParams = original.getParameters();
+            ImplementationBridgeHelpers.SqlQuerySpecHelper.SqlQuerySpecAccessor sqlQuerySpecAccessor =
+                ImplementationBridgeHelpers.SqlQuerySpecHelper.getSqlQuerySpecAccessor();
+            if (originalParams != null) {
+                for (SqlParameter p : originalParams) {
+                    clonedParams.add(sqlQuerySpecAccessor.cloneSqlParameter(p));
+                }
+            }
+            SqlQuerySpec clonedQuery = new SqlQuerySpec(original.getQueryText(), clonedParams);
+            sqlQuerySpecAccessor.applySerializerToParameters(clonedQuery, candidateSerializer);
+            initParams.setQuery(clonedQuery);
+        }
+
         if (hybridSearchQueryInfo != null) {
             return PipelinedDocumentQueryExecutionContext.createHybridAsyncCore(
                 diagnosticsClientContext,

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/ModelBridgeInternal.java

@@ -5,6 +5,7 @@
 
 import com.azure.cosmos.ConsistencyLevel;
 import com.azure.cosmos.CosmosDiagnostics;
+import com.azure.cosmos.CosmosItemSerializer;
 import com.azure.cosmos.implementation.ClientEncryptionKey;
 import com.azure.cosmos.implementation.Conflict;
 import com.azure.cosmos.implementation.CosmosPagedFluxOptions;
@@ -761,5 +762,6 @@ public static void  initializeAllAccessors() {
         CosmosClientTelemetryConfig.initialize();
         CosmosContainerIdentity.initialize();
         PriorityLevel.initialize();
+        SqlQuerySpec.initialize();
     }
 }

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/SqlParameter.java

@@ -3,16 +3,21 @@
 
 package com.azure.cosmos.models;
 
+import com.azure.cosmos.CosmosItemSerializer;
 import com.azure.cosmos.implementation.JsonSerializable;
 import com.fasterxml.jackson.databind.node.ObjectNode;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 import java.util.Objects;
 
 /**
  * Represents a SQL parameter in the SqlQuerySpec used for queries in the Azure Cosmos DB database service.
  */
 public final class SqlParameter {
+    private static final Logger LOGGER = LoggerFactory.getLogger(SqlParameter.class);
     private JsonSerializable jsonSerializable;
+    private Object rawValue;
 
     /**
      * Initializes a new instance of the SqlParameter class.
@@ -29,6 +34,7 @@ public SqlParameter() {
     SqlParameter(ObjectNode objectNode) {
 
         this.jsonSerializable = new JsonSerializable(objectNode);
+        this.rawValue = this.jsonSerializable.getObject("value", Object.class);
     }
 
     /**
@@ -84,13 +90,60 @@ public <T> T getValue(Class<T> classType) {
      * @return the SqlParameter.
      */
     public SqlParameter setValue(Object value) {
+        this.rawValue = value;
         this.jsonSerializable.set(
             "value",
             value
         );
         return this;
     }
 
+    /**
+     * Returns the original raw value before any serialization.
+     * Package-private — used internally to clone parameters while preserving the
+     * original Java type (e.g. Instant) that would otherwise be lost when reading
+     * from the JSON property bag.
+     */
+    Object getRawValue() {
+        return this.rawValue;
+    }
+
+    /**
+     * Re-serializes the parameter value using the given custom item serializer.
+     * This is called internally during query execution to ensure parameter values
+     * are serialized consistently with document values when a custom serializer is configured.
+     *
+     * @param serializer the custom item serializer to apply.
+     */
+    void applySerializer(CosmosItemSerializer serializer) {
+        if (this.rawValue != null && serializer != null) {
+            try {
+                this.jsonSerializable.set("value", this.rawValue, serializer, true);
+            } catch (UnsupportedOperationException e) {
+                // Some serializer implementations (e.g. Spark connector) only implement
+                // deserialize() and stub serialize() as unimplemented. Fall back to the
+                // default serialization that was already applied via setValue().
+                LOGGER.debug(
+                    "Custom serializer '{}' does not support serialize(); "
+                        + "falling back to default serialization for SqlParameter value.",
+                    serializer.getClass().getSimpleName(),
+                    e);
+            } catch (Error e) {
+                // Scala's ??? operator throws scala.NotImplementedError (extends Error).
+                // Only swallow that specific case; re-throw all other Errors.
+                if ("scala.NotImplementedError".equals(e.getClass().getName())) {
+                    LOGGER.debug(
+                        "Custom serializer '{}' does not support serialize(); "
+                            + "falling back to default serialization for SqlParameter value.",
+                        serializer.getClass().getSimpleName(),
+                        e);
+                } else {
+                    throw e;
+                }
+            }
+        }
+    }
+
     void populatePropertyBag() {
         this.jsonSerializable.populatePropertyBag();
     }

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/SqlQuerySpec.java

@@ -3,6 +3,8 @@
 
 package com.azure.cosmos.models;
 
+import com.azure.cosmos.CosmosItemSerializer;
+import com.azure.cosmos.implementation.ImplementationBridgeHelpers;
 import com.azure.cosmos.implementation.JsonSerializable;
 import com.fasterxml.jackson.databind.node.ObjectNode;
 
@@ -134,5 +136,41 @@ void populatePropertyBag() {
         }
     }
 
+    /**
+     * Re-serializes all parameter values using the given custom item serializer.
+     * Called internally during query execution to ensure parameter values are serialized
+     * consistently with document values when a custom serializer is configured.
+     *
+     * @param serializer the custom item serializer to apply.
+     */
+    void applySerializerToParameters(CosmosItemSerializer serializer) {
+        List<SqlParameter> params = this.getParameters();
+        if (serializer != null && params != null) {
+            for (SqlParameter param : params) {
+                param.applySerializer(serializer);
+            }
+        }
+    }
+
     JsonSerializable getJsonSerializable() { return this.jsonSerializable; }
+
+    static void initialize() {
+        ImplementationBridgeHelpers.SqlQuerySpecHelper.setSqlQuerySpecAccessor(
+            new ImplementationBridgeHelpers.SqlQuerySpecHelper.SqlQuerySpecAccessor() {
+                @Override
+                public void applySerializerToParameters(SqlQuerySpec sqlQuerySpec, CosmosItemSerializer serializer) {
+                    if (sqlQuerySpec != null) {
+                        sqlQuerySpec.applySerializerToParameters(serializer);
+                    }
+                }
+
+                @Override
+                public SqlParameter cloneSqlParameter(SqlParameter original) {
+                    return new SqlParameter(original.getName(), original.getRawValue());
+                }
+            }
+        );
+    }
+
+    static { initialize(); }
 }