Subpartitioning: Fixes bug for queries on subpartitioned containers

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

@@ -9,6 +9,7 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionContext
     using System.Linq;
     using System.Threading;
     using System.Threading.Tasks;
+    using global::Azure;
     using Microsoft.Azure.Cosmos;
     using Microsoft.Azure.Cosmos.CosmosElements;
     using Microsoft.Azure.Cosmos.Pagination;
@@ -27,6 +28,7 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionContext
     using Microsoft.Azure.Cosmos.SqlObjects;
     using Microsoft.Azure.Cosmos.SqlObjects.Visitors;
     using Microsoft.Azure.Cosmos.Tracing;
+    using Microsoft.Azure.Documents;
 
     internal static class CosmosQueryExecutionContextFactory
     {
@@ -211,12 +213,12 @@ private static async Task<TryCatch<IQueryPipelineStage>> TryCreateCoreContextAsy
 
                                 // Only thing that matters is that we target the correct range.
                                 Documents.PartitionKeyDefinition partitionKeyDefinition = GetPartitionKeyDefinition(inputParameters, containerQueryProperties);
-                                List<Documents.PartitionKeyRange> targetRanges = await cosmosQueryContext.QueryClient.GetTargetPartitionKeyRangesByEpkStringAsync(
+                                List<PartitionKeyRange> targetRanges = await cosmosQueryContext.QueryClient.GetTargetPartitionKeyRangesAsync(
                                     cosmosQueryContext.ResourceLink,
                                     containerQueryProperties.ResourceId,
-                                    containerQueryProperties.EffectivePartitionKeyString,
+                                    containerQueryProperties.EffectivePartitionKeyRanges,
                                     forceRefresh: false,
-                                    createQueryPipelineTrace);
+                                    trace);
 
                                 return CosmosQueryExecutionContextFactory.TryCreatePassthroughQueryExecutionContext(
                                     documentContainer,
@@ -635,12 +637,12 @@ private static async Task<PartitionedQueryExecutionInfo> GetPartitionedQueryExec
             ITrace trace)
         {
             List<Documents.PartitionKeyRange> targetRanges;
-            if (containerQueryProperties.EffectivePartitionKeyString != null)
+            if (containerQueryProperties.EffectivePartitionKeyRanges != null)
             {
-                targetRanges = await queryClient.GetTargetPartitionKeyRangesByEpkStringAsync(
+                targetRanges = await queryClient.GetTargetPartitionKeyRangesAsync(
                     resourceLink,
                     containerQueryProperties.ResourceId,
-                    containerQueryProperties.EffectivePartitionKeyString,
+                    containerQueryProperties.EffectivePartitionKeyRanges,
                     forceRefresh: false,
                     trace);
             }
@@ -771,14 +773,13 @@ private static Documents.PartitionKeyDefinition GetPartitionKeyDefinition(InputP
             else
             {
                 Documents.PartitionKeyDefinition partitionKeyDefinition = GetPartitionKeyDefinition(inputParameters, containerQueryProperties);
-                if (inputParameters.PartitionKey != null)
+                if (inputParameters.PartitionKey.HasValue)
                 {
                     Debug.Assert(partitionKeyDefinition != null, "CosmosQueryExecutionContextFactory Assert!", "PartitionKeyDefinition cannot be null if partitionKey is defined");
-
-                    targetRanges = await cosmosQueryContext.QueryClient.GetTargetPartitionKeyRangesByEpkStringAsync(
+                    targetRanges = await cosmosQueryContext.QueryClient.GetTargetPartitionKeyRangesAsync(
                         cosmosQueryContext.ResourceLink,
                         containerQueryProperties.ResourceId,
-                        containerQueryProperties.EffectivePartitionKeyString,
+                        containerQueryProperties.EffectivePartitionKeyRanges,
                         forceRefresh: false,
                         trace);
                 }
@@ -801,6 +802,24 @@ private static Documents.PartitionKeyDefinition GetPartitionKeyDefinition(InputP
             return null;
         }
 
+        private static Documents.Routing.PartitionKeyInternal GetNoneValue(
+            PartitionKeyDefinition partitionKeyDefinition)
+        {
+            if (partitionKeyDefinition == null)
+            {
+                throw new ArgumentNullException($"{nameof(partitionKeyDefinition)}");
+            }
+
+            if (partitionKeyDefinition.Paths.Count == 0 || (partitionKeyDefinition.IsSystemKey == true))
+            {
+                return Documents.Routing.PartitionKeyInternal.Empty;
+            }
+            else
+            {
+                return Documents.Routing.PartitionKeyInternal.Undefined;
+            }
+        }
+
         public sealed class InputParameters
         {
             private const int DefaultMaxConcurrency = 0;
@@ -816,7 +835,7 @@ public InputParameters(
                 int? maxConcurrency,
                 int? maxItemCount,
                 int? maxBufferedItemCount,
-                PartitionKey? partitionKey,
+                Cosmos.PartitionKey? partitionKey,
                 IReadOnlyDictionary<string, object> properties,
                 PartitionedQueryExecutionInfo partitionedQueryExecutionInfo,
                 ExecutionEnvironment? executionEnvironment,
@@ -866,7 +885,7 @@ public InputParameters(
             public int MaxConcurrency { get; }
             public int MaxItemCount { get; }
             public int MaxBufferedItemCount { get; }
-            public PartitionKey? PartitionKey { get; }
+            public Cosmos.PartitionKey? PartitionKey { get; }
             public IReadOnlyDictionary<string, object> Properties { get; }
             public PartitionedQueryExecutionInfo PartitionedQueryExecutionInfo { get; }
             public ExecutionEnvironment ExecutionEnvironment { get; }

Microsoft.Azure.Cosmos/src/Query/Core/QueryClient/ContainerQueryProperties.cs

@@ -4,24 +4,26 @@
 
 namespace Microsoft.Azure.Cosmos.Query.Core.QueryClient
 {
+    using System.Collections.Generic;
     using Microsoft.Azure.Documents;
+    using Microsoft.Azure.Documents.Routing;
 
     internal readonly struct ContainerQueryProperties
     {
         public ContainerQueryProperties(
             string resourceId,
-            string effectivePartitionKeyString,
+            List<Range<string>> effectivePartitionKeyRanges,
             PartitionKeyDefinition partitionKeyDefinition,
             Cosmos.GeospatialType geospatialType)
         {
             this.ResourceId = resourceId;
-            this.EffectivePartitionKeyString = effectivePartitionKeyString;
+            this.EffectivePartitionKeyRanges = effectivePartitionKeyRanges;
             this.PartitionKeyDefinition = partitionKeyDefinition;
             this.GeospatialType = geospatialType;
         }
 
         public string ResourceId { get; }
-        public string EffectivePartitionKeyString { get; }
+        public List<Range<string>> EffectivePartitionKeyRanges { get; }
         public PartitionKeyDefinition PartitionKeyDefinition { get; }
         public Cosmos.GeospatialType GeospatialType { get; }
     }

Microsoft.Azure.Cosmos/src/Query/v3Query/CosmosQueryClientCore.cs

@@ -63,17 +63,23 @@ public override async Task<ContainerQueryProperties> GetCachedContainerQueryProp
                 trace,
                 cancellationToken);
 
-            string effectivePartitionKeyString = null;
+            List<Range<string>> effectivePartitionKeyRange = null;
             if (partitionKey != null)
             {
                 // Dis-ambiguate the NonePK if used 
                 PartitionKeyInternal partitionKeyInternal = partitionKey.Value.IsNone ? containerProperties.GetNoneValue() : partitionKey.Value.InternalKey;
-                effectivePartitionKeyString = partitionKeyInternal.GetEffectivePartitionKeyString(containerProperties.PartitionKey);
+                effectivePartitionKeyRange.Add(PartitionKeyInternal.GetEffectivePartitionKeyRange(
+                        containerProperties.PartitionKey,
+                        new Range<PartitionKeyInternal>(
+                            min: partitionKeyInternal,
+                            max: partitionKeyInternal,
+                            isMinInclusive: true,
+                            isMaxInclusive: true)));
             }
 
             return new ContainerQueryProperties(
                 containerProperties.ResourceId,
-                effectivePartitionKeyString,
+                effectivePartitionKeyRange,
                 containerProperties.PartitionKey,
                 containerProperties.GeospatialConfig.GeospatialType);
         }

Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CosmosQueryClientCoreTest.cs

@@ -51,7 +51,7 @@ public async Task TryGetOverlappingRangesAsyncTest()
 
             Assert.IsNotNull(containerProperties);
             Assert.IsNotNull(containerProperties.ResourceId);
-            Assert.IsNotNull(containerProperties.EffectivePartitionKeyString);
+            Assert.IsNotNull(containerProperties.EffectivePartitionKeyRanges);
 
             IReadOnlyList<Documents.PartitionKeyRange> pkRange = await this.queryClientCore.TryGetOverlappingRangesAsync(
                 collectionResourceId: containerProperties.ResourceId,