Query : Fixes ODE and ORDER BY issue by cherry picking to hotfix bran…

…ch (#4498) * Fix bug in OrderByCrossPartitionQueryPipelineStage to ensure that failures in creating the inner pipeline stage are bubbled up to the higher stages (#4419) * Query: Fixes Persisted continuationToken issue by turning off ODE by default (#4413) * Turn of ODE by default. * Updated tests * Update * Test fix * Added test for query resumption with ODE token with partition split container (failure case). * Update. * Update Benchmark * Addressed comments. --------- Co-authored-by: neildsh <[email protected]>
Azure · May 16, 2024 · db45a83 · db45a83
1 parent 8a59688
commit db45a83
Show file tree

Hide file tree

Showing 15 changed files with 3,737 additions and 3,696 deletions.
diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs
@@ -758,19 +758,14 @@ private static Documents.PartitionKeyDefinition GetPartitionKeyDefinition(InputP
             ITrace trace)
         {
             bool clientDisableOptimisticDirectExecution = await cosmosQueryContext.QueryClient.GetClientDisableOptimisticDirectExecutionAsync();
+            bool isOdeContinuationToken = inputParameters.InitialUserContinuationToken != null &&
+                OptimisticDirectExecutionContinuationToken.IsOptimisticDirectExecutionContinuationToken(inputParameters.InitialUserContinuationToken);
 
             // Use the Ode code path only if ClientDisableOptimisticDirectExecution is false and EnableOptimisticDirectExecution is true
-            if (clientDisableOptimisticDirectExecution || !inputParameters.EnableOptimisticDirectExecution)
+            // But allow the query using ODE pipeline if it's earlier roundtrips are made using ODE continuation token.
+            if ((clientDisableOptimisticDirectExecution || !inputParameters.EnableOptimisticDirectExecution) &&
+                !isOdeContinuationToken)
             {
-                if (inputParameters.InitialUserContinuationToken != null
-                          && OptimisticDirectExecutionContinuationToken.IsOptimisticDirectExecutionContinuationToken(inputParameters.InitialUserContinuationToken))
-                {
-                    string errorMessage = "Execution of this query using the supplied continuation token requires EnableOptimisticDirectExecution to be set in QueryRequestOptions. " +
-                        "If the error persists after that, contact system administrator.";
-
-                    throw new MalformedContinuationTokenException($"{errorMessage} Continuation Token: {inputParameters.InitialUserContinuationToken}");
-                }
-
                 return null;
             }
 
@@ -817,6 +812,11 @@ private static Documents.PartitionKeyDefinition GetPartitionKeyDefinition(InputP
                 return targetRanges.Single();
             }
 
+            if (isOdeContinuationToken)
+            {
+                throw new InvalidOperationException("Execution of this query cannot resume using Optimistic Direct Execution continuation token due to partition split. Please restart the query without the continuation token.");
+            }
+
             return null;
         }
 

diff --git a/...src/Query/Core/Pipeline/CrossPartition/OrderBy/OrderByCrossPartitionQueryPipelineStage.cs b/...src/Query/Core/Pipeline/CrossPartition/OrderBy/OrderByCrossPartitionQueryPipelineStage.cs
@@ -160,7 +160,7 @@ public async ValueTask<bool> MoveNextAsync(ITrace trace, CancellationToken cance
             }
 
             TryCatch<bool> hasNext = await this.inner.TryAsync(pipelineStage => pipelineStage.MoveNextAsync(trace, cancellationToken));
-            return hasNext.Succeeded && hasNext.Result;
+            return hasNext.Failed || hasNext.Result;
         }
 
         public ValueTask DisposeAsync()

diff --git a/Microsoft.Azure.Cosmos/src/RequestOptions/QueryRequestOptions.cs b/Microsoft.Azure.Cosmos/src/RequestOptions/QueryRequestOptions.cs
@@ -50,7 +50,7 @@ public class QueryRequestOptions : RequestOptions
         /// <value>
         /// Direct (optimistic) execution offers improved performance for several kinds of queries such as a single partition streaming query.
         /// </value>
