diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs index 042c8a7606..ccc8168983 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs @@ -248,6 +248,8 @@ private static async Task> TryCreateFromPartitione inputParameters.InitialFeedRange, trace); + Debug.Assert(targetRanges != null, $"{nameof(CosmosQueryExecutionContextFactory)} Assert!", "targetRanges != null"); + TryCatch tryCreatePipelineStage; Documents.PartitionKeyRange targetRange = await TryGetTargetRangeOptimisticDirectExecutionAsync( inputParameters, @@ -270,7 +272,7 @@ private static async Task> TryCreateFromPartitione } else { - bool singleLogicalPartitionKeyQuery = inputParameters.PartitionKey.HasValue + bool singleLogicalPartitionKeyQuery = (inputParameters.PartitionKey.HasValue && targetRanges.Count == 1) || ((partitionedQueryExecutionInfo.QueryRanges.Count == 1) && partitionedQueryExecutionInfo.QueryRanges[0].IsSingleValue); bool serverStreamingQuery = !partitionedQueryExecutionInfo.QueryInfo.HasAggregates diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/BaselineTest/TestBaseline/SubpartitionTests.TestQueriesOnSplitContainer.xml b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/BaselineTest/TestBaseline/SubpartitionTests.TestQueriesOnSplitContainer.xml index 192bfd9a20..c21f7e90d7 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/BaselineTest/TestBaseline/SubpartitionTests.TestQueriesOnSplitContainer.xml +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/BaselineTest/TestBaseline/SubpartitionTests.TestQueriesOnSplitContainer.xml @@ -57,4 +57,62 @@ {"id":"2","value2":"97"}]]> + + + SELECT ORDER BY with ODE + + True + + + + + + + + SELECT ORDER BY without ODE + + False + + + + + \ No newline at end of file diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/InMemoryContainer.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/InMemoryContainer.cs index 4c0a62b72d..f126b0970d 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/InMemoryContainer.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/InMemoryContainer.cs @@ -52,11 +52,13 @@ internal class InMemoryContainer : IMonadicDocumentContainer private Dictionary cachedPartitionKeyRangeIdToHashRange; private readonly bool createSplitForMultiHashAtSecondlevel; private readonly bool resolvePartitionsBasedOnPrefix; + private readonly QueryRequestOptions queryRequestOptions; public InMemoryContainer( PartitionKeyDefinition partitionKeyDefinition, bool createSplitForMultiHashAtSecondlevel = false, - bool resolvePartitionsBasedOnPrefix = false) + bool resolvePartitionsBasedOnPrefix = false, + QueryRequestOptions queryRequestOptions = null) { this.partitionKeyDefinition = partitionKeyDefinition ?? throw new ArgumentNullException(nameof(partitionKeyDefinition)); PartitionKeyHashRange fullRange = new PartitionKeyHashRange(startInclusive: null, endExclusive: new PartitionKeyHash(Cosmos.UInt128.MaxValue)); @@ -76,6 +78,7 @@ public InMemoryContainer( this.parentToChildMapping = new Dictionary(); this.createSplitForMultiHashAtSecondlevel = createSplitForMultiHashAtSecondlevel; this.resolvePartitionsBasedOnPrefix = resolvePartitionsBasedOnPrefix; + this.queryRequestOptions = queryRequestOptions; } public Task>> MonadicGetFeedRangesAsync( @@ -512,7 +515,7 @@ public virtual Task> MonadicQueryAsync( } List documents = new List(); - foreach (Record record in records.Where(r => IsRecordWithinFeedRange(r, feedRangeState.FeedRange, this.partitionKeyDefinition))) + foreach (Record record in records.Where(r => IsRecordWithinFeedRange(r, feedRangeState.FeedRange, this.partitionKeyDefinition) && IsRecordWithinQueryPartition(r, this.queryRequestOptions, this.partitionKeyDefinition))) { CosmosObject document = ConvertRecordToCosmosElement(record); documents.Add(CosmosObject.Create(document)); @@ -716,6 +719,26 @@ public virtual Task> MonadicQueryAsync( } } + private bool IsRecordWithinQueryPartition(Record record, QueryRequestOptions queryRequestOptions, PartitionKeyDefinition partitionKeyDefinition) + { + if(queryRequestOptions?.PartitionKey == null) + { + return true; + } + + IList partitionKey = GetPartitionKeysFromObjectModel(queryRequestOptions.PartitionKey.Value); + IList partitionKeyFromRecord = GetPartitionKeysFromPayload(record.Payload, partitionKeyDefinition); + if (partitionKeyDefinition.Kind == PartitionKind.MultiHash) + { + PartitionKeyHash partitionKeyHash = GetHashFromPartitionKeys(partitionKey, partitionKeyDefinition); + PartitionKeyHash partitionKeyFromRecordHash = GetHashFromPartitionKeys(partitionKeyFromRecord, partitionKeyDefinition); + + return partitionKeyHash.Equals(partitionKeyFromRecordHash) || partitionKeyFromRecordHash.Value.StartsWith(partitionKeyHash.Value); + } + + return partitionKey.SequenceEqual(partitionKeyFromRecord); + } + public Task> MonadicChangeFeedAsync( FeedRangeState feedRangeState, ChangeFeedPaginationOptions changeFeedPaginationOptions, diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/SubpartitionTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/SubpartitionTests.cs index 8253e973e0..1605c09a1a 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/SubpartitionTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/SubpartitionTests.cs @@ -35,9 +35,12 @@ public void TestQueriesOnSplitContainer() { List inputs = new List { - new SubpartitionTestInput("SELECT", query: @"SELECT c.id, c.value2 FROM c", ode: true), - new SubpartitionTestInput("SELECT without ODE", query: @"SELECT c.id, c.value2 FROM c", ode: false), + new SubpartitionTestInput(description: "SELECT", query: @"SELECT c.id, c.value2 FROM c", ode: true), + new SubpartitionTestInput(description: "SELECT without ODE", query: @"SELECT c.id, c.value2 FROM c", ode: false), + 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), + 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), }; + this.ExecuteTestSuite(inputs); } @@ -65,16 +68,19 @@ public async Task VerifyTestFrameworkSupportsPartitionSplit() public override SubpartitionTestOutput ExecuteTest(SubpartitionTestInput input) { - IMonadicDocumentContainer monadicDocumentContainer = CreateSplitDocumentContainerAsync(DocumentCount).Result; + QueryRequestOptions queryRequestOptions = new QueryRequestOptions() + { + PartitionKey = new PartitionKeyBuilder().Add(SplitPartitionKey.ToString()).Build() + }; + + IMonadicDocumentContainer monadicDocumentContainer = CreateSplitDocumentContainerAsync(DocumentCount, queryRequestOptions).Result; DocumentContainer documentContainer = new DocumentContainer(monadicDocumentContainer); + TryCatch _ = monadicDocumentContainer.MonadicRefreshProviderAsync(NoOpTrace.Singleton, cancellationToken: default).Result; + List containerRanges = documentContainer.GetFeedRangesAsync(NoOpTrace.Singleton, cancellationToken: default).Result; List documents = new List(); - QueryRequestOptions queryRequestOptions = new QueryRequestOptions() - { - PartitionKey = new PartitionKeyBuilder().Add(SplitPartitionKey.ToString()).Build() - }; (CosmosQueryExecutionContextFactory.InputParameters inputParameters, CosmosQueryContextCore cosmosQueryContextCore) = - CreateInputParamsAndQueryContext(input, queryRequestOptions); + CreateInputParamsAndQueryContext(input, queryRequestOptions, containerRanges); IQueryPipelineStage queryPipelineStage = CosmosQueryExecutionContextFactory.Create( documentContainer, cosmosQueryContextCore, @@ -92,10 +98,10 @@ public override SubpartitionTestOutput ExecuteTest(SubpartitionTestInput input) documents.AddRange(tryGetPage.Result.Documents); } - return new SubpartitionTestOutput(documents); + return new SubpartitionTestOutput(documents, input.SortResults); } - private static Tuple CreateInputParamsAndQueryContext(SubpartitionTestInput input, QueryRequestOptions queryRequestOptions) + private static Tuple CreateInputParamsAndQueryContext(SubpartitionTestInput input, QueryRequestOptions queryRequestOptions, IReadOnlyList containerRanges) { string query = input.Query; CosmosElement continuationToken = null; @@ -134,10 +140,20 @@ public override SubpartitionTestOutput ExecuteTest(SubpartitionTestInput input) isNonStreamingOrderByQueryFeatureDisabled: queryRequestOptions.IsNonStreamingOrderByQueryFeatureDisabled, testInjections: queryRequestOptions.TestSettings); + List targetPkRanges = new(); + foreach (FeedRangeEpk feedRangeEpk in containerRanges) + { + targetPkRanges.Add(new PartitionKeyRange + { + MinInclusive = feedRangeEpk.Range.Min, + MaxExclusive = feedRangeEpk.Range.Max, + }); + } + string databaseId = "db1234"; string resourceLink = $"dbs/{databaseId}/colls"; CosmosQueryContextCore cosmosQueryContextCore = new CosmosQueryContextCore( - client: new TestCosmosQueryClient(queryPartitionProvider), + client: new TestCosmosQueryClient(queryPartitionProvider, targetPkRanges), resourceTypeEnum: Documents.ResourceType.Document, operationType: Documents.OperationType.Query, resourceType: typeof(QueryResponseCore), @@ -215,20 +231,20 @@ internal static PartitionKeyDefinition CreatePartitionKeyDefinition() return partitionKeyDefinition; } - private static async Task CreateSplitDocumentContainerAsync(int numItems) + private static async Task CreateSplitDocumentContainerAsync(int numItems, QueryRequestOptions queryRequestOptions) { PartitionKeyDefinition partitionKeyDefinition = CreatePartitionKeyDefinition(); - InMemoryContainer inMemoryContainer = await CreateSplitInMemoryDocumentContainerAsync(numItems, partitionKeyDefinition); + InMemoryContainer inMemoryContainer = await CreateSplitInMemoryDocumentContainerAsync(numItems, partitionKeyDefinition, queryRequestOptions); DocumentContainer documentContainer = new DocumentContainer(inMemoryContainer); return documentContainer; } - private static async Task CreateSplitInMemoryDocumentContainerAsync(int numItems, PartitionKeyDefinition partitionKeyDefinition) + private static async Task CreateSplitInMemoryDocumentContainerAsync(int numItems, PartitionKeyDefinition partitionKeyDefinition, QueryRequestOptions queryRequestOptions = null) { - InMemoryContainer inMemoryContainer = new InMemoryContainer(partitionKeyDefinition, createSplitForMultiHashAtSecondlevel: true, resolvePartitionsBasedOnPrefix: true); + InMemoryContainer inMemoryContainer = new InMemoryContainer(partitionKeyDefinition, createSplitForMultiHashAtSecondlevel: true, resolvePartitionsBasedOnPrefix: true, queryRequestOptions: queryRequestOptions); for (int i = 0; i < numItems; i++) { - CosmosObject item = CosmosObject.Parse($"{{\"id\" : \"{i % 5}\", \"value1\" : \"{Guid.NewGuid()}\", \"value2\" : \"{i}\" }}"); + CosmosObject item = CosmosObject.Parse($"{{\"id\" : \"{i % 5}\", \"value1\" : \"{Guid.NewGuid()}\", \"value2\" : \"{i}\", \"intVal\" : {(numItems/2) - i} }}"); while (true) { TryCatch monadicCreateRecord = await inMemoryContainer.MonadicCreateItemAsync(item, cancellationToken: default); @@ -243,13 +259,16 @@ private static async Task CreateSplitInMemoryDocumentContaine return inMemoryContainer; } + internal class TestCosmosQueryClient : CosmosQueryClient { private readonly QueryPartitionProvider queryPartitionProvider; + private readonly IReadOnlyList targetPartitionKeyRanges; - public TestCosmosQueryClient(QueryPartitionProvider queryPartitionProvider) + public TestCosmosQueryClient(QueryPartitionProvider queryPartitionProvider, IEnumerable targetPartitionKeyRanges) { this.queryPartitionProvider = queryPartitionProvider; + this.targetPartitionKeyRanges = targetPartitionKeyRanges.ToList(); } public override Action OnExecuteScalarQueryCallback => throw new NotImplementedException(); @@ -322,14 +341,7 @@ public override Task> GetTargetPartitionKeyRangeByFeedRa public override Task> GetTargetPartitionKeyRangesAsync(string resourceLink, string collectionResourceId, IReadOnlyList> providedRanges, bool forceRefresh, ITrace trace) { - return Task.FromResult(new List - { - new PartitionKeyRange() - { - MinInclusive = Documents.Routing.PartitionKeyInternal.MinimumInclusiveEffectivePartitionKey, - MaxExclusive = Documents.Routing.PartitionKeyInternal.MaximumExclusiveEffectivePartitionKey - } - }); + return Task.FromResult(this.targetPartitionKeyRanges.ToList()); } public override Task> TryGetOverlappingRangesAsync(string collectionResourceId, Documents.Routing.Range range, bool forceRefresh = false) @@ -351,17 +363,20 @@ public override async Task> TryGetPartit public class SubpartitionTestInput : BaselineTestInput { - public SubpartitionTestInput(string description, string query, bool ode) + public SubpartitionTestInput(string description, string query, bool ode, bool sortResults = true) :base(description) { this.Query = query; this.ODE = ode; + this.SortResults = sortResults; } internal string Query { get; } internal bool ODE { get; } + internal bool SortResults { get; } + public override void SerializeAsXml(XmlWriter xmlWriter) { xmlWriter.WriteElementString("Description", this.Description); @@ -375,17 +390,25 @@ public override void SerializeAsXml(XmlWriter xmlWriter) public class SubpartitionTestOutput : BaselineTestOutput { private readonly List documents; + private readonly bool sortResults; - internal SubpartitionTestOutput(IReadOnlyList documents) + internal SubpartitionTestOutput(IReadOnlyList documents, bool sortResults) { this.documents = documents.ToList(); + this.sortResults = sortResults; } public override void SerializeAsXml(XmlWriter xmlWriter) { xmlWriter.WriteStartElement("Documents"); - string content = string.Join($",{Environment.NewLine}", - this.documents.Select(doc => doc.ToString()).OrderBy(serializedDoc => serializedDoc)); + + IEnumerable lines = this.documents.Select(doc => doc.ToString()); + if(this.sortResults) + { + lines = lines.OrderBy(serializedDoc => serializedDoc); + } + + string content = string.Join($",{Environment.NewLine}", lines); xmlWriter.WriteCData(content); xmlWriter.WriteEndElement(); }