From f01caeeb339690b63abb0d9ce3a799fbe2998dba Mon Sep 17 00:00:00 2001
From: neildsh <35383880+neildsh@users.noreply.github.com>
Date: Mon, 31 Oct 2022 14:12:32 -0700
Subject: [PATCH] Query: Fixes performance regression on target partition on
 some ORDER BY queries with continuation (#3525)

* Revert performance regression caused by https://github.com/Azure/azure-cosmos-dotnet-v3/pull/1289/

* Remove irrelevant comment

* Add a test for validating formatted filters for the target partition
---
 ...OrderByCrossPartitionQueryPipelineStage.cs |  4 +-
 ...ByCrossPartitionQueryPipelineStageTests.cs | 57 +++++++++++++++++++
 2 files changed, 58 insertions(+), 3 deletions(-)

diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/OrderBy/OrderByCrossPartitionQueryPipelineStage.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/OrderBy/OrderByCrossPartitionQueryPipelineStage.cs
index a660f29d0b..bbd3d7be4e 100644
--- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/OrderBy/OrderByCrossPartitionQueryPipelineStage.cs
+++ b/Microsoft.Azure.Cosmos/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/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/OrderByCrossPartitionQueryPipelineStageTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/OrderByCrossPartitionQueryPipelineStageTests.cs
index 448f63f02c..52062c8afb 100644
--- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/OrderByCrossPartitionQueryPipelineStageTests.cs
+++ b/Microsoft.Azure.Cosmos/tests/Microsoft.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()
         {