Skip to content

Commit 8c2c4c8

Browse files
NaluTripicianSrinikhilReddy
NaluTripician
and
SrinikhilReddy
authored
Subpartitioning: Fixes bug for queries on subpartitioned containers (#3934) (#4085)
* 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]>
1 parent aff544b commit 8c2c4c8

16 files changed

+471
-141
lines changed

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

Lines changed: 51 additions & 43 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionContext
99
using System.Linq;
1010
using System.Threading;
1111
using System.Threading.Tasks;
12+
using global::Azure;
1213
using Microsoft.Azure.Cosmos;
1314
using Microsoft.Azure.Cosmos.CosmosElements;
1415
using Microsoft.Azure.Cosmos.Pagination;
@@ -27,6 +28,7 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionContext
2728
using Microsoft.Azure.Cosmos.SqlObjects;
2829
using Microsoft.Azure.Cosmos.SqlObjects.Visitors;
2930
using Microsoft.Azure.Cosmos.Tracing;
31+
using Microsoft.Azure.Documents.Routing;
3032

3133
internal static class CosmosQueryExecutionContextFactory
3234
{
@@ -211,10 +213,10 @@ private static async Task<TryCatch<IQueryPipelineStage>> TryCreateCoreContextAsy
211213

212214
// Only thing that matters is that we target the correct range.
213215
Documents.PartitionKeyDefinition partitionKeyDefinition = GetPartitionKeyDefinition(inputParameters, containerQueryProperties);
214-
List<Documents.PartitionKeyRange> targetRanges = await cosmosQueryContext.QueryClient.GetTargetPartitionKeyRangesByEpkStringAsync(
216+
List<Documents.PartitionKeyRange> targetRanges = await cosmosQueryContext.QueryClient.GetTargetPartitionKeyRangesAsync(
215217
cosmosQueryContext.ResourceLink,
216218
containerQueryProperties.ResourceId,
217-
containerQueryProperties.EffectivePartitionKeyString,
219+
containerQueryProperties.EffectiveRangesForPartitionKey,
218220
forceRefresh: false,
219221
createQueryPipelineTrace);
220222

@@ -635,67 +637,54 @@ private static async Task<PartitionedQueryExecutionInfo> GetPartitionedQueryExec
635637
ITrace trace)
636638
{
637639
List<Documents.PartitionKeyRange> targetRanges;
638-
if (containerQueryProperties.EffectivePartitionKeyString != null)
640+
if (containerQueryProperties.EffectiveRangesForPartitionKey != null)
639641
{
640-
targetRanges = await queryClient.GetTargetPartitionKeyRangesByEpkStringAsync(
642+
targetRanges = await queryClient.GetTargetPartitionKeyRangesAsync(
641643
resourceLink,
642644
containerQueryProperties.ResourceId,
643-
containerQueryProperties.EffectivePartitionKeyString,
645+
containerQueryProperties.EffectiveRangesForPartitionKey,
644646
forceRefresh: false,
645647
trace);
646648
}
647649
else if (TryGetEpkProperty(properties, out string effectivePartitionKeyString))
648650
{
649-
targetRanges = await queryClient.GetTargetPartitionKeyRangesByEpkStringAsync(
651+
//Note that here we have no way to consume the EPK string as there is no way to convert
652+
//the string to the partition key type to evaulate the number of components which needs to be done for the
653+
//multihahs methods/classes. This is particually important for queries with prefix partition key.
654+
//the EPK sting header is only for internal use but this needs to be fixed in the future.
655+
List<Range<string>> effectiveRanges = new List<Range<string>>
656+
{ Range<string>.GetPointRange(effectivePartitionKeyString) };
657+
658+
targetRanges = await queryClient.GetTargetPartitionKeyRangesAsync(
650659
resourceLink,
651660
containerQueryProperties.ResourceId,
652-
effectivePartitionKeyString,
661+
effectiveRanges,
653662
forceRefresh: false,
654663
trace);
655664
}
656665
else if (feedRangeInternal != null)
657666
{
658667
targetRanges = await queryClient.GetTargetPartitionKeyRangeByFeedRangeAsync(
659-
resourceLink,
660-
containerQueryProperties.ResourceId,
661-
containerQueryProperties.PartitionKeyDefinition,
662-
feedRangeInternal,
663-
forceRefresh: false,
664-
trace);
668+
resourceLink,
669+
containerQueryProperties.ResourceId,
670+
containerQueryProperties.PartitionKeyDefinition,
671+
feedRangeInternal,
672+
forceRefresh: false,
673+
trace);
665674
}
666675
else
667676
{
668-
targetRanges = await queryClient.GetTargetPartitionKeyRangesAsync(
669-
resourceLink,
670-
containerQueryProperties.ResourceId,
671-
partitionedQueryExecutionInfo.QueryRanges,
672-
forceRefresh: false,
673-
trace);
677+
targetRanges = await queryClient.GetTargetPartitionKeyRangesAsync(
678+
resourceLink,
679+
containerQueryProperties.ResourceId,
680+
partitionedQueryExecutionInfo.QueryRanges,
681+
forceRefresh: false,
682+
trace);
674683
}
675684

676685
return targetRanges;
677686
}
678687

679-
private static void SetTestInjectionPipelineType(InputParameters inputParameters, string pipelineType)
680-
{
681-
TestInjections.ResponseStats responseStats = inputParameters?.TestInjections?.Stats;
682-
if (responseStats != null)
683-
{
684-
if (pipelineType == OptimisticDirectExecution)
685-
{
686-
responseStats.PipelineType = TestInjections.PipelineType.OptimisticDirectExecution;
687-
}
688-
else if (pipelineType == Specialized)
689-
{
690-
responseStats.PipelineType = TestInjections.PipelineType.Specialized;
691-
}
692-
else
693-
{
694-
responseStats.PipelineType = TestInjections.PipelineType.Passthrough;
695-
}
696-
}
697-
}
698-
699688
private static bool TryGetEpkProperty(
700689
IReadOnlyDictionary<string, object> properties,
701690
out string effectivePartitionKeyString)
@@ -718,6 +707,26 @@ private static bool TryGetEpkProperty(
718707
return false;
719708
}
720709

710+
private static void SetTestInjectionPipelineType(InputParameters inputParameters, string pipelineType)
711+
{
712+
TestInjections.ResponseStats responseStats = inputParameters?.TestInjections?.Stats;
713+
if (responseStats != null)
714+
{
715+
if (pipelineType == OptimisticDirectExecution)
716+
{
717+
responseStats.PipelineType = TestInjections.PipelineType.OptimisticDirectExecution;
718+
}
719+
else if (pipelineType == Specialized)
720+
{
721+
responseStats.PipelineType = TestInjections.PipelineType.Specialized;
722+
}
723+
else
724+
{
725+
responseStats.PipelineType = TestInjections.PipelineType.Passthrough;
726+
}
727+
}
728+
}
729+
721730
private static Documents.PartitionKeyDefinition GetPartitionKeyDefinition(InputParameters inputParameters, ContainerQueryProperties containerQueryProperties)
722731
{
723732
//todo:elasticcollections this may rely on information from collection cache which is outdated
@@ -771,14 +780,13 @@ private static Documents.PartitionKeyDefinition GetPartitionKeyDefinition(InputP
771780
else
772781
{
773782
Documents.PartitionKeyDefinition partitionKeyDefinition = GetPartitionKeyDefinition(inputParameters, containerQueryProperties);
774-
if (inputParameters.PartitionKey != null)
783+
if (inputParameters.PartitionKey.HasValue)
775784
{
776785
Debug.Assert(partitionKeyDefinition != null, "CosmosQueryExecutionContextFactory Assert!", "PartitionKeyDefinition cannot be null if partitionKey is defined");
777-
778-
targetRanges = await cosmosQueryContext.QueryClient.GetTargetPartitionKeyRangesByEpkStringAsync(
786+
targetRanges = await cosmosQueryContext.QueryClient.GetTargetPartitionKeyRangesAsync(
779787
cosmosQueryContext.ResourceLink,
780788
containerQueryProperties.ResourceId,
781-
containerQueryProperties.EffectivePartitionKeyString,
789+
containerQueryProperties.EffectiveRangesForPartitionKey,
782790
forceRefresh: false,
783791
trace);
784792
}

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

Lines changed: 8 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -4,24 +4,29 @@
44

55
namespace Microsoft.Azure.Cosmos.Query.Core.QueryClient
66
{
7+
using System.Collections.Generic;
78
using Microsoft.Azure.Documents;
9+
using Microsoft.Azure.Documents.Routing;
810

911
internal readonly struct ContainerQueryProperties
1012
{
1113
public ContainerQueryProperties(
1214
string resourceId,
13-
string effectivePartitionKeyString,
15+
IReadOnlyList<Range<string>> effectivePartitionKeyRanges,
1416
PartitionKeyDefinition partitionKeyDefinition,
1517
Cosmos.GeospatialType geospatialType)
1618
{
1719
this.ResourceId = resourceId;
18-
this.EffectivePartitionKeyString = effectivePartitionKeyString;
20+
this.EffectiveRangesForPartitionKey = effectivePartitionKeyRanges;
1921
this.PartitionKeyDefinition = partitionKeyDefinition;
2022
this.GeospatialType = geospatialType;
2123
}
2224

2325
public string ResourceId { get; }
24-
public string EffectivePartitionKeyString { get; }
26+
27+
//A PartitionKey has one range when it is a full PartitionKey value.
28+
//It can span many it is a prefix PartitionKey for a sub-partitioned container.
29+
public IReadOnlyList<Range<string>> EffectiveRangesForPartitionKey { get; }
2530
public PartitionKeyDefinition PartitionKeyDefinition { get; }
2631
public Cosmos.GeospatialType GeospatialType { get; }
2732
}

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

Lines changed: 1 addition & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -76,13 +76,6 @@ public abstract Task<PartitionedQueryExecutionInfo> ExecuteQueryPlanRequestAsync
7676

7777
public abstract void ClearSessionTokenCache(string collectionFullName);
7878

79-
public abstract Task<List<Documents.PartitionKeyRange>> GetTargetPartitionKeyRangesByEpkStringAsync(
80-
string resourceLink,
81-
string collectionResourceId,
82-
string effectivePartitionKeyString,
83-
bool forceRefresh,
84-
ITrace trace);
85-
8679
public abstract Task<List<Documents.PartitionKeyRange>> GetTargetPartitionKeyRangeByFeedRangeAsync(
8780
string resourceLink,
8881
string collectionResourceId,
@@ -94,7 +87,7 @@ public abstract Task<PartitionedQueryExecutionInfo> ExecuteQueryPlanRequestAsync
9487
public abstract Task<List<Documents.PartitionKeyRange>> GetTargetPartitionKeyRangesAsync(
9588
string resourceLink,
9689
string collectionResourceId,
97-
List<Documents.Routing.Range<string>> providedRanges,
90+
IReadOnlyList<Documents.Routing.Range<string>> providedRanges,
9891
bool forceRefresh,
9992
ITrace trace);
10093

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

Lines changed: 13 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -63,17 +63,26 @@ public override async Task<ContainerQueryProperties> GetCachedContainerQueryProp
6363
trace,
6464
cancellationToken);
6565

66-
string effectivePartitionKeyString = null;
66+
List<Range<string>> effectivePartitionKeyRange = null;
6767
if (partitionKey != null)
6868
{
6969
// Dis-ambiguate the NonePK if used
7070
PartitionKeyInternal partitionKeyInternal = partitionKey.Value.IsNone ? containerProperties.GetNoneValue() : partitionKey.Value.InternalKey;
71-
effectivePartitionKeyString = partitionKeyInternal.GetEffectivePartitionKeyString(containerProperties.PartitionKey);
71+
effectivePartitionKeyRange = new List<Range<string>>
72+
{
73+
PartitionKeyInternal.GetEffectivePartitionKeyRange(
74+
containerProperties.PartitionKey,
75+
new Range<PartitionKeyInternal>(
76+
min: partitionKeyInternal,
77+
max: partitionKeyInternal,
78+
isMinInclusive: true,
79+
isMaxInclusive: true))
80+
};
7281
}
7382

7483
return new ContainerQueryProperties(
7584
containerProperties.ResourceId,
76-
effectivePartitionKeyString,
85+
effectivePartitionKeyRange,
7786
containerProperties.PartitionKey,
7887
containerProperties.GeospatialConfig.GeospatialType);
7988
}
@@ -200,24 +209,6 @@ public override async Task<PartitionedQueryExecutionInfo> ExecuteQueryPlanReques
200209
return partitionedQueryExecutionInfo;
201210
}
202211

203-
public override Task<List<PartitionKeyRange>> GetTargetPartitionKeyRangesByEpkStringAsync(
204-
string resourceLink,
205-
string collectionResourceId,
206-
string effectivePartitionKeyString,
207-
bool forceRefresh,
208-
ITrace trace)
209-
{
210-
return this.GetTargetPartitionKeyRangesAsync(
211-
resourceLink,
212-
collectionResourceId,
213-
new List<Range<string>>
214-
{
215-
Range<string>.GetPointRange(effectivePartitionKeyString)
216-
},
217-
forceRefresh,
218-
trace);
219-
}
220-
221212
public override async Task<List<PartitionKeyRange>> GetTargetPartitionKeyRangeByFeedRangeAsync(
222213
string resourceLink,
223214
string collectionResourceId,
@@ -243,7 +234,7 @@ public override async Task<List<PartitionKeyRange>> GetTargetPartitionKeyRangeBy
243234
public override async Task<List<PartitionKeyRange>> GetTargetPartitionKeyRangesAsync(
244235
string resourceLink,
245236
string collectionResourceId,
246-
List<Range<string>> providedRanges,
237+
IReadOnlyList<Range<string>> providedRanges,
247238
bool forceRefresh,
248239
ITrace trace)
249240
{

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

Lines changed: 26 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@
55
namespace Microsoft.Azure.Cosmos.Routing
66
{
77
using System;
8-
using System.Runtime.CompilerServices;
8+
using System.Collections.Generic;
99
using System.Runtime.InteropServices;
1010
using System.Text;
1111
using Microsoft.Azure.Documents.Routing;
@@ -35,12 +35,34 @@ namespace Microsoft.Azure.Cosmos.Routing
3535
/// </example>
3636
internal readonly struct PartitionKeyHash : IComparable<PartitionKeyHash>, IEquatable<PartitionKeyHash>
3737
{
38+
private readonly IReadOnlyList<UInt128> values;
39+
3840
public PartitionKeyHash(UInt128 value)
41+
: this(new UInt128[] { value })
3942
{
40-
this.Value = value;
4143
}
4244

43-
public UInt128 Value { get; }
45+
public PartitionKeyHash(UInt128[] values)
46+
{
47+
StringBuilder stringBuilder = new StringBuilder();
48+
foreach (UInt128 value in values)
49+
{
50+
if (stringBuilder.Length > 0)
51+
{
52+
stringBuilder.Append('-');
53+
}
54+
stringBuilder.Append(value.ToString());
55+
}
56+
57+
this.Value = stringBuilder.ToString();
58+
this.values = values;
59+
}
60+
61+
public readonly static PartitionKeyHash None = new PartitionKeyHash(0);
62+
63+
public string Value { get; }
64+
65+
internal readonly IReadOnlyList<UInt128> HashValues => this.values;
4466

4567
public int CompareTo(PartitionKeyHash other)
4668
{
@@ -66,7 +88,7 @@ public override bool Equals(object obj)
6688

6789
public override int GetHashCode() => this.Value.GetHashCode();
6890

69-
public override string ToString() => this.Value.ToString();
91+
public override string ToString() => this.Value;
7092

7193
public static bool TryParse(string value, out PartitionKeyHash parsedValue)
7294
{

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -181,4 +181,4 @@ public override string ToString()
181181
return stringBuilder.ToString();
182182
}
183183
}
184-
}
184+
}

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

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -132,9 +132,9 @@ public static CreateOutcome TryCreate(
132132
{
133133
if (partitionKeyHashRange.StartInclusive.HasValue)
134134
{
135-
if (partitionKeyHashRange.StartInclusive.Value.Value < minStart)
135+
if (partitionKeyHashRange.StartInclusive.Value.HashValues[0] < minStart)
136136
{
137-
minStart = partitionKeyHashRange.StartInclusive.Value.Value;
137+
minStart = partitionKeyHashRange.StartInclusive.Value.HashValues[0];
138138
}
139139
}
140140
else
@@ -144,18 +144,18 @@ public static CreateOutcome TryCreate(
144144

145145
if (partitionKeyHashRange.EndExclusive.HasValue)
146146
{
147-
if (partitionKeyHashRange.EndExclusive.Value.Value > maxEnd)
147+
if (partitionKeyHashRange.EndExclusive.Value.HashValues[0] > maxEnd)
148148
{
149-
maxEnd = partitionKeyHashRange.EndExclusive.Value.Value;
149+
maxEnd = partitionKeyHashRange.EndExclusive.Value.HashValues[0];
150150
}
151151
}
152152
else
153153
{
154154
maxEnd = UInt128.MaxValue;
155155
}
156156

157-
UInt128 width = partitionKeyHashRange.EndExclusive.GetValueOrDefault(new PartitionKeyHash(UInt128.MaxValue)).Value
158-
- partitionKeyHashRange.StartInclusive.GetValueOrDefault(new PartitionKeyHash(UInt128.MinValue)).Value;
157+
UInt128 width = partitionKeyHashRange.EndExclusive.GetValueOrDefault(new PartitionKeyHash(UInt128.MaxValue)).HashValues[0]
158+
- partitionKeyHashRange.StartInclusive.GetValueOrDefault(new PartitionKeyHash(UInt128.MinValue)).HashValues[0];
159159
sumOfWidth += width;
160160
if (sumOfWidth < width)
161161
{
@@ -223,4 +223,4 @@ public enum CreateOutcome
223223
Success,
224224
}
225225
}
226-
}
226+
}

0 commit comments

Comments
 (0)