Skip to content

Commit 657606b

Browse files
neildshneildsh
committed
Plumb a flag for disabling the hybrid search query plan optimization
1 parent 34b719b commit 657606b

22 files changed

Lines changed: 128 additions & 93 deletions

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

Lines changed: 8 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -584,7 +584,7 @@ private static async Task<PartitionedQueryExecutionInfo> GetPartitionedQueryExec
584584
inputParameters.SqlQuerySpec,
585585
cosmosQueryContext.ResourceLink,
586586
inputParameters.PartitionKey,
587-
inputParameters.IsNonStreamingOrderByQueryFeatureDisabled,
587+
inputParameters.IsHybridSearchQueryPlanOptimizationDisabled,
588588
trace,
589589
cancellationToken);
590590
}
@@ -601,6 +601,7 @@ private static async Task<PartitionedQueryExecutionInfo> GetPartitionedQueryExec
601601
inputParameters.PartitionKey != null,
602602
containerQueryProperties.GeospatialType,
603603
cosmosQueryContext.UseSystemPrefix,
604+
inputParameters.IsHybridSearchQueryPlanOptimizationDisabled,
604605
trace,
605606
cancellationToken);
606607
}
@@ -832,7 +833,7 @@ private InputParameters(
832833
PartitionedQueryExecutionInfo partitionedQueryExecutionInfo,
833834
bool returnResultsInDeterministicOrder,
834835
bool enableOptimisticDirectExecution,
835-
bool isNonStreamingOrderByQueryFeatureDisabled,
836+
bool isHybridSearchQueryPlanOptimizationDisabled,
836837
bool enableDistributedQueryGatewayMode,
837838
TestInjections testInjections)
838839
{
@@ -847,7 +848,7 @@ private InputParameters(
847848
this.PartitionedQueryExecutionInfo = partitionedQueryExecutionInfo;
848849
this.ReturnResultsInDeterministicOrder = returnResultsInDeterministicOrder;
849850
this.EnableOptimisticDirectExecution = enableOptimisticDirectExecution;
850-
this.IsNonStreamingOrderByQueryFeatureDisabled = isNonStreamingOrderByQueryFeatureDisabled;
851+
this.IsHybridSearchQueryPlanOptimizationDisabled = isHybridSearchQueryPlanOptimizationDisabled;
851852
this.EnableDistributedQueryGatewayMode = enableDistributedQueryGatewayMode;
852853
this.TestInjections = testInjections;
853854
}
@@ -864,7 +865,7 @@ public static InputParameters Create(
864865
PartitionedQueryExecutionInfo partitionedQueryExecutionInfo,
865866
bool? returnResultsInDeterministicOrder,
866867
bool enableOptimisticDirectExecution,
867-
bool isNonStreamingOrderByQueryFeatureDisabled,
868+
bool isHybridSearchQueryPlanOptimizationDisabled,
868869
bool enableDistributedQueryGatewayMode,
869870
TestInjections testInjections)
870871
{
@@ -903,7 +904,7 @@ public static InputParameters Create(
903904
partitionedQueryExecutionInfo: partitionedQueryExecutionInfo,
904905
returnResultsInDeterministicOrder: returnResultsInDeterministicOrder.GetValueOrDefault(InputParameters.DefaultReturnResultsInDeterministicOrder),
905906
enableOptimisticDirectExecution: enableOptimisticDirectExecution,
906-
isNonStreamingOrderByQueryFeatureDisabled: isNonStreamingOrderByQueryFeatureDisabled,
907+
isHybridSearchQueryPlanOptimizationDisabled: isHybridSearchQueryPlanOptimizationDisabled,
907908
enableDistributedQueryGatewayMode: enableDistributedQueryGatewayMode,
908909
testInjections: testInjections);
909910
}
@@ -920,7 +921,7 @@ public static InputParameters Create(
920921
public bool ReturnResultsInDeterministicOrder { get; }
921922
public TestInjections TestInjections { get; }
922923
public bool EnableOptimisticDirectExecution { get; }
923-
public bool IsNonStreamingOrderByQueryFeatureDisabled { get; }
924+
public bool IsHybridSearchQueryPlanOptimizationDisabled { get; }
924925
public bool EnableDistributedQueryGatewayMode { get; }
925926

926927
public InputParameters WithContinuationToken(CosmosElement token)
@@ -937,7 +938,7 @@ public InputParameters WithContinuationToken(CosmosElement token)
937938
this.PartitionedQueryExecutionInfo,
938939
this.ReturnResultsInDeterministicOrder,
939940
this.EnableOptimisticDirectExecution,
940-
this.IsNonStreamingOrderByQueryFeatureDisabled,
941+
this.IsHybridSearchQueryPlanOptimizationDisabled,
941942
this.EnableDistributedQueryGatewayMode,
942943
this.TestInjections);
943944
}

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -48,6 +48,7 @@ public abstract Task<TryCatch<PartitionedQueryExecutionInfo>> TryGetPartitionedQ
4848
bool hasLogicalPartitionKey,
4949
bool allowDCount,
5050
bool useSystemPrefix,
51+
bool isHybridSearchQueryPlanOptimizationDisabled,
5152
Cosmos.GeospatialType geospatialType,
5253
CancellationToken cancellationToken);
5354

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

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -127,6 +127,7 @@ public TryCatch<PartitionedQueryExecutionInfo> TryGetPartitionedQueryExecutionIn
127127
bool hasLogicalPartitionKey,
128128
bool allowDCount,
129129
bool useSystemPrefix,
130+
bool hybridSearchSkipOrderByRewrite,
130131
GeospatialType geospatialType)
131132
{
132133
TryCatch<PartitionedQueryExecutionInfoInternal> tryGetInternalQueryInfo = this.TryGetPartitionedQueryExecutionInfoInternal(
@@ -139,6 +140,7 @@ public TryCatch<PartitionedQueryExecutionInfo> TryGetPartitionedQueryExecutionIn
139140
hasLogicalPartitionKey: hasLogicalPartitionKey,
140141
allowDCount: allowDCount,
141142
useSystemPrefix: useSystemPrefix,
143+
hybridSearchSkipOrderByRewrite: hybridSearchSkipOrderByRewrite,
142144
geospatialType: geospatialType);
143145

144146
if (!tryGetInternalQueryInfo.Succeeded)
@@ -190,6 +192,7 @@ internal TryCatch<PartitionedQueryExecutionInfoInternal> TryGetPartitionedQueryE
190192
bool hasLogicalPartitionKey,
191193
bool allowDCount,
192194
bool useSystemPrefix,
195+
bool hybridSearchSkipOrderByRewrite,
193196
GeospatialType geospatialType)
194197
{
195198
if (querySpecJsonString == null || partitionKeyDefinition == null)
@@ -242,7 +245,8 @@ internal TryCatch<PartitionedQueryExecutionInfoInternal> TryGetPartitionedQueryE
242245
bHasLogicalPartitionKey = Convert.ToInt32(hasLogicalPartitionKey),
243246
bIsContinuationExpected = Convert.ToInt32(isContinuationExpected),
244247
bRequireFormattableOrderByQuery = Convert.ToInt32(requireFormattableOrderByQuery),
245-
bUseSystemPrefix = Convert.ToInt32(useSystemPrefix),
248+
bUseSystemPrefix = Convert.ToInt32(useSystemPrefix),
249+
bHybridSearchSkipOrderByRewrite = Convert.ToInt32(!hybridSearchSkipOrderByRewrite),
246250
eGeospatialType = Convert.ToInt32(geospatialType),
247251
ePartitionKind = Convert.ToInt32(partitionKind)
248252
};

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

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@ public async Task<TryCatch<PartitionedQueryExecutionInfo>> TryGetQueryPlanAsync(
2929
VectorEmbeddingPolicy vectorEmbeddingPolicy,
3030
bool hasLogicalPartitionKey,
3131
bool useSystemPrefix,
32+
bool isHybridSearchQueryPlanOptimizationDisabled,
3233
GeospatialType geospatialType,
3334
CancellationToken cancellationToken)
3435
{
@@ -51,6 +52,7 @@ public async Task<TryCatch<PartitionedQueryExecutionInfo>> TryGetQueryPlanAsync(
5152
vectorEmbeddingPolicy,
5253
hasLogicalPartitionKey,
5354
useSystemPrefix,
55+
isHybridSearchQueryPlanOptimizationDisabled,
5456
geospatialType,
5557
cancellationToken);
5658
if (!tryGetQueryInfo.Succeeded)
@@ -68,6 +70,7 @@ private Task<TryCatch<PartitionedQueryExecutionInfo>> TryGetQueryInfoAsync(
6870
VectorEmbeddingPolicy vectorEmbeddingPolicy,
6971
bool hasLogicalPartitionKey,
7072
bool useSystemPrefix,
73+
bool isHybridSearchQueryPlanOptimizationDisabled,
7174
Cosmos.GeospatialType geospatialType,
7275
CancellationToken cancellationToken = default)
7376
{
@@ -83,7 +86,8 @@ private Task<TryCatch<PartitionedQueryExecutionInfo>> TryGetQueryInfoAsync(
8386
allowNonValueAggregateQuery: true,
8487
hasLogicalPartitionKey: hasLogicalPartitionKey,
8588
allowDCount: true,
86-
useSystemPrefix: useSystemPrefix,
89+
useSystemPrefix: useSystemPrefix,
90+
isHybridSearchQueryPlanOptimizationDisabled: isHybridSearchQueryPlanOptimizationDisabled,
8791
geospatialType: geospatialType,
8892
cancellationToken: cancellationToken);
8993
}

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

Lines changed: 27 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -28,36 +28,38 @@ internal static class QueryPlanRetriever
2828
| QueryFeatures.OrderBy
2929
| QueryFeatures.Top
3030
| QueryFeatures.NonValueAggregate
31-
| QueryFeatures.DCount
32-
| QueryFeatures.NonStreamingOrderBy
33-
| QueryFeatures.CountIf
34-
| QueryFeatures.HybridSearch
35-
| QueryFeatures.WeightedRankFusion;
36-
37-
private static readonly QueryFeatures SupportedQueryFeaturesWithoutNonStreamingOrderBy =
38-
SupportedQueryFeatures & (~QueryFeatures.NonStreamingOrderBy);
39-
40-
private static readonly string SupportedQueryFeaturesString = SupportedQueryFeatures.ToString();
41-
42-
private static readonly string SupportedQueryFeaturesWithoutNonStreamingOrderByString =
43-
SupportedQueryFeaturesWithoutNonStreamingOrderBy.ToString();
44-
45-
private static string GetSupportedQueryFeaturesString(bool isNonStreamingOrderByQueryFeatureDisabled)
46-
{
47-
return isNonStreamingOrderByQueryFeatureDisabled ?
48-
SupportedQueryFeaturesWithoutNonStreamingOrderByString :
49-
SupportedQueryFeaturesString;
31+
| QueryFeatures.DCount
32+
| QueryFeatures.NonStreamingOrderBy
33+
| QueryFeatures.CountIf
34+
| QueryFeatures.HybridSearch
35+
| QueryFeatures.WeightedRankFusion
36+
| QueryFeatures.HybridSearchSkipOrderByRewrite;
37+
38+
private static readonly QueryFeatures SupportedQueryFeaturesWithHybridSearchQueryPlanOptimizationDisabled =
39+
SupportedQueryFeatures & (~QueryFeatures.HybridSearchSkipOrderByRewrite);
40+
41+
private static readonly string SupportedQueryFeaturesString = SupportedQueryFeatures.ToString();
42+
43+
private static readonly string SupportedQueryFeaturesWithHybridSearchQueryPlanOptimizationDisabledString =
44+
SupportedQueryFeaturesWithHybridSearchQueryPlanOptimizationDisabled.ToString();
45+
46+
private static string GetSupportedQueryFeaturesString(bool isHybridSearchQueryPlanOptimizationDisabled)
47+
{
48+
return isHybridSearchQueryPlanOptimizationDisabled ?
49+
SupportedQueryFeaturesWithHybridSearchQueryPlanOptimizationDisabledString :
50+
SupportedQueryFeaturesString;
5051
}
5152

5253
public static async Task<PartitionedQueryExecutionInfo> GetQueryPlanWithServiceInteropAsync(
5354
CosmosQueryClient queryClient,
5455
SqlQuerySpec sqlQuerySpec,
5556
Documents.ResourceType resourceType,
56-
PartitionKeyDefinition partitionKeyDefinition,
57+
PartitionKeyDefinition partitionKeyDefinition,
5758
VectorEmbeddingPolicy vectorEmbeddingPolicy,
5859
bool hasLogicalPartitionKey,
5960
GeospatialType geospatialType,
6061
bool useSystemPrefix,
62+
bool isHybridSearchQueryPlanOptimizationDisabled,
6163
ITrace trace,
6264
CancellationToken cancellationToken = default)
6365
{
@@ -85,10 +87,11 @@ public static async Task<PartitionedQueryExecutionInfo> GetQueryPlanWithServiceI
8587
TryCatch<PartitionedQueryExecutionInfo> tryGetQueryPlan = await queryPlanHandler.TryGetQueryPlanAsync(
8688
sqlQuerySpec,
8789
resourceType,
88-
partitionKeyDefinition,
90+
partitionKeyDefinition,
8991
vectorEmbeddingPolicy,
9092
hasLogicalPartitionKey,
9193
useSystemPrefix,
94+
isHybridSearchQueryPlanOptimizationDisabled,
9295
geospatialType,
9396
cancellationToken);
9497

@@ -112,8 +115,8 @@ public static Task<PartitionedQueryExecutionInfo> GetQueryPlanThroughGatewayAsyn
112115
CosmosQueryContext queryContext,
113116
SqlQuerySpec sqlQuerySpec,
114117
string resourceLink,
115-
PartitionKey? partitionKey,
116-
bool isNonStreamingOrderByQueryFeatureDisabled,
118+
PartitionKey? partitionKey,
119+
bool isHybridSearchQueryPlanOptimizationDisabled,
117120
ITrace trace,
118121
CancellationToken cancellationToken = default)
119122
{
@@ -149,7 +152,7 @@ public static Task<PartitionedQueryExecutionInfo> GetQueryPlanThroughGatewayAsyn
149152
OperationType.QueryPlan,
150153
sqlQuerySpec,
151154
partitionKey,
152-
GetSupportedQueryFeaturesString(isNonStreamingOrderByQueryFeatureDisabled),
155+
GetSupportedQueryFeaturesString(isHybridSearchQueryPlanOptimizationDisabled),
153156
trace,
154157
cancellationToken);
155158
}

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

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -186,7 +186,8 @@ public async Task<PartitionedQueryExecutionInfo> GetPartitionedQueryExecutionInf
186186
isContinuationExpected: isContinuationExpected,
187187
allowNonValueAggregateQuery: allowNonValueAggregateQuery,
188188
hasLogicalPartitionKey: hasLogicalPartitionKey,
189-
allowDCount: allowDCount,
189+
allowDCount: allowDCount,
190+
hybridSearchSkipOrderByRewrite: false,
190191
geospatialType: geospatialType,
191192
useSystemPrefix: false);
192193
if (!tryGetPartitionedQueryExecutionInfo.Succeeded)

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

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -101,6 +101,7 @@ public override async Task<TryCatch<PartitionedQueryExecutionInfo>> TryGetPartit
101101
bool hasLogicalPartitionKey,
102102
bool allowDCount,
103103
bool useSystemPrefix,
104+
bool isHybridSearchQueryPlanOptimizationDisabled,
104105
Cosmos.GeospatialType geospatialType,
105106
CancellationToken cancellationToken)
106107
{
@@ -126,6 +127,7 @@ public override async Task<TryCatch<PartitionedQueryExecutionInfo>> TryGetPartit
126127
hasLogicalPartitionKey: hasLogicalPartitionKey,
127128
allowDCount: allowDCount,
128129
useSystemPrefix: useSystemPrefix,
130+
hybridSearchSkipOrderByRewrite: !isHybridSearchQueryPlanOptimizationDisabled,
129131
geospatialType: geospatialType);
130132
}
131133

