Subpartitioning: Fixes bug for queries on subpartitioned containers (#3934)

* initial fix, needs testing on prod

* test fix

* clean up pr

* query rework

* refactors previous changes

* requested changes and bug fixes

* nits

* requested changes

* bug fixes

* start of test

* added test

* nit: changed name of EffectivePartitionKeyRanges to EffectiveRangesForPartitionKey

* Address code comments

* Address code comments

* saving work

* addresses code comments

* nit, spacing

* PartitionKeyHash fixes

* Fixes bugs in tests

* Removed bad method, added additional test coverage

* Removed EffectivePartitionKeyString use

* test fix

* requested changes

* Requested changes

* fixed test

* Test fix

* Added comment

---------

Co-authored-by: SrinikhilReddy <[email protected]>

Author: NaluTripician

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.Routing;
 
     internal static class CosmosQueryExecutionContextFactory
     {
@@ -211,10 +213,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.EffectiveRangesForPartitionKey,
                                     forceRefresh: false,
                                     createQueryPipelineTrace);
 
@@ -635,67 +637,54 @@ private static async Task<PartitionedQueryExecutionInfo> GetPartitionedQueryExec
             ITrace trace)
         {
             List<Documents.PartitionKeyRange> targetRanges;
-            if (containerQueryProperties.EffectivePartitionKeyString != null)
+            if (containerQueryProperties.EffectiveRangesForPartitionKey != null)
             {
-                targetRanges = await queryClient.GetTargetPartitionKeyRangesByEpkStringAsync(
+                targetRanges = await queryClient.GetTargetPartitionKeyRangesAsync(
                     resourceLink,
                     containerQueryProperties.ResourceId,
-                    containerQueryProperties.EffectivePartitionKeyString,
+                    containerQueryProperties.EffectiveRangesForPartitionKey,
                     forceRefresh: false,
                     trace);
             }
             else if (TryGetEpkProperty(properties, out string effectivePartitionKeyString))
             {
-                targetRanges = await queryClient.GetTargetPartitionKeyRangesByEpkStringAsync(
+                //Note that here we have no way to consume the EPK string as there is no way to convert
+                //the string to the partition key type to evaulate the number of components which needs to be done for the 
+                //multihahs methods/classes. This is particually important for queries with prefix partition key. 
+                //the EPK sting header is only for internal use but this needs to be fixed in the future.
+                List<Range<string>> effectiveRanges = new List<Range<string>>
+                    { Range<string>.GetPointRange(effectivePartitionKeyString) };
+
+                targetRanges = await queryClient.GetTargetPartitionKeyRangesAsync(
                     resourceLink,
                     containerQueryProperties.ResourceId,
-                    effectivePartitionKeyString,
+                    effectiveRanges,
                     forceRefresh: false,
                     trace);
             }
             else if (feedRangeInternal != null)
             {
                 targetRanges = await queryClient.GetTargetPartitionKeyRangeByFeedRangeAsync(
-                    resourceLink,
-                    containerQueryProperties.ResourceId,
-                    containerQueryProperties.PartitionKeyDefinition,
-                    feedRangeInternal,
-                    forceRefresh: false,
-                    trace);
+                                    resourceLink,
+                                    containerQueryProperties.ResourceId,
+                                    containerQueryProperties.PartitionKeyDefinition,
+                                    feedRangeInternal,
+                                    forceRefresh: false,
+                                    trace);
             }
             else
             {
-                targetRanges = await queryClient.GetTargetPartitionKeyRangesAsync(
-                    resourceLink,
-                    containerQueryProperties.ResourceId,
-                    partitionedQueryExecutionInfo.QueryRanges,
-                    forceRefresh: false,
-                    trace);
+                    targetRanges = await queryClient.GetTargetPartitionKeyRangesAsync(
+                                    resourceLink,
+                                    containerQueryProperties.ResourceId,
+                                    partitionedQueryExecutionInfo.QueryRanges,
+                                    forceRefresh: false,
+                                    trace);
             }
 
             return targetRanges;
         }
 
-        private static void SetTestInjectionPipelineType(InputParameters inputParameters, string pipelineType)
-        {
-            TestInjections.ResponseStats responseStats = inputParameters?.TestInjections?.Stats;
-            if (responseStats != null)
-            {
-                if (pipelineType == OptimisticDirectExecution)
-                {
-                    responseStats.PipelineType = TestInjections.PipelineType.OptimisticDirectExecution;
-                }
-                else if (pipelineType == Specialized)
-                {
-                    responseStats.PipelineType = TestInjections.PipelineType.Specialized;
-                }
-                else 
-                {
-                    responseStats.PipelineType = TestInjections.PipelineType.Passthrough;
-                }
-            }
-        }
-
         private static bool TryGetEpkProperty(
             IReadOnlyDictionary<string, object> properties,
             out string effectivePartitionKeyString)
@@ -718,6 +707,26 @@ private static bool TryGetEpkProperty(
             return false;
         }
 
+        private static void SetTestInjectionPipelineType(InputParameters inputParameters, string pipelineType)
+        {
+            TestInjections.ResponseStats responseStats = inputParameters?.TestInjections?.Stats;
+            if (responseStats != null)
+            {
+                if (pipelineType == OptimisticDirectExecution)
+                {
+                    responseStats.PipelineType = TestInjections.PipelineType.OptimisticDirectExecution;
+                }
+                else if (pipelineType == Specialized)
+                {
+                    responseStats.PipelineType = TestInjections.PipelineType.Specialized;
+                }
+                else 
+                {
+                    responseStats.PipelineType = TestInjections.PipelineType.Passthrough;
+                }
+            }
+        }
+
         private static Documents.PartitionKeyDefinition GetPartitionKeyDefinition(InputParameters inputParameters, ContainerQueryProperties containerQueryProperties)
         {
             //todo:elasticcollections this may rely on information from collection cache which is outdated
@@ -781,14 +790,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.EffectiveRangesForPartitionKey,
                         forceRefresh: false,
                         trace);
                 }

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

