Query: Fixes performance regression on target partition on some ORDER…

… BY queries with continuation (#3525) * Revert performance regression caused by #1289 * Remove irrelevant comment * Add a test for validating formatted filters for the target partition
Azure · Nov 3, 2022 · f01caee · f01caee
1 parent 63fe067
commit f01caee
Show file tree

Hide file tree

Showing 2 changed files with 58 additions and 3 deletions.
diff --git a/...src/Query/Core/Pipeline/CrossPartition/OrderBy/OrderByCrossPartitionQueryPipelineStage.cs b/...src/Query/Core/Pipeline/CrossPartition/OrderBy/OrderByCrossPartitionQueryPipelineStage.cs
@@ -1027,9 +1027,7 @@ private static (string leftFilter, string targetFilter, string rightFilter) GetF
                 }
             }
 
-            // For the target filter we can make an optimization to just return "true",
-            // since we already have the backend continuation token to resume with.
-            return (left.ToString(), TrueFilter, right.ToString());
+            return (left.ToString(), target.ToString(), right.ToString());
         }
 
         private static async Task<TryCatch<(bool doneFiltering, int itemsLeftToSkip, TryCatch<OrderByQueryPage> monadicQueryByPage)>> FilterNextAsync(

diff --git a/...crosoft.Azure.Cosmos.Tests/Query/Pipeline/OrderByCrossPartitionQueryPipelineStageTests.cs b/...crosoft.Azure.Cosmos.Tests/Query/Pipeline/OrderByCrossPartitionQueryPipelineStageTests.cs
@@ -278,6 +278,63 @@ public void MonadicCreate_MultipleOrderByContinuationToken()
             }
         }
 
+        [TestMethod]
+        public async Task TestFormattedFiltersForTargetPartitionWithContinuationTokenAsync()
+        {
+            QueryPage emptyPage = new QueryPage(
+                documents: new List<CosmosElement>(),
+                requestCharge: 0,
+                activityId: string.Empty,
+                responseLengthInBytes: 0,
+                cosmosQueryExecutionInfo: default,
+                disallowContinuationTokenMessage: default,
+                additionalHeaders: default,
+                state: default);
+
+            string expectedQuerySpec = "SELECT * FROM c WHERE ( c._ts >= 1665482200 OR IS_STRING(c._ts) OR IS_ARRAY(c._ts) OR IS_OBJECT(c._ts) ) ORDER BY c._ts";
+            Mock<IDocumentContainer> mockContainer = new Mock<IDocumentContainer>(MockBehavior.Strict);
+            mockContainer
+                .Setup(
+                c => c.MonadicQueryAsync(
+                    It.Is<SqlQuerySpec>(sqlQuerySpec => expectedQuerySpec.Equals(sqlQuerySpec.QueryText)),
+                    It.IsAny<FeedRangeState<QueryState>>(),
+                    It.IsAny<QueryPaginationOptions>(),
+                    NoOpTrace.Singleton,
+                    default))
+                .ReturnsAsync(TryCatch<QueryPage>.FromResult(emptyPage));
+
+            string continuationToken = @"[{""compositeToken"":{""token"":null,""range"":{""min"":""A"",""max"":""B""}},""orderByItems"":[{""item"":1665482200}],""rid"":""64kUAPYyHHk6XgIAAADACQ=="",""skipCount"":1,""filter"":""( c._ts >= 1665482198 OR IS_STRING(c._ts) OR IS_ARRAY(c._ts) OR IS_OBJECT(c._ts) )""}]";
+
+            IReadOnlyList<FeedRangeEpk> targetRanges = new List<FeedRangeEpk>()
+            {
+                new FeedRangeEpk(new Range<string>(min: "A", max: "B", isMinInclusive: true, isMaxInclusive: false)),
+                new FeedRangeEpk(new Range<string>(min: "B", max: "C", isMinInclusive: true, isMaxInclusive: false))
+            };
+
+            TryCatch<IQueryPipelineStage> monadicCreate = OrderByCrossPartitionQueryPipelineStage.MonadicCreate(
+                documentContainer: mockContainer.Object,
+                sqlQuerySpec: new SqlQuerySpec("SELECT * FROM c WHERE {documentdb-formattableorderbyquery-filter} ORDER BY c._ts"),
+                targetRanges: targetRanges,
+                partitionKey: null,
+                orderByColumns: new List<OrderByColumn>()
+                {
+                    new OrderByColumn("c._ts", SortOrder.Ascending)
+                },
+                queryPaginationOptions: new QueryPaginationOptions(pageSizeHint: 1),
+                maxConcurrency: 0,
+                cancellationToken: default,
+                continuationToken: CosmosElement.Parse(continuationToken));
+            Assert.IsTrue(monadicCreate.Succeeded);
+
+            IQueryPipelineStage queryPipelineStage = monadicCreate.Result;
+            for (int i = 0; i < targetRanges.Count; ++i)
+            {
+                Assert.IsTrue(await queryPipelineStage.MoveNextAsync(NoOpTrace.Singleton));
+            }
+
+            Assert.IsFalse(await queryPipelineStage.MoveNextAsync(NoOpTrace.Singleton));
+        }
+
         [TestMethod]
         public async Task TestDrainFully_StartFromBeginingAsync_NoDocuments()
         {