-        public bool EnableOptimisticDirectExecution { get; set; } = ConfigurationManager.IsOptimisticDirectExecutionEnabled(defaultValue: true);
+        public bool EnableOptimisticDirectExecution { get; set; } = ConfigurationManager.IsOptimisticDirectExecutionEnabled(defaultValue: false);
 
         /// <summary>
         /// Gets or sets the maximum number of items that can be buffered client side during 

diff --git a/...s.EmulatorTests/BaselineTest/TestBaseline/EndToEndTraceWriterBaselineTests.QueryAsync.xml b/...s.EmulatorTests/BaselineTest/TestBaseline/EndToEndTraceWriterBaselineTests.QueryAsync.xml
diff --git a/...mulatorTests/BaselineTest/TestBaseline/EndToEndTraceWriterBaselineTests.ReadManyAsync.xml b/...mulatorTests/BaselineTest/TestBaseline/EndToEndTraceWriterBaselineTests.ReadManyAsync.xml
diff --git a/...latorTests/BaselineTest/TestBaseline/LinqGeneralBaselineTests.TestDistinctTranslation.xml b/...latorTests/BaselineTest/TestBaseline/LinqGeneralBaselineTests.TestDistinctTranslation.xml
@@ -404,7 +404,7 @@ JOIN (
     WHERE (LENGTH(v2["FamilyName"]) > 10) 
     ORDER BY v2 ASC
 ]]></SqlQuery>
-      <ErrorMessage><![CDATA[Status Code: BadRequest","message":"Message: {"Errors":["Order-by over correlated collections is not supported."]}rn]]></ErrorMessage>
+      <ErrorMessage><![CDATA[Status Code: BadRequest]]></ErrorMessage>
     </Output>
   </Result>
   <Result>
@@ -423,7 +423,7 @@ JOIN (
     WHERE ((LENGTH(v2["FamilyName"]) > 10) AND (LENGTH(v2["FamilyName"]) < 20)) 
     ORDER BY v2 ASC
 ]]></SqlQuery>
-      <ErrorMessage><![CDATA[Status Code: BadRequest","message":"Message: {"Errors":["Order-by over correlated collections is not supported."]}rn]]></ErrorMessage>
+      <ErrorMessage><![CDATA[Status Code: BadRequest]]></ErrorMessage>
     </Output>
   </Result>
   <Result>
@@ -442,7 +442,7 @@ JOIN (
     WHERE ((LENGTH(v2["FamilyName"]) > 10) AND (LENGTH(v2["FamilyName"]) < 20)) 
     ORDER BY v2 ASC
 ]]></SqlQuery>
-      <ErrorMessage><![CDATA[Status Code: BadRequest","message":"Message: {"Errors":["Order-by over correlated collections is not supported."]}rn]]></ErrorMessage>
+      <ErrorMessage><![CDATA[Status Code: BadRequest]]></ErrorMessage>
     </Output>
   </Result>
   <Result>
@@ -482,7 +482,7 @@ JOIN (
     WHERE (LENGTH(v2["FamilyName"]) > 10) 
     ORDER BY v2 ASC
 ]]></SqlQuery>
-      <ErrorMessage><![CDATA[Status Code: BadRequest","message":"Message: {"Errors":["Order-by over correlated collections is not supported."]}rn]]></ErrorMessage>
+      <ErrorMessage><![CDATA[Status Code: BadRequest]]></ErrorMessage>
     </Output>
   </Result>
   <Result>
@@ -559,7 +559,7 @@ JOIN (
     WHERE (LENGTH(v2["FamilyName"]) > 10) 
     ORDER BY v2 ASC
 ]]></SqlQuery>
-      <ErrorMessage><![CDATA[Status Code: BadRequest","message":"Message: {"Errors":["Order-by over correlated collections is not supported."]}rn]]></ErrorMessage>
+      <ErrorMessage><![CDATA[Status Code: BadRequest]]></ErrorMessage>
     </Output>
   </Result>
   <Result>
@@ -578,7 +578,7 @@ JOIN (
     WHERE (LENGTH(v2["FamilyName"]) > 10) 
     ORDER BY v2 ASC
 ]]></SqlQuery>
