[Internal] Query: Adds single physical partition check for Optimistic…

…DirectExecution queries (#3699) * Added single physical partition check for Ode queries. Updated test infrastructure for Ode emulator tests too. * Refactored emulator tests to have all the test cases at the top of the file * Updated TryGetTargetRangeOptimisticDirectExecutionAsync() * Uodated logic on how many times CreateIngestQueryDeleteAsync() gets called * Added debug asserts for partitionKeyDefinition * Added pageSizeOptions parameter in CreateInput() in EmulatorTests * Fixed comments
Azure · Feb 17, 2023 · 67e1a90 · 67e1a90
1 parent 209e473
commit 67e1a90
Show file tree

Hide file tree

Showing 5 changed files with 218 additions and 253 deletions.
diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs
@@ -755,22 +755,9 @@ private static Documents.PartitionKeyDefinition GetPartitionKeyDefinition(InputP
         {
             if (!inputParameters.EnableOptimisticDirectExecution) return null;
 
-            // case 1: Is query going to a single partition
-            bool hasPartitionKey = inputParameters.PartitionKey.HasValue
-                && inputParameters.PartitionKey != PartitionKey.Null
-                && inputParameters.PartitionKey != PartitionKey.None;
-
-            // case 2: does query execution plan have a single query range
-            bool hasQueryRanges = partitionedQueryExecutionInfo != null 
-                && partitionedQueryExecutionInfo.QueryRanges.Count == 1 
-                && partitionedQueryExecutionInfo.QueryRanges[0].IsSingleValue;
-
-            if (!hasPartitionKey && !hasQueryRanges) return null;
-
-            //TODO: does collection have only one physical partition
-
-            List<Documents.PartitionKeyRange> targetRanges = new List<Documents.PartitionKeyRange>();
+            Debug.Assert(containerQueryProperties.ResourceId != null, "CosmosQueryExecutionContextFactory Assert!", "Container ResourceId cannot be null!");
 
+            List<Documents.PartitionKeyRange> targetRanges;
             if (partitionedQueryExecutionInfo != null)
             {
                 targetRanges = await CosmosQueryExecutionContextFactory.GetTargetPartitionKeyRangesAsync(
@@ -785,15 +772,26 @@ private static Documents.PartitionKeyDefinition GetPartitionKeyDefinition(InputP
             else
             {
                 Documents.PartitionKeyDefinition partitionKeyDefinition = GetPartitionKeyDefinition(inputParameters, containerQueryProperties);
-                if (partitionKeyDefinition != null && containerQueryProperties.ResourceId != null && inputParameters.PartitionKey != null)
+                if (inputParameters.PartitionKey != null)
                 {
+                    Debug.Assert(partitionKeyDefinition != null, "CosmosQueryExecutionContextFactory Assert!", "PartitionKeyDefinition cannot be null if partitionKey is defined");
+
                     targetRanges = await cosmosQueryContext.QueryClient.GetTargetPartitionKeyRangesByEpkStringAsync(
                         cosmosQueryContext.ResourceLink,
                         containerQueryProperties.ResourceId,
                         inputParameters.PartitionKey.Value.InternalKey.GetEffectivePartitionKeyString(partitionKeyDefinition),
                         forceRefresh: false,
                         trace);
                 }
+                else
+                {
+                    targetRanges = await cosmosQueryContext.QueryClient.GetTargetPartitionKeyRangesAsync(
+                        cosmosQueryContext.ResourceLink,
+                        containerQueryProperties.ResourceId,
+                        new List<Documents.Routing.Range<string>> { FeedRangeEpk.FullRange.Range },
+                        forceRefresh: false,
+                        trace);
+                }
             }
 
             if (targetRanges.Count == 1)