@@ -4,24 +4,29 @@
 
 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,
+            IReadOnlyList<Range<string>> effectivePartitionKeyRanges,
             PartitionKeyDefinition partitionKeyDefinition,
             Cosmos.GeospatialType geospatialType)
         {
             this.ResourceId = resourceId;
-            this.EffectivePartitionKeyString = effectivePartitionKeyString;
+            this.EffectiveRangesForPartitionKey = effectivePartitionKeyRanges;
             this.PartitionKeyDefinition = partitionKeyDefinition;
             this.GeospatialType = geospatialType;
         }
 
         public string ResourceId { get; }
-        public string EffectivePartitionKeyString { get; }
+
+        //A PartitionKey has one range when it is a full PartitionKey value.
+        //It can span many  it is a prefix PartitionKey for a sub-partitioned container.
+        public IReadOnlyList<Range<string>> EffectiveRangesForPartitionKey { get; }
         public PartitionKeyDefinition PartitionKeyDefinition { get; }
         public Cosmos.GeospatialType GeospatialType { get; }
     }

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

@@ -76,13 +76,6 @@ public abstract Task<PartitionedQueryExecutionInfo> ExecuteQueryPlanRequestAsync
 
         public abstract void ClearSessionTokenCache(string collectionFullName);
 
-        public abstract Task<List<Documents.PartitionKeyRange>> GetTargetPartitionKeyRangesByEpkStringAsync(
-            string resourceLink,
-            string collectionResourceId,
-            string effectivePartitionKeyString,
-            bool forceRefresh,
-            ITrace trace);
-
         public abstract Task<List<Documents.PartitionKeyRange>> GetTargetPartitionKeyRangeByFeedRangeAsync(
             string resourceLink,
             string collectionResourceId,
@@ -94,7 +87,7 @@ public abstract Task<PartitionedQueryExecutionInfo> ExecuteQueryPlanRequestAsync
         public abstract Task<List<Documents.PartitionKeyRange>> GetTargetPartitionKeyRangesAsync(
             string resourceLink,
             string collectionResourceId,
-            List<Documents.Routing.Range<string>> providedRanges,
+            IReadOnlyList<Documents.Routing.Range<string>> providedRanges,
             bool forceRefresh,
             ITrace trace);
 

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);
         }
@@ -200,24 +209,6 @@ public override async Task<PartitionedQueryExecutionInfo> ExecuteQueryPlanReques
             return partitionedQueryExecutionInfo;
         }
 