-      <ErrorMessage><![CDATA[Status Code: BadRequest","message":"Message: {"Errors":["Order-by over correlated collections is not supported."]}rn]]></ErrorMessage>
+      <ErrorMessage><![CDATA[Status Code: BadRequest]]></ErrorMessage>
     </Output>
   </Result>
   <Result>
@@ -597,7 +597,7 @@ JOIN (
     WHERE (LENGTH(v2["FamilyName"]) > 10) 
     ORDER BY v2["GivenName"]["Length"] ASC
 ]]></SqlQuery>
-      <ErrorMessage><![CDATA[Status Code: BadRequest","message":"Message: {"Errors":["Order-by over correlated collections is not supported."]}rn]]></ErrorMessage>
+      <ErrorMessage><![CDATA[Status Code: BadRequest]]></ErrorMessage>
     </Output>
   </Result>
   <Result>
@@ -616,7 +616,7 @@ JOIN (
     WHERE (LENGTH(v2["FamilyName"]) > 10) 
     ORDER BY v2["GivenName"]["Length"] ASC
 ]]></SqlQuery>
-      <ErrorMessage><![CDATA[Status Code: BadRequest","message":"Message: {"Errors":["Order-by over correlated collections is not supported."]}rn]]></ErrorMessage>
+      <ErrorMessage><![CDATA[Status Code: BadRequest]]></ErrorMessage>
     </Output>
   </Result>
   <Result>
@@ -635,7 +635,7 @@ JOIN (
     WHERE (LENGTH(v2["FamilyName"]) > 10) 
     ORDER BY v2 ASC
 ]]></SqlQuery>
-      <ErrorMessage><![CDATA[Status Code: BadRequest","message":"Message: {"Errors":["Order-by over correlated collections is not supported."]}rn]]></ErrorMessage>
+      <ErrorMessage><![CDATA[Status Code: BadRequest]]></ErrorMessage>
     </Output>
   </Result>
   <Result>
@@ -654,7 +654,7 @@ JOIN (
     WHERE (LENGTH(v2["FamilyName"]) > 10) 
     ORDER BY v2 ASC
 ]]></SqlQuery>
-      <ErrorMessage><![CDATA[Status Code: BadRequest","message":"Message: {"Errors":["Order-by over correlated collections is not supported."]}rn]]></ErrorMessage>
+      <ErrorMessage><![CDATA[Status Code: BadRequest]]></ErrorMessage>
     </Output>
   </Result>
   <Result>
@@ -689,7 +689,7 @@ JOIN (
     JOIN v0 IN root["Parents"]) AS v2 
     ORDER BY v2["FamilyName"] ASC
 ]]></SqlQuery>
-      <ErrorMessage><![CDATA[Status Code: BadRequest","message":"Message: {"Errors":["Order-by over correlated collections is not supported."]}rn]]></ErrorMessage>
+      <ErrorMessage><![CDATA[Status Code: BadRequest]]></ErrorMessage>
     </Output>
   </Result>
   <Result>
@@ -707,7 +707,7 @@ JOIN (
     JOIN v0 IN root["Parents"]) AS v2 
     ORDER BY v2["FamilyName"] ASC
 ]]></SqlQuery>
-      <ErrorMessage><![CDATA[Status Code: BadRequest","message":"Message: {"Errors":["Order-by over correlated collections is not supported."]}rn]]></ErrorMessage>
+      <ErrorMessage><![CDATA[Status Code: BadRequest]]></ErrorMessage>
     </Output>
   </Result>
   <Result>
@@ -773,7 +773,7 @@ JOIN (
     JOIN p0 IN root["Parents"]) AS v1 
     ORDER BY v1 ASC
 ]]></SqlQuery>
-      <ErrorMessage><![CDATA[Status Code: BadRequest","message":"Message: {"Errors":["Order-by over correlated collections is not supported."]}rn]]></ErrorMessage>
+      <ErrorMessage><![CDATA[Status Code: BadRequest]]></ErrorMessage>
     </Output>
   </Result>
   <Result>
@@ -792,7 +792,7 @@ JOIN (
     WHERE (LENGTH(v1) > 10) 
     ORDER BY v1 ASC
 ]]></SqlQuery>
