Subpartitioning: Fixes bug for queries on subpartitioned containers

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

@@ -211,10 +211,10 @@ 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<Documents.PartitionKeyRange> targetRanges = await cosmosQueryContext.QueryClient.GetTargetPartitionKeyRangesAsync(
                                     cosmosQueryContext.ResourceLink,
                                     containerQueryProperties.ResourceId,
-                                    containerQueryProperties.EffectivePartitionKeyString,
+                                    containerQueryProperties.EffectivePartitionKeyRanges,
                                     forceRefresh: false,
                                     createQueryPipelineTrace);
 
@@ -635,12 +635,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 +771,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);
                 }

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,26 @@ 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 = new List<Range<string>>
+                {
+                    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,

Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs

@@ -959,9 +959,16 @@ public override Task ForceRefreshCollectionCacheAsync(string collectionLink, Can
 
         public override Task<ContainerQueryProperties> GetCachedContainerQueryPropertiesAsync(string containerLink, Cosmos.PartitionKey? partitionKey, ITrace trace, CancellationToken cancellationToken)
         {
-            return Task.FromResult(new ContainerQueryProperties(
+           return Task.FromResult(new ContainerQueryProperties(
                 "test",
-                WFConstants.BackendHeaders.EffectivePartitionKeyString,
+                new List<Range<string>>
+                { 
+                    new Range<string>(
+                        PartitionKeyInternal.MinimumInclusiveEffectivePartitionKey,
+                        PartitionKeyInternal.MaximumExclusiveEffectivePartitionKey,
+                        true,
+                        true)
+                },
                 new PartitionKeyDefinition(),
                 Cosmos.GeospatialType.Geometry));
         }