Skip to content

Commit f1816f6

Browse files
adityasaadityasa
committed
Query : Fixes ORDER BY query issue when partial partition key is specified with hierarchical partition (#4507)
* Initial commit * Initial commit * Update.
1 parent 618d120 commit f1816f6

File tree

4 files changed

+138
-32
lines changed

4 files changed

+138
-32
lines changed

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

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -248,6 +248,8 @@ private static async Task<TryCatch<IQueryPipelineStage>> TryCreateFromPartitione
248248
inputParameters.InitialFeedRange,
249249
trace);
250250

251+
Debug.Assert(targetRanges != null, $"{nameof(CosmosQueryExecutionContextFactory)} Assert!", "targetRanges != null");
252+
251253
TryCatch<IQueryPipelineStage> tryCreatePipelineStage;
252254
Documents.PartitionKeyRange targetRange = await TryGetTargetRangeOptimisticDirectExecutionAsync(
253255
inputParameters,
@@ -270,7 +272,7 @@ private static async Task<TryCatch<IQueryPipelineStage>> TryCreateFromPartitione
270272
}
271273
else
272274
{
273-
bool singleLogicalPartitionKeyQuery = inputParameters.PartitionKey.HasValue
275+
bool singleLogicalPartitionKeyQuery = (inputParameters.PartitionKey.HasValue && targetRanges.Count == 1)
274276
|| ((partitionedQueryExecutionInfo.QueryRanges.Count == 1)
275277
&& partitionedQueryExecutionInfo.QueryRanges[0].IsSingleValue);
276278
bool serverStreamingQuery = !partitionedQueryExecutionInfo.QueryInfo.HasAggregates

‎Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/BaselineTest/TestBaseline/SubpartitionTests.TestQueriesOnSplitContainer.xml

Lines changed: 58 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -57,4 +57,62 @@
5757
{"id":"2","value2":"97"}]]></Documents>
5858
</Output>
5959
</Result>
60+
<Result>
61+
<Input>
62+
<Description>SELECT ORDER BY with ODE</Description>
63+
<Query><![CDATA[SELECT c.id, c.value2, c.intVal FROM c ORDER BY c.intVal]]></Query>
64+
<ODE>True</ODE>
65+
</Input>
66+
<Output>
67+
<Documents><![CDATA[{"id":"2","value2":"97","intVal":-47},
68+
{"id":"2","value2":"92","intVal":-42},
69+
{"id":"2","value2":"87","intVal":-37},
70+
{"id":"2","value2":"82","intVal":-32},
71+
{"id":"2","value2":"77","intVal":-27},
72+
{"id":"2","value2":"72","intVal":-22},
73+
{"id":"2","value2":"67","intVal":-17},
74+
{"id":"2","value2":"62","intVal":-12},
75+
{"id":"2","value2":"57","intVal":-7},
76+
{"id":"2","value2":"52","intVal":-2},
77+
{"id":"2","value2":"47","intVal":3},
78+
{"id":"2","value2":"42","intVal":8},
79+
{"id":"2","value2":"37","intVal":13},
80+
{"id":"2","value2":"32","intVal":18},
81+
{"id":"2","value2":"27","intVal":23},
82+
{"id":"2","value2":"22","intVal":28},
83+
{"id":"2","value2":"17","intVal":33},
84+
{"id":"2","value2":"12","intVal":38},
85+
{"id":"2","value2":"7","intVal":43},
86+
{"id":"2","value2":"2","intVal":48}]]></Documents>
87+
</Output>
88+
</Result>
89+
<Result>
90+
<Input>
91+
<Description>SELECT ORDER BY without ODE</Description>
92+
<Query><![CDATA[SELECT c.id, c.value2, c.intVal FROM c ORDER BY c.intVal]]></Query>
93+
<ODE>False</ODE>
94+
</Input>
95+
<Output>
96+
<Documents><![CDATA[{"id":"2","value2":"97","intVal":-47},
97+
{"id":"2","value2":"92","intVal":-42},
98+
{"id":"2","value2":"87","intVal":-37},
99+
{"id":"2","value2":"82","intVal":-32},
100+
{"id":"2","value2":"77","intVal":-27},
101+
{"id":"2","value2":"72","intVal":-22},
102+
{"id":"2","value2":"67","intVal":-17},
103+
{"id":"2","value2":"62","intVal":-12},
104+
{"id":"2","value2":"57","intVal":-7},
105+
{"id":"2","value2":"52","intVal":-2},
106+
{"id":"2","value2":"47","intVal":3},
107+
{"id":"2","value2":"42","intVal":8},
108+
{"id":"2","value2":"37","intVal":13},
109+
{"id":"2","value2":"32","intVal":18},
110+
{"id":"2","value2":"27","intVal":23},
111+
{"id":"2","value2":"22","intVal":28},
112+
{"id":"2","value2":"17","intVal":33},
113+
{"id":"2","value2":"12","intVal":38},
114+
{"id":"2","value2":"7","intVal":43},
115+
{"id":"2","value2":"2","intVal":48}]]></Documents>
116+
</Output>
117+
</Result>
60118
</Results>

