Query: Adds an environment variable for disabling the hybrid search query plan optimization

Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs

@@ -584,7 +584,7 @@ private static async Task<PartitionedQueryExecutionInfo> GetPartitionedQueryExec
                     inputParameters.SqlQuerySpec,
                     cosmosQueryContext.ResourceLink,
                     inputParameters.PartitionKey,
-                    inputParameters.IsNonStreamingOrderByQueryFeatureDisabled,
+                    inputParameters.IsHybridSearchQueryPlanOptimizationDisabled,
                     trace,
                     cancellationToken);
             }
@@ -601,6 +601,7 @@ private static async Task<PartitionedQueryExecutionInfo> GetPartitionedQueryExec
                     inputParameters.PartitionKey != null,
                     containerQueryProperties.GeospatialType,
                     cosmosQueryContext.UseSystemPrefix,
+                    inputParameters.IsHybridSearchQueryPlanOptimizationDisabled,
                     trace,
                     cancellationToken);
             }
@@ -832,7 +833,7 @@ private InputParameters(
                 PartitionedQueryExecutionInfo partitionedQueryExecutionInfo,
                 bool returnResultsInDeterministicOrder,
                 bool enableOptimisticDirectExecution,
-                bool isNonStreamingOrderByQueryFeatureDisabled,
+                bool isHybridSearchQueryPlanOptimizationDisabled,
                 bool enableDistributedQueryGatewayMode,
                 TestInjections testInjections)
             {
@@ -847,7 +848,7 @@ private InputParameters(
                 this.PartitionedQueryExecutionInfo = partitionedQueryExecutionInfo;
                 this.ReturnResultsInDeterministicOrder = returnResultsInDeterministicOrder;
                 this.EnableOptimisticDirectExecution = enableOptimisticDirectExecution;
-                this.IsNonStreamingOrderByQueryFeatureDisabled = isNonStreamingOrderByQueryFeatureDisabled;
+                this.IsHybridSearchQueryPlanOptimizationDisabled = isHybridSearchQueryPlanOptimizationDisabled;
                 this.EnableDistributedQueryGatewayMode = enableDistributedQueryGatewayMode;
                 this.TestInjections = testInjections;
             }
@@ -864,7 +865,7 @@ public static InputParameters Create(
                 PartitionedQueryExecutionInfo partitionedQueryExecutionInfo,
                 bool? returnResultsInDeterministicOrder,
                 bool enableOptimisticDirectExecution,
-                bool isNonStreamingOrderByQueryFeatureDisabled,
+                bool isHybridSearchQueryPlanOptimizationDisabled,
                 bool enableDistributedQueryGatewayMode,
                 TestInjections testInjections)
             {
@@ -903,7 +904,7 @@ public static InputParameters Create(
                     partitionedQueryExecutionInfo: partitionedQueryExecutionInfo,
                     returnResultsInDeterministicOrder: returnResultsInDeterministicOrder.GetValueOrDefault(InputParameters.DefaultReturnResultsInDeterministicOrder),
                     enableOptimisticDirectExecution: enableOptimisticDirectExecution,
-                    isNonStreamingOrderByQueryFeatureDisabled: isNonStreamingOrderByQueryFeatureDisabled,
+                    isHybridSearchQueryPlanOptimizationDisabled: isHybridSearchQueryPlanOptimizationDisabled,
                     enableDistributedQueryGatewayMode: enableDistributedQueryGatewayMode,
                     testInjections: testInjections);
             }
@@ -920,7 +921,7 @@ public static InputParameters Create(
             public bool ReturnResultsInDeterministicOrder { get; }
             public TestInjections TestInjections { get; }
             public bool EnableOptimisticDirectExecution { get; }
-            public bool IsNonStreamingOrderByQueryFeatureDisabled { get; }
+            public bool IsHybridSearchQueryPlanOptimizationDisabled { get; }
             public bool EnableDistributedQueryGatewayMode { get; }
 
             public InputParameters WithContinuationToken(CosmosElement token)
@@ -937,7 +938,7 @@ public InputParameters WithContinuationToken(CosmosElement token)
                     this.PartitionedQueryExecutionInfo,
                     this.ReturnResultsInDeterministicOrder,
                     this.EnableOptimisticDirectExecution,
-                    this.IsNonStreamingOrderByQueryFeatureDisabled,
+                    this.IsHybridSearchQueryPlanOptimizationDisabled,
                     this.EnableDistributedQueryGatewayMode,
                     this.TestInjections);
             }