-        public override Task<List<PartitionKeyRange>> GetTargetPartitionKeyRangesByEpkStringAsync(
-            string resourceLink,
-            string collectionResourceId,
-            string effectivePartitionKeyString,
-            bool forceRefresh,
-            ITrace trace)
-        {
-            return this.GetTargetPartitionKeyRangesAsync(
-                resourceLink,
-                collectionResourceId,
-                new List<Range<string>>
-                {
-                    Range<string>.GetPointRange(effectivePartitionKeyString)
-                },
-                forceRefresh,
-                trace);
-        }
-
         public override async Task<List<PartitionKeyRange>> GetTargetPartitionKeyRangeByFeedRangeAsync(
             string resourceLink,
             string collectionResourceId,
@@ -243,7 +234,7 @@ public override async Task<List<PartitionKeyRange>> GetTargetPartitionKeyRangeBy
         public override async Task<List<PartitionKeyRange>> GetTargetPartitionKeyRangesAsync(
             string resourceLink,
             string collectionResourceId,
-            List<Range<string>> providedRanges,
+            IReadOnlyList<Range<string>> providedRanges,
             bool forceRefresh,
             ITrace trace)
         {

Microsoft.Azure.Cosmos/src/Routing/PartitionKeyHash.cs

@@ -5,7 +5,7 @@
 namespace Microsoft.Azure.Cosmos.Routing
 {
     using System;
-    using System.Runtime.CompilerServices;
+    using System.Collections.Generic;
     using System.Runtime.InteropServices;
     using System.Text;
     using Microsoft.Azure.Documents.Routing;
@@ -35,12 +35,34 @@ namespace Microsoft.Azure.Cosmos.Routing
     /// </example>
     internal readonly struct PartitionKeyHash : IComparable<PartitionKeyHash>, IEquatable<PartitionKeyHash>
     {
+        private readonly IReadOnlyList<UInt128> values;
+
         public PartitionKeyHash(UInt128 value)
+            : this(new UInt128[] { value })
         {
-            this.Value = value;
         }
 
-        public UInt128 Value { get; }
+        public PartitionKeyHash(UInt128[] values)
+        {
+            StringBuilder stringBuilder = new StringBuilder();
+            foreach (UInt128 value in values)
+            {
+                if (stringBuilder.Length > 0)
+                {
+                    stringBuilder.Append('-');
+                }
+                stringBuilder.Append(value.ToString());
+            }
+
+            this.Value = stringBuilder.ToString();
+            this.values = values;
+        }
+
+        public readonly static PartitionKeyHash None = new PartitionKeyHash(0);
+
+        public string Value { get; }
+
+        internal readonly IReadOnlyList<UInt128> HashValues => this.values;
 
         public int CompareTo(PartitionKeyHash other)
         {
@@ -66,7 +88,7 @@ public override bool Equals(object obj)
 
         public override int GetHashCode() => this.Value.GetHashCode();
 
-        public override string ToString() => this.Value.ToString();
+        public override string ToString() => this.Value;
 
         public static bool TryParse(string value, out PartitionKeyHash parsedValue)
         {

Microsoft.Azure.Cosmos/src/Routing/PartitionKeyHashRange.cs

@@ -181,4 +181,4 @@ public override string ToString()
             return stringBuilder.ToString();
         }
     }
-}
+}

Microsoft.Azure.Cosmos/src/Routing/PartitionKeyHashRanges.cs

@@ -132,9 +132,9 @@ public static CreateOutcome TryCreate(
             {
                 if (partitionKeyHashRange.StartInclusive.HasValue)
                 {
-                    if (partitionKeyHashRange.StartInclusive.Value.Value < minStart)
+                    if (partitionKeyHashRange.StartInclusive.Value.HashValues[0] < minStart)
                     {
-                        minStart = partitionKeyHashRange.StartInclusive.Value.Value;
+                        minStart = partitionKeyHashRange.StartInclusive.Value.HashValues[0];
                     }
                 }
                 else
@@ -144,18 +144,18 @@ public static CreateOutcome TryCreate(
 
                 if (partitionKeyHashRange.EndExclusive.HasValue)
                 {
-                    if (partitionKeyHashRange.EndExclusive.Value.Value > maxEnd)
+                    if (partitionKeyHashRange.EndExclusive.Value.HashValues[0] > maxEnd)
                     {
-                        maxEnd = partitionKeyHashRange.EndExclusive.Value.Value;
+                        maxEnd = partitionKeyHashRange.EndExclusive.Value.HashValues[0];
                     }
                 }
                 else
                 {
                     maxEnd = UInt128.MaxValue;
                 }
 
-                UInt128 width = partitionKeyHashRange.EndExclusive.GetValueOrDefault(new PartitionKeyHash(UInt128.MaxValue)).Value
-                    - partitionKeyHashRange.StartInclusive.GetValueOrDefault(new PartitionKeyHash(UInt128.MinValue)).Value;
+                UInt128 width = partitionKeyHashRange.EndExclusive.GetValueOrDefault(new PartitionKeyHash(UInt128.MaxValue)).HashValues[0]
+                    - partitionKeyHashRange.StartInclusive.GetValueOrDefault(new PartitionKeyHash(UInt128.MinValue)).HashValues[0];
                 sumOfWidth += width;
                 if (sumOfWidth < width)
                 {
@@ -223,4 +223,4 @@ public enum CreateOutcome
             Success,
         }
     }
-}
+}