diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs index e7a65b7457..dd6b2c45c7 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs @@ -138,7 +138,7 @@ private static async Task> TryCreateCoreContextAsy cosmosQueryContext.ResourceLink, inputParameters.PartitionKey, createQueryPipelineTrace, - cancellationToken); + cancellationToken); cosmosQueryContext.ContainerResourceId = containerQueryProperties.ResourceId; Documents.PartitionKeyRange targetRange = await TryGetTargetRangeOptimisticDirectExecutionAsync( @@ -205,6 +205,7 @@ private static async Task> TryCreateCoreContextAsy documentContainer, inputParameters, targetRanges, + containerQueryProperties, cancellationToken); } } @@ -295,11 +296,12 @@ private static async Task> TryCreateFromPartitione documentContainer, inputParameters, targetRanges, + containerQueryProperties, cancellationToken); } else { - tryCreatePipelineStage = TryCreateSpecializedDocumentQueryExecutionContext(documentContainer, cosmosQueryContext, inputParameters, targetRanges, partitionedQueryExecutionInfo, cancellationToken); + tryCreatePipelineStage = TryCreateSpecializedDocumentQueryExecutionContext(documentContainer, cosmosQueryContext, inputParameters, targetRanges, containerQueryProperties, partitionedQueryExecutionInfo, cancellationToken); } } @@ -359,6 +361,7 @@ private static async Task> TryCreateSinglePartitio cosmosQueryContext, inputParameters, targetRanges, + containerQueryProperties, partitionedQueryExecutionInfo, cancellationToken); } @@ -382,6 +385,7 @@ private static TryCatch TryCreateSpecializedDocumentQueryEx CosmosQueryContext cosmosQueryContext, InputParameters inputParameters, List targetRanges, + ContainerQueryProperties containerQueryProperties, PartitionedQueryExecutionInfo partitionedQueryExecutionInfo, CancellationToken cancellationToken) { @@ -418,6 +422,7 @@ private static TryCatch TryCreateSpecializedDocumentQueryEx inputParameters, partitionedQueryExecutionInfo, targetRanges, + containerQueryProperties, cancellationToken); } @@ -450,6 +455,7 @@ private static async Task> TryCreateSpecializedDoc cosmosQueryContext, inputParameters, targetRanges, + containerQueryProperties, partitionedQueryExecutionInfo, cancellationToken); } @@ -467,6 +473,7 @@ private static TryCatch TryCreateOptimisticDirectExecutionC documentContainer: documentContainer, inputParameters: inputParameters, targetRange: new FeedRangeEpk(targetRange.ToRange()), + containerQueryProperties: containerQueryProperties, fallbackQueryPipelineStageFactory: (continuationToken) => { // In fallback scenario, the Specialized pipeline is always invoked @@ -488,6 +495,7 @@ private static TryCatch TryCreatePassthroughQueryExecutionC DocumentContainer documentContainer, InputParameters inputParameters, List targetRanges, + ContainerQueryProperties containerQueryProperties, CancellationToken cancellationToken) { // Return a parallel context, since we still want to be able to handle splits and concurrency / buffering. @@ -505,8 +513,9 @@ private static TryCatch TryCreatePassthroughQueryExecutionC queryPaginationOptions: new QueryPaginationOptions( pageSizeHint: inputParameters.MaxItemCount), partitionKey: inputParameters.PartitionKey, - prefetchPolicy: PrefetchPolicy.PrefetchSinglePage, + containerQueryProperties: containerQueryProperties, maxConcurrency: inputParameters.MaxConcurrency, + prefetchPolicy: PrefetchPolicy.PrefetchSinglePage, cancellationToken: cancellationToken, continuationToken: inputParameters.InitialUserContinuationToken); } @@ -516,7 +525,8 @@ private static TryCatch TryCreateSpecializedDocumentQueryEx CosmosQueryContext cosmosQueryContext, InputParameters inputParameters, PartitionedQueryExecutionInfo partitionedQueryExecutionInfo, - List targetRanges, + List targetRanges, + ContainerQueryProperties containerQueryProperties, CancellationToken cancellationToken) { QueryInfo queryInfo = partitionedQueryExecutionInfo.QueryInfo; @@ -576,6 +586,7 @@ private static TryCatch TryCreateSpecializedDocumentQueryEx queryInfo: partitionedQueryExecutionInfo.QueryInfo, queryPaginationOptions: new QueryPaginationOptions( pageSizeHint: (int)optimalPageSize), + containerQueryProperties: containerQueryProperties, maxConcurrency: inputParameters.MaxConcurrency, requestContinuationToken: inputParameters.InitialUserContinuationToken, requestCancellationToken: cancellationToken); diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/Parallel/ParallelCrossPartitionQueryPipelineStage.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/Parallel/ParallelCrossPartitionQueryPipelineStage.cs index b84a36870e..a471f22bb1 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/Parallel/ParallelCrossPartitionQueryPipelineStage.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/Parallel/ParallelCrossPartitionQueryPipelineStage.cs @@ -14,6 +14,7 @@ namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.CrossPartition.Parallel using Microsoft.Azure.Cosmos.Query.Core.Exceptions; using Microsoft.Azure.Cosmos.Query.Core.Monads; using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Pagination; + using Microsoft.Azure.Cosmos.Query.Core.QueryClient; using Microsoft.Azure.Cosmos.Tracing; using static Microsoft.Azure.Cosmos.Query.Core.Pipeline.CrossPartition.PartitionMapper; @@ -137,6 +138,7 @@ public static TryCatch MonadicCreate( IReadOnlyList targetRanges, Cosmos.PartitionKey? partitionKey, QueryPaginationOptions queryPaginationOptions, + ContainerQueryProperties containerQueryProperties, int maxConcurrency, PrefetchPolicy prefetchPolicy, CosmosElement continuationToken, @@ -162,7 +164,7 @@ public static TryCatch MonadicCreate( CrossPartitionRangePageAsyncEnumerator crossPartitionPageEnumerator = new CrossPartitionRangePageAsyncEnumerator( feedRangeProvider: documentContainer, - createPartitionRangeEnumerator: ParallelCrossPartitionQueryPipelineStage.MakeCreateFunction(documentContainer, sqlQuerySpec, queryPaginationOptions, partitionKey, cancellationToken), + createPartitionRangeEnumerator: ParallelCrossPartitionQueryPipelineStage.MakeCreateFunction(documentContainer, sqlQuerySpec, queryPaginationOptions, partitionKey, containerQueryProperties, cancellationToken), comparer: Comparer.Singleton, maxConcurrency: maxConcurrency, prefetchPolicy: prefetchPolicy, @@ -249,12 +251,14 @@ private static CreatePartitionRangePageAsyncEnumerator Ma SqlQuerySpec sqlQuerySpec, QueryPaginationOptions queryPaginationOptions, Cosmos.PartitionKey? partitionKey, + ContainerQueryProperties containerQueryProperties, CancellationToken cancellationToken) => (FeedRangeState feedRangeState) => new QueryPartitionRangePageAsyncEnumerator( queryDataSource, sqlQuerySpec, feedRangeState, partitionKey, queryPaginationOptions, + containerQueryProperties, cancellationToken); public void SetCancellationToken(CancellationToken cancellationToken) diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/Parallel/QueryPartitionRangePageAsyncEnumerator.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/Parallel/QueryPartitionRangePageAsyncEnumerator.cs index e2918a366e..bc38d91570 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/Parallel/QueryPartitionRangePageAsyncEnumerator.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/Parallel/QueryPartitionRangePageAsyncEnumerator.cs @@ -10,6 +10,7 @@ namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.CrossPartition.Parallel using Microsoft.Azure.Cosmos.Pagination; using Microsoft.Azure.Cosmos.Query.Core.Monads; using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Pagination; + using Microsoft.Azure.Cosmos.Query.Core.QueryClient; using Microsoft.Azure.Cosmos.Tracing; internal sealed class QueryPartitionRangePageAsyncEnumerator : PartitionRangePageAsyncEnumerator @@ -17,6 +18,7 @@ internal sealed class QueryPartitionRangePageAsyncEnumerator : PartitionRangePag private readonly IQueryDataSource queryDataSource; private readonly SqlQuerySpec sqlQuerySpec; private readonly QueryPaginationOptions queryPaginationOptions; + private readonly ContainerQueryProperties containerQueryProperties; private readonly Cosmos.PartitionKey? partitionKey; public QueryPartitionRangePageAsyncEnumerator( @@ -25,6 +27,7 @@ public QueryPartitionRangePageAsyncEnumerator( FeedRangeState feedRangeState, Cosmos.PartitionKey? partitionKey, QueryPaginationOptions queryPaginationOptions, + ContainerQueryProperties containerQueryProperties, CancellationToken cancellationToken) : base(feedRangeState, cancellationToken) { @@ -32,6 +35,7 @@ public QueryPartitionRangePageAsyncEnumerator( this.sqlQuerySpec = sqlQuerySpec ?? throw new ArgumentNullException(nameof(sqlQuerySpec)); this.queryPaginationOptions = queryPaginationOptions; this.partitionKey = partitionKey; + this.containerQueryProperties = containerQueryProperties; } public override ValueTask DisposeAsync() => default; @@ -43,9 +47,7 @@ protected override Task> GetNextPageAsync(ITrace trace, Canc throw new ArgumentNullException(nameof(trace)); } - // We sadly need to check the partition key, since a user can set a partition key in the request options with a different continuation token. - // In the future the partition filtering and continuation information needs to be a tightly bounded contract (like cross feed range state). - FeedRangeInternal feedRange = this.partitionKey.HasValue ? new FeedRangePartitionKey(this.partitionKey.Value) : this.FeedRangeState.FeedRange; + FeedRangeInternal feedRange = this.LimitFeedRangeToSinglePartition(); return this.queryDataSource.MonadicQueryAsync( sqlQuerySpec: this.sqlQuerySpec, feedRangeState: new FeedRangeState(feedRange, this.FeedRangeState.State), @@ -53,5 +55,84 @@ protected override Task> GetNextPageAsync(ITrace trace, Canc trace: trace, cancellationToken); } + + /// + /// Updates the FeedRange to limit the scope of this enumerator to single physical partition. + /// Generally speaking, a subpartitioned container can experience split partition at any level of hierarchical partition key. + /// This could cause a situation where more than one physical partition contains the data for a partial partition key. + /// Currently, enumerator instantiation does not honor physical partition boundary and allocates entire epk range which could spans across multiple physical partitions to the enumerator. + /// Since such an epk range does not exist at the container level, Service generates a GoneException. + /// This method restrics the range of each container by shrinking the ends of the range so that they do not span across physical partition. + /// + private FeedRangeInternal LimitFeedRangeToSinglePartition() + { + // We sadly need to check the partition key, since a user can set a partition key in the request options with a different continuation token. + // In the future the partition filtering and continuation information needs to be a tightly bounded contract (like cross feed range state). + FeedRangeInternal feedRange = this.FeedRangeState.FeedRange; + + if (feedRange is FeedRangeEpk feedRangeEpk && this.partitionKey.HasValue) + { + if (this.containerQueryProperties.EffectiveRangesForPartitionKey == null || + this.containerQueryProperties.EffectiveRangesForPartitionKey.Count == 0) + { + throw new InvalidOperationException( + "EffectiveRangesForPartitionKey should be populated when PK is specified in request options."); + } + + foreach (Documents.Routing.Range epkForPartitionKey in + this.containerQueryProperties.EffectiveRangesForPartitionKey) + { + if (Documents.Routing.Range.CheckOverlapping( + feedRangeEpk.Range, + epkForPartitionKey)) + { + if (!feedRangeEpk.Range.Equals(epkForPartitionKey)) + { + String overlappingMin; + bool minInclusive; + String overlappingMax; + bool maxInclusive; + + if (Documents.Routing.Range.MinComparer.Instance.Compare( + epkForPartitionKey, + feedRangeEpk.Range) < 0) + { + overlappingMin = feedRangeEpk.Range.Min; + minInclusive = feedRangeEpk.Range.IsMinInclusive; + } + else + { + overlappingMin = epkForPartitionKey.Min; + minInclusive = epkForPartitionKey.IsMinInclusive; + } + + if (Documents.Routing.Range.MaxComparer.Instance.Compare( + epkForPartitionKey, + feedRangeEpk.Range) > 0) + { + overlappingMax = feedRangeEpk.Range.Max; + maxInclusive = feedRangeEpk.Range.IsMaxInclusive; + } + else + { + overlappingMax = epkForPartitionKey.Max; + maxInclusive = epkForPartitionKey.IsMaxInclusive; + } + + feedRange = new FeedRangeEpk( + new Documents.Routing.Range( + overlappingMin, + overlappingMax, + minInclusive, + maxInclusive)); + } + + break; + } + } + } + + return feedRange; + } } } \ No newline at end of file diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/OptimisticDirectExecution/OptimisticDirectExecutionQueryPipelineStage.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/OptimisticDirectExecution/OptimisticDirectExecutionQueryPipelineStage.cs index 7e571d0a41..94e42b09fd 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/OptimisticDirectExecution/OptimisticDirectExecutionQueryPipelineStage.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/OptimisticDirectExecution/OptimisticDirectExecutionQueryPipelineStage.cs @@ -18,6 +18,7 @@ namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.OptimisticDirectExecutionQu using Microsoft.Azure.Cosmos.Query.Core.Pipeline.CrossPartition; using Microsoft.Azure.Cosmos.Query.Core.Pipeline.CrossPartition.Parallel; using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Pagination; + using Microsoft.Azure.Cosmos.Query.Core.QueryClient; using Microsoft.Azure.Cosmos.Tracing; using Microsoft.Azure.Documents; @@ -142,6 +143,7 @@ public static TryCatch MonadicCreate( DocumentContainer documentContainer, CosmosQueryExecutionContextFactory.InputParameters inputParameters, FeedRangeEpk targetRange, + ContainerQueryProperties containerQueryProperties, FallbackQueryPipelineStageFactory fallbackQueryPipelineStageFactory, CancellationToken cancellationToken) { @@ -152,6 +154,7 @@ public static TryCatch MonadicCreate( targetRange: targetRange, queryPaginationOptions: paginationOptions, partitionKey: inputParameters.PartitionKey, + containerQueryProperties: containerQueryProperties, continuationToken: inputParameters.InitialUserContinuationToken, cancellationToken: cancellationToken); @@ -247,6 +250,7 @@ public static TryCatch MonadicCreate( FeedRangeEpk targetRange, Cosmos.PartitionKey? partitionKey, QueryPaginationOptions queryPaginationOptions, + ContainerQueryProperties containerQueryProperties, CosmosElement continuationToken, CancellationToken cancellationToken) { @@ -283,6 +287,7 @@ public static TryCatch MonadicCreate( feedRangeState, partitionKey, queryPaginationOptions, + containerQueryProperties, cancellationToken); OptimisticDirectExecutionQueryPipelineImpl stage = new OptimisticDirectExecutionQueryPipelineImpl(partitionPageEnumerator); diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/PipelineFactory.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/PipelineFactory.cs index a119d5866a..256b517afa 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/PipelineFactory.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/PipelineFactory.cs @@ -20,6 +20,7 @@ namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Pagination; using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Skip; using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Take; + using Microsoft.Azure.Cosmos.Query.Core.QueryClient; using Microsoft.Azure.Cosmos.Query.Core.QueryPlan; internal static class PipelineFactory @@ -32,6 +33,7 @@ public static TryCatch MonadicCreate( PartitionKey? partitionKey, QueryInfo queryInfo, QueryPaginationOptions queryPaginationOptions, + ContainerQueryProperties containerQueryProperties, int maxConcurrency, CosmosElement requestContinuationToken, CancellationToken requestCancellationToken) @@ -89,6 +91,7 @@ public static TryCatch MonadicCreate( targetRanges: targetRanges, queryPaginationOptions: queryPaginationOptions, partitionKey: partitionKey, + containerQueryProperties: containerQueryProperties, prefetchPolicy: prefetchPolicy, maxConcurrency: maxConcurrency, continuationToken: continuationToken, 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 new file mode 100644 index 0000000000..dc4890e3c1 --- /dev/null +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/BaselineTest/TestBaseline/SubpartitionTests.TestQueriesOnSplitContainer.xml @@ -0,0 +1,27 @@ + + + + + + + + \ No newline at end of file diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Microsoft.Azure.Cosmos.Tests.csproj b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Microsoft.Azure.Cosmos.Tests.csproj index 81d0b6797a..c9ad45d555 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Microsoft.Azure.Cosmos.Tests.csproj +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Microsoft.Azure.Cosmos.Tests.csproj @@ -31,6 +31,7 @@ + @@ -349,6 +350,9 @@ + + PreserveNewest + PreserveNewest 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 ad17a6bac6..40730400f9 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 @@ -8,15 +8,18 @@ namespace Microsoft.Azure.Cosmos.Tests.Pagination using System.Collections; using System.Collections.Generic; using System.Collections.Immutable; + using Debug = System.Diagnostics.Debug; using System.IO; using System.Linq; using System.Reflection; + using System.Text; using System.Threading; using System.Threading.Tasks; using Microsoft.Azure.Cosmos; using Microsoft.Azure.Cosmos.ChangeFeed.Pagination; using Microsoft.Azure.Cosmos.CosmosElements; using Microsoft.Azure.Cosmos.CosmosElements.Numbers; + using Microsoft.Azure.Cosmos.Handlers; using Microsoft.Azure.Cosmos.Json; using Microsoft.Azure.Cosmos.Pagination; using Microsoft.Azure.Cosmos.Query.Core; @@ -35,6 +38,7 @@ namespace Microsoft.Azure.Cosmos.Tests.Pagination using static Microsoft.Azure.Cosmos.Query.Core.SqlQueryResumeFilter; using ResourceIdentifier = Cosmos.Pagination.ResourceIdentifier; using UInt128 = UInt128; + using Microsoft.Azure.Documents.Routing; // Collection useful for mocking requests and repartitioning (splits / merge). internal class InMemoryContainer : IMonadicDocumentContainer @@ -46,9 +50,13 @@ internal class InMemoryContainer : IMonadicDocumentContainer private PartitionKeyHashRangeDictionary> partitionedChanges; private Dictionary partitionKeyRangeIdToHashRange; private Dictionary cachedPartitionKeyRangeIdToHashRange; + private readonly bool createSplitForMultiHashAtSecondlevel; + private readonly bool resolvePartitionsBasedOnPrefix; public InMemoryContainer( - PartitionKeyDefinition partitionKeyDefinition) + PartitionKeyDefinition partitionKeyDefinition, + bool createSplitForMultiHashAtSecondlevel = false, + bool resolvePartitionsBasedOnPrefix = false) { this.partitionKeyDefinition = partitionKeyDefinition ?? throw new ArgumentNullException(nameof(partitionKeyDefinition)); PartitionKeyHashRange fullRange = new PartitionKeyHashRange(startInclusive: null, endExclusive: new PartitionKeyHash(Cosmos.UInt128.MaxValue)); @@ -66,6 +74,8 @@ public InMemoryContainer( { 0, fullRange } }; this.parentToChildMapping = new Dictionary(); + this.createSplitForMultiHashAtSecondlevel = createSplitForMultiHashAtSecondlevel; + this.resolvePartitionsBasedOnPrefix = resolvePartitionsBasedOnPrefix; } public Task>> MonadicGetFeedRangesAsync( @@ -471,7 +481,10 @@ public virtual Task> MonadicQueryAsync( using (ITrace childTrace = trace.StartChild("Query Transport", TraceComponent.Transport, TraceLevel.Info)) { - TryCatch monadicPartitionKeyRangeId = this.MonadicGetPartitionKeyRangeIdFromFeedRange(feedRangeState.FeedRange); + FeedRange feedRange = this.resolvePartitionsBasedOnPrefix ? + ResolveFeedRangeBasedOnPrefixContainer(feedRangeState.FeedRange, this.partitionKeyDefinition) : + feedRangeState.FeedRange; + TryCatch monadicPartitionKeyRangeId = this.MonadicGetPartitionKeyRangeIdFromFeedRange(feedRange); if (monadicPartitionKeyRangeId.Failed) { return Task.FromResult(TryCatch.FromException(monadicPartitionKeyRangeId.Exception)); @@ -941,6 +954,29 @@ public Task MonadicSplitAsync( return Task.FromResult(TryCatch.FromResult()); } + internal static FeedRange ResolveFeedRangeBasedOnPrefixContainer( + FeedRange feedRange, + PartitionKeyDefinition partitionKeyDefinition) + { + if (feedRange is FeedRangePartitionKey feedRangePartitionKey) + { + if (partitionKeyDefinition != null && partitionKeyDefinition.Kind == PartitionKind.MultiHash + && feedRangePartitionKey.PartitionKey.InternalKey?.Components?.Count < partitionKeyDefinition.Paths?.Count) + { + PartitionKeyHash partitionKeyHash = feedRangePartitionKey.PartitionKey.InternalKey.Components[0] switch + { + null => PartitionKeyHash.V2.HashUndefined(), + StringPartitionKeyComponent stringPartitionKey => PartitionKeyHash.V2.Hash((string)stringPartitionKey.ToObject()), + NumberPartitionKeyComponent numberPartitionKey => PartitionKeyHash.V2.Hash(Number64.ToDouble(numberPartitionKey.Value)), + _ => throw new ArgumentOutOfRangeException(), + }; + feedRange = new FeedRangeEpk(new Documents.Routing.Range(min: partitionKeyHash.Value, max: partitionKeyHash.Value + "-FF", isMinInclusive:true, isMaxInclusive: false)); + } + } + + return feedRange; + } + public Task MonadicMergeAsync( FeedRangeInternal feedRange1, FeedRangeInternal feedRange2, @@ -1335,7 +1371,7 @@ private static bool IsRecordWithinFeedRange( } } - private TryCatch MonadicGetPartitionKeyRangeIdFromFeedRange(FeedRange feedRange) + internal TryCatch MonadicGetPartitionKeyRangeIdFromFeedRange(FeedRange feedRange) { int partitionKeyRangeId; if (feedRange is FeedRangeEpk feedRangeEpk) @@ -1404,12 +1440,118 @@ private TryCatch MonadicGetPartitionKeyRangeIdFromFeedRange(FeedRange feedR private static PartitionKeyHashRange FeedRangeEpkToHashRange(FeedRangeEpk feedRangeEpk) { - PartitionKeyHash? start = feedRangeEpk.Range.Min == string.Empty ? (PartitionKeyHash?)null : PartitionKeyHash.Parse(feedRangeEpk.Range.Min); - PartitionKeyHash? end = feedRangeEpk.Range.Max == string.Empty || feedRangeEpk.Range.Max == "FF" ? (PartitionKeyHash?)null : PartitionKeyHash.Parse(feedRangeEpk.Range.Max); + PartitionKeyHash? start = + feedRangeEpk.Range.Min == string.Empty ? + (PartitionKeyHash?)null : + FromHashString(feedRangeEpk.Range.Min); + PartitionKeyHash? end = + feedRangeEpk.Range.Max == string.Empty || feedRangeEpk.Range.Max == "FF" ? + (PartitionKeyHash?)null : + FromHashString(feedRangeEpk.Range.Max); PartitionKeyHashRange hashRange = new PartitionKeyHashRange(start, end); return hashRange; } + /// + /// Creates a partition key hash from a rangeHash value. Supports if the rangeHash is over a hierarchical partition key. + /// + private static PartitionKeyHash FromHashString(string rangeHash) + { + List hashes = new(); + foreach(string hashComponent in GetHashComponents(rangeHash)) + { + // Hash FF has a special meaning in CosmosDB stack. It represents the max range which needs to be correctly represented for UInt128 parsing. + string value = hashComponent.Equals("FF", StringComparison.OrdinalIgnoreCase) ? + "FF-FF-FF-FF-FF-FF-FF-FF-FF-FF-FF-FF-FF-FF-FF-FF" : + hashComponent; + + bool success = UInt128.TryParse(value, out UInt128 uInt128); + Debug.Assert(success, "InMemoryContainer Assert!", "UInt128 parsing must succeed"); + hashes.Add(uInt128); + } + + return new PartitionKeyHash(hashes.ToArray()); + } + + /// + /// PartitionKeyHash.Parse requires a UInt128 parse-able string which itself requires hyphens to be present between subsequent byte values. + /// The hash values generated by rest of the (test) code may or may not honor this. + /// Furthermore, in case of hierarchical partitions, the hash values are concatenated together and therefore need to be broken into separate segments for parsing each one individually. + /// + /// + /// + private static IEnumerable GetHashComponents(string rangeValue) + { + int start = 0; + + while (start < rangeValue.Length) + { + string uInt128Segment = FixupUInt128(rangeValue, ref start); + yield return uInt128Segment; + } + } + + private static string FixupUInt128(string buffer, ref int start) + { + string result; + if (buffer.Length <= start + 2) + { + result = buffer.Substring(start); + start = buffer.Length; + } + else + { + StringBuilder stringBuilder = new StringBuilder(); + int index = start; + bool done = false; + int count = 0; + while (!done) + { + Debug.Assert(buffer[index] != '-', "InMemoryContainer Assert!", "First character of a chunk cannot be a hyphen"); + stringBuilder.Append(buffer[index]); + index++; + + Debug.Assert(index < buffer.Length, "InMemoryContainer Assert!", "At least 2 characters must be found in a chunk"); + Debug.Assert(buffer[index] != '-', "InMemoryContainer Assert!", "Second character of a chunk cannot be a hyphen"); + stringBuilder.Append(buffer[index]); + index++; + + if ((index < buffer.Length) && (buffer[index] == '-')) + { + index++; + } + + count++; + done = count == 16 || (index >= buffer.Length); + + if (!done) + { + stringBuilder.Append('-'); + } + } + + start = index; + + result = stringBuilder.ToString(); + Debug.Assert( + result.Length >= 2, + "InMemoryContainer Assert!", + "At least 1 byte must be present in hash value"); + Debug.Assert( + result[0] != '-' && result[result.Length - 1] != '-', + "InMemoryContainer Assert!", + "Hyphens should NOT be present at the start of end of the string"); + Debug.Assert( + Enumerable + .Range(1, result.Length - 1) + .All(i => (i % 3 == 2) == (result[i] == '-')), + "InMemoryContainer Assert!", + "Hyphens should be (only) present after every subsequent byte value"); + } + + return result; + } + private static FeedRangeEpk HashRangeToFeedRangeEpk(PartitionKeyHashRange hashRange) { return new FeedRangeEpk( @@ -1439,7 +1581,7 @@ private PartitionKeyHash ComputeMedianSplitPointAmongDocumentsInPKRange(Partitio // For MultiHash Collection, split at top level to ensure documents for top level key exist across partitions // after split - if (medianPkHash.HashValues.Count > 1) + if (medianPkHash.HashValues.Count > 1 && !this.createSplitForMultiHashAtSecondlevel) { return new PartitionKeyHash(medianPkHash.HashValues[0]); } diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs index bfbdda891c..611cdc719b 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs @@ -1255,18 +1255,18 @@ public override Task ForceRefreshCollectionCacheAsync(string collectionLink, Can public override Task GetCachedContainerQueryPropertiesAsync(string containerLink, Cosmos.PartitionKey? partitionKey, ITrace trace, CancellationToken cancellationToken) { - return Task.FromResult(new ContainerQueryProperties( - "test", - new List> - { - new Range( - PartitionKeyInternal.MinimumInclusiveEffectivePartitionKey, - PartitionKeyInternal.MaximumExclusiveEffectivePartitionKey, - true, - true) - }, - new PartitionKeyDefinition(), - Cosmos.GeospatialType.Geometry)); + return Task.FromResult(new ContainerQueryProperties( + "test", + new List> + { + new Range( + PartitionKeyInternal.MinimumInclusiveEffectivePartitionKey, + PartitionKeyInternal.MaximumExclusiveEffectivePartitionKey, + true, + true) + }, + new PartitionKeyDefinition(), + Cosmos.GeospatialType.Geometry)); } public override async Task GetClientDisableOptimisticDirectExecutionAsync() diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/FactoryTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/FactoryTests.cs index 21f87f2014..49d298d7b9 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/FactoryTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/FactoryTests.cs @@ -28,6 +28,7 @@ public void TestCreate() sqlQuerySpec: new SqlQuerySpec("SELECT * FROM c"), targetRanges: new List() { FeedRangeEpk.FullRange }, partitionKey: null, + containerQueryProperties: new Cosmos.Query.Core.QueryClient.ContainerQueryProperties(), queryInfo: new QueryInfo() { }, queryPaginationOptions: new QueryPaginationOptions(pageSizeHint: 10), maxConcurrency: 10, diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/FullPipelineTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/FullPipelineTests.cs index 97598ea59f..13f4c6091b 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/FullPipelineTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/FullPipelineTests.cs @@ -590,6 +590,7 @@ private static async Task CreatePipelineAsync( partitionKey: null, GetQueryPlan(query), queryPaginationOptions: new QueryPaginationOptions(pageSizeHint: pageSize), + containerQueryProperties: new ContainerQueryProperties(), maxConcurrency: 10, requestCancellationToken: default, requestContinuationToken: state); diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/ParallelCrossPartitionQueryPipelineStageTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/ParallelCrossPartitionQueryPipelineStageTests.cs index e772996748..69c49112b9 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/ParallelCrossPartitionQueryPipelineStageTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/ParallelCrossPartitionQueryPipelineStageTests.cs @@ -36,6 +36,7 @@ public void MonadicCreate_NullContinuationToken() targetRanges: new List() { FeedRangeEpk.FullRange }, queryPaginationOptions: new QueryPaginationOptions(pageSizeHint: 10), partitionKey: null, + containerQueryProperties: new Cosmos.Query.Core.QueryClient.ContainerQueryProperties(), maxConcurrency: 10, prefetchPolicy: PrefetchPolicy.PrefetchSinglePage, cancellationToken: default, @@ -54,6 +55,7 @@ public void MonadicCreate_NonCosmosArrayContinuationToken() targetRanges: new List() { FeedRangeEpk.FullRange }, queryPaginationOptions: new QueryPaginationOptions(pageSizeHint: 10), partitionKey: null, + containerQueryProperties: new Cosmos.Query.Core.QueryClient.ContainerQueryProperties(), maxConcurrency: 10, prefetchPolicy: PrefetchPolicy.PrefetchSinglePage, cancellationToken: default, @@ -73,6 +75,7 @@ public void MonadicCreate_EmptyArrayContinuationToken() targetRanges: new List() { FeedRangeEpk.FullRange }, queryPaginationOptions: new QueryPaginationOptions(pageSizeHint: 10), partitionKey: null, + containerQueryProperties: new Cosmos.Query.Core.QueryClient.ContainerQueryProperties(), maxConcurrency: 10, prefetchPolicy: PrefetchPolicy.PrefetchSinglePage, cancellationToken: default, @@ -92,6 +95,7 @@ public void MonadicCreate_NonParallelContinuationToken() targetRanges: new List() { FeedRangeEpk.FullRange }, queryPaginationOptions: new QueryPaginationOptions(pageSizeHint: 10), partitionKey: null, + containerQueryProperties: new Cosmos.Query.Core.QueryClient.ContainerQueryProperties(), maxConcurrency: 10, prefetchPolicy: PrefetchPolicy.PrefetchSinglePage, cancellationToken: default, @@ -115,6 +119,7 @@ public void MonadicCreate_SingleParallelContinuationToken() targetRanges: new List() { new FeedRangeEpk(new Documents.Routing.Range(min: "A", max: "B", isMinInclusive: true, isMaxInclusive: false)) }, queryPaginationOptions: new QueryPaginationOptions(pageSizeHint: 10), partitionKey: null, + containerQueryProperties: new Cosmos.Query.Core.QueryClient.ContainerQueryProperties(), maxConcurrency: 10, prefetchPolicy: PrefetchPolicy.PrefetchSinglePage, cancellationToken: default, @@ -145,6 +150,7 @@ public void MonadicCreate_MultipleParallelContinuationToken() }, queryPaginationOptions: new QueryPaginationOptions(pageSizeHint: 10), partitionKey: null, + containerQueryProperties: new Cosmos.Query.Core.QueryClient.ContainerQueryProperties(), maxConcurrency: 10, prefetchPolicy: PrefetchPolicy.PrefetchSinglePage, cancellationToken: default, @@ -186,6 +192,7 @@ async Task CreatePipelineStateAsync(IDocumentContainer docu cancellationToken: default), queryPaginationOptions: new QueryPaginationOptions(pageSizeHint: 10), partitionKey: null, + containerQueryProperties: new Cosmos.Query.Core.QueryClient.ContainerQueryProperties(), maxConcurrency: 10, prefetchPolicy: aggressivePrefetch ? PrefetchPolicy.PrefetchAll : PrefetchPolicy.PrefetchSinglePage, cancellationToken: default, diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/QueryPartitionRangePageEnumeratorTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/QueryPartitionRangePageEnumeratorTests.cs index 4672fb17f5..9380aecd23 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/QueryPartitionRangePageEnumeratorTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/QueryPartitionRangePageEnumeratorTests.cs @@ -100,6 +100,7 @@ public async Task TestSplitAsync() sqlQuerySpec: new Cosmos.Query.Core.SqlQuerySpec("SELECT * FROM c"), feedRangeState: feedRangeState, partitionKey: null, + containerQueryProperties: new Cosmos.Query.Core.QueryClient.ContainerQueryProperties(), queryPaginationOptions: new QueryPaginationOptions(pageSizeHint: 10), cancellationToken: default), trace: NoOpTrace.Singleton); @@ -143,6 +144,7 @@ protected override IAsyncEnumerable> CreateEnumerable( sqlQuerySpec: new Cosmos.Query.Core.SqlQuerySpec("SELECT * FROM c"), feedRangeState: feedRangeState, partitionKey: null, + containerQueryProperties: new Cosmos.Query.Core.QueryClient.ContainerQueryProperties(), queryPaginationOptions: new QueryPaginationOptions(pageSizeHint: 10), cancellationToken: default), trace: NoOpTrace.Singleton); @@ -166,6 +168,7 @@ protected override Task>> CreateEnumeratorA sqlQuerySpec: new Cosmos.Query.Core.SqlQuerySpec("SELECT * FROM c"), feedRangeState: new FeedRangeState(ranges[0], state), partitionKey: null, + containerQueryProperties: new Cosmos.Query.Core.QueryClient.ContainerQueryProperties(), queryPaginationOptions: new QueryPaginationOptions(pageSizeHint: 10), cancellationToken: cancellationToken), trace: NoOpTrace.Singleton); 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 new file mode 100644 index 0000000000..99d3f4cac3 --- /dev/null +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/SubpartitionTests.cs @@ -0,0 +1,382 @@ +namespace Microsoft.Azure.Cosmos.Tests.Query +{ + using System; + using System.Collections.Generic; + using System.IO; + using System.Linq; + using System.Threading; + using System.Threading.Tasks; + using System.Xml; + using Microsoft.Azure.Cosmos.CosmosElements; + using Microsoft.Azure.Cosmos.Pagination; + using Microsoft.Azure.Cosmos.Query; + using Microsoft.Azure.Cosmos.Query.Core; + using Microsoft.Azure.Cosmos.Query.Core.ExecutionContext; + using Microsoft.Azure.Cosmos.Query.Core.Monads; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Pagination; + using Microsoft.Azure.Cosmos.Query.Core.QueryClient; + using Microsoft.Azure.Cosmos.Query.Core.QueryPlan; + using Microsoft.Azure.Cosmos.Routing; + using Microsoft.Azure.Cosmos.Test.BaselineTest; + using Microsoft.Azure.Cosmos.Tests.Pagination; + using Microsoft.Azure.Cosmos.Tracing; + using Microsoft.Azure.Documents; + using Microsoft.VisualStudio.TestTools.UnitTesting; + + [TestClass] + public class SubpartitionTests : BaselineTests + { + private const int DocumentCount = 100; + private const int SplitPartitionKey = 2; + + [TestMethod] + public void TestQueriesOnSplitContainer() + { + this.ExecuteTestSuite(new List { new SubpartitionTestInput("Test Queries on Split container") }); + } + + /// + /// The test is a baseline for mock framework which splits the container at the top level of a hierarchical partition key. + /// After split, it is expected that more than one physical partitions contain data for some value of a top level path of partition key. + /// Please note that this does NOT occur in a single-partition key scenario where all data for a given value of a partition key + /// is contained within single physical partition. + /// This situation is known to create issues, especially while running queries due to inconsistent handling of FeedRangePartitionKey and FeedRangeEpk in the SDK stack. + /// Test framework's behavior in being able to replicate this situation is critical to for ensuring that tests provide sufficient protection against regressions. + /// + [TestMethod] + public async Task VerifyTestFrameworkSupportsPartitionSplit() + { + PartitionKeyDefinition partitionKeyDefinition = CreatePartitionKeyDefinition(); + InMemoryContainer inMemoryContainer = await CreateSplitInMemoryDocumentContainerAsync(DocumentCount, partitionKeyDefinition); + Cosmos.PartitionKey partitionKey = new Cosmos.PartitionKeyBuilder().Add(SplitPartitionKey.ToString()).Build(); + FeedRangePartitionKey feedRangePartitionKey = new FeedRangePartitionKey(partitionKey); + FeedRangeEpk feedRangeEpk = InMemoryContainer.ResolveFeedRangeBasedOnPrefixContainer(feedRangePartitionKey, partitionKeyDefinition) as FeedRangeEpk; + Assert.IsNotNull(feedRangeEpk); + TryCatch pkRangeId = inMemoryContainer.MonadicGetPartitionKeyRangeIdFromFeedRange(feedRangeEpk); + Assert.IsTrue(pkRangeId.Failed, $"Expected to fail for partition key {SplitPartitionKey}"); + Assert.IsTrue(pkRangeId.Exception.InnerException.Message.StartsWith("Epk Range: [B5-D7-B7-26-D6-EA-DB-11-F1-EF-AD-92-12-15-D6-60,B5-D7-B7-26-D6-EA-DB-11-F1-EF-AD-92-12-15-D6-60-FF) is gone."), "Gone exception is expected!"); + } + + public SubpartitionTestOutput ExecuteTest2(SubpartitionTestInput input) + { + return new SubpartitionTestOutput(new List()); + } + + public override SubpartitionTestOutput ExecuteTest(SubpartitionTestInput input) + { + IMonadicDocumentContainer monadicDocumentContainer = CreateSplitDocumentContainerAsync(DocumentCount).Result; + DocumentContainer documentContainer = new DocumentContainer(monadicDocumentContainer); + + List documents = new List(); + QueryRequestOptions queryRequestOptions = new QueryRequestOptions() + { + PartitionKey = new PartitionKeyBuilder().Add(SplitPartitionKey.ToString()).Build() + }; + (CosmosQueryExecutionContextFactory.InputParameters inputParameters, CosmosQueryContextCore cosmosQueryContextCore) = + CreateInputParamsAndQueryContext(queryRequestOptions); + IQueryPipelineStage queryPipelineStage = CosmosQueryExecutionContextFactory.Create( + documentContainer, + cosmosQueryContextCore, + inputParameters, + NoOpTrace.Singleton); + while (queryPipelineStage.MoveNextAsync(NoOpTrace.Singleton).Result) + { + TryCatch tryGetPage = queryPipelineStage.Current; + + if (tryGetPage.Failed) + { + Assert.Fail("Unexpected error. Gone Exception should not reach till here"); + } + + documents.AddRange(tryGetPage.Result.Documents); + } + + return new SubpartitionTestOutput(documents); + } + + private static Tuple CreateInputParamsAndQueryContext(QueryRequestOptions queryRequestOptions, bool clientDisableOde = false) + { + string query = @"SELECT c.id, c.value2 FROM c"; + CosmosElement continuationToken = null; + PartitionKeyDefinition partitionKeyDefinition = new PartitionKeyDefinition() + { + Paths = new System.Collections.ObjectModel.Collection() + { + "/id", + "/value1", + "/value2" + }, + Kind = PartitionKind.MultiHash, + Version = PartitionKeyDefinitionVersion.V2, + }; + + CosmosSerializerCore serializerCore = new(); + using StreamReader streamReader = new(serializerCore.ToStreamSqlQuerySpec(new SqlQuerySpec(query), Documents.ResourceType.Document)); + string sqlQuerySpecJsonString = streamReader.ReadToEnd(); + + (PartitionedQueryExecutionInfo partitionedQueryExecutionInfo, QueryPartitionProvider queryPartitionProvider) = GetPartitionedQueryExecutionInfoAndPartitionProvider(sqlQuerySpecJsonString, partitionKeyDefinition, clientDisableOde); + CosmosQueryExecutionContextFactory.InputParameters inputParameters = new CosmosQueryExecutionContextFactory.InputParameters( + sqlQuerySpec: new SqlQuerySpec(query), + initialUserContinuationToken: continuationToken, + initialFeedRange: null, + maxConcurrency: queryRequestOptions.MaxConcurrency, + maxItemCount: queryRequestOptions.MaxItemCount, + maxBufferedItemCount: queryRequestOptions.MaxBufferedItemCount, + partitionKey: queryRequestOptions.PartitionKey, + properties: new Dictionary() { { "x-ms-query-partitionkey-definition", partitionKeyDefinition } }, + partitionedQueryExecutionInfo: null, + executionEnvironment: null, + returnResultsInDeterministicOrder: null, + enableOptimisticDirectExecution: queryRequestOptions.EnableOptimisticDirectExecution, + testInjections: queryRequestOptions.TestSettings); + + string databaseId = "db1234"; + string resourceLink = $"dbs/{databaseId}/colls"; + CosmosQueryContextCore cosmosQueryContextCore = new CosmosQueryContextCore( + client: new TestCosmosQueryClient(queryPartitionProvider), + resourceTypeEnum: Documents.ResourceType.Document, + operationType: Documents.OperationType.Query, + resourceType: typeof(QueryResponseCore), + resourceLink: resourceLink, + isContinuationExpected: true, + allowNonValueAggregateQuery: true, + useSystemPrefix: false, + correlatedActivityId: Guid.NewGuid()); + + return Tuple.Create(inputParameters, cosmosQueryContextCore); + } + + internal static Tuple GetPartitionedQueryExecutionInfoAndPartitionProvider(string querySpecJsonString, PartitionKeyDefinition pkDefinition, bool clientDisableOde = false) + { + QueryPartitionProvider queryPartitionProvider = CreateCustomQueryPartitionProvider("clientDisableOptimisticDirectExecution", clientDisableOde.ToString().ToLower()); + TryCatch tryGetQueryPlan = queryPartitionProvider.TryGetPartitionedQueryExecutionInfo( + querySpecJsonString: querySpecJsonString, + partitionKeyDefinition: pkDefinition, + requireFormattableOrderByQuery: true, + isContinuationExpected: true, + allowNonValueAggregateQuery: true, + hasLogicalPartitionKey: false, + allowDCount: true, + useSystemPrefix: false, + geospatialType: Cosmos.GeospatialType.Geography); + + PartitionedQueryExecutionInfo partitionedQueryExecutionInfo = tryGetQueryPlan.Succeeded ? tryGetQueryPlan.Result : throw tryGetQueryPlan.Exception; + return Tuple.Create(partitionedQueryExecutionInfo, queryPartitionProvider); + } + + private static QueryPartitionProvider CreateCustomQueryPartitionProvider(string key, string value) + { + Dictionary queryEngineConfiguration = new Dictionary() + { + {"maxSqlQueryInputLength", 262144}, + {"maxJoinsPerSqlQuery", 5}, + {"maxLogicalAndPerSqlQuery", 2000}, + {"maxLogicalOrPerSqlQuery", 2000}, + {"maxUdfRefPerSqlQuery", 10}, + {"maxInExpressionItemsCount", 16000}, + {"queryMaxGroupByTableCellCount", 500000 }, + {"queryMaxInMemorySortDocumentCount", 500}, + {"maxQueryRequestTimeoutFraction", 0.90}, + {"sqlAllowNonFiniteNumbers", false}, + {"sqlAllowAggregateFunctions", true}, + {"sqlAllowSubQuery", true}, + {"sqlAllowScalarSubQuery", true}, + {"allowNewKeywords", true}, + {"sqlAllowLike", true}, + {"sqlAllowGroupByClause", true}, + {"maxSpatialQueryCells", 12}, + {"spatialMaxGeometryPointCount", 256}, + {"sqlDisableQueryILOptimization", false}, + {"sqlDisableFilterPlanOptimization", false}, + {"clientDisableOptimisticDirectExecution", false} + }; + + queryEngineConfiguration[key] = bool.TryParse(value, out bool boolValue) ? boolValue : value; + + return new QueryPartitionProvider(queryEngineConfiguration); + } + + private static PartitionKeyDefinition CreatePartitionKeyDefinition() + { + PartitionKeyDefinition partitionKeyDefinition = new PartitionKeyDefinition() + { + Paths = new System.Collections.ObjectModel.Collection() + { + "/id", + "/value1", + "/value2" + }, + Kind = PartitionKind.MultiHash, + Version = PartitionKeyDefinitionVersion.V2, + }; + + return partitionKeyDefinition; + } + + private static async Task CreateSplitDocumentContainerAsync(int numItems) + { + PartitionKeyDefinition partitionKeyDefinition = CreatePartitionKeyDefinition(); + InMemoryContainer inMemoryContainer = await CreateSplitInMemoryDocumentContainerAsync(numItems, partitionKeyDefinition); + DocumentContainer documentContainer = new DocumentContainer(inMemoryContainer); + return documentContainer; + } + + private static async Task CreateSplitInMemoryDocumentContainerAsync(int numItems, PartitionKeyDefinition partitionKeyDefinition) + { + InMemoryContainer inMemoryContainer = new InMemoryContainer(partitionKeyDefinition, createSplitForMultiHashAtSecondlevel: true, resolvePartitionsBasedOnPrefix: true); + for (int i = 0; i < numItems; i++) + { + // Insert an item + CosmosObject item = CosmosObject.Parse($"{{\"id\" : \"{i % 5}\", \"value1\" : \"{Guid.NewGuid()}\", \"value2\" : \"{i}\" }}"); + while (true) + { + TryCatch monadicCreateRecord = await inMemoryContainer.MonadicCreateItemAsync(item, cancellationToken: default); + if (monadicCreateRecord.Succeeded) + { + break; + } + } + } + + await inMemoryContainer.MonadicSplitAsync(FeedRangeEpk.FullRange, cancellationToken: default); + + return inMemoryContainer; + } + internal class TestCosmosQueryClient : CosmosQueryClient + { + private readonly QueryPartitionProvider queryPartitionProvider; + + public TestCosmosQueryClient(QueryPartitionProvider queryPartitionProvider) + { + this.queryPartitionProvider = queryPartitionProvider; + } + + public override Action OnExecuteScalarQueryCallback => throw new NotImplementedException(); + + public override bool BypassQueryParsing() + { + return false; + } + + public override void ClearSessionTokenCache(string collectionFullName) + { + throw new NotImplementedException(); + } + + public override Task> ExecuteItemQueryAsync(string resourceUri, ResourceType resourceType, OperationType operationType, Cosmos.FeedRange feedRange, QueryRequestOptions requestOptions, AdditionalRequestHeaders additionalRequestHeaders, SqlQuerySpec sqlQuerySpec, string continuationToken, int pageSize, ITrace trace, CancellationToken cancellationToken) + { + throw new NotImplementedException(); + } + + public override Task ExecuteQueryPlanRequestAsync(string resourceUri, ResourceType resourceType, OperationType operationType, SqlQuerySpec sqlQuerySpec, Cosmos.PartitionKey? partitionKey, string supportedQueryFeatures, Guid clientQueryCorrelationId, ITrace trace, CancellationToken cancellationToken) + { + return Task.FromResult(new PartitionedQueryExecutionInfo()); + } + + public override Task ForceRefreshCollectionCacheAsync(string collectionLink, CancellationToken cancellationToken) + { + throw new NotImplementedException(); + } + + public override Task GetCachedContainerQueryPropertiesAsync(string containerLink, Cosmos.PartitionKey? partitionKey, ITrace trace, CancellationToken cancellationToken) + { + List hashes = new(); + foreach (Documents.Routing.IPartitionKeyComponent component in partitionKey.Value.InternalKey.Components) + { + PartitionKeyHash partitionKeyHash = component switch + { + null => PartitionKeyHash.V2.HashUndefined(), + Documents.Routing.StringPartitionKeyComponent stringPartitionKey => PartitionKeyHash.V2.Hash((string)stringPartitionKey.ToObject()), + Documents.Routing.NumberPartitionKeyComponent numberPartitionKey => PartitionKeyHash.V2.Hash(Number64.ToDouble(numberPartitionKey.Value)), + _ => throw new ArgumentOutOfRangeException(), + }; + hashes.Add(partitionKeyHash.Value); + } + + string min = string.Join(string.Empty, hashes); + string max = min + "-FF"; + return Task.FromResult(new ContainerQueryProperties( + "test", + new List> + { + new Documents.Routing.Range( + min, + max, + true, + true) + }, + new PartitionKeyDefinition(), + Cosmos.GeospatialType.Geometry)); + } + + public override async Task GetClientDisableOptimisticDirectExecutionAsync() + { + return this.queryPartitionProvider.ClientDisableOptimisticDirectExecution; + } + + public override Task> GetTargetPartitionKeyRangeByFeedRangeAsync(string resourceLink, string collectionResourceId, PartitionKeyDefinition partitionKeyDefinition, FeedRangeInternal feedRangeInternal, bool forceRefresh, ITrace trace) + { + throw new NotImplementedException(); + } + + 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 + } + }); + } + + public override Task> TryGetOverlappingRangesAsync(string collectionResourceId, Documents.Routing.Range range, bool forceRefresh = false) + { + throw new NotImplementedException(); + } + + public override async Task> TryGetPartitionedQueryExecutionInfoAsync(SqlQuerySpec sqlQuerySpec, ResourceType resourceType, PartitionKeyDefinition partitionKeyDefinition, bool requireFormattableOrderByQuery, bool isContinuationExpected, bool allowNonValueAggregateQuery, bool hasLogicalPartitionKey, bool allowDCount, bool useSystemPrefix, Cosmos.GeospatialType geospatialType, CancellationToken cancellationToken) + { + CosmosSerializerCore serializerCore = new(); + using StreamReader streamReader = new(serializerCore.ToStreamSqlQuerySpec(sqlQuerySpec, Documents.ResourceType.Document)); + string sqlQuerySpecJsonString = streamReader.ReadToEnd(); + + (PartitionedQueryExecutionInfo partitionedQueryExecutionInfo, QueryPartitionProvider queryPartitionProvider) = OptimisticDirectExecutionQueryBaselineTests.GetPartitionedQueryExecutionInfoAndPartitionProvider(sqlQuerySpecJsonString, partitionKeyDefinition); + return TryCatch.FromResult(partitionedQueryExecutionInfo); + } + } + } + + public class SubpartitionTestInput : BaselineTestInput + { + public SubpartitionTestInput(string description) + :base(description) + { + } + + public override void SerializeAsXml(XmlWriter xmlWriter) + { + } + } + + public class SubpartitionTestOutput : BaselineTestOutput + { + private readonly List documents; + + internal SubpartitionTestOutput(IReadOnlyList documents) + { + this.documents = documents.ToList(); + } + + public override void SerializeAsXml(XmlWriter xmlWriter) + { + xmlWriter.WriteStartElement("Documents"); + string content = string.Join($",{Environment.NewLine}", + this.documents.Select(doc => doc.ToString()).OrderBy(serializedDoc => serializedDoc)); + xmlWriter.WriteCData(content); + xmlWriter.WriteEndElement(); + } + } +} diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Tracing/TraceWriterBaselineTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Tracing/TraceWriterBaselineTests.cs index 2d69341182..fc3324a5da 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Tracing/TraceWriterBaselineTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Tracing/TraceWriterBaselineTests.cs @@ -758,6 +758,7 @@ private static IQueryPipelineStage CreatePipeline(IDocumentContainer documentCon partitionKey: null, GetQueryPlan(query), new QueryPaginationOptions(pageSizeHint: pageSize), + containerQueryProperties: new Cosmos.Query.Core.QueryClient.ContainerQueryProperties(), maxConcurrency: 10, requestCancellationToken: default, requestContinuationToken: state);