Microsoft.Azure.Cosmos/src/Query/Core/QueryClient/CosmosQueryClient.cs

@@ -48,6 +48,7 @@ public abstract Task<TryCatch<PartitionedQueryExecutionInfo>> TryGetPartitionedQ
             bool hasLogicalPartitionKey,
             bool allowDCount,
             bool useSystemPrefix,
+            bool isHybridSearchQueryPlanOptimizationDisabled,
             Cosmos.GeospatialType geospatialType,
             CancellationToken cancellationToken);
 

Microsoft.Azure.Cosmos/src/Query/Core/QueryPlan/QueryPartitionProvider.cs

@@ -127,6 +127,7 @@ public TryCatch<PartitionedQueryExecutionInfo> TryGetPartitionedQueryExecutionIn
             bool hasLogicalPartitionKey,
             bool allowDCount,
             bool useSystemPrefix,
+            bool hybridSearchSkipOrderByRewrite,
             GeospatialType geospatialType)
         {
             TryCatch<PartitionedQueryExecutionInfoInternal> tryGetInternalQueryInfo = this.TryGetPartitionedQueryExecutionInfoInternal(
@@ -139,6 +140,7 @@ public TryCatch<PartitionedQueryExecutionInfo> TryGetPartitionedQueryExecutionIn
                 hasLogicalPartitionKey: hasLogicalPartitionKey,
                 allowDCount: allowDCount,
                 useSystemPrefix: useSystemPrefix,
+                hybridSearchSkipOrderByRewrite: hybridSearchSkipOrderByRewrite,
                 geospatialType: geospatialType);
 
             if (!tryGetInternalQueryInfo.Succeeded)
@@ -190,6 +192,7 @@ internal TryCatch<PartitionedQueryExecutionInfoInternal> TryGetPartitionedQueryE
             bool hasLogicalPartitionKey,
             bool allowDCount,
             bool useSystemPrefix,
+            bool hybridSearchSkipOrderByRewrite,
             GeospatialType geospatialType)
         {
             if (querySpecJsonString == null || partitionKeyDefinition == null)
@@ -242,7 +245,8 @@ internal TryCatch<PartitionedQueryExecutionInfoInternal> TryGetPartitionedQueryE
                         bHasLogicalPartitionKey = Convert.ToInt32(hasLogicalPartitionKey),
                         bIsContinuationExpected = Convert.ToInt32(isContinuationExpected),
                         bRequireFormattableOrderByQuery = Convert.ToInt32(requireFormattableOrderByQuery),
-                        bUseSystemPrefix = Convert.ToInt32(useSystemPrefix),
+                        bUseSystemPrefix = Convert.ToInt32(useSystemPrefix),
+                        bHybridSearchSkipOrderByRewrite = Convert.ToInt32(!hybridSearchSkipOrderByRewrite),
                         eGeospatialType = Convert.ToInt32(geospatialType),
                         ePartitionKind = Convert.ToInt32(partitionKind)
                     };

Microsoft.Azure.Cosmos/src/Query/Core/QueryPlan/QueryPlanHandler.cs

