Query: Fixes to not use passthrough context for collections with HPK

Microsoft.Azure.Cosmos/src/Handler/RequestMessage.cs

@@ -385,9 +385,13 @@ private bool AssertPartitioningPropertiesAndHeaders()
 
             bool partitionKeyRangeIdExists = !string.IsNullOrEmpty(this.Headers.PartitionKeyRangeId);
             if (partitionKeyRangeIdExists)
-            {
+            {
+                OperationType operationType = this.OperationType;
                 // Assert operation type is not write
-                if (this.OperationType != OperationType.Query && this.OperationType != OperationType.ReadFeed && this.OperationType != OperationType.Batch)
+                if (operationType != OperationType.Query
+                    && operationType != OperationType.QueryPlan
+                    && operationType != OperationType.ReadFeed 
+                    && operationType != OperationType.Batch)
                 {
                     throw new ArgumentOutOfRangeException(RMResources.UnexpectedPartitionKeyRangeId);
                 }

Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs

@@ -191,7 +191,7 @@ private static async Task<TryCatch<IQueryPipelineStage>> TryCreateCoreContextAsy
                     // then try seeing if we can execute as a passthrough using client side only logic.
                     // This is to short circuit the need to go to the gateway to get the query plan.
                     if (cosmosQueryContext.QueryClient.BypassQueryParsing()
-                        && inputParameters.PartitionKey.HasValue)
+                        && inputParameters.PartitionKey.HasValue && containerQueryProperties.PartitionKeyDefinition.Paths.Count <= 1)
                     {
                         bool parsed;
                         SqlQuery sqlQuery;
@@ -301,9 +301,10 @@ private static async Task<TryCatch<IQueryPipelineStage>> TryCreateFromPartitione
                 }
                 else
                 {
-                    bool singleLogicalPartitionKeyQuery = (inputParameters.PartitionKey.HasValue && targetRanges.Count == 1)
+                    bool singleLogicalPartitionKeyQuery = ((inputParameters.PartitionKey.HasValue && targetRanges.Count == 1)
                         || ((partitionedQueryExecutionInfo.QueryRanges.Count == 1)
-                        && partitionedQueryExecutionInfo.QueryRanges[0].IsSingleValue);
+                        && partitionedQueryExecutionInfo.QueryRanges[0].IsSingleValue))
+                        && containerQueryProperties.PartitionKeyDefinition.Paths.Count <= 1;
                     bool serverStreamingQuery = !partitionedQueryExecutionInfo.QueryInfo.HasAggregates
                         && !partitionedQueryExecutionInfo.QueryInfo.HasDistinct
                         && !partitionedQueryExecutionInfo.QueryInfo.HasNonStreamingOrderBy

Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/Query/BypassQueryParsingTests.cs

@@ -2,44 +2,116 @@ namespace Microsoft.Azure.Cosmos.EmulatorTests.Query
 {
     using System;
     using System.Collections.Generic;
+    using System.Collections.ObjectModel;
     using System.Linq;
     using System.Threading.Tasks;
     using Microsoft.Azure.Cosmos.CosmosElements;
+    using Microsoft.Azure.Cosmos.Query.Core;
     using Microsoft.VisualStudio.TestTools.UnitTesting;
 
     [TestClass]
     [TestCategory("Query")]
     public sealed class BypassQueryParsingTests : QueryTestsBase
     {
-        [TestMethod]
-        public async Task TestBypassQueryParsingWithNonePartitionKey()
+        private const int DocumentCount = 400;
+
+        private static readonly Documents.PartitionKeyDefinition HierarchicalPartitionKeyDefinition = new Documents.PartitionKeyDefinition
         {
-            int documentCount = 400;
+            Paths = new Collection<string> { "/nullField", "/numberField" },
+            Kind = Documents.PartitionKind.MultiHash,
+            Version = Documents.PartitionKeyDefinitionVersion.V2
+        };
 
-            string query = "SELECT VALUE r.numberField FROM r";
-            IReadOnlyList<string> expectedOutput = Enumerable.Range(0, documentCount).Select(i => i.ToString()).ToList();
+        private static readonly Documents.PartitionKeyDefinition SimplePartitionKeyDefinition = new Documents.PartitionKeyDefinition
+        {
+            Paths = new Collection<string> { "/UndefinedField" },
+            Kind = Documents.PartitionKind.Hash
+        };
 
-            await this.ValidateQueryBypassWithNonePartitionKey(documentCount, query, expectedOutput);
+        [TestMethod]
+        [TestCategory("Query")]
+        public async Task TestBypassQueryParsing()
+        {
+            IReadOnlyList<QueryTestCase> testCases = new List<QueryTestCase>
+            {
+                new(
+                    PartitionKey.None,
+                    "SELECT VALUE r.numberField FROM r",
+                    Enumerable.Range(0, DocumentCount).Select(i => i.ToString()).ToList(),
+                    TestInjections.PipelineType.Passthrough),
+                new(
+                    PartitionKey.None,
+                    @"SELECT VALUE { """" : r.numberField } FROM r",
+                    Enumerable.Range(0, DocumentCount).Select(i => String.Format("{{\"\":{0}}}", i)).ToList(),
+                    TestInjections.PipelineType.Passthrough),
+            };
+
+            await this.ValidateQueryBypass(
+                ConnectionModes.Direct | ConnectionModes.Gateway,
+                CollectionTypes.NonPartitioned | CollectionTypes.SinglePartition | CollectionTypes.MultiPartition,
+                SimplePartitionKeyDefinition,
+                testCases);
         }
 
         [TestMethod]
         [TestCategory("Query")]
-        public async Task TestBypassQueryParsingWithNonePartitionKeyEmptyPropertyName()
+        public async Task TestBypassQueryParsingWithHPK()
         {
-            int documentCount = 400;
-
-            string query = @"SELECT VALUE { """" : r.numberField } FROM r";
-            IReadOnlyList<string> expectedOutput = Enumerable.Range(0, documentCount).Select(i => String.Format("{{\"\":{0}}}", i)).ToList();
+            PartitionKey partialPartitionKey = new PartitionKeyBuilder().AddNullValue().Build();
 
-            await this.ValidateQueryBypassWithNonePartitionKey(documentCount, query, expectedOutput);
+            IReadOnlyList<QueryTestCase> testCases = new List<QueryTestCase>
+            {
+                new(
+                    partialPartitionKey,
+                    "SELECT VALUE r.numberField FROM r",
+                    Enumerable.Range(0, DocumentCount).Select(i => i.ToString()).ToList(),
+                    TestInjections.PipelineType.Passthrough), // Passthrough because it is a client streaming query
+                new(
+                    partialPartitionKey,
+                    $"SELECT TOP {DocumentCount} VALUE r.numberField FROM r",
+                    Enumerable.Range(0, DocumentCount).Select(i => i.ToString()).ToList(),
+                    TestInjections.PipelineType.Specialized),
+                new(
+                    partialPartitionKey,
+                    @"SELECT VALUE { """" : r.numberField } FROM r",
+                    Enumerable.Range(0, DocumentCount).Select(i => String.Format("{{\"\":{0}}}", i)).ToList(),
+                    TestInjections.PipelineType.Passthrough), // Passthrough because it is a client streaming query
+                new(
+                    partialPartitionKey,
+                    "SELECT VALUE r.numberField FROM r ORDER BY r.numberField",
+                    Enumerable.Range(0, DocumentCount).Select(i => i.ToString()).ToList(),
+                    TestInjections.PipelineType.Specialized),
+            };
+
+            await this.ValidateQueryBypass(
+                ConnectionModes.Gateway,
+                CollectionTypes.MultiPartition,
+                HierarchicalPartitionKeyDefinition,
+                testCases);
         }
 
-        private async Task ValidateQueryBypassWithNonePartitionKey(int documentCount, string query, IReadOnlyList<string> expectedOutput)
+        private Task ValidateQueryBypass(ConnectionModes connectionModes, CollectionTypes collectionTypes, Documents.PartitionKeyDefinition partitionKeyDefinition, IReadOnlyList<QueryTestCase> testCases)
         {
-            QueryRequestOptions feedOptions = new QueryRequestOptions { PartitionKey = PartitionKey.None };
+            IReadOnlyList<string> documents = CreateDocuments(DocumentCount);
 
-            async Task ImplementationAsync(Container container, IReadOnlyList<CosmosObject> documents)
+            return this.CreateIngestQueryDeleteAsync(
+                connectionModes,
+                collectionTypes,
+                documents,
+                query: (container, _) => RunTestsAsync(container, testCases),
+                partitionKeyDefinition);
+        }
+
+        private static async Task RunTestsAsync(Container container, IReadOnlyList<QueryTestCase> testCases)
+        {
+            foreach (QueryTestCase testCase in testCases)
             {
+                QueryRequestOptions feedOptions = new QueryRequestOptions
+                {
+                    PartitionKey = testCase.PartitionKey,
+                    TestSettings = new TestInjections(simulate429s: false, simulateEmptyPages: false, responseStats: new())
+                };
+
                 ContainerInternal containerCore = container as ContainerInlineCore;
 
                 MockCosmosQueryClient cosmosQueryClientCore = new MockCosmosQueryClient(
@@ -53,20 +125,19 @@ async Task ImplementationAsync(Container container, IReadOnlyList<CosmosObject>
                     containerCore.Id,
                     cosmosQueryClientCore);
 
-                List<CosmosElement> items = await RunQueryAsync(containerWithBypassParsing, query, feedOptions);
+                List<CosmosElement> items = await RunQueryAsync(containerWithBypassParsing, testCase.Query, feedOptions);
                 string[] actualOutput = items.Select(x => x.ToString()).ToArray();
 
-                Assert.IsTrue(expectedOutput.SequenceEqual(actualOutput));
-            }
+                if(!testCase.ExpectedOutput.SequenceEqual(actualOutput))
+                {
+                    System.Diagnostics.Trace.WriteLine($"Expected: [{string.Join(", ", testCase.ExpectedOutput)}]");
+                    System.Diagnostics.Trace.WriteLine($"Actual:   [{string.Join(", ", actualOutput)}]");
 
-            IReadOnlyList<string> documents = CreateDocuments(documentCount);
+                    Assert.Fail("Query results do not match expected results.");
+                }
 
-            await this.CreateIngestQueryDeleteAsync(
-                ConnectionModes.Direct | ConnectionModes.Gateway,
-                CollectionTypes.NonPartitioned | CollectionTypes.SinglePartition | CollectionTypes.MultiPartition,
-                documents,
-                ImplementationAsync,
-                "/undefinedPartitionKey");
+                Assert.AreEqual(testCase.ExpectedPipelineType, feedOptions.TestSettings.Stats.PipelineType);
+            }
         }
 
         private static IReadOnlyList<string> CreateDocuments(int documentCount)
@@ -80,5 +151,24 @@ private static IReadOnlyList<string> CreateDocuments(int documentCount)
 
             return documents;
         }
+
+        private sealed class QueryTestCase
+        {
+            public PartitionKey PartitionKey { get; }
+
+            public string Query { get; }
+
+            public IReadOnlyList<string> ExpectedOutput { get; }
+
+            public TestInjections.PipelineType ExpectedPipelineType { get; }
+
+            public QueryTestCase(PartitionKey partitionKey, string query, IReadOnlyList<string> expectedOutput, TestInjections.PipelineType expectedPipelineType)
+            {
+                this.PartitionKey = partitionKey;
+                this.Query = query ?? throw new ArgumentNullException(nameof(query));
+                this.ExpectedOutput = expectedOutput ?? throw new ArgumentNullException(nameof(expectedOutput));
+                this.ExpectedPipelineType = expectedPipelineType;
+            }
+        }
     }
 }