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>