@@ -288,7 +290,7 @@ public override async Task<List<PartitionKeyRange>> GetTargetPartitionKeyRangesA
288290

289291
public override bool BypassQueryParsing()
290292
{
291-
return CustomTypeExtensions.ByPassQueryParsing();
293+
return true; // CustomTypeExtensions.ByPassQueryParsing();
292294
}
293295

294296
public override void ClearSessionTokenCache(string collectionFullName)

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -149,7 +149,7 @@ public static QueryIterator Create(
149149
partitionedQueryExecutionInfo: partitionedQueryExecutionInfo,
150150
returnResultsInDeterministicOrder: queryRequestOptions.ReturnResultsInDeterministicOrder,
151151
enableOptimisticDirectExecution: queryRequestOptions.EnableOptimisticDirectExecution,
152-
isNonStreamingOrderByQueryFeatureDisabled: queryRequestOptions.IsNonStreamingOrderByQueryFeatureDisabled,
152+
isHybridSearchQueryPlanOptimizationDisabled: queryRequestOptions.IsHybridSearchQueryPlanOptimizationDisabled,
153153
enableDistributedQueryGatewayMode: queryRequestOptions.EnableDistributedQueryGatewayMode && (clientContext.ClientOptions.ConnectionMode == ConnectionMode.Gateway),
154154
testInjections: queryRequestOptions.TestSettings);
155155

0 commit comments

Comments
 (0)