@@ -29,6 +29,7 @@ public async Task<TryCatch<PartitionedQueryExecutionInfo>> TryGetQueryPlanAsync(
             VectorEmbeddingPolicy vectorEmbeddingPolicy,
             bool hasLogicalPartitionKey,
             bool useSystemPrefix,
+            bool isHybridSearchQueryPlanOptimizationDisabled,
             GeospatialType geospatialType,
             CancellationToken cancellationToken)
         {
@@ -51,6 +52,7 @@ public async Task<TryCatch<PartitionedQueryExecutionInfo>> TryGetQueryPlanAsync(
                 vectorEmbeddingPolicy,
                 hasLogicalPartitionKey,
                 useSystemPrefix,
+                isHybridSearchQueryPlanOptimizationDisabled,
                 geospatialType,
                 cancellationToken);
             if (!tryGetQueryInfo.Succeeded)
@@ -68,6 +70,7 @@ private Task<TryCatch<PartitionedQueryExecutionInfo>> TryGetQueryInfoAsync(
             VectorEmbeddingPolicy vectorEmbeddingPolicy,
             bool hasLogicalPartitionKey,
             bool useSystemPrefix,
+            bool isHybridSearchQueryPlanOptimizationDisabled,
             Cosmos.GeospatialType geospatialType,
             CancellationToken cancellationToken = default)
         {
@@ -83,7 +86,8 @@ private Task<TryCatch<PartitionedQueryExecutionInfo>> TryGetQueryInfoAsync(
                 allowNonValueAggregateQuery: true,
                 hasLogicalPartitionKey: hasLogicalPartitionKey,
                 allowDCount: true,
-                useSystemPrefix: useSystemPrefix,
+                useSystemPrefix: useSystemPrefix,
+                isHybridSearchQueryPlanOptimizationDisabled: isHybridSearchQueryPlanOptimizationDisabled,
                 geospatialType: geospatialType,
                 cancellationToken: cancellationToken);
         }

Microsoft.Azure.Cosmos/src/Query/Core/QueryPlan/QueryPlanRetriever.cs

@@ -28,36 +28,38 @@ internal static class QueryPlanRetriever
             | QueryFeatures.OrderBy
             | QueryFeatures.Top
             | QueryFeatures.NonValueAggregate
-            | QueryFeatures.DCount
-            | QueryFeatures.NonStreamingOrderBy
-            | QueryFeatures.CountIf
-            | QueryFeatures.HybridSearch
-            | QueryFeatures.WeightedRankFusion;
-
-        private static readonly QueryFeatures SupportedQueryFeaturesWithoutNonStreamingOrderBy =
-            SupportedQueryFeatures & (~QueryFeatures.NonStreamingOrderBy);
-
-        private static readonly string SupportedQueryFeaturesString = SupportedQueryFeatures.ToString();
-
-        private static readonly string SupportedQueryFeaturesWithoutNonStreamingOrderByString =
-            SupportedQueryFeaturesWithoutNonStreamingOrderBy.ToString();
-
-        private static string GetSupportedQueryFeaturesString(bool isNonStreamingOrderByQueryFeatureDisabled)
-        {
-            return isNonStreamingOrderByQueryFeatureDisabled ?
-                SupportedQueryFeaturesWithoutNonStreamingOrderByString :
-                SupportedQueryFeaturesString;
+            | QueryFeatures.DCount
+            | QueryFeatures.NonStreamingOrderBy
+            | QueryFeatures.CountIf
+            | QueryFeatures.HybridSearch
+            | QueryFeatures.WeightedRankFusion
+            | QueryFeatures.HybridSearchSkipOrderByRewrite;
+
+        private static readonly QueryFeatures SupportedQueryFeaturesWithHybridSearchQueryPlanOptimizationDisabled =
+            SupportedQueryFeatures & (~QueryFeatures.HybridSearchSkipOrderByRewrite);
+
+        private static readonly string SupportedQueryFeaturesString = SupportedQueryFeatures.ToString();
+
+        private static readonly string SupportedQueryFeaturesWithHybridSearchQueryPlanOptimizationDisabledString =
+            SupportedQueryFeaturesWithHybridSearchQueryPlanOptimizationDisabled.ToString();
+
+        private static string GetSupportedQueryFeaturesString(bool isHybridSearchQueryPlanOptimizationDisabled)
+        {
+            return isHybridSearchQueryPlanOptimizationDisabled ?
+                SupportedQueryFeaturesWithHybridSearchQueryPlanOptimizationDisabledString :
+                SupportedQueryFeaturesString;
         }
 
         public static async Task<PartitionedQueryExecutionInfo> GetQueryPlanWithServiceInteropAsync(
             CosmosQueryClient queryClient,
             SqlQuerySpec sqlQuerySpec,
             Documents.ResourceType resourceType,
-            PartitionKeyDefinition partitionKeyDefinition,
+            PartitionKeyDefinition partitionKeyDefinition,
             VectorEmbeddingPolicy vectorEmbeddingPolicy,
             bool hasLogicalPartitionKey,
             GeospatialType geospatialType,
             bool useSystemPrefix,
+            bool isHybridSearchQueryPlanOptimizationDisabled,
             ITrace trace,
             CancellationToken cancellationToken = default)
         {
@@ -85,10 +87,11 @@ public static async Task<PartitionedQueryExecutionInfo> GetQueryPlanWithServiceI
                 TryCatch<PartitionedQueryExecutionInfo> tryGetQueryPlan = await queryPlanHandler.TryGetQueryPlanAsync(
                     sqlQuerySpec,
                     resourceType,
-                    partitionKeyDefinition,
+                    partitionKeyDefinition,
                     vectorEmbeddingPolicy,
                     hasLogicalPartitionKey,
                     useSystemPrefix,
+                    isHybridSearchQueryPlanOptimizationDisabled,
                     geospatialType,
                     cancellationToken);
 
@@ -112,8 +115,8 @@ public static Task<PartitionedQueryExecutionInfo> GetQueryPlanThroughGatewayAsyn
             CosmosQueryContext queryContext,
             SqlQuerySpec sqlQuerySpec,
             string resourceLink,
-            PartitionKey? partitionKey,
-            bool isNonStreamingOrderByQueryFeatureDisabled,
+            PartitionKey? partitionKey,
+            bool isHybridSearchQueryPlanOptimizationDisabled,
             ITrace trace,
             CancellationToken cancellationToken = default)
         {
@@ -149,7 +152,7 @@ public static Task<PartitionedQueryExecutionInfo> GetQueryPlanThroughGatewayAsyn
                     OperationType.QueryPlan,
                     sqlQuerySpec,
                     partitionKey,
-                    GetSupportedQueryFeaturesString(isNonStreamingOrderByQueryFeatureDisabled),
+                    GetSupportedQueryFeaturesString(isHybridSearchQueryPlanOptimizationDisabled),
                     trace,
                     cancellationToken);
             }

Microsoft.Azure.Cosmos/src/Query/v2Query/DocumentQueryExecutionContextBase.cs

@@ -186,7 +186,8 @@ public async Task<PartitionedQueryExecutionInfo> GetPartitionedQueryExecutionInf
                 isContinuationExpected: isContinuationExpected,
                 allowNonValueAggregateQuery: allowNonValueAggregateQuery,
                 hasLogicalPartitionKey: hasLogicalPartitionKey,
-                allowDCount: allowDCount,
+                allowDCount: allowDCount,
+                hybridSearchSkipOrderByRewrite: false,
                 geospatialType: geospatialType,
                 useSystemPrefix: false);
             if (!tryGetPartitionedQueryExecutionInfo.Succeeded)