-      <ErrorMessage><![CDATA[Status Code: BadRequest","message":"Message: {"Errors":["Order-by over correlated collections is not supported."]}rn]]></ErrorMessage>
+      <ErrorMessage><![CDATA[Status Code: BadRequest]]></ErrorMessage>
     </Output>
   </Result>
   <Result>

diff --git a/...ulatorTests/BaselineTest/TestBaseline/LinqGeneralBaselineTests.TestGroupByTranslation.xml b/...ulatorTests/BaselineTest/TestBaseline/LinqGeneralBaselineTests.TestGroupByTranslation.xml
@@ -320,6 +320,7 @@ SELECT VALUE root["id"]
 FROM root 
 GROUP BY root["id"]  
 ORDER BY root["id"] DESC]]></SqlQuery>
+      <ErrorMessage><![CDATA[Status Code: BadRequest]]></ErrorMessage>
     </Output>
   </Result>
   <Result>
@@ -334,6 +335,7 @@ FROM root
 WHERE (root["id"] != "a") 
 GROUP BY root["id"]  
 ORDER BY root["id"] ASC]]></SqlQuery>
+      <ErrorMessage><![CDATA[Status Code: BadRequest]]></ErrorMessage>
     </Output>
   </Result>
   <Result>

diff --git a/...mulatorTests/BaselineTest/TestBaseline/LinqGeneralBaselineTests.TestThenByTranslation.xml b/...mulatorTests/BaselineTest/TestBaseline/LinqGeneralBaselineTests.TestThenByTranslation.xml
@@ -119,7 +119,7 @@ JOIN (
 SELECT VALUE root 
 FROM root 
 ORDER BY root["FamilyId"] ASC, root["FamilyId"] ASC]]></SqlQuery>
-      <ErrorMessage><![CDATA[Status Code: BadRequest","message":"Message: {"Errors":["The order by query does not have a corresponding composite index that it can be served from."]}rn]]></ErrorMessage>
+      <ErrorMessage><![CDATA[Status Code: BadRequest]]></ErrorMessage>
     </Output>
   </Result>
   <Result>
@@ -132,7 +132,7 @@ ORDER BY root["FamilyId"] ASC, root["FamilyId"] ASC]]></SqlQuery>
 SELECT VALUE root 
 FROM root 
 ORDER BY root["FamilyId"] ASC, root["FamilyId"] DESC]]></SqlQuery>
-      <ErrorMessage><![CDATA[Status Code: BadRequest","message":"Message: {"Errors":["The order by query does not have a corresponding composite index that it can be served from."]}rn]]></ErrorMessage>
+      <ErrorMessage><![CDATA[Status Code: BadRequest]]></ErrorMessage>
     </Output>
   </Result>
   <Result>
@@ -388,6 +388,7 @@ FROM (
         JOIN c0 IN root["Children"]) AS v1) AS r1 
         ORDER BY r1["FamilyId"] ASC, r1["FamilyNumber"] ASC
 ]]></SqlQuery>
+      <ErrorMessage><![CDATA[Underlying object does not have an '_rid' or '__sys_rid' field.]]></ErrorMessage>
     </Output>
   </Result>
   <Result>
@@ -434,6 +435,7 @@ FROM (
         WHERE (p0["GivenName"]["Length"] > 5)) AS v0) AS r0 
         ORDER BY r0["FamilyId"] ASC, r0["FamilyNumber"] ASC
 ]]></SqlQuery>
+      <ErrorMessage><![CDATA[Underlying object does not have an '_rid' or '__sys_rid' field.]]></ErrorMessage>
     </Output>
   </Result>
   <Result>
@@ -584,7 +586,7 @@ FROM root
 JOIN f0 IN root["Records"]["Transactions"] 
 WHERE (ARRAY_LENGTH(root["Children"]) > 0) 
 ORDER BY f0["Type"] ASC, f0["Amount"] ASC]]></SqlQuery>
-      <ErrorMessage><![CDATA[Status Code: BadRequest","message":"Message: {"Errors":["Order-by over correlated collections is not supported."]}rn]]></ErrorMessage>
+      <ErrorMessage><![CDATA[Status Code: BadRequest]]></ErrorMessage>
     </Output>
   </Result>
   <Result>