‎Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/InMemoryContainer.cs

Lines changed: 25 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -52,11 +52,13 @@ internal class InMemoryContainer : IMonadicDocumentContainer
5252
private Dictionary<int, PartitionKeyHashRange> cachedPartitionKeyRangeIdToHashRange;
5353
private readonly bool createSplitForMultiHashAtSecondlevel;
5454
private readonly bool resolvePartitionsBasedOnPrefix;
55+
private readonly QueryRequestOptions queryRequestOptions;
5556

5657
public InMemoryContainer(
5758
PartitionKeyDefinition partitionKeyDefinition,
5859
bool createSplitForMultiHashAtSecondlevel = false,
59-
bool resolvePartitionsBasedOnPrefix = false)
60+
bool resolvePartitionsBasedOnPrefix = false,
61+
QueryRequestOptions queryRequestOptions = null)
6062
{
6163
this.partitionKeyDefinition = partitionKeyDefinition ?? throw new ArgumentNullException(nameof(partitionKeyDefinition));
6264
PartitionKeyHashRange fullRange = new PartitionKeyHashRange(startInclusive: null, endExclusive: new PartitionKeyHash(Cosmos.UInt128.MaxValue));
@@ -76,6 +78,7 @@ public InMemoryContainer(
7678
this.parentToChildMapping = new Dictionary<int, (int, int)>();
7779
this.createSplitForMultiHashAtSecondlevel = createSplitForMultiHashAtSecondlevel;
7880
this.resolvePartitionsBasedOnPrefix = resolvePartitionsBasedOnPrefix;
81+
this.queryRequestOptions = queryRequestOptions;
7982
}
8083

8184
public Task<TryCatch<List<FeedRangeEpk>>> MonadicGetFeedRangesAsync(
@@ -512,7 +515,7 @@ public virtual Task<TryCatch<QueryPage>> MonadicQueryAsync(
512515
}
513516

514517
List<CosmosObject> documents = new List<CosmosObject>();
515-
foreach (Record record in records.Where(r => IsRecordWithinFeedRange(r, feedRangeState.FeedRange, this.partitionKeyDefinition)))
518+
foreach (Record record in records.Where(r => IsRecordWithinFeedRange(r, feedRangeState.FeedRange, this.partitionKeyDefinition) && IsRecordWithinQueryPartition(r, this.queryRequestOptions, this.partitionKeyDefinition)))
516519
{
517520
CosmosObject document = ConvertRecordToCosmosElement(record);
518521
documents.Add(CosmosObject.Create(document));
@@ -716,6 +719,26 @@ public virtual Task<TryCatch<QueryPage>> MonadicQueryAsync(
716719
}
717720
}
718721

722+
private bool IsRecordWithinQueryPartition(Record record, QueryRequestOptions queryRequestOptions, PartitionKeyDefinition partitionKeyDefinition)
723+
{
724+
if(queryRequestOptions?.PartitionKey == null)
725+
{
726+
return true;
727+
}
728+
729+
IList<CosmosElement> partitionKey = GetPartitionKeysFromObjectModel(queryRequestOptions.PartitionKey.Value);
730+
IList<CosmosElement> partitionKeyFromRecord = GetPartitionKeysFromPayload(record.Payload, partitionKeyDefinition);
731+
if (partitionKeyDefinition.Kind == PartitionKind.MultiHash)
732+
{
733+
PartitionKeyHash partitionKeyHash = GetHashFromPartitionKeys(partitionKey, partitionKeyDefinition);
734+
PartitionKeyHash partitionKeyFromRecordHash = GetHashFromPartitionKeys(partitionKeyFromRecord, partitionKeyDefinition);
735+
736+
return partitionKeyHash.Equals(partitionKeyFromRecordHash) || partitionKeyFromRecordHash.Value.StartsWith(partitionKeyHash.Value);
737+
}
738+
739+
return partitionKey.SequenceEqual(partitionKeyFromRecord);
740+
}
741+
719742
public Task<TryCatch<ChangeFeedPage>> MonadicChangeFeedAsync(
720743
FeedRangeState<ChangeFeedState> feedRangeState,
721744
ChangeFeedPaginationOptions changeFeedPaginationOptions,

‎Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/SubpartitionTests.cs

Lines changed: 52 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -35,9 +35,12 @@ public void TestQueriesOnSplitContainer()
3535
{
3636
List<SubpartitionTestInput> inputs = new List<SubpartitionTestInput>
3737
{
38-
new SubpartitionTestInput("SELECT", query: @"SELECT c.id, c.value2 FROM c", ode: true),
39-
new SubpartitionTestInput("SELECT without ODE", query: @"SELECT c.id, c.value2 FROM c", ode: false),
38+
new SubpartitionTestInput(description: "SELECT", query: @"SELECT c.id, c.value2 FROM c", ode: true),
39+
new SubpartitionTestInput(description: "SELECT without ODE", query: @"SELECT c.id, c.value2 FROM c", ode: false),
40+
new SubpartitionTestInput(description: "SELECT ORDER BY with ODE", query: @"SELECT c.id, c.value2, c.intVal FROM c ORDER BY c.intVal", ode: true, sortResults: false),
41+
new SubpartitionTestInput(description: "SELECT ORDER BY without ODE", query: @"SELECT c.id, c.value2, c.intVal FROM c ORDER BY c.intVal", ode: false, sortResults: false),
4042
};
43+
4144
this.ExecuteTestSuite(inputs);
4245
}
4346

@@ -65,16 +68,19 @@ public async Task VerifyTestFrameworkSupportsPartitionSplit()
6568

6669
public override SubpartitionTestOutput ExecuteTest(SubpartitionTestInput input)
6770
{
68-
IMonadicDocumentContainer monadicDocumentContainer = CreateSplitDocumentContainerAsync(DocumentCount).Result;
71+
QueryRequestOptions queryRequestOptions = new QueryRequestOptions()
72+
{
73+
PartitionKey = new PartitionKeyBuilder().Add(SplitPartitionKey.ToString()).Build()
74+
};
75+
76+
IMonadicDocumentContainer monadicDocumentContainer = CreateSplitDocumentContainerAsync(DocumentCount, queryRequestOptions).Result;
6977
DocumentContainer documentContainer = new DocumentContainer(monadicDocumentContainer);
78+
TryCatch _ = monadicDocumentContainer.MonadicRefreshProviderAsync(NoOpTrace.Singleton, cancellationToken: default).Result;
79+
List<FeedRangeEpk> containerRanges = documentContainer.GetFeedRangesAsync(NoOpTrace.Singleton, cancellationToken: default).Result;
7080

7181
List<CosmosElement> documents = new List<CosmosElement>();
72-
QueryRequestOptions queryRequestOptions = new QueryRequestOptions()
73-
{
74-
PartitionKey = new PartitionKeyBuilder().Add(SplitPartitionKey.ToString()).Build()
75-
};
7682
(CosmosQueryExecutionContextFactory.InputParameters inputParameters, CosmosQueryContextCore cosmosQueryContextCore) =
77-
CreateInputParamsAndQueryContext(input, queryRequestOptions);
83+
CreateInputParamsAndQueryContext(input, queryRequestOptions, containerRanges);
7884
IQueryPipelineStage queryPipelineStage = CosmosQueryExecutionContextFactory.Create(
7985
documentContainer,
8086
cosmosQueryContextCore,
@@ -92,10 +98,10 @@ public override SubpartitionTestOutput ExecuteTest(SubpartitionTestInput input)
9298
documents.AddRange(tryGetPage.Result.Documents);
9399
}
94100

95-
return new SubpartitionTestOutput(documents);
101+
return new SubpartitionTestOutput(documents, input.SortResults);
96102
}
97103

98-
private static Tuple<CosmosQueryExecutionContextFactory.InputParameters, CosmosQueryContextCore> CreateInputParamsAndQueryContext(SubpartitionTestInput input, QueryRequestOptions queryRequestOptions)
104+
private static Tuple<CosmosQueryExecutionContextFactory.InputParameters, CosmosQueryContextCore> CreateInputParamsAndQueryContext(SubpartitionTestInput input, QueryRequestOptions queryRequestOptions, IReadOnlyList<FeedRangeEpk> containerRanges)
99105
{
100106
string query = input.Query;
101107
CosmosElement continuationToken = null;
@@ -134,10 +140,20 @@ public override SubpartitionTestOutput ExecuteTest(SubpartitionTestInput input)
134140
isNonStreamingOrderByQueryFeatureDisabled: queryRequestOptions.IsNonStreamingOrderByQueryFeatureDisabled,
135141
testInjections: queryRequestOptions.TestSettings);
136142

143+
List<PartitionKeyRange> targetPkRanges = new();
144+
foreach (FeedRangeEpk feedRangeEpk in containerRanges)
145+
{
146+
targetPkRanges.Add(new PartitionKeyRange
147+
{
148+
MinInclusive = feedRangeEpk.Range.Min,
149+
MaxExclusive = feedRangeEpk.Range.Max,
150+
});
151+
}
152+
137153
string databaseId = "db1234";
138154
string resourceLink = $"dbs/{databaseId}/colls";
139155
CosmosQueryContextCore cosmosQueryContextCore = new CosmosQueryContextCore(
140-
client: new TestCosmosQueryClient(queryPartitionProvider),
156+
client: new TestCosmosQueryClient(queryPartitionProvider, targetPkRanges),
141157
resourceTypeEnum: Documents.ResourceType.Document,
142158
operationType: Documents.OperationType.Query,
143159
resourceType: typeof(QueryResponseCore),
@@ -215,20 +231,20 @@ internal static PartitionKeyDefinition CreatePartitionKeyDefinition()
215231
return partitionKeyDefinition;
216232
}
217233

218-
private static async Task<IDocumentContainer> CreateSplitDocumentContainerAsync(int numItems)
234+
private static async Task<IDocumentContainer> CreateSplitDocumentContainerAsync(int numItems, QueryRequestOptions queryRequestOptions)
219235
{
220236
PartitionKeyDefinition partitionKeyDefinition = CreatePartitionKeyDefinition();
221-
InMemoryContainer inMemoryContainer = await CreateSplitInMemoryDocumentContainerAsync(numItems, partitionKeyDefinition);
237+
InMemoryContainer inMemoryContainer = await CreateSplitInMemoryDocumentContainerAsync(numItems, partitionKeyDefinition, queryRequestOptions);
222238
DocumentContainer documentContainer = new DocumentContainer(inMemoryContainer);
223239
return documentContainer;
224240
}
225241

226-
private static async Task<InMemoryContainer> CreateSplitInMemoryDocumentContainerAsync(int numItems, PartitionKeyDefinition partitionKeyDefinition)
242+
private static async Task<InMemoryContainer> CreateSplitInMemoryDocumentContainerAsync(int numItems, PartitionKeyDefinition partitionKeyDefinition, QueryRequestOptions queryRequestOptions = null)
227243
{
228-
InMemoryContainer inMemoryContainer = new InMemoryContainer(partitionKeyDefinition, createSplitForMultiHashAtSecondlevel: true, resolvePartitionsBasedOnPrefix: true);
244+
InMemoryContainer inMemoryContainer = new InMemoryContainer(partitionKeyDefinition, createSplitForMultiHashAtSecondlevel: true, resolvePartitionsBasedOnPrefix: true, queryRequestOptions: queryRequestOptions);
229245
for (int i = 0; i < numItems; i++)
230246
{
231-
CosmosObject item = CosmosObject.Parse($"{{\"id\" : \"{i % 5}\", \"value1\" : \"{Guid.NewGuid()}\", \"value2\" : \"{i}\" }}");
247+
CosmosObject item = CosmosObject.Parse($"{{\"id\" : \"{i % 5}\", \"value1\" : \"{Guid.NewGuid()}\", \"value2\" : \"{i}\", \"intVal\" : {(numItems/2) - i} }}");
232248
while (true)
233249
{
234250
TryCatch<Record> monadicCreateRecord = await inMemoryContainer.MonadicCreateItemAsync(item, cancellationToken: default);
@@ -243,13 +259,16 @@ private static async Task<InMemoryContainer> CreateSplitInMemoryDocumentContaine
243259

244260
return inMemoryContainer;
245261
}
262+
246263
internal class TestCosmosQueryClient : CosmosQueryClient
247264
{
248265
private readonly QueryPartitionProvider queryPartitionProvider;
266+
private readonly IReadOnlyList<PartitionKeyRange> targetPartitionKeyRanges;
249267

250-
public TestCosmosQueryClient(QueryPartitionProvider queryPartitionProvider)
268+
public TestCosmosQueryClient(QueryPartitionProvider queryPartitionProvider, IEnumerable<PartitionKeyRange> targetPartitionKeyRanges)
251269
{
252270
this.queryPartitionProvider = queryPartitionProvider;
271+
this.targetPartitionKeyRanges = targetPartitionKeyRanges.ToList();
253272
}
254273

255274
public override Action<IQueryable> OnExecuteScalarQueryCallback => throw new NotImplementedException();
@@ -322,14 +341,7 @@ public override Task<List<PartitionKeyRange>> GetTargetPartitionKeyRangeByFeedRa
322341

323342
public override Task<List<PartitionKeyRange>> GetTargetPartitionKeyRangesAsync(string resourceLink, string collectionResourceId, IReadOnlyList<Documents.Routing.Range<string>> providedRanges, bool forceRefresh, ITrace trace)
324343
{
325-
return Task.FromResult(new List<PartitionKeyRange>
326-
{
327-
new PartitionKeyRange()
328-
{
329-
MinInclusive = Documents.Routing.PartitionKeyInternal.MinimumInclusiveEffectivePartitionKey,
330-
MaxExclusive = Documents.Routing.PartitionKeyInternal.MaximumExclusiveEffectivePartitionKey
331-
}
332-
});
344+
return Task.FromResult(this.targetPartitionKeyRanges.ToList());
333345
}
334346

335347
public override Task<IReadOnlyList<PartitionKeyRange>> TryGetOverlappingRangesAsync(string collectionResourceId, Documents.Routing.Range<string> range, bool forceRefresh = false)
@@ -351,17 +363,20 @@ public override async Task<TryCatch<PartitionedQueryExecutionInfo>> TryGetPartit
351363

352364
public class SubpartitionTestInput : BaselineTestInput
353365
{
354-
public SubpartitionTestInput(string description, string query, bool ode)
366+
public SubpartitionTestInput(string description, string query, bool ode, bool sortResults = true)
355367
:base(description)
356368
{
357369
this.Query = query;
358370
this.ODE = ode;
371+
this.SortResults = sortResults;
359372
}
360373

361374
internal string Query { get; }
362375

363376
internal bool ODE { get; }
364377

378+
internal bool SortResults { get; }
379+
365380
public override void SerializeAsXml(XmlWriter xmlWriter)
366381
{
367382
xmlWriter.WriteElementString("Description", this.Description);
@@ -375,17 +390,25 @@ public override void SerializeAsXml(XmlWriter xmlWriter)
375390
public class SubpartitionTestOutput : BaselineTestOutput
376391
{
377392
private readonly List<CosmosElement> documents;
393+
private readonly bool sortResults;
378394

379-
internal SubpartitionTestOutput(IReadOnlyList<CosmosElement> documents)
395+
internal SubpartitionTestOutput(IReadOnlyList<CosmosElement> documents, bool sortResults)
380396
{
381397
this.documents = documents.ToList();
398+
this.sortResults = sortResults;
382399
}
383400

384401
public override void SerializeAsXml(XmlWriter xmlWriter)
385402
{
386403
xmlWriter.WriteStartElement("Documents");
387-
string content = string.Join($",{Environment.NewLine}",
388-
this.documents.Select(doc => doc.ToString()).OrderBy(serializedDoc => serializedDoc));
404+
405+
IEnumerable<string> lines = this.documents.Select(doc => doc.ToString());
406+
if(this.sortResults)
407+
{
408+
lines = lines.OrderBy(serializedDoc => serializedDoc);
409+
}
410+
411+
string content = string.Join($",{Environment.NewLine}", lines);
389412
xmlWriter.WriteCData(content);
390413
xmlWriter.WriteEndElement();
391414
}

0 commit comments

Comments
 (0)