Microsoft.Azure.Cosmos/src/Query/v3Query/CosmosQueryClientCore.cs

@@ -101,6 +101,7 @@ public override async Task<TryCatch<PartitionedQueryExecutionInfo>> TryGetPartit
             bool hasLogicalPartitionKey,
             bool allowDCount,
             bool useSystemPrefix,
+            bool isHybridSearchQueryPlanOptimizationDisabled,
             Cosmos.GeospatialType geospatialType,
             CancellationToken cancellationToken)
         {
@@ -126,6 +127,7 @@ public override async Task<TryCatch<PartitionedQueryExecutionInfo>> TryGetPartit
                 hasLogicalPartitionKey: hasLogicalPartitionKey,
                 allowDCount: allowDCount,
                 useSystemPrefix: useSystemPrefix,
+                hybridSearchSkipOrderByRewrite: !isHybridSearchQueryPlanOptimizationDisabled,
                 geospatialType: geospatialType);
         }
 

Microsoft.Azure.Cosmos/src/Query/v3Query/QueryIterator.cs

@@ -149,7 +149,7 @@ public static QueryIterator Create(
                 partitionedQueryExecutionInfo: partitionedQueryExecutionInfo,
                 returnResultsInDeterministicOrder: queryRequestOptions.ReturnResultsInDeterministicOrder,
                 enableOptimisticDirectExecution: queryRequestOptions.EnableOptimisticDirectExecution,
-                isNonStreamingOrderByQueryFeatureDisabled: queryRequestOptions.IsNonStreamingOrderByQueryFeatureDisabled,
+                isHybridSearchQueryPlanOptimizationDisabled: queryRequestOptions.IsHybridSearchQueryPlanOptimizationDisabled,
                 enableDistributedQueryGatewayMode: queryRequestOptions.EnableDistributedQueryGatewayMode && (clientContext.ClientOptions.ConnectionMode == ConnectionMode.Gateway),
                 testInjections: queryRequestOptions.TestSettings);