Subpartitioning: Fixes bug for queries on subpartitioned containers

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

@@ -7,6 +7,7 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionContext
     using System.Collections.Generic;
     using System.Diagnostics;
     using System.Linq;
+    using System.Runtime.CompilerServices;
     using System.Threading;
     using System.Threading.Tasks;
     using Microsoft.Azure.Cosmos;
@@ -211,10 +212,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 +636,47 @@ 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(
-                    resourceLink,
-                    containerQueryProperties.ResourceId,
-                    effectivePartitionKeyString,
-                    forceRefresh: false,
-                    trace);
-            }
-            else if (feedRangeInternal != null)
-            {
-                targetRanges = await queryClient.GetTargetPartitionKeyRangeByFeedRangeAsync(
+                List<Documents.Routing.Range<string>> effectiveRanges = new List<Documents.Routing.Range<string>>
+                    { new Documents.Routing.Range<string>(effectivePartitionKeyString, effectivePartitionKeyString, true, true) };
+                targetRanges = await queryClient.GetTargetPartitionKeyRangesAsync(
                     resourceLink,
                     containerQueryProperties.ResourceId,
-                    containerQueryProperties.PartitionKeyDefinition,
-                    feedRangeInternal,
+                    effectiveRanges,
                     forceRefresh: false,
                     trace);
             }
             else
             {
-                targetRanges = await queryClient.GetTargetPartitionKeyRangesAsync(
-                    resourceLink,
-                    containerQueryProperties.ResourceId,
-                    partitionedQueryExecutionInfo.QueryRanges,
-                    forceRefresh: false,
-                    trace);
+                targetRanges = feedRangeInternal != null
+                    ? await queryClient.GetTargetPartitionKeyRangeByFeedRangeAsync(
+                                    resourceLink,
+                                    containerQueryProperties.ResourceId,
+                                    containerQueryProperties.PartitionKeyDefinition,
+                                    feedRangeInternal,
+                                    forceRefresh: false,
+                                    trace)
+                    : 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 +699,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 +782,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

@@ -35,12 +35,30 @@ namespace Microsoft.Azure.Cosmos.Routing
     /// </example>
     internal readonly struct PartitionKeyHash : IComparable<PartitionKeyHash>, IEquatable<PartitionKeyHash>
     {
+        private readonly 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)
+            {
+                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 UInt128[] HashValues => this.values;
 
         public int CompareTo(PartitionKeyHash other)
         {
@@ -66,7 +84,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/PartitionKeyHashRangeSplitterAndMerger.cs

@@ -39,8 +39,8 @@ public static SplitOutcome TrySplitRange(PartitionKeyHashRange partitionKeyHashR
                 return SplitOutcome.NumRangesNeedsToBeGreaterThanZero;
             }
 
-            UInt128 actualEnd = partitionKeyHashRange.EndExclusive.HasValue ? partitionKeyHashRange.EndExclusive.Value.Value : UInt128.MaxValue;
-            UInt128 actualStart = partitionKeyHashRange.StartInclusive.HasValue ? partitionKeyHashRange.StartInclusive.Value.Value : UInt128.MinValue;
+            UInt128 actualEnd = partitionKeyHashRange.EndExclusive.HasValue ? partitionKeyHashRange.EndExclusive.Value.HashValues[0] : UInt128.MaxValue;
+            UInt128 actualStart = partitionKeyHashRange.StartInclusive.HasValue ? partitionKeyHashRange.StartInclusive.Value.HashValues[0] : UInt128.MinValue;
             UInt128 rangeLength = actualEnd - actualStart;
             if (rangeLength < rangeCount)
             {
@@ -67,7 +67,7 @@ public static SplitOutcome TrySplitRange(PartitionKeyHashRange partitionKeyHashR
             for (int i = 1; i < rangeCount - 1; i++)
             {
                 PartitionKeyHash start = new PartitionKeyHash(actualStart + (childRangeLength * i));
-                PartitionKeyHash end = new PartitionKeyHash(start.Value + childRangeLength);
+                PartitionKeyHash end = new PartitionKeyHash(start.HashValues[0] + childRangeLength);
                 childRanges.Add(new PartitionKeyHashRange(start, end));
             }
 
@@ -82,6 +82,39 @@ public static SplitOutcome TrySplitRange(PartitionKeyHashRange partitionKeyHashR
             return SplitOutcome.Success;
         }
 
+        public static PartitionKeyHashRanges SplitRange(PartitionKeyHashRange partitionKeyHashRange, PartitionKeyHash explicitSplitPoint)
+        {
+            SplitOutcome splitOutcome = PartitionKeyHashRangeSplitterAndMerger.TrySplitRange(
+                partitionKeyHashRange,
+                explicitSplitPoint,
+                out PartitionKeyHashRanges splitRanges);
+
+            return splitOutcome switch
+            {
+                SplitOutcome.Success => splitRanges,
+                _ => throw new ArgumentOutOfRangeException($"Unknown {nameof(SplitOutcome)}: {splitOutcome}."),
+            };
+        }
+
+        public static SplitOutcome TrySplitRange(PartitionKeyHashRange partitionKeyHashRange, PartitionKeyHash explicitSplitPoint, out PartitionKeyHashRanges splitRanges)
+        {
+            if (explicitSplitPoint < (partitionKeyHashRange.StartInclusive != null ? partitionKeyHashRange.StartInclusive : PartitionKeyHash.None) ||
+                explicitSplitPoint > partitionKeyHashRange.EndExclusive)
+            {
+                splitRanges = default;
+                return SplitOutcome.RangeNotWideEnough;
+            }
+
+            IList<PartitionKeyHashRange> childRanges = new List<PartitionKeyHashRange>
+            {
+                new PartitionKeyHashRange(partitionKeyHashRange.StartInclusive, explicitSplitPoint),
+                new PartitionKeyHashRange(explicitSplitPoint, partitionKeyHashRange.EndExclusive)
+            };
+
+            splitRanges = PartitionKeyHashRanges.Create(childRanges);
+            return SplitOutcome.Success;
+        }
+
         public static PartitionKeyHashRange MergeRanges(PartitionKeyHashRanges partitionedSortedEffectiveRanges)
         {
             if (partitionedSortedEffectiveRanges == null)
@@ -119,4 +152,4 @@ public enum SplitOutcome
             RangeNotWideEnough,
         }
     }
-}
+}