From 86a0dc40363b05fc76490ebb044893b15b6593fe Mon Sep 17 00:00:00 2001 From: Brandon Chong Date: Mon, 22 Jun 2020 19:33:52 -0700 Subject: [PATCH 01/85] rename --- .../src/FeedRange/FeedRanges/FeedRangeEPK.cs | 8 +++---- .../FeedRange/FeedRanges/FeedRangeVisitor.cs | 2 +- .../src/Resource/Container/ContainerCore.cs | 4 ++-- .../FeedIterators/FeedRangeIteratorCore.cs | 4 ++-- .../QueryResponses/ChangeFeedIteratorCore.cs | 6 ++--- .../FeedRange/FeedRangeInternalConverter.cs | 6 ++--- .../FeedToken/FeedRangeTests.cs | 2 +- .../FeedRange/FeedRangeContinuationTests.cs | 2 +- .../FeedRange/FeedRangeTests.cs | 12 +++++----- .../ReadFeedTokenIteratorCoreTests.cs | 24 +++++++++---------- 10 files changed, 35 insertions(+), 35 deletions(-) diff --git a/Microsoft.Azure.Cosmos/src/FeedRange/FeedRanges/FeedRangeEPK.cs b/Microsoft.Azure.Cosmos/src/FeedRange/FeedRanges/FeedRangeEPK.cs index 32a4467493..df16de872c 100644 --- a/Microsoft.Azure.Cosmos/src/FeedRange/FeedRanges/FeedRangeEPK.cs +++ b/Microsoft.Azure.Cosmos/src/FeedRange/FeedRanges/FeedRangeEPK.cs @@ -13,20 +13,20 @@ namespace Microsoft.Azure.Cosmos /// /// FeedRange that represents an effective partition key range. /// - internal sealed class FeedRangeEPK : FeedRangeInternal + internal sealed class FeedRangeEpk : FeedRangeInternal { public Documents.Routing.Range Range { get; } - public static FeedRangeEPK ForFullRange() + public static FeedRangeEpk ForFullRange() { - return new FeedRangeEPK(new Documents.Routing.Range( + return new FeedRangeEpk(new Documents.Routing.Range( Documents.Routing.PartitionKeyInternal.MinimumInclusiveEffectivePartitionKey, Documents.Routing.PartitionKeyInternal.MaximumExclusiveEffectivePartitionKey, isMinInclusive: true, isMaxInclusive: false)); } - public FeedRangeEPK(Documents.Routing.Range range) + public FeedRangeEpk(Documents.Routing.Range range) { this.Range = range; } diff --git a/Microsoft.Azure.Cosmos/src/FeedRange/FeedRanges/FeedRangeVisitor.cs b/Microsoft.Azure.Cosmos/src/FeedRange/FeedRanges/FeedRangeVisitor.cs index 1910aa5d02..51c15ab806 100644 --- a/Microsoft.Azure.Cosmos/src/FeedRange/FeedRanges/FeedRangeVisitor.cs +++ b/Microsoft.Azure.Cosmos/src/FeedRange/FeedRanges/FeedRangeVisitor.cs @@ -27,7 +27,7 @@ public void Visit(FeedRangePartitionKeyRange feedRange) ChangeFeedRequestOptions.FillPartitionKeyRangeId(this.request, feedRange.PartitionKeyRangeId); } - public void Visit(FeedRangeEPK feedRange) + public void Visit(FeedRangeEpk feedRange) { // No-op since the range is defined by the composite continuation token } diff --git a/Microsoft.Azure.Cosmos/src/Resource/Container/ContainerCore.cs b/Microsoft.Azure.Cosmos/src/Resource/Container/ContainerCore.cs index e87dccd915..e699339a23 100644 --- a/Microsoft.Azure.Cosmos/src/Resource/Container/ContainerCore.cs +++ b/Microsoft.Azure.Cosmos/src/Resource/Container/ContainerCore.cs @@ -241,10 +241,10 @@ public async Task> GetFeedRangesAsync( isMinInclusive: true, isMaxInclusive: false), forceRefresh: true); - List feedTokens = new List(partitionKeyRanges.Count); + List feedTokens = new List(partitionKeyRanges.Count); foreach (PartitionKeyRange partitionKeyRange in partitionKeyRanges) { - feedTokens.Add(new FeedRangeEPK(partitionKeyRange.ToRange())); + feedTokens.Add(new FeedRangeEpk(partitionKeyRange.ToRange())); } return feedTokens; diff --git a/Microsoft.Azure.Cosmos/src/Resource/FeedIterators/FeedRangeIteratorCore.cs b/Microsoft.Azure.Cosmos/src/Resource/FeedIterators/FeedRangeIteratorCore.cs index 17658ad4f6..ca6ef06dc4 100644 --- a/Microsoft.Azure.Cosmos/src/Resource/FeedIterators/FeedRangeIteratorCore.cs +++ b/Microsoft.Azure.Cosmos/src/Resource/FeedIterators/FeedRangeIteratorCore.cs @@ -49,7 +49,7 @@ public static FeedRangeIteratorCore Create( } // Backward compatible with old format - feedRangeInternal = FeedRangeEPK.ForFullRange(); + feedRangeInternal = FeedRangeEpk.ForFullRange(); feedRangeContinuation = new FeedRangeCompositeContinuation( string.Empty, feedRangeInternal, @@ -65,7 +65,7 @@ public static FeedRangeIteratorCore Create( return new FeedRangeIteratorCore(containerCore, feedRangeContinuation, options, resourceType, queryDefinition); } - feedRangeInternal = feedRangeInternal ?? FeedRangeEPK.ForFullRange(); + feedRangeInternal = feedRangeInternal ?? FeedRangeEpk.ForFullRange(); return new FeedRangeIteratorCore(containerCore, feedRangeInternal, options, resourceType, queryDefinition); } diff --git a/Microsoft.Azure.Cosmos/src/Resource/QueryResponses/ChangeFeedIteratorCore.cs b/Microsoft.Azure.Cosmos/src/Resource/QueryResponses/ChangeFeedIteratorCore.cs index fbd8dc2a63..3331f4c612 100644 --- a/Microsoft.Azure.Cosmos/src/Resource/QueryResponses/ChangeFeedIteratorCore.cs +++ b/Microsoft.Azure.Cosmos/src/Resource/QueryResponses/ChangeFeedIteratorCore.cs @@ -48,7 +48,7 @@ public static ChangeFeedIteratorCore Create( } } - feedRangeInternal = feedRangeInternal ?? FeedRangeEPK.ForFullRange(); + feedRangeInternal = feedRangeInternal ?? FeedRangeEpk.ForFullRange(); return new ChangeFeedIteratorCore(container, feedRangeInternal, changeFeedRequestOptions); } @@ -234,7 +234,7 @@ private async Task InitializeFeedContinuationAsync(CancellationToken cancellatio { IReadOnlyList pkRanges = await partitionKeyRangeCache.TryGetOverlappingRangesAsync( collectionRid: this.lazyContainerRid.Result.Result, - range: (this.FeedRangeInternal as FeedRangeEPK).Range, + range: (this.FeedRangeInternal as FeedRangeEpk).Range, forceRefresh: false); ranges = pkRanges.Select(pkRange => pkRange.ToRange()).ToList(); } @@ -254,7 +254,7 @@ private async Task InitializeFeedContinuationAsync(CancellationToken cancellatio partitionKeyDefinition: null); // Override the original PKRangeId based FeedRange - this.FeedRangeInternal = new FeedRangeEPK(effectiveRanges[0]); + this.FeedRangeInternal = new FeedRangeEpk(effectiveRanges[0]); this.FeedRangeContinuation = new FeedRangeCompositeContinuation( containerRid: this.lazyContainerRid.Result.Result, feedRange: this.FeedRangeInternal, diff --git a/Microsoft.Azure.Cosmos/src/Serializer/FeedRange/FeedRangeInternalConverter.cs b/Microsoft.Azure.Cosmos/src/Serializer/FeedRange/FeedRangeInternalConverter.cs index 60afd93da1..a806671646 100644 --- a/Microsoft.Azure.Cosmos/src/Serializer/FeedRange/FeedRangeInternalConverter.cs +++ b/Microsoft.Azure.Cosmos/src/Serializer/FeedRange/FeedRangeInternalConverter.cs @@ -18,7 +18,7 @@ internal sealed class FeedRangeInternalConverter : JsonConverter public override bool CanConvert(Type objectType) { - return objectType == typeof(FeedRangeEPK) + return objectType == typeof(FeedRangeEpk) || objectType == typeof(FeedRangePartitionKey) || objectType == typeof(FeedRangePartitionKeyRange); } @@ -63,7 +63,7 @@ public static FeedRangeInternal ReadJObject( try { Documents.Routing.Range completeRange = (Documents.Routing.Range)rangeJsonConverter.ReadJson(rangeJToken.CreateReader(), typeof(Documents.Routing.Range), null, serializer); - return new FeedRangeEPK(completeRange); + return new FeedRangeEpk(completeRange); } catch (JsonSerializationException) { @@ -94,7 +94,7 @@ public static void WriteJObject( object value, JsonSerializer serializer) { - if (value is FeedRangeEPK feedRangeEPK) + if (value is FeedRangeEpk feedRangeEPK) { writer.WritePropertyName(FeedRangeInternalConverter.RangePropertyName); rangeJsonConverter.WriteJson(writer, feedRangeEPK.Range, serializer); diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/FeedToken/FeedRangeTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/FeedToken/FeedRangeTests.cs index f8acd0e65a..c44f12ae44 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/FeedToken/FeedRangeTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/FeedToken/FeedRangeTests.cs @@ -45,7 +45,7 @@ public async Task FeedRange_EPK_Serialization() List tokens = new List(); foreach (FeedRange range in ranges) { - FeedRangeEPK feedRangeEPK = range as FeedRangeEPK; + FeedRangeEpk feedRangeEPK = range as FeedRangeEpk; FeedRangeCompositeContinuation feedRangeCompositeContinuation = new FeedRangeCompositeContinuation(containerRid, feedRangeEPK, new List>() { feedRangeEPK.Range }, continuation); tokens.Add(feedRangeCompositeContinuation); serializations.Add(feedRangeCompositeContinuation.ToString()); diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/FeedRange/FeedRangeContinuationTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/FeedRange/FeedRangeContinuationTests.cs index c541d02ffd..cc8c8e7252 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/FeedRange/FeedRangeContinuationTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/FeedRange/FeedRangeContinuationTests.cs @@ -45,7 +45,7 @@ public void FeedRangeCompositeContinuation_TryParse() new Documents.Routing.Range("A", "B", true, false), new Documents.Routing.Range("D", "E", true, false), }; - FeedRangeInternal feedRangeInternal = new FeedRangeEPK(new Documents.Routing.Range("A", "E", true, false)); + FeedRangeInternal feedRangeInternal = new FeedRangeEpk(new Documents.Routing.Range("A", "E", true, false)); FeedRangeCompositeContinuation token = new FeedRangeCompositeContinuation(containerRid, feedRangeInternal, keyRanges); Assert.IsTrue(FeedRangeContinuation.TryParse(token.ToString(), out _)); Assert.IsFalse(FeedRangeContinuation.TryParse("whatever", out _)); diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/FeedRange/FeedRangeTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/FeedRange/FeedRangeTests.cs index 9ae01625b6..3436b1fe90 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/FeedRange/FeedRangeTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/FeedRange/FeedRangeTests.cs @@ -21,7 +21,7 @@ public class FeedRangeTests public void FeedRangeEPK_Range() { Documents.Routing.Range range = new Documents.Routing.Range("AA", "BB", true, false); - FeedRangeEPK feedRangeEPK = new FeedRangeEPK(range); + FeedRangeEpk feedRangeEPK = new FeedRangeEpk(range); Assert.AreEqual(range, feedRangeEPK.Range); } @@ -45,7 +45,7 @@ public void FeedRangePKRangeId_PKRange() public async Task FeedRangeEPK_GetEffectiveRangesAsync() { Documents.Routing.Range range = new Documents.Routing.Range("AA", "BB", true, false); - FeedRangeEPK feedRangeEPK = new FeedRangeEPK(range); + FeedRangeEpk feedRangeEPK = new FeedRangeEpk(range); List> ranges = await feedRangeEPK.GetEffectiveRangesAsync(Mock.Of(), null, null); Assert.AreEqual(1, ranges.Count); Assert.AreEqual(range, ranges[0]); @@ -125,7 +125,7 @@ public async Task FeedRangeEPK_GetPartitionKeyRangesAsync() .Setup(f => f.TryGetOverlappingRangesAsync(It.IsAny(), It.Is>(s => s == range), It.IsAny())) .ReturnsAsync(new List() { partitionKeyRange }); - FeedRangeEPK feedRangeEPK = new FeedRangeEPK(range); + FeedRangeEpk feedRangeEPK = new FeedRangeEpk(range); IEnumerable pkRanges = await feedRangeEPK.GetPartitionKeyRangesAsync(routingProvider, null, null, default(CancellationToken)); Assert.AreEqual(1, pkRanges.Count()); Assert.AreEqual(partitionKeyRange.Id, pkRanges.First()); @@ -164,7 +164,7 @@ public async Task FeedRangePKRangeId_GetPartitionKeyRangesAsync() public void FeedRangeEPK_RequestVisitor() { Documents.Routing.Range range = new Documents.Routing.Range("AA", "BB", true, false); - FeedRangeEPK feedRange = new FeedRangeEPK(range); + FeedRangeEpk feedRange = new FeedRangeEpk(range); RequestMessage requestMessage = new RequestMessage(); FeedRangeVisitor feedRangeVisitor = new FeedRangeVisitor(requestMessage); feedRange.Accept(feedRangeVisitor); @@ -199,9 +199,9 @@ public void FeedRangePK_RequestVisitor() public void FeedRangeEPK_ToJsonFromJson() { Documents.Routing.Range range = new Documents.Routing.Range("AA", "BB", true, false); - FeedRangeEPK feedRangeEPK = new FeedRangeEPK(range); + FeedRangeEpk feedRangeEPK = new FeedRangeEpk(range); string representation = feedRangeEPK.ToJsonString(); - FeedRangeEPK feedRangeEPKDeserialized = Cosmos.FeedRange.FromJsonString(representation) as FeedRangeEPK; + FeedRangeEpk feedRangeEPKDeserialized = Cosmos.FeedRange.FromJsonString(representation) as FeedRangeEpk; Assert.IsNotNull(feedRangeEPKDeserialized); Assert.AreEqual(feedRangeEPK.Range.Min, feedRangeEPKDeserialized.Range.Min); Assert.AreEqual(feedRangeEPK.Range.Max, feedRangeEPKDeserialized.Range.Max); diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/FeedRange/ReadFeedTokenIteratorCoreTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/FeedRange/ReadFeedTokenIteratorCoreTests.cs index 974ac10430..c2e4ef8f6d 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/FeedRange/ReadFeedTokenIteratorCoreTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/FeedRange/ReadFeedTokenIteratorCoreTests.cs @@ -30,9 +30,9 @@ public void ReadFeedIteratorCore_HasMoreResultsDefault() public void ReadFeedIteratorCore_Create_Default() { FeedRangeIteratorCore feedTokenIterator = FeedRangeIteratorCore.Create(Mock.Of(), null, null, null); - FeedRangeEPK defaultRange = feedTokenIterator.FeedRangeInternal as FeedRangeEPK; - Assert.AreEqual(FeedRangeEPK.ForFullRange().Range.Min, defaultRange.Range.Min); - Assert.AreEqual(FeedRangeEPK.ForFullRange().Range.Max, defaultRange.Range.Max); + FeedRangeEpk defaultRange = feedTokenIterator.FeedRangeInternal as FeedRangeEpk; + Assert.AreEqual(FeedRangeEpk.ForFullRange().Range.Min, defaultRange.Range.Min); + Assert.AreEqual(FeedRangeEpk.ForFullRange().Range.Max, defaultRange.Range.Max); Assert.IsNull(feedTokenIterator.FeedRangeContinuation); } @@ -40,7 +40,7 @@ public void ReadFeedIteratorCore_Create_Default() public void ReadFeedIteratorCore_Create_WithRange() { Documents.Routing.Range range = new Documents.Routing.Range("A", "B", true, false); - FeedRangeEPK feedRangeEPK = new FeedRangeEPK(range); + FeedRangeEpk feedRangeEPK = new FeedRangeEpk(range); FeedRangeIteratorCore feedTokenIterator = FeedRangeIteratorCore.Create(Mock.Of(), feedRangeEPK, null, null); Assert.AreEqual(feedRangeEPK, feedTokenIterator.FeedRangeInternal); Assert.IsNull(feedTokenIterator.FeedRangeContinuation); @@ -52,9 +52,9 @@ public void ReadFeedIteratorCore_Create_WithContinuation() string continuation = Guid.NewGuid().ToString(); FeedRangeIteratorCore feedTokenIterator = FeedRangeIteratorCore.Create(Mock.Of(), null, continuation, null); - FeedRangeEPK defaultRange = feedTokenIterator.FeedRangeInternal as FeedRangeEPK; - Assert.AreEqual(FeedRangeEPK.ForFullRange().Range.Min, defaultRange.Range.Min); - Assert.AreEqual(FeedRangeEPK.ForFullRange().Range.Max, defaultRange.Range.Max); + FeedRangeEpk defaultRange = feedTokenIterator.FeedRangeInternal as FeedRangeEpk; + Assert.AreEqual(FeedRangeEpk.ForFullRange().Range.Min, defaultRange.Range.Min); + Assert.AreEqual(FeedRangeEpk.ForFullRange().Range.Max, defaultRange.Range.Max); Assert.IsNotNull(feedTokenIterator.FeedRangeContinuation); Assert.AreEqual(continuation, feedTokenIterator.FeedRangeContinuation.GetContinuation()); } @@ -64,12 +64,12 @@ public void ReadFeedIteratorCore_Create_WithFeedContinuation() { string continuation = Guid.NewGuid().ToString(); - FeedRangeEPK feedRangeEPK = FeedRangeEPK.ForFullRange(); + FeedRangeEpk feedRangeEPK = FeedRangeEpk.ForFullRange(); FeedRangeCompositeContinuation feedRangeSimpleContinuation = new FeedRangeCompositeContinuation(Guid.NewGuid().ToString(), feedRangeEPK, new List>() { feedRangeEPK.Range }, continuation); FeedRangeIteratorCore feedTokenIterator = FeedRangeIteratorCore.Create(Mock.Of(), null, feedRangeSimpleContinuation.ToString(), null); - FeedRangeEPK defaultRange = feedTokenIterator.FeedRangeInternal as FeedRangeEPK; - Assert.AreEqual(FeedRangeEPK.ForFullRange().Range.Min, defaultRange.Range.Min); - Assert.AreEqual(FeedRangeEPK.ForFullRange().Range.Max, defaultRange.Range.Max); + FeedRangeEpk defaultRange = feedTokenIterator.FeedRangeInternal as FeedRangeEpk; + Assert.AreEqual(FeedRangeEpk.ForFullRange().Range.Min, defaultRange.Range.Min); + Assert.AreEqual(FeedRangeEpk.ForFullRange().Range.Max, defaultRange.Range.Max); Assert.IsNotNull(feedTokenIterator.FeedRangeContinuation); Assert.AreEqual(continuation, feedTokenIterator.FeedRangeContinuation.GetContinuation()); } @@ -477,7 +477,7 @@ public async Task ReadFeedIteratorCore_WithNoInitialState_ReadNextAsync() Assert.IsTrue(FeedRangeContinuation.TryParse(response.ContinuationToken, out FeedRangeContinuation parsedToken)); FeedRangeCompositeContinuation feedRangeCompositeContinuation = parsedToken as FeedRangeCompositeContinuation; - FeedRangeEPK feedTokenEPKRange = feedRangeCompositeContinuation.FeedRange as FeedRangeEPK; + FeedRangeEpk feedTokenEPKRange = feedRangeCompositeContinuation.FeedRange as FeedRangeEpk; // Assert that a FeedToken for the entire range is used Assert.AreEqual(Documents.Routing.PartitionKeyInternal.MinimumInclusiveEffectivePartitionKey, feedTokenEPKRange.Range.Min); Assert.AreEqual(Documents.Routing.PartitionKeyInternal.MaximumExclusiveEffectivePartitionKey, feedTokenEPKRange.Range.Max); From 567b8ff240b12bfbd9367131079d432a25325e75 Mon Sep 17 00:00:00 2001 From: Brandon Chong Date: Mon, 22 Jun 2020 19:34:10 -0700 Subject: [PATCH 02/85] drafted out abstract classes --- .../CrossPartitionRangePaginator.cs | 97 +++++++++++++++++++ .../src/Pagination/FeedRangeProvider.cs | 55 +++++++++++ Microsoft.Azure.Cosmos/src/Pagination/Page.cs | 14 +++ .../src/Pagination/PartitionRangePaginator.cs | 24 +++++ .../src/Pagination/State.cs | 10 ++ 5 files changed, 200 insertions(+) create mode 100644 Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionRangePaginator.cs create mode 100644 Microsoft.Azure.Cosmos/src/Pagination/FeedRangeProvider.cs create mode 100644 Microsoft.Azure.Cosmos/src/Pagination/Page.cs create mode 100644 Microsoft.Azure.Cosmos/src/Pagination/PartitionRangePaginator.cs create mode 100644 Microsoft.Azure.Cosmos/src/Pagination/State.cs diff --git a/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionRangePaginator.cs b/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionRangePaginator.cs new file mode 100644 index 0000000000..91f111f915 --- /dev/null +++ b/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionRangePaginator.cs @@ -0,0 +1,97 @@ +// ------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +// ------------------------------------------------------------ + +namespace Microsoft.Azure.Cosmos.Pagination +{ + using System; + using System.Collections.Generic; + using System.Net; + using System.Threading; + using System.Threading.Tasks; + using Microsoft.Azure.Cosmos.Query.Core.Collections; + using Microsoft.Azure.Cosmos.Query.Core.Monads; + + /// + /// Coordinates draining pages from multiple , while maintaining a global sort order and handling repartitioning (splits, merge). + /// + internal abstract class CrossPartitionRangePaginator + { + private readonly FeedRangeProvider feedRangeProvider; + private readonly Func createPartitionRangePaginator; + private readonly PriorityQueue paginators; + + public CrossPartitionRangePaginator( + FeedRangeProvider feedRangeProvider, + Func createPartitionRangePaginator, + IEnumerable paginators, + IComparer comparer) + { + this.feedRangeProvider = feedRangeProvider ?? throw new ArgumentNullException(nameof(feedRangeProvider)); + this.createPartitionRangePaginator = createPartitionRangePaginator ?? throw new ArgumentNullException(nameof(createPartitionRangePaginator)); + + if (paginators == null) + { + throw new ArgumentNullException(nameof(paginators)); + } + + if (comparer == null) + { + throw new ArgumentNullException(nameof(paginators)); + } + + this.paginators = new PriorityQueue(paginators, comparer); + } + + public Page CurrentPage { get; set; } + + public async Task TryMoveNextPageAsync(CancellationToken cancellationToken = default) + { + cancellationToken.ThrowIfCancellationRequested(); + PartitionRangePaginator currentPaginator = this.paginators.Dequeue(); + TryCatch tryMoveNextPageAsync = await currentPaginator.TryMoveNextPageAsync(cancellationToken); + if (tryMoveNextPageAsync.Failed) + { + Exception exception = tryMoveNextPageAsync.Exception; + if (IsSplitException(exception)) + { + // Handle split + IEnumerable childRanges = await this.feedRangeProvider.GetChildRangeAsync(currentPaginator.FeedRange, cancellationToken); + foreach (FeedRange childRange in childRanges) + { + PartitionRangePaginator childPaginator = this.createPartitionRangePaginator(childRange, currentPaginator.GetState()); + this.paginators.Enqueue(childPaginator); + } + + // Recursively retry + return await this.TryMoveNextPageAsync(cancellationToken); + } + + if (IsMergeException(exception)) + { + throw new NotImplementedException(); + } + + return tryMoveNextPageAsync; + } + + this.CurrentPage = currentPaginator.CurrentPage; + this.paginators.Enqueue(currentPaginator); + + return TryCatch.FromResult(); + } + + private static bool IsSplitException(Exception exeception) + { + return exeception is CosmosException cosmosException + && cosmosException.StatusCode == HttpStatusCode.Gone + && cosmosException.SubStatusCode == (int)Documents.SubStatusCodes.PartitionKeyRangeGone; + } + + private static bool IsMergeException(Exception exception) + { + // TODO: code this out + return false; + } + } +} diff --git a/Microsoft.Azure.Cosmos/src/Pagination/FeedRangeProvider.cs b/Microsoft.Azure.Cosmos/src/Pagination/FeedRangeProvider.cs new file mode 100644 index 0000000000..1e6e556236 --- /dev/null +++ b/Microsoft.Azure.Cosmos/src/Pagination/FeedRangeProvider.cs @@ -0,0 +1,55 @@ +// ------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +// ------------------------------------------------------------ + +namespace Microsoft.Azure.Cosmos.Pagination +{ + using System; + using System.Collections.Generic; + using System.Threading; + using System.Threading.Tasks; + using Microsoft.Azure.Cosmos.Query.Core.QueryClient; + + internal sealed class FeedRangeProvider + { + private readonly CosmosQueryClient cosmosQueryClient; + private readonly string collectionRid; + + public FeedRangeProvider(CosmosQueryClient cosmosQueryClient, string collectionRid) + { + this.cosmosQueryClient = cosmosQueryClient ?? throw new ArgumentNullException(nameof(cosmosQueryClient)); + this.collectionRid = collectionRid ?? throw new ArgumentNullException(nameof(collectionRid)); + } + + public async Task> GetChildRangeAsync( + FeedRange feedRange, + CancellationToken cancellationToken = default) + { + cancellationToken.ThrowIfCancellationRequested(); + + if (feedRange == null) + { + throw new ArgumentNullException(nameof(feedRange)); + } + + if (!(feedRange is FeedRangeEpk feedRangeEPK)) + { + throw new ArgumentOutOfRangeException(nameof(feedRange)); + } + + IReadOnlyList replacementRanges = await this.cosmosQueryClient.TryGetOverlappingRangesAsync( + this.collectionRid, + new Documents.Routing.Range(feedRangeEPK.Range.Min, feedRangeEPK.Range.Max, isMaxInclusive: true, isMinInclusive: false), + forceRefresh: true); + + List childFeedRanges = new List(replacementRanges.Count); + + foreach (Documents.PartitionKeyRange replacementRange in replacementRanges) + { + childFeedRanges.Add(new FeedRangeEpk(replacementRange.ToRange())); + } + + return childFeedRanges; + } + } +} diff --git a/Microsoft.Azure.Cosmos/src/Pagination/Page.cs b/Microsoft.Azure.Cosmos/src/Pagination/Page.cs new file mode 100644 index 0000000000..24a9cd0e36 --- /dev/null +++ b/Microsoft.Azure.Cosmos/src/Pagination/Page.cs @@ -0,0 +1,14 @@ +// ------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +// ------------------------------------------------------------ + +namespace Microsoft.Azure.Cosmos.Pagination +{ + using System; + using System.Collections.Generic; + using System.Text; + + internal abstract class Page + { + } +} diff --git a/Microsoft.Azure.Cosmos/src/Pagination/PartitionRangePaginator.cs b/Microsoft.Azure.Cosmos/src/Pagination/PartitionRangePaginator.cs new file mode 100644 index 0000000000..98d7f353b0 --- /dev/null +++ b/Microsoft.Azure.Cosmos/src/Pagination/PartitionRangePaginator.cs @@ -0,0 +1,24 @@ +// ------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +// ------------------------------------------------------------ + +namespace Microsoft.Azure.Cosmos.Pagination +{ + using System.Threading; + using System.Threading.Tasks; + using Microsoft.Azure.Cosmos.Query.Core.Monads; + + /// + /// Has the ability to page through a partition range. + /// + internal abstract class PartitionRangePaginator + { + public FeedRange FeedRange { get; } + + public Page CurrentPage { get; } + + public abstract State GetState(); + + public abstract Task TryMoveNextPageAsync(CancellationToken cancellationToken = default); + } +} diff --git a/Microsoft.Azure.Cosmos/src/Pagination/State.cs b/Microsoft.Azure.Cosmos/src/Pagination/State.cs new file mode 100644 index 0000000000..f62edbf2eb --- /dev/null +++ b/Microsoft.Azure.Cosmos/src/Pagination/State.cs @@ -0,0 +1,10 @@ +namespace Microsoft.Azure.Cosmos.Pagination +{ + using System; + using System.Collections.Generic; + using System.Text; + + internal abstract class State + { + } +} From 1ad7e9cae81789c09b8909f5eed4904e6e29b3d4 Mon Sep 17 00:00:00 2001 From: Brandon Chong Date: Tue, 23 Jun 2020 11:10:43 -0700 Subject: [PATCH 03/85] drafted out basic in memory collection --- .../src/Routing/PartitionKeyHash.cs | 12 + .../src/Routing/PartitionKeyHashRange.cs | 7 + .../PartitionKeyHashRangeDictionary.cs | 89 ++++++++ .../src/Routing/PartitionKeyHashRanges.cs | 35 +-- .../InMemoryCollection.cs | 208 ++++++++++++++++++ 5 files changed, 326 insertions(+), 25 deletions(-) create mode 100644 Microsoft.Azure.Cosmos/src/Routing/PartitionKeyHashRangeDictionary.cs create mode 100644 Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/InMemoryCollection.cs diff --git a/Microsoft.Azure.Cosmos/src/Routing/PartitionKeyHash.cs b/Microsoft.Azure.Cosmos/src/Routing/PartitionKeyHash.cs index 6554fbb708..2633236978 100644 --- a/Microsoft.Azure.Cosmos/src/Routing/PartitionKeyHash.cs +++ b/Microsoft.Azure.Cosmos/src/Routing/PartitionKeyHash.cs @@ -198,5 +198,17 @@ private static PartitionKeyHash Hash(ReadOnlySpan bytesForHashing) return new PartitionKeyHash(hash); } } + + public static bool operator ==(PartitionKeyHash left, PartitionKeyHash right) => left.Equals(right); + + public static bool operator !=(PartitionKeyHash left, PartitionKeyHash right) => !(left == right); + + public static bool operator <(PartitionKeyHash left, PartitionKeyHash right) => left.CompareTo(right) < 0; + + public static bool operator <=(PartitionKeyHash left, PartitionKeyHash right) => left.CompareTo(right) <= 0; + + public static bool operator >(PartitionKeyHash left, PartitionKeyHash right) => left.CompareTo(right) > 0; + + public static bool operator >=(PartitionKeyHash left, PartitionKeyHash right) => left.CompareTo(right) >= 0; } } diff --git a/Microsoft.Azure.Cosmos/src/Routing/PartitionKeyHashRange.cs b/Microsoft.Azure.Cosmos/src/Routing/PartitionKeyHashRange.cs index 23fcea8c02..a0aecb5676 100644 --- a/Microsoft.Azure.Cosmos/src/Routing/PartitionKeyHashRange.cs +++ b/Microsoft.Azure.Cosmos/src/Routing/PartitionKeyHashRange.cs @@ -29,6 +29,13 @@ public PartitionKeyHashRange(PartitionKeyHash? startInclusive, PartitionKeyHash? public PartitionKeyHash? EndExclusive { get; } + public bool Contains(PartitionKeyHash partitionKeyHash) + { + bool rangeStartsBefore = !this.StartInclusive.HasValue || (this.StartInclusive.Value <= partitionKeyHash); + bool rangeEndsAfter = !this.EndExclusive.HasValue || (partitionKeyHash <= this.EndExclusive.Value); + return rangeStartsBefore && rangeEndsAfter; + } + public int CompareTo(PartitionKeyHashRange other) { // Provide a total sort order by first comparing on the start and then going to the end. diff --git a/Microsoft.Azure.Cosmos/src/Routing/PartitionKeyHashRangeDictionary.cs b/Microsoft.Azure.Cosmos/src/Routing/PartitionKeyHashRangeDictionary.cs new file mode 100644 index 0000000000..7282868dac --- /dev/null +++ b/Microsoft.Azure.Cosmos/src/Routing/PartitionKeyHashRangeDictionary.cs @@ -0,0 +1,89 @@ +// ------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +// ------------------------------------------------------------ + +namespace Microsoft.Azure.Cosmos.Routing +{ + using System; + using System.Collections.Generic; + + internal sealed class PartitionKeyHashRangeDictionary + { + private readonly SortedDictionary dictionary; + + public PartitionKeyHashRangeDictionary(PartitionKeyHashRanges partitionKeyHashRanges) + { + if (partitionKeyHashRanges == null) + { + throw new ArgumentNullException(nameof(partitionKeyHashRanges)); + } + + this.dictionary = new SortedDictionary(); + foreach (PartitionKeyHashRange partitionKeyHashRange in partitionKeyHashRanges) + { + this.dictionary.Add(partitionKeyHashRange, (false, default)); + } + } + + public bool TryGetValue(PartitionKeyHash partitionKeyHash, out T value) + { + if (!this.TryGetContainingRange(partitionKeyHash, out PartitionKeyHashRange range)) + { + value = default; + return false; + } + + if (!this.dictionary.TryGetValue(range, out (bool valueSet, T value) nullableValue)) + { + value = default; + return false; + } + + if (!nullableValue.valueSet) + { + value = default; + return false; + } + + value = nullableValue.value; + return true; + } + + public T this[PartitionKeyHash key] + { + get + { + if (!this.TryGetValue(key, out T value)) + { + throw new KeyNotFoundException(); + } + + return value; + } + set + { + if (!this.TryGetContainingRange(key, out PartitionKeyHashRange range)) + { + throw new NotSupportedException("Dictionary does not support adding new elements."); + } + + this.dictionary[range] = (true, value); + } + } + + private bool TryGetContainingRange(PartitionKeyHash partitionKeyHash, out PartitionKeyHashRange range) + { + foreach (PartitionKeyHashRange candidateRange in this.dictionary.Keys) + { + if (candidateRange.Contains(partitionKeyHash)) + { + range = candidateRange; + return true; + } + } + + range = default; + return false; + } + } +} diff --git a/Microsoft.Azure.Cosmos/src/Routing/PartitionKeyHashRanges.cs b/Microsoft.Azure.Cosmos/src/Routing/PartitionKeyHashRanges.cs index 0ae0e37cdb..22e15f8aaf 100644 --- a/Microsoft.Azure.Cosmos/src/Routing/PartitionKeyHashRanges.cs +++ b/Microsoft.Azure.Cosmos/src/Routing/PartitionKeyHashRanges.cs @@ -34,32 +34,17 @@ public static PartitionKeyHashRanges Create(IEnumerable p partitionKeyHashRanges, out PartitionKeyHashRanges partitionedSortedEffectiveRanges); - switch (createStatus) + return createStatus switch { - case CreateOutcome.DuplicatePartitionKeyRange: - throw new ArgumentException($"{nameof(partitionKeyHashRanges)} must not have duplicate values."); - - case CreateOutcome.EmptyPartitionKeyRange: - throw new ArgumentException($"{nameof(partitionKeyHashRanges)} must not have an empty range."); - - case CreateOutcome.NoPartitionKeyRanges: - throw new ArgumentException($"{nameof(partitionKeyHashRanges)} must not be empty."); - - case CreateOutcome.NullPartitionKeyRanges: - throw new ArgumentNullException(nameof(partitionKeyHashRanges)); - - case CreateOutcome.RangesAreNotContiguous: - throw new ArgumentException($"{nameof(partitionKeyHashRanges)} must have contiguous ranges."); - - case CreateOutcome.RangesOverlap: - throw new ArgumentException($"{nameof(partitionKeyHashRanges)} must not overlapping ranges."); - - case CreateOutcome.Success: - return partitionedSortedEffectiveRanges; - - default: - throw new ArgumentOutOfRangeException($"Unknown {nameof(CreateOutcome)}: {createStatus}."); - } + CreateOutcome.DuplicatePartitionKeyRange => throw new ArgumentException($"{nameof(partitionKeyHashRanges)} must not have duplicate values."), + CreateOutcome.EmptyPartitionKeyRange => throw new ArgumentException($"{nameof(partitionKeyHashRanges)} must not have an empty range."), + CreateOutcome.NoPartitionKeyRanges => throw new ArgumentException($"{nameof(partitionKeyHashRanges)} must not be empty."), + CreateOutcome.NullPartitionKeyRanges => throw new ArgumentNullException(nameof(partitionKeyHashRanges)), + CreateOutcome.RangesAreNotContiguous => throw new ArgumentException($"{nameof(partitionKeyHashRanges)} must have contiguous ranges."), + CreateOutcome.RangesOverlap => throw new ArgumentException($"{nameof(partitionKeyHashRanges)} must not overlapping ranges."), + CreateOutcome.Success => partitionedSortedEffectiveRanges, + _ => throw new ArgumentOutOfRangeException($"Unknown {nameof(CreateOutcome)}: {createStatus}."), + }; } public static CreateOutcome TryCreate( diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/InMemoryCollection.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/InMemoryCollection.cs new file mode 100644 index 0000000000..a2009049bc --- /dev/null +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/InMemoryCollection.cs @@ -0,0 +1,208 @@ +//------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +//------------------------------------------------------------ + +namespace Microsoft.Azure.Cosmos.Tests +{ + using System; + using System.Collections; + using System.Collections.Generic; + using Microsoft.Azure.Cosmos.CosmosElements; + using Microsoft.Azure.Cosmos.Routing; + using Microsoft.Azure.Documents; + + // Collection useful for mocking requests and repartitioning (splits / merge). + internal sealed class InMemoryContainer + { + private readonly PartitionKeyHashRangeDictionary partitionedRecords; + private readonly PartitionKeyDefinition partitionKeyDefinition; + + public InMemoryContainer(PartitionKeyDefinition partitionKeyDefinition) + { + PartitionKeyHashRange fullRange = new PartitionKeyHashRange(startInclusive: null, endExclusive: null); + PartitionKeyHashRanges partitionKeyHashRanges = PartitionKeyHashRanges.Create(new PartitionKeyHashRange[] { fullRange }); + this.partitionedRecords = new PartitionKeyHashRangeDictionary(partitionKeyHashRanges); + this.partitionKeyDefinition = partitionKeyDefinition ?? throw new ArgumentNullException(nameof(partitionKeyDefinition)); + } + + public Record CreateItem(CosmosObject payload) + { + if (payload == null) + { + throw new ArgumentNullException(nameof(payload)); + } + + PartitionKeyHash partitionKeyHash = GetHashFromPayload(payload, this.partitionKeyDefinition); + if (!this.partitionedRecords.TryGetValue(partitionKeyHash, out Records records)) + { + this.partitionedRecords[partitionKeyHash] = new Records(); + } + + return records.Add(payload); + } + + public bool TryReadItem(CosmosElement partitionKey, Guid identifier, out Record record) + { + PartitionKeyHash partitionKeyHash = GetHashFromPartitionKey(partitionKey, this.partitionKeyDefinition); + if (!this.partitionedRecords.TryGetValue(partitionKeyHash, out Records records)) + { + record = default; + return false; + } + + foreach (Record candidate in records) + { + if (candidate.Identifier == identifier) + { + record = candidate; + return true; + } + } + + record = default; + return false; + } + + private static PartitionKeyHash GetHashFromPayload(CosmosObject payload, PartitionKeyDefinition partitionKeyDefinition) + { + // Restrict the partition key definition for now to keep things simple + if (partitionKeyDefinition.Kind != PartitionKind.Hash) + { + throw new ArgumentOutOfRangeException("Can only support hash partitioning"); + } + + if (partitionKeyDefinition.Version != PartitionKeyDefinitionVersion.V2) + { + throw new ArgumentOutOfRangeException("Can only support hash v2"); + } + + if (partitionKeyDefinition.Paths.Count != 1) + { + throw new ArgumentOutOfRangeException("Can only support a single partition key path."); + } + + string[] tokens = partitionKeyDefinition.Paths[0].Split("/"); + + CosmosElement partitionKey = payload; + foreach (string token in tokens) + { + if (partitionKey != default) + { + if (!payload.TryGetValue(token, out partitionKey)) + { + partitionKey = default; + } + } + } + + return GetHashFromPartitionKey(partitionKey, partitionKeyDefinition); + } + + private static PartitionKeyHash GetHashFromPartitionKey(CosmosElement partitionKey, PartitionKeyDefinition partitionKeyDefinition) + { + // Restrict the partition key definition for now to keep things simple + if (partitionKeyDefinition.Kind != PartitionKind.Hash) + { + throw new ArgumentOutOfRangeException("Can only support hash partitioning"); + } + + if (partitionKeyDefinition.Version != PartitionKeyDefinitionVersion.V2) + { + throw new ArgumentOutOfRangeException("Can only support hash v2"); + } + + if (partitionKeyDefinition.Paths.Count != 1) + { + throw new ArgumentOutOfRangeException("Can only support a single partition key path."); + } + + PartitionKeyHash partitionKeyHash; + switch (partitionKey) + { + case null: + partitionKeyHash = PartitionKeyHash.V2.HashUndefined(); + break; + + case CosmosString stringPartitionKey: + partitionKeyHash = PartitionKeyHash.V2.Hash(stringPartitionKey.Value); + break; + + case CosmosNumber numberPartitionKey: + partitionKeyHash = PartitionKeyHash.V2.Hash(Number64.ToDouble(numberPartitionKey.Value)); + break; + + case CosmosBoolean cosmosBoolean: + partitionKeyHash = PartitionKeyHash.V2.Hash(cosmosBoolean.Value); + break; + + case CosmosNull _: + partitionKeyHash = PartitionKeyHash.V2.HashNull(); + break; + + default: + throw new ArgumentOutOfRangeException(); + } + + return partitionKeyHash; + } + + public sealed class Record + { + private Record(long resourceIdentifier, long timestamp, Guid identifier, CosmosObject payload) + { + this.ResourceIdentifier = resourceIdentifier < 0 ? throw new ArgumentOutOfRangeException(nameof(resourceIdentifier)) : resourceIdentifier; + this.Timestamp = timestamp < 0 ? throw new ArgumentOutOfRangeException(nameof(timestamp)) : timestamp; + this.Identifier = identifier; + this.Payload = payload ?? throw new ArgumentNullException(nameof(payload)); + } + + public long ResourceIdentifier { get; } + + public long Timestamp { get; } + + public Guid Identifier { get; } + + public CosmosObject Payload { get; } + + public static Record Create(long previousResourceIdentifier, CosmosObject payload) + { + return new Record(previousResourceIdentifier + 1, DateTime.UtcNow.Ticks, Guid.NewGuid(), payload); + } + } + + private sealed class Records : IReadOnlyList + { + private readonly List storage; + + public Records() + { + this.storage = new List(); + } + + public Record this[int index] => this.storage[index]; + + public int Count => this.storage.Count; + + public IEnumerator GetEnumerator() => this.storage.GetEnumerator(); + + IEnumerator IEnumerable.GetEnumerator() => this.storage.GetEnumerator(); + + public Record Add(CosmosObject payload) + { + long previousResourceId; + if (this.Count == 0) + { + previousResourceId = 0; + } + else + { + previousResourceId = this.storage[this.storage.Count - 1].ResourceIdentifier; + } + + Record record = Record.Create(previousResourceId, payload); + this.storage.Add(record); + return record; + } + } + } +} From 3e8f9f6001fe75d7ff6e79df4e15b7d44f340bd1 Mon Sep 17 00:00:00 2001 From: Brandon Chong Date: Tue, 23 Jun 2020 13:54:15 -0700 Subject: [PATCH 04/85] got basic crud test working --- .../src/Pagination/State.cs | 6 ++- .../src/Routing/PartitionKeyRangeCache.cs | 2 +- Microsoft.Azure.Cosmos/src/UInt128.cs | 12 ++--- .../InMemoryCollection.cs | 10 ++-- .../InMemoryCollectionTests.cs | 48 +++++++++++++++++++ .../Microsoft.Azure.Cosmos.Tests.csproj | 4 ++ 6 files changed, 70 insertions(+), 12 deletions(-) create mode 100644 Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/InMemoryCollectionTests.cs diff --git a/Microsoft.Azure.Cosmos/src/Pagination/State.cs b/Microsoft.Azure.Cosmos/src/Pagination/State.cs index f62edbf2eb..f91cf1278c 100644 --- a/Microsoft.Azure.Cosmos/src/Pagination/State.cs +++ b/Microsoft.Azure.Cosmos/src/Pagination/State.cs @@ -1,4 +1,8 @@ -namespace Microsoft.Azure.Cosmos.Pagination +// ------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +// ------------------------------------------------------------ + +namespace Microsoft.Azure.Cosmos.Pagination { using System; using System.Collections.Generic; diff --git a/Microsoft.Azure.Cosmos/src/Routing/PartitionKeyRangeCache.cs b/Microsoft.Azure.Cosmos/src/Routing/PartitionKeyRangeCache.cs index 81e66d8508..22b58a5ee3 100644 --- a/Microsoft.Azure.Cosmos/src/Routing/PartitionKeyRangeCache.cs +++ b/Microsoft.Azure.Cosmos/src/Routing/PartitionKeyRangeCache.cs @@ -170,7 +170,7 @@ private async Task GetRoutingMapForCollectionAsync( RetryOptions retryOptions = new RetryOptions(); using (DocumentServiceResponse response = await BackoffRetryUtility.ExecuteAsync( - () => ExecutePartitionKeyRangeReadChangeFeedAsync(collectionRid, headers), + () => this.ExecutePartitionKeyRangeReadChangeFeedAsync(collectionRid, headers), new ResourceThrottleRetryPolicy(retryOptions.MaxRetryAttemptsOnThrottledRequests, retryOptions.MaxRetryWaitTimeInSeconds), cancellationToken)) { diff --git a/Microsoft.Azure.Cosmos/src/UInt128.cs b/Microsoft.Azure.Cosmos/src/UInt128.cs index 804fd38846..d11b5c9025 100644 --- a/Microsoft.Azure.Cosmos/src/UInt128.cs +++ b/Microsoft.Azure.Cosmos/src/UInt128.cs @@ -421,12 +421,12 @@ public int CompareTo(object value) return 1; } - if (value is UInt128) + if (!(value is UInt128 uint128)) { - return this.CompareTo((UInt128)value); + throw new ArgumentException("Value must be a UInt128."); } - throw new ArgumentException("Value must be a UInt128."); + return this.CompareTo(uint128); } /// @@ -465,12 +465,12 @@ public override bool Equals(object obj) return true; } - if (obj is UInt128) + if (!(obj is UInt128 uint128)) { - return this.Equals((UInt128)obj); + return false; } - return false; + return this.Equals(uint128); } /// diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/InMemoryCollection.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/InMemoryCollection.cs index a2009049bc..d06c84d776 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/InMemoryCollection.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/InMemoryCollection.cs @@ -7,17 +7,18 @@ namespace Microsoft.Azure.Cosmos.Tests using System; using System.Collections; using System.Collections.Generic; + using System.Linq; using Microsoft.Azure.Cosmos.CosmosElements; using Microsoft.Azure.Cosmos.Routing; using Microsoft.Azure.Documents; // Collection useful for mocking requests and repartitioning (splits / merge). - internal sealed class InMemoryContainer + internal sealed class InMemoryCollection { private readonly PartitionKeyHashRangeDictionary partitionedRecords; private readonly PartitionKeyDefinition partitionKeyDefinition; - public InMemoryContainer(PartitionKeyDefinition partitionKeyDefinition) + public InMemoryCollection(PartitionKeyDefinition partitionKeyDefinition) { PartitionKeyHashRange fullRange = new PartitionKeyHashRange(startInclusive: null, endExclusive: null); PartitionKeyHashRanges partitionKeyHashRanges = PartitionKeyHashRanges.Create(new PartitionKeyHashRange[] { fullRange }); @@ -35,7 +36,8 @@ public Record CreateItem(CosmosObject payload) PartitionKeyHash partitionKeyHash = GetHashFromPayload(payload, this.partitionKeyDefinition); if (!this.partitionedRecords.TryGetValue(partitionKeyHash, out Records records)) { - this.partitionedRecords[partitionKeyHash] = new Records(); + records = new Records(); + this.partitionedRecords[partitionKeyHash] = records; } return records.Add(payload); @@ -81,7 +83,7 @@ private static PartitionKeyHash GetHashFromPayload(CosmosObject payload, Partiti throw new ArgumentOutOfRangeException("Can only support a single partition key path."); } - string[] tokens = partitionKeyDefinition.Paths[0].Split("/"); + IEnumerable tokens = partitionKeyDefinition.Paths[0].Split("/").Skip(1); CosmosElement partitionKey = payload; foreach (string token in tokens) diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/InMemoryCollectionTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/InMemoryCollectionTests.cs new file mode 100644 index 0000000000..dfd24a6fdd --- /dev/null +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/InMemoryCollectionTests.cs @@ -0,0 +1,48 @@ +//------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +//------------------------------------------------------------ + +namespace Microsoft.Azure.Cosmos.Tests +{ + using System; + using Microsoft.Azure.Cosmos.CosmosElements; + using Microsoft.Azure.Cosmos.CosmosElements.Numbers; + using Microsoft.Azure.Documents; + using Microsoft.VisualStudio.TestTools.UnitTesting; + + [TestClass] + public class InMemoryCollectionTests + { + [TestMethod] + public void TestCrud() + { + PartitionKeyDefinition partitionKeyDefinition = new PartitionKeyDefinition() + { + Paths = new System.Collections.ObjectModel.Collection() + { + "/pk" + }, + Kind = PartitionKind.Hash, + Version = PartitionKeyDefinitionVersion.V2, + }; + + InMemoryCollection inMemoryCollection = new InMemoryCollection(partitionKeyDefinition); + + // Insert an item + CosmosObject item = CosmosObject.Parse("{\"pk\" : 42 }"); + InMemoryCollection.Record record = inMemoryCollection.CreateItem(item); + Assert.IsNotNull(record); + Assert.AreNotEqual(Guid.Empty, record.Identifier); + Assert.AreEqual(1, record.ResourceIdentifier); + + // Try to read it back + Assert.IsTrue( + inMemoryCollection.TryReadItem( + partitionKey: CosmosNumber64.Create(42), + record.Identifier, + out InMemoryCollection.Record readRecord)); + + Assert.AreEqual(item.ToString(), readRecord.Payload.ToString()); + } + } +} 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 522c34ab9c..ed847de63e 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 @@ -239,6 +239,10 @@ + + + + true true From 2f3fef01afb23fc352dcb3e3b9d363ad4b0a38d6 Mon Sep 17 00:00:00 2001 From: Brandon Chong Date: Tue, 23 Jun 2020 16:31:14 -0700 Subject: [PATCH 05/85] implemented split and readfeed --- .../PartitionKeyHashRangeDictionary.cs | 29 ++++- .../InMemoryCollection.cs | 91 +++++++++++++- .../InMemoryCollectionTests.cs | 113 ++++++++++++++++++ 3 files changed, 229 insertions(+), 4 deletions(-) diff --git a/Microsoft.Azure.Cosmos/src/Routing/PartitionKeyHashRangeDictionary.cs b/Microsoft.Azure.Cosmos/src/Routing/PartitionKeyHashRangeDictionary.cs index 7282868dac..44c920cc28 100644 --- a/Microsoft.Azure.Cosmos/src/Routing/PartitionKeyHashRangeDictionary.cs +++ b/Microsoft.Azure.Cosmos/src/Routing/PartitionKeyHashRangeDictionary.cs @@ -33,7 +33,12 @@ public bool TryGetValue(PartitionKeyHash partitionKeyHash, out T value) return false; } - if (!this.dictionary.TryGetValue(range, out (bool valueSet, T value) nullableValue)) + return this.TryGetValue(range, out value); + } + + public bool TryGetValue(PartitionKeyHashRange partitionKeyHashRange, out T value) + { + if (!this.dictionary.TryGetValue(partitionKeyHashRange, out (bool valueSet, T value) nullableValue)) { value = default; return false; @@ -71,6 +76,28 @@ public T this[PartitionKeyHash key] } } + public T this[PartitionKeyHashRange key] + { + get + { + if (!this.TryGetValue(key, out T value)) + { + throw new KeyNotFoundException(); + } + + return value; + } + set + { + if (!this.TryGetValue(key, out _)) + { + throw new NotSupportedException("Dictionary does not support adding new elements."); + } + + this.dictionary[key] = (true, value); + } + } + private bool TryGetContainingRange(PartitionKeyHash partitionKeyHash, out PartitionKeyHashRange range) { foreach (PartitionKeyHashRange candidateRange in this.dictionary.Keys) diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/InMemoryCollection.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/InMemoryCollection.cs index d06c84d776..df35084c9f 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/InMemoryCollection.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/InMemoryCollection.cs @@ -11,19 +11,26 @@ namespace Microsoft.Azure.Cosmos.Tests using Microsoft.Azure.Cosmos.CosmosElements; using Microsoft.Azure.Cosmos.Routing; using Microsoft.Azure.Documents; + using Moq; // Collection useful for mocking requests and repartitioning (splits / merge). internal sealed class InMemoryCollection { - private readonly PartitionKeyHashRangeDictionary partitionedRecords; private readonly PartitionKeyDefinition partitionKeyDefinition; + private PartitionKeyHashRangeDictionary partitionedRecords; + private Dictionary partitionKeyRangeIdToHashRange; + public InMemoryCollection(PartitionKeyDefinition partitionKeyDefinition) { PartitionKeyHashRange fullRange = new PartitionKeyHashRange(startInclusive: null, endExclusive: null); PartitionKeyHashRanges partitionKeyHashRanges = PartitionKeyHashRanges.Create(new PartitionKeyHashRange[] { fullRange }); this.partitionedRecords = new PartitionKeyHashRangeDictionary(partitionKeyHashRanges); this.partitionKeyDefinition = partitionKeyDefinition ?? throw new ArgumentNullException(nameof(partitionKeyDefinition)); + this.partitionKeyRangeIdToHashRange = new Dictionary() + { + { 0, fullRange } + }; } public Record CreateItem(CosmosObject payload) @@ -54,7 +61,12 @@ record = default; foreach (Record candidate in records) { - if (candidate.Identifier == identifier) + bool identifierMatches = candidate.Identifier == identifier; + + CosmosElement candidatePartitionKey = GetPartitionKeyFromPayload(candidate.Payload, this.partitionKeyDefinition); + bool partitionKeyMatches = CosmosElementEqualityComparer.Value.Equals(candidatePartitionKey, partitionKey); + + if (identifierMatches && partitionKeyMatches) { record = candidate; return true; @@ -65,7 +77,80 @@ record = default; return false; } + public bool TryReadFeed(int partitionKeyRangeId, int pageIndex, int pageSize, out List page) + { + if (!this.partitionKeyRangeIdToHashRange.TryGetValue(partitionKeyRangeId, out PartitionKeyHashRange range)) + { + page = default; + return false; + } + + if (!this.partitionedRecords.TryGetValue(range, out Records records)) + { + throw new InvalidOperationException("failed to find the range."); + } + + page = records.Skip(pageIndex * pageSize).Take(pageSize).ToList(); + return true; + } + + public IReadOnlyDictionary PartitionKeyRangeFeedReed() => this.partitionKeyRangeIdToHashRange; + + public void Split(int partitionKeyRangeId) + { + // Get the current range and records + if (!this.partitionKeyRangeIdToHashRange.TryGetValue(partitionKeyRangeId, out PartitionKeyHashRange parentRange)) + { + throw new InvalidOperationException("Failed to find the range."); + } + + if (!this.partitionedRecords.TryGetValue(parentRange, out Records records)) + { + throw new InvalidOperationException("failed to find the range."); + } + + int maxPartitionKeyRangeId = this.partitionKeyRangeIdToHashRange.Keys.Max(); + + // Split the range space + PartitionKeyHashRanges partitionKeyHashRanges = PartitionKeyHashRangeSplitterAndMerger.SplitRange(parentRange, 2); + + // Update the partition routing map + Dictionary newPartitionKeyRangeIdToHashRange = new Dictionary() + { + { maxPartitionKeyRangeId + 1, partitionKeyHashRanges.First() }, + { maxPartitionKeyRangeId + 2, partitionKeyHashRanges.Last() }, + }; + + // Copy over the partitioned records (minus the parent range) + PartitionKeyHashRangeDictionary newPartitionedRecords = new PartitionKeyHashRangeDictionary( + PartitionKeyHashRanges.Create( + newPartitionKeyRangeIdToHashRange.Values)); + + foreach (PartitionKeyHashRange range in this.partitionKeyRangeIdToHashRange.Values) + { + if (!range.Equals(parentRange)) + { + newPartitionedRecords[range] = this.partitionedRecords[range]; + } + } + + this.partitionedRecords = newPartitionedRecords; + this.partitionKeyRangeIdToHashRange = newPartitionKeyRangeIdToHashRange; + + // Rehash the records in the parent range + foreach (Record record in records) + { + this.CreateItem(record.Payload); + } + } + private static PartitionKeyHash GetHashFromPayload(CosmosObject payload, PartitionKeyDefinition partitionKeyDefinition) + { + CosmosElement partitionKey = GetPartitionKeyFromPayload(payload, partitionKeyDefinition); + return GetHashFromPartitionKey(partitionKey, partitionKeyDefinition); + } + + private static CosmosElement GetPartitionKeyFromPayload(CosmosObject payload, PartitionKeyDefinition partitionKeyDefinition) { // Restrict the partition key definition for now to keep things simple if (partitionKeyDefinition.Kind != PartitionKind.Hash) @@ -97,7 +182,7 @@ private static PartitionKeyHash GetHashFromPayload(CosmosObject payload, Partiti } } - return GetHashFromPartitionKey(partitionKey, partitionKeyDefinition); + return partitionKey; } private static PartitionKeyHash GetHashFromPartitionKey(CosmosElement partitionKey, PartitionKeyDefinition partitionKeyDefinition) diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/InMemoryCollectionTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/InMemoryCollectionTests.cs index dfd24a6fdd..e63346c86e 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/InMemoryCollectionTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/InMemoryCollectionTests.cs @@ -5,6 +5,8 @@ namespace Microsoft.Azure.Cosmos.Tests { using System; + using System.Collections.Generic; + using System.Linq; using Microsoft.Azure.Cosmos.CosmosElements; using Microsoft.Azure.Cosmos.CosmosElements.Numbers; using Microsoft.Azure.Documents; @@ -44,5 +46,116 @@ public void TestCrud() Assert.AreEqual(item.ToString(), readRecord.Payload.ToString()); } + + [TestMethod] + public void TestPartitionKey() + { + PartitionKeyDefinition partitionKeyDefinition = new PartitionKeyDefinition() + { + Paths = new System.Collections.ObjectModel.Collection() + { + "/pk" + }, + Kind = PartitionKind.Hash, + Version = PartitionKeyDefinitionVersion.V2, + }; + + InMemoryCollection inMemoryCollection = new InMemoryCollection(partitionKeyDefinition); + + // Insert an item + CosmosObject item1 = CosmosObject.Parse("{\"pk\" : 42 }"); + InMemoryCollection.Record record1 = inMemoryCollection.CreateItem(item1); + + // Insert into another partition key + CosmosObject item2 = CosmosObject.Parse("{\"pk\" : 1337 }"); + InMemoryCollection.Record record2 = inMemoryCollection.CreateItem(item2); + + // Try to read back an id with wrong pk + Assert.IsFalse( + inMemoryCollection.TryReadItem( + partitionKey: item1["pk"], + record2.Identifier, + out InMemoryCollection.Record _)); + } + + [TestMethod] + public void UndefinedPartitionKey() + { + PartitionKeyDefinition partitionKeyDefinition = new PartitionKeyDefinition() + { + Paths = new System.Collections.ObjectModel.Collection() + { + "/pk" + }, + Kind = PartitionKind.Hash, + Version = PartitionKeyDefinitionVersion.V2, + }; + + InMemoryCollection inMemoryCollection = new InMemoryCollection(partitionKeyDefinition); + + // Insert an item + CosmosObject item = CosmosObject.Parse("{}"); + InMemoryCollection.Record record = inMemoryCollection.CreateItem(item); + + // Try to read back an id with wrong pk + Assert.IsTrue( + inMemoryCollection.TryReadItem( + partitionKey: null, + record.Identifier, + out InMemoryCollection.Record _)); + } + + [TestMethod] + public void Split() + { + PartitionKeyDefinition partitionKeyDefinition = new PartitionKeyDefinition() + { + Paths = new System.Collections.ObjectModel.Collection() + { + "/pk" + }, + Kind = PartitionKind.Hash, + Version = PartitionKeyDefinitionVersion.V2, + }; + + InMemoryCollection inMemoryCollection = new InMemoryCollection(partitionKeyDefinition); + + Assert.AreEqual(1, inMemoryCollection.PartitionKeyRangeFeedReed().Count); + + int numItemsToInsert = 10; + for (int i = 0; i < numItemsToInsert; i++) + { + // Insert an item + CosmosObject item = CosmosObject.Parse($"{{\"pk\" : {i} }}"); + inMemoryCollection.CreateItem(item); + } + + inMemoryCollection.Split(partitionKeyRangeId: 0); + + Assert.AreEqual(2, inMemoryCollection.PartitionKeyRangeFeedReed().Count); + + int AssertChildPartition(int partitionKeyRangeId) + { + Assert.IsTrue(inMemoryCollection.TryReadFeed( + partitionKeyRangeId: partitionKeyRangeId, + pageIndex: 0, + pageSize: 100, + out List partitionRecords)); + + List values = new List(); + foreach (InMemoryCollection.Record record in partitionRecords) + { + values.Add(Number64.ToLong((record.Payload["pk"] as CosmosNumber).Value)); + } + + List sortedValues = values.OrderBy(x => x).ToList(); + Assert.IsTrue(values.SequenceEqual(sortedValues)); + + return values.Count; + } + + int count = AssertChildPartition(partitionKeyRangeId: 1) + AssertChildPartition(partitionKeyRangeId: 2); + Assert.AreEqual(numItemsToInsert, count); + } } } From f229619f7452a3637f1f041fb56199d21e3c3a6e Mon Sep 17 00:00:00 2001 From: Brandon Chong Date: Tue, 23 Jun 2020 21:47:20 -0700 Subject: [PATCH 06/85] switching tracks to try asyncenumerable --- .../src/Pagination/PartitionRangePaginator.cs | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/Microsoft.Azure.Cosmos/src/Pagination/PartitionRangePaginator.cs b/Microsoft.Azure.Cosmos/src/Pagination/PartitionRangePaginator.cs index 98d7f353b0..f469cea37a 100644 --- a/Microsoft.Azure.Cosmos/src/Pagination/PartitionRangePaginator.cs +++ b/Microsoft.Azure.Cosmos/src/Pagination/PartitionRangePaginator.cs @@ -7,15 +7,18 @@ namespace Microsoft.Azure.Cosmos.Pagination using System.Threading; using System.Threading.Tasks; using Microsoft.Azure.Cosmos.Query.Core.Monads; + using Microsoft.Azure.Cosmos.Routing; /// /// Has the ability to page through a partition range. /// internal abstract class PartitionRangePaginator { - public FeedRange FeedRange { get; } + public PartitionKeyHashRange Range { get; } - public Page CurrentPage { get; } + public Page CurrentPage { get; protected set; } + + public bool HasMoreResults { get; protected set; } public abstract State GetState(); From 66f5db0c2d88afce9bb95922b7536da353424ed5 Mon Sep 17 00:00:00 2001 From: Brandon Chong Date: Wed, 24 Jun 2020 17:25:38 -0700 Subject: [PATCH 07/85] started playing with async enumerator --- Microsoft.Azure.Cosmos/src/Microsoft.Azure.Cosmos.csproj | 5 +++-- .../src/Pagination/CrossPartitionAsyncPaginator.cs | 9 +++++++++ .../src/Pagination/PartitionRangePaginator.cs | 9 +++++---- 3 files changed, 17 insertions(+), 6 deletions(-) create mode 100644 Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionAsyncPaginator.cs diff --git a/Microsoft.Azure.Cosmos/src/Microsoft.Azure.Cosmos.csproj b/Microsoft.Azure.Cosmos/src/Microsoft.Azure.Cosmos.csproj index 81f97fd6da..8ccc1a5e51 100644 --- a/Microsoft.Azure.Cosmos/src/Microsoft.Azure.Cosmos.csproj +++ b/Microsoft.Azure.Cosmos/src/Microsoft.Azure.Cosmos.csproj @@ -86,6 +86,7 @@ + @@ -93,8 +94,8 @@ - - + + diff --git a/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionAsyncPaginator.cs b/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionAsyncPaginator.cs new file mode 100644 index 0000000000..969a6b385a --- /dev/null +++ b/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionAsyncPaginator.cs @@ -0,0 +1,9 @@ +namespace Microsoft.Azure.Cosmos.Pagination +{ + using System.Collections.Generic; + using Microsoft.Azure.Cosmos.Query.Core.Monads; + + internal sealed class CrossPartitionAsyncPaginator : IAsyncEnumerator> + { + } +} diff --git a/Microsoft.Azure.Cosmos/src/Pagination/PartitionRangePaginator.cs b/Microsoft.Azure.Cosmos/src/Pagination/PartitionRangePaginator.cs index f469cea37a..13afc0c6c3 100644 --- a/Microsoft.Azure.Cosmos/src/Pagination/PartitionRangePaginator.cs +++ b/Microsoft.Azure.Cosmos/src/Pagination/PartitionRangePaginator.cs @@ -4,6 +4,7 @@ namespace Microsoft.Azure.Cosmos.Pagination { + using System.Collections.Generic; using System.Threading; using System.Threading.Tasks; using Microsoft.Azure.Cosmos.Query.Core.Monads; @@ -12,16 +13,16 @@ namespace Microsoft.Azure.Cosmos.Pagination /// /// Has the ability to page through a partition range. /// - internal abstract class PartitionRangePaginator + internal abstract class PartitionRangePaginator : IAsyncEnumerator> { public PartitionKeyHashRange Range { get; } - public Page CurrentPage { get; protected set; } + public TryCatch Current { get; protected set; } - public bool HasMoreResults { get; protected set; } + public abstract ValueTask DisposeAsync(); public abstract State GetState(); - public abstract Task TryMoveNextPageAsync(CancellationToken cancellationToken = default); + public abstract ValueTask MoveNextAsync(); } } From 157617d3c4b0a79deebd63b7947c7a18e941b0f8 Mon Sep 17 00:00:00 2001 From: Brandon Chong Date: Wed, 24 Jun 2020 18:19:15 -0700 Subject: [PATCH 08/85] added full IAsyncEnumerable interface --- .../CrossPartitionAsyncPaginator.cs | 9 --- .../CrossPartitionRangePageEnumerable.cs | 46 ++++++++++++++++ ...s => CrossPartitionRangePageEnumerator.cs} | 55 +++++++++++-------- ...tor.cs => PartitionRangePageEnumerator.cs} | 5 +- .../src/Routing/PartitionKeyHash.cs | 5 ++ 5 files changed, 84 insertions(+), 36 deletions(-) delete mode 100644 Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionAsyncPaginator.cs create mode 100644 Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionRangePageEnumerable.cs rename Microsoft.Azure.Cosmos/src/Pagination/{CrossPartitionRangePaginator.cs => CrossPartitionRangePageEnumerator.cs} (60%) rename Microsoft.Azure.Cosmos/src/Pagination/{PartitionRangePaginator.cs => PartitionRangePageEnumerator.cs} (81%) diff --git a/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionAsyncPaginator.cs b/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionAsyncPaginator.cs deleted file mode 100644 index 969a6b385a..0000000000 --- a/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionAsyncPaginator.cs +++ /dev/null @@ -1,9 +0,0 @@ -namespace Microsoft.Azure.Cosmos.Pagination -{ - using System.Collections.Generic; - using Microsoft.Azure.Cosmos.Query.Core.Monads; - - internal sealed class CrossPartitionAsyncPaginator : IAsyncEnumerator> - { - } -} diff --git a/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionRangePageEnumerable.cs b/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionRangePageEnumerable.cs new file mode 100644 index 0000000000..9716ff987b --- /dev/null +++ b/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionRangePageEnumerable.cs @@ -0,0 +1,46 @@ +namespace Microsoft.Azure.Cosmos.Pagination +{ + using System; + using System.Collections.Generic; + using System.Linq; + using System.Threading; + using Microsoft.Azure.Cosmos.Query.Core.Monads; + + internal sealed class CrossPartitionRangePageEnumerable : IAsyncEnumerable> + { + private readonly IEnumerable<(FeedRange, State)> rangeAndStates; + private readonly Func createPartitionRangePaginator; + private readonly IComparer comparer; + private readonly FeedRangeProvider feedRangeProvider; + + public CrossPartitionRangePageEnumerable( + IEnumerable<(FeedRange, State)> rangeAndStates, + Func createPartitionRangePaginator, + IComparer comparer, + FeedRangeProvider feedRangeProvider) + { + this.rangeAndStates = rangeAndStates ?? throw new ArgumentNullException(nameof(rangeAndStates)); + this.createPartitionRangePaginator = createPartitionRangePaginator ?? throw new ArgumentNullException(nameof(createPartitionRangePaginator)); + this.comparer = comparer ?? throw new ArgumentNullException(nameof(comparer)); + this.feedRangeProvider = feedRangeProvider ?? throw new ArgumentNullException(nameof(comparer)); + } + + public IAsyncEnumerator> GetAsyncEnumerator(CancellationToken cancellationToken = default) + { + cancellationToken.ThrowIfCancellationRequested(); + + List enumerators = new List(this.rangeAndStates.Count()); + foreach ((FeedRange range, State state) in this.rangeAndStates) + { + PartitionRangePageEnumerator enumerator = this.createPartitionRangePaginator(range, state); + enumerators.Add(enumerator); + } + + return new CrossPartitionRangePageEnumerator( + this.feedRangeProvider, + this.createPartitionRangePaginator, + enumerators, + this.comparer); + } + } +} diff --git a/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionRangePaginator.cs b/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionRangePageEnumerator.cs similarity index 60% rename from Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionRangePaginator.cs rename to Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionRangePageEnumerator.cs index 91f111f915..97a8a17370 100644 --- a/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionRangePaginator.cs +++ b/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionRangePageEnumerator.cs @@ -7,25 +7,24 @@ namespace Microsoft.Azure.Cosmos.Pagination using System; using System.Collections.Generic; using System.Net; - using System.Threading; using System.Threading.Tasks; using Microsoft.Azure.Cosmos.Query.Core.Collections; using Microsoft.Azure.Cosmos.Query.Core.Monads; /// - /// Coordinates draining pages from multiple , while maintaining a global sort order and handling repartitioning (splits, merge). + /// Coordinates draining pages from multiple , while maintaining a global sort order and handling repartitioning (splits, merge). /// - internal abstract class CrossPartitionRangePaginator + internal sealed class CrossPartitionRangePageEnumerator : IAsyncEnumerator> { private readonly FeedRangeProvider feedRangeProvider; - private readonly Func createPartitionRangePaginator; - private readonly PriorityQueue paginators; + private readonly Func createPartitionRangePaginator; + private readonly PriorityQueue paginators; - public CrossPartitionRangePaginator( + public CrossPartitionRangePageEnumerator( FeedRangeProvider feedRangeProvider, - Func createPartitionRangePaginator, - IEnumerable paginators, - IComparer comparer) + Func createPartitionRangePaginator, + IEnumerable paginators, + IComparer comparer) { this.feedRangeProvider = feedRangeProvider ?? throw new ArgumentNullException(nameof(feedRangeProvider)); this.createPartitionRangePaginator = createPartitionRangePaginator ?? throw new ArgumentNullException(nameof(createPartitionRangePaginator)); @@ -40,45 +39,53 @@ public CrossPartitionRangePaginator( throw new ArgumentNullException(nameof(paginators)); } - this.paginators = new PriorityQueue(paginators, comparer); + this.paginators = new PriorityQueue(paginators, comparer); } - public Page CurrentPage { get; set; } + public TryCatch Current { get; private set; } - public async Task TryMoveNextPageAsync(CancellationToken cancellationToken = default) + public async ValueTask MoveNextAsync() { - cancellationToken.ThrowIfCancellationRequested(); - PartitionRangePaginator currentPaginator = this.paginators.Dequeue(); - TryCatch tryMoveNextPageAsync = await currentPaginator.TryMoveNextPageAsync(cancellationToken); - if (tryMoveNextPageAsync.Failed) + PartitionRangePageEnumerator currentPaginator = this.paginators.Dequeue(); + bool movedNext = await currentPaginator.MoveNextAsync(); + if (!movedNext) { - Exception exception = tryMoveNextPageAsync.Exception; + return false; + } + + if (currentPaginator.Current.Failed) + { + // Check if it's a retryable exception. + Exception exception = currentPaginator.Current.Exception; if (IsSplitException(exception)) { // Handle split - IEnumerable childRanges = await this.feedRangeProvider.GetChildRangeAsync(currentPaginator.FeedRange, cancellationToken); + IEnumerable childRanges = await this.feedRangeProvider.GetChildRangeAsync(currentPaginator.Range); foreach (FeedRange childRange in childRanges) { - PartitionRangePaginator childPaginator = this.createPartitionRangePaginator(childRange, currentPaginator.GetState()); + PartitionRangePageEnumerator childPaginator = this.createPartitionRangePaginator(childRange, currentPaginator.GetState()); this.paginators.Enqueue(childPaginator); } // Recursively retry - return await this.TryMoveNextPageAsync(cancellationToken); + return await this.MoveNextAsync(); } if (IsMergeException(exception)) { throw new NotImplementedException(); } - - return tryMoveNextPageAsync; } - this.CurrentPage = currentPaginator.CurrentPage; + this.Current = currentPaginator.Current; this.paginators.Enqueue(currentPaginator); - return TryCatch.FromResult(); + return true; + } + + public ValueTask DisposeAsync() + { + throw new NotImplementedException(); } private static bool IsSplitException(Exception exeception) diff --git a/Microsoft.Azure.Cosmos/src/Pagination/PartitionRangePaginator.cs b/Microsoft.Azure.Cosmos/src/Pagination/PartitionRangePageEnumerator.cs similarity index 81% rename from Microsoft.Azure.Cosmos/src/Pagination/PartitionRangePaginator.cs rename to Microsoft.Azure.Cosmos/src/Pagination/PartitionRangePageEnumerator.cs index 13afc0c6c3..ef368c1a84 100644 --- a/Microsoft.Azure.Cosmos/src/Pagination/PartitionRangePaginator.cs +++ b/Microsoft.Azure.Cosmos/src/Pagination/PartitionRangePageEnumerator.cs @@ -5,7 +5,6 @@ namespace Microsoft.Azure.Cosmos.Pagination { using System.Collections.Generic; - using System.Threading; using System.Threading.Tasks; using Microsoft.Azure.Cosmos.Query.Core.Monads; using Microsoft.Azure.Cosmos.Routing; @@ -13,9 +12,9 @@ namespace Microsoft.Azure.Cosmos.Pagination /// /// Has the ability to page through a partition range. /// - internal abstract class PartitionRangePaginator : IAsyncEnumerator> + internal abstract class PartitionRangePageEnumerator : IAsyncEnumerator> { - public PartitionKeyHashRange Range { get; } + public FeedRange Range { get; } public TryCatch Current { get; protected set; } diff --git a/Microsoft.Azure.Cosmos/src/Routing/PartitionKeyHash.cs b/Microsoft.Azure.Cosmos/src/Routing/PartitionKeyHash.cs index 2633236978..2b6384c701 100644 --- a/Microsoft.Azure.Cosmos/src/Routing/PartitionKeyHash.cs +++ b/Microsoft.Azure.Cosmos/src/Routing/PartitionKeyHash.cs @@ -71,6 +71,11 @@ public override int GetHashCode() return this.Value.GetHashCode(); } + public string ToHexString() + { + throw new NotImplementedException(); + } + public static class V1 { private const int MaxStringLength = 100; From 4ce4a13c890d449ae9559701e1ae74493fb29e3e Mon Sep 17 00:00:00 2001 From: Brandon Chong Date: Thu, 25 Jun 2020 13:32:44 -0700 Subject: [PATCH 09/85] made the single partition range enumerator more plug and play --- .../PartitionRangePageEnumerator.cs | 25 ++++++++++++++++--- 1 file changed, 21 insertions(+), 4 deletions(-) diff --git a/Microsoft.Azure.Cosmos/src/Pagination/PartitionRangePageEnumerator.cs b/Microsoft.Azure.Cosmos/src/Pagination/PartitionRangePageEnumerator.cs index ef368c1a84..17cc7f50bd 100644 --- a/Microsoft.Azure.Cosmos/src/Pagination/PartitionRangePageEnumerator.cs +++ b/Microsoft.Azure.Cosmos/src/Pagination/PartitionRangePageEnumerator.cs @@ -5,6 +5,7 @@ namespace Microsoft.Azure.Cosmos.Pagination { using System.Collections.Generic; + using System.Threading; using System.Threading.Tasks; using Microsoft.Azure.Cosmos.Query.Core.Monads; using Microsoft.Azure.Cosmos.Routing; @@ -14,14 +15,30 @@ namespace Microsoft.Azure.Cosmos.Pagination /// internal abstract class PartitionRangePageEnumerator : IAsyncEnumerator> { + private bool hasStarted; + public FeedRange Range { get; } - public TryCatch Current { get; protected set; } + public TryCatch Current { get; private set; } - public abstract ValueTask DisposeAsync(); + public State State { get; private set; } + + public async ValueTask MoveNextAsync() + { + if (this.hasStarted && (this.State == default)) + { + return false; + } - public abstract State GetState(); + (TryCatch page, State state) = await this.GetNextPageAsync(this.State); + this.State = state; + this.Current = page; - public abstract ValueTask MoveNextAsync(); + return true; + } + + public abstract Task<(TryCatch, State)> GetNextPageAsync(State state, CancellationToken cancellationToken = default); + + public abstract ValueTask DisposeAsync(); } } From 4e62c7f238ef9c3a6a207b70bb8cbbd14d588f94 Mon Sep 17 00:00:00 2001 From: Brandon Chong Date: Thu, 25 Jun 2020 18:21:04 -0700 Subject: [PATCH 10/85] got in memory partition range enumerator working --- .../CrossPartitionRangePageEnumerable.cs | 6 +- .../CrossPartitionRangePageEnumerator.cs | 2 +- .../PartitionRangePageEnumerable.cs | 35 +++ .../PartitionRangePageEnumerator.cs | 13 +- .../InMemoryCollection.cs | 73 +++--- .../InMemoryCollectionTests.cs | 13 +- .../Microsoft.Azure.Cosmos.Tests.csproj | 5 +- ...emoryCollectionPartitionRangeEnumerator.cs | 91 ++++++++ ...CollectionPartitionRangeEnumeratorTests.cs | 221 ++++++++++++++++++ 9 files changed, 410 insertions(+), 49 deletions(-) create mode 100644 Microsoft.Azure.Cosmos/src/Pagination/PartitionRangePageEnumerable.cs create mode 100644 Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/InMemoryCollectionPartitionRangeEnumerator.cs create mode 100644 Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/InMemoryCollectionPartitionRangeEnumeratorTests.cs diff --git a/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionRangePageEnumerable.cs b/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionRangePageEnumerable.cs index 9716ff987b..49662bc03f 100644 --- a/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionRangePageEnumerable.cs +++ b/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionRangePageEnumerable.cs @@ -1,4 +1,8 @@ -namespace Microsoft.Azure.Cosmos.Pagination +// ------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +// ------------------------------------------------------------ + +namespace Microsoft.Azure.Cosmos.Pagination { using System; using System.Collections.Generic; diff --git a/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionRangePageEnumerator.cs b/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionRangePageEnumerator.cs index 97a8a17370..dbdc0f4f17 100644 --- a/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionRangePageEnumerator.cs +++ b/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionRangePageEnumerator.cs @@ -63,7 +63,7 @@ public async ValueTask MoveNextAsync() IEnumerable childRanges = await this.feedRangeProvider.GetChildRangeAsync(currentPaginator.Range); foreach (FeedRange childRange in childRanges) { - PartitionRangePageEnumerator childPaginator = this.createPartitionRangePaginator(childRange, currentPaginator.GetState()); + PartitionRangePageEnumerator childPaginator = this.createPartitionRangePaginator(childRange, currentPaginator.State); this.paginators.Enqueue(childPaginator); } diff --git a/Microsoft.Azure.Cosmos/src/Pagination/PartitionRangePageEnumerable.cs b/Microsoft.Azure.Cosmos/src/Pagination/PartitionRangePageEnumerable.cs new file mode 100644 index 0000000000..2b7592d0f8 --- /dev/null +++ b/Microsoft.Azure.Cosmos/src/Pagination/PartitionRangePageEnumerable.cs @@ -0,0 +1,35 @@ +// ------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +// ------------------------------------------------------------ + +namespace Microsoft.Azure.Cosmos.Pagination +{ + using System; + using System.Collections.Generic; + using System.Threading; + using Microsoft.Azure.Cosmos.Query.Core.Monads; + + internal sealed class PartitionRangePageEnumerable : IAsyncEnumerable> + { + private readonly FeedRange range; + private readonly State state; + private readonly Func createPartitionRangePaginator; + + public PartitionRangePageEnumerable( + FeedRange range, + State state, + Func createPartitionRangePaginator) + { + this.range = range ?? throw new ArgumentNullException(nameof(range)); + this.state = state ?? throw new ArgumentNullException(nameof(state)); + this.createPartitionRangePaginator = createPartitionRangePaginator ?? throw new ArgumentNullException(nameof(createPartitionRangePaginator)); + } + + public IAsyncEnumerator> GetAsyncEnumerator(CancellationToken cancellationToken = default) + { + cancellationToken.ThrowIfCancellationRequested(); + + return this.createPartitionRangePaginator(this.range, this.state); + } + } +} diff --git a/Microsoft.Azure.Cosmos/src/Pagination/PartitionRangePageEnumerator.cs b/Microsoft.Azure.Cosmos/src/Pagination/PartitionRangePageEnumerator.cs index 17cc7f50bd..43a44a0c0b 100644 --- a/Microsoft.Azure.Cosmos/src/Pagination/PartitionRangePageEnumerator.cs +++ b/Microsoft.Azure.Cosmos/src/Pagination/PartitionRangePageEnumerator.cs @@ -8,7 +8,6 @@ namespace Microsoft.Azure.Cosmos.Pagination using System.Threading; using System.Threading.Tasks; using Microsoft.Azure.Cosmos.Query.Core.Monads; - using Microsoft.Azure.Cosmos.Routing; /// /// Has the ability to page through a partition range. @@ -17,6 +16,12 @@ internal abstract class PartitionRangePageEnumerator : IAsyncEnumerator Current { get; private set; } @@ -30,14 +35,16 @@ public async ValueTask MoveNextAsync() return false; } - (TryCatch page, State state) = await this.GetNextPageAsync(this.State); + this.hasStarted = true; + + (TryCatch page, State state) = await this.GetNextPageAsync(); this.State = state; this.Current = page; return true; } - public abstract Task<(TryCatch, State)> GetNextPageAsync(State state, CancellationToken cancellationToken = default); + public abstract Task<(TryCatch, State)> GetNextPageAsync(CancellationToken cancellationToken = default); public abstract ValueTask DisposeAsync(); } diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/InMemoryCollection.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/InMemoryCollection.cs index df35084c9f..e590bdfca9 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/InMemoryCollection.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/InMemoryCollection.cs @@ -9,6 +9,7 @@ namespace Microsoft.Azure.Cosmos.Tests using System.Collections.Generic; using System.Linq; using Microsoft.Azure.Cosmos.CosmosElements; + using Microsoft.Azure.Cosmos.Query.Core.Monads; using Microsoft.Azure.Cosmos.Routing; using Microsoft.Azure.Documents; using Moq; @@ -77,12 +78,17 @@ record = default; return false; } - public bool TryReadFeed(int partitionKeyRangeId, int pageIndex, int pageSize, out List page) + public TryCatch> ReadFeed(int partitionKeyRangeId, long resourceIndentifer, int pageSize) { if (!this.partitionKeyRangeIdToHashRange.TryGetValue(partitionKeyRangeId, out PartitionKeyHashRange range)) { - page = default; - return false; + return TryCatch>.FromException( + new CosmosException( + message: $"PartitionKeyRangeId {partitionKeyRangeId} is gone", + statusCode: System.Net.HttpStatusCode.Gone, + subStatusCode: (int)SubStatusCodes.PartitionKeyRangeGone, + activityId: Guid.NewGuid().ToString(), + requestCharge: default)); } if (!this.partitionedRecords.TryGetValue(range, out Records records)) @@ -90,8 +96,12 @@ public bool TryReadFeed(int partitionKeyRangeId, int pageIndex, int pageSize, ou throw new InvalidOperationException("failed to find the range."); } - page = records.Skip(pageIndex * pageSize).Take(pageSize).ToList(); - return true; + List page = records + .Where(record => record.ResourceIdentifier > resourceIndentifer) + .Take(pageSize) + .ToList(); + + return TryCatch>.FromResult(page); } public IReadOnlyDictionary PartitionKeyRangeFeedReed() => this.partitionKeyRangeIdToHashRange; @@ -104,7 +114,7 @@ public void Split(int partitionKeyRangeId) throw new InvalidOperationException("Failed to find the range."); } - if (!this.partitionedRecords.TryGetValue(parentRange, out Records records)) + if (!this.partitionedRecords.TryGetValue(parentRange, out Records parentRecords)) { throw new InvalidOperationException("failed to find the range."); } @@ -138,9 +148,16 @@ public void Split(int partitionKeyRangeId) this.partitionKeyRangeIdToHashRange = newPartitionKeyRangeIdToHashRange; // Rehash the records in the parent range - foreach (Record record in records) + foreach (Record record in parentRecords) { - this.CreateItem(record.Payload); + PartitionKeyHash partitionKeyHash = GetHashFromPayload(record.Payload, this.partitionKeyDefinition); + if (!this.partitionedRecords.TryGetValue(partitionKeyHash, out Records records)) + { + records = new Records(); + this.partitionedRecords[partitionKeyHash] = records; + } + + records.Add(record); } } @@ -203,33 +220,15 @@ private static PartitionKeyHash GetHashFromPartitionKey(CosmosElement partitionK throw new ArgumentOutOfRangeException("Can only support a single partition key path."); } - PartitionKeyHash partitionKeyHash; - switch (partitionKey) + PartitionKeyHash partitionKeyHash = partitionKey switch { - case null: - partitionKeyHash = PartitionKeyHash.V2.HashUndefined(); - break; - - case CosmosString stringPartitionKey: - partitionKeyHash = PartitionKeyHash.V2.Hash(stringPartitionKey.Value); - break; - - case CosmosNumber numberPartitionKey: - partitionKeyHash = PartitionKeyHash.V2.Hash(Number64.ToDouble(numberPartitionKey.Value)); - break; - - case CosmosBoolean cosmosBoolean: - partitionKeyHash = PartitionKeyHash.V2.Hash(cosmosBoolean.Value); - break; - - case CosmosNull _: - partitionKeyHash = PartitionKeyHash.V2.HashNull(); - break; - - default: - throw new ArgumentOutOfRangeException(); - } - + null => PartitionKeyHash.V2.HashUndefined(), + CosmosString stringPartitionKey => PartitionKeyHash.V2.Hash(stringPartitionKey.Value), + CosmosNumber numberPartitionKey => PartitionKeyHash.V2.Hash(Number64.ToDouble(numberPartitionKey.Value)), + CosmosBoolean cosmosBoolean => PartitionKeyHash.V2.Hash(cosmosBoolean.Value), + CosmosNull _ => PartitionKeyHash.V2.HashNull(), + _ => throw new ArgumentOutOfRangeException(), + }; return partitionKeyHash; } @@ -290,6 +289,12 @@ public Record Add(CosmosObject payload) this.storage.Add(record); return record; } + + public Record Add(Record record) + { + this.storage.Add(record); + return record; + } } } } diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/InMemoryCollectionTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/InMemoryCollectionTests.cs index e63346c86e..36a9698047 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/InMemoryCollectionTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/InMemoryCollectionTests.cs @@ -9,6 +9,7 @@ namespace Microsoft.Azure.Cosmos.Tests using System.Linq; using Microsoft.Azure.Cosmos.CosmosElements; using Microsoft.Azure.Cosmos.CosmosElements.Numbers; + using Microsoft.Azure.Cosmos.Query.Core.Monads; using Microsoft.Azure.Documents; using Microsoft.VisualStudio.TestTools.UnitTesting; @@ -136,14 +137,14 @@ public void Split() int AssertChildPartition(int partitionKeyRangeId) { - Assert.IsTrue(inMemoryCollection.TryReadFeed( + TryCatch> tryGetPartitionRecords = inMemoryCollection.ReadFeed( partitionKeyRangeId: partitionKeyRangeId, - pageIndex: 0, - pageSize: 100, - out List partitionRecords)); + resourceIndentifer: 0, + pageSize: 100); + tryGetPartitionRecords.ThrowIfFailed(); List values = new List(); - foreach (InMemoryCollection.Record record in partitionRecords) + foreach (InMemoryCollection.Record record in tryGetPartitionRecords.Result) { values.Add(Number64.ToLong((record.Payload["pk"] as CosmosNumber).Value)); } @@ -156,6 +157,6 @@ int AssertChildPartition(int partitionKeyRangeId) int count = AssertChildPartition(partitionKeyRangeId: 1) + AssertChildPartition(partitionKeyRangeId: 2); Assert.AreEqual(numItemsToInsert, count); - } + } } } 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 ed847de63e..59baef2af6 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 @@ -9,6 +9,7 @@ false false Microsoft.Azure.Cosmos.Tests + 8.0 @@ -239,10 +240,6 @@ - - - - true true diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/InMemoryCollectionPartitionRangeEnumerator.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/InMemoryCollectionPartitionRangeEnumerator.cs new file mode 100644 index 0000000000..9759db97a3 --- /dev/null +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/InMemoryCollectionPartitionRangeEnumerator.cs @@ -0,0 +1,91 @@ +//------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +//------------------------------------------------------------ + +namespace Microsoft.Azure.Cosmos.Tests.Pagination +{ + using System; + using System.Collections.Generic; + using System.Linq; + using System.Threading; + using System.Threading.Tasks; + using Microsoft.Azure.Cosmos.Pagination; + using Microsoft.Azure.Cosmos.Query.Core.Monads; + using Microsoft.Azure.Documents; + + internal sealed class InMemoryCollectionPartitionRangeEnumerator : PartitionRangePageEnumerator + { + private readonly InMemoryCollection inMemoryCollection; + private readonly int pageSize; + private readonly int partitionKeyRangeId; + + public InMemoryCollectionPartitionRangeEnumerator(InMemoryCollection inMemoryCollection, int partitionKeyRangeId, int pageSize, State state = null) + : base(new FeedRangePartitionKeyRange(partitionKeyRangeId.ToString()), state ?? new InMemoryCollectionState(resourceIdentifier: 0)) + { + this.inMemoryCollection = inMemoryCollection ?? throw new ArgumentNullException(nameof(inMemoryCollection)); + this.partitionKeyRangeId = partitionKeyRangeId; + this.pageSize = pageSize; + + if (state != null) + { + if (!(state is InMemoryCollectionState _)) + { + throw new ArgumentOutOfRangeException(nameof(state)); + } + } + } + + public override ValueTask DisposeAsync() + { + // Do Nothing + return default; + } + + public override Task<(TryCatch, State)> GetNextPageAsync(CancellationToken cancellationToken = default) + { + cancellationToken.ThrowIfCancellationRequested(); + + TryCatch> tryReadPage = this.inMemoryCollection.ReadFeed( + partitionKeyRangeId: this.partitionKeyRangeId, + resourceIndentifer: ((InMemoryCollectionState)this.State).ResourceIdentifier, + pageSize: this.pageSize); + if (tryReadPage.Failed) + { + return Task.FromResult((TryCatch.FromException(tryReadPage.Exception), this.State)); + } + + if (tryReadPage.Result.Count == 0) + { + InMemoryCollectionPage emptyPage = new InMemoryCollectionPage(new List()); + InMemoryCollectionState nullContinuation = default; + return Task.FromResult((TryCatch.FromResult(emptyPage), (State)nullContinuation)); + } + + InMemoryCollectionPage page = new InMemoryCollectionPage(tryReadPage.Result); + State inMemoryCollectionState = new InMemoryCollectionState(page.Records.Last().ResourceIdentifier); + (TryCatch, State) tryPageAndState = (TryCatch.FromResult(page), inMemoryCollectionState); + + return Task.FromResult(tryPageAndState); + } + + private sealed class InMemoryCollectionState : State + { + public InMemoryCollectionState(long resourceIdentifier) + { + this.ResourceIdentifier = resourceIdentifier; + } + + public long ResourceIdentifier { get; } + } + + public sealed class InMemoryCollectionPage : Page + { + public InMemoryCollectionPage(List records) + { + this.Records = records; + } + + public List Records { get; } + } + } +} diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/InMemoryCollectionPartitionRangeEnumeratorTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/InMemoryCollectionPartitionRangeEnumeratorTests.cs new file mode 100644 index 0000000000..02cd34f6cb --- /dev/null +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/InMemoryCollectionPartitionRangeEnumeratorTests.cs @@ -0,0 +1,221 @@ +//------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +//------------------------------------------------------------ + +namespace Microsoft.Azure.Cosmos.Tests.Pagination +{ + using System.Security.Cryptography.X509Certificates; + using Microsoft.VisualStudio.TestTools.UnitTesting; + using Microsoft.Azure.Documents; + using Microsoft.Azure.Cosmos.CosmosElements; + using Microsoft.Azure.Cosmos.Pagination; + using System.Collections.Generic; + using Microsoft.Azure.Cosmos.Query.Core.Monads; + using System.Threading; + using System; + using System.Threading.Tasks; + using Microsoft.Azure.Cosmos.Serialization.HybridRow.RecordIO; + using System.Linq; + + [TestClass] + public class InMemoryCollectionPartitionRangeEnumeratorTests + { + [TestMethod] + public async Task TestDrainFullyAsync() + { + InMemoryCollection inMemoryCollection = CreateInMemoryCollection(); + IAsyncEnumerable> enumerable = new PageEnumeratorToEnumerableAdaptor(() => + { + return new InMemoryCollectionPartitionRangeEnumerator( + inMemoryCollection, + partitionKeyRangeId: 0, + pageSize: 10); + }); + + List records = new List(); + await foreach (TryCatch tryGetPage in enumerable) + { + tryGetPage.ThrowIfFailed(); + + if (!(tryGetPage.Result is InMemoryCollectionPartitionRangeEnumerator.InMemoryCollectionPage page)) + { + throw new InvalidCastException(); + } + + records.AddRange(page.Records); + } + + Assert.AreEqual(100, new HashSet(records.Select(record => record.ResourceIdentifier)).Count); + } + + [TestMethod] + public async Task TestResumingFromState() + { + InMemoryCollection inMemoryCollection = CreateInMemoryCollection(); + InMemoryCollectionPartitionRangeEnumerator enumerator = new InMemoryCollectionPartitionRangeEnumerator( + inMemoryCollection, + partitionKeyRangeId: 0, + pageSize: 10); + + List records = new List(); + State state = default; + + // Drain a couple of iterations + for(int i = 0; i < 3; i++) + { + await enumerator.MoveNextAsync(); + + TryCatch tryGetPage = enumerator.Current; + tryGetPage.ThrowIfFailed(); + + if (!(tryGetPage.Result is InMemoryCollectionPartitionRangeEnumerator.InMemoryCollectionPage page)) + { + throw new InvalidCastException(); + } + + records.AddRange(page.Records); + state = enumerator.State; + } + + // Resume from state + enumerator = new InMemoryCollectionPartitionRangeEnumerator( + inMemoryCollection, + partitionKeyRangeId: 0, + pageSize: 10, + state: state); + + IAsyncEnumerable> enumerable = new PageEnumeratorToEnumerableAdaptor(() => enumerator); + await foreach (TryCatch tryGetPage in enumerable) + { + tryGetPage.ThrowIfFailed(); + + if (!(tryGetPage.Result is InMemoryCollectionPartitionRangeEnumerator.InMemoryCollectionPage page)) + { + throw new InvalidCastException(); + } + + records.AddRange(page.Records); + } + + Assert.AreEqual(100, new HashSet(records.Select(record => record.ResourceIdentifier)).Count); + } + + [TestMethod] + public async Task TestSplitAsync() + { + InMemoryCollection inMemoryCollection = CreateInMemoryCollection(); + InMemoryCollectionPartitionRangeEnumerator enumerator = new InMemoryCollectionPartitionRangeEnumerator( + inMemoryCollection, + partitionKeyRangeId: 0, + pageSize: 10); + + List records = new List(); + State state = default; + + // Drain a couple of iterations + for (int i = 0; i < 3; i++) + { + await enumerator.MoveNextAsync(); + + TryCatch tryGetPage = enumerator.Current; + tryGetPage.ThrowIfFailed(); + + if (!(tryGetPage.Result is InMemoryCollectionPartitionRangeEnumerator.InMemoryCollectionPage page)) + { + throw new InvalidCastException(); + } + + records.AddRange(page.Records); + state = enumerator.State; + } + + inMemoryCollection.Split(partitionKeyRangeId: 0); + + // Try To read from the partition that is gone. + await enumerator.MoveNextAsync(); + Assert.IsTrue(enumerator.Current.Failed); + + // Resume on the children using the parent continuaiton token + InMemoryCollectionPartitionRangeEnumerator enumerator1 = new InMemoryCollectionPartitionRangeEnumerator( + inMemoryCollection, + partitionKeyRangeId: 1, + pageSize: 10, + state: state); + + InMemoryCollectionPartitionRangeEnumerator enumerator2 = new InMemoryCollectionPartitionRangeEnumerator( + inMemoryCollection, + partitionKeyRangeId: 2, + pageSize: 10, + state: state); + + IAsyncEnumerable> enumerable; + enumerable = new PageEnumeratorToEnumerableAdaptor(() => enumerator1); + await foreach (TryCatch tryGetPage in enumerable) + { + tryGetPage.ThrowIfFailed(); + + if (!(tryGetPage.Result is InMemoryCollectionPartitionRangeEnumerator.InMemoryCollectionPage page)) + { + throw new InvalidCastException(); + } + + records.AddRange(page.Records); + } + + enumerable = new PageEnumeratorToEnumerableAdaptor(() => enumerator2); + await foreach (TryCatch tryGetPage in enumerable) + { + tryGetPage.ThrowIfFailed(); + + if (!(tryGetPage.Result is InMemoryCollectionPartitionRangeEnumerator.InMemoryCollectionPage page)) + { + throw new InvalidCastException(); + } + + records.AddRange(page.Records); + } + + Assert.AreEqual(100, new HashSet(records.Select(record => record.ResourceIdentifier)).Count); + } + + private static InMemoryCollection CreateInMemoryCollection() + { + PartitionKeyDefinition partitionKeyDefinition = new PartitionKeyDefinition() + { + Paths = new System.Collections.ObjectModel.Collection() + { + "/pk" + }, + Kind = PartitionKind.Hash, + Version = PartitionKeyDefinitionVersion.V2, + }; + + InMemoryCollection inMemoryCollection = new InMemoryCollection(partitionKeyDefinition); + + int numItemsToInsert = 100; + for (int i = 0; i < numItemsToInsert; i++) + { + // Insert an item + CosmosObject item = CosmosObject.Parse($"{{\"pk\" : {i} }}"); + inMemoryCollection.CreateItem(item); + } + + return inMemoryCollection; + } + + private sealed class PageEnumeratorToEnumerableAdaptor : IAsyncEnumerable> + { + private readonly Func>> factory; + + public PageEnumeratorToEnumerableAdaptor(Func>> factory) + { + this.factory = factory; + } + + public IAsyncEnumerator> GetAsyncEnumerator(CancellationToken cancellationToken = default) + { + return this.factory(); + } + } + } +} From 79e3d32986db8c9f9c8b85dfc59d932f78e971d7 Mon Sep 17 00:00:00 2001 From: Brandon Chong Date: Thu, 25 Jun 2020 22:35:46 -0700 Subject: [PATCH 11/85] cleaned up code --- .../CreatePartitionRangeEnumerator.cs | 8 + .../CrossPartitionRangePageEnumerable.cs | 10 +- .../CrossPartitionRangePageEnumerator.cs | 8 +- .../PartitionRangePageEnumerable.cs | 10 +- ...CollectionPartitionRangeEnumeratorTests.cs | 181 ++++++++---------- 5 files changed, 105 insertions(+), 112 deletions(-) create mode 100644 Microsoft.Azure.Cosmos/src/Pagination/CreatePartitionRangeEnumerator.cs diff --git a/Microsoft.Azure.Cosmos/src/Pagination/CreatePartitionRangeEnumerator.cs b/Microsoft.Azure.Cosmos/src/Pagination/CreatePartitionRangeEnumerator.cs new file mode 100644 index 0000000000..134b69e188 --- /dev/null +++ b/Microsoft.Azure.Cosmos/src/Pagination/CreatePartitionRangeEnumerator.cs @@ -0,0 +1,8 @@ +// ------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +// ------------------------------------------------------------ + +namespace Microsoft.Azure.Cosmos.Pagination +{ + internal delegate PartitionRangePageEnumerator CreatePartitionRangePageEnumerator(FeedRange feedRange, State state); +} diff --git a/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionRangePageEnumerable.cs b/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionRangePageEnumerable.cs index 49662bc03f..7187dc8d91 100644 --- a/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionRangePageEnumerable.cs +++ b/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionRangePageEnumerable.cs @@ -13,18 +13,18 @@ namespace Microsoft.Azure.Cosmos.Pagination internal sealed class CrossPartitionRangePageEnumerable : IAsyncEnumerable> { private readonly IEnumerable<(FeedRange, State)> rangeAndStates; - private readonly Func createPartitionRangePaginator; + private readonly CreatePartitionRangePageEnumerator createPartitionRangeEnumerator; private readonly IComparer comparer; private readonly FeedRangeProvider feedRangeProvider; public CrossPartitionRangePageEnumerable( IEnumerable<(FeedRange, State)> rangeAndStates, - Func createPartitionRangePaginator, + CreatePartitionRangePageEnumerator createPartitionRangeEnumerator, IComparer comparer, FeedRangeProvider feedRangeProvider) { this.rangeAndStates = rangeAndStates ?? throw new ArgumentNullException(nameof(rangeAndStates)); - this.createPartitionRangePaginator = createPartitionRangePaginator ?? throw new ArgumentNullException(nameof(createPartitionRangePaginator)); + this.createPartitionRangeEnumerator = createPartitionRangeEnumerator ?? throw new ArgumentNullException(nameof(createPartitionRangeEnumerator)); this.comparer = comparer ?? throw new ArgumentNullException(nameof(comparer)); this.feedRangeProvider = feedRangeProvider ?? throw new ArgumentNullException(nameof(comparer)); } @@ -36,13 +36,13 @@ public IAsyncEnumerator> GetAsyncEnumerator(CancellationToken can List enumerators = new List(this.rangeAndStates.Count()); foreach ((FeedRange range, State state) in this.rangeAndStates) { - PartitionRangePageEnumerator enumerator = this.createPartitionRangePaginator(range, state); + PartitionRangePageEnumerator enumerator = this.createPartitionRangeEnumerator(range, state); enumerators.Add(enumerator); } return new CrossPartitionRangePageEnumerator( this.feedRangeProvider, - this.createPartitionRangePaginator, + this.createPartitionRangeEnumerator, enumerators, this.comparer); } diff --git a/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionRangePageEnumerator.cs b/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionRangePageEnumerator.cs index dbdc0f4f17..a86b6b6adc 100644 --- a/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionRangePageEnumerator.cs +++ b/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionRangePageEnumerator.cs @@ -17,17 +17,17 @@ namespace Microsoft.Azure.Cosmos.Pagination internal sealed class CrossPartitionRangePageEnumerator : IAsyncEnumerator> { private readonly FeedRangeProvider feedRangeProvider; - private readonly Func createPartitionRangePaginator; + private readonly CreatePartitionRangePageEnumerator createPartitionRangeEnumerator; private readonly PriorityQueue paginators; public CrossPartitionRangePageEnumerator( FeedRangeProvider feedRangeProvider, - Func createPartitionRangePaginator, + CreatePartitionRangePageEnumerator createPartitionRangeEnumerator, IEnumerable paginators, IComparer comparer) { this.feedRangeProvider = feedRangeProvider ?? throw new ArgumentNullException(nameof(feedRangeProvider)); - this.createPartitionRangePaginator = createPartitionRangePaginator ?? throw new ArgumentNullException(nameof(createPartitionRangePaginator)); + this.createPartitionRangeEnumerator = createPartitionRangeEnumerator ?? throw new ArgumentNullException(nameof(createPartitionRangeEnumerator)); if (paginators == null) { @@ -63,7 +63,7 @@ public async ValueTask MoveNextAsync() IEnumerable childRanges = await this.feedRangeProvider.GetChildRangeAsync(currentPaginator.Range); foreach (FeedRange childRange in childRanges) { - PartitionRangePageEnumerator childPaginator = this.createPartitionRangePaginator(childRange, currentPaginator.State); + PartitionRangePageEnumerator childPaginator = this.createPartitionRangeEnumerator(childRange, currentPaginator.State); this.paginators.Enqueue(childPaginator); } diff --git a/Microsoft.Azure.Cosmos/src/Pagination/PartitionRangePageEnumerable.cs b/Microsoft.Azure.Cosmos/src/Pagination/PartitionRangePageEnumerable.cs index 2b7592d0f8..d1c686323b 100644 --- a/Microsoft.Azure.Cosmos/src/Pagination/PartitionRangePageEnumerable.cs +++ b/Microsoft.Azure.Cosmos/src/Pagination/PartitionRangePageEnumerable.cs @@ -13,23 +13,23 @@ internal sealed class PartitionRangePageEnumerable : IAsyncEnumerable createPartitionRangePaginator; + private readonly CreatePartitionRangePageEnumerator createPartitionRangeEnumerator; public PartitionRangePageEnumerable( FeedRange range, State state, - Func createPartitionRangePaginator) + CreatePartitionRangePageEnumerator createPartitionRangeEnumerator) { this.range = range ?? throw new ArgumentNullException(nameof(range)); - this.state = state ?? throw new ArgumentNullException(nameof(state)); - this.createPartitionRangePaginator = createPartitionRangePaginator ?? throw new ArgumentNullException(nameof(createPartitionRangePaginator)); + this.state = state; + this.createPartitionRangeEnumerator = createPartitionRangeEnumerator ?? throw new ArgumentNullException(nameof(createPartitionRangeEnumerator)); } public IAsyncEnumerator> GetAsyncEnumerator(CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); - return this.createPartitionRangePaginator(this.range, this.state); + return this.createPartitionRangeEnumerator(this.range, this.state); } } } diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/InMemoryCollectionPartitionRangeEnumeratorTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/InMemoryCollectionPartitionRangeEnumeratorTests.cs index 02cd34f6cb..e3f3b353b9 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/InMemoryCollectionPartitionRangeEnumeratorTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/InMemoryCollectionPartitionRangeEnumeratorTests.cs @@ -4,17 +4,14 @@ namespace Microsoft.Azure.Cosmos.Tests.Pagination { - using System.Security.Cryptography.X509Certificates; + using System; + using System.Threading.Tasks; using Microsoft.VisualStudio.TestTools.UnitTesting; using Microsoft.Azure.Documents; using Microsoft.Azure.Cosmos.CosmosElements; using Microsoft.Azure.Cosmos.Pagination; using System.Collections.Generic; using Microsoft.Azure.Cosmos.Query.Core.Monads; - using System.Threading; - using System; - using System.Threading.Tasks; - using Microsoft.Azure.Cosmos.Serialization.HybridRow.RecordIO; using System.Linq; [TestClass] @@ -23,45 +20,36 @@ public class InMemoryCollectionPartitionRangeEnumeratorTests [TestMethod] public async Task TestDrainFullyAsync() { - InMemoryCollection inMemoryCollection = CreateInMemoryCollection(); - IAsyncEnumerable> enumerable = new PageEnumeratorToEnumerableAdaptor(() => - { - return new InMemoryCollectionPartitionRangeEnumerator( + int numItems = 100; + InMemoryCollection inMemoryCollection = CreateInMemoryCollection(numItems); + PartitionRangePageEnumerable enumerable = new PartitionRangePageEnumerable( + range: new FeedRangePartitionKeyRange("0"), + state: default, + (range, state) => new InMemoryCollectionPartitionRangeEnumerator( inMemoryCollection, - partitionKeyRangeId: 0, - pageSize: 10); - }); - - List records = new List(); - await foreach (TryCatch tryGetPage in enumerable) - { - tryGetPage.ThrowIfFailed(); - - if (!(tryGetPage.Result is InMemoryCollectionPartitionRangeEnumerator.InMemoryCollectionPage page)) - { - throw new InvalidCastException(); - } + partitionKeyRangeId: int.Parse(((FeedRangePartitionKeyRange)range).PartitionKeyRangeId), + pageSize: 10, + state: state)); - records.AddRange(page.Records); - } - - Assert.AreEqual(100, new HashSet(records.Select(record => record.ResourceIdentifier)).Count); + HashSet resourceIdentifiers = await GetResourceIdentifiersAsync(enumerable); + Assert.AreEqual(numItems, resourceIdentifiers.Count); } [TestMethod] public async Task TestResumingFromState() { - InMemoryCollection inMemoryCollection = CreateInMemoryCollection(); + int numItems = 100; + InMemoryCollection inMemoryCollection = CreateInMemoryCollection(numItems); InMemoryCollectionPartitionRangeEnumerator enumerator = new InMemoryCollectionPartitionRangeEnumerator( - inMemoryCollection, - partitionKeyRangeId: 0, - pageSize: 10); + inMemoryCollection, + partitionKeyRangeId: 0, + pageSize: 10); List records = new List(); State state = default; // Drain a couple of iterations - for(int i = 0; i < 3; i++) + for (int i = 0; i < 3; i++) { await enumerator.MoveNextAsync(); @@ -84,30 +72,28 @@ public async Task TestResumingFromState() pageSize: 10, state: state); - IAsyncEnumerable> enumerable = new PageEnumeratorToEnumerableAdaptor(() => enumerator); - await foreach (TryCatch tryGetPage in enumerable) - { - tryGetPage.ThrowIfFailed(); - - if (!(tryGetPage.Result is InMemoryCollectionPartitionRangeEnumerator.InMemoryCollectionPage page)) - { - throw new InvalidCastException(); - } - - records.AddRange(page.Records); - } + PartitionRangePageEnumerable enumerable = new PartitionRangePageEnumerable( + range: new FeedRangePartitionKeyRange("0"), + state: state, + (range, state) => new InMemoryCollectionPartitionRangeEnumerator( + inMemoryCollection, + partitionKeyRangeId: int.Parse(((FeedRangePartitionKeyRange)range).PartitionKeyRangeId), + pageSize: 10, + state: state)); + HashSet resourceIdentifiers = await GetResourceIdentifiersAsync(enumerable); - Assert.AreEqual(100, new HashSet(records.Select(record => record.ResourceIdentifier)).Count); + Assert.AreEqual(numItems, records.Count + resourceIdentifiers.Count); } [TestMethod] public async Task TestSplitAsync() { - InMemoryCollection inMemoryCollection = CreateInMemoryCollection(); + int numItems = 100; + InMemoryCollection inMemoryCollection = CreateInMemoryCollection(numItems); InMemoryCollectionPartitionRangeEnumerator enumerator = new InMemoryCollectionPartitionRangeEnumerator( - inMemoryCollection, - partitionKeyRangeId: 0, - pageSize: 10); + inMemoryCollection, + partitionKeyRangeId: 0, + pageSize: 10); List records = new List(); State state = default; @@ -136,49 +122,33 @@ public async Task TestSplitAsync() Assert.IsTrue(enumerator.Current.Failed); // Resume on the children using the parent continuaiton token - InMemoryCollectionPartitionRangeEnumerator enumerator1 = new InMemoryCollectionPartitionRangeEnumerator( - inMemoryCollection, - partitionKeyRangeId: 1, - pageSize: 10, - state: state); - - InMemoryCollectionPartitionRangeEnumerator enumerator2 = new InMemoryCollectionPartitionRangeEnumerator( - inMemoryCollection, - partitionKeyRangeId: 2, - pageSize: 10, - state: state); - - IAsyncEnumerable> enumerable; - enumerable = new PageEnumeratorToEnumerableAdaptor(() => enumerator1); - await foreach (TryCatch tryGetPage in enumerable) - { - tryGetPage.ThrowIfFailed(); - - if (!(tryGetPage.Result is InMemoryCollectionPartitionRangeEnumerator.InMemoryCollectionPage page)) - { - throw new InvalidCastException(); - } - - records.AddRange(page.Records); - } - - enumerable = new PageEnumeratorToEnumerableAdaptor(() => enumerator2); - await foreach (TryCatch tryGetPage in enumerable) - { - tryGetPage.ThrowIfFailed(); - - if (!(tryGetPage.Result is InMemoryCollectionPartitionRangeEnumerator.InMemoryCollectionPage page)) - { - throw new InvalidCastException(); - } + PartitionRangePageEnumerable enumerable1 = new PartitionRangePageEnumerable( + range: new FeedRangePartitionKeyRange("1"), + state: state, + (range, state) => new InMemoryCollectionPartitionRangeEnumerator( + inMemoryCollection, + partitionKeyRangeId: int.Parse(((FeedRangePartitionKeyRange)range).PartitionKeyRangeId), + pageSize: 10, + state: state)); + HashSet resourceIdentifiers1 = await GetResourceIdentifiersAsync(enumerable1); + + PartitionRangePageEnumerable enumerable2 = new PartitionRangePageEnumerable( + range: new FeedRangePartitionKeyRange("2"), + state: state, + (range, state) => new InMemoryCollectionPartitionRangeEnumerator( + inMemoryCollection, + partitionKeyRangeId: int.Parse(((FeedRangePartitionKeyRange)range).PartitionKeyRangeId), + pageSize: 10, + state: state)); + HashSet resourceIdentifiers2 = await GetResourceIdentifiersAsync(enumerable2); - records.AddRange(page.Records); - } + List commonAmongChildren = resourceIdentifiers1.Intersect(resourceIdentifiers2).ToList(); + Assert.AreEqual(0, commonAmongChildren.Count); - Assert.AreEqual(100, new HashSet(records.Select(record => record.ResourceIdentifier)).Count); + Assert.AreEqual(numItems, records.Count + resourceIdentifiers1.Count + resourceIdentifiers2.Count); } - private static InMemoryCollection CreateInMemoryCollection() + private static InMemoryCollection CreateInMemoryCollection(int numItems) { PartitionKeyDefinition partitionKeyDefinition = new PartitionKeyDefinition() { @@ -192,8 +162,7 @@ private static InMemoryCollection CreateInMemoryCollection() InMemoryCollection inMemoryCollection = new InMemoryCollection(partitionKeyDefinition); - int numItemsToInsert = 100; - for (int i = 0; i < numItemsToInsert; i++) + for (int i = 0; i < numItems; i++) { // Insert an item CosmosObject item = CosmosObject.Parse($"{{\"pk\" : {i} }}"); @@ -203,19 +172,35 @@ private static InMemoryCollection CreateInMemoryCollection() return inMemoryCollection; } - private sealed class PageEnumeratorToEnumerableAdaptor : IAsyncEnumerable> + private static async Task> GetResourceIdentifiersAsync( + PartitionRangePageEnumerable enumerable, + int? numIterations = default) { - private readonly Func>> factory; - - public PageEnumeratorToEnumerableAdaptor(Func>> factory) + int iterationNumber = 0; + HashSet resourceIdentifiers = new HashSet(); + await foreach (TryCatch tryGetPage in enumerable) { - this.factory = factory; - } + if (numIterations.HasValue && iterationNumber >= numIterations) + { + break; + } - public IAsyncEnumerator> GetAsyncEnumerator(CancellationToken cancellationToken = default) - { - return this.factory(); + tryGetPage.ThrowIfFailed(); + + if (!(tryGetPage.Result is InMemoryCollectionPartitionRangeEnumerator.InMemoryCollectionPage page)) + { + throw new InvalidCastException(); + } + + foreach (InMemoryCollection.Record record in page.Records) + { + resourceIdentifiers.Add(record.ResourceIdentifier); + } + + iterationNumber++; } + + return resourceIdentifiers; } } } From fddb47fdb39686f2a2945c18507a51c482cb1488 Mon Sep 17 00:00:00 2001 From: Brandon Chong Date: Fri, 26 Jun 2020 12:43:18 -0700 Subject: [PATCH 12/85] got cross partition working --- .../CrossPartitionRangePageEnumerable.cs | 23 +- .../CrossPartitionRangePageEnumerator.cs | 96 ++++++- .../src/Pagination/FeedRangeProvider.cs | 7 +- .../src/Pagination/IFeedRangeProvider.cs | 19 ++ .../PartitionRangePageEnumerator.cs | 4 +- .../src/Query/Core/AsyncLazy.cs | 2 +- .../PartitionKeyHashRangeDictionary.cs | 2 +- .../InMemoryCollection.cs | 22 +- .../InMemoryCollectionTests.cs | 75 +++++ ...CollectionPartitionRangeEnumeratorTests.cs | 268 ++++++++++++++++++ .../InMemoryCollectionFeedRangeProvider.cs | 55 ++++ ...CollectionPartitionRangeEnumeratorTests.cs | 121 ++++---- 12 files changed, 587 insertions(+), 107 deletions(-) create mode 100644 Microsoft.Azure.Cosmos/src/Pagination/IFeedRangeProvider.cs create mode 100644 Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/CrossPartitionInMemoryCollectionPartitionRangeEnumeratorTests.cs create mode 100644 Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/InMemoryCollectionFeedRangeProvider.cs diff --git a/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionRangePageEnumerable.cs b/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionRangePageEnumerable.cs index 7187dc8d91..3cc5678894 100644 --- a/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionRangePageEnumerable.cs +++ b/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionRangePageEnumerable.cs @@ -12,39 +12,32 @@ namespace Microsoft.Azure.Cosmos.Pagination internal sealed class CrossPartitionRangePageEnumerable : IAsyncEnumerable> { - private readonly IEnumerable<(FeedRange, State)> rangeAndStates; + private readonly State state; private readonly CreatePartitionRangePageEnumerator createPartitionRangeEnumerator; private readonly IComparer comparer; - private readonly FeedRangeProvider feedRangeProvider; + private readonly IFeedRangeProvider feedRangeProvider; public CrossPartitionRangePageEnumerable( - IEnumerable<(FeedRange, State)> rangeAndStates, + IFeedRangeProvider feedRangeProvider, CreatePartitionRangePageEnumerator createPartitionRangeEnumerator, IComparer comparer, - FeedRangeProvider feedRangeProvider) + State state = default) { - this.rangeAndStates = rangeAndStates ?? throw new ArgumentNullException(nameof(rangeAndStates)); + this.feedRangeProvider = feedRangeProvider ?? throw new ArgumentNullException(nameof(comparer)); this.createPartitionRangeEnumerator = createPartitionRangeEnumerator ?? throw new ArgumentNullException(nameof(createPartitionRangeEnumerator)); this.comparer = comparer ?? throw new ArgumentNullException(nameof(comparer)); - this.feedRangeProvider = feedRangeProvider ?? throw new ArgumentNullException(nameof(comparer)); + this.state = state; } public IAsyncEnumerator> GetAsyncEnumerator(CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); - List enumerators = new List(this.rangeAndStates.Count()); - foreach ((FeedRange range, State state) in this.rangeAndStates) - { - PartitionRangePageEnumerator enumerator = this.createPartitionRangeEnumerator(range, state); - enumerators.Add(enumerator); - } - return new CrossPartitionRangePageEnumerator( this.feedRangeProvider, this.createPartitionRangeEnumerator, - enumerators, - this.comparer); + this.comparer, + this.state); } } } diff --git a/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionRangePageEnumerator.cs b/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionRangePageEnumerator.cs index a86b6b6adc..38835cdde9 100644 --- a/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionRangePageEnumerator.cs +++ b/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionRangePageEnumerator.cs @@ -7,7 +7,9 @@ namespace Microsoft.Azure.Cosmos.Pagination using System; using System.Collections.Generic; using System.Net; + using System.Threading; using System.Threading.Tasks; + using Microsoft.Azure.Cosmos.Query.Core; using Microsoft.Azure.Cosmos.Query.Core.Collections; using Microsoft.Azure.Cosmos.Query.Core.Monads; @@ -16,37 +18,85 @@ namespace Microsoft.Azure.Cosmos.Pagination /// internal sealed class CrossPartitionRangePageEnumerator : IAsyncEnumerator> { - private readonly FeedRangeProvider feedRangeProvider; + private readonly IFeedRangeProvider feedRangeProvider; private readonly CreatePartitionRangePageEnumerator createPartitionRangeEnumerator; - private readonly PriorityQueue paginators; + private readonly AsyncLazy> lazyEnumerators; + private readonly State originalState; public CrossPartitionRangePageEnumerator( - FeedRangeProvider feedRangeProvider, + IFeedRangeProvider feedRangeProvider, CreatePartitionRangePageEnumerator createPartitionRangeEnumerator, - IEnumerable paginators, - IComparer comparer) + IComparer comparer, + State state = default) { this.feedRangeProvider = feedRangeProvider ?? throw new ArgumentNullException(nameof(feedRangeProvider)); this.createPartitionRangeEnumerator = createPartitionRangeEnumerator ?? throw new ArgumentNullException(nameof(createPartitionRangeEnumerator)); - if (paginators == null) + if (comparer == null) { - throw new ArgumentNullException(nameof(paginators)); + throw new ArgumentNullException(nameof(comparer)); } - if (comparer == null) + this.originalState = state; + + this.lazyEnumerators = new AsyncLazy>(async (CancellationToken token) => { - throw new ArgumentNullException(nameof(paginators)); - } + List<(FeedRange, State)> rangeAndStates; + if (state == default) + { + // Fan out to all partitions with default state + IEnumerable ranges = await feedRangeProvider.GetFeedRangesAsync(token); + + rangeAndStates = new List<(FeedRange, State)>(); + foreach (FeedRange range in ranges) + { + rangeAndStates.Add((range, default)); + } + } + else + { + if (!(state is CrossPartitionState crossPartitionState)) + { + throw new ArgumentOutOfRangeException(nameof(state)); + } + + rangeAndStates = crossPartitionState.Value; + } - this.paginators = new PriorityQueue(paginators, comparer); + PriorityQueue enumerators = new PriorityQueue(comparer); + foreach ((FeedRange range, State rangeState) in rangeAndStates) + { + PartitionRangePageEnumerator enumerator = createPartitionRangeEnumerator(range, rangeState); + enumerators.Enqueue(enumerator); + } + + return enumerators; + }); } public TryCatch Current { get; private set; } + public State GetState() + { + if (!this.lazyEnumerators.ValueInitialized) + { + return this.originalState; + } + + PriorityQueue enumerators = this.lazyEnumerators.Result; + List<(FeedRange, State)> feedRangeAndStates = new List<(FeedRange, State)>(enumerators.Count); + foreach (PartitionRangePageEnumerator enumerator in enumerators) + { + feedRangeAndStates.Add((enumerator.Range, enumerator.State)); + } + + return new CrossPartitionState(feedRangeAndStates); + } + public async ValueTask MoveNextAsync() { - PartitionRangePageEnumerator currentPaginator = this.paginators.Dequeue(); + PriorityQueue enumerators = await this.lazyEnumerators.GetValueAsync(); + PartitionRangePageEnumerator currentPaginator = enumerators.Dequeue(); bool movedNext = await currentPaginator.MoveNextAsync(); if (!movedNext) { @@ -57,6 +107,11 @@ public async ValueTask MoveNextAsync() { // Check if it's a retryable exception. Exception exception = currentPaginator.Current.Exception; + while (exception.InnerException != null) + { + exception = exception.InnerException; + } + if (IsSplitException(exception)) { // Handle split @@ -64,7 +119,7 @@ public async ValueTask MoveNextAsync() foreach (FeedRange childRange in childRanges) { PartitionRangePageEnumerator childPaginator = this.createPartitionRangeEnumerator(childRange, currentPaginator.State); - this.paginators.Enqueue(childPaginator); + enumerators.Enqueue(childPaginator); } // Recursively retry @@ -78,14 +133,15 @@ public async ValueTask MoveNextAsync() } this.Current = currentPaginator.Current; - this.paginators.Enqueue(currentPaginator); + enumerators.Enqueue(currentPaginator); return true; } public ValueTask DisposeAsync() { - throw new NotImplementedException(); + // Do Nothing. + return default; } private static bool IsSplitException(Exception exeception) @@ -100,5 +156,15 @@ private static bool IsMergeException(Exception exception) // TODO: code this out return false; } + + private sealed class CrossPartitionState : State + { + public CrossPartitionState(List<(FeedRange, State)> value) + { + this.Value = value; + } + + public List<(FeedRange, State)> Value { get; } + } } } diff --git a/Microsoft.Azure.Cosmos/src/Pagination/FeedRangeProvider.cs b/Microsoft.Azure.Cosmos/src/Pagination/FeedRangeProvider.cs index 1e6e556236..a879b2a25d 100644 --- a/Microsoft.Azure.Cosmos/src/Pagination/FeedRangeProvider.cs +++ b/Microsoft.Azure.Cosmos/src/Pagination/FeedRangeProvider.cs @@ -10,7 +10,7 @@ namespace Microsoft.Azure.Cosmos.Pagination using System.Threading.Tasks; using Microsoft.Azure.Cosmos.Query.Core.QueryClient; - internal sealed class FeedRangeProvider + internal sealed class FeedRangeProvider : IFeedRangeProvider { private readonly CosmosQueryClient cosmosQueryClient; private readonly string collectionRid; @@ -51,5 +51,10 @@ public async Task> GetChildRangeAsync( return childFeedRanges; } + + public Task> GetFeedRangesAsync( + CancellationToken cancellationToken = default) => this.GetChildRangeAsync( + FeedRangeEpk.ForFullRange(), + cancellationToken); } } diff --git a/Microsoft.Azure.Cosmos/src/Pagination/IFeedRangeProvider.cs b/Microsoft.Azure.Cosmos/src/Pagination/IFeedRangeProvider.cs new file mode 100644 index 0000000000..5a5cf416ce --- /dev/null +++ b/Microsoft.Azure.Cosmos/src/Pagination/IFeedRangeProvider.cs @@ -0,0 +1,19 @@ +// ------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +// ------------------------------------------------------------ + +namespace Microsoft.Azure.Cosmos.Pagination +{ + using System.Collections.Generic; + using System.Threading; + using System.Threading.Tasks; + + internal interface IFeedRangeProvider + { + public Task> GetChildRangeAsync( + FeedRange feedRange, + CancellationToken cancellationToken = default); + + public Task> GetFeedRangesAsync(CancellationToken cancellationToken = default); + } +} diff --git a/Microsoft.Azure.Cosmos/src/Pagination/PartitionRangePageEnumerator.cs b/Microsoft.Azure.Cosmos/src/Pagination/PartitionRangePageEnumerator.cs index 43a44a0c0b..3f39ba1ea6 100644 --- a/Microsoft.Azure.Cosmos/src/Pagination/PartitionRangePageEnumerator.cs +++ b/Microsoft.Azure.Cosmos/src/Pagination/PartitionRangePageEnumerator.cs @@ -28,9 +28,11 @@ protected PartitionRangePageEnumerator(FeedRange range, State state = null) public State State { get; private set; } + public bool HasMoreResults => !this.hasStarted || (this.State != default); + public async ValueTask MoveNextAsync() { - if (this.hasStarted && (this.State == default)) + if (!this.HasMoreResults) { return false; } diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/AsyncLazy.cs b/Microsoft.Azure.Cosmos/src/Query/Core/AsyncLazy.cs index 388311f48a..ecac17e9ad 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/AsyncLazy.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/AsyncLazy.cs @@ -20,7 +20,7 @@ public AsyncLazy(Func> valueFactory) public bool ValueInitialized { get; private set; } - public async Task GetValueAsync(CancellationToken cancellationToken) + public async Task GetValueAsync(CancellationToken cancellationToken = default) { // Note that this class is not thread safe. // if the valueFactory has side effects than this will have issues. diff --git a/Microsoft.Azure.Cosmos/src/Routing/PartitionKeyHashRangeDictionary.cs b/Microsoft.Azure.Cosmos/src/Routing/PartitionKeyHashRangeDictionary.cs index 44c920cc28..227af57257 100644 --- a/Microsoft.Azure.Cosmos/src/Routing/PartitionKeyHashRangeDictionary.cs +++ b/Microsoft.Azure.Cosmos/src/Routing/PartitionKeyHashRangeDictionary.cs @@ -89,7 +89,7 @@ public T this[PartitionKeyHashRange key] } set { - if (!this.TryGetValue(key, out _)) + if (!this.dictionary.TryGetValue(key, out _)) { throw new NotSupportedException("Dictionary does not support adding new elements."); } diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/InMemoryCollection.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/InMemoryCollection.cs index e590bdfca9..38b2057cc2 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/InMemoryCollection.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/InMemoryCollection.cs @@ -18,6 +18,7 @@ namespace Microsoft.Azure.Cosmos.Tests internal sealed class InMemoryCollection { private readonly PartitionKeyDefinition partitionKeyDefinition; + private readonly Dictionary parentToChildMapping; private PartitionKeyHashRangeDictionary partitionedRecords; private Dictionary partitionKeyRangeIdToHashRange; @@ -27,11 +28,13 @@ public InMemoryCollection(PartitionKeyDefinition partitionKeyDefinition) PartitionKeyHashRange fullRange = new PartitionKeyHashRange(startInclusive: null, endExclusive: null); PartitionKeyHashRanges partitionKeyHashRanges = PartitionKeyHashRanges.Create(new PartitionKeyHashRange[] { fullRange }); this.partitionedRecords = new PartitionKeyHashRangeDictionary(partitionKeyHashRanges); + this.partitionedRecords[fullRange] = new Records(); this.partitionKeyDefinition = partitionKeyDefinition ?? throw new ArgumentNullException(nameof(partitionKeyDefinition)); this.partitionKeyRangeIdToHashRange = new Dictionary() { { 0, fullRange } }; + this.parentToChildMapping = new Dictionary(); } public Record CreateItem(CosmosObject payload) @@ -125,16 +128,29 @@ public void Split(int partitionKeyRangeId) PartitionKeyHashRanges partitionKeyHashRanges = PartitionKeyHashRangeSplitterAndMerger.SplitRange(parentRange, 2); // Update the partition routing map + this.parentToChildMapping[partitionKeyRangeId] = (maxPartitionKeyRangeId + 1, maxPartitionKeyRangeId + 2); Dictionary newPartitionKeyRangeIdToHashRange = new Dictionary() { { maxPartitionKeyRangeId + 1, partitionKeyHashRanges.First() }, { maxPartitionKeyRangeId + 2, partitionKeyHashRanges.Last() }, }; + foreach (KeyValuePair kvp in this.partitionKeyRangeIdToHashRange) + { + int oldRangeId = kvp.Key; + PartitionKeyHashRange oldRange = kvp.Value; + if (!oldRange.Equals(parentRange)) + { + newPartitionKeyRangeIdToHashRange[oldRangeId] = oldRange; + } + } + // Copy over the partitioned records (minus the parent range) PartitionKeyHashRangeDictionary newPartitionedRecords = new PartitionKeyHashRangeDictionary( - PartitionKeyHashRanges.Create( - newPartitionKeyRangeIdToHashRange.Values)); + PartitionKeyHashRanges.Create(newPartitionKeyRangeIdToHashRange.Values)); + + newPartitionedRecords[partitionKeyHashRanges.First()] = new Records(); + newPartitionedRecords[partitionKeyHashRanges.Last()] = new Records(); foreach (PartitionKeyHashRange range in this.partitionKeyRangeIdToHashRange.Values) { @@ -161,6 +177,8 @@ public void Split(int partitionKeyRangeId) } } + public (int, int) GetChildRanges(int partitionKeyRangeId) => this.parentToChildMapping[partitionKeyRangeId]; + private static PartitionKeyHash GetHashFromPayload(CosmosObject payload, PartitionKeyDefinition partitionKeyDefinition) { CosmosElement partitionKey = GetPartitionKeyFromPayload(payload, partitionKeyDefinition); diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/InMemoryCollectionTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/InMemoryCollectionTests.cs index 36a9698047..c9559feaef 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/InMemoryCollectionTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/InMemoryCollectionTests.cs @@ -134,6 +134,9 @@ public void Split() inMemoryCollection.Split(partitionKeyRangeId: 0); Assert.AreEqual(2, inMemoryCollection.PartitionKeyRangeFeedReed().Count); + (int leftChild, int rightChild) = inMemoryCollection.GetChildRanges(partitionKeyRangeId: 0); + Assert.AreEqual(1, leftChild); + Assert.AreEqual(2, rightChild); int AssertChildPartition(int partitionKeyRangeId) { @@ -157,6 +160,78 @@ int AssertChildPartition(int partitionKeyRangeId) int count = AssertChildPartition(partitionKeyRangeId: 1) + AssertChildPartition(partitionKeyRangeId: 2); Assert.AreEqual(numItemsToInsert, count); + } + + [TestMethod] + public void MultiSplit() + { + PartitionKeyDefinition partitionKeyDefinition = new PartitionKeyDefinition() + { + Paths = new System.Collections.ObjectModel.Collection() + { + "/pk" + }, + Kind = PartitionKind.Hash, + Version = PartitionKeyDefinitionVersion.V2, + }; + + InMemoryCollection inMemoryCollection = new InMemoryCollection(partitionKeyDefinition); + + int numItemsToInsert = 10; + for (int i = 0; i < numItemsToInsert; i++) + { + // Insert an item + CosmosObject item = CosmosObject.Parse($"{{\"pk\" : {i} }}"); + inMemoryCollection.CreateItem(item); } + + Assert.AreEqual(1, inMemoryCollection.PartitionKeyRangeFeedReed().Count); + + inMemoryCollection.Split(partitionKeyRangeId: 0); + + Assert.AreEqual(2, inMemoryCollection.PartitionKeyRangeFeedReed().Count); + (int leftChild, int rightChild) = inMemoryCollection.GetChildRanges(partitionKeyRangeId: 0); + Assert.AreEqual(1, leftChild); + Assert.AreEqual(2, rightChild); + + inMemoryCollection.Split(partitionKeyRangeId: 1); + inMemoryCollection.Split(partitionKeyRangeId: 2); + + + Assert.AreEqual(4, inMemoryCollection.PartitionKeyRangeFeedReed().Count); + (int leftChild1, int rightChild1) = inMemoryCollection.GetChildRanges(partitionKeyRangeId: 1); + Assert.AreEqual(3, leftChild1); + Assert.AreEqual(4, rightChild1); + + (int leftChild2, int rightChild2) = inMemoryCollection.GetChildRanges(partitionKeyRangeId: 2); + Assert.AreEqual(5, leftChild2); + Assert.AreEqual(6, rightChild2); + + int AssertChildPartition(int partitionKeyRangeId) + { + TryCatch> tryGetPartitionRecords = inMemoryCollection.ReadFeed( + partitionKeyRangeId: partitionKeyRangeId, + resourceIndentifer: 0, + pageSize: 100); + tryGetPartitionRecords.ThrowIfFailed(); + + List values = new List(); + foreach (InMemoryCollection.Record record in tryGetPartitionRecords.Result) + { + values.Add(Number64.ToLong((record.Payload["pk"] as CosmosNumber).Value)); + } + + List sortedValues = values.OrderBy(x => x).ToList(); + Assert.IsTrue(values.SequenceEqual(sortedValues)); + + return values.Count; + } + + int count = AssertChildPartition(partitionKeyRangeId: 3) + + AssertChildPartition(partitionKeyRangeId: 4) + + AssertChildPartition(partitionKeyRangeId: 5) + + AssertChildPartition(partitionKeyRangeId: 6); + Assert.AreEqual(numItemsToInsert, count); + } } } diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/CrossPartitionInMemoryCollectionPartitionRangeEnumeratorTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/CrossPartitionInMemoryCollectionPartitionRangeEnumeratorTests.cs new file mode 100644 index 0000000000..1479fcd2e8 --- /dev/null +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/CrossPartitionInMemoryCollectionPartitionRangeEnumeratorTests.cs @@ -0,0 +1,268 @@ +//------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +//------------------------------------------------------------ + +namespace Microsoft.Azure.Cosmos.Tests.Pagination +{ + using System; + using System.Collections.Generic; + using System.Linq; + using System.Threading.Tasks; + using Microsoft.Azure.Cosmos.CosmosElements; + using Microsoft.Azure.Cosmos.Pagination; + using Microsoft.Azure.Cosmos.Query.Core.Monads; + using Microsoft.Azure.Documents; + using Microsoft.VisualStudio.TestTools.UnitTesting; + + [TestClass] + public class CrossPartitionInMemoryCollectionPartitionRangeEnumeratorTests + { + [TestMethod] + public async Task TestDrainFullyAsync() + { + int numItems = 1000; + InMemoryCollection inMemoryCollection = CreateInMemoryCollection(numItems); + + CrossPartitionRangePageEnumerable enumerable = new CrossPartitionRangePageEnumerable( + feedRangeProvider: new InMemoryCollectionFeedRangeProvider(inMemoryCollection), + createPartitionRangeEnumerator: (range, state) => new InMemoryCollectionPartitionRangeEnumerator( + inMemoryCollection, + partitionKeyRangeId: int.Parse(((FeedRangePartitionKeyRange)range).PartitionKeyRangeId), + pageSize: 10, + state: state), + comparer: PartitionRangePageEnumeratorComparer.Singleton); + + HashSet identifiers = await DrainFullyAsync(enumerable); + Assert.AreEqual(numItems, identifiers.Count); + } + + [TestMethod] + public async Task TestResumingFromStateAsync() + { + int numItems = 1000; + InMemoryCollection inMemoryCollection = CreateInMemoryCollection(numItems); + + IFeedRangeProvider feedRangeProvider = new InMemoryCollectionFeedRangeProvider(inMemoryCollection); + CreatePartitionRangePageEnumerator createEnumerator = (Cosmos.FeedRange range, State state) => new InMemoryCollectionPartitionRangeEnumerator( + inMemoryCollection, + partitionKeyRangeId: int.Parse(((FeedRangePartitionKeyRange)range).PartitionKeyRangeId), + pageSize: 10, + state: state); + + CrossPartitionRangePageEnumerator enumerator = new CrossPartitionRangePageEnumerator( + feedRangeProvider: feedRangeProvider, + createPartitionRangeEnumerator: createEnumerator, + comparer: PartitionRangePageEnumeratorComparer.Singleton); + + (HashSet firstDrainResults, State state) = await PartialDrainAsync(enumerator, numIterations: 3); + + // Resume from state + CrossPartitionRangePageEnumerable enumerable = new CrossPartitionRangePageEnumerable( + feedRangeProvider: feedRangeProvider, + createPartitionRangeEnumerator: createEnumerator, + comparer: PartitionRangePageEnumeratorComparer.Singleton, + state: state); + + HashSet secondDrainResults = await DrainFullyAsync(enumerable); + Assert.AreEqual(numItems, firstDrainResults.Count + secondDrainResults.Count); + } + + [TestMethod] + public async Task TestSplitWithResumeContinuationAsync() + { + int numItems = 1000; + InMemoryCollection inMemoryCollection = CreateInMemoryCollection(numItems); + + IFeedRangeProvider feedRangeProvider = new InMemoryCollectionFeedRangeProvider(inMemoryCollection); + CreatePartitionRangePageEnumerator createEnumerator = (Cosmos.FeedRange range, State state) => new InMemoryCollectionPartitionRangeEnumerator( + inMemoryCollection, + partitionKeyRangeId: int.Parse(((FeedRangePartitionKeyRange)range).PartitionKeyRangeId), + pageSize: 10, + state: state); + + CrossPartitionRangePageEnumerator enumerator = new CrossPartitionRangePageEnumerator( + feedRangeProvider: feedRangeProvider, + createPartitionRangeEnumerator: createEnumerator, + comparer: PartitionRangePageEnumeratorComparer.Singleton); + + (HashSet firstDrainResults, State state) = await PartialDrainAsync(enumerator, numIterations: 3); + + int minPartitionKeyRangeId = inMemoryCollection.PartitionKeyRangeFeedReed().Keys.Min(); + int maxPartitionKeyRangeId = inMemoryCollection.PartitionKeyRangeFeedReed().Keys.Max(); + // Split the partition we were reading from + inMemoryCollection.Split(minPartitionKeyRangeId); + + // And a partition we have let to read from + inMemoryCollection.Split((minPartitionKeyRangeId + maxPartitionKeyRangeId) / 2); + + // Resume from state + CrossPartitionRangePageEnumerable enumerable = new CrossPartitionRangePageEnumerable( + feedRangeProvider: feedRangeProvider, + createPartitionRangeEnumerator: createEnumerator, + comparer: PartitionRangePageEnumeratorComparer.Singleton, + state: state); + + HashSet secondDrainResults = await DrainFullyAsync(enumerable); + Assert.AreEqual(numItems, firstDrainResults.Count + secondDrainResults.Count); + } + + [TestMethod] + public async Task TestSplitWithDuringDrainAsync() + { + int numItems = 1000; + InMemoryCollection inMemoryCollection = CreateInMemoryCollection(numItems); + + IFeedRangeProvider feedRangeProvider = new InMemoryCollectionFeedRangeProvider(inMemoryCollection); + CreatePartitionRangePageEnumerator createEnumerator = (Cosmos.FeedRange range, State state) => new InMemoryCollectionPartitionRangeEnumerator( + inMemoryCollection, + partitionKeyRangeId: int.Parse(((FeedRangePartitionKeyRange)range).PartitionKeyRangeId), + pageSize: 10, + state: state); + + CrossPartitionRangePageEnumerable enumerable = new CrossPartitionRangePageEnumerable( + feedRangeProvider: feedRangeProvider, + createPartitionRangeEnumerator: createEnumerator, + comparer: PartitionRangePageEnumeratorComparer.Singleton); + + HashSet identifiers = new HashSet(); + Random random = new Random(); + await foreach (TryCatch tryGetPage in enumerable) + { + if (random.Next() % 2 == 0) + { + List partitionKeyRangeIds = inMemoryCollection.PartitionKeyRangeFeedReed().Keys.ToList(); + int randomIdToSplit = partitionKeyRangeIds[random.Next(0, partitionKeyRangeIds.Count)]; + inMemoryCollection.Split(randomIdToSplit); + } + + tryGetPage.ThrowIfFailed(); + + if (!(tryGetPage.Result is InMemoryCollectionPartitionRangeEnumerator.InMemoryCollectionPage page)) + { + throw new InvalidCastException(); + } + + foreach (InMemoryCollection.Record record in page.Records) + { + identifiers.Add(record.Identifier); + } + } + + Assert.AreEqual(numItems, identifiers.Count); + } + + private static InMemoryCollection CreateInMemoryCollection(int numItems) + { + PartitionKeyDefinition partitionKeyDefinition = new PartitionKeyDefinition() + { + Paths = new System.Collections.ObjectModel.Collection() + { + "/pk" + }, + Kind = PartitionKind.Hash, + Version = PartitionKeyDefinitionVersion.V2, + }; + + InMemoryCollection inMemoryCollection = new InMemoryCollection(partitionKeyDefinition); + + inMemoryCollection.Split(partitionKeyRangeId: 0); + + inMemoryCollection.Split(partitionKeyRangeId: 1); + inMemoryCollection.Split(partitionKeyRangeId: 2); + + inMemoryCollection.Split(partitionKeyRangeId: 3); + inMemoryCollection.Split(partitionKeyRangeId: 4); + inMemoryCollection.Split(partitionKeyRangeId: 5); + inMemoryCollection.Split(partitionKeyRangeId: 6); + + for (int i = 0; i < numItems; i++) + { + // Insert an item + CosmosObject item = CosmosObject.Parse($"{{\"pk\" : {i} }}"); + inMemoryCollection.CreateItem(item); + } + + return inMemoryCollection; + } + + private static async Task> DrainFullyAsync(CrossPartitionRangePageEnumerable enumerable) + { + HashSet identifiers = new HashSet(); + await foreach (TryCatch tryGetPage in enumerable) + { + tryGetPage.ThrowIfFailed(); + + if (!(tryGetPage.Result is InMemoryCollectionPartitionRangeEnumerator.InMemoryCollectionPage page)) + { + throw new InvalidCastException(); + } + + foreach (InMemoryCollection.Record record in page.Records) + { + identifiers.Add(record.Identifier); + } + } + + return identifiers; + } + + private static async Task<(HashSet, State)> PartialDrainAsync( + CrossPartitionRangePageEnumerator enumerator, + int numIterations) + { + HashSet identifiers = new HashSet(); + State state = default; + + // Drain a couple of iterations + for (int i = 0; i < numIterations; i++) + { + await enumerator.MoveNextAsync(); + + TryCatch tryGetPage = enumerator.Current; + tryGetPage.ThrowIfFailed(); + + if (!(tryGetPage.Result is InMemoryCollectionPartitionRangeEnumerator.InMemoryCollectionPage page)) + { + throw new InvalidCastException(); + } + + foreach (InMemoryCollection.Record record in page.Records) + { + identifiers.Add(record.Identifier); + } + + state = enumerator.GetState(); + } + + return (identifiers, state); + } + + private sealed class PartitionRangePageEnumeratorComparer : IComparer + { + public static readonly PartitionRangePageEnumeratorComparer Singleton = new PartitionRangePageEnumeratorComparer(); + + public int Compare(PartitionRangePageEnumerator partitionRangePageEnumerator1, PartitionRangePageEnumerator partitionRangePageEnumerator2) + { + if (object.ReferenceEquals(partitionRangePageEnumerator1, partitionRangePageEnumerator2)) + { + return 0; + } + + if (partitionRangePageEnumerator1.HasMoreResults && !partitionRangePageEnumerator2.HasMoreResults) + { + return -1; + } + + if (!partitionRangePageEnumerator1.HasMoreResults && partitionRangePageEnumerator2.HasMoreResults) + { + return 1; + } + + // Either both don't have results or both do. + return string.CompareOrdinal( + ((FeedRangePartitionKeyRange)partitionRangePageEnumerator1.Range).PartitionKeyRangeId, + ((FeedRangePartitionKeyRange)partitionRangePageEnumerator2.Range).PartitionKeyRangeId); + } + } + } +} diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/InMemoryCollectionFeedRangeProvider.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/InMemoryCollectionFeedRangeProvider.cs new file mode 100644 index 0000000000..c7134016db --- /dev/null +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/InMemoryCollectionFeedRangeProvider.cs @@ -0,0 +1,55 @@ +namespace Microsoft.Azure.Cosmos.Tests.Pagination +{ + using System; + using System.Collections.Generic; + using System.Threading; + using System.Threading.Tasks; + using Microsoft.Azure.Cosmos; + using Microsoft.Azure.Cosmos.Pagination; + + internal sealed class InMemoryCollectionFeedRangeProvider : IFeedRangeProvider + { + private readonly InMemoryCollection inMemoryCollection; + + public InMemoryCollectionFeedRangeProvider(InMemoryCollection inMemoryCollection) + { + this.inMemoryCollection = inMemoryCollection ?? throw new ArgumentNullException(nameof(inMemoryCollection)); + } + + public Task> GetChildRangeAsync( + FeedRange feedRange, + CancellationToken cancellationToken = default) + { + cancellationToken.ThrowIfCancellationRequested(); + + if (!(feedRange is FeedRangePartitionKeyRange feedRangePartitionKeyRange)) + { + throw new ArgumentOutOfRangeException(nameof(feedRange)); + } + + int partitionKeyRangeId = int.Parse(feedRangePartitionKeyRange.PartitionKeyRangeId); + (int leftChild, int rightChild) = this.inMemoryCollection.GetChildRanges(partitionKeyRangeId); + + return Task.FromResult( + (IEnumerable)new List() + { + new FeedRangePartitionKeyRange(leftChild.ToString()), + new FeedRangePartitionKeyRange(rightChild.ToString()), + }); + } + + public Task> GetFeedRangesAsync(CancellationToken cancellationToken = default) + { + cancellationToken.ThrowIfCancellationRequested(); + + List ranges = new List(); + foreach (int partitionKeyRangeId in this.inMemoryCollection.PartitionKeyRangeFeedReed().Keys) + { + FeedRange range = new FeedRangePartitionKeyRange(partitionKeyRangeId.ToString()); + ranges.Add(range); + } + + return Task.FromResult((IEnumerable)ranges); + } + } +} diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/InMemoryCollectionPartitionRangeEnumeratorTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/InMemoryCollectionPartitionRangeEnumeratorTests.cs index e3f3b353b9..0914dc7160 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/InMemoryCollectionPartitionRangeEnumeratorTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/InMemoryCollectionPartitionRangeEnumeratorTests.cs @@ -31,12 +31,12 @@ public async Task TestDrainFullyAsync() pageSize: 10, state: state)); - HashSet resourceIdentifiers = await GetResourceIdentifiersAsync(enumerable); - Assert.AreEqual(numItems, resourceIdentifiers.Count); + HashSet identifiers = await DrainFullyAsync(enumerable); + Assert.AreEqual(numItems, identifiers.Count); } [TestMethod] - public async Task TestResumingFromState() + public async Task TestResumingFromStateAsync() { int numItems = 100; InMemoryCollection inMemoryCollection = CreateInMemoryCollection(numItems); @@ -45,25 +45,7 @@ public async Task TestResumingFromState() partitionKeyRangeId: 0, pageSize: 10); - List records = new List(); - State state = default; - - // Drain a couple of iterations - for (int i = 0; i < 3; i++) - { - await enumerator.MoveNextAsync(); - - TryCatch tryGetPage = enumerator.Current; - tryGetPage.ThrowIfFailed(); - - if (!(tryGetPage.Result is InMemoryCollectionPartitionRangeEnumerator.InMemoryCollectionPage page)) - { - throw new InvalidCastException(); - } - - records.AddRange(page.Records); - state = enumerator.State; - } + (HashSet firstDrainResults, State state) = await PartialDrainAsync(enumerator, numIterations: 3); // Resume from state enumerator = new InMemoryCollectionPartitionRangeEnumerator( @@ -80,9 +62,9 @@ public async Task TestResumingFromState() partitionKeyRangeId: int.Parse(((FeedRangePartitionKeyRange)range).PartitionKeyRangeId), pageSize: 10, state: state)); - HashSet resourceIdentifiers = await GetResourceIdentifiersAsync(enumerable); + HashSet secondDrainResults = await DrainFullyAsync(enumerable); - Assert.AreEqual(numItems, records.Count + resourceIdentifiers.Count); + Assert.AreEqual(numItems, firstDrainResults.Count + secondDrainResults.Count); } [TestMethod] @@ -95,26 +77,9 @@ public async Task TestSplitAsync() partitionKeyRangeId: 0, pageSize: 10); - List records = new List(); - State state = default; - - // Drain a couple of iterations - for (int i = 0; i < 3; i++) - { - await enumerator.MoveNextAsync(); - - TryCatch tryGetPage = enumerator.Current; - tryGetPage.ThrowIfFailed(); - - if (!(tryGetPage.Result is InMemoryCollectionPartitionRangeEnumerator.InMemoryCollectionPage page)) - { - throw new InvalidCastException(); - } - - records.AddRange(page.Records); - state = enumerator.State; - } + (HashSet parentIdentifiers, State state) = await PartialDrainAsync(enumerator, numIterations: 3); + // Split the partition inMemoryCollection.Split(partitionKeyRangeId: 0); // Try To read from the partition that is gone. @@ -122,30 +87,23 @@ public async Task TestSplitAsync() Assert.IsTrue(enumerator.Current.Failed); // Resume on the children using the parent continuaiton token - PartitionRangePageEnumerable enumerable1 = new PartitionRangePageEnumerable( - range: new FeedRangePartitionKeyRange("1"), - state: state, - (range, state) => new InMemoryCollectionPartitionRangeEnumerator( - inMemoryCollection, - partitionKeyRangeId: int.Parse(((FeedRangePartitionKeyRange)range).PartitionKeyRangeId), - pageSize: 10, - state: state)); - HashSet resourceIdentifiers1 = await GetResourceIdentifiersAsync(enumerable1); - - PartitionRangePageEnumerable enumerable2 = new PartitionRangePageEnumerable( - range: new FeedRangePartitionKeyRange("2"), + HashSet childIdentifiers = new HashSet(); + foreach (int partitionKeyRangeId in new int[] { 1, 2 }) + { + PartitionRangePageEnumerable enumerable1 = new PartitionRangePageEnumerable( + range: new FeedRangePartitionKeyRange(partitionKeyRangeId.ToString()), state: state, (range, state) => new InMemoryCollectionPartitionRangeEnumerator( inMemoryCollection, partitionKeyRangeId: int.Parse(((FeedRangePartitionKeyRange)range).PartitionKeyRangeId), pageSize: 10, state: state)); - HashSet resourceIdentifiers2 = await GetResourceIdentifiersAsync(enumerable2); + HashSet resourceIdentifiers = await DrainFullyAsync(enumerable1); - List commonAmongChildren = resourceIdentifiers1.Intersect(resourceIdentifiers2).ToList(); - Assert.AreEqual(0, commonAmongChildren.Count); + childIdentifiers.UnionWith(resourceIdentifiers); + } - Assert.AreEqual(numItems, records.Count + resourceIdentifiers1.Count + resourceIdentifiers2.Count); + Assert.AreEqual(numItems, parentIdentifiers.Count + childIdentifiers.Count); } private static InMemoryCollection CreateInMemoryCollection(int numItems) @@ -172,19 +130,42 @@ private static InMemoryCollection CreateInMemoryCollection(int numItems) return inMemoryCollection; } - private static async Task> GetResourceIdentifiersAsync( - PartitionRangePageEnumerable enumerable, - int? numIterations = default) + private static async Task<(HashSet, State)> PartialDrainAsync( + PartitionRangePageEnumerator enumerator, + int numIterations) { - int iterationNumber = 0; - HashSet resourceIdentifiers = new HashSet(); - await foreach (TryCatch tryGetPage in enumerable) + HashSet identifiers = new HashSet(); + State state = default; + + // Drain a couple of iterations + for (int i = 0; i < numIterations; i++) { - if (numIterations.HasValue && iterationNumber >= numIterations) + await enumerator.MoveNextAsync(); + + TryCatch tryGetPage = enumerator.Current; + tryGetPage.ThrowIfFailed(); + + if (!(tryGetPage.Result is InMemoryCollectionPartitionRangeEnumerator.InMemoryCollectionPage page)) + { + throw new InvalidCastException(); + } + + foreach(InMemoryCollection.Record record in page.Records) { - break; + identifiers.Add(record.Identifier); } + state = enumerator.State; + } + + return (identifiers, state); + } + + private static async Task> DrainFullyAsync(PartitionRangePageEnumerable enumerable) + { + HashSet identifiers = new HashSet(); + await foreach (TryCatch tryGetPage in enumerable) + { tryGetPage.ThrowIfFailed(); if (!(tryGetPage.Result is InMemoryCollectionPartitionRangeEnumerator.InMemoryCollectionPage page)) @@ -194,13 +175,11 @@ private static async Task> GetResourceIdentifiersAsync( foreach (InMemoryCollection.Record record in page.Records) { - resourceIdentifiers.Add(record.ResourceIdentifier); + identifiers.Add(record.Identifier); } - - iterationNumber++; } - return resourceIdentifiers; + return identifiers; } } } From 1a618c60d57f1272c8940faa591ee88770ba958d Mon Sep 17 00:00:00 2001 From: Brandon Chong Date: Fri, 26 Jun 2020 14:40:21 -0700 Subject: [PATCH 13/85] resolved iteration comments --- .../src/Pagination/CrossPartitionRangePageEnumerator.cs | 2 +- Microsoft.Azure.Cosmos/src/Query/Core/AsyncLazy.cs | 2 +- Microsoft.Azure.Cosmos/src/Routing/PartitionKeyHash.cs | 5 ----- 3 files changed, 2 insertions(+), 7 deletions(-) diff --git a/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionRangePageEnumerator.cs b/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionRangePageEnumerator.cs index 38835cdde9..76baf73cca 100644 --- a/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionRangePageEnumerator.cs +++ b/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionRangePageEnumerator.cs @@ -95,7 +95,7 @@ public State GetState() public async ValueTask MoveNextAsync() { - PriorityQueue enumerators = await this.lazyEnumerators.GetValueAsync(); + PriorityQueue enumerators = await this.lazyEnumerators.GetValueAsync(cancellationToken: default); PartitionRangePageEnumerator currentPaginator = enumerators.Dequeue(); bool movedNext = await currentPaginator.MoveNextAsync(); if (!movedNext) diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/AsyncLazy.cs b/Microsoft.Azure.Cosmos/src/Query/Core/AsyncLazy.cs index ecac17e9ad..388311f48a 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/AsyncLazy.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/AsyncLazy.cs @@ -20,7 +20,7 @@ public AsyncLazy(Func> valueFactory) public bool ValueInitialized { get; private set; } - public async Task GetValueAsync(CancellationToken cancellationToken = default) + public async Task GetValueAsync(CancellationToken cancellationToken) { // Note that this class is not thread safe. // if the valueFactory has side effects than this will have issues. diff --git a/Microsoft.Azure.Cosmos/src/Routing/PartitionKeyHash.cs b/Microsoft.Azure.Cosmos/src/Routing/PartitionKeyHash.cs index 2b6384c701..2633236978 100644 --- a/Microsoft.Azure.Cosmos/src/Routing/PartitionKeyHash.cs +++ b/Microsoft.Azure.Cosmos/src/Routing/PartitionKeyHash.cs @@ -71,11 +71,6 @@ public override int GetHashCode() return this.Value.GetHashCode(); } - public string ToHexString() - { - throw new NotImplementedException(); - } - public static class V1 { private const int MaxStringLength = 100; From 3bf1fa44e981293a785959da39801ce374155189 Mon Sep 17 00:00:00 2001 From: Brandon Chong Date: Mon, 29 Jun 2020 12:07:56 -0700 Subject: [PATCH 14/85] started this query integration work, but need to make minor edits --- Microsoft.Azure.Cosmos/src/Pagination/Page.cs | 6 +- ...ackendQueryPartitionRangePageEnumerator.cs | 73 +++++++++++++++++++ .../ItemProducers/IQueryDataSource.cs | 19 +++++ .../ItemProducers/QueryPage.cs | 38 ++++++++++ .../QueryPartitionRangePageEnumerator.cs | 19 +++++ .../src/Query/Core/QueryResponseFactory.cs | 2 - .../Query/v3Query/CosmosQueryClientCore.cs | 43 +++++++---- 7 files changed, 178 insertions(+), 22 deletions(-) create mode 100644 Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/ItemProducers/BackendQueryPartitionRangePageEnumerator.cs create mode 100644 Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/ItemProducers/IQueryDataSource.cs create mode 100644 Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/ItemProducers/QueryPage.cs create mode 100644 Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/ItemProducers/QueryPartitionRangePageEnumerator.cs diff --git a/Microsoft.Azure.Cosmos/src/Pagination/Page.cs b/Microsoft.Azure.Cosmos/src/Pagination/Page.cs index 24a9cd0e36..beb4ff9c59 100644 --- a/Microsoft.Azure.Cosmos/src/Pagination/Page.cs +++ b/Microsoft.Azure.Cosmos/src/Pagination/Page.cs @@ -4,11 +4,9 @@ namespace Microsoft.Azure.Cosmos.Pagination { - using System; - using System.Collections.Generic; - using System.Text; - internal abstract class Page { + + public State State { get; } } } diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/ItemProducers/BackendQueryPartitionRangePageEnumerator.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/ItemProducers/BackendQueryPartitionRangePageEnumerator.cs new file mode 100644 index 0000000000..e3ea967a5f --- /dev/null +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/ItemProducers/BackendQueryPartitionRangePageEnumerator.cs @@ -0,0 +1,73 @@ +namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.ItemProducers +{ + using System; + using System.Threading; + using System.Threading.Tasks; + using Microsoft.Azure.Cosmos.Pagination; + using Microsoft.Azure.Cosmos.Query.Core.Monads; + using Microsoft.Azure.Cosmos.Query.Core.QueryClient; + using Microsoft.Azure.Documents; + + internal sealed class BackendQueryPartitionRangePageEnumerator : QueryPartitionRangePageEnumerator + { + private readonly CosmosQueryContext cosmosQueryContext; + private readonly int pageSize; + + public BackendQueryPartitionRangePageEnumerator( + CosmosQueryContext cosmosQueryContext, + SqlQuerySpec sqlQuerySpec, + FeedRange feedRange, + int pageSize, + State state = default) + : base(sqlQuerySpec, feedRange, state) + { + this.cosmosQueryContext = cosmosQueryContext ?? throw new ArgumentNullException(nameof(cosmosQueryContext)); + this.pageSize = pageSize; + + if (state != default) + { + if (!(state is QueryState)) + { + throw new ArgumentOutOfRangeException(nameof(state)); + } + } + + if (!(feedRange is FeedRangePartitionKeyRange)) + { + throw new ArgumentOutOfRangeException(nameof(feedRange)); + } + } + + public override ValueTask DisposeAsync() + { + return default; + } + + public override async Task<(TryCatch, State)> GetNextPageAsync(CancellationToken cancellationToken = default) + { + cancellationToken.ThrowIfCancellationRequested(); + + QueryResponseCore queryResponse = await this.cosmosQueryContext.ExecuteQueryAsync( + querySpecForInit: this.sqlQuerySpec, + continuationToken: ((QueryState)this.State).ContinuationToken, + partitionKeyRange: new PartitionKeyRangeIdentity( + this.cosmosQueryContext.ContainerResourceId, + ((FeedRangePartitionKeyRange)this.Range).PartitionKeyRangeId), + isContinuationExpected: this.cosmosQueryContext.IsContinuationExpected, + pageSize: this.pageSize, + cancellationToken: cancellationToken); + + + } + + private sealed class QueryState : State + { + public QueryState(string continuationToken) + { + this.ContinuationToken = continuationToken; + } + + public string ContinuationToken { get; } + } + } +} diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/ItemProducers/IQueryDataSource.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/ItemProducers/IQueryDataSource.cs new file mode 100644 index 0000000000..82f5cda6af --- /dev/null +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/ItemProducers/IQueryDataSource.cs @@ -0,0 +1,19 @@ +// ------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +// ------------------------------------------------------------ + +namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.ItemProducers +{ + using System; + using System.Collections.Generic; + using System.Text; + using System.Threading.Tasks; + using Microsoft.Azure.Cosmos.Query.Core.Monads; + + internal interface IQueryDataSource + { + public Task> ExecuteQueryAsync( + SqlQuerySpec sqlQuerySpec, + ) + } +} diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/ItemProducers/QueryPage.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/ItemProducers/QueryPage.cs new file mode 100644 index 0000000000..8f4b276619 --- /dev/null +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/ItemProducers/QueryPage.cs @@ -0,0 +1,38 @@ +// ------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +// ------------------------------------------------------------ + +namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.ItemProducers +{ + using System; + using System.Collections.Generic; + using Microsoft.Azure.Cosmos.CosmosElements; + using Microsoft.Azure.Cosmos.Query.Core.QueryClient; + + internal sealed class QueryPage + { + public QueryPage( + IReadOnlyList documents, + double requestCharge, + string activityId, + long responseLengthInBytes, + CosmosQueryExecutionInfo cosmosQueryExecutionInfo) + { + this.Documents = documents ?? throw new ArgumentNullException(nameof(documents)); + this.RequestCharge = requestCharge < 0 ? throw new ArgumentOutOfRangeException(nameof(requestCharge)) : requestCharge; + this.ActivityId = activityId; + this.ResponseLengthInBytes = responseLengthInBytes < 0 ? throw new ArgumentOutOfRangeException(nameof(responseLengthInBytes)) : responseLengthInBytes; + this.CosmosQueryExecutionInfo = cosmosQueryExecutionInfo; + } + + public IReadOnlyList Documents { get; } + + public double RequestCharge { get; } + + public string ActivityId { get; } + + public long ResponseLengthInBytes { get; } + + public CosmosQueryExecutionInfo CosmosQueryExecutionInfo { get; } + } +} diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/ItemProducers/QueryPartitionRangePageEnumerator.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/ItemProducers/QueryPartitionRangePageEnumerator.cs new file mode 100644 index 0000000000..03ddbe2673 --- /dev/null +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/ItemProducers/QueryPartitionRangePageEnumerator.cs @@ -0,0 +1,19 @@ +namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.ItemProducers +{ + using System; + using Microsoft.Azure.Cosmos.Pagination; + + internal abstract class QueryPartitionRangePageEnumerator : PartitionRangePageEnumerator + { + protected readonly SqlQuerySpec sqlQuerySpec; + + public QueryPartitionRangePageEnumerator( + SqlQuerySpec sqlQuerySpec, + FeedRange feedRange, + State state = default) + : base(feedRange, state) + { + this.sqlQuerySpec = sqlQuerySpec ?? throw new ArgumentNullException(nameof(sqlQuerySpec)); + } + } +} diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/QueryResponseFactory.cs b/Microsoft.Azure.Cosmos/src/Query/Core/QueryResponseFactory.cs index ac0b4b346b..e4b5f93ce0 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/QueryResponseFactory.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/QueryResponseFactory.cs @@ -5,8 +5,6 @@ namespace Microsoft.Azure.Cosmos.Query.Core { using System; - using System.Diagnostics; - using System.Runtime.CompilerServices; using Microsoft.Azure.Cosmos.Query.Core.Exceptions; using Microsoft.Azure.Cosmos.Query.Core.Monads; using Microsoft.Azure.Cosmos.Query.Core.QueryClient; diff --git a/Microsoft.Azure.Cosmos/src/Query/v3Query/CosmosQueryClientCore.cs b/Microsoft.Azure.Cosmos/src/Query/v3Query/CosmosQueryClientCore.cs index 7244195ba3..0790a3f0bb 100644 --- a/Microsoft.Azure.Cosmos/src/Query/v3Query/CosmosQueryClientCore.cs +++ b/Microsoft.Azure.Cosmos/src/Query/v3Query/CosmosQueryClientCore.cs @@ -16,6 +16,7 @@ namespace Microsoft.Azure.Cosmos using Microsoft.Azure.Cosmos.Diagnostics; using Microsoft.Azure.Cosmos.Json; using Microsoft.Azure.Cosmos.Query.Core; + using Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.ItemProducers; using Microsoft.Azure.Cosmos.Query.Core.Monads; using Microsoft.Azure.Cosmos.Query.Core.QueryClient; using Microsoft.Azure.Cosmos.Query.Core.QueryPlan; @@ -281,7 +282,7 @@ public override void ClearSessionTokenCache(string collectionFullName) sessionContainer.ClearTokenByCollectionFullname(collectionFullName); } - private static QueryResponseCore GetCosmosElementResponse( + private static TryCatch GetCosmosElementResponse( Guid clientQueryCorrelationId, QueryRequestOptions requestOptions, ResourceType resourceType, @@ -301,12 +302,22 @@ private static QueryResponseCore GetCosmosElementResponse( if (!cosmosResponseMessage.IsSuccessStatusCode) { - return QueryResponseCore.CreateFailure( - statusCode: cosmosResponseMessage.StatusCode, - subStatusCodes: cosmosResponseMessage.Headers.SubStatusCode, - cosmosException: cosmosResponseMessage.CosmosException, - requestCharge: cosmosResponseMessage.Headers.RequestCharge, - activityId: cosmosResponseMessage.Headers.ActivityId); + CosmosException exception; + if (cosmosResponseMessage.CosmosException != null) + { + exception = cosmosResponseMessage.CosmosException; + } + else + { + exception = new CosmosException( + cosmosResponseMessage.ErrorMessage, + cosmosResponseMessage.StatusCode, + (int)cosmosResponseMessage.Headers.SubStatusCode, + cosmosResponseMessage.Headers.ActivityId, + cosmosResponseMessage.Headers.RequestCharge); + } + + return TryCatch.FromException(exception); } if (!(cosmosResponseMessage.Content is MemoryStream memoryStream)) @@ -316,7 +327,7 @@ private static QueryResponseCore GetCosmosElementResponse( } long responseLengthBytes = memoryStream.Length; - CosmosArray cosmosArray = CosmosQueryClientCore.ParseElementsFromRestStream( + CosmosArray documents = CosmosQueryClientCore.ParseElementsFromRestStream( memoryStream, resourceType, requestOptions.CosmosSerializationFormatOptions); @@ -331,14 +342,14 @@ private static QueryResponseCore GetCosmosElementResponse( cosmosQueryExecutionInfo = default; } - return QueryResponseCore.CreateSuccess( - result: cosmosArray, - requestCharge: cosmosResponseMessage.Headers.RequestCharge, - activityId: cosmosResponseMessage.Headers.ActivityId, - responseLengthBytes: responseLengthBytes, - disallowContinuationTokenMessage: null, - continuationToken: cosmosResponseMessage.Headers.ContinuationToken, - cosmosQueryExecutionInfo: cosmosQueryExecutionInfo); + QueryPage response = new QueryPage( + documents, + cosmosResponseMessage.Headers.RequestCharge, + cosmosResponseMessage.Headers.ActivityId, + responseLengthBytes, + cosmosQueryExecutionInfo); + + return TryCatch.FromResult(response); } } From 7ab5e11f0e35183186a4e58f4152333b3fc646c6 Mon Sep 17 00:00:00 2001 From: Brandon Chong Date: Mon, 29 Jun 2020 13:14:46 -0700 Subject: [PATCH 15/85] made the continuation token part of the page --- .../CrossPartitionRangePageEnumerator.cs | 44 +++++++++++-------- Microsoft.Azure.Cosmos/src/Pagination/Page.cs | 10 +++-- .../PartitionRangePageEnumerator.cs | 10 +++-- ...CollectionPartitionRangeEnumeratorTests.cs | 23 ++++++++-- ...emoryCollectionPartitionRangeEnumerator.cs | 19 ++++---- 5 files changed, 66 insertions(+), 40 deletions(-) diff --git a/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionRangePageEnumerator.cs b/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionRangePageEnumerator.cs index 76baf73cca..26c432f49f 100644 --- a/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionRangePageEnumerator.cs +++ b/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionRangePageEnumerator.cs @@ -76,23 +76,6 @@ public CrossPartitionRangePageEnumerator( public TryCatch Current { get; private set; } - public State GetState() - { - if (!this.lazyEnumerators.ValueInitialized) - { - return this.originalState; - } - - PriorityQueue enumerators = this.lazyEnumerators.Result; - List<(FeedRange, State)> feedRangeAndStates = new List<(FeedRange, State)>(enumerators.Count); - foreach (PartitionRangePageEnumerator enumerator in enumerators) - { - feedRangeAndStates.Add((enumerator.Range, enumerator.State)); - } - - return new CrossPartitionState(feedRangeAndStates); - } - public async ValueTask MoveNextAsync() { PriorityQueue enumerators = await this.lazyEnumerators.GetValueAsync(cancellationToken: default); @@ -132,9 +115,23 @@ public async ValueTask MoveNextAsync() } } - this.Current = currentPaginator.Current; enumerators.Enqueue(currentPaginator); + TryCatch backendPage = currentPaginator.Current; + if (backendPage.Failed) + { + this.Current = backendPage; + return true; + } + + List<(FeedRange, State)> feedRangeAndStates = new List<(FeedRange, State)>(enumerators.Count); + foreach (PartitionRangePageEnumerator enumerator in enumerators) + { + feedRangeAndStates.Add((enumerator.Range, enumerator.State)); + } + + CrossPartitionState crossPartitionState = new CrossPartitionState(feedRangeAndStates); + this.Current = TryCatch.FromResult(new CrossPartitionPage(backendPage.Result, crossPartitionState)); return true; } @@ -157,6 +154,17 @@ private static bool IsMergeException(Exception exception) return false; } + public sealed class CrossPartitionPage : Page + { + public CrossPartitionPage(Page backendEndPage, State state) + : base(state) + { + this.Page = backendEndPage; + } + + public Page Page { get; } + } + private sealed class CrossPartitionState : State { public CrossPartitionState(List<(FeedRange, State)> value) diff --git a/Microsoft.Azure.Cosmos/src/Pagination/Page.cs b/Microsoft.Azure.Cosmos/src/Pagination/Page.cs index 24a9cd0e36..d169af0208 100644 --- a/Microsoft.Azure.Cosmos/src/Pagination/Page.cs +++ b/Microsoft.Azure.Cosmos/src/Pagination/Page.cs @@ -4,11 +4,13 @@ namespace Microsoft.Azure.Cosmos.Pagination { - using System; - using System.Collections.Generic; - using System.Text; - internal abstract class Page { + protected Page(State state) + { + this.State = state; + } + + public State State { get; } } } diff --git a/Microsoft.Azure.Cosmos/src/Pagination/PartitionRangePageEnumerator.cs b/Microsoft.Azure.Cosmos/src/Pagination/PartitionRangePageEnumerator.cs index 3f39ba1ea6..9120085d44 100644 --- a/Microsoft.Azure.Cosmos/src/Pagination/PartitionRangePageEnumerator.cs +++ b/Microsoft.Azure.Cosmos/src/Pagination/PartitionRangePageEnumerator.cs @@ -39,14 +39,16 @@ public async ValueTask MoveNextAsync() this.hasStarted = true; - (TryCatch page, State state) = await this.GetNextPageAsync(); - this.State = state; - this.Current = page; + this.Current = await this.GetNextPageAsync(); + if (this.Current.Succeeded) + { + this.State = this.Current.Result.State; + } return true; } - public abstract Task<(TryCatch, State)> GetNextPageAsync(CancellationToken cancellationToken = default); + public abstract Task> GetNextPageAsync(CancellationToken cancellationToken = default); public abstract ValueTask DisposeAsync(); } diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/CrossPartitionInMemoryCollectionPartitionRangeEnumeratorTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/CrossPartitionInMemoryCollectionPartitionRangeEnumeratorTests.cs index 1479fcd2e8..d3f8e6143f 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/CrossPartitionInMemoryCollectionPartitionRangeEnumeratorTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/CrossPartitionInMemoryCollectionPartitionRangeEnumeratorTests.cs @@ -137,7 +137,12 @@ public async Task TestSplitWithDuringDrainAsync() tryGetPage.ThrowIfFailed(); - if (!(tryGetPage.Result is InMemoryCollectionPartitionRangeEnumerator.InMemoryCollectionPage page)) + if (!(tryGetPage.Result is CrossPartitionRangePageEnumerator.CrossPartitionPage crossPartitionPage)) + { + throw new InvalidCastException(); + } + + if (!(crossPartitionPage.Page is InMemoryCollectionPartitionRangeEnumerator.InMemoryCollectionPage page)) { throw new InvalidCastException(); } @@ -192,7 +197,12 @@ private static async Task> DrainFullyAsync(CrossPartitionRangePage { tryGetPage.ThrowIfFailed(); - if (!(tryGetPage.Result is InMemoryCollectionPartitionRangeEnumerator.InMemoryCollectionPage page)) + if (!(tryGetPage.Result is CrossPartitionRangePageEnumerator.CrossPartitionPage crossPartitionPage)) + { + throw new InvalidCastException(); + } + + if (!(crossPartitionPage.Page is InMemoryCollectionPartitionRangeEnumerator.InMemoryCollectionPage page)) { throw new InvalidCastException(); } @@ -221,7 +231,12 @@ private static async Task> DrainFullyAsync(CrossPartitionRangePage TryCatch tryGetPage = enumerator.Current; tryGetPage.ThrowIfFailed(); - if (!(tryGetPage.Result is InMemoryCollectionPartitionRangeEnumerator.InMemoryCollectionPage page)) + if (!(tryGetPage.Result is CrossPartitionRangePageEnumerator.CrossPartitionPage crossPartitionPage)) + { + throw new InvalidCastException(); + } + + if (!(crossPartitionPage.Page is InMemoryCollectionPartitionRangeEnumerator.InMemoryCollectionPage page)) { throw new InvalidCastException(); } @@ -231,7 +246,7 @@ private static async Task> DrainFullyAsync(CrossPartitionRangePage identifiers.Add(record.Identifier); } - state = enumerator.GetState(); + state = tryGetPage.Result.State; } return (identifiers, state); diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/InMemoryCollectionPartitionRangeEnumerator.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/InMemoryCollectionPartitionRangeEnumerator.cs index 9759db97a3..877201b257 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/InMemoryCollectionPartitionRangeEnumerator.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/InMemoryCollectionPartitionRangeEnumerator.cs @@ -41,7 +41,7 @@ public override ValueTask DisposeAsync() return default; } - public override Task<(TryCatch, State)> GetNextPageAsync(CancellationToken cancellationToken = default) + public override Task> GetNextPageAsync(CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); @@ -51,21 +51,19 @@ public override ValueTask DisposeAsync() pageSize: this.pageSize); if (tryReadPage.Failed) { - return Task.FromResult((TryCatch.FromException(tryReadPage.Exception), this.State)); + return Task.FromResult(TryCatch.FromException(tryReadPage.Exception)); } if (tryReadPage.Result.Count == 0) { - InMemoryCollectionPage emptyPage = new InMemoryCollectionPage(new List()); - InMemoryCollectionState nullContinuation = default; - return Task.FromResult((TryCatch.FromResult(emptyPage), (State)nullContinuation)); + InMemoryCollectionPage emptyPage = new InMemoryCollectionPage(new List(), state: default); + return Task.FromResult(TryCatch.FromResult(emptyPage)); } - InMemoryCollectionPage page = new InMemoryCollectionPage(tryReadPage.Result); - State inMemoryCollectionState = new InMemoryCollectionState(page.Records.Last().ResourceIdentifier); - (TryCatch, State) tryPageAndState = (TryCatch.FromResult(page), inMemoryCollectionState); + State inMemoryCollectionState = new InMemoryCollectionState(tryReadPage.Result.Last().ResourceIdentifier); + InMemoryCollectionPage page = new InMemoryCollectionPage(tryReadPage.Result, inMemoryCollectionState); - return Task.FromResult(tryPageAndState); + return Task.FromResult(TryCatch.FromResult(page)); } private sealed class InMemoryCollectionState : State @@ -80,7 +78,8 @@ public InMemoryCollectionState(long resourceIdentifier) public sealed class InMemoryCollectionPage : Page { - public InMemoryCollectionPage(List records) + public InMemoryCollectionPage(List records, State state) + : base(state) { this.Records = records; } From 595f6adc1fe3e102ecead2a27b3006f89cc1313d Mon Sep 17 00:00:00 2001 From: Brandon Chong Date: Mon, 29 Jun 2020 16:13:28 -0700 Subject: [PATCH 16/85] added 429 tests --- .../InMemoryCollection.cs | 33 ++++- ...CollectionPartitionRangeEnumeratorTests.cs | 124 +++++++++++++++++- ...CollectionPartitionRangeEnumeratorTests.cs | 110 +++++++++++++++- 3 files changed, 259 insertions(+), 8 deletions(-) diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/InMemoryCollection.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/InMemoryCollection.cs index 38b2057cc2..683b046537 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/InMemoryCollection.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/InMemoryCollection.cs @@ -12,24 +12,34 @@ namespace Microsoft.Azure.Cosmos.Tests using Microsoft.Azure.Cosmos.Query.Core.Monads; using Microsoft.Azure.Cosmos.Routing; using Microsoft.Azure.Documents; - using Moq; // Collection useful for mocking requests and repartitioning (splits / merge). internal sealed class InMemoryCollection { + private static readonly CosmosException RequestRateTooLargeException = new CosmosException( + message: "Request Rate Too Large", + statusCode: (System.Net.HttpStatusCode)429, + subStatusCode: default, + activityId: Guid.NewGuid().ToString(), + requestCharge: default); private readonly PartitionKeyDefinition partitionKeyDefinition; private readonly Dictionary parentToChildMapping; + private readonly FailureConfigs failureConfigs; + private readonly Random random; private PartitionKeyHashRangeDictionary partitionedRecords; private Dictionary partitionKeyRangeIdToHashRange; - public InMemoryCollection(PartitionKeyDefinition partitionKeyDefinition) + public InMemoryCollection(PartitionKeyDefinition partitionKeyDefinition, FailureConfigs failureConfigs = default) { + this.partitionKeyDefinition = partitionKeyDefinition ?? throw new ArgumentNullException(nameof(partitionKeyDefinition)); + this.failureConfigs = failureConfigs; + this.random = new Random(); + PartitionKeyHashRange fullRange = new PartitionKeyHashRange(startInclusive: null, endExclusive: null); PartitionKeyHashRanges partitionKeyHashRanges = PartitionKeyHashRanges.Create(new PartitionKeyHashRange[] { fullRange }); this.partitionedRecords = new PartitionKeyHashRangeDictionary(partitionKeyHashRanges); this.partitionedRecords[fullRange] = new Records(); - this.partitionKeyDefinition = partitionKeyDefinition ?? throw new ArgumentNullException(nameof(partitionKeyDefinition)); this.partitionKeyRangeIdToHashRange = new Dictionary() { { 0, fullRange } @@ -83,6 +93,11 @@ record = default; public TryCatch> ReadFeed(int partitionKeyRangeId, long resourceIndentifer, int pageSize) { + if (this.Return429()) + { + return TryCatch>.FromException(RequestRateTooLargeException); + } + if (!this.partitionKeyRangeIdToHashRange.TryGetValue(partitionKeyRangeId, out PartitionKeyHashRange range)) { return TryCatch>.FromException( @@ -179,6 +194,8 @@ public void Split(int partitionKeyRangeId) public (int, int) GetChildRanges(int partitionKeyRangeId) => this.parentToChildMapping[partitionKeyRangeId]; + public bool Return429() => (this.failureConfigs != null) && this.failureConfigs.Inject429s && ((this.random.Next() % 2) == 0); + private static PartitionKeyHash GetHashFromPayload(CosmosObject payload, PartitionKeyDefinition partitionKeyDefinition) { CosmosElement partitionKey = GetPartitionKeyFromPayload(payload, partitionKeyDefinition); @@ -274,6 +291,16 @@ public static Record Create(long previousResourceIdentifier, CosmosObject payloa } } + public sealed class FailureConfigs + { + public FailureConfigs(bool inject429s) + { + this.Inject429s = inject429s; + } + + public bool Inject429s { get; } + } + private sealed class Records : IReadOnlyList { private readonly List storage; diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/CrossPartitionInMemoryCollectionPartitionRangeEnumeratorTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/CrossPartitionInMemoryCollectionPartitionRangeEnumeratorTests.cs index d3f8e6143f..00f2dcd85d 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/CrossPartitionInMemoryCollectionPartitionRangeEnumeratorTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/CrossPartitionInMemoryCollectionPartitionRangeEnumeratorTests.cs @@ -156,7 +156,127 @@ public async Task TestSplitWithDuringDrainAsync() Assert.AreEqual(numItems, identifiers.Count); } - private static InMemoryCollection CreateInMemoryCollection(int numItems) + [TestMethod] + public async Task Test429sAsync() + { + int numItems = 100; + InMemoryCollection inMemoryCollection = CreateInMemoryCollection(numItems, new InMemoryCollection.FailureConfigs(inject429s: true)); + CrossPartitionRangePageEnumerable enumerable = new CrossPartitionRangePageEnumerable( + feedRangeProvider: new InMemoryCollectionFeedRangeProvider(inMemoryCollection), + createPartitionRangeEnumerator: (range, state) => new InMemoryCollectionPartitionRangeEnumerator( + inMemoryCollection, + partitionKeyRangeId: int.Parse(((FeedRangePartitionKeyRange)range).PartitionKeyRangeId), + pageSize: 10, + state: state), + comparer: PartitionRangePageEnumeratorComparer.Singleton); + + HashSet identifiers = new HashSet(); + await foreach (TryCatch tryGetPage in enumerable) + { + if (tryGetPage.Failed) + { + Exception exception = tryGetPage.Exception; + while (exception.InnerException != null) + { + exception = exception.InnerException; + } + + if (!((exception is CosmosException cosmosException) && (cosmosException.StatusCode == (System.Net.HttpStatusCode)429))) + { + throw tryGetPage.Exception; + } + } + else + { + if (!(tryGetPage.Result is CrossPartitionRangePageEnumerator.CrossPartitionPage crossPartitionPage)) + { + throw new InvalidCastException(); + } + + if (!(crossPartitionPage.Page is InMemoryCollectionPartitionRangeEnumerator.InMemoryCollectionPage page)) + { + throw new InvalidCastException(); + } + + foreach (InMemoryCollection.Record record in page.Records) + { + identifiers.Add(record.Identifier); + } + } + } + + Assert.AreEqual(numItems, identifiers.Count); + } + + [TestMethod] + public async Task Test429sWithContinuationsAsync() + { + int numItems = 100; + InMemoryCollection inMemoryCollection = CreateInMemoryCollection(numItems, new InMemoryCollection.FailureConfigs(inject429s: true)); + + IFeedRangeProvider feedRangeProvider = new InMemoryCollectionFeedRangeProvider(inMemoryCollection); + CreatePartitionRangePageEnumerator createEnumerator = (Cosmos.FeedRange range, State state) => new InMemoryCollectionPartitionRangeEnumerator( + inMemoryCollection, + partitionKeyRangeId: int.Parse(((FeedRangePartitionKeyRange)range).PartitionKeyRangeId), + pageSize: 10, + state: state); + + CrossPartitionRangePageEnumerator enumerator = new CrossPartitionRangePageEnumerator( + feedRangeProvider: feedRangeProvider, + createPartitionRangeEnumerator: createEnumerator, + comparer: PartitionRangePageEnumeratorComparer.Singleton); + + HashSet identifiers = new HashSet(); + State state = default; + + while (await enumerator.MoveNextAsync()) + { + TryCatch tryGetPage = enumerator.Current; + if (tryGetPage.Failed) + { + Exception exception = tryGetPage.Exception; + while (exception.InnerException != null) + { + exception = exception.InnerException; + } + + if (!((exception is CosmosException cosmosException) && (cosmosException.StatusCode == (System.Net.HttpStatusCode)429))) + { + throw tryGetPage.Exception; + } + + // Create a new enumerator from that state to simulate when the user want's to start resume later from a continuation token. + enumerator = new CrossPartitionRangePageEnumerator( + feedRangeProvider: feedRangeProvider, + createPartitionRangeEnumerator: createEnumerator, + comparer: PartitionRangePageEnumeratorComparer.Singleton, + state: state); + } + else + { + if (!(tryGetPage.Result is CrossPartitionRangePageEnumerator.CrossPartitionPage crossPartitionPage)) + { + throw new InvalidCastException(); + } + + if (!(crossPartitionPage.Page is InMemoryCollectionPartitionRangeEnumerator.InMemoryCollectionPage page)) + { + throw new InvalidCastException(); + } + + foreach (InMemoryCollection.Record record in page.Records) + { + identifiers.Add(record.Identifier); + } + + state = tryGetPage.Result.State; + } + } + + Assert.AreEqual(numItems, identifiers.Count); + } + + private static InMemoryCollection CreateInMemoryCollection(int numItems, InMemoryCollection.FailureConfigs failureConfigs = default) { PartitionKeyDefinition partitionKeyDefinition = new PartitionKeyDefinition() { @@ -168,7 +288,7 @@ private static InMemoryCollection CreateInMemoryCollection(int numItems) Version = PartitionKeyDefinitionVersion.V2, }; - InMemoryCollection inMemoryCollection = new InMemoryCollection(partitionKeyDefinition); + InMemoryCollection inMemoryCollection = new InMemoryCollection(partitionKeyDefinition, failureConfigs); inMemoryCollection.Split(partitionKeyRangeId: 0); diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/InMemoryCollectionPartitionRangeEnumeratorTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/InMemoryCollectionPartitionRangeEnumeratorTests.cs index 0914dc7160..d9f967138f 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/InMemoryCollectionPartitionRangeEnumeratorTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/InMemoryCollectionPartitionRangeEnumeratorTests.cs @@ -13,6 +13,7 @@ namespace Microsoft.Azure.Cosmos.Tests.Pagination using System.Collections.Generic; using Microsoft.Azure.Cosmos.Query.Core.Monads; using System.Linq; + using System.Security.Policy; [TestClass] public class InMemoryCollectionPartitionRangeEnumeratorTests @@ -106,7 +107,110 @@ public async Task TestSplitAsync() Assert.AreEqual(numItems, parentIdentifiers.Count + childIdentifiers.Count); } - private static InMemoryCollection CreateInMemoryCollection(int numItems) + [TestMethod] + public async Task Test429sAsync() + { + int numItems = 100; + InMemoryCollection inMemoryCollection = CreateInMemoryCollection(numItems, new InMemoryCollection.FailureConfigs(inject429s: true)); + PartitionRangePageEnumerable enumerable = new PartitionRangePageEnumerable( + range: new FeedRangePartitionKeyRange("0"), + state: default, + (range, state) => new InMemoryCollectionPartitionRangeEnumerator( + inMemoryCollection, + partitionKeyRangeId: int.Parse(((FeedRangePartitionKeyRange)range).PartitionKeyRangeId), + pageSize: 10, + state: state)); + + HashSet identifiers = new HashSet(); + await foreach (TryCatch tryGetPage in enumerable) + { + if (tryGetPage.Failed) + { + Exception exception = tryGetPage.Exception; + while (exception.InnerException != null) + { + exception = exception.InnerException; + } + + if (!((exception is CosmosException cosmosException) && (cosmosException.StatusCode == (System.Net.HttpStatusCode)429))) + { + throw tryGetPage.Exception; + } + } + else + { + if (!(tryGetPage.Result is InMemoryCollectionPartitionRangeEnumerator.InMemoryCollectionPage page)) + { + throw new InvalidCastException(); + } + + foreach (InMemoryCollection.Record record in page.Records) + { + identifiers.Add(record.Identifier); + } + } + } + + Assert.AreEqual(numItems, identifiers.Count); + } + + [TestMethod] + public async Task Test429sWithContinuationsAsync() + { + int numItems = 100; + InMemoryCollection inMemoryCollection = CreateInMemoryCollection(numItems, new InMemoryCollection.FailureConfigs(inject429s: true)); + + HashSet identifiers = new HashSet(); + State state = default; + + InMemoryCollectionPartitionRangeEnumerator enumerator = new InMemoryCollectionPartitionRangeEnumerator( + inMemoryCollection, + partitionKeyRangeId: 0, + pageSize: 10); + + while (await enumerator.MoveNextAsync()) + { + TryCatch tryGetPage = enumerator.Current; + if (tryGetPage.Failed) + { + Exception exception = tryGetPage.Exception; + while (exception.InnerException != null) + { + exception = exception.InnerException; + } + + if (!((exception is CosmosException cosmosException) && (cosmosException.StatusCode == (System.Net.HttpStatusCode)429))) + { + throw tryGetPage.Exception; + } + + // Create a new enumerator from that state to simulate when the user want's to start resume later from a continuation token. + enumerator = new InMemoryCollectionPartitionRangeEnumerator( + inMemoryCollection, + partitionKeyRangeId: 0, + pageSize: 10, + state: state); + } + else + { + if (!(tryGetPage.Result is InMemoryCollectionPartitionRangeEnumerator.InMemoryCollectionPage page)) + { + throw new InvalidCastException(); + } + + foreach (InMemoryCollection.Record record in page.Records) + { + identifiers.Add(record.Identifier); + } + + state = tryGetPage.Result.State; + } + } + + Assert.AreEqual(numItems, identifiers.Count); + } + + private static InMemoryCollection CreateInMemoryCollection(int numItems, InMemoryCollection.FailureConfigs failureConfigs = default) { PartitionKeyDefinition partitionKeyDefinition = new PartitionKeyDefinition() { @@ -118,7 +222,7 @@ private static InMemoryCollection CreateInMemoryCollection(int numItems) Version = PartitionKeyDefinitionVersion.V2, }; - InMemoryCollection inMemoryCollection = new InMemoryCollection(partitionKeyDefinition); + InMemoryCollection inMemoryCollection = new InMemoryCollection(partitionKeyDefinition, failureConfigs); for (int i = 0; i < numItems; i++) { @@ -150,7 +254,7 @@ private static InMemoryCollection CreateInMemoryCollection(int numItems) throw new InvalidCastException(); } - foreach(InMemoryCollection.Record record in page.Records) + foreach (InMemoryCollection.Record record in page.Records) { identifiers.Add(record.Identifier); } From 0b4a808cd2a4328fab5faa7819b6e0798adcb3e6 Mon Sep 17 00:00:00 2001 From: Brandon Chong Date: Mon, 29 Jun 2020 17:53:50 -0700 Subject: [PATCH 17/85] refactored tests to have more common code --- ...CollectionPartitionRangeEnumeratorTests.cs | 257 +++--------------- ...emoryCollectionPartitionRangeEnumerator.cs | 1 - ...CollectionPartitionRangeEnumeratorTests.cs | 214 ++++----------- ...CollectionPartitionRangeEnumeratorTests.cs | 113 ++++++++ 4 files changed, 199 insertions(+), 386 deletions(-) create mode 100644 Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/SinglePartitionInMemoryCollectionPartitionRangeEnumeratorTests.cs diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/CrossPartitionInMemoryCollectionPartitionRangeEnumeratorTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/CrossPartitionInMemoryCollectionPartitionRangeEnumeratorTests.cs index 00f2dcd85d..66c92af832 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/CrossPartitionInMemoryCollectionPartitionRangeEnumeratorTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/CrossPartitionInMemoryCollectionPartitionRangeEnumeratorTests.cs @@ -15,63 +15,13 @@ namespace Microsoft.Azure.Cosmos.Tests.Pagination using Microsoft.VisualStudio.TestTools.UnitTesting; [TestClass] - public class CrossPartitionInMemoryCollectionPartitionRangeEnumeratorTests + public class CrossPartitionInMemoryCollectionPartitionRangeEnumeratorTests : InMemoryCollectionPartitionRangeEnumeratorTests { - [TestMethod] - public async Task TestDrainFullyAsync() - { - int numItems = 1000; - InMemoryCollection inMemoryCollection = CreateInMemoryCollection(numItems); - - CrossPartitionRangePageEnumerable enumerable = new CrossPartitionRangePageEnumerable( - feedRangeProvider: new InMemoryCollectionFeedRangeProvider(inMemoryCollection), - createPartitionRangeEnumerator: (range, state) => new InMemoryCollectionPartitionRangeEnumerator( - inMemoryCollection, - partitionKeyRangeId: int.Parse(((FeedRangePartitionKeyRange)range).PartitionKeyRangeId), - pageSize: 10, - state: state), - comparer: PartitionRangePageEnumeratorComparer.Singleton); - - HashSet identifiers = await DrainFullyAsync(enumerable); - Assert.AreEqual(numItems, identifiers.Count); - } - - [TestMethod] - public async Task TestResumingFromStateAsync() - { - int numItems = 1000; - InMemoryCollection inMemoryCollection = CreateInMemoryCollection(numItems); - - IFeedRangeProvider feedRangeProvider = new InMemoryCollectionFeedRangeProvider(inMemoryCollection); - CreatePartitionRangePageEnumerator createEnumerator = (Cosmos.FeedRange range, State state) => new InMemoryCollectionPartitionRangeEnumerator( - inMemoryCollection, - partitionKeyRangeId: int.Parse(((FeedRangePartitionKeyRange)range).PartitionKeyRangeId), - pageSize: 10, - state: state); - - CrossPartitionRangePageEnumerator enumerator = new CrossPartitionRangePageEnumerator( - feedRangeProvider: feedRangeProvider, - createPartitionRangeEnumerator: createEnumerator, - comparer: PartitionRangePageEnumeratorComparer.Singleton); - - (HashSet firstDrainResults, State state) = await PartialDrainAsync(enumerator, numIterations: 3); - - // Resume from state - CrossPartitionRangePageEnumerable enumerable = new CrossPartitionRangePageEnumerable( - feedRangeProvider: feedRangeProvider, - createPartitionRangeEnumerator: createEnumerator, - comparer: PartitionRangePageEnumeratorComparer.Singleton, - state: state); - - HashSet secondDrainResults = await DrainFullyAsync(enumerable); - Assert.AreEqual(numItems, firstDrainResults.Count + secondDrainResults.Count); - } - [TestMethod] public async Task TestSplitWithResumeContinuationAsync() { int numItems = 1000; - InMemoryCollection inMemoryCollection = CreateInMemoryCollection(numItems); + InMemoryCollection inMemoryCollection = this.CreateInMemoryCollection(numItems); IFeedRangeProvider feedRangeProvider = new InMemoryCollectionFeedRangeProvider(inMemoryCollection); CreatePartitionRangePageEnumerator createEnumerator = (Cosmos.FeedRange range, State state) => new InMemoryCollectionPartitionRangeEnumerator( @@ -85,7 +35,7 @@ public async Task TestSplitWithResumeContinuationAsync() createPartitionRangeEnumerator: createEnumerator, comparer: PartitionRangePageEnumeratorComparer.Singleton); - (HashSet firstDrainResults, State state) = await PartialDrainAsync(enumerator, numIterations: 3); + (HashSet firstDrainResults, State state) = await this.PartialDrainAsync(enumerator, numIterations: 3); int minPartitionKeyRangeId = inMemoryCollection.PartitionKeyRangeFeedReed().Keys.Min(); int maxPartitionKeyRangeId = inMemoryCollection.PartitionKeyRangeFeedReed().Keys.Max(); @@ -102,7 +52,7 @@ public async Task TestSplitWithResumeContinuationAsync() comparer: PartitionRangePageEnumeratorComparer.Singleton, state: state); - HashSet secondDrainResults = await DrainFullyAsync(enumerable); + HashSet secondDrainResults = await this.DrainFullyAsync(enumerable); Assert.AreEqual(numItems, firstDrainResults.Count + secondDrainResults.Count); } @@ -110,7 +60,7 @@ public async Task TestSplitWithResumeContinuationAsync() public async Task TestSplitWithDuringDrainAsync() { int numItems = 1000; - InMemoryCollection inMemoryCollection = CreateInMemoryCollection(numItems); + InMemoryCollection inMemoryCollection = this.CreateInMemoryCollection(numItems); IFeedRangeProvider feedRangeProvider = new InMemoryCollectionFeedRangeProvider(inMemoryCollection); CreatePartitionRangePageEnumerator createEnumerator = (Cosmos.FeedRange range, State state) => new InMemoryCollectionPartitionRangeEnumerator( @@ -156,127 +106,22 @@ public async Task TestSplitWithDuringDrainAsync() Assert.AreEqual(numItems, identifiers.Count); } - [TestMethod] - public async Task Test429sAsync() + internal override List GetRecordsFromPage(Page page) { - int numItems = 100; - InMemoryCollection inMemoryCollection = CreateInMemoryCollection(numItems, new InMemoryCollection.FailureConfigs(inject429s: true)); - CrossPartitionRangePageEnumerable enumerable = new CrossPartitionRangePageEnumerable( - feedRangeProvider: new InMemoryCollectionFeedRangeProvider(inMemoryCollection), - createPartitionRangeEnumerator: (range, state) => new InMemoryCollectionPartitionRangeEnumerator( - inMemoryCollection, - partitionKeyRangeId: int.Parse(((FeedRangePartitionKeyRange)range).PartitionKeyRangeId), - pageSize: 10, - state: state), - comparer: PartitionRangePageEnumeratorComparer.Singleton); - - HashSet identifiers = new HashSet(); - await foreach (TryCatch tryGetPage in enumerable) + if (!(page is CrossPartitionRangePageEnumerator.CrossPartitionPage crossPartitionPage)) { - if (tryGetPage.Failed) - { - Exception exception = tryGetPage.Exception; - while (exception.InnerException != null) - { - exception = exception.InnerException; - } - - if (!((exception is CosmosException cosmosException) && (cosmosException.StatusCode == (System.Net.HttpStatusCode)429))) - { - throw tryGetPage.Exception; - } - } - else - { - if (!(tryGetPage.Result is CrossPartitionRangePageEnumerator.CrossPartitionPage crossPartitionPage)) - { - throw new InvalidCastException(); - } - - if (!(crossPartitionPage.Page is InMemoryCollectionPartitionRangeEnumerator.InMemoryCollectionPage page)) - { - throw new InvalidCastException(); - } - - foreach (InMemoryCollection.Record record in page.Records) - { - identifiers.Add(record.Identifier); - } - } + throw new InvalidCastException(); } - Assert.AreEqual(numItems, identifiers.Count); - } - - [TestMethod] - public async Task Test429sWithContinuationsAsync() - { - int numItems = 100; - InMemoryCollection inMemoryCollection = CreateInMemoryCollection(numItems, new InMemoryCollection.FailureConfigs(inject429s: true)); - - IFeedRangeProvider feedRangeProvider = new InMemoryCollectionFeedRangeProvider(inMemoryCollection); - CreatePartitionRangePageEnumerator createEnumerator = (Cosmos.FeedRange range, State state) => new InMemoryCollectionPartitionRangeEnumerator( - inMemoryCollection, - partitionKeyRangeId: int.Parse(((FeedRangePartitionKeyRange)range).PartitionKeyRangeId), - pageSize: 10, - state: state); - - CrossPartitionRangePageEnumerator enumerator = new CrossPartitionRangePageEnumerator( - feedRangeProvider: feedRangeProvider, - createPartitionRangeEnumerator: createEnumerator, - comparer: PartitionRangePageEnumeratorComparer.Singleton); - - HashSet identifiers = new HashSet(); - State state = default; - - while (await enumerator.MoveNextAsync()) + if (!(crossPartitionPage.Page is InMemoryCollectionPartitionRangeEnumerator.InMemoryCollectionPage inMemoryCollectionPage)) { - TryCatch tryGetPage = enumerator.Current; - if (tryGetPage.Failed) - { - Exception exception = tryGetPage.Exception; - while (exception.InnerException != null) - { - exception = exception.InnerException; - } - - if (!((exception is CosmosException cosmosException) && (cosmosException.StatusCode == (System.Net.HttpStatusCode)429))) - { - throw tryGetPage.Exception; - } - - // Create a new enumerator from that state to simulate when the user want's to start resume later from a continuation token. - enumerator = new CrossPartitionRangePageEnumerator( - feedRangeProvider: feedRangeProvider, - createPartitionRangeEnumerator: createEnumerator, - comparer: PartitionRangePageEnumeratorComparer.Singleton, - state: state); - } - else - { - if (!(tryGetPage.Result is CrossPartitionRangePageEnumerator.CrossPartitionPage crossPartitionPage)) - { - throw new InvalidCastException(); - } - - if (!(crossPartitionPage.Page is InMemoryCollectionPartitionRangeEnumerator.InMemoryCollectionPage page)) - { - throw new InvalidCastException(); - } - - foreach (InMemoryCollection.Record record in page.Records) - { - identifiers.Add(record.Identifier); - } - - state = tryGetPage.Result.State; - } + throw new InvalidCastException(); } - Assert.AreEqual(numItems, identifiers.Count); + return inMemoryCollectionPage.Records; } - private static InMemoryCollection CreateInMemoryCollection(int numItems, InMemoryCollection.FailureConfigs failureConfigs = default) + internal override InMemoryCollection CreateInMemoryCollection(int numItems, InMemoryCollection.FailureConfigs failureConfigs = null) { PartitionKeyDefinition partitionKeyDefinition = new PartitionKeyDefinition() { @@ -310,66 +155,38 @@ private static InMemoryCollection CreateInMemoryCollection(int numItems, InMemor return inMemoryCollection; } - private static async Task> DrainFullyAsync(CrossPartitionRangePageEnumerable enumerable) + internal override IAsyncEnumerable> CreateEnumerable(InMemoryCollection inMemoryCollection, State state = null) { - HashSet identifiers = new HashSet(); - await foreach (TryCatch tryGetPage in enumerable) - { - tryGetPage.ThrowIfFailed(); - - if (!(tryGetPage.Result is CrossPartitionRangePageEnumerator.CrossPartitionPage crossPartitionPage)) - { - throw new InvalidCastException(); - } - - if (!(crossPartitionPage.Page is InMemoryCollectionPartitionRangeEnumerator.InMemoryCollectionPage page)) - { - throw new InvalidCastException(); - } - - foreach (InMemoryCollection.Record record in page.Records) - { - identifiers.Add(record.Identifier); - } - } + IFeedRangeProvider feedRangeProvider = new InMemoryCollectionFeedRangeProvider(inMemoryCollection); + CreatePartitionRangePageEnumerator createEnumerator = (Cosmos.FeedRange range, State state) => new InMemoryCollectionPartitionRangeEnumerator( + inMemoryCollection, + partitionKeyRangeId: int.Parse(((FeedRangePartitionKeyRange)range).PartitionKeyRangeId), + pageSize: 10, + state: state); - return identifiers; + return new CrossPartitionRangePageEnumerable( + feedRangeProvider: feedRangeProvider, + createPartitionRangeEnumerator: createEnumerator, + comparer: PartitionRangePageEnumeratorComparer.Singleton, + state: state); } - private static async Task<(HashSet, State)> PartialDrainAsync( - CrossPartitionRangePageEnumerator enumerator, - int numIterations) + internal override IAsyncEnumerator> CreateEnumerator(InMemoryCollection inMemoryCollection, State state = null) { - HashSet identifiers = new HashSet(); - State state = default; - - // Drain a couple of iterations - for (int i = 0; i < numIterations; i++) - { - await enumerator.MoveNextAsync(); - - TryCatch tryGetPage = enumerator.Current; - tryGetPage.ThrowIfFailed(); - - if (!(tryGetPage.Result is CrossPartitionRangePageEnumerator.CrossPartitionPage crossPartitionPage)) - { - throw new InvalidCastException(); - } - - if (!(crossPartitionPage.Page is InMemoryCollectionPartitionRangeEnumerator.InMemoryCollectionPage page)) - { - throw new InvalidCastException(); - } - - foreach (InMemoryCollection.Record record in page.Records) - { - identifiers.Add(record.Identifier); - } + IFeedRangeProvider feedRangeProvider = new InMemoryCollectionFeedRangeProvider(inMemoryCollection); + CreatePartitionRangePageEnumerator createEnumerator = (Cosmos.FeedRange range, State state) => new InMemoryCollectionPartitionRangeEnumerator( + inMemoryCollection, + partitionKeyRangeId: int.Parse(((FeedRangePartitionKeyRange)range).PartitionKeyRangeId), + pageSize: 10, + state: state); - state = tryGetPage.Result.State; - } + CrossPartitionRangePageEnumerator enumerator = new CrossPartitionRangePageEnumerator( + feedRangeProvider: feedRangeProvider, + createPartitionRangeEnumerator: createEnumerator, + comparer: PartitionRangePageEnumeratorComparer.Singleton, + state: state); - return (identifiers, state); + return enumerator; } private sealed class PartitionRangePageEnumeratorComparer : IComparer diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/InMemoryCollectionPartitionRangeEnumerator.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/InMemoryCollectionPartitionRangeEnumerator.cs index 877201b257..79a2edf2d2 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/InMemoryCollectionPartitionRangeEnumerator.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/InMemoryCollectionPartitionRangeEnumerator.cs @@ -11,7 +11,6 @@ namespace Microsoft.Azure.Cosmos.Tests.Pagination using System.Threading.Tasks; using Microsoft.Azure.Cosmos.Pagination; using Microsoft.Azure.Cosmos.Query.Core.Monads; - using Microsoft.Azure.Documents; internal sealed class InMemoryCollectionPartitionRangeEnumerator : PartitionRangePageEnumerator { diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/InMemoryCollectionPartitionRangeEnumeratorTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/InMemoryCollectionPartitionRangeEnumeratorTests.cs index d9f967138f..6836b2b1ef 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/InMemoryCollectionPartitionRangeEnumeratorTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/InMemoryCollectionPartitionRangeEnumeratorTests.cs @@ -1,125 +1,46 @@ -//------------------------------------------------------------ -// Copyright (c) Microsoft Corporation. All rights reserved. -//------------------------------------------------------------ - -namespace Microsoft.Azure.Cosmos.Tests.Pagination +namespace Microsoft.Azure.Cosmos.Tests.Pagination { using System; + using System.Collections.Generic; using System.Threading.Tasks; - using Microsoft.VisualStudio.TestTools.UnitTesting; - using Microsoft.Azure.Documents; - using Microsoft.Azure.Cosmos.CosmosElements; using Microsoft.Azure.Cosmos.Pagination; - using System.Collections.Generic; using Microsoft.Azure.Cosmos.Query.Core.Monads; - using System.Linq; - using System.Security.Policy; + using Microsoft.VisualStudio.TestTools.UnitTesting; - [TestClass] - public class InMemoryCollectionPartitionRangeEnumeratorTests + public abstract class InMemoryCollectionPartitionRangeEnumeratorTests { [TestMethod] public async Task TestDrainFullyAsync() { - int numItems = 100; - InMemoryCollection inMemoryCollection = CreateInMemoryCollection(numItems); - PartitionRangePageEnumerable enumerable = new PartitionRangePageEnumerable( - range: new FeedRangePartitionKeyRange("0"), - state: default, - (range, state) => new InMemoryCollectionPartitionRangeEnumerator( - inMemoryCollection, - partitionKeyRangeId: int.Parse(((FeedRangePartitionKeyRange)range).PartitionKeyRangeId), - pageSize: 10, - state: state)); - - HashSet identifiers = await DrainFullyAsync(enumerable); + int numItems = 1000; + InMemoryCollection inMemoryCollection = this.CreateInMemoryCollection(numItems); + IAsyncEnumerable> enumerable = this.CreateEnumerable(inMemoryCollection); + HashSet identifiers = await this.DrainFullyAsync(enumerable); Assert.AreEqual(numItems, identifiers.Count); } [TestMethod] public async Task TestResumingFromStateAsync() { - int numItems = 100; - InMemoryCollection inMemoryCollection = CreateInMemoryCollection(numItems); - InMemoryCollectionPartitionRangeEnumerator enumerator = new InMemoryCollectionPartitionRangeEnumerator( - inMemoryCollection, - partitionKeyRangeId: 0, - pageSize: 10); - - (HashSet firstDrainResults, State state) = await PartialDrainAsync(enumerator, numIterations: 3); - - // Resume from state - enumerator = new InMemoryCollectionPartitionRangeEnumerator( - inMemoryCollection, - partitionKeyRangeId: 0, - pageSize: 10, - state: state); - - PartitionRangePageEnumerable enumerable = new PartitionRangePageEnumerable( - range: new FeedRangePartitionKeyRange("0"), - state: state, - (range, state) => new InMemoryCollectionPartitionRangeEnumerator( - inMemoryCollection, - partitionKeyRangeId: int.Parse(((FeedRangePartitionKeyRange)range).PartitionKeyRangeId), - pageSize: 10, - state: state)); - HashSet secondDrainResults = await DrainFullyAsync(enumerable); - - Assert.AreEqual(numItems, firstDrainResults.Count + secondDrainResults.Count); - } - - [TestMethod] - public async Task TestSplitAsync() - { - int numItems = 100; - InMemoryCollection inMemoryCollection = CreateInMemoryCollection(numItems); - InMemoryCollectionPartitionRangeEnumerator enumerator = new InMemoryCollectionPartitionRangeEnumerator( - inMemoryCollection, - partitionKeyRangeId: 0, - pageSize: 10); - - (HashSet parentIdentifiers, State state) = await PartialDrainAsync(enumerator, numIterations: 3); + int numItems = 1000; + InMemoryCollection inMemoryCollection = this.CreateInMemoryCollection(numItems); - // Split the partition - inMemoryCollection.Split(partitionKeyRangeId: 0); + IAsyncEnumerator> enumerator = this.CreateEnumerator(inMemoryCollection); + (HashSet firstDrainResults, State state) = await this.PartialDrainAsync(enumerator, numIterations: 3); - // Try To read from the partition that is gone. - await enumerator.MoveNextAsync(); - Assert.IsTrue(enumerator.Current.Failed); + IAsyncEnumerable> enumerable = this.CreateEnumerable(inMemoryCollection, state); + HashSet secondDrainResults = await this.DrainFullyAsync(enumerable); - // Resume on the children using the parent continuaiton token - HashSet childIdentifiers = new HashSet(); - foreach (int partitionKeyRangeId in new int[] { 1, 2 }) - { - PartitionRangePageEnumerable enumerable1 = new PartitionRangePageEnumerable( - range: new FeedRangePartitionKeyRange(partitionKeyRangeId.ToString()), - state: state, - (range, state) => new InMemoryCollectionPartitionRangeEnumerator( - inMemoryCollection, - partitionKeyRangeId: int.Parse(((FeedRangePartitionKeyRange)range).PartitionKeyRangeId), - pageSize: 10, - state: state)); - HashSet resourceIdentifiers = await DrainFullyAsync(enumerable1); - - childIdentifiers.UnionWith(resourceIdentifiers); - } - - Assert.AreEqual(numItems, parentIdentifiers.Count + childIdentifiers.Count); + Assert.AreEqual(numItems, firstDrainResults.Count + secondDrainResults.Count); } [TestMethod] public async Task Test429sAsync() { int numItems = 100; - InMemoryCollection inMemoryCollection = CreateInMemoryCollection(numItems, new InMemoryCollection.FailureConfigs(inject429s: true)); - PartitionRangePageEnumerable enumerable = new PartitionRangePageEnumerable( - range: new FeedRangePartitionKeyRange("0"), - state: default, - (range, state) => new InMemoryCollectionPartitionRangeEnumerator( - inMemoryCollection, - partitionKeyRangeId: int.Parse(((FeedRangePartitionKeyRange)range).PartitionKeyRangeId), - pageSize: 10, - state: state)); + InMemoryCollection inMemoryCollection = this.CreateInMemoryCollection(numItems, new InMemoryCollection.FailureConfigs(inject429s: true)); + + IAsyncEnumerable> enumerable = this.CreateEnumerable(inMemoryCollection); HashSet identifiers = new HashSet(); await foreach (TryCatch tryGetPage in enumerable) @@ -139,12 +60,8 @@ public async Task Test429sAsync() } else { - if (!(tryGetPage.Result is InMemoryCollectionPartitionRangeEnumerator.InMemoryCollectionPage page)) - { - throw new InvalidCastException(); - } - - foreach (InMemoryCollection.Record record in page.Records) + List records = this.GetRecordsFromPage(tryGetPage.Result); + foreach (InMemoryCollection.Record record in records) { identifiers.Add(record.Identifier); } @@ -158,16 +75,13 @@ public async Task Test429sAsync() public async Task Test429sWithContinuationsAsync() { int numItems = 100; - InMemoryCollection inMemoryCollection = CreateInMemoryCollection(numItems, new InMemoryCollection.FailureConfigs(inject429s: true)); + InMemoryCollection inMemoryCollection = this.CreateInMemoryCollection(numItems, new InMemoryCollection.FailureConfigs(inject429s: true)); + + IAsyncEnumerator> enumerator = this.CreateEnumerator(inMemoryCollection); HashSet identifiers = new HashSet(); State state = default; - InMemoryCollectionPartitionRangeEnumerator enumerator = new InMemoryCollectionPartitionRangeEnumerator( - inMemoryCollection, - partitionKeyRangeId: 0, - pageSize: 10); - while (await enumerator.MoveNextAsync()) { TryCatch tryGetPage = enumerator.Current; @@ -185,20 +99,12 @@ public async Task Test429sWithContinuationsAsync() } // Create a new enumerator from that state to simulate when the user want's to start resume later from a continuation token. - enumerator = new InMemoryCollectionPartitionRangeEnumerator( - inMemoryCollection, - partitionKeyRangeId: 0, - pageSize: 10, - state: state); + enumerator = this.CreateEnumerator(inMemoryCollection, state); } else { - if (!(tryGetPage.Result is InMemoryCollectionPartitionRangeEnumerator.InMemoryCollectionPage page)) - { - throw new InvalidCastException(); - } - - foreach (InMemoryCollection.Record record in page.Records) + List records = this.GetRecordsFromPage(tryGetPage.Result); + foreach (InMemoryCollection.Record record in records) { identifiers.Add(record.Identifier); } @@ -210,32 +116,34 @@ public async Task Test429sWithContinuationsAsync() Assert.AreEqual(numItems, identifiers.Count); } - private static InMemoryCollection CreateInMemoryCollection(int numItems, InMemoryCollection.FailureConfigs failureConfigs = default) + internal abstract List GetRecordsFromPage(Page page); + + internal abstract InMemoryCollection CreateInMemoryCollection(int numItems, InMemoryCollection.FailureConfigs failureConfigs = default); + + internal abstract IAsyncEnumerable> CreateEnumerable(InMemoryCollection inMemoryCollection, State state = null); + + internal abstract IAsyncEnumerator> CreateEnumerator(InMemoryCollection inMemoryCollection, State state = null); + + internal async Task> DrainFullyAsync(IAsyncEnumerable> enumerable) { - PartitionKeyDefinition partitionKeyDefinition = new PartitionKeyDefinition() + HashSet identifiers = new HashSet(); + await foreach (TryCatch tryGetPage in enumerable) { - Paths = new System.Collections.ObjectModel.Collection() - { - "/pk" - }, - Kind = PartitionKind.Hash, - Version = PartitionKeyDefinitionVersion.V2, - }; + tryGetPage.ThrowIfFailed(); - InMemoryCollection inMemoryCollection = new InMemoryCollection(partitionKeyDefinition, failureConfigs); + List records = this.GetRecordsFromPage(tryGetPage.Result); - for (int i = 0; i < numItems; i++) - { - // Insert an item - CosmosObject item = CosmosObject.Parse($"{{\"pk\" : {i} }}"); - inMemoryCollection.CreateItem(item); + foreach (InMemoryCollection.Record record in records) + { + identifiers.Add(record.Identifier); + } } - return inMemoryCollection; + return identifiers; } - private static async Task<(HashSet, State)> PartialDrainAsync( - PartitionRangePageEnumerator enumerator, + internal async Task<(HashSet, State)> PartialDrainAsync( + IAsyncEnumerator> enumerator, int numIterations) { HashSet identifiers = new HashSet(); @@ -249,41 +157,17 @@ private static InMemoryCollection CreateInMemoryCollection(int numItems, InMemor TryCatch tryGetPage = enumerator.Current; tryGetPage.ThrowIfFailed(); - if (!(tryGetPage.Result is InMemoryCollectionPartitionRangeEnumerator.InMemoryCollectionPage page)) - { - throw new InvalidCastException(); - } + List records = this.GetRecordsFromPage(tryGetPage.Result); - foreach (InMemoryCollection.Record record in page.Records) + foreach (InMemoryCollection.Record record in records) { identifiers.Add(record.Identifier); } - state = enumerator.State; + state = tryGetPage.Result.State; } return (identifiers, state); } - - private static async Task> DrainFullyAsync(PartitionRangePageEnumerable enumerable) - { - HashSet identifiers = new HashSet(); - await foreach (TryCatch tryGetPage in enumerable) - { - tryGetPage.ThrowIfFailed(); - - if (!(tryGetPage.Result is InMemoryCollectionPartitionRangeEnumerator.InMemoryCollectionPage page)) - { - throw new InvalidCastException(); - } - - foreach (InMemoryCollection.Record record in page.Records) - { - identifiers.Add(record.Identifier); - } - } - - return identifiers; - } } } diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/SinglePartitionInMemoryCollectionPartitionRangeEnumeratorTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/SinglePartitionInMemoryCollectionPartitionRangeEnumeratorTests.cs new file mode 100644 index 0000000000..d456ee8d37 --- /dev/null +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/SinglePartitionInMemoryCollectionPartitionRangeEnumeratorTests.cs @@ -0,0 +1,113 @@ +//------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +//------------------------------------------------------------ + +namespace Microsoft.Azure.Cosmos.Tests.Pagination +{ + using System; + using System.Threading.Tasks; + using Microsoft.VisualStudio.TestTools.UnitTesting; + using Microsoft.Azure.Documents; + using Microsoft.Azure.Cosmos.CosmosElements; + using Microsoft.Azure.Cosmos.Pagination; + using System.Collections.Generic; + using Microsoft.Azure.Cosmos.Query.Core.Monads; + + [TestClass] + public class SinglePartitionInMemoryCollectionPartitionRangeEnumeratorTests : InMemoryCollectionPartitionRangeEnumeratorTests + { + [TestMethod] + public async Task TestSplitAsync() + { + int numItems = 100; + InMemoryCollection inMemoryCollection = this.CreateInMemoryCollection(numItems); + InMemoryCollectionPartitionRangeEnumerator enumerator = new InMemoryCollectionPartitionRangeEnumerator( + inMemoryCollection, + partitionKeyRangeId: 0, + pageSize: 10); + + (HashSet parentIdentifiers, State state) = await this.PartialDrainAsync(enumerator, numIterations: 3); + + // Split the partition + inMemoryCollection.Split(partitionKeyRangeId: 0); + + // Try To read from the partition that is gone. + await enumerator.MoveNextAsync(); + Assert.IsTrue(enumerator.Current.Failed); + + // Resume on the children using the parent continuaiton token + HashSet childIdentifiers = new HashSet(); + foreach (int partitionKeyRangeId in new int[] { 1, 2 }) + { + PartitionRangePageEnumerable enumerable1 = new PartitionRangePageEnumerable( + range: new FeedRangePartitionKeyRange(partitionKeyRangeId.ToString()), + state: state, + (range, state) => new InMemoryCollectionPartitionRangeEnumerator( + inMemoryCollection, + partitionKeyRangeId: int.Parse(((FeedRangePartitionKeyRange)range).PartitionKeyRangeId), + pageSize: 10, + state: state)); + HashSet resourceIdentifiers = await this.DrainFullyAsync(enumerable1); + + childIdentifiers.UnionWith(resourceIdentifiers); + } + + Assert.AreEqual(numItems, parentIdentifiers.Count + childIdentifiers.Count); + } + + internal override List GetRecordsFromPage(Page page) + { + if (!(page is InMemoryCollectionPartitionRangeEnumerator.InMemoryCollectionPage inMemoryCollectionPage)) + { + throw new InvalidCastException(); + } + + return inMemoryCollectionPage.Records; + } + + internal override InMemoryCollection CreateInMemoryCollection(int numItems, InMemoryCollection.FailureConfigs failureConfigs = null) + { + PartitionKeyDefinition partitionKeyDefinition = new PartitionKeyDefinition() + { + Paths = new System.Collections.ObjectModel.Collection() + { + "/pk" + }, + Kind = PartitionKind.Hash, + Version = PartitionKeyDefinitionVersion.V2, + }; + + InMemoryCollection inMemoryCollection = new InMemoryCollection(partitionKeyDefinition, failureConfigs); + + for (int i = 0; i < numItems; i++) + { + // Insert an item + CosmosObject item = CosmosObject.Parse($"{{\"pk\" : {i} }}"); + inMemoryCollection.CreateItem(item); + } + + return inMemoryCollection; + } + + internal override IAsyncEnumerable> CreateEnumerable(InMemoryCollection inMemoryCollection, State state = null) + { + return new PartitionRangePageEnumerable( + range: new FeedRangePartitionKeyRange("0"), + state: state, + (range, state) => new InMemoryCollectionPartitionRangeEnumerator( + inMemoryCollection, + partitionKeyRangeId: int.Parse(((FeedRangePartitionKeyRange)range).PartitionKeyRangeId), + pageSize: 10, + state: state)); + } + + internal override IAsyncEnumerator> CreateEnumerator(InMemoryCollection inMemoryCollection, State state = null) + { + return new InMemoryCollectionPartitionRangeEnumerator( + inMemoryCollection, + partitionKeyRangeId: 0, + pageSize: 10, + state: state); + } + } +} From 2e1b8276d21f2eeba2ecb86e9fcccc1e44dc7703 Mon Sep 17 00:00:00 2001 From: Brandon Chong Date: Mon, 29 Jun 2020 18:22:54 -0700 Subject: [PATCH 18/85] added an empty page test --- .../InMemoryCollection.cs | 25 ++++++++++++++---- .../InMemoryCollectionTests.cs | 9 ++++--- ...emoryCollectionPartitionRangeEnumerator.cs | 14 ++++------ ...CollectionPartitionRangeEnumeratorTests.cs | 26 +++++++++++++++++-- 4 files changed, 54 insertions(+), 20 deletions(-) diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/InMemoryCollection.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/InMemoryCollection.cs index 683b046537..023426bc2d 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/InMemoryCollection.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/InMemoryCollection.cs @@ -91,16 +91,21 @@ record = default; return false; } - public TryCatch> ReadFeed(int partitionKeyRangeId, long resourceIndentifer, int pageSize) + public TryCatch<(List, long?)> ReadFeed(int partitionKeyRangeId, long resourceIndentifer, int pageSize) { if (this.Return429()) { - return TryCatch>.FromException(RequestRateTooLargeException); + return TryCatch<(List, long?)>.FromException(RequestRateTooLargeException); + } + + if (this.ReturnEmptyPage()) + { + return TryCatch<(List, long?)>.FromResult((new List(), resourceIndentifer)); } if (!this.partitionKeyRangeIdToHashRange.TryGetValue(partitionKeyRangeId, out PartitionKeyHashRange range)) { - return TryCatch>.FromException( + return TryCatch<(List, long?)>.FromException( new CosmosException( message: $"PartitionKeyRangeId {partitionKeyRangeId} is gone", statusCode: System.Net.HttpStatusCode.Gone, @@ -119,7 +124,12 @@ public TryCatch> ReadFeed(int partitionKeyRangeId, long resourceInd .Take(pageSize) .ToList(); - return TryCatch>.FromResult(page); + if (page.Count == 0) + { + return TryCatch<(List, long?)>.FromResult((page, null)); + } + + return TryCatch<(List, long?)>.FromResult((page, page.Last().ResourceIdentifier)); } public IReadOnlyDictionary PartitionKeyRangeFeedReed() => this.partitionKeyRangeIdToHashRange; @@ -196,6 +206,8 @@ public void Split(int partitionKeyRangeId) public bool Return429() => (this.failureConfigs != null) && this.failureConfigs.Inject429s && ((this.random.Next() % 2) == 0); + public bool ReturnEmptyPage() => (this.failureConfigs != null) && this.failureConfigs.InjectEmptyPages && ((this.random.Next() % 2) == 0); + private static PartitionKeyHash GetHashFromPayload(CosmosObject payload, PartitionKeyDefinition partitionKeyDefinition) { CosmosElement partitionKey = GetPartitionKeyFromPayload(payload, partitionKeyDefinition); @@ -293,12 +305,15 @@ public static Record Create(long previousResourceIdentifier, CosmosObject payloa public sealed class FailureConfigs { - public FailureConfigs(bool inject429s) + public FailureConfigs(bool inject429s, bool injectEmptyPages) { this.Inject429s = inject429s; + this.InjectEmptyPages = injectEmptyPages; } public bool Inject429s { get; } + + public bool InjectEmptyPages { get; } } private sealed class Records : IReadOnlyList diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/InMemoryCollectionTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/InMemoryCollectionTests.cs index c9559feaef..14c62fecdc 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/InMemoryCollectionTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/InMemoryCollectionTests.cs @@ -10,6 +10,7 @@ namespace Microsoft.Azure.Cosmos.Tests using Microsoft.Azure.Cosmos.CosmosElements; using Microsoft.Azure.Cosmos.CosmosElements.Numbers; using Microsoft.Azure.Cosmos.Query.Core.Monads; + using Microsoft.Azure.Cosmos.Serialization.HybridRow.RecordIO; using Microsoft.Azure.Documents; using Microsoft.VisualStudio.TestTools.UnitTesting; @@ -140,14 +141,14 @@ public void Split() int AssertChildPartition(int partitionKeyRangeId) { - TryCatch> tryGetPartitionRecords = inMemoryCollection.ReadFeed( + TryCatch<(List records, long? continuation)> tryGetPartitionRecords = inMemoryCollection.ReadFeed( partitionKeyRangeId: partitionKeyRangeId, resourceIndentifer: 0, pageSize: 100); tryGetPartitionRecords.ThrowIfFailed(); List values = new List(); - foreach (InMemoryCollection.Record record in tryGetPartitionRecords.Result) + foreach (InMemoryCollection.Record record in tryGetPartitionRecords.Result.records) { values.Add(Number64.ToLong((record.Payload["pk"] as CosmosNumber).Value)); } @@ -209,14 +210,14 @@ public void MultiSplit() int AssertChildPartition(int partitionKeyRangeId) { - TryCatch> tryGetPartitionRecords = inMemoryCollection.ReadFeed( + TryCatch<(List records, long? continuation)> tryGetPartitionRecords = inMemoryCollection.ReadFeed( partitionKeyRangeId: partitionKeyRangeId, resourceIndentifer: 0, pageSize: 100); tryGetPartitionRecords.ThrowIfFailed(); List values = new List(); - foreach (InMemoryCollection.Record record in tryGetPartitionRecords.Result) + foreach (InMemoryCollection.Record record in tryGetPartitionRecords.Result.records) { values.Add(Number64.ToLong((record.Payload["pk"] as CosmosNumber).Value)); } diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/InMemoryCollectionPartitionRangeEnumerator.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/InMemoryCollectionPartitionRangeEnumerator.cs index 79a2edf2d2..81716d68b9 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/InMemoryCollectionPartitionRangeEnumerator.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/InMemoryCollectionPartitionRangeEnumerator.cs @@ -11,6 +11,7 @@ namespace Microsoft.Azure.Cosmos.Tests.Pagination using System.Threading.Tasks; using Microsoft.Azure.Cosmos.Pagination; using Microsoft.Azure.Cosmos.Query.Core.Monads; + using Microsoft.Azure.Cosmos.Serialization.HybridRow.RecordIO; internal sealed class InMemoryCollectionPartitionRangeEnumerator : PartitionRangePageEnumerator { @@ -44,23 +45,18 @@ public override Task> GetNextPageAsync(CancellationToken cancella { cancellationToken.ThrowIfCancellationRequested(); - TryCatch> tryReadPage = this.inMemoryCollection.ReadFeed( + TryCatch<(List records, long? continuation)> tryReadPage = this.inMemoryCollection.ReadFeed( partitionKeyRangeId: this.partitionKeyRangeId, resourceIndentifer: ((InMemoryCollectionState)this.State).ResourceIdentifier, pageSize: this.pageSize); + if (tryReadPage.Failed) { return Task.FromResult(TryCatch.FromException(tryReadPage.Exception)); } - if (tryReadPage.Result.Count == 0) - { - InMemoryCollectionPage emptyPage = new InMemoryCollectionPage(new List(), state: default); - return Task.FromResult(TryCatch.FromResult(emptyPage)); - } - - State inMemoryCollectionState = new InMemoryCollectionState(tryReadPage.Result.Last().ResourceIdentifier); - InMemoryCollectionPage page = new InMemoryCollectionPage(tryReadPage.Result, inMemoryCollectionState); + State inMemoryCollectionState = tryReadPage.Result.continuation.HasValue ? new InMemoryCollectionState(tryReadPage.Result.continuation.Value) : default; + InMemoryCollectionPage page = new InMemoryCollectionPage(tryReadPage.Result.records, inMemoryCollectionState); return Task.FromResult(TryCatch.FromResult(page)); } diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/InMemoryCollectionPartitionRangeEnumeratorTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/InMemoryCollectionPartitionRangeEnumeratorTests.cs index 6836b2b1ef..fda7921f78 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/InMemoryCollectionPartitionRangeEnumeratorTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/InMemoryCollectionPartitionRangeEnumeratorTests.cs @@ -38,7 +38,11 @@ public async Task TestResumingFromStateAsync() public async Task Test429sAsync() { int numItems = 100; - InMemoryCollection inMemoryCollection = this.CreateInMemoryCollection(numItems, new InMemoryCollection.FailureConfigs(inject429s: true)); + InMemoryCollection inMemoryCollection = this.CreateInMemoryCollection( + numItems, + new InMemoryCollection.FailureConfigs( + inject429s: true, + injectEmptyPages: false)); IAsyncEnumerable> enumerable = this.CreateEnumerable(inMemoryCollection); @@ -75,7 +79,11 @@ public async Task Test429sAsync() public async Task Test429sWithContinuationsAsync() { int numItems = 100; - InMemoryCollection inMemoryCollection = this.CreateInMemoryCollection(numItems, new InMemoryCollection.FailureConfigs(inject429s: true)); + InMemoryCollection inMemoryCollection = this.CreateInMemoryCollection( + numItems, + new InMemoryCollection.FailureConfigs( + inject429s: true, + injectEmptyPages: false)); IAsyncEnumerator> enumerator = this.CreateEnumerator(inMemoryCollection); @@ -116,6 +124,20 @@ public async Task Test429sWithContinuationsAsync() Assert.AreEqual(numItems, identifiers.Count); } + [TestMethod] + public async Task TestEmptyPages() + { + int numItems = 100; + InMemoryCollection inMemoryCollection = this.CreateInMemoryCollection( + numItems, + new InMemoryCollection.FailureConfigs( + inject429s: false, + injectEmptyPages: true)); + IAsyncEnumerable> enumerable = this.CreateEnumerable(inMemoryCollection); + HashSet identifiers = await this.DrainFullyAsync(enumerable); + Assert.AreEqual(numItems, identifiers.Count); + } + internal abstract List GetRecordsFromPage(Page page); internal abstract InMemoryCollection CreateInMemoryCollection(int numItems, InMemoryCollection.FailureConfigs failureConfigs = default); From b099bd0fa15943d1aee2c6d393828b1efa2f3e0a Mon Sep 17 00:00:00 2001 From: Brandon Chong Date: Tue, 30 Jun 2020 13:11:33 -0700 Subject: [PATCH 19/85] added basic query support --- .../InMemoryCollection.cs | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/InMemoryCollection.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/InMemoryCollection.cs index 023426bc2d..9d1f20a395 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/InMemoryCollection.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/InMemoryCollection.cs @@ -9,6 +9,7 @@ namespace Microsoft.Azure.Cosmos.Tests using System.Collections.Generic; using System.Linq; using Microsoft.Azure.Cosmos.CosmosElements; + using Microsoft.Azure.Cosmos.Query.Core; using Microsoft.Azure.Cosmos.Query.Core.Monads; using Microsoft.Azure.Cosmos.Routing; using Microsoft.Azure.Documents; @@ -132,6 +133,17 @@ record = default; return TryCatch<(List, long?)>.FromResult((page, page.Last().ResourceIdentifier)); } + public TryCatch<(List, long?)> Query(SqlQuerySpec sqlQuerySpec, int partitionKeyRangeId, long resourceIdentifier, int pageSize) + { + if (sqlQuerySpec == null) + { + throw new ArgumentNullException(nameof(sqlQuerySpec)); + } + + // for now just always do a "SELECT * FROM c" query + return this.ReadFeed(partitionKeyRangeId, resourceIdentifier, pageSize); + } + public IReadOnlyDictionary PartitionKeyRangeFeedReed() => this.partitionKeyRangeIdToHashRange; public void Split(int partitionKeyRangeId) From 254f3f8bdda0f5ce95476d155ab2d986d17aa469 Mon Sep 17 00:00:00 2001 From: Brandon Chong Date: Tue, 30 Jun 2020 17:46:42 -0700 Subject: [PATCH 20/85] made code templatized --- .../CreatePartitionRangeEnumerator.cs | 4 +- .../src/Pagination/CrossPartitionPage.cs | 19 + .../CrossPartitionRangePageEnumerable.cs | 21 +- .../CrossPartitionRangePageEnumerator.cs | 90 ++--- .../src/Pagination/CrossPartitionState.cs | 20 ++ Microsoft.Azure.Cosmos/src/Pagination/Page.cs | 7 +- .../PartitionRangePageEnumerable.cs | 14 +- .../PartitionRangePageEnumerator.cs | 12 +- ...CollectionPartitionRangeEnumeratorTests.cs | 333 +++++++++--------- .../Pagination/InMemoryCollectionPage.cs | 20 ++ ...emoryCollectionPartitionRangeEnumerator.cs | 47 +-- ...CollectionPartitionRangeEnumeratorTests.cs | 42 +-- .../Pagination/InMemoryCollectionState.cs | 18 + ...CollectionPartitionRangeEnumeratorTests.cs | 181 ++++++---- 14 files changed, 461 insertions(+), 367 deletions(-) create mode 100644 Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionPage.cs create mode 100644 Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionState.cs create mode 100644 Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/InMemoryCollectionPage.cs create mode 100644 Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/InMemoryCollectionState.cs diff --git a/Microsoft.Azure.Cosmos/src/Pagination/CreatePartitionRangeEnumerator.cs b/Microsoft.Azure.Cosmos/src/Pagination/CreatePartitionRangeEnumerator.cs index 134b69e188..74d79fd5c5 100644 --- a/Microsoft.Azure.Cosmos/src/Pagination/CreatePartitionRangeEnumerator.cs +++ b/Microsoft.Azure.Cosmos/src/Pagination/CreatePartitionRangeEnumerator.cs @@ -4,5 +4,7 @@ namespace Microsoft.Azure.Cosmos.Pagination { - internal delegate PartitionRangePageEnumerator CreatePartitionRangePageEnumerator(FeedRange feedRange, State state); + internal delegate PartitionRangePageEnumerator CreatePartitionRangePageEnumerator(FeedRange feedRange, TState state) + where TPage : Page + where TState : State; } diff --git a/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionPage.cs b/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionPage.cs new file mode 100644 index 0000000000..fab0a7823e --- /dev/null +++ b/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionPage.cs @@ -0,0 +1,19 @@ +// ------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +// ------------------------------------------------------------ + +namespace Microsoft.Azure.Cosmos.Pagination +{ + internal sealed class CrossPartitionPage : Page> + where TBackendPage : Page + where TBackendState : State + { + public CrossPartitionPage(TBackendPage backendEndPage, CrossPartitionState state) + : base(state) + { + this.Page = backendEndPage; + } + + public TBackendPage Page { get; } + } +} diff --git a/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionRangePageEnumerable.cs b/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionRangePageEnumerable.cs index 3cc5678894..01f991397e 100644 --- a/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionRangePageEnumerable.cs +++ b/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionRangePageEnumerable.cs @@ -6,22 +6,23 @@ namespace Microsoft.Azure.Cosmos.Pagination { using System; using System.Collections.Generic; - using System.Linq; using System.Threading; using Microsoft.Azure.Cosmos.Query.Core.Monads; - internal sealed class CrossPartitionRangePageEnumerable : IAsyncEnumerable> + internal sealed class CrossPartitionRangePageEnumerable : IAsyncEnumerable>> + where TPage : Page + where TState : State { - private readonly State state; - private readonly CreatePartitionRangePageEnumerator createPartitionRangeEnumerator; - private readonly IComparer comparer; + private readonly CrossPartitionState state; + private readonly CreatePartitionRangePageEnumerator createPartitionRangeEnumerator; + private readonly IComparer> comparer; private readonly IFeedRangeProvider feedRangeProvider; public CrossPartitionRangePageEnumerable( IFeedRangeProvider feedRangeProvider, - CreatePartitionRangePageEnumerator createPartitionRangeEnumerator, - IComparer comparer, - State state = default) + CreatePartitionRangePageEnumerator createPartitionRangeEnumerator, + IComparer> comparer, + CrossPartitionState state = default) { this.feedRangeProvider = feedRangeProvider ?? throw new ArgumentNullException(nameof(comparer)); this.createPartitionRangeEnumerator = createPartitionRangeEnumerator ?? throw new ArgumentNullException(nameof(createPartitionRangeEnumerator)); @@ -29,11 +30,11 @@ public CrossPartitionRangePageEnumerable( this.state = state; } - public IAsyncEnumerator> GetAsyncEnumerator(CancellationToken cancellationToken = default) + public IAsyncEnumerator>> GetAsyncEnumerator(CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); - return new CrossPartitionRangePageEnumerator( + return new CrossPartitionRangePageEnumerator( this.feedRangeProvider, this.createPartitionRangeEnumerator, this.comparer, diff --git a/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionRangePageEnumerator.cs b/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionRangePageEnumerator.cs index 26c432f49f..eddeef7a35 100644 --- a/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionRangePageEnumerator.cs +++ b/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionRangePageEnumerator.cs @@ -14,20 +14,21 @@ namespace Microsoft.Azure.Cosmos.Pagination using Microsoft.Azure.Cosmos.Query.Core.Monads; /// - /// Coordinates draining pages from multiple , while maintaining a global sort order and handling repartitioning (splits, merge). + /// Coordinates draining pages from multiple , while maintaining a global sort order and handling repartitioning (splits, merge). /// - internal sealed class CrossPartitionRangePageEnumerator : IAsyncEnumerator> + internal sealed class CrossPartitionRangePageEnumerator : IAsyncEnumerator>> + where TPage : Page + where TState : State { private readonly IFeedRangeProvider feedRangeProvider; - private readonly CreatePartitionRangePageEnumerator createPartitionRangeEnumerator; - private readonly AsyncLazy> lazyEnumerators; - private readonly State originalState; + private readonly CreatePartitionRangePageEnumerator createPartitionRangeEnumerator; + private readonly AsyncLazy>> lazyEnumerators; public CrossPartitionRangePageEnumerator( IFeedRangeProvider feedRangeProvider, - CreatePartitionRangePageEnumerator createPartitionRangeEnumerator, - IComparer comparer, - State state = default) + CreatePartitionRangePageEnumerator createPartitionRangeEnumerator, + IComparer> comparer, + CrossPartitionState state = default) { this.feedRangeProvider = feedRangeProvider ?? throw new ArgumentNullException(nameof(feedRangeProvider)); this.createPartitionRangeEnumerator = createPartitionRangeEnumerator ?? throw new ArgumentNullException(nameof(createPartitionRangeEnumerator)); @@ -37,36 +38,31 @@ public CrossPartitionRangePageEnumerator( throw new ArgumentNullException(nameof(comparer)); } - this.originalState = state; - - this.lazyEnumerators = new AsyncLazy>(async (CancellationToken token) => + this.lazyEnumerators = new AsyncLazy>>(async (CancellationToken token) => { - List<(FeedRange, State)> rangeAndStates; - if (state == default) + IReadOnlyList<(FeedRange, TState)> rangeAndStates; + if (state != default) + { + rangeAndStates = state.Value; + } + else { // Fan out to all partitions with default state IEnumerable ranges = await feedRangeProvider.GetFeedRangesAsync(token); - rangeAndStates = new List<(FeedRange, State)>(); + List<(FeedRange, TState)> rangesAndStatesBuilder = new List<(FeedRange, TState)>(); foreach (FeedRange range in ranges) { - rangeAndStates.Add((range, default)); - } - } - else - { - if (!(state is CrossPartitionState crossPartitionState)) - { - throw new ArgumentOutOfRangeException(nameof(state)); + rangesAndStatesBuilder.Add((range, default)); } - rangeAndStates = crossPartitionState.Value; + rangeAndStates = rangesAndStatesBuilder; } - PriorityQueue enumerators = new PriorityQueue(comparer); - foreach ((FeedRange range, State rangeState) in rangeAndStates) + PriorityQueue> enumerators = new PriorityQueue>(comparer); + foreach ((FeedRange range, TState rangeState) in rangeAndStates) { - PartitionRangePageEnumerator enumerator = createPartitionRangeEnumerator(range, rangeState); + PartitionRangePageEnumerator enumerator = createPartitionRangeEnumerator(range, rangeState); enumerators.Enqueue(enumerator); } @@ -74,12 +70,12 @@ public CrossPartitionRangePageEnumerator( }); } - public TryCatch Current { get; private set; } + public TryCatch> Current { get; private set; } public async ValueTask MoveNextAsync() { - PriorityQueue enumerators = await this.lazyEnumerators.GetValueAsync(cancellationToken: default); - PartitionRangePageEnumerator currentPaginator = enumerators.Dequeue(); + PriorityQueue> enumerators = await this.lazyEnumerators.GetValueAsync(cancellationToken: default); + PartitionRangePageEnumerator currentPaginator = enumerators.Dequeue(); bool movedNext = await currentPaginator.MoveNextAsync(); if (!movedNext) { @@ -101,7 +97,7 @@ public async ValueTask MoveNextAsync() IEnumerable childRanges = await this.feedRangeProvider.GetChildRangeAsync(currentPaginator.Range); foreach (FeedRange childRange in childRanges) { - PartitionRangePageEnumerator childPaginator = this.createPartitionRangeEnumerator(childRange, currentPaginator.State); + PartitionRangePageEnumerator childPaginator = this.createPartitionRangeEnumerator(childRange, currentPaginator.State); enumerators.Enqueue(childPaginator); } @@ -117,21 +113,22 @@ public async ValueTask MoveNextAsync() enumerators.Enqueue(currentPaginator); - TryCatch backendPage = currentPaginator.Current; + TryCatch backendPage = currentPaginator.Current; if (backendPage.Failed) { - this.Current = backendPage; + this.Current = TryCatch>.FromException(backendPage.Exception); return true; } - List<(FeedRange, State)> feedRangeAndStates = new List<(FeedRange, State)>(enumerators.Count); - foreach (PartitionRangePageEnumerator enumerator in enumerators) + List<(FeedRange, TState)> feedRangeAndStates = new List<(FeedRange, TState)>(enumerators.Count); + foreach (PartitionRangePageEnumerator enumerator in enumerators) { feedRangeAndStates.Add((enumerator.Range, enumerator.State)); } - CrossPartitionState crossPartitionState = new CrossPartitionState(feedRangeAndStates); - this.Current = TryCatch.FromResult(new CrossPartitionPage(backendPage.Result, crossPartitionState)); + CrossPartitionState crossPartitionState = new CrossPartitionState(feedRangeAndStates); + this.Current = TryCatch>.FromResult( + new CrossPartitionPage(backendPage.Result, crossPartitionState)); return true; } @@ -153,26 +150,5 @@ private static bool IsMergeException(Exception exception) // TODO: code this out return false; } - - public sealed class CrossPartitionPage : Page - { - public CrossPartitionPage(Page backendEndPage, State state) - : base(state) - { - this.Page = backendEndPage; - } - - public Page Page { get; } - } - - private sealed class CrossPartitionState : State - { - public CrossPartitionState(List<(FeedRange, State)> value) - { - this.Value = value; - } - - public List<(FeedRange, State)> Value { get; } - } } } diff --git a/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionState.cs b/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionState.cs new file mode 100644 index 0000000000..338652a341 --- /dev/null +++ b/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionState.cs @@ -0,0 +1,20 @@ +// ------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +// ------------------------------------------------------------ + +namespace Microsoft.Azure.Cosmos.Pagination +{ + using System; + using System.Collections.Generic; + + internal sealed class CrossPartitionState : State + where TState : State + { + public CrossPartitionState(IReadOnlyList<(FeedRange, TState)> value) + { + this.Value = value ?? throw new ArgumentNullException(nameof(value)); + } + + public IReadOnlyList<(FeedRange, TState)> Value { get; } + } +} diff --git a/Microsoft.Azure.Cosmos/src/Pagination/Page.cs b/Microsoft.Azure.Cosmos/src/Pagination/Page.cs index d169af0208..cb2306ac52 100644 --- a/Microsoft.Azure.Cosmos/src/Pagination/Page.cs +++ b/Microsoft.Azure.Cosmos/src/Pagination/Page.cs @@ -4,13 +4,14 @@ namespace Microsoft.Azure.Cosmos.Pagination { - internal abstract class Page + internal abstract class Page + where TState : State { - protected Page(State state) + protected Page(TState state) { this.State = state; } - public State State { get; } + public TState State { get; } } } diff --git a/Microsoft.Azure.Cosmos/src/Pagination/PartitionRangePageEnumerable.cs b/Microsoft.Azure.Cosmos/src/Pagination/PartitionRangePageEnumerable.cs index d1c686323b..f9ca6d40d9 100644 --- a/Microsoft.Azure.Cosmos/src/Pagination/PartitionRangePageEnumerable.cs +++ b/Microsoft.Azure.Cosmos/src/Pagination/PartitionRangePageEnumerable.cs @@ -9,23 +9,25 @@ namespace Microsoft.Azure.Cosmos.Pagination using System.Threading; using Microsoft.Azure.Cosmos.Query.Core.Monads; - internal sealed class PartitionRangePageEnumerable : IAsyncEnumerable> + internal sealed class PartitionRangePageEnumerable : IAsyncEnumerable> + where TPage : Page + where TState : State { private readonly FeedRange range; - private readonly State state; - private readonly CreatePartitionRangePageEnumerator createPartitionRangeEnumerator; + private readonly TState state; + private readonly CreatePartitionRangePageEnumerator createPartitionRangeEnumerator; public PartitionRangePageEnumerable( FeedRange range, - State state, - CreatePartitionRangePageEnumerator createPartitionRangeEnumerator) + TState state, + CreatePartitionRangePageEnumerator createPartitionRangeEnumerator) { this.range = range ?? throw new ArgumentNullException(nameof(range)); this.state = state; this.createPartitionRangeEnumerator = createPartitionRangeEnumerator ?? throw new ArgumentNullException(nameof(createPartitionRangeEnumerator)); } - public IAsyncEnumerator> GetAsyncEnumerator(CancellationToken cancellationToken = default) + public IAsyncEnumerator> GetAsyncEnumerator(CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); diff --git a/Microsoft.Azure.Cosmos/src/Pagination/PartitionRangePageEnumerator.cs b/Microsoft.Azure.Cosmos/src/Pagination/PartitionRangePageEnumerator.cs index 9120085d44..d7e22984d9 100644 --- a/Microsoft.Azure.Cosmos/src/Pagination/PartitionRangePageEnumerator.cs +++ b/Microsoft.Azure.Cosmos/src/Pagination/PartitionRangePageEnumerator.cs @@ -12,11 +12,13 @@ namespace Microsoft.Azure.Cosmos.Pagination /// /// Has the ability to page through a partition range. /// - internal abstract class PartitionRangePageEnumerator : IAsyncEnumerator> + internal abstract class PartitionRangePageEnumerator : IAsyncEnumerator> + where TPage : Page + where TState : State { private bool hasStarted; - protected PartitionRangePageEnumerator(FeedRange range, State state = null) + protected PartitionRangePageEnumerator(FeedRange range, TState state = null) { this.Range = range; this.State = state; @@ -24,9 +26,9 @@ protected PartitionRangePageEnumerator(FeedRange range, State state = null) public FeedRange Range { get; } - public TryCatch Current { get; private set; } + public TryCatch Current { get; private set; } - public State State { get; private set; } + public TState State { get; private set; } public bool HasMoreResults => !this.hasStarted || (this.State != default); @@ -48,7 +50,7 @@ public async ValueTask MoveNextAsync() return true; } - public abstract Task> GetNextPageAsync(CancellationToken cancellationToken = default); + public abstract Task> GetNextPageAsync(CancellationToken cancellationToken = default); public abstract ValueTask DisposeAsync(); } diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/CrossPartitionInMemoryCollectionPartitionRangeEnumeratorTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/CrossPartitionInMemoryCollectionPartitionRangeEnumeratorTests.cs index 66c92af832..d90b900a55 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/CrossPartitionInMemoryCollectionPartitionRangeEnumeratorTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/CrossPartitionInMemoryCollectionPartitionRangeEnumeratorTests.cs @@ -15,205 +15,218 @@ namespace Microsoft.Azure.Cosmos.Tests.Pagination using Microsoft.VisualStudio.TestTools.UnitTesting; [TestClass] - public class CrossPartitionInMemoryCollectionPartitionRangeEnumeratorTests : InMemoryCollectionPartitionRangeEnumeratorTests + public sealed class CrossPartitionInMemoryCollectionPartitionRangeEnumeratorTests { [TestMethod] - public async Task TestSplitWithResumeContinuationAsync() + public async Task Test429sAsync() { - int numItems = 1000; - InMemoryCollection inMemoryCollection = this.CreateInMemoryCollection(numItems); - - IFeedRangeProvider feedRangeProvider = new InMemoryCollectionFeedRangeProvider(inMemoryCollection); - CreatePartitionRangePageEnumerator createEnumerator = (Cosmos.FeedRange range, State state) => new InMemoryCollectionPartitionRangeEnumerator( - inMemoryCollection, - partitionKeyRangeId: int.Parse(((FeedRangePartitionKeyRange)range).PartitionKeyRangeId), - pageSize: 10, - state: state); - - CrossPartitionRangePageEnumerator enumerator = new CrossPartitionRangePageEnumerator( - feedRangeProvider: feedRangeProvider, - createPartitionRangeEnumerator: createEnumerator, - comparer: PartitionRangePageEnumeratorComparer.Singleton); - - (HashSet firstDrainResults, State state) = await this.PartialDrainAsync(enumerator, numIterations: 3); - - int minPartitionKeyRangeId = inMemoryCollection.PartitionKeyRangeFeedReed().Keys.Min(); - int maxPartitionKeyRangeId = inMemoryCollection.PartitionKeyRangeFeedReed().Keys.Max(); - // Split the partition we were reading from - inMemoryCollection.Split(minPartitionKeyRangeId); - - // And a partition we have let to read from - inMemoryCollection.Split((minPartitionKeyRangeId + maxPartitionKeyRangeId) / 2); - - // Resume from state - CrossPartitionRangePageEnumerable enumerable = new CrossPartitionRangePageEnumerable( - feedRangeProvider: feedRangeProvider, - createPartitionRangeEnumerator: createEnumerator, - comparer: PartitionRangePageEnumeratorComparer.Singleton, - state: state); - - HashSet secondDrainResults = await this.DrainFullyAsync(enumerable); - Assert.AreEqual(numItems, firstDrainResults.Count + secondDrainResults.Count); + Implementation implementation = new Implementation(); + await implementation.Test429sAsync(); } [TestMethod] - public async Task TestSplitWithDuringDrainAsync() + public async Task Test429sWithContinuationsAsync() { - int numItems = 1000; - InMemoryCollection inMemoryCollection = this.CreateInMemoryCollection(numItems); - - IFeedRangeProvider feedRangeProvider = new InMemoryCollectionFeedRangeProvider(inMemoryCollection); - CreatePartitionRangePageEnumerator createEnumerator = (Cosmos.FeedRange range, State state) => new InMemoryCollectionPartitionRangeEnumerator( - inMemoryCollection, - partitionKeyRangeId: int.Parse(((FeedRangePartitionKeyRange)range).PartitionKeyRangeId), - pageSize: 10, - state: state); - - CrossPartitionRangePageEnumerable enumerable = new CrossPartitionRangePageEnumerable( - feedRangeProvider: feedRangeProvider, - createPartitionRangeEnumerator: createEnumerator, - comparer: PartitionRangePageEnumeratorComparer.Singleton); - - HashSet identifiers = new HashSet(); - Random random = new Random(); - await foreach (TryCatch tryGetPage in enumerable) - { - if (random.Next() % 2 == 0) - { - List partitionKeyRangeIds = inMemoryCollection.PartitionKeyRangeFeedReed().Keys.ToList(); - int randomIdToSplit = partitionKeyRangeIds[random.Next(0, partitionKeyRangeIds.Count)]; - inMemoryCollection.Split(randomIdToSplit); - } + Implementation implementation = new Implementation(); + await implementation.Test429sWithContinuationsAsync(); + } - tryGetPage.ThrowIfFailed(); + [TestMethod] + public async Task TestDrainFullyAsync() + { + Implementation implementation = new Implementation(); + await implementation.TestDrainFullyAsync(); + } - if (!(tryGetPage.Result is CrossPartitionRangePageEnumerator.CrossPartitionPage crossPartitionPage)) - { - throw new InvalidCastException(); - } + [TestMethod] + public async Task TestEmptyPages() + { + Implementation implementation = new Implementation(); + await implementation.TestEmptyPages(); + } - if (!(crossPartitionPage.Page is InMemoryCollectionPartitionRangeEnumerator.InMemoryCollectionPage page)) - { - throw new InvalidCastException(); - } + [TestMethod] + public async Task TestResumingFromStateAsync() + { + Implementation implementation = new Implementation(); + await implementation.TestResumingFromStateAsync(); + } - foreach (InMemoryCollection.Record record in page.Records) - { - identifiers.Add(record.Identifier); - } - } + [TestMethod] + public async Task TestSplitWithDuringDrainAsync() + { + Implementation implementation = new Implementation(); + await implementation.TestSplitWithDuringDrainAsync(); + } - Assert.AreEqual(numItems, identifiers.Count); + [TestMethod] + public async Task TestSplitWithResumeContinuationAsync() + { + Implementation implementation = new Implementation(); + await implementation.TestSplitWithResumeContinuationAsync(); } - internal override List GetRecordsFromPage(Page page) + private sealed class Implementation : InMemoryCollectionPartitionRangeEnumeratorTests, CrossPartitionState> { - if (!(page is CrossPartitionRangePageEnumerator.CrossPartitionPage crossPartitionPage)) + [TestMethod] + public async Task TestSplitWithResumeContinuationAsync() { - throw new InvalidCastException(); + int numItems = 1000; + InMemoryCollection inMemoryCollection = this.CreateInMemoryCollection(numItems); + IAsyncEnumerator>> enumerator = this.CreateEnumerator(inMemoryCollection); + + (HashSet firstDrainResults, CrossPartitionState state) = await this.PartialDrainAsync(enumerator, numIterations: 3); + + int minPartitionKeyRangeId = inMemoryCollection.PartitionKeyRangeFeedReed().Keys.Min(); + int maxPartitionKeyRangeId = inMemoryCollection.PartitionKeyRangeFeedReed().Keys.Max(); + // Split the partition we were reading from + inMemoryCollection.Split(minPartitionKeyRangeId); + + // And a partition we have let to read from + inMemoryCollection.Split((minPartitionKeyRangeId + maxPartitionKeyRangeId) / 2); + + // Resume from state + IAsyncEnumerable>> enumerable = this.CreateEnumerable(inMemoryCollection, state); + + HashSet secondDrainResults = await this.DrainFullyAsync(enumerable); + Assert.AreEqual(numItems, firstDrainResults.Count + secondDrainResults.Count); } - if (!(crossPartitionPage.Page is InMemoryCollectionPartitionRangeEnumerator.InMemoryCollectionPage inMemoryCollectionPage)) + [TestMethod] + public async Task TestSplitWithDuringDrainAsync() { - throw new InvalidCastException(); - } + int numItems = 1000; + InMemoryCollection inMemoryCollection = this.CreateInMemoryCollection(numItems); + IAsyncEnumerable>> enumerable = this.CreateEnumerable(inMemoryCollection); - return inMemoryCollectionPage.Records; - } + HashSet identifiers = new HashSet(); + Random random = new Random(); + await foreach (TryCatch> tryGetPage in enumerable) + { + if (random.Next() % 2 == 0) + { + List partitionKeyRangeIds = inMemoryCollection.PartitionKeyRangeFeedReed().Keys.ToList(); + int randomIdToSplit = partitionKeyRangeIds[random.Next(0, partitionKeyRangeIds.Count)]; + inMemoryCollection.Split(randomIdToSplit); + } + + tryGetPage.ThrowIfFailed(); + + List records = this.GetRecordsFromPage(tryGetPage.Result); + foreach (InMemoryCollection.Record record in records) + { + identifiers.Add(record.Identifier); + } + } - internal override InMemoryCollection CreateInMemoryCollection(int numItems, InMemoryCollection.FailureConfigs failureConfigs = null) - { - PartitionKeyDefinition partitionKeyDefinition = new PartitionKeyDefinition() + Assert.AreEqual(numItems, identifiers.Count); + } + + internal override InMemoryCollection CreateInMemoryCollection(int numItems, InMemoryCollection.FailureConfigs failureConfigs = null) { - Paths = new System.Collections.ObjectModel.Collection() + PartitionKeyDefinition partitionKeyDefinition = new PartitionKeyDefinition() + { + Paths = new System.Collections.ObjectModel.Collection() { "/pk" }, - Kind = PartitionKind.Hash, - Version = PartitionKeyDefinitionVersion.V2, - }; - - InMemoryCollection inMemoryCollection = new InMemoryCollection(partitionKeyDefinition, failureConfigs); + Kind = PartitionKind.Hash, + Version = PartitionKeyDefinitionVersion.V2, + }; - inMemoryCollection.Split(partitionKeyRangeId: 0); + InMemoryCollection inMemoryCollection = new InMemoryCollection(partitionKeyDefinition, failureConfigs); - inMemoryCollection.Split(partitionKeyRangeId: 1); - inMemoryCollection.Split(partitionKeyRangeId: 2); + inMemoryCollection.Split(partitionKeyRangeId: 0); - inMemoryCollection.Split(partitionKeyRangeId: 3); - inMemoryCollection.Split(partitionKeyRangeId: 4); - inMemoryCollection.Split(partitionKeyRangeId: 5); - inMemoryCollection.Split(partitionKeyRangeId: 6); + inMemoryCollection.Split(partitionKeyRangeId: 1); + inMemoryCollection.Split(partitionKeyRangeId: 2); - for (int i = 0; i < numItems; i++) - { - // Insert an item - CosmosObject item = CosmosObject.Parse($"{{\"pk\" : {i} }}"); - inMemoryCollection.CreateItem(item); - } + inMemoryCollection.Split(partitionKeyRangeId: 3); + inMemoryCollection.Split(partitionKeyRangeId: 4); + inMemoryCollection.Split(partitionKeyRangeId: 5); + inMemoryCollection.Split(partitionKeyRangeId: 6); - return inMemoryCollection; - } + for (int i = 0; i < numItems; i++) + { + // Insert an item + CosmosObject item = CosmosObject.Parse($"{{\"pk\" : {i} }}"); + inMemoryCollection.CreateItem(item); + } - internal override IAsyncEnumerable> CreateEnumerable(InMemoryCollection inMemoryCollection, State state = null) - { - IFeedRangeProvider feedRangeProvider = new InMemoryCollectionFeedRangeProvider(inMemoryCollection); - CreatePartitionRangePageEnumerator createEnumerator = (Cosmos.FeedRange range, State state) => new InMemoryCollectionPartitionRangeEnumerator( - inMemoryCollection, - partitionKeyRangeId: int.Parse(((FeedRangePartitionKeyRange)range).PartitionKeyRangeId), - pageSize: 10, - state: state); - - return new CrossPartitionRangePageEnumerable( - feedRangeProvider: feedRangeProvider, - createPartitionRangeEnumerator: createEnumerator, - comparer: PartitionRangePageEnumeratorComparer.Singleton, - state: state); - } + return inMemoryCollection; + } - internal override IAsyncEnumerator> CreateEnumerator(InMemoryCollection inMemoryCollection, State state = null) - { - IFeedRangeProvider feedRangeProvider = new InMemoryCollectionFeedRangeProvider(inMemoryCollection); - CreatePartitionRangePageEnumerator createEnumerator = (Cosmos.FeedRange range, State state) => new InMemoryCollectionPartitionRangeEnumerator( - inMemoryCollection, - partitionKeyRangeId: int.Parse(((FeedRangePartitionKeyRange)range).PartitionKeyRangeId), - pageSize: 10, - state: state); - - CrossPartitionRangePageEnumerator enumerator = new CrossPartitionRangePageEnumerator( - feedRangeProvider: feedRangeProvider, - createPartitionRangeEnumerator: createEnumerator, - comparer: PartitionRangePageEnumeratorComparer.Singleton, - state: state); - - return enumerator; - } + internal override IAsyncEnumerable>> CreateEnumerable( + InMemoryCollection inMemoryCollection, + CrossPartitionState state = null) + { + IFeedRangeProvider feedRangeProvider = new InMemoryCollectionFeedRangeProvider(inMemoryCollection); + PartitionRangePageEnumerator createEnumerator(Cosmos.FeedRange range, InMemoryCollectionState state) => new InMemoryCollectionPartitionRangeEnumerator( + inMemoryCollection, + partitionKeyRangeId: int.Parse(((FeedRangePartitionKeyRange)range).PartitionKeyRangeId), + pageSize: 10, + state: state); + + return new CrossPartitionRangePageEnumerable( + feedRangeProvider: feedRangeProvider, + createPartitionRangeEnumerator: createEnumerator, + comparer: PartitionRangePageEnumeratorComparer.Singleton, + state: state); + } - private sealed class PartitionRangePageEnumeratorComparer : IComparer - { - public static readonly PartitionRangePageEnumeratorComparer Singleton = new PartitionRangePageEnumeratorComparer(); + internal override IAsyncEnumerator>> CreateEnumerator( + InMemoryCollection inMemoryCollection, + CrossPartitionState state = null) + { + IFeedRangeProvider feedRangeProvider = new InMemoryCollectionFeedRangeProvider(inMemoryCollection); + PartitionRangePageEnumerator createEnumerator(Cosmos.FeedRange range, InMemoryCollectionState state) => new InMemoryCollectionPartitionRangeEnumerator( + inMemoryCollection, + partitionKeyRangeId: int.Parse(((FeedRangePartitionKeyRange)range).PartitionKeyRangeId), + pageSize: 10, + state: state); + + CrossPartitionRangePageEnumerator enumerator = new CrossPartitionRangePageEnumerator( + feedRangeProvider: feedRangeProvider, + createPartitionRangeEnumerator: createEnumerator, + comparer: PartitionRangePageEnumeratorComparer.Singleton, + state: state); + + return enumerator; + } - public int Compare(PartitionRangePageEnumerator partitionRangePageEnumerator1, PartitionRangePageEnumerator partitionRangePageEnumerator2) + internal override List GetRecordsFromPage(CrossPartitionPage page) { - if (object.ReferenceEquals(partitionRangePageEnumerator1, partitionRangePageEnumerator2)) - { - return 0; - } + return page.Page.Records; + } - if (partitionRangePageEnumerator1.HasMoreResults && !partitionRangePageEnumerator2.HasMoreResults) - { - return -1; - } + private sealed class PartitionRangePageEnumeratorComparer : IComparer> + { + public static readonly PartitionRangePageEnumeratorComparer Singleton = new PartitionRangePageEnumeratorComparer(); - if (!partitionRangePageEnumerator1.HasMoreResults && partitionRangePageEnumerator2.HasMoreResults) + public int Compare( + PartitionRangePageEnumerator partitionRangePageEnumerator1, + PartitionRangePageEnumerator partitionRangePageEnumerator2) { - return 1; + if (object.ReferenceEquals(partitionRangePageEnumerator1, partitionRangePageEnumerator2)) + { + return 0; + } + + if (partitionRangePageEnumerator1.HasMoreResults && !partitionRangePageEnumerator2.HasMoreResults) + { + return -1; + } + + if (!partitionRangePageEnumerator1.HasMoreResults && partitionRangePageEnumerator2.HasMoreResults) + { + return 1; + } + + // Either both don't have results or both do. + return string.CompareOrdinal( + ((FeedRangePartitionKeyRange)partitionRangePageEnumerator1.Range).PartitionKeyRangeId, + ((FeedRangePartitionKeyRange)partitionRangePageEnumerator2.Range).PartitionKeyRangeId); } - - // Either both don't have results or both do. - return string.CompareOrdinal( - ((FeedRangePartitionKeyRange)partitionRangePageEnumerator1.Range).PartitionKeyRangeId, - ((FeedRangePartitionKeyRange)partitionRangePageEnumerator2.Range).PartitionKeyRangeId); } } } diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/InMemoryCollectionPage.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/InMemoryCollectionPage.cs new file mode 100644 index 0000000000..aea233388c --- /dev/null +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/InMemoryCollectionPage.cs @@ -0,0 +1,20 @@ +//------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +//------------------------------------------------------------ + +namespace Microsoft.Azure.Cosmos.Tests.Pagination +{ + using System.Collections.Generic; + using Microsoft.Azure.Cosmos.Pagination; + + internal sealed class InMemoryCollectionPage : Page + { + public InMemoryCollectionPage(List records, InMemoryCollectionState state) + : base(state) + { + this.Records = records; + } + + public List Records { get; } + } +} diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/InMemoryCollectionPartitionRangeEnumerator.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/InMemoryCollectionPartitionRangeEnumerator.cs index 81716d68b9..de09150abd 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/InMemoryCollectionPartitionRangeEnumerator.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/InMemoryCollectionPartitionRangeEnumerator.cs @@ -6,33 +6,27 @@ namespace Microsoft.Azure.Cosmos.Tests.Pagination { using System; using System.Collections.Generic; - using System.Linq; using System.Threading; using System.Threading.Tasks; using Microsoft.Azure.Cosmos.Pagination; using Microsoft.Azure.Cosmos.Query.Core.Monads; - using Microsoft.Azure.Cosmos.Serialization.HybridRow.RecordIO; - internal sealed class InMemoryCollectionPartitionRangeEnumerator : PartitionRangePageEnumerator + internal sealed class InMemoryCollectionPartitionRangeEnumerator : PartitionRangePageEnumerator { private readonly InMemoryCollection inMemoryCollection; private readonly int pageSize; private readonly int partitionKeyRangeId; - public InMemoryCollectionPartitionRangeEnumerator(InMemoryCollection inMemoryCollection, int partitionKeyRangeId, int pageSize, State state = null) + public InMemoryCollectionPartitionRangeEnumerator( + InMemoryCollection inMemoryCollection, + int partitionKeyRangeId, + int pageSize, + InMemoryCollectionState state = null) : base(new FeedRangePartitionKeyRange(partitionKeyRangeId.ToString()), state ?? new InMemoryCollectionState(resourceIdentifier: 0)) { this.inMemoryCollection = inMemoryCollection ?? throw new ArgumentNullException(nameof(inMemoryCollection)); this.partitionKeyRangeId = partitionKeyRangeId; this.pageSize = pageSize; - - if (state != null) - { - if (!(state is InMemoryCollectionState _)) - { - throw new ArgumentOutOfRangeException(nameof(state)); - } - } } public override ValueTask DisposeAsync() @@ -41,7 +35,7 @@ public override ValueTask DisposeAsync() return default; } - public override Task> GetNextPageAsync(CancellationToken cancellationToken = default) + public override Task> GetNextPageAsync(CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); @@ -52,34 +46,13 @@ public override Task> GetNextPageAsync(CancellationToken cancella if (tryReadPage.Failed) { - return Task.FromResult(TryCatch.FromException(tryReadPage.Exception)); + return Task.FromResult(TryCatch.FromException(tryReadPage.Exception)); } - State inMemoryCollectionState = tryReadPage.Result.continuation.HasValue ? new InMemoryCollectionState(tryReadPage.Result.continuation.Value) : default; + InMemoryCollectionState inMemoryCollectionState = tryReadPage.Result.continuation.HasValue ? new InMemoryCollectionState(tryReadPage.Result.continuation.Value) : default; InMemoryCollectionPage page = new InMemoryCollectionPage(tryReadPage.Result.records, inMemoryCollectionState); - return Task.FromResult(TryCatch.FromResult(page)); - } - - private sealed class InMemoryCollectionState : State - { - public InMemoryCollectionState(long resourceIdentifier) - { - this.ResourceIdentifier = resourceIdentifier; - } - - public long ResourceIdentifier { get; } - } - - public sealed class InMemoryCollectionPage : Page - { - public InMemoryCollectionPage(List records, State state) - : base(state) - { - this.Records = records; - } - - public List Records { get; } + return Task.FromResult(TryCatch.FromResult(page)); } } } diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/InMemoryCollectionPartitionRangeEnumeratorTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/InMemoryCollectionPartitionRangeEnumeratorTests.cs index fda7921f78..189c8ee6dd 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/InMemoryCollectionPartitionRangeEnumeratorTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/InMemoryCollectionPartitionRangeEnumeratorTests.cs @@ -7,14 +7,16 @@ using Microsoft.Azure.Cosmos.Query.Core.Monads; using Microsoft.VisualStudio.TestTools.UnitTesting; - public abstract class InMemoryCollectionPartitionRangeEnumeratorTests + internal abstract class InMemoryCollectionPartitionRangeEnumeratorTests + where TPage : Page + where TState : State { [TestMethod] public async Task TestDrainFullyAsync() { int numItems = 1000; InMemoryCollection inMemoryCollection = this.CreateInMemoryCollection(numItems); - IAsyncEnumerable> enumerable = this.CreateEnumerable(inMemoryCollection); + IAsyncEnumerable> enumerable = this.CreateEnumerable(inMemoryCollection); HashSet identifiers = await this.DrainFullyAsync(enumerable); Assert.AreEqual(numItems, identifiers.Count); } @@ -25,10 +27,10 @@ public async Task TestResumingFromStateAsync() int numItems = 1000; InMemoryCollection inMemoryCollection = this.CreateInMemoryCollection(numItems); - IAsyncEnumerator> enumerator = this.CreateEnumerator(inMemoryCollection); - (HashSet firstDrainResults, State state) = await this.PartialDrainAsync(enumerator, numIterations: 3); + IAsyncEnumerator> enumerator = this.CreateEnumerator(inMemoryCollection); + (HashSet firstDrainResults, TState state) = await this.PartialDrainAsync(enumerator, numIterations: 3); - IAsyncEnumerable> enumerable = this.CreateEnumerable(inMemoryCollection, state); + IAsyncEnumerable> enumerable = this.CreateEnumerable(inMemoryCollection, state); HashSet secondDrainResults = await this.DrainFullyAsync(enumerable); Assert.AreEqual(numItems, firstDrainResults.Count + secondDrainResults.Count); @@ -44,10 +46,10 @@ public async Task Test429sAsync() inject429s: true, injectEmptyPages: false)); - IAsyncEnumerable> enumerable = this.CreateEnumerable(inMemoryCollection); + IAsyncEnumerable> enumerable = this.CreateEnumerable(inMemoryCollection); HashSet identifiers = new HashSet(); - await foreach (TryCatch tryGetPage in enumerable) + await foreach (TryCatch tryGetPage in enumerable) { if (tryGetPage.Failed) { @@ -85,14 +87,14 @@ public async Task Test429sWithContinuationsAsync() inject429s: true, injectEmptyPages: false)); - IAsyncEnumerator> enumerator = this.CreateEnumerator(inMemoryCollection); + IAsyncEnumerator> enumerator = this.CreateEnumerator(inMemoryCollection); HashSet identifiers = new HashSet(); - State state = default; + TState state = default; while (await enumerator.MoveNextAsync()) { - TryCatch tryGetPage = enumerator.Current; + TryCatch tryGetPage = enumerator.Current; if (tryGetPage.Failed) { Exception exception = tryGetPage.Exception; @@ -133,23 +135,23 @@ public async Task TestEmptyPages() new InMemoryCollection.FailureConfigs( inject429s: false, injectEmptyPages: true)); - IAsyncEnumerable> enumerable = this.CreateEnumerable(inMemoryCollection); + IAsyncEnumerable> enumerable = this.CreateEnumerable(inMemoryCollection); HashSet identifiers = await this.DrainFullyAsync(enumerable); Assert.AreEqual(numItems, identifiers.Count); } - internal abstract List GetRecordsFromPage(Page page); + internal abstract List GetRecordsFromPage(TPage page); internal abstract InMemoryCollection CreateInMemoryCollection(int numItems, InMemoryCollection.FailureConfigs failureConfigs = default); - internal abstract IAsyncEnumerable> CreateEnumerable(InMemoryCollection inMemoryCollection, State state = null); + internal abstract IAsyncEnumerable> CreateEnumerable(InMemoryCollection inMemoryCollection, TState state = null); - internal abstract IAsyncEnumerator> CreateEnumerator(InMemoryCollection inMemoryCollection, State state = null); + internal abstract IAsyncEnumerator> CreateEnumerator(InMemoryCollection inMemoryCollection, TState state = null); - internal async Task> DrainFullyAsync(IAsyncEnumerable> enumerable) + internal async Task> DrainFullyAsync(IAsyncEnumerable> enumerable) { HashSet identifiers = new HashSet(); - await foreach (TryCatch tryGetPage in enumerable) + await foreach (TryCatch tryGetPage in enumerable) { tryGetPage.ThrowIfFailed(); @@ -164,19 +166,19 @@ internal async Task> DrainFullyAsync(IAsyncEnumerable, State)> PartialDrainAsync( - IAsyncEnumerator> enumerator, + internal async Task<(HashSet, TState)> PartialDrainAsync( + IAsyncEnumerator> enumerator, int numIterations) { HashSet identifiers = new HashSet(); - State state = default; + TState state = default; // Drain a couple of iterations for (int i = 0; i < numIterations; i++) { await enumerator.MoveNextAsync(); - TryCatch tryGetPage = enumerator.Current; + TryCatch tryGetPage = enumerator.Current; tryGetPage.ThrowIfFailed(); List records = this.GetRecordsFromPage(tryGetPage.Result); diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/InMemoryCollectionState.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/InMemoryCollectionState.cs new file mode 100644 index 0000000000..45c0d9f402 --- /dev/null +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/InMemoryCollectionState.cs @@ -0,0 +1,18 @@ +//------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +//------------------------------------------------------------ + +namespace Microsoft.Azure.Cosmos.Tests.Pagination +{ + using Microsoft.Azure.Cosmos.Pagination; + + internal sealed class InMemoryCollectionState : State + { + public InMemoryCollectionState(long resourceIdentifier) + { + this.ResourceIdentifier = resourceIdentifier; + } + + public long ResourceIdentifier { get; } + } +} diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/SinglePartitionInMemoryCollectionPartitionRangeEnumeratorTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/SinglePartitionInMemoryCollectionPartitionRangeEnumeratorTests.cs index d456ee8d37..c803c4eff6 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/SinglePartitionInMemoryCollectionPartitionRangeEnumeratorTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/SinglePartitionInMemoryCollectionPartitionRangeEnumeratorTests.cs @@ -14,100 +14,145 @@ namespace Microsoft.Azure.Cosmos.Tests.Pagination using Microsoft.Azure.Cosmos.Query.Core.Monads; [TestClass] - public class SinglePartitionInMemoryCollectionPartitionRangeEnumeratorTests : InMemoryCollectionPartitionRangeEnumeratorTests + public sealed class SinglePartitionInMemoryCollectionPartitionRangeEnumeratorTests { [TestMethod] - public async Task TestSplitAsync() + public async Task Test429sAsync() { - int numItems = 100; - InMemoryCollection inMemoryCollection = this.CreateInMemoryCollection(numItems); - InMemoryCollectionPartitionRangeEnumerator enumerator = new InMemoryCollectionPartitionRangeEnumerator( - inMemoryCollection, - partitionKeyRangeId: 0, - pageSize: 10); - - (HashSet parentIdentifiers, State state) = await this.PartialDrainAsync(enumerator, numIterations: 3); + Implementation implementation = new Implementation(); + await implementation.Test429sAsync(); + } - // Split the partition - inMemoryCollection.Split(partitionKeyRangeId: 0); + [TestMethod] + public async Task Test429sWithContinuationsAsync() + { + Implementation implementation = new Implementation(); + await implementation.Test429sWithContinuationsAsync(); + } - // Try To read from the partition that is gone. - await enumerator.MoveNextAsync(); - Assert.IsTrue(enumerator.Current.Failed); + [TestMethod] + public async Task TestDrainFullyAsync() + { + Implementation implementation = new Implementation(); + await implementation.TestDrainFullyAsync(); + } - // Resume on the children using the parent continuaiton token - HashSet childIdentifiers = new HashSet(); - foreach (int partitionKeyRangeId in new int[] { 1, 2 }) - { - PartitionRangePageEnumerable enumerable1 = new PartitionRangePageEnumerable( - range: new FeedRangePartitionKeyRange(partitionKeyRangeId.ToString()), - state: state, - (range, state) => new InMemoryCollectionPartitionRangeEnumerator( - inMemoryCollection, - partitionKeyRangeId: int.Parse(((FeedRangePartitionKeyRange)range).PartitionKeyRangeId), - pageSize: 10, - state: state)); - HashSet resourceIdentifiers = await this.DrainFullyAsync(enumerable1); + [TestMethod] + public async Task TestEmptyPages() + { + Implementation implementation = new Implementation(); + await implementation.TestEmptyPages(); + } - childIdentifiers.UnionWith(resourceIdentifiers); - } + [TestMethod] + public async Task TestResumingFromStateAsync() + { + Implementation implementation = new Implementation(); + await implementation.TestResumingFromStateAsync(); + } - Assert.AreEqual(numItems, parentIdentifiers.Count + childIdentifiers.Count); + [TestMethod] + public async Task TestSplitAsync() + { + Implementation implementation = new Implementation(); + await implementation.TestSplitAsync(); } - internal override List GetRecordsFromPage(Page page) + [TestClass] + private sealed class Implementation : InMemoryCollectionPartitionRangeEnumeratorTests { - if (!(page is InMemoryCollectionPartitionRangeEnumerator.InMemoryCollectionPage inMemoryCollectionPage)) + [TestMethod] + public async Task TestSplitAsync() { - throw new InvalidCastException(); + int numItems = 100; + InMemoryCollection inMemoryCollection = this.CreateInMemoryCollection(numItems); + InMemoryCollectionPartitionRangeEnumerator enumerator = new InMemoryCollectionPartitionRangeEnumerator( + inMemoryCollection, + partitionKeyRangeId: 0, + pageSize: 10); + + (HashSet parentIdentifiers, InMemoryCollectionState state) = await this.PartialDrainAsync(enumerator, numIterations: 3); + + // Split the partition + inMemoryCollection.Split(partitionKeyRangeId: 0); + + // Try To read from the partition that is gone. + await enumerator.MoveNextAsync(); + Assert.IsTrue(enumerator.Current.Failed); + + // Resume on the children using the parent continuaiton token + HashSet childIdentifiers = new HashSet(); + foreach (int partitionKeyRangeId in new int[] { 1, 2 }) + { + PartitionRangePageEnumerable enumerable = new PartitionRangePageEnumerable( + range: new FeedRangePartitionKeyRange(partitionKeyRangeId.ToString()), + state: state, + (range, state) => new InMemoryCollectionPartitionRangeEnumerator( + inMemoryCollection, + partitionKeyRangeId: int.Parse(((FeedRangePartitionKeyRange)range).PartitionKeyRangeId), + pageSize: 10, + state: state)); + HashSet resourceIdentifiers = await this.DrainFullyAsync(enumerable); + + childIdentifiers.UnionWith(resourceIdentifiers); + } + + Assert.AreEqual(numItems, parentIdentifiers.Count + childIdentifiers.Count); } - return inMemoryCollectionPage.Records; - } + internal override List GetRecordsFromPage(InMemoryCollectionPage page) + { + return page.Records; + } - internal override InMemoryCollection CreateInMemoryCollection(int numItems, InMemoryCollection.FailureConfigs failureConfigs = null) - { - PartitionKeyDefinition partitionKeyDefinition = new PartitionKeyDefinition() + internal override InMemoryCollection CreateInMemoryCollection(int numItems, InMemoryCollection.FailureConfigs failureConfigs = null) { - Paths = new System.Collections.ObjectModel.Collection() + PartitionKeyDefinition partitionKeyDefinition = new PartitionKeyDefinition() + { + Paths = new System.Collections.ObjectModel.Collection() { "/pk" }, - Kind = PartitionKind.Hash, - Version = PartitionKeyDefinitionVersion.V2, - }; + Kind = PartitionKind.Hash, + Version = PartitionKeyDefinitionVersion.V2, + }; - InMemoryCollection inMemoryCollection = new InMemoryCollection(partitionKeyDefinition, failureConfigs); + InMemoryCollection inMemoryCollection = new InMemoryCollection(partitionKeyDefinition, failureConfigs); - for (int i = 0; i < numItems; i++) - { - // Insert an item - CosmosObject item = CosmosObject.Parse($"{{\"pk\" : {i} }}"); - inMemoryCollection.CreateItem(item); + for (int i = 0; i < numItems; i++) + { + // Insert an item + CosmosObject item = CosmosObject.Parse($"{{\"pk\" : {i} }}"); + inMemoryCollection.CreateItem(item); + } + + return inMemoryCollection; } - return inMemoryCollection; - } + internal override IAsyncEnumerable> CreateEnumerable( + InMemoryCollection inMemoryCollection, + InMemoryCollectionState state = null) + { + return new PartitionRangePageEnumerable( + range: new FeedRangePartitionKeyRange("0"), + state: state, + (range, state) => new InMemoryCollectionPartitionRangeEnumerator( + inMemoryCollection, + partitionKeyRangeId: int.Parse(((FeedRangePartitionKeyRange)range).PartitionKeyRangeId), + pageSize: 10, + state: state)); + } - internal override IAsyncEnumerable> CreateEnumerable(InMemoryCollection inMemoryCollection, State state = null) - { - return new PartitionRangePageEnumerable( - range: new FeedRangePartitionKeyRange("0"), - state: state, - (range, state) => new InMemoryCollectionPartitionRangeEnumerator( + internal override IAsyncEnumerator> CreateEnumerator( + InMemoryCollection inMemoryCollection, + InMemoryCollectionState state = null) + { + return new InMemoryCollectionPartitionRangeEnumerator( inMemoryCollection, - partitionKeyRangeId: int.Parse(((FeedRangePartitionKeyRange)range).PartitionKeyRangeId), + partitionKeyRangeId: 0, pageSize: 10, - state: state)); - } - - internal override IAsyncEnumerator> CreateEnumerator(InMemoryCollection inMemoryCollection, State state = null) - { - return new InMemoryCollectionPartitionRangeEnumerator( - inMemoryCollection, - partitionKeyRangeId: 0, - pageSize: 10, - state: state); + state: state); + } } } } From c7909b99d6f80c67f521e9c81c2e38baf35ff4d0 Mon Sep 17 00:00:00 2001 From: Brandon Chong Date: Wed, 1 Jul 2020 10:46:10 -0700 Subject: [PATCH 21/85] drafted all parallel query enumerator --- .../CrossPartitionRangePageEnumerator.cs | 2 +- ...rallelCrossPartitionQueryPageEnumerator.cs | 67 +++++++++++++++++++ 2 files changed, 68 insertions(+), 1 deletion(-) create mode 100644 Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/Parallel/ParallelCrossPartitionQueryPageEnumerator.cs diff --git a/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionRangePageEnumerator.cs b/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionRangePageEnumerator.cs index eddeef7a35..ceae19d3c9 100644 --- a/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionRangePageEnumerator.cs +++ b/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionRangePageEnumerator.cs @@ -16,7 +16,7 @@ namespace Microsoft.Azure.Cosmos.Pagination /// /// Coordinates draining pages from multiple , while maintaining a global sort order and handling repartitioning (splits, merge). /// - internal sealed class CrossPartitionRangePageEnumerator : IAsyncEnumerator>> + internal abstract class CrossPartitionRangePageEnumerator : IAsyncEnumerator>> where TPage : Page where TState : State { diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/Parallel/ParallelCrossPartitionQueryPageEnumerator.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/Parallel/ParallelCrossPartitionQueryPageEnumerator.cs new file mode 100644 index 0000000000..ab1884c52a --- /dev/null +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/Parallel/ParallelCrossPartitionQueryPageEnumerator.cs @@ -0,0 +1,67 @@ +// ------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +// ------------------------------------------------------------ + +namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.Parallel +{ + using System.Collections.Generic; + using Microsoft.Azure.Cosmos.Pagination; + using Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.ItemProducers; + + internal sealed class ParallelCrossPartitionQueryPageEnumerator : CrossPartitionRangePageEnumerator + { + public ParallelCrossPartitionQueryPageEnumerator( + IFeedRangeProvider feedRangeProvider, + IQueryDataSource queryDataSource, + SqlQuerySpec sqlQuerySpec, + int pageSize, + CrossPartitionState state = default) + : base( + feedRangeProvider: feedRangeProvider, + createPartitionRangeEnumerator: ParallelCrossPartitionQueryPageEnumerator.MakeCreateFunction(queryDataSource, sqlQuerySpec, pageSize), + comparer: Comparer.Singleton, + state: state) + { + } + + public static CreatePartitionRangePageEnumerator MakeCreateFunction( + IQueryDataSource queryDataSource, + SqlQuerySpec sqlQuerySpec, + int pageSize) => (FeedRange range, QueryState state) => new QueryPartitionRangePageEnumerator( + queryDataSource, + sqlQuerySpec, + range, + pageSize, + state); + + private sealed class Comparer : IComparer> + { + public static readonly Comparer Singleton = new Comparer(); + + public int Compare( + PartitionRangePageEnumerator partitionRangePageEnumerator1, + PartitionRangePageEnumerator partitionRangePageEnumerator2) + { + if (object.ReferenceEquals(partitionRangePageEnumerator1, partitionRangePageEnumerator2)) + { + return 0; + } + + if (partitionRangePageEnumerator1.HasMoreResults && !partitionRangePageEnumerator2.HasMoreResults) + { + return -1; + } + + if (!partitionRangePageEnumerator1.HasMoreResults && partitionRangePageEnumerator2.HasMoreResults) + { + return 1; + } + + // Either both don't have results or both do. + return string.CompareOrdinal( + ((FeedRangePartitionKeyRange)partitionRangePageEnumerator1.Range).PartitionKeyRangeId, + ((FeedRangePartitionKeyRange)partitionRangePageEnumerator2.Range).PartitionKeyRangeId); + } + } + } +} From 51fb6a80384bb00530934c1e54b7a4029dd60e3f Mon Sep 17 00:00:00 2001 From: Brandon Chong Date: Thu, 2 Jul 2020 11:08:37 -0700 Subject: [PATCH 22/85] got aggregates working --- .../src/ClientRetryPolicy.cs | 3 +- .../CrossPartitionRangePageEnumerator.cs | 2 +- ...ggregateDocumentQueryExecutionComponent.cs | 40 +--- ...yDocumentQueryExecutionComponent.Client.cs | 1 + ...DocumentQueryExecutionComponent.Compute.cs | 1 + .../GroupByDocumentQueryExecutionComponent.cs | 6 +- .../ItemProducers/ItemProducer.cs | 3 +- .../Aggregate/AggregateOperator.cs | 2 +- .../AggregateQueryPipelineStage.Client.cs} | 85 ++++---- .../AggregateQueryPipelineStage.Compute.cs} | 183 ++++++++---------- .../Aggregate/AggregateQueryPipelineStage.cs | 136 +++++++++++++ .../Aggregate/Aggregators/AggregateItem.cs | 2 +- .../Aggregators/AverageAggregator.cs | 4 +- .../Aggregate/Aggregators/CountAggregator.cs | 5 +- .../Aggregate/Aggregators/IAggregator.cs | 4 +- .../Aggregate/Aggregators/MinMaxAggregator.cs | 4 +- .../Aggregators/SingleGroupAggregator.cs | 6 +- .../Aggregate/Aggregators/SumAggregator.cs | 5 +- .../Pipeline/FinishedQueryPipelineStage.cs | 30 +++ .../Core/Pipeline/IQueryPipelineStage.cs | 16 ++ .../Partitions}/BackendQueryDataSource.cs | 2 +- .../Partitions}/IQueryDataSource.cs | 4 +- ...rallelCrossPartitionQueryPageEnumerator.cs | 21 +- .../QueryPartitionRangePageEnumerator.cs | 8 +- .../ItemProducers => Pipeline}/QueryPage.cs | 4 +- .../Core/Pipeline/QueryPipelineStageBase.cs | 50 +++++ .../ItemProducers => Pipeline}/QueryState.cs | 11 +- .../Core/QueryClient/CosmosQueryClient.cs | 1 + .../Core/QueryClient/CosmosQueryContext.cs | 1 + .../src/Query/Core/QueryPlan/QueryInfo.cs | 2 +- .../Query/v3Query/CosmosQueryClientCore.cs | 13 +- .../Query/v3Query/CosmosQueryContextCore.cs | 1 + .../Query/MockCosmosQueryClient.cs | 1 + .../CosmosQueryUnitTests.cs | 1 + .../InMemoryCollectionQueryDataSource.cs | 5 +- .../AggregateQueryPipelineStageTests.cs | 127 ++++++++++++ .../Query/Pipeline/EnumerableStage.cs | 27 +++ .../Query/Pipeline/MockQueryPipelineStage.cs | 40 ++++ .../QueryPartitionRangePageEnumeratorTests.cs | 2 + .../Query/QueryPlanBaselineTests.cs | 5 +- 40 files changed, 630 insertions(+), 234 deletions(-) rename Microsoft.Azure.Cosmos/src/Query/Core/{ExecutionComponent => Pipeline}/Aggregate/AggregateOperator.cs (81%) rename Microsoft.Azure.Cosmos/src/Query/Core/{ExecutionComponent/Aggregate/AggregateDocumentQueryExecutionComponent.Client.cs => Pipeline/Aggregate/AggregateQueryPipelineStage.Client.cs} (50%) rename Microsoft.Azure.Cosmos/src/Query/Core/{ExecutionComponent/Aggregate/AggregateDocumentQueryExecutionComponent.Compute.cs => Pipeline/Aggregate/AggregateQueryPipelineStage.Compute.cs} (52%) create mode 100644 Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Aggregate/AggregateQueryPipelineStage.cs rename Microsoft.Azure.Cosmos/src/Query/Core/{ExecutionComponent => Pipeline}/Aggregate/Aggregators/AggregateItem.cs (94%) rename Microsoft.Azure.Cosmos/src/Query/Core/{ExecutionComponent => Pipeline}/Aggregate/Aggregators/AverageAggregator.cs (98%) rename Microsoft.Azure.Cosmos/src/Query/Core/{ExecutionComponent => Pipeline}/Aggregate/Aggregators/CountAggregator.cs (94%) rename Microsoft.Azure.Cosmos/src/Query/Core/{ExecutionComponent => Pipeline}/Aggregate/Aggregators/IAggregator.cs (87%) rename Microsoft.Azure.Cosmos/src/Query/Core/{ExecutionComponent => Pipeline}/Aggregate/Aggregators/MinMaxAggregator.cs (99%) rename Microsoft.Azure.Cosmos/src/Query/Core/{ExecutionComponent => Pipeline}/Aggregate/Aggregators/SingleGroupAggregator.cs (98%) rename Microsoft.Azure.Cosmos/src/Query/Core/{ExecutionComponent => Pipeline}/Aggregate/Aggregators/SumAggregator.cs (95%) create mode 100644 Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/FinishedQueryPipelineStage.cs create mode 100644 Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/IQueryPipelineStage.cs rename Microsoft.Azure.Cosmos/src/Query/Core/{ExecutionContext/ItemProducers => Pipeline/Partitions}/BackendQueryDataSource.cs (95%) rename Microsoft.Azure.Cosmos/src/Query/Core/{ExecutionContext/ItemProducers => Pipeline/Partitions}/IQueryDataSource.cs (88%) rename Microsoft.Azure.Cosmos/src/Query/Core/{ExecutionContext/Parallel => Pipeline/Partitions}/ParallelCrossPartitionQueryPageEnumerator.cs (74%) rename Microsoft.Azure.Cosmos/src/Query/Core/{ExecutionContext/ItemProducers => Pipeline/Partitions}/QueryPartitionRangePageEnumerator.cs (87%) rename Microsoft.Azure.Cosmos/src/Query/Core/{ExecutionContext/ItemProducers => Pipeline}/QueryPage.cs (95%) create mode 100644 Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/QueryPipelineStageBase.cs rename Microsoft.Azure.Cosmos/src/Query/Core/{ExecutionContext/ItemProducers => Pipeline}/QueryState.cs (53%) create mode 100644 Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/AggregateQueryPipelineStageTests.cs create mode 100644 Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/EnumerableStage.cs create mode 100644 Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/MockQueryPipelineStage.cs diff --git a/Microsoft.Azure.Cosmos/src/ClientRetryPolicy.cs b/Microsoft.Azure.Cosmos/src/ClientRetryPolicy.cs index 328dc2f767..3b7d7a4071 100644 --- a/Microsoft.Azure.Cosmos/src/ClientRetryPolicy.cs +++ b/Microsoft.Azure.Cosmos/src/ClientRetryPolicy.cs @@ -65,8 +65,7 @@ public async Task ShouldRetryAsync( this.retryContext = null; // Received Connection error (HttpRequestException), initiate the endpoint rediscovery - HttpRequestException httpException = exception as HttpRequestException; - if (httpException != null) + if (exception is HttpRequestException httpException) { DefaultTrace.TraceWarning("Endpoint not reachable. Refresh cache and retry"); return await this.ShouldRetryOnEndpointFailureAsync(this.isReadRequest, false); diff --git a/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionRangePageEnumerator.cs b/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionRangePageEnumerator.cs index ceae19d3c9..eddeef7a35 100644 --- a/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionRangePageEnumerator.cs +++ b/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionRangePageEnumerator.cs @@ -16,7 +16,7 @@ namespace Microsoft.Azure.Cosmos.Pagination /// /// Coordinates draining pages from multiple , while maintaining a global sort order and handling repartitioning (splits, merge). /// - internal abstract class CrossPartitionRangePageEnumerator : IAsyncEnumerator>> + internal sealed class CrossPartitionRangePageEnumerator : IAsyncEnumerator>> where TPage : Page where TState : State { diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/AggregateDocumentQueryExecutionComponent.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/AggregateDocumentQueryExecutionComponent.cs index 3d10d24a65..1c00bd5f73 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/AggregateDocumentQueryExecutionComponent.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/AggregateDocumentQueryExecutionComponent.cs @@ -7,10 +7,10 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.Aggregate using System.Collections.Generic; using System.Threading.Tasks; using Microsoft.Azure.Cosmos.CosmosElements; - using Microsoft.Azure.Cosmos.Query.Core.ContinuationTokens; - using Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.Aggregate.Aggregators; using Microsoft.Azure.Cosmos.Query.Core.ExecutionContext; using Microsoft.Azure.Cosmos.Query.Core.Monads; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Aggregate; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Aggregate.Aggregators; /// /// Execution component that is able to aggregate local aggregates from multiple continuations and partitions. @@ -56,7 +56,7 @@ protected AggregateDocumentQueryExecutionComponent( this.isValueAggregateQuery = isValueAggregateQuery; } - public static async Task> TryCreateAsync( + public static Task> TryCreateAsync( ExecutionEnvironment executionEnvironment, IReadOnlyList aggregates, IReadOnlyDictionary aliasToAggregateType, @@ -65,39 +65,7 @@ public static async Task> TryCreateAs CosmosElement continuationToken, Func>> tryCreateSourceAsync) { - if (tryCreateSourceAsync == null) - { - throw new ArgumentNullException(nameof(tryCreateSourceAsync)); - } - - TryCatch tryCreateAggregate; - switch (executionEnvironment) - { - case ExecutionEnvironment.Client: - tryCreateAggregate = await ClientAggregateDocumentQueryExecutionComponent.TryCreateAsync( - aggregates, - aliasToAggregateType, - orderedAliases, - hasSelectValue, - continuationToken, - tryCreateSourceAsync); - break; - - case ExecutionEnvironment.Compute: - tryCreateAggregate = await ComputeAggregateDocumentQueryExecutionComponent.TryCreateAsync( - aggregates, - aliasToAggregateType, - orderedAliases, - hasSelectValue, - continuationToken, - tryCreateSourceAsync); - break; - - default: - throw new ArgumentException($"Unknown {nameof(ExecutionEnvironment)}: {executionEnvironment}."); - } - - return tryCreateAggregate; + return default; } /// diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/GroupBy/GroupByDocumentQueryExecutionComponent.Client.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/GroupBy/GroupByDocumentQueryExecutionComponent.Client.cs index bd4ca6b7da..b83c90c257 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/GroupBy/GroupByDocumentQueryExecutionComponent.Client.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/GroupBy/GroupByDocumentQueryExecutionComponent.Client.cs @@ -10,6 +10,7 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.GroupBy using Microsoft.Azure.Cosmos.CosmosElements; using Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.Aggregate; using Microsoft.Azure.Cosmos.Query.Core.Monads; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Aggregate; using Microsoft.Azure.Cosmos.Query.Core.QueryClient; internal abstract partial class GroupByDocumentQueryExecutionComponent : DocumentQueryExecutionComponentBase diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/GroupBy/GroupByDocumentQueryExecutionComponent.Compute.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/GroupBy/GroupByDocumentQueryExecutionComponent.Compute.cs index e0dd91565d..b7b76112a4 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/GroupBy/GroupByDocumentQueryExecutionComponent.Compute.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/GroupBy/GroupByDocumentQueryExecutionComponent.Compute.cs @@ -15,6 +15,7 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.GroupBy using Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.Aggregate; using Microsoft.Azure.Cosmos.Query.Core.Metrics; using Microsoft.Azure.Cosmos.Query.Core.Monads; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Aggregate; using Microsoft.Azure.Cosmos.Query.Core.QueryClient; internal abstract partial class GroupByDocumentQueryExecutionComponent : DocumentQueryExecutionComponentBase diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/GroupBy/GroupByDocumentQueryExecutionComponent.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/GroupBy/GroupByDocumentQueryExecutionComponent.cs index d3aeb0157d..f526b5e788 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/GroupBy/GroupByDocumentQueryExecutionComponent.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/GroupBy/GroupByDocumentQueryExecutionComponent.cs @@ -9,14 +9,12 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.GroupBy using System.Linq; using System.Threading.Tasks; using Microsoft.Azure.Cosmos.CosmosElements; - using Microsoft.Azure.Cosmos.Json; - using Microsoft.Azure.Cosmos.Query.Core.ContinuationTokens; using Microsoft.Azure.Cosmos.Query.Core.Exceptions; - using Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.Aggregate; - using Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.Aggregate.Aggregators; using Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.Distinct; using Microsoft.Azure.Cosmos.Query.Core.ExecutionContext; using Microsoft.Azure.Cosmos.Query.Core.Monads; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Aggregate; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Aggregate.Aggregators; /// /// Query execution component that groups groupings across continuations and pages. diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/ItemProducers/ItemProducer.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/ItemProducers/ItemProducer.cs index dff5397577..1717278fb6 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/ItemProducers/ItemProducer.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/ItemProducers/ItemProducer.cs @@ -14,6 +14,7 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.ItemProducers using Microsoft.Azure.Cosmos.Query.Core.Collections; using Microsoft.Azure.Cosmos.Query.Core.Metrics; using Microsoft.Azure.Cosmos.Query.Core.Monads; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline; using Microsoft.Azure.Cosmos.Query.Core.QueryClient; using Microsoft.Azure.Cosmos.Resource.CosmosExceptions; using Microsoft.Azure.Documents; @@ -285,7 +286,7 @@ public async Task BufferMoreDocumentsAsync(CancellationToken token) activityId: tryGetQueryPage.Result.ActivityId, responseLengthBytes: tryGetQueryPage.Result.ResponseLengthInBytes, disallowContinuationTokenMessage: default, - continuationToken: tryGetQueryPage.Result.State.ContinuationToken, + continuationToken: ((CosmosString)tryGetQueryPage.Result.State.Value).Value, cosmosQueryExecutionInfo: tryGetQueryPage.Result.CosmosQueryExecutionInfo); } else diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/AggregateOperator.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Aggregate/AggregateOperator.cs similarity index 81% rename from Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/AggregateOperator.cs rename to Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Aggregate/AggregateOperator.cs index 1c3300bdf9..fdd52e516b 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/AggregateOperator.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Aggregate/AggregateOperator.cs @@ -2,7 +2,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. //------------------------------------------------------------ -namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.Aggregate +namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.Aggregate { internal enum AggregateOperator { diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/AggregateDocumentQueryExecutionComponent.Client.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Aggregate/AggregateQueryPipelineStage.Client.cs similarity index 50% rename from Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/AggregateDocumentQueryExecutionComponent.Client.cs rename to Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Aggregate/AggregateQueryPipelineStage.Client.cs index 3a807806b4..ed688d6a3d 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/AggregateDocumentQueryExecutionComponent.Client.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Aggregate/AggregateQueryPipelineStage.Client.cs @@ -1,26 +1,23 @@ // ------------------------------------------------------------ // Copyright (c) Microsoft Corporation. All rights reserved. // ------------------------------------------------------------ -namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.Aggregate + +namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.Aggregate { using System; using System.Collections.Generic; using System.Threading; using System.Threading.Tasks; using Microsoft.Azure.Cosmos.CosmosElements; - using Microsoft.Azure.Cosmos.Diagnostics; - using Microsoft.Azure.Cosmos.Json; - using Microsoft.Azure.Cosmos.Query.Core.ContinuationTokens; - using Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.Aggregate.Aggregators; using Microsoft.Azure.Cosmos.Query.Core.Monads; - using Microsoft.Azure.Cosmos.Query.Core.QueryClient; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Aggregate.Aggregators; - internal abstract partial class AggregateDocumentQueryExecutionComponent : DocumentQueryExecutionComponentBase + internal abstract partial class AggregateQueryPipelineStage : QueryPipelineStageBase { - private sealed class ClientAggregateDocumentQueryExecutionComponent : AggregateDocumentQueryExecutionComponent + private sealed class ClientAggregateQueryPipelineStage : AggregateQueryPipelineStage { - private ClientAggregateDocumentQueryExecutionComponent( - IDocumentQueryExecutionComponent source, + private ClientAggregateQueryPipelineStage( + IQueryPipelineStage source, SingleGroupAggregator singleGroupAggregator, bool isValueAggregateQuery) : base(source, singleGroupAggregator, isValueAggregateQuery) @@ -28,13 +25,13 @@ private ClientAggregateDocumentQueryExecutionComponent( // all the work is done in the base constructor. } - public static async Task> TryCreateAsync( + public static async Task> TryCreateAsync( IReadOnlyList aggregates, IReadOnlyDictionary aliasToAggregateType, IReadOnlyList orderedAliases, bool hasSelectValue, CosmosElement continuationToken, - Func>> tryCreateSourceAsync) + Func>> tryCreateSourceAsync) { if (tryCreateSourceAsync == null) { @@ -47,25 +44,26 @@ public static async Task> TryCreateAs orderedAliases, hasSelectValue, continuationToken: null); + if (tryCreateSingleGroupAggregator.Failed) + { + return TryCatch.FromException(tryCreateSingleGroupAggregator.Exception); + } - if (!tryCreateSingleGroupAggregator.Succeeded) + TryCatch tryCreateSource = await tryCreateSourceAsync(continuationToken); + if (tryCreateSource.Failed) { - return TryCatch.FromException(tryCreateSingleGroupAggregator.Exception); + return tryCreateSource; } - return (await tryCreateSourceAsync(continuationToken)) - .Try((source) => - { - return new ClientAggregateDocumentQueryExecutionComponent( - source, - tryCreateSingleGroupAggregator.Result, - hasSelectValue); - }); + ClientAggregateQueryPipelineStage stage = new ClientAggregateQueryPipelineStage( + tryCreateSource.Result, + tryCreateSingleGroupAggregator.Result, + hasSelectValue); + + return TryCatch.FromResult(stage); } - public override async Task DrainAsync( - int maxElements, - CancellationToken cancellationToken) + protected override async Task> GetNextPageAsync(CancellationToken cancellationToken) { cancellationToken.ThrowIfCancellationRequested(); @@ -75,21 +73,25 @@ public override async Task DrainAsync( double requestCharge = 0; long responseLengthBytes = 0; - while (!this.Source.IsDone) + while (this.inputStage.HasMoreResults) { - QueryResponseCore sourceResponse = await this.Source.DrainAsync(int.MaxValue, cancellationToken); - if (!sourceResponse.IsSuccess) + await this.inputStage.MoveNextAsync(); + TryCatch tryGetPageFromSource = this.inputStage.Current; + + if (tryGetPageFromSource.Failed) { - return sourceResponse; + return tryGetPageFromSource; } - requestCharge += sourceResponse.RequestCharge; - responseLengthBytes += sourceResponse.ResponseLengthBytes; + QueryPage sourcePage = tryGetPageFromSource.Result; + + requestCharge += sourcePage.RequestCharge; + responseLengthBytes += sourcePage.ResponseLengthInBytes; - foreach (CosmosElement element in sourceResponse.CosmosElements) + foreach (CosmosElement element in sourcePage.Documents) { RewrittenAggregateProjections rewrittenAggregateProjections = new RewrittenAggregateProjections( - this.isValueAggregateQuery, + this.isValueQuery, element); this.singleGroupAggregator.AddValues(rewrittenAggregateProjections.Payload); } @@ -102,18 +104,15 @@ public override async Task DrainAsync( finalResult.Add(aggregationResult); } - return QueryResponseCore.CreateSuccess( - result: finalResult, - continuationToken: null, - activityId: null, - disallowContinuationTokenMessage: null, + QueryPage queryPage = new QueryPage( + documents: finalResult, requestCharge: requestCharge, - responseLengthBytes: responseLengthBytes); - } + activityId: default, + responseLengthInBytes: responseLengthBytes, + cosmosQueryExecutionInfo: default, + state: default); - public override CosmosElement GetCosmosElementContinuationToken() - { - throw new NotImplementedException(); + return TryCatch.FromResult(queryPage); } } } diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/AggregateDocumentQueryExecutionComponent.Compute.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Aggregate/AggregateQueryPipelineStage.Compute.cs similarity index 52% rename from Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/AggregateDocumentQueryExecutionComponent.Compute.cs rename to Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Aggregate/AggregateQueryPipelineStage.Compute.cs index 5ba00a7fa1..f1a97e3268 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/AggregateDocumentQueryExecutionComponent.Compute.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Aggregate/AggregateQueryPipelineStage.Compute.cs @@ -1,27 +1,28 @@ // ------------------------------------------------------------ // Copyright (c) Microsoft Corporation. All rights reserved. // ------------------------------------------------------------ -namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.Aggregate + +namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.Aggregate { using System; using System.Collections.Generic; - using System.Runtime.CompilerServices; using System.Threading; using System.Threading.Tasks; using Microsoft.Azure.Cosmos.CosmosElements; using Microsoft.Azure.Cosmos.Query.Core.Exceptions; - using Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.Aggregate.Aggregators; using Microsoft.Azure.Cosmos.Query.Core.Monads; - using Microsoft.Azure.Cosmos.Query.Core.QueryClient; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Aggregate.Aggregators; - internal abstract partial class AggregateDocumentQueryExecutionComponent : DocumentQueryExecutionComponentBase + internal abstract partial class AggregateQueryPipelineStage : QueryPipelineStageBase { - private sealed class ComputeAggregateDocumentQueryExecutionComponent : AggregateDocumentQueryExecutionComponent + private static readonly IReadOnlyList EmptyResults = new List().AsReadOnly(); + + private sealed class ComputeAggregateQueryPipelineStage : AggregateQueryPipelineStage { - private static readonly IReadOnlyList EmptyResults = new List().AsReadOnly(); + private static readonly CosmosString DoneSourceToken = CosmosString.Create("DONE"); - private ComputeAggregateDocumentQueryExecutionComponent( - IDocumentQueryExecutionComponent source, + private ComputeAggregateQueryPipelineStage( + IQueryPipelineStage source, SingleGroupAggregator singleGroupAggregator, bool isValueAggregateQuery) : base(source, singleGroupAggregator, isValueAggregateQuery) @@ -29,21 +30,24 @@ private ComputeAggregateDocumentQueryExecutionComponent( // all the work is done in the base constructor. } - public static async Task> TryCreateAsync( + public static async Task> TryCreateAsync( IReadOnlyList aggregates, IReadOnlyDictionary aliasToAggregateType, IReadOnlyList orderedAliases, bool hasSelectValue, - CosmosElement requestContinuation, - Func>> tryCreateSourceAsync) + CosmosElement continuationToken, + Func>> tryCreateSourceAsync) { AggregateContinuationToken aggregateContinuationToken; - if (requestContinuation != null) + if (continuationToken != null) { - if (!AggregateContinuationToken.TryCreateFromCosmosElement(requestContinuation, out aggregateContinuationToken)) + if (!AggregateContinuationToken.TryCreateFromCosmosElement( + continuationToken, + out aggregateContinuationToken)) { - return TryCatch.FromException( - new MalformedContinuationTokenException($"Malfomed {nameof(AggregateContinuationToken)}: '{requestContinuation}'")); + return TryCatch.FromException( + new MalformedContinuationTokenException( + $"Malfomed {nameof(AggregateContinuationToken)}: '{continuationToken}'")); } } else @@ -57,117 +61,96 @@ public static async Task> TryCreateAs orderedAliases, hasSelectValue, aggregateContinuationToken.SingleGroupAggregatorContinuationToken); + if (tryCreateSingleGroupAggregator.Failed) + { + return TryCatch.FromException(tryCreateSingleGroupAggregator.Exception); + } - if (!tryCreateSingleGroupAggregator.Succeeded) + TryCatch tryCreateSource; + if (aggregateContinuationToken.SourceContinuationToken is CosmosString stringToken && (stringToken.Value == DoneSourceToken.Value)) { - return TryCatch.FromException( - tryCreateSingleGroupAggregator.Exception); + tryCreateSource = TryCatch.FromResult(FinishedQueryPipelineStage.Value); + } + else + { + tryCreateSource = await tryCreateSourceAsync(aggregateContinuationToken.SourceContinuationToken); } - return (await tryCreateSourceAsync(aggregateContinuationToken.SourceContinuationToken)) - .Try((source) => - { - return new ComputeAggregateDocumentQueryExecutionComponent( - source, - tryCreateSingleGroupAggregator.Result, - hasSelectValue); - }); + if (tryCreateSource.Failed) + { + return tryCreateSource; + } + + ComputeAggregateQueryPipelineStage stage = new ComputeAggregateQueryPipelineStage( + tryCreateSource.Result, + tryCreateSingleGroupAggregator.Result, + hasSelectValue); + + return TryCatch.FromResult(stage); } - public override async Task DrainAsync( - int maxElements, - CancellationToken cancellationToken) + protected override async Task> GetNextPageAsync(CancellationToken cancellationToken) { cancellationToken.ThrowIfCancellationRequested(); // Draining aggregates is broken down into two stages - QueryResponseCore response; - if (!this.Source.IsDone) + if (!this.inputStage.HasMoreResults) { - // Stage 1: - // Drain the aggregates fully from all continuations and all partitions - // And return empty pages in the meantime. - QueryResponseCore sourceResponse = await this.Source.DrainAsync(int.MaxValue, cancellationToken); - if (!sourceResponse.IsSuccess) + // Stage 2: + // Return the final page after draining. + List finalResult = new List(); + CosmosElement aggregationResult = this.singleGroupAggregator.GetResult(); + if (aggregationResult != null) { - return sourceResponse; + finalResult.Add(aggregationResult); } - foreach (CosmosElement element in sourceResponse.CosmosElements) - { - RewrittenAggregateProjections rewrittenAggregateProjections = new RewrittenAggregateProjections( - this.isValueAggregateQuery, - element); - this.singleGroupAggregator.AddValues(rewrittenAggregateProjections.Payload); - } + QueryPage finalPage = new QueryPage( + documents: finalResult, + requestCharge: default, + activityId: default, + responseLengthInBytes: default, + cosmosQueryExecutionInfo: default, + state: default); - if (this.Source.IsDone) - { - response = this.GetFinalResponse(); - } - else - { - response = this.GetEmptyPage(sourceResponse); - } + return TryCatch.FromResult(finalPage); } - else - { - // Stage 2: - // Return the final page after draining. - response = this.GetFinalResponse(); - } - - return response; - } - private QueryResponseCore GetFinalResponse() - { - List finalResult = new List(); - CosmosElement aggregationResult = this.singleGroupAggregator.GetResult(); - if (aggregationResult != null) + // Stage 1: + // Drain the aggregates fully from all continuations and all partitions + // And return empty pages in the meantime. + await this.inputStage.MoveNextAsync(); + TryCatch tryGetSourcePage = this.inputStage.Current; + if (tryGetSourcePage.Failed) { - finalResult.Add(aggregationResult); + return tryGetSourcePage; } - QueryResponseCore response = QueryResponseCore.CreateSuccess( - result: finalResult, - requestCharge: 0, - activityId: null, - responseLengthBytes: 0, - disallowContinuationTokenMessage: null, - continuationToken: null); - - return response; - } - - private QueryResponseCore GetEmptyPage(QueryResponseCore sourceResponse) - { - // We need to give empty pages until the results are fully drained. - QueryResponseCore response = QueryResponseCore.CreateSuccess( - result: EmptyResults, - requestCharge: sourceResponse.RequestCharge, - activityId: sourceResponse.ActivityId, - responseLengthBytes: sourceResponse.ResponseLengthBytes, - disallowContinuationTokenMessage: null, - continuationToken: sourceResponse.ContinuationToken); - - return response; - } - - public override CosmosElement GetCosmosElementContinuationToken() - { - if (this.IsDone) + QueryPage sourcePage = tryGetSourcePage.Result; + foreach (CosmosElement element in sourcePage.Documents) { - return default; + RewrittenAggregateProjections rewrittenAggregateProjections = new RewrittenAggregateProjections( + this.isValueQuery, + element); + this.singleGroupAggregator.AddValues(rewrittenAggregateProjections.Payload); } AggregateContinuationToken aggregateContinuationToken = new AggregateContinuationToken( singleGroupAggregatorContinuationToken: this.singleGroupAggregator.GetCosmosElementContinuationToken(), - sourceContinuationToken: this.Source.GetCosmosElementContinuationToken()); - return AggregateContinuationToken.ToCosmosElement(aggregateContinuationToken); + sourceContinuationToken: sourcePage.State != null ? sourcePage.State.Value : DoneSourceToken); + QueryState queryState = new QueryState(AggregateContinuationToken.ToCosmosElement(aggregateContinuationToken)); + QueryPage emptyPage = new QueryPage( + documents: EmptyResults, + requestCharge: sourcePage.RequestCharge, + activityId: sourcePage.ActivityId, + responseLengthInBytes: sourcePage.ResponseLengthInBytes, + cosmosQueryExecutionInfo: sourcePage.CosmosQueryExecutionInfo, + state: queryState); + + return TryCatch.FromResult(emptyPage); } - private readonly struct AggregateContinuationToken + private sealed class AggregateContinuationToken { private const string SourceTokenName = "SourceToken"; private const string AggregationTokenName = "AggregationToken"; diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Aggregate/AggregateQueryPipelineStage.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Aggregate/AggregateQueryPipelineStage.cs new file mode 100644 index 0000000000..273d44bcb0 --- /dev/null +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Aggregate/AggregateQueryPipelineStage.cs @@ -0,0 +1,136 @@ +// ------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +// ------------------------------------------------------------ + +namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.Aggregate +{ + using System; + using System.Collections.Generic; + using System.Threading.Tasks; + using Microsoft.Azure.Cosmos.CosmosElements; + using Microsoft.Azure.Cosmos.Query.Core.ExecutionContext; + using Microsoft.Azure.Cosmos.Query.Core.Monads; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Aggregate.Aggregators; + + /// + /// Stage that is able to aggregate local aggregates from multiple continuations and partitions. + /// At a high level aggregates queries only return a "partial" aggregate. + /// "partial" means that the result is only valid for that one continuation (and one partition). + /// For example suppose you have the query "SELECT COUNT(1) FROM c" and you have a single partition collection, + /// then you will get one count for each continuation of the query. + /// If you wanted the true result for this query, then you will have to take the sum of all continuations. + /// The reason why we have multiple continuations is because for a long running query we have to break up the results into multiple continuations. + /// Fortunately all the aggregates can be aggregated across continuations and partitions. + /// + internal abstract partial class AggregateQueryPipelineStage : QueryPipelineStageBase + { + /// + /// This class does most of the work, since a query like: + /// + /// SELECT VALUE AVG(c.age) + /// FROM c + /// + /// is really just an aggregation on a single grouping (the whole collection). + /// + private readonly SingleGroupAggregator singleGroupAggregator; + + /// + /// We need to keep track of whether the projection has the 'VALUE' keyword. + /// + private readonly bool isValueQuery; + + /// + /// Initializes a new instance of the AggregateDocumentQueryExecutionComponent class. + /// + /// The source component that will supply the local aggregates from multiple continuations and partitions. + /// The single group aggregator that we will feed results into. + /// Whether or not the query has the 'VALUE' keyword. + /// This constructor is private since there is some async initialization that needs to happen in CreateAsync(). + public AggregateQueryPipelineStage( + IQueryPipelineStage source, + SingleGroupAggregator singleGroupAggregator, + bool isValueQuery) + : base(source) + { + this.singleGroupAggregator = singleGroupAggregator ?? throw new ArgumentNullException(nameof(singleGroupAggregator)); + this.isValueQuery = isValueQuery; + } + + public static async Task> TryCreateAsync( + ExecutionEnvironment executionEnvironment, + IReadOnlyList aggregates, + IReadOnlyDictionary aliasToAggregateType, + IReadOnlyList orderedAliases, + bool hasSelectValue, + CosmosElement continuationToken, + Func>> tryCreateSourceAsync) + { + if (tryCreateSourceAsync == null) + { + throw new ArgumentNullException(nameof(tryCreateSourceAsync)); + } + + TryCatch tryCreateAggregate = executionEnvironment switch + { + ExecutionEnvironment.Client => await ClientAggregateQueryPipelineStage.TryCreateAsync( + aggregates, + aliasToAggregateType, + orderedAliases, + hasSelectValue, + continuationToken, + tryCreateSourceAsync), + ExecutionEnvironment.Compute => await ComputeAggregateQueryPipelineStage.TryCreateAsync( + aggregates, + aliasToAggregateType, + orderedAliases, + hasSelectValue, + continuationToken, + tryCreateSourceAsync), + _ => throw new ArgumentException($"Unknown {nameof(ExecutionEnvironment)}: {executionEnvironment}."), + }; + + return tryCreateAggregate; + } + + /// + /// Struct for getting the payload out of the rewritten projection. + /// + private readonly struct RewrittenAggregateProjections + { + public RewrittenAggregateProjections(bool isValueAggregateQuery, CosmosElement raw) + { + if (raw == null) + { + throw new ArgumentNullException(nameof(raw)); + } + + if (isValueAggregateQuery) + { + // SELECT VALUE [{"item": {"sum": SUM(c.blah), "count": COUNT(c.blah)}}] + if (!(raw is CosmosArray aggregates)) + { + throw new ArgumentException($"{nameof(RewrittenAggregateProjections)} was not an array for a value aggregate query. Type is: {raw.Type}"); + } + + this.Payload = aggregates[0]; + } + else + { + if (!(raw is CosmosObject cosmosObject)) + { + throw new ArgumentException($"{nameof(raw)} must not be an object."); + } + + if (!cosmosObject.TryGetValue("payload", out CosmosElement cosmosPayload)) + { + throw new InvalidOperationException($"Underlying object does not have an 'payload' field."); + } + + this.Payload = cosmosPayload ?? throw new ArgumentException($"{nameof(RewrittenAggregateProjections)} does not have a 'payload' property."); + } + } + + public CosmosElement Payload { get; } + } + } +} diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/AggregateItem.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Aggregate/Aggregators/AggregateItem.cs similarity index 94% rename from Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/AggregateItem.cs rename to Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Aggregate/Aggregators/AggregateItem.cs index 47683a57c7..4f476b74dd 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/AggregateItem.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Aggregate/Aggregators/AggregateItem.cs @@ -2,7 +2,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. //------------------------------------------------------------ -namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.Aggregate.Aggregators +namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.Aggregate.Aggregators { using System; using Microsoft.Azure.Cosmos.CosmosElements; diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/AverageAggregator.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Aggregate/Aggregators/AverageAggregator.cs similarity index 98% rename from Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/AverageAggregator.cs rename to Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Aggregate/Aggregators/AverageAggregator.cs index 1501e35d6d..7ad3ea040f 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/AverageAggregator.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Aggregate/Aggregators/AverageAggregator.cs @@ -1,13 +1,13 @@ //------------------------------------------------------------ // Copyright (c) Microsoft Corporation. All rights reserved. //------------------------------------------------------------ -namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.Aggregate.Aggregators + +namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.Aggregate.Aggregators { using System; using System.Collections.Generic; using Microsoft.Azure.Cosmos.CosmosElements; using Microsoft.Azure.Cosmos.CosmosElements.Numbers; - using Microsoft.Azure.Cosmos.Json; using Microsoft.Azure.Cosmos.Query.Core.Exceptions; using Microsoft.Azure.Cosmos.Query.Core.Monads; diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/CountAggregator.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Aggregate/Aggregators/CountAggregator.cs similarity index 94% rename from Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/CountAggregator.cs rename to Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Aggregate/Aggregators/CountAggregator.cs index 84851aaf66..696e44a1b5 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/CountAggregator.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Aggregate/Aggregators/CountAggregator.cs @@ -1,14 +1,13 @@ //------------------------------------------------------------ // Copyright (c) Microsoft Corporation. All rights reserved. //------------------------------------------------------------ -namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.Aggregate.Aggregators + +namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.Aggregate.Aggregators { using System; using System.Globalization; using Microsoft.Azure.Cosmos.CosmosElements; using Microsoft.Azure.Cosmos.CosmosElements.Numbers; - using Microsoft.Azure.Cosmos.Json; - using Microsoft.Azure.Cosmos.Query.Core.ContinuationTokens; using Microsoft.Azure.Cosmos.Query.Core.Exceptions; using Microsoft.Azure.Cosmos.Query.Core.Monads; diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/IAggregator.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Aggregate/Aggregators/IAggregator.cs similarity index 87% rename from Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/IAggregator.cs rename to Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Aggregate/Aggregators/IAggregator.cs index dace3bf6a7..c264163845 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/IAggregator.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Aggregate/Aggregators/IAggregator.cs @@ -1,10 +1,10 @@ //------------------------------------------------------------ // Copyright (c) Microsoft Corporation. All rights reserved. //------------------------------------------------------------ -namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.Aggregate.Aggregators + +namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.Aggregate.Aggregators { using Microsoft.Azure.Cosmos.CosmosElements; - using Microsoft.Azure.Cosmos.Json; /// /// Interface for all aggregators that are used to aggregate across continuation and partition boundaries. diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/MinMaxAggregator.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Aggregate/Aggregators/MinMaxAggregator.cs similarity index 99% rename from Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/MinMaxAggregator.cs rename to Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Aggregate/Aggregators/MinMaxAggregator.cs index 5f77dbba77..25cce653dd 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/MinMaxAggregator.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Aggregate/Aggregators/MinMaxAggregator.cs @@ -1,12 +1,12 @@ //------------------------------------------------------------ // Copyright (c) Microsoft Corporation. All rights reserved. //------------------------------------------------------------ -namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.Aggregate.Aggregators + +namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.Aggregate.Aggregators { using System; using System.Collections.Generic; using Microsoft.Azure.Cosmos.CosmosElements; - using Microsoft.Azure.Cosmos.Json; using Microsoft.Azure.Cosmos.Query.Core.Exceptions; using Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.ItemProducers; using Microsoft.Azure.Cosmos.Query.Core.Monads; diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/SingleGroupAggregator.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Aggregate/Aggregators/SingleGroupAggregator.cs similarity index 98% rename from Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/SingleGroupAggregator.cs rename to Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Aggregate/Aggregators/SingleGroupAggregator.cs index 95a567226d..fe62c167c4 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/SingleGroupAggregator.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Aggregate/Aggregators/SingleGroupAggregator.cs @@ -1,16 +1,16 @@ //------------------------------------------------------------ // Copyright (c) Microsoft Corporation. All rights reserved. //------------------------------------------------------------ -namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.Aggregate.Aggregators + +namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.Aggregate.Aggregators { using System; using System.Collections.Generic; using System.Linq; using Microsoft.Azure.Cosmos.CosmosElements; - using Microsoft.Azure.Cosmos.Json; - using Microsoft.Azure.Cosmos.Query.Core.ContinuationTokens; using Microsoft.Azure.Cosmos.Query.Core.Exceptions; using Microsoft.Azure.Cosmos.Query.Core.Monads; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Aggregate; /// /// Aggregates all the projections for a single grouping. diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/SumAggregator.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Aggregate/Aggregators/SumAggregator.cs similarity index 95% rename from Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/SumAggregator.cs rename to Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Aggregate/Aggregators/SumAggregator.cs index 433a3535c6..2536cf25b9 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/SumAggregator.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Aggregate/Aggregators/SumAggregator.cs @@ -1,14 +1,13 @@ //------------------------------------------------------------ // Copyright (c) Microsoft Corporation. All rights reserved. //------------------------------------------------------------ -namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.Aggregate.Aggregators + +namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.Aggregate.Aggregators { using System; using System.Globalization; using Microsoft.Azure.Cosmos.CosmosElements; using Microsoft.Azure.Cosmos.CosmosElements.Numbers; - using Microsoft.Azure.Cosmos.Json; - using Microsoft.Azure.Cosmos.Query.Core.ContinuationTokens; using Microsoft.Azure.Cosmos.Query.Core.Exceptions; using Microsoft.Azure.Cosmos.Query.Core.Monads; diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/FinishedQueryPipelineStage.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/FinishedQueryPipelineStage.cs new file mode 100644 index 0000000000..1728479894 --- /dev/null +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/FinishedQueryPipelineStage.cs @@ -0,0 +1,30 @@ +// ------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +// ------------------------------------------------------------ + +namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline +{ + using System; + using System.Threading; + using System.Threading.Tasks; + using Microsoft.Azure.Cosmos.Query.Core.Monads; + + internal sealed class FinishedQueryPipelineStage : IQueryPipelineStage + { + public static readonly FinishedQueryPipelineStage Value = new FinishedQueryPipelineStage(); + + private FinishedQueryPipelineStage() + { + } + + public bool HasMoreResults => false; + + public TryCatch Current => default; + + public ValueTask DisposeAsync() => default; + + public Task> GetNextPageAsync(CancellationToken cancellationToken) => throw new NotSupportedException(); + + public ValueTask MoveNextAsync() => new ValueTask(false); + } +} diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/IQueryPipelineStage.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/IQueryPipelineStage.cs new file mode 100644 index 0000000000..f5ad3230e0 --- /dev/null +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/IQueryPipelineStage.cs @@ -0,0 +1,16 @@ +// ------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +// ------------------------------------------------------------ + +namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline +{ + using System.Collections.Generic; + using System.Threading; + using System.Threading.Tasks; + using Microsoft.Azure.Cosmos.Query.Core.Monads; + + internal interface IQueryPipelineStage : IAsyncEnumerator> + { + bool HasMoreResults { get; } + } +} diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/ItemProducers/BackendQueryDataSource.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Partitions/BackendQueryDataSource.cs similarity index 95% rename from Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/ItemProducers/BackendQueryDataSource.cs rename to Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Partitions/BackendQueryDataSource.cs index b426657ac8..ced7eb235d 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/ItemProducers/BackendQueryDataSource.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Partitions/BackendQueryDataSource.cs @@ -2,7 +2,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // ------------------------------------------------------------ -namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.ItemProducers +namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.Partitions { using System; using System.Threading; diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/ItemProducers/IQueryDataSource.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Partitions/IQueryDataSource.cs similarity index 88% rename from Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/ItemProducers/IQueryDataSource.cs rename to Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Partitions/IQueryDataSource.cs index b5835a879e..5ac3ee6380 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/ItemProducers/IQueryDataSource.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Partitions/IQueryDataSource.cs @@ -2,7 +2,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // ------------------------------------------------------------ -namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.ItemProducers +namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.Partitions { using System.Threading; using System.Threading.Tasks; @@ -17,4 +17,4 @@ public Task> ExecuteQueryAsync( int pageSize, CancellationToken cancellationToken); } -} +} \ No newline at end of file diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/Parallel/ParallelCrossPartitionQueryPageEnumerator.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Partitions/ParallelCrossPartitionQueryPageEnumerator.cs similarity index 74% rename from Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/Parallel/ParallelCrossPartitionQueryPageEnumerator.cs rename to Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Partitions/ParallelCrossPartitionQueryPageEnumerator.cs index ab1884c52a..37f9ad2c89 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/Parallel/ParallelCrossPartitionQueryPageEnumerator.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Partitions/ParallelCrossPartitionQueryPageEnumerator.cs @@ -2,29 +2,28 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // ------------------------------------------------------------ -namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.Parallel +namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.Partitions { using System.Collections.Generic; using Microsoft.Azure.Cosmos.Pagination; - using Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.ItemProducers; - internal sealed class ParallelCrossPartitionQueryPageEnumerator : CrossPartitionRangePageEnumerator + internal static class ParallelCrossPartitionQueryPageEnumerator { - public ParallelCrossPartitionQueryPageEnumerator( + public static CrossPartitionRangePageEnumerator Create( IFeedRangeProvider feedRangeProvider, IQueryDataSource queryDataSource, SqlQuerySpec sqlQuerySpec, int pageSize, CrossPartitionState state = default) - : base( - feedRangeProvider: feedRangeProvider, - createPartitionRangeEnumerator: ParallelCrossPartitionQueryPageEnumerator.MakeCreateFunction(queryDataSource, sqlQuerySpec, pageSize), - comparer: Comparer.Singleton, - state: state) { + return new CrossPartitionRangePageEnumerator( + feedRangeProvider, + ParallelCrossPartitionQueryPageEnumerator.MakeCreateFunction(queryDataSource, sqlQuerySpec, pageSize), + Comparer.Singleton, + state: state); } - public static CreatePartitionRangePageEnumerator MakeCreateFunction( + private static CreatePartitionRangePageEnumerator MakeCreateFunction( IQueryDataSource queryDataSource, SqlQuerySpec sqlQuerySpec, int pageSize) => (FeedRange range, QueryState state) => new QueryPartitionRangePageEnumerator( @@ -64,4 +63,4 @@ public int Compare( } } } -} +} \ No newline at end of file diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/ItemProducers/QueryPartitionRangePageEnumerator.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Partitions/QueryPartitionRangePageEnumerator.cs similarity index 87% rename from Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/ItemProducers/QueryPartitionRangePageEnumerator.cs rename to Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Partitions/QueryPartitionRangePageEnumerator.cs index c519ed52ab..dbbc5db103 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/ItemProducers/QueryPartitionRangePageEnumerator.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Partitions/QueryPartitionRangePageEnumerator.cs @@ -2,13 +2,15 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // ------------------------------------------------------------ -namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.ItemProducers +namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.Partitions { using System; using System.Threading; using System.Threading.Tasks; + using Microsoft.Azure.Cosmos.CosmosElements; using Microsoft.Azure.Cosmos.Pagination; using Microsoft.Azure.Cosmos.Query.Core.Monads; + using Microsoft.Azure.Cosmos.SqlObjects; internal sealed class QueryPartitionRangePageEnumerator : PartitionRangePageEnumerator { @@ -36,11 +38,11 @@ public QueryPartitionRangePageEnumerator( public override Task> GetNextPageAsync(CancellationToken cancellationToken) => this.queryDataSource.ExecuteQueryAsync( sqlQuerySpec: this.sqlQuerySpec, - continuationToken: this.State?.ContinuationToken, + continuationToken: this.State.Value != null ? ((CosmosString)this.State.Value).Value : null, partitionKeyRangeId: int.Parse(((FeedRangePartitionKeyRange)this.Range).PartitionKeyRangeId), pageSize: this.pageSize, cancellationToken); public override ValueTask DisposeAsync() => default; } -} +} \ No newline at end of file diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/ItemProducers/QueryPage.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/QueryPage.cs similarity index 95% rename from Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/ItemProducers/QueryPage.cs rename to Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/QueryPage.cs index 881dc1f65b..6410266f09 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/ItemProducers/QueryPage.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/QueryPage.cs @@ -2,7 +2,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // ------------------------------------------------------------ -namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.ItemProducers +namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline { using System; using System.Collections.Generic; @@ -38,4 +38,4 @@ public QueryPage( public CosmosQueryExecutionInfo CosmosQueryExecutionInfo { get; } } -} +} \ No newline at end of file diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/QueryPipelineStageBase.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/QueryPipelineStageBase.cs new file mode 100644 index 0000000000..1864ee7e99 --- /dev/null +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/QueryPipelineStageBase.cs @@ -0,0 +1,50 @@ +// ------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +// ------------------------------------------------------------ + +namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline +{ + using System; + using System.Threading; + using System.Threading.Tasks; + using Microsoft.Azure.Cosmos.Query.Core.Monads; + + internal abstract class QueryPipelineStageBase : IQueryPipelineStage + { + protected readonly IQueryPipelineStage inputStage; + private bool hasStarted; + + protected QueryPipelineStageBase(IQueryPipelineStage inputStage) + { + this.inputStage = inputStage ?? throw new ArgumentNullException(nameof(inputStage)); + } + + public TryCatch Current { get; private set; } + + public QueryState State { get; protected set; } + + public bool HasMoreResults => !this.hasStarted || (this.State != default); + + public ValueTask DisposeAsync() => this.inputStage.DisposeAsync(); + + protected abstract Task> GetNextPageAsync(CancellationToken cancellationToken); + + public async ValueTask MoveNextAsync() + { + if (!this.HasMoreResults) + { + return false; + } + + this.hasStarted = true; + + this.Current = await this.GetNextPageAsync(default); + if (this.Current.Succeeded) + { + this.State = this.Current.Result.State; + } + + return true; + } + } +} diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/ItemProducers/QueryState.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/QueryState.cs similarity index 53% rename from Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/ItemProducers/QueryState.cs rename to Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/QueryState.cs index 3e84f4b3ea..b822fc8ede 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/ItemProducers/QueryState.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/QueryState.cs @@ -2,18 +2,19 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // ------------------------------------------------------------ -namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.ItemProducers +namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline { using System; + using Microsoft.Azure.Cosmos.CosmosElements; using Microsoft.Azure.Cosmos.Pagination; internal sealed class QueryState : State { - public QueryState(string continuationToken) + public QueryState(CosmosElement value) { - this.ContinuationToken = continuationToken ?? throw new ArgumentNullException(nameof(continuationToken)); + this.Value = value ?? throw new ArgumentNullException(nameof(value)); } - public string ContinuationToken { get; } + public CosmosElement Value { get; } } -} +} \ No newline at end of file diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/QueryClient/CosmosQueryClient.cs b/Microsoft.Azure.Cosmos/src/Query/Core/QueryClient/CosmosQueryClient.cs index f9820b5457..cbe8861991 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/QueryClient/CosmosQueryClient.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/QueryClient/CosmosQueryClient.cs @@ -12,6 +12,7 @@ namespace Microsoft.Azure.Cosmos.Query.Core.QueryClient using Microsoft.Azure.Cosmos.Diagnostics; using Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.ItemProducers; using Microsoft.Azure.Cosmos.Query.Core.Monads; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline; using Microsoft.Azure.Cosmos.Query.Core.QueryPlan; internal abstract class CosmosQueryClient diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/QueryClient/CosmosQueryContext.cs b/Microsoft.Azure.Cosmos/src/Query/Core/QueryClient/CosmosQueryContext.cs index 6e9c963e6d..17701a7113 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/QueryClient/CosmosQueryContext.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/QueryClient/CosmosQueryContext.cs @@ -8,6 +8,7 @@ namespace Microsoft.Azure.Cosmos.Query.Core.QueryClient using System.Threading.Tasks; using Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.ItemProducers; using Microsoft.Azure.Cosmos.Query.Core.Monads; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline; using Microsoft.Azure.Cosmos.Query.Core.QueryPlan; using OperationType = Documents.OperationType; using PartitionKeyRangeIdentity = Documents.PartitionKeyRangeIdentity; diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/QueryPlan/QueryInfo.cs b/Microsoft.Azure.Cosmos/src/Query/Core/QueryPlan/QueryInfo.cs index fd59b1d3f6..9bd52eb443 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/QueryPlan/QueryInfo.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/QueryPlan/QueryInfo.cs @@ -6,9 +6,9 @@ namespace Microsoft.Azure.Cosmos.Query.Core.QueryPlan { using System.Collections.Generic; using System.Linq; - using Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.Aggregate; using Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.Distinct; using Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.OrderBy; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Aggregate; using Newtonsoft.Json; using Newtonsoft.Json.Converters; diff --git a/Microsoft.Azure.Cosmos/src/Query/v3Query/CosmosQueryClientCore.cs b/Microsoft.Azure.Cosmos/src/Query/v3Query/CosmosQueryClientCore.cs index 2edc0b91fb..01e828f8f5 100644 --- a/Microsoft.Azure.Cosmos/src/Query/v3Query/CosmosQueryClientCore.cs +++ b/Microsoft.Azure.Cosmos/src/Query/v3Query/CosmosQueryClientCore.cs @@ -18,6 +18,7 @@ namespace Microsoft.Azure.Cosmos using Microsoft.Azure.Cosmos.Query.Core; using Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.ItemProducers; using Microsoft.Azure.Cosmos.Query.Core.Monads; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline; using Microsoft.Azure.Cosmos.Query.Core.QueryClient; using Microsoft.Azure.Cosmos.Query.Core.QueryPlan; using Microsoft.Azure.Cosmos.Routing; @@ -342,13 +343,23 @@ private static TryCatch GetCosmosElementResponse( cosmosQueryExecutionInfo = default; } + QueryState queryState; + if (cosmosResponseMessage.Headers.ContinuationToken != null) + { + queryState = new QueryState(CosmosString.Create(cosmosResponseMessage.Headers.ContinuationToken)); + } + else + { + queryState = default; + } + QueryPage response = new QueryPage( documents, cosmosResponseMessage.Headers.RequestCharge, cosmosResponseMessage.Headers.ActivityId, responseLengthBytes, cosmosQueryExecutionInfo, - new QueryState(cosmosResponseMessage.Headers.ContinuationToken)); + queryState); return TryCatch.FromResult(response); } diff --git a/Microsoft.Azure.Cosmos/src/Query/v3Query/CosmosQueryContextCore.cs b/Microsoft.Azure.Cosmos/src/Query/v3Query/CosmosQueryContextCore.cs index 6aae5a6b03..7522abc00f 100644 --- a/Microsoft.Azure.Cosmos/src/Query/v3Query/CosmosQueryContextCore.cs +++ b/Microsoft.Azure.Cosmos/src/Query/v3Query/CosmosQueryContextCore.cs @@ -11,6 +11,7 @@ namespace Microsoft.Azure.Cosmos.Query using Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.ItemProducers; using Microsoft.Azure.Cosmos.Query.Core.Metrics; using Microsoft.Azure.Cosmos.Query.Core.Monads; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline; using Microsoft.Azure.Cosmos.Query.Core.QueryClient; using Microsoft.Azure.Cosmos.Query.Core.QueryPlan; using Microsoft.Azure.Documents; diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/Query/MockCosmosQueryClient.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/Query/MockCosmosQueryClient.cs index fd08d199ce..ab310b0d77 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/Query/MockCosmosQueryClient.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/Query/MockCosmosQueryClient.cs @@ -11,6 +11,7 @@ using Microsoft.Azure.Cosmos.Diagnostics; using Microsoft.Azure.Cosmos.Query.Core.Monads; using Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.ItemProducers; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline; /// /// A helper that forces the SDK to use the gateway or the service interop for the query plan diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/CosmosQueryUnitTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/CosmosQueryUnitTests.cs index 492409158b..d7b67dc566 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/CosmosQueryUnitTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/CosmosQueryUnitTests.cs @@ -23,6 +23,7 @@ namespace Microsoft.Azure.Cosmos.Tests using Microsoft.Azure.Cosmos.Query.Core.ExecutionContext; using Microsoft.Azure.Cosmos.Query.Core.Metrics; using Microsoft.Azure.Cosmos.Query.Core.Monads; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Aggregate; using Microsoft.Azure.Cosmos.Query.Core.QueryClient; using Microsoft.Azure.Cosmos.Query.Core.QueryPlan; using Microsoft.Azure.Cosmos.Resource.CosmosExceptions; diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/InMemoryCollectionQueryDataSource.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/InMemoryCollectionQueryDataSource.cs index 487d7dbc00..ec57143c24 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/InMemoryCollectionQueryDataSource.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/InMemoryCollectionQueryDataSource.cs @@ -11,8 +11,9 @@ namespace Microsoft.Azure.Cosmos.Tests.Query using Microsoft.Azure.Cosmos.CosmosElements; using Microsoft.Azure.Cosmos.CosmosElements.Numbers; using Microsoft.Azure.Cosmos.Query.Core; - using Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.ItemProducers; using Microsoft.Azure.Cosmos.Query.Core.Monads; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Partitions; internal sealed class InMemoryCollectionQueryDataSource : IQueryDataSource { @@ -64,7 +65,7 @@ public Task> ExecuteQueryAsync( activityId: Guid.NewGuid().ToString(), responseLengthInBytes: 1337, cosmosQueryExecutionInfo: default, - state: tryExecuteQuery.Result.resourceIdentifer.HasValue ? new QueryState(tryExecuteQuery.Result.resourceIdentifer.Value.ToString()) : null); + state: tryExecuteQuery.Result.resourceIdentifer.HasValue ? new QueryState(CosmosString.Create(tryExecuteQuery.Result.resourceIdentifer.Value.ToString())) : null); return Task.FromResult(TryCatch.FromResult(queryPage)); } diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/AggregateQueryPipelineStageTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/AggregateQueryPipelineStageTests.cs new file mode 100644 index 0000000000..8da2f7b7c5 --- /dev/null +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/AggregateQueryPipelineStageTests.cs @@ -0,0 +1,127 @@ +//------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +//------------------------------------------------------------ + +namespace Microsoft.Azure.Cosmos.Tests.Query.Pipeline +{ + using System; + using System.Collections.Generic; + using System.Linq; + using System.Threading.Tasks; + using Microsoft.Azure.Cosmos.CosmosElements; + 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.Aggregate; + using Microsoft.VisualStudio.TestTools.UnitTesting; + + [TestClass] + public class AggregateQueryPipelineStageTests + { + [TestMethod] + public async Task SinglePageAsync() + { + IReadOnlyList> pages = new List>() + { + new List() + { + CosmosElement.Parse("{\"payload\": {\"$1\": {\"item\": 42}}}") + } + }; + + List elements = await AggregateQueryPipelineStageTests.CreateAndDrain( + pages: pages, + executionEnvironment: ExecutionEnvironment.Compute, + aggregates: new List() { AggregateOperator.Sum }, + aliasToAggregateType: new Dictionary() { { "$1", AggregateOperator.Sum } }, + orderedAliases: new List() { "$1" }, + hasSelectValue: false, + continuationToken: null); + + Assert.AreEqual(1, elements.Count); + Assert.AreEqual(42, Number64.ToLong(((elements[0] as CosmosObject)["$1"] as CosmosNumber).Value)); + } + + [TestMethod] + public async Task MultiplePagesAsync() + { + long[] values = new long[] { 42, 1337 }; + IReadOnlyList> pages = values + .Select(value => new List() + { + CosmosElement.Parse($"{{\"payload\": {{\"$1\": {{\"item\": {value}}}}}}}") + }) + .ToList(); + + List elements = await AggregateQueryPipelineStageTests.CreateAndDrain( + pages: pages, + executionEnvironment: ExecutionEnvironment.Compute, + aggregates: new List() { AggregateOperator.Sum }, + aliasToAggregateType: new Dictionary() { { "$1", AggregateOperator.Sum } }, + orderedAliases: new List() { "$1" }, + hasSelectValue: false, + continuationToken: null); + + Assert.AreEqual(1, elements.Count); + Assert.AreEqual(values.Sum(), Number64.ToLong(((elements[0] as CosmosObject)["$1"] as CosmosNumber).Value)); + } + + [TestMethod] + public async Task UndefinedSinglePageAsync() + { + IReadOnlyList> pages = new List>() + { + new List() + { + CosmosElement.Parse("{\"payload\": {\"$1\": {}}}") + } + }; + + List elements = await AggregateQueryPipelineStageTests.CreateAndDrain( + pages: pages, + executionEnvironment: ExecutionEnvironment.Compute, + aggregates: new List() { AggregateOperator.Sum }, + aliasToAggregateType: new Dictionary() { { "$1", AggregateOperator.Sum } }, + orderedAliases: new List() { "$1" }, + hasSelectValue: false, + continuationToken: null); + + Assert.AreEqual(1, elements.Count); + Assert.AreEqual(0, (elements[0] as CosmosObject).Keys.Count()); + } + + private static async Task> CreateAndDrain( + IReadOnlyList> pages, + ExecutionEnvironment executionEnvironment, + IReadOnlyList aggregates, + IReadOnlyDictionary aliasToAggregateType, + IReadOnlyList orderedAliases, + bool hasSelectValue, + CosmosElement continuationToken) + { + IQueryPipelineStage source = new MockQueryPipelineStage(pages); + + TryCatch tryCreateAggregateQueryPipelineStage = await AggregateQueryPipelineStage.TryCreateAsync( + executionEnvironment: executionEnvironment, + aggregates: aggregates, + aliasToAggregateType: aliasToAggregateType, + orderedAliases: orderedAliases, + hasSelectValue: hasSelectValue, + continuationToken: continuationToken, + tryCreateSourceAsync: (CosmosElement continuationToken) => Task.FromResult(TryCatch.FromResult(source))); + Assert.IsTrue(tryCreateAggregateQueryPipelineStage.Succeeded); + + IQueryPipelineStage aggregateQueryPipelineStage = tryCreateAggregateQueryPipelineStage.Result; + + List elements = new List(); + await foreach (TryCatch page in new EnumerableStage(aggregateQueryPipelineStage)) + { + page.ThrowIfFailed(); + + elements.AddRange(page.Result.Documents); + } + + return elements; + } + } +} diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/EnumerableStage.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/EnumerableStage.cs new file mode 100644 index 0000000000..0620fa309a --- /dev/null +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/EnumerableStage.cs @@ -0,0 +1,27 @@ +//------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +//------------------------------------------------------------ + +namespace Microsoft.Azure.Cosmos.Tests.Query.Pipeline +{ + using System; + using System.Collections.Generic; + using System.Threading; + using Microsoft.Azure.Cosmos.Query.Core.Monads; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline; + + internal sealed class EnumerableStage : IAsyncEnumerable> + { + private readonly IQueryPipelineStage stage; + + public EnumerableStage(IQueryPipelineStage stage) + { + this.stage = stage ?? throw new ArgumentNullException(nameof(stage)); + } + + public IAsyncEnumerator> GetAsyncEnumerator(CancellationToken cancellationToken = default) + { + return this.stage; + } + } +} diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/MockQueryPipelineStage.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/MockQueryPipelineStage.cs new file mode 100644 index 0000000000..c97f92eaad --- /dev/null +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/MockQueryPipelineStage.cs @@ -0,0 +1,40 @@ +//------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +//------------------------------------------------------------ + +namespace Microsoft.Azure.Cosmos.Tests.Query.Pipeline +{ + using System; + using System.Collections.Generic; + using System.Threading; + using System.Threading.Tasks; + using Microsoft.Azure.Cosmos.CosmosElements; + using Microsoft.Azure.Cosmos.Query.Core.Monads; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline; + + internal sealed class MockQueryPipelineStage : QueryPipelineStageBase + { + private readonly IReadOnlyList> pages; + private int pageIndex; + + public MockQueryPipelineStage(IReadOnlyList> pages) + : base(FinishedQueryPipelineStage.Value) + { + this.pages = pages ?? throw new ArgumentNullException(nameof(pages)); + } + + protected override Task> GetNextPageAsync(CancellationToken cancellationToken) + { + IReadOnlyList documents = this.pages[this.pageIndex++]; + QueryState state = (this.pageIndex == this.pages.Count) ? null : new QueryState(CosmosString.Create(this.pageIndex.ToString())); + QueryPage page = new QueryPage( + documents: documents, + requestCharge: default, + activityId: Guid.NewGuid().ToString(), + responseLengthInBytes: default, + cosmosQueryExecutionInfo: default, + state: state); + return Task.FromResult(TryCatch.FromResult(page)); + } + } +} 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 280980dafe..bff9aac39b 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 @@ -7,6 +7,8 @@ using Microsoft.Azure.Cosmos.Pagination; using Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.ItemProducers; using Microsoft.Azure.Cosmos.Query.Core.Monads; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Partitions; using Microsoft.Azure.Cosmos.Tests.Pagination; using Microsoft.Azure.Documents; using Microsoft.VisualStudio.TestTools.UnitTesting; diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/QueryPlanBaselineTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/QueryPlanBaselineTests.cs index c3ec1ff418..1757105f7f 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/QueryPlanBaselineTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/QueryPlanBaselineTests.cs @@ -15,6 +15,7 @@ using Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.Aggregate; using Microsoft.Azure.Cosmos.Query.Core.Monads; using System.Linq; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Aggregate; /// /// Tests for . @@ -1332,9 +1333,9 @@ public override QueryPlanBaselineTestOutput ExecuteTest(QueryPlanBaselineTestInp input.SqlQuerySpec, input.PartitionKeyDefinition, requireFormattableOrderByQuery: true, - isContinuationExpected: true, + isContinuationExpected: false, allowNonValueAggregateQuery: true, - hasLogicalPartitionKey: true); + hasLogicalPartitionKey: false); if (info.Failed) { From 647af75388d913c10df12aaba29b78ecdf29683d Mon Sep 17 00:00:00 2001 From: Brandon Chong Date: Thu, 2 Jul 2020 14:14:05 -0700 Subject: [PATCH 23/85] got skip working --- .../PartitionRangePageEnumerator.cs | 4 +- .../SkipDocumentQueryExecutionComponent.cs | 31 +----- .../ItemProducers/ItemProducer.cs | 2 +- .../AggregateQueryPipelineStage.Client.cs | 1 + .../AggregateQueryPipelineStage.Compute.cs | 2 + .../QueryPartitionRangePageEnumerator.cs | 2 +- .../src/Query/Core/Pipeline/QueryPage.cs | 4 + .../Core/Pipeline/QueryPipelineStageBase.cs | 4 +- .../Skip/SkipQueryPipelineStage.Client.cs} | 105 ++++++++++-------- .../Skip/SkipQueryPipelineStage.Compute.cs} | 105 ++++++++++-------- .../Pipeline/Skip/SkipQueryPipelineStage.cs | 54 +++++++++ .../Query/v3Query/CosmosQueryClientCore.cs | 1 + .../InMemoryCollectionQueryDataSource.cs | 1 + .../Query/Pipeline/MockQueryPipelineStage.cs | 1 + .../Pipeline/SkipQueryPipelineStageTests.cs | 68 ++++++++++++ 15 files changed, 254 insertions(+), 131 deletions(-) rename Microsoft.Azure.Cosmos/src/Query/Core/{ExecutionComponent/SkipTake/SkipDocumentQueryExecutionComponent.Client.cs => Pipeline/Skip/SkipQueryPipelineStage.Client.cs} (65%) rename Microsoft.Azure.Cosmos/src/Query/Core/{ExecutionComponent/SkipTake/SkipDocumentQueryExecutionComponent.Compute.cs => Pipeline/Skip/SkipQueryPipelineStage.Compute.cs} (63%) create mode 100644 Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Skip/SkipQueryPipelineStage.cs create mode 100644 Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/SkipQueryPipelineStageTests.cs diff --git a/Microsoft.Azure.Cosmos/src/Pagination/PartitionRangePageEnumerator.cs b/Microsoft.Azure.Cosmos/src/Pagination/PartitionRangePageEnumerator.cs index 2ac2d15fff..aa35740874 100644 --- a/Microsoft.Azure.Cosmos/src/Pagination/PartitionRangePageEnumerator.cs +++ b/Microsoft.Azure.Cosmos/src/Pagination/PartitionRangePageEnumerator.cs @@ -39,14 +39,14 @@ public async ValueTask MoveNextAsync() return false; } - this.hasStarted = true; - this.Current = await this.GetNextPageAsync(default); if (this.Current.Succeeded) { this.State = this.Current.Result.State; } + this.hasStarted = true; + return true; } diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/SkipDocumentQueryExecutionComponent.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/SkipDocumentQueryExecutionComponent.cs index 0153f6a1e6..d07a69bf63 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/SkipDocumentQueryExecutionComponent.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/SkipDocumentQueryExecutionComponent.cs @@ -32,36 +32,7 @@ public static Task> TryCreateAsync( CosmosElement continuationToken, Func>> tryCreateSourceAsync) { - Task> tryCreate; - switch (executionEnvironment) - { - case ExecutionEnvironment.Client: - tryCreate = ClientSkipDocumentQueryExecutionComponent.TryCreateAsync( - offsetCount, - continuationToken, - tryCreateSourceAsync); - break; - - case ExecutionEnvironment.Compute: - tryCreate = ComputeSkipDocumentQueryExecutionComponent.TryCreateAsync( - offsetCount, - continuationToken, - tryCreateSourceAsync); - break; - - default: - throw new ArgumentException($"Unknown {nameof(ExecutionEnvironment)}: {executionEnvironment}"); - } - - return tryCreate; - } - - public override bool IsDone - { - get - { - return this.Source.IsDone; - } + return default; } } } \ No newline at end of file diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/ItemProducers/ItemProducer.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/ItemProducers/ItemProducer.cs index 1717278fb6..2ad20bebe2 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/ItemProducers/ItemProducer.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/ItemProducers/ItemProducer.cs @@ -286,7 +286,7 @@ public async Task BufferMoreDocumentsAsync(CancellationToken token) activityId: tryGetQueryPage.Result.ActivityId, responseLengthBytes: tryGetQueryPage.Result.ResponseLengthInBytes, disallowContinuationTokenMessage: default, - continuationToken: ((CosmosString)tryGetQueryPage.Result.State.Value).Value, + continuationToken: tryGetQueryPage.Result.State == null ? null : ((CosmosString)tryGetQueryPage.Result.State.Value).Value, cosmosQueryExecutionInfo: tryGetQueryPage.Result.CosmosQueryExecutionInfo); } else diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Aggregate/AggregateQueryPipelineStage.Client.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Aggregate/AggregateQueryPipelineStage.Client.cs index ed688d6a3d..fa8ff94db9 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Aggregate/AggregateQueryPipelineStage.Client.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Aggregate/AggregateQueryPipelineStage.Client.cs @@ -110,6 +110,7 @@ protected override async Task> GetNextPageAsync(Cancellation activityId: default, responseLengthInBytes: responseLengthBytes, cosmosQueryExecutionInfo: default, + disallowContinuationTokenMessage: default, state: default); return TryCatch.FromResult(queryPage); diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Aggregate/AggregateQueryPipelineStage.Compute.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Aggregate/AggregateQueryPipelineStage.Compute.cs index f1a97e3268..1aa6ab0e67 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Aggregate/AggregateQueryPipelineStage.Compute.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Aggregate/AggregateQueryPipelineStage.Compute.cs @@ -111,6 +111,7 @@ protected override async Task> GetNextPageAsync(Cancellation activityId: default, responseLengthInBytes: default, cosmosQueryExecutionInfo: default, + disallowContinuationTokenMessage: default, state: default); return TryCatch.FromResult(finalPage); @@ -145,6 +146,7 @@ protected override async Task> GetNextPageAsync(Cancellation activityId: sourcePage.ActivityId, responseLengthInBytes: sourcePage.ResponseLengthInBytes, cosmosQueryExecutionInfo: sourcePage.CosmosQueryExecutionInfo, + disallowContinuationTokenMessage: sourcePage.DisallowContinuationTokenMessage, state: queryState); return TryCatch.FromResult(emptyPage); diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Partitions/QueryPartitionRangePageEnumerator.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Partitions/QueryPartitionRangePageEnumerator.cs index dbbc5db103..8025ef702c 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Partitions/QueryPartitionRangePageEnumerator.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Partitions/QueryPartitionRangePageEnumerator.cs @@ -38,7 +38,7 @@ public QueryPartitionRangePageEnumerator( public override Task> GetNextPageAsync(CancellationToken cancellationToken) => this.queryDataSource.ExecuteQueryAsync( sqlQuerySpec: this.sqlQuerySpec, - continuationToken: this.State.Value != null ? ((CosmosString)this.State.Value).Value : null, + continuationToken: this.State == null ? null : ((CosmosString)this.State.Value).Value, partitionKeyRangeId: int.Parse(((FeedRangePartitionKeyRange)this.Range).PartitionKeyRangeId), pageSize: this.pageSize, cancellationToken); diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/QueryPage.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/QueryPage.cs index 6410266f09..8f0b4c27c8 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/QueryPage.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/QueryPage.cs @@ -18,6 +18,7 @@ public QueryPage( string activityId, long responseLengthInBytes, CosmosQueryExecutionInfo cosmosQueryExecutionInfo, + string disallowContinuationTokenMessage, QueryState state) : base(state) { @@ -26,6 +27,7 @@ public QueryPage( this.ActivityId = activityId; this.ResponseLengthInBytes = responseLengthInBytes < 0 ? throw new ArgumentOutOfRangeException(nameof(responseLengthInBytes)) : responseLengthInBytes; this.CosmosQueryExecutionInfo = cosmosQueryExecutionInfo; + this.DisallowContinuationTokenMessage = disallowContinuationTokenMessage; } public IReadOnlyList Documents { get; } @@ -37,5 +39,7 @@ public QueryPage( public long ResponseLengthInBytes { get; } public CosmosQueryExecutionInfo CosmosQueryExecutionInfo { get; } + + public string DisallowContinuationTokenMessage { get; } } } \ No newline at end of file diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/QueryPipelineStageBase.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/QueryPipelineStageBase.cs index 1864ee7e99..28267a894d 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/QueryPipelineStageBase.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/QueryPipelineStageBase.cs @@ -36,14 +36,14 @@ public async ValueTask MoveNextAsync() return false; } - this.hasStarted = true; - this.Current = await this.GetNextPageAsync(default); if (this.Current.Succeeded) { this.State = this.Current.Result.State; } + this.hasStarted = true; + return true; } } diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/SkipDocumentQueryExecutionComponent.Client.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Skip/SkipQueryPipelineStage.Client.cs similarity index 65% rename from Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/SkipDocumentQueryExecutionComponent.Client.cs rename to Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Skip/SkipQueryPipelineStage.Client.cs index cceb370cc3..b86ec80bc8 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/SkipDocumentQueryExecutionComponent.Client.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Skip/SkipQueryPipelineStage.Client.cs @@ -1,33 +1,34 @@ -//------------------------------------------------------------ +// ------------------------------------------------------------ // Copyright (c) Microsoft Corporation. All rights reserved. -//------------------------------------------------------------ -namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.SkipTake +// ------------------------------------------------------------ + +namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.Skip { using System; using System.Collections.Generic; using System.Linq; + using System.Text; using System.Threading; using System.Threading.Tasks; using Microsoft.Azure.Cosmos.CosmosElements; using Microsoft.Azure.Cosmos.Query.Core.Exceptions; using Microsoft.Azure.Cosmos.Query.Core.Monads; - using Microsoft.Azure.Cosmos.Query.Core.QueryClient; using Newtonsoft.Json; - internal abstract partial class SkipDocumentQueryExecutionComponent : DocumentQueryExecutionComponentBase + internal abstract partial class SkipQueryPipelineStage : QueryPipelineStageBase { - private sealed class ClientSkipDocumentQueryExecutionComponent : SkipDocumentQueryExecutionComponent + private sealed class ClientSkipQueryPipelineStage : SkipQueryPipelineStage { - private ClientSkipDocumentQueryExecutionComponent(IDocumentQueryExecutionComponent source, long skipCount) + private ClientSkipQueryPipelineStage(IQueryPipelineStage source, long skipCount) : base(source, skipCount) { // Work is done in base constructor. } - public static async Task> TryCreateAsync( + public static async Task> TryCreateAsync( int offsetCount, CosmosElement continuationToken, - Func>> tryCreateSourceAsync) + Func>> tryCreateSourceAsync) { if (tryCreateSourceAsync == null) { @@ -39,8 +40,9 @@ public static async Task> TryCreateAs { if (!OffsetContinuationToken.TryParse(continuationToken.ToString(), out offsetContinuationToken)) { - return TryCatch.FromException( - new MalformedContinuationTokenException($"Invalid {nameof(SkipDocumentQueryExecutionComponent)}: {continuationToken}.")); + return TryCatch.FromException( + new MalformedContinuationTokenException( + $"Invalid {nameof(SkipQueryPipelineStage)}: {continuationToken}.")); } } else @@ -50,8 +52,9 @@ public static async Task> TryCreateAs if (offsetContinuationToken.Offset > offsetCount) { - return TryCatch.FromException( - new MalformedContinuationTokenException("offset count in continuation token can not be greater than the offsetcount in the query.")); + return TryCatch.FromException( + new MalformedContinuationTokenException( + "offset count in continuation token can not be greater than the offsetcount in the query.")); } CosmosElement sourceToken; @@ -60,7 +63,7 @@ public static async Task> TryCreateAs TryCatch tryParse = CosmosElement.Monadic.Parse(offsetContinuationToken.SourceToken); if (tryParse.Failed) { - return TryCatch.FromException( + return TryCatch.FromException( new MalformedContinuationTokenException( message: $"source token: '{offsetContinuationToken.SourceToken ?? ""}' is not valid.", innerException: tryParse.Exception)); @@ -73,51 +76,61 @@ public static async Task> TryCreateAs sourceToken = null; } - return (await tryCreateSourceAsync(sourceToken)) - .Try((source) => new ClientSkipDocumentQueryExecutionComponent( - source, - offsetContinuationToken.Offset)); + TryCatch tryCreateSource = await tryCreateSourceAsync(sourceToken); + if (tryCreateSource.Failed) + { + return tryCreateSource; + } + + IQueryPipelineStage stage = new ClientSkipQueryPipelineStage( + tryCreateSource.Result, + offsetContinuationToken.Offset); + + return TryCatch.FromResult(stage); } - public override async Task DrainAsync(int maxElements, CancellationToken token) + protected override async Task> GetNextPageAsync(CancellationToken cancellationToken) { - token.ThrowIfCancellationRequested(); - QueryResponseCore sourcePage = await base.DrainAsync(maxElements, token); - if (!sourcePage.IsSuccess) + cancellationToken.ThrowIfCancellationRequested(); + + await this.inputStage.MoveNextAsync(); + TryCatch tryGetSourcePage = this.inputStage.Current; + if (tryGetSourcePage.Failed) { - return sourcePage; + return tryGetSourcePage; } - // skip the documents but keep all the other headers - IReadOnlyList documentsAfterSkip = sourcePage.CosmosElements.Skip(this.skipCount).ToList(); + QueryPage sourcePage = tryGetSourcePage.Result; - int numberOfDocumentsSkipped = sourcePage.CosmosElements.Count() - documentsAfterSkip.Count(); + // Skip the documents but keep all the other headers + IReadOnlyList documentsAfterSkip = sourcePage.Documents.Skip(this.skipCount).ToList(); + + int numberOfDocumentsSkipped = sourcePage.Documents.Count - documentsAfterSkip.Count; this.skipCount -= numberOfDocumentsSkipped; - string updatedContinuationToken; + QueryState state; if (sourcePage.DisallowContinuationTokenMessage == null) { - updatedContinuationToken = new OffsetContinuationToken( + string token = new OffsetContinuationToken( offset: this.skipCount, - sourceToken: sourcePage.ContinuationToken).ToString(); + sourceToken: sourcePage.State != null ? ((CosmosString)sourcePage.State.Value).Value : null).ToString(); + state = new QueryState(CosmosString.Create(token)); } else { - updatedContinuationToken = null; + state = null; } - return QueryResponseCore.CreateSuccess( - result: documentsAfterSkip, - continuationToken: updatedContinuationToken, - disallowContinuationTokenMessage: sourcePage.DisallowContinuationTokenMessage, - activityId: sourcePage.ActivityId, + QueryPage queryPage = new QueryPage( + documents: documentsAfterSkip, requestCharge: sourcePage.RequestCharge, - responseLengthBytes: sourcePage.ResponseLengthBytes); - } + activityId: sourcePage.ActivityId, + responseLengthInBytes: sourcePage.ResponseLengthInBytes, + cosmosQueryExecutionInfo: sourcePage.CosmosQueryExecutionInfo, + disallowContinuationTokenMessage: sourcePage.DisallowContinuationTokenMessage, + state: state); - public override CosmosElement GetCosmosElementContinuationToken() - { - throw new NotImplementedException(); + return TryCatch.FromResult(queryPage); } /// @@ -145,19 +158,13 @@ public OffsetContinuationToken(int offset, string sourceToken) /// The number of items to skip in the query. /// [JsonProperty("offset")] - public int Offset - { - get; - } + public int Offset { get; } /// /// Gets the continuation token for the source component of the query. /// [JsonProperty("sourceToken")] - public string SourceToken - { - get; - } + public string SourceToken { get; } /// /// Tries to parse out the OffsetContinuationToken. @@ -195,4 +202,4 @@ public override string ToString() } } } -} \ No newline at end of file +} diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/SkipDocumentQueryExecutionComponent.Compute.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Skip/SkipQueryPipelineStage.Compute.cs similarity index 63% rename from Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/SkipDocumentQueryExecutionComponent.Compute.cs rename to Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Skip/SkipQueryPipelineStage.Compute.cs index 7479565e4d..0c5f56c4b1 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/SkipDocumentQueryExecutionComponent.Compute.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Skip/SkipQueryPipelineStage.Compute.cs @@ -1,36 +1,33 @@ -//------------------------------------------------------------ +// ------------------------------------------------------------ // Copyright (c) Microsoft Corporation. All rights reserved. -//------------------------------------------------------------ -namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.SkipTake +// ------------------------------------------------------------ + +namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.Skip { using System; using System.Collections.Generic; using System.Linq; - using System.Runtime.CompilerServices; using System.Threading; using System.Threading.Tasks; using Microsoft.Azure.Cosmos.CosmosElements; using Microsoft.Azure.Cosmos.CosmosElements.Numbers; - using Microsoft.Azure.Cosmos.Json; using Microsoft.Azure.Cosmos.Query.Core.Exceptions; using Microsoft.Azure.Cosmos.Query.Core.Monads; - using Microsoft.Azure.Cosmos.Query.Core.QueryClient; - using Microsoft.Azure.Documents; - internal abstract partial class SkipDocumentQueryExecutionComponent : DocumentQueryExecutionComponentBase + internal abstract partial class SkipQueryPipelineStage : QueryPipelineStageBase { - private sealed class ComputeSkipDocumentQueryExecutionComponent : SkipDocumentQueryExecutionComponent + private sealed class ComputeSkipQueryPipelineStage : SkipQueryPipelineStage { - private ComputeSkipDocumentQueryExecutionComponent(IDocumentQueryExecutionComponent source, long skipCount) + private ComputeSkipQueryPipelineStage(IQueryPipelineStage source, long skipCount) : base(source, skipCount) { // Work is done in base constructor. } - public static async Task> TryCreateAsync( + public static async Task> TryCreateAsync( int offsetCount, CosmosElement continuationToken, - Func>> tryCreateSourceAsync) + Func>> tryCreateSourceAsync) { if (tryCreateSourceAsync == null) { @@ -43,8 +40,8 @@ public static async Task> TryCreateAs (bool parsed, OffsetContinuationToken parsedToken) = OffsetContinuationToken.TryParse(continuationToken); if (!parsed) { - return TryCatch.FromException( - new MalformedContinuationTokenException($"Invalid {nameof(SkipDocumentQueryExecutionComponent)}: {continuationToken}.")); + return TryCatch.FromException( + new MalformedContinuationTokenException($"Invalid {nameof(SkipQueryPipelineStage)}: {continuationToken}.")); } offsetContinuationToken = parsedToken; @@ -56,51 +53,67 @@ public static async Task> TryCreateAs if (offsetContinuationToken.Offset > offsetCount) { - return TryCatch.FromException( - new MalformedContinuationTokenException("offset count in continuation token can not be greater than the offsetcount in the query.")); + return TryCatch.FromException( + new MalformedContinuationTokenException( + "offset count in continuation token can not be greater than the offsetcount in the query.")); } - return (await tryCreateSourceAsync(offsetContinuationToken.SourceToken)) - .Try((source) => new ComputeSkipDocumentQueryExecutionComponent( - source, - offsetContinuationToken.Offset)); + TryCatch tryCreateSource = await tryCreateSourceAsync(offsetContinuationToken.SourceToken); + if (tryCreateSource.Failed) + { + return tryCreateSource; + } + + IQueryPipelineStage stage = new ComputeSkipQueryPipelineStage( + tryCreateSource.Result, + offsetContinuationToken.Offset); + + return TryCatch.FromResult(stage); } - public override async Task DrainAsync(int maxElements, CancellationToken token) + protected override async Task> GetNextPageAsync(CancellationToken cancellationToken) { - token.ThrowIfCancellationRequested(); - QueryResponseCore sourcePage = await base.DrainAsync(maxElements, token); - if (!sourcePage.IsSuccess) + cancellationToken.ThrowIfCancellationRequested(); + + await this.inputStage.MoveNextAsync(); + TryCatch tryGetSourcePage = this.inputStage.Current; + if (tryGetSourcePage.Failed) { - return sourcePage; + return tryGetSourcePage; } - // skip the documents but keep all the other headers - IReadOnlyList documentsAfterSkip = sourcePage.CosmosElements.Skip(this.skipCount).ToList(); + QueryPage sourcePage = tryGetSourcePage.Result; - int numberOfDocumentsSkipped = sourcePage.CosmosElements.Count() - documentsAfterSkip.Count(); - this.skipCount -= numberOfDocumentsSkipped; + // Skip the documents but keep all the other headers + IReadOnlyList documentsAfterSkip = sourcePage.Documents.Skip(this.skipCount).ToList(); - return QueryResponseCore.CreateSuccess( - result: documentsAfterSkip, - continuationToken: null, - disallowContinuationTokenMessage: DocumentQueryExecutionComponentBase.UseCosmosElementContinuationTokenInstead, - activityId: sourcePage.ActivityId, - requestCharge: sourcePage.RequestCharge, - responseLengthBytes: sourcePage.ResponseLengthBytes); - } + int numberOfDocumentsSkipped = sourcePage.Documents.Count() - documentsAfterSkip.Count(); + this.skipCount -= numberOfDocumentsSkipped; - public override CosmosElement GetCosmosElementContinuationToken() - { - if (this.IsDone) + QueryState state; + if (sourcePage.State == null) { - return default; + state = default; } + else + { + OffsetContinuationToken offsetContinuationToken = new OffsetContinuationToken( + offset: this.skipCount, + sourceToken: sourcePage.State.Value); + + state = new QueryState(OffsetContinuationToken.ToCosmosElement(offsetContinuationToken)); + } + + QueryPage queryPage = new QueryPage( + documents: documentsAfterSkip, + requestCharge: sourcePage.RequestCharge, + activityId: sourcePage.ActivityId, + responseLengthInBytes: sourcePage.ResponseLengthInBytes, + cosmosQueryExecutionInfo: sourcePage.CosmosQueryExecutionInfo, + disallowContinuationTokenMessage: sourcePage.DisallowContinuationTokenMessage, + state: state); - OffsetContinuationToken offsetContinuationToken = new OffsetContinuationToken( - offset: this.skipCount, - sourceToken: this.Source.GetCosmosElementContinuationToken()); - return OffsetContinuationToken.ToCosmosElement(offsetContinuationToken); + return TryCatch.FromResult(queryPage); } /// @@ -190,4 +203,4 @@ public static (bool parsed, OffsetContinuationToken offsetContinuationToken) Try } } } -} \ No newline at end of file +} diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Skip/SkipQueryPipelineStage.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Skip/SkipQueryPipelineStage.cs new file mode 100644 index 0000000000..3448218e5b --- /dev/null +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Skip/SkipQueryPipelineStage.cs @@ -0,0 +1,54 @@ +// ------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +// ------------------------------------------------------------ + +namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.Skip +{ + using System; + using System.Collections.Generic; + using System.Text; + using System.Threading.Tasks; + using Microsoft.Azure.Cosmos.CosmosElements; + using Microsoft.Azure.Cosmos.Query.Core.ExecutionContext; + using Microsoft.Azure.Cosmos.Query.Core.Monads; + + internal abstract partial class SkipQueryPipelineStage : QueryPipelineStageBase + { + private int skipCount; + + protected SkipQueryPipelineStage( + IQueryPipelineStage source, + long skipCount) + : base(source) + { + if (skipCount > int.MaxValue) + { + throw new ArgumentOutOfRangeException(nameof(skipCount)); + } + + this.skipCount = (int)skipCount; + } + + public static Task> TryCreateAsync( + ExecutionEnvironment executionEnvironment, + int offsetCount, + CosmosElement continuationToken, + Func>> tryCreateSourceAsync) + { + Task> tryCreate = executionEnvironment switch + { + ExecutionEnvironment.Client => ClientSkipQueryPipelineStage.TryCreateAsync( + offsetCount, + continuationToken, + tryCreateSourceAsync), + ExecutionEnvironment.Compute => ComputeSkipQueryPipelineStage.TryCreateAsync( + offsetCount, + continuationToken, + tryCreateSourceAsync), + _ => throw new ArgumentException($"Unknown {nameof(ExecutionEnvironment)}: {executionEnvironment}"), + }; + + return tryCreate; + } + } +} diff --git a/Microsoft.Azure.Cosmos/src/Query/v3Query/CosmosQueryClientCore.cs b/Microsoft.Azure.Cosmos/src/Query/v3Query/CosmosQueryClientCore.cs index 01e828f8f5..1038deea1e 100644 --- a/Microsoft.Azure.Cosmos/src/Query/v3Query/CosmosQueryClientCore.cs +++ b/Microsoft.Azure.Cosmos/src/Query/v3Query/CosmosQueryClientCore.cs @@ -359,6 +359,7 @@ private static TryCatch GetCosmosElementResponse( cosmosResponseMessage.Headers.ActivityId, responseLengthBytes, cosmosQueryExecutionInfo, + disallowContinuationTokenMessage: null, queryState); return TryCatch.FromResult(response); diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/InMemoryCollectionQueryDataSource.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/InMemoryCollectionQueryDataSource.cs index ec57143c24..2f934f9e68 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/InMemoryCollectionQueryDataSource.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/InMemoryCollectionQueryDataSource.cs @@ -65,6 +65,7 @@ public Task> ExecuteQueryAsync( activityId: Guid.NewGuid().ToString(), responseLengthInBytes: 1337, cosmosQueryExecutionInfo: default, + disallowContinuationTokenMessage: default, state: tryExecuteQuery.Result.resourceIdentifer.HasValue ? new QueryState(CosmosString.Create(tryExecuteQuery.Result.resourceIdentifer.Value.ToString())) : null); return Task.FromResult(TryCatch.FromResult(queryPage)); diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/MockQueryPipelineStage.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/MockQueryPipelineStage.cs index c97f92eaad..dbbea34dc3 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/MockQueryPipelineStage.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/MockQueryPipelineStage.cs @@ -33,6 +33,7 @@ protected override Task> GetNextPageAsync(CancellationToken activityId: Guid.NewGuid().ToString(), responseLengthInBytes: default, cosmosQueryExecutionInfo: default, + disallowContinuationTokenMessage: default, state: state); return Task.FromResult(TryCatch.FromResult(page)); } diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/SkipQueryPipelineStageTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/SkipQueryPipelineStageTests.cs new file mode 100644 index 0000000000..baa87ed9a3 --- /dev/null +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/SkipQueryPipelineStageTests.cs @@ -0,0 +1,68 @@ +namespace Microsoft.Azure.Cosmos.Tests.Query.Pipeline +{ + using System; + using System.Collections.Generic; + using System.Linq; + using System.Threading.Tasks; + using Microsoft.Azure.Cosmos.CosmosElements; + 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.Skip; + using Microsoft.VisualStudio.TestTools.UnitTesting; + + [TestClass] + public sealed class SkipQueryPipelineStageTests + { + [TestMethod] + public async Task SanityTests() + { + long[] values = new long[] { 42, 1337 }; + IReadOnlyList> pages = values + .Select(value => new List() + { + CosmosElement.Parse($"{{\"item\": {value}}}") + }) + .ToList(); + + foreach (int offsetCount in new int[] { 0, 1, values.Length, 2 * values.Length }) + { + List elements = await SkipQueryPipelineStageTests.CreateAndDrainAsync( + pages: pages, + executionEnvironment: ExecutionEnvironment.Compute, + offsetCount: offsetCount, + continuationToken: null); + + Assert.AreEqual(Math.Max(values.Length - offsetCount, 0), elements.Count); + } + } + + private static async Task> CreateAndDrainAsync( + IReadOnlyList> pages, + ExecutionEnvironment executionEnvironment, + int offsetCount, + CosmosElement continuationToken) + { + IQueryPipelineStage source = new MockQueryPipelineStage(pages); + + TryCatch tryCreateSkipQueryPipelineStage = await SkipQueryPipelineStage.TryCreateAsync( + executionEnvironment: executionEnvironment, + offsetCount: offsetCount, + continuationToken: continuationToken, + tryCreateSourceAsync: (CosmosElement continuationToken) => Task.FromResult(TryCatch.FromResult(source))); + Assert.IsTrue(tryCreateSkipQueryPipelineStage.Succeeded); + + IQueryPipelineStage aggregateQueryPipelineStage = tryCreateSkipQueryPipelineStage.Result; + + List elements = new List(); + await foreach (TryCatch page in new EnumerableStage(aggregateQueryPipelineStage)) + { + page.ThrowIfFailed(); + + elements.AddRange(page.Result.Documents); + } + + return elements; + } + } +} From 4b599b5e1e8ea7abfed425152ca149f1fbf10ead Mon Sep 17 00:00:00 2001 From: Brandon Chong Date: Thu, 2 Jul 2020 16:44:09 -0700 Subject: [PATCH 24/85] adding take stage --- .../TakeDocumentQueryExecutionComponent.cs | 76 ++------- .../PipelinedDocumentQueryExecutionContext.cs | 4 +- .../Take/TakeQueryPipelineStage.Client.cs} | 144 ++++++++++-------- .../Take/TakeQueryPipelineStage.Compute.cs} | 117 +++++++++----- .../Pipeline/Take/TakeQueryPipelineStage.cs | 59 +++++++ .../CosmosQueryUnitTests.cs | 4 +- .../Pipeline/TakeQueryPipelineStageTests.cs | 68 +++++++++ 7 files changed, 299 insertions(+), 173 deletions(-) rename Microsoft.Azure.Cosmos/src/Query/Core/{ExecutionComponent/SkipTake/TakeDocumentQueryExecutionComponent.Client.cs => Pipeline/Take/TakeQueryPipelineStage.Client.cs} (71%) rename Microsoft.Azure.Cosmos/src/Query/Core/{ExecutionComponent/SkipTake/TakeDocumentQueryExecutionComponent.Compute.cs => Pipeline/Take/TakeQueryPipelineStage.Compute.cs} (54%) create mode 100644 Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Take/TakeQueryPipelineStage.cs create mode 100644 Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/TakeQueryPipelineStageTests.cs diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/TakeDocumentQueryExecutionComponent.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/TakeDocumentQueryExecutionComponent.cs index 48a76be5f7..7c9a56c714 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/TakeDocumentQueryExecutionComponent.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/TakeDocumentQueryExecutionComponent.cs @@ -1,92 +1,38 @@ //------------------------------------------------------------ // Copyright (c) Microsoft Corporation. All rights reserved. //------------------------------------------------------------ + namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.SkipTake { using System; using System.Threading.Tasks; using Microsoft.Azure.Cosmos.CosmosElements; + using Microsoft.Azure.Cosmos.Query.Core.ContinuationTokens; using Microsoft.Azure.Cosmos.Query.Core.ExecutionContext; using Microsoft.Azure.Cosmos.Query.Core.Monads; internal abstract partial class TakeDocumentQueryExecutionComponent : DocumentQueryExecutionComponentBase { - private int takeCount; + private int skipCount; - protected TakeDocumentQueryExecutionComponent( - IDocumentQueryExecutionComponent source, - int takeCount) + protected TakeDocumentQueryExecutionComponent(IDocumentQueryExecutionComponent source, long skipCount) : base(source) { - this.takeCount = takeCount; - } - - public static Task> TryCreateLimitDocumentQueryExecutionComponentAsync( - ExecutionEnvironment executionEnvironment, - int limitCount, - CosmosElement requestContinuationToken, - Func>> tryCreateSourceAsync) - { - Task> tryCreateComponentAsync; - switch (executionEnvironment) + if (skipCount > int.MaxValue) { - case ExecutionEnvironment.Client: - tryCreateComponentAsync = ClientTakeDocumentQueryExecutionComponent.TryCreateLimitDocumentQueryExecutionComponentAsync( - limitCount, - requestContinuationToken, - tryCreateSourceAsync); - break; - - case ExecutionEnvironment.Compute: - tryCreateComponentAsync = ComputeTakeDocumentQueryExecutionComponent.TryCreateAsync( - limitCount, - requestContinuationToken, - tryCreateSourceAsync); - break; - - default: - throw new ArgumentOutOfRangeException($"Unknown {nameof(ExecutionEnvironment)}: {executionEnvironment}."); + throw new ArgumentOutOfRangeException(nameof(skipCount)); } - return tryCreateComponentAsync; + this.skipCount = (int)skipCount; } - public static Task> TryCreateTopDocumentQueryExecutionComponentAsync( + public static Task> TryCreateAsync( ExecutionEnvironment executionEnvironment, - int topCount, - CosmosElement requestContinuationToken, + int offsetCount, + CosmosElement continuationToken, Func>> tryCreateSourceAsync) { - Task> tryCreateComponentAsync; - switch (executionEnvironment) - { - case ExecutionEnvironment.Client: - tryCreateComponentAsync = ClientTakeDocumentQueryExecutionComponent.TryCreateTopDocumentQueryExecutionComponentAsync( - topCount, - requestContinuationToken, - tryCreateSourceAsync); - break; - - case ExecutionEnvironment.Compute: - tryCreateComponentAsync = ComputeTakeDocumentQueryExecutionComponent.TryCreateAsync( - topCount, - requestContinuationToken, - tryCreateSourceAsync); - break; - - default: - throw new ArgumentOutOfRangeException($"Unknown {nameof(ExecutionEnvironment)}: {executionEnvironment}."); - } - - return tryCreateComponentAsync; - } - - public override bool IsDone - { - get - { - return this.Source.IsDone || this.takeCount <= 0; - } + return default; } } } \ No newline at end of file diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/PipelinedDocumentQueryExecutionContext.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/PipelinedDocumentQueryExecutionContext.cs index 83ebb489e5..a2e7f4a227 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/PipelinedDocumentQueryExecutionContext.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/PipelinedDocumentQueryExecutionContext.cs @@ -243,7 +243,7 @@ Task> tryCreateParallelComponentAsync Func>> tryCreateSourceAsync = tryCreatePipelineAsync; tryCreatePipelineAsync = async (continuationToken) => { - return await TakeDocumentQueryExecutionComponent.TryCreateLimitDocumentQueryExecutionComponentAsync( + return await TakeDocumentQueryExecutionComponent.TryCreateAsync( executionEnvironment, queryInfo.Limit.Value, continuationToken, @@ -256,7 +256,7 @@ Task> tryCreateParallelComponentAsync Func>> tryCreateSourceAsync = tryCreatePipelineAsync; tryCreatePipelineAsync = async (continuationToken) => { - return await TakeDocumentQueryExecutionComponent.TryCreateTopDocumentQueryExecutionComponentAsync( + return await TakeDocumentQueryExecutionComponent.TryCreateAsync( executionEnvironment, queryInfo.Top.Value, continuationToken, diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/TakeDocumentQueryExecutionComponent.Client.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Take/TakeQueryPipelineStage.Client.cs similarity index 71% rename from Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/TakeDocumentQueryExecutionComponent.Client.cs rename to Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Take/TakeQueryPipelineStage.Client.cs index 4bc5a7eb73..9ffcb2c6ec 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/TakeDocumentQueryExecutionComponent.Client.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Take/TakeQueryPipelineStage.Client.cs @@ -1,7 +1,8 @@ //------------------------------------------------------------ // Copyright (c) Microsoft Corporation. All rights reserved. //------------------------------------------------------------ -namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.SkipTake + +namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.Take { using System; using System.Collections.Generic; @@ -11,25 +12,27 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.SkipTake using Microsoft.Azure.Cosmos.CosmosElements; using Microsoft.Azure.Cosmos.Query.Core.Exceptions; using Microsoft.Azure.Cosmos.Query.Core.Monads; - using Microsoft.Azure.Cosmos.Query.Core.QueryClient; using Newtonsoft.Json; - internal abstract partial class TakeDocumentQueryExecutionComponent : DocumentQueryExecutionComponentBase + internal abstract partial class TakeQueryPipelineStage : QueryPipelineStageBase { - private sealed class ClientTakeDocumentQueryExecutionComponent : TakeDocumentQueryExecutionComponent + private sealed class ClientTakeQueryPipelineStage : TakeQueryPipelineStage { private readonly TakeEnum takeEnum; - private ClientTakeDocumentQueryExecutionComponent(IDocumentQueryExecutionComponent source, int takeCount, TakeEnum takeEnum) + private ClientTakeQueryPipelineStage( + IQueryPipelineStage source, + int takeCount, + TakeEnum takeEnum) : base(source, takeCount) { this.takeEnum = takeEnum; } - public static async Task> TryCreateLimitDocumentQueryExecutionComponentAsync( + public static async Task> TryCreateLimitStageAsync( int limitCount, CosmosElement requestContinuationToken, - Func>> tryCreateSourceAsync) + Func>> tryCreateSourceAsync) { if (limitCount < 0) { @@ -46,8 +49,9 @@ public static async Task> TryCreateLi { if (!LimitContinuationToken.TryParse(requestContinuationToken.ToString(), out limitContinuationToken)) { - return TryCatch.FromException( - new MalformedContinuationTokenException($"Malformed {nameof(LimitContinuationToken)}: {requestContinuationToken}.")); + return TryCatch.FromException( + new MalformedContinuationTokenException( + $"Malformed {nameof(LimitContinuationToken)}: {requestContinuationToken}.")); } } else @@ -57,8 +61,9 @@ public static async Task> TryCreateLi if (limitContinuationToken.Limit > limitCount) { - return TryCatch.FromException( - new MalformedContinuationTokenException($"{nameof(LimitContinuationToken.Limit)} in {nameof(LimitContinuationToken)}: {requestContinuationToken}: {limitContinuationToken.Limit} can not be greater than the limit count in the query: {limitCount}.")); + return TryCatch.FromException( + new MalformedContinuationTokenException( + $"{nameof(LimitContinuationToken.Limit)} in {nameof(LimitContinuationToken)}: {requestContinuationToken}: {limitContinuationToken.Limit} can not be greater than the limit count in the query: {limitCount}.")); } CosmosElement sourceToken; @@ -67,7 +72,7 @@ public static async Task> TryCreateLi TryCatch tryParse = CosmosElement.Monadic.Parse(limitContinuationToken.SourceToken); if (tryParse.Failed) { - return TryCatch.FromException( + return TryCatch.FromException( new MalformedContinuationTokenException( message: $"Malformed {nameof(LimitContinuationToken)}: {requestContinuationToken}.", innerException: tryParse.Exception)); @@ -80,17 +85,24 @@ public static async Task> TryCreateLi sourceToken = null; } - return (await tryCreateSourceAsync(sourceToken)) - .Try((source) => new ClientTakeDocumentQueryExecutionComponent( - source, + TryCatch tryCreateSource = await tryCreateSourceAsync(sourceToken); + if (tryCreateSource.Failed) + { + return tryCreateSource; + } + + IQueryPipelineStage stage = new ClientTakeQueryPipelineStage( + tryCreateSource.Result, limitContinuationToken.Limit, - TakeEnum.Limit)); + TakeEnum.Limit); + + return TryCatch.FromResult(stage); } - public static async Task> TryCreateTopDocumentQueryExecutionComponentAsync( + public static async Task> TryCreateTopStageAsync( int topCount, CosmosElement requestContinuationToken, - Func>> tryCreateSourceAsync) + Func>> tryCreateSourceAsync) { if (topCount < 0) { @@ -107,8 +119,9 @@ public static async Task> TryCreateTo { if (!TopContinuationToken.TryParse(requestContinuationToken.ToString(), out topContinuationToken)) { - return TryCatch.FromException( - new MalformedContinuationTokenException($"Malformed {nameof(LimitContinuationToken)}: {requestContinuationToken}.")); + return TryCatch.FromException( + new MalformedContinuationTokenException( + $"Malformed {nameof(LimitContinuationToken)}: {requestContinuationToken}.")); } } else @@ -118,8 +131,9 @@ public static async Task> TryCreateTo if (topContinuationToken.Top > topCount) { - return TryCatch.FromException( - new MalformedContinuationTokenException($"{nameof(TopContinuationToken.Top)} in {nameof(TopContinuationToken)}: {requestContinuationToken}: {topContinuationToken.Top} can not be greater than the top count in the query: {topCount}.")); + return TryCatch.FromException( + new MalformedContinuationTokenException( + $"{nameof(TopContinuationToken.Top)} in {nameof(TopContinuationToken)}: {requestContinuationToken}: {topContinuationToken.Top} can not be greater than the top count in the query: {topCount}.")); } CosmosElement sourceToken; @@ -128,7 +142,7 @@ public static async Task> TryCreateTo TryCatch tryParse = CosmosElement.Monadic.Parse(topContinuationToken.SourceToken); if (tryParse.Failed) { - return TryCatch.FromException( + return TryCatch.FromException( new MalformedContinuationTokenException( message: $"{nameof(TopContinuationToken.SourceToken)} in {nameof(TopContinuationToken)}: {requestContinuationToken}: {topContinuationToken.SourceToken ?? ""} was malformed.", innerException: tryParse.Exception)); @@ -141,63 +155,67 @@ public static async Task> TryCreateTo sourceToken = null; } - return (await tryCreateSourceAsync(sourceToken)) - .Try((source) => new ClientTakeDocumentQueryExecutionComponent( - source, + TryCatch tryCreateSource = await tryCreateSourceAsync(sourceToken); + if (tryCreateSource.Failed) + { + return tryCreateSource; + } + + IQueryPipelineStage stage = new ClientTakeQueryPipelineStage( + tryCreateSource.Result, topContinuationToken.Top, - TakeEnum.Top)); + TakeEnum.Top); + + return TryCatch.FromResult(stage); } - public override async Task DrainAsync(int maxElements, CancellationToken token) + protected override async Task> GetNextPageAsync(CancellationToken cancellationToken) { - token.ThrowIfCancellationRequested(); - QueryResponseCore sourcePage = await base.DrainAsync(maxElements, token); - if (!sourcePage.IsSuccess) + cancellationToken.ThrowIfCancellationRequested(); + + await this.inputStage.MoveNextAsync(); + TryCatch tryGetSourcePage = this.inputStage.Current; + if (tryGetSourcePage.Failed) { - return sourcePage; + return tryGetSourcePage; } - List takedDocuments = sourcePage.CosmosElements.Take(this.takeCount).ToList(); + QueryPage sourcePage = tryGetSourcePage.Result; + + List takedDocuments = sourcePage.Documents.Take(this.takeCount).ToList(); this.takeCount -= takedDocuments.Count; - string updatedContinuationToken; - if (!this.IsDone && (sourcePage.DisallowContinuationTokenMessage == null)) + QueryState state; + if ((sourcePage.State != null) && (sourcePage.DisallowContinuationTokenMessage == null)) { - switch (this.takeEnum) + string updatedContinuationToken = this.takeEnum switch { - case TakeEnum.Limit: - updatedContinuationToken = new LimitContinuationToken( - limit: this.takeCount, - sourceToken: sourcePage.ContinuationToken).ToString(); - break; - - case TakeEnum.Top: - updatedContinuationToken = new TopContinuationToken( - top: this.takeCount, - sourceToken: sourcePage.ContinuationToken).ToString(); - break; - - default: - throw new ArgumentOutOfRangeException($"Unknown {nameof(TakeEnum)}: {this.takeEnum}."); - } + TakeEnum.Limit => new LimitContinuationToken( + limit: this.takeCount, + sourceToken: ((CosmosString)sourcePage.State.Value).Value).ToString(), + TakeEnum.Top => new TopContinuationToken( + top: this.takeCount, + sourceToken: ((CosmosString)sourcePage.State.Value).Value).ToString(), + _ => throw new ArgumentOutOfRangeException($"Unknown {nameof(TakeEnum)}: {this.takeEnum}."), + }; + + state = new QueryState(CosmosString.Create(updatedContinuationToken)); } else { - updatedContinuationToken = null; + state = null; } - return QueryResponseCore.CreateSuccess( - result: takedDocuments, - continuationToken: updatedContinuationToken, - disallowContinuationTokenMessage: sourcePage.DisallowContinuationTokenMessage, - activityId: sourcePage.ActivityId, + QueryPage queryPage = new QueryPage( + documents: takedDocuments, requestCharge: sourcePage.RequestCharge, - responseLengthBytes: sourcePage.ResponseLengthBytes); - } + activityId: sourcePage.ActivityId, + responseLengthInBytes: sourcePage.ResponseLengthInBytes, + cosmosQueryExecutionInfo: sourcePage.CosmosQueryExecutionInfo, + disallowContinuationTokenMessage: sourcePage.DisallowContinuationTokenMessage, + state: state); - public override CosmosElement GetCosmosElementContinuationToken() - { - throw new NotImplementedException(); + return TryCatch.FromResult(queryPage); } private enum TakeEnum @@ -354,4 +372,4 @@ public override string ToString() } } } -} \ No newline at end of file +} diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/TakeDocumentQueryExecutionComponent.Compute.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Take/TakeQueryPipelineStage.Compute.cs similarity index 54% rename from Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/TakeDocumentQueryExecutionComponent.Compute.cs rename to Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Take/TakeQueryPipelineStage.Compute.cs index 6028d43afb..1d0693ae42 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/TakeDocumentQueryExecutionComponent.Compute.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Take/TakeQueryPipelineStage.Compute.cs @@ -1,7 +1,8 @@ -//------------------------------------------------------------ +// ------------------------------------------------------------ // Copyright (c) Microsoft Corporation. All rights reserved. -//------------------------------------------------------------ -namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.SkipTake +// ------------------------------------------------------------ + +namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.Take { using System; using System.Collections.Generic; @@ -12,22 +13,40 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.SkipTake using Microsoft.Azure.Cosmos.CosmosElements.Numbers; using Microsoft.Azure.Cosmos.Query.Core.Exceptions; using Microsoft.Azure.Cosmos.Query.Core.Monads; - using Microsoft.Azure.Cosmos.Query.Core.QueryClient; + using Newtonsoft.Json; - internal abstract partial class TakeDocumentQueryExecutionComponent : DocumentQueryExecutionComponentBase + internal abstract partial class TakeQueryPipelineStage : QueryPipelineStageBase { - private sealed class ComputeTakeDocumentQueryExecutionComponent : TakeDocumentQueryExecutionComponent + private sealed class ComputeTakeQueryPipelineStage : TakeQueryPipelineStage { - private ComputeTakeDocumentQueryExecutionComponent(IDocumentQueryExecutionComponent source, int takeCount) + private ComputeTakeQueryPipelineStage( + IQueryPipelineStage source, + int takeCount) : base(source, takeCount) { // Work is done in the base class. } - public static async Task> TryCreateAsync( + public static Task> TryCreateLimitStageAsync( + int takeCount, + CosmosElement requestContinuationToken, + Func>> tryCreateSourceAsync) => ComputeTakeQueryPipelineStage.TryCreateAsync( + takeCount, + requestContinuationToken, + tryCreateSourceAsync); + + public static Task> TryCreateTopStageAsync( + int takeCount, + CosmosElement requestContinuationToken, + Func>> tryCreateSourceAsync) => ComputeTakeQueryPipelineStage.TryCreateAsync( + takeCount, + requestContinuationToken, + tryCreateSourceAsync); + + private static async Task> TryCreateAsync( int takeCount, CosmosElement requestContinuationToken, - Func>> tryCreateSourceAsync) + Func>> tryCreateSourceAsync) { if (takeCount < 0) { @@ -44,8 +63,9 @@ public static async Task> TryCreateAs { if (!TakeContinuationToken.TryParse(requestContinuationToken, out takeContinuationToken)) { - return TryCatch.FromException( - new MalformedContinuationTokenException($"Malformed {nameof(TakeContinuationToken)}: {requestContinuationToken}.")); + return TryCatch.FromException( + new MalformedContinuationTokenException( + $"Malformed {nameof(TakeContinuationToken)}: {requestContinuationToken}.")); } } else @@ -55,48 +75,63 @@ public static async Task> TryCreateAs if (takeContinuationToken.TakeCount > takeCount) { - return TryCatch.FromException( - new MalformedContinuationTokenException($"{nameof(TakeContinuationToken.TakeCount)} in {nameof(TakeContinuationToken)}: {requestContinuationToken}: {takeContinuationToken.TakeCount} can not be greater than the limit count in the query: {takeCount}.")); + return TryCatch.FromException( + new MalformedContinuationTokenException( + $"{nameof(TakeContinuationToken.TakeCount)} in {nameof(TakeContinuationToken)}: {requestContinuationToken}: {takeContinuationToken.TakeCount} can not be greater than the limit count in the query: {takeCount}.")); + } + + TryCatch tryCreateSource = await tryCreateSourceAsync(takeContinuationToken.SourceToken); + if (tryCreateSource.Failed) + { + return tryCreateSource; } - return (await tryCreateSourceAsync(takeContinuationToken.SourceToken)) - .Try((source) => new ComputeTakeDocumentQueryExecutionComponent( - source, - takeContinuationToken.TakeCount)); + IQueryPipelineStage stage = new ComputeTakeQueryPipelineStage( + tryCreateSource.Result, + takeContinuationToken.TakeCount); + + return TryCatch.FromResult(stage); } - public override async Task DrainAsync(int maxElements, CancellationToken token) + protected override async Task> GetNextPageAsync(CancellationToken cancellationToken) { - token.ThrowIfCancellationRequested(); - QueryResponseCore sourcePage = await base.DrainAsync(maxElements, token); - if (!sourcePage.IsSuccess) + cancellationToken.ThrowIfCancellationRequested(); + + await this.inputStage.MoveNextAsync(); + TryCatch tryGetSourcePage = this.inputStage.Current; + if (tryGetSourcePage.Failed) { - return sourcePage; + return tryGetSourcePage; } - List takedDocuments = sourcePage.CosmosElements.Take(this.takeCount).ToList(); - this.takeCount -= takedDocuments.Count; + QueryPage sourcePage = tryGetSourcePage.Result; - return QueryResponseCore.CreateSuccess( - result: takedDocuments, - continuationToken: null, - disallowContinuationTokenMessage: DocumentQueryExecutionComponentBase.UseCosmosElementContinuationTokenInstead, - activityId: sourcePage.ActivityId, - requestCharge: sourcePage.RequestCharge, - responseLengthBytes: sourcePage.ResponseLengthBytes); - } + List takedDocuments = sourcePage.Documents.Take(this.takeCount).ToList(); + this.takeCount -= takedDocuments.Count; - public override CosmosElement GetCosmosElementContinuationToken() - { - if (this.IsDone) + QueryState queryState; + if (sourcePage.State != null) { - return default; + TakeContinuationToken takeContinuationToken = new TakeContinuationToken( + takeCount: this.takeCount, + sourceToken: sourcePage.State.Value); + queryState = new QueryState(TakeContinuationToken.ToCosmosElement(takeContinuationToken)); } + else + { + queryState = default; + } + + QueryPage queryPage = new QueryPage( + documents: takedDocuments, + requestCharge: sourcePage.RequestCharge, + activityId: sourcePage.ActivityId, + responseLengthInBytes: sourcePage.ResponseLengthInBytes, + cosmosQueryExecutionInfo: sourcePage.CosmosQueryExecutionInfo, + disallowContinuationTokenMessage: sourcePage.DisallowContinuationTokenMessage, + state: queryState); - TakeContinuationToken takeContinuationToken = new TakeContinuationToken( - takeCount: this.takeCount, - sourceToken: this.Source.GetCosmosElementContinuationToken()); - return TakeContinuationToken.ToCosmosElement(takeContinuationToken); + return TryCatch.FromResult(queryPage); } private readonly struct TakeContinuationToken @@ -170,4 +205,4 @@ public static bool TryParse(CosmosElement value, out TakeContinuationToken takeC } } } -} \ No newline at end of file +} diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Take/TakeQueryPipelineStage.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Take/TakeQueryPipelineStage.cs new file mode 100644 index 0000000000..401891eee8 --- /dev/null +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Take/TakeQueryPipelineStage.cs @@ -0,0 +1,59 @@ +// ------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +// ------------------------------------------------------------ + +namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.Take +{ + using System; + using System.Threading.Tasks; + using Microsoft.Azure.Cosmos.CosmosElements; + using Microsoft.Azure.Cosmos.Query.Core.ExecutionContext; + using Microsoft.Azure.Cosmos.Query.Core.Monads; + + internal abstract partial class TakeQueryPipelineStage : QueryPipelineStageBase + { + private int takeCount; + + protected TakeQueryPipelineStage( + IQueryPipelineStage source, + int takeCount) + : base(source) + { + this.takeCount = takeCount; + } + + public static Task> TryCreateLimitStageAsync( + ExecutionEnvironment executionEnvironment, + int limitCount, + CosmosElement requestContinuationToken, + Func>> tryCreateSourceAsync) => executionEnvironment switch + { + ExecutionEnvironment.Client => ClientTakeQueryPipelineStage.TryCreateLimitStageAsync( + limitCount, + requestContinuationToken, + tryCreateSourceAsync), + ExecutionEnvironment.Compute => ComputeTakeQueryPipelineStage.TryCreateLimitStageAsync( + limitCount, + requestContinuationToken, + tryCreateSourceAsync), + _ => throw new ArgumentOutOfRangeException($"Unknown {nameof(ExecutionEnvironment)}: {executionEnvironment}."), + }; + + public static Task> TryCreateTopStageAsync( + ExecutionEnvironment executionEnvironment, + int limitCount, + CosmosElement requestContinuationToken, + Func>> tryCreateSourceAsync) => executionEnvironment switch + { + ExecutionEnvironment.Client => ClientTakeQueryPipelineStage.TryCreateTopStageAsync( + limitCount, + requestContinuationToken, + tryCreateSourceAsync), + ExecutionEnvironment.Compute => ComputeTakeQueryPipelineStage.TryCreateTopStageAsync( + limitCount, + requestContinuationToken, + tryCreateSourceAsync), + _ => throw new ArgumentOutOfRangeException($"Unknown {nameof(ExecutionEnvironment)}: {executionEnvironment}."), + }; + } +} diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/CosmosQueryUnitTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/CosmosQueryUnitTests.cs index d7b67dc566..4207a3da4d 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/CosmosQueryUnitTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/CosmosQueryUnitTests.cs @@ -331,13 +331,13 @@ public async Task TestCosmosQueryPartitionKeyDefinition() null, func)).Result); - components.Add((await TakeDocumentQueryExecutionComponent.TryCreateLimitDocumentQueryExecutionComponentAsync( + components.Add((await TakeDocumentQueryExecutionComponent.TryCreateAsync( ExecutionEnvironment.Client, 5, null, func)).Result); - components.Add((await TakeDocumentQueryExecutionComponent.TryCreateTopDocumentQueryExecutionComponentAsync( + components.Add((await TakeDocumentQueryExecutionComponent.TryCreateAsync( ExecutionEnvironment.Client, 5, null, diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/TakeQueryPipelineStageTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/TakeQueryPipelineStageTests.cs new file mode 100644 index 0000000000..fcbc04eda2 --- /dev/null +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/TakeQueryPipelineStageTests.cs @@ -0,0 +1,68 @@ +namespace Microsoft.Azure.Cosmos.Tests.Query.Pipeline +{ + using System; + using System.Collections.Generic; + using System.Linq; + using System.Threading.Tasks; + using Microsoft.Azure.Cosmos.CosmosElements; + 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.Take; + using Microsoft.VisualStudio.TestTools.UnitTesting; + + [TestClass] + public sealed class TakeQueryPipelineStageTests + { + [TestMethod] + public async Task SanityTests() + { + long[] values = new long[] { 42, 1337 }; + IReadOnlyList> pages = values + .Select(value => new List() + { + CosmosElement.Parse($"{{\"item\": {value}}}") + }) + .ToList(); + + foreach (int takeCount in new int[] { 0, 1, values.Length, 2 * values.Length }) + { + List elements = await TakeQueryPipelineStageTests.CreateAndDrainAsync( + pages: pages, + executionEnvironment: ExecutionEnvironment.Compute, + takeCount: takeCount, + continuationToken: null); + + Assert.AreEqual(Math.Min(takeCount, values.Length), elements.Count); + } + } + + private static async Task> CreateAndDrainAsync( + IReadOnlyList> pages, + ExecutionEnvironment executionEnvironment, + int takeCount, + CosmosElement continuationToken) + { + IQueryPipelineStage source = new MockQueryPipelineStage(pages); + + TryCatch tryCreateSkipQueryPipelineStage = await TakeQueryPipelineStage.TryCreateLimitStageAsync( + executionEnvironment: executionEnvironment, + limitCount: takeCount, + requestContinuationToken: continuationToken, + tryCreateSourceAsync: (CosmosElement continuationToken) => Task.FromResult(TryCatch.FromResult(source))); + Assert.IsTrue(tryCreateSkipQueryPipelineStage.Succeeded); + + IQueryPipelineStage takeQueryPipelineStage = tryCreateSkipQueryPipelineStage.Result; + + List elements = new List(); + await foreach (TryCatch page in new EnumerableStage(takeQueryPipelineStage)) + { + page.ThrowIfFailed(); + + elements.AddRange(page.Result.Documents); + } + + return elements; + } + } +} From 0fd47fd71a29b8dda26cda782668c9fabf840e2c Mon Sep 17 00:00:00 2001 From: Brandon Chong Date: Thu, 2 Jul 2020 17:33:06 -0700 Subject: [PATCH 25/85] added distinct --- ...DistinctDocumentQueryExecutionComponent.cs | 45 +------- .../GroupByDocumentQueryExecutionComponent.cs | 2 +- .../CosmosQueryExecutionContextFactory.cs | 2 +- .../ItemProducers/ItemComparer.cs | 3 +- .../PipelinedDocumentQueryExecutionContext.cs | 1 + .../Distinct/DistinctHash.cs | 2 +- .../DistinctMap.OrderedDistinctMap.cs | 3 +- .../DistinctMap.UnorderedDistinctMap.cs | 5 +- .../Distinct/DistinctMap.cs | 3 +- .../DistinctQueryPipelineStage.Client.cs} | 109 ++++++++---------- .../DistinctQueryPipelineStage.Compute.cs} | 96 +++++++-------- .../Distinct/DistinctQueryPipelineStage.cs | 56 +++++++++ .../Distinct/DistinctQueryType.cs | 3 +- .../src/Query/Core/QueryPlan/QueryInfo.cs | 2 +- .../CosmosQueryUnitTests.cs | 1 + .../Query/DistinctHashBaselineTests.cs | 3 +- .../DistinctQueryPipelineStageTests.cs | 65 +++++++++++ 17 files changed, 239 insertions(+), 162 deletions(-) rename Microsoft.Azure.Cosmos/src/Query/Core/{ExecutionComponent => Pipeline}/Distinct/DistinctHash.cs (99%) rename Microsoft.Azure.Cosmos/src/Query/Core/{ExecutionComponent => Pipeline}/Distinct/DistinctMap.OrderedDistinctMap.cs (98%) rename Microsoft.Azure.Cosmos/src/Query/Core/{ExecutionComponent => Pipeline}/Distinct/DistinctMap.UnorderedDistinctMap.cs (99%) rename Microsoft.Azure.Cosmos/src/Query/Core/{ExecutionComponent => Pipeline}/Distinct/DistinctMap.cs (97%) rename Microsoft.Azure.Cosmos/src/Query/Core/{ExecutionComponent/Distinct/DistinctDocumentQueryExecutionComponent.Client.cs => Pipeline/Distinct/DistinctQueryPipelineStage.Client.cs} (68%) rename Microsoft.Azure.Cosmos/src/Query/Core/{ExecutionComponent/Distinct/DistinctDocumentQueryExecutionComponent.Compute.cs => Pipeline/Distinct/DistinctQueryPipelineStage.Compute.cs} (62%) create mode 100644 Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Distinct/DistinctQueryPipelineStage.cs rename Microsoft.Azure.Cosmos/src/Query/Core/{ExecutionComponent => Pipeline}/Distinct/DistinctQueryType.cs (93%) create mode 100644 Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/DistinctQueryPipelineStageTests.cs diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctDocumentQueryExecutionComponent.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctDocumentQueryExecutionComponent.cs index 09e35f81c7..71c7ad185b 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctDocumentQueryExecutionComponent.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctDocumentQueryExecutionComponent.cs @@ -6,23 +6,10 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.Distinct using System; using System.Threading.Tasks; using Microsoft.Azure.Cosmos.CosmosElements; - using Microsoft.Azure.Cosmos.Json; - using Microsoft.Azure.Cosmos.Query.Core.ContinuationTokens; using Microsoft.Azure.Cosmos.Query.Core.ExecutionContext; using Microsoft.Azure.Cosmos.Query.Core.Monads; - using Newtonsoft.Json; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Distinct; - /// - /// Distinct queries return documents that are distinct with a page. - /// This means that documents are not guaranteed to be distinct across continuations and partitions. - /// The reasoning for this is because the backend treats each continuation of a query as a separate request - /// and partitions are not aware of each other. - /// The solution is that the client keeps a running hash set of all the documents it has already seen, - /// so that when it encounters a duplicate document from another continuation it will not be emitted to the user. - /// The only problem is that if the user chooses to go through the continuation token API for DocumentQuery instead - /// of while(HasMoreResults) ExecuteNextAsync, then will see duplicates across continuations. - /// There is no workaround for that use case, since the continuation token will have to include all the documents seen. - /// internal abstract partial class DistinctDocumentQueryExecutionComponent : DocumentQueryExecutionComponentBase { /// @@ -38,39 +25,13 @@ protected DistinctDocumentQueryExecutionComponent( this.distinctMap = distinctMap ?? throw new ArgumentNullException(nameof(distinctMap)); } - public static async Task> TryCreateAsync( + public static Task> TryCreateAsync( ExecutionEnvironment executionEnvironment, CosmosElement requestContinuation, Func>> tryCreateSourceAsync, DistinctQueryType distinctQueryType) { - if (tryCreateSourceAsync == null) - { - throw new ArgumentNullException(nameof(tryCreateSourceAsync)); - } - - TryCatch tryCreateDistinctDocumentQueryExecutionComponent; - switch (executionEnvironment) - { - case ExecutionEnvironment.Client: - tryCreateDistinctDocumentQueryExecutionComponent = await ClientDistinctDocumentQueryExecutionComponent.TryCreateAsync( - requestContinuation, - tryCreateSourceAsync, - distinctQueryType); - break; - - case ExecutionEnvironment.Compute: - tryCreateDistinctDocumentQueryExecutionComponent = await ComputeDistinctDocumentQueryExecutionComponent.TryCreateAsync( - requestContinuation, - tryCreateSourceAsync, - distinctQueryType); - break; - - default: - throw new ArgumentException($"Unknown {nameof(ExecutionEnvironment)}: {executionEnvironment}."); - } - - return tryCreateDistinctDocumentQueryExecutionComponent; + return default; } } } diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/GroupBy/GroupByDocumentQueryExecutionComponent.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/GroupBy/GroupByDocumentQueryExecutionComponent.cs index f526b5e788..36707da09a 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/GroupBy/GroupByDocumentQueryExecutionComponent.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/GroupBy/GroupByDocumentQueryExecutionComponent.cs @@ -10,11 +10,11 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.GroupBy using System.Threading.Tasks; using Microsoft.Azure.Cosmos.CosmosElements; using Microsoft.Azure.Cosmos.Query.Core.Exceptions; - using Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.Distinct; using Microsoft.Azure.Cosmos.Query.Core.ExecutionContext; using Microsoft.Azure.Cosmos.Query.Core.Monads; using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Aggregate; using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Aggregate.Aggregators; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Distinct; /// /// Query execution component that groups groupings across continuations and pages. diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosQueryExecutionContextFactory.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosQueryExecutionContextFactory.cs index 532d36e6f1..16ca62bd30 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosQueryExecutionContextFactory.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosQueryExecutionContextFactory.cs @@ -13,8 +13,8 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionContext using Microsoft.Azure.Cosmos.Query.Core; using Microsoft.Azure.Cosmos.Query.Core.ContinuationTokens; using Microsoft.Azure.Cosmos.Query.Core.Exceptions; - using Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.Distinct; using Microsoft.Azure.Cosmos.Query.Core.Monads; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Distinct; using Microsoft.Azure.Cosmos.Query.Core.QueryClient; using Microsoft.Azure.Cosmos.Query.Core.QueryPlan; using Microsoft.Azure.Cosmos.SqlObjects; diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/ItemProducers/ItemComparer.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/ItemProducers/ItemComparer.cs index a735d71bec..0427dd7e06 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/ItemProducers/ItemComparer.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/ItemProducers/ItemComparer.cs @@ -7,8 +7,7 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.ItemProducers using System.Collections.Generic; using Microsoft.Azure.Cosmos.CosmosElements; using Microsoft.Azure.Cosmos.Json; - using Microsoft.Azure.Cosmos.Query.Core; - using Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.Distinct; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Distinct; /// /// Utility class used to compare all items that we get back from a query. diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/PipelinedDocumentQueryExecutionContext.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/PipelinedDocumentQueryExecutionContext.cs index a2e7f4a227..29c7bace17 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/PipelinedDocumentQueryExecutionContext.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/PipelinedDocumentQueryExecutionContext.cs @@ -19,6 +19,7 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionContext using Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.OrderBy; using Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.Parallel; using Microsoft.Azure.Cosmos.Query.Core.Monads; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Distinct; using Microsoft.Azure.Cosmos.Query.Core.QueryClient; using Microsoft.Azure.Cosmos.Query.Core.QueryPlan; using Microsoft.Azure.Documents.Collections; diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctHash.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Distinct/DistinctHash.cs similarity index 99% rename from Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctHash.cs rename to Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Distinct/DistinctHash.cs index eb738f18ef..a0e7a6fc5d 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctHash.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Distinct/DistinctHash.cs @@ -2,7 +2,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // ------------------------------------------------------------ -namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.Distinct +namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.Distinct { using System; using System.Collections.Generic; diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctMap.OrderedDistinctMap.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Distinct/DistinctMap.OrderedDistinctMap.cs similarity index 98% rename from Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctMap.OrderedDistinctMap.cs rename to Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Distinct/DistinctMap.OrderedDistinctMap.cs index da6e632d4c..fd159cfc29 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctMap.OrderedDistinctMap.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Distinct/DistinctMap.OrderedDistinctMap.cs @@ -1,7 +1,8 @@ //------------------------------------------------------------ // Copyright (c) Microsoft Corporation. All rights reserved. //------------------------------------------------------------ -namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.Distinct + +namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.Distinct { using System; using Microsoft.Azure.Cosmos.CosmosElements; diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctMap.UnorderedDistinctMap.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Distinct/DistinctMap.UnorderedDistinctMap.cs similarity index 99% rename from Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctMap.UnorderedDistinctMap.cs rename to Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Distinct/DistinctMap.UnorderedDistinctMap.cs index 50c19ff862..51ac06b927 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctMap.UnorderedDistinctMap.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Distinct/DistinctMap.UnorderedDistinctMap.cs @@ -1,7 +1,8 @@ //------------------------------------------------------------ // Copyright (c) Microsoft Corporation. All rights reserved. //------------------------------------------------------------ -namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.Distinct + +namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.Distinct { using System; using System.Collections.Generic; @@ -10,8 +11,6 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.Distinct using System.Text; using Microsoft.Azure.Cosmos.CosmosElements; using Microsoft.Azure.Cosmos.CosmosElements.Numbers; - using Microsoft.Azure.Cosmos.Json; - using Microsoft.Azure.Cosmos.Query.Core; using Microsoft.Azure.Cosmos.Query.Core.Exceptions; using Microsoft.Azure.Cosmos.Query.Core.Monads; diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctMap.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Distinct/DistinctMap.cs similarity index 97% rename from Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctMap.cs rename to Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Distinct/DistinctMap.cs index 4267bc9fff..2621f5d393 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctMap.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Distinct/DistinctMap.cs @@ -1,7 +1,8 @@ //------------------------------------------------------------ // Copyright (c) Microsoft Corporation. All rights reserved. //------------------------------------------------------------ -namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.Distinct + +namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.Distinct { using System; using Microsoft.Azure.Cosmos.CosmosElements; diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctDocumentQueryExecutionComponent.Client.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Distinct/DistinctQueryPipelineStage.Client.cs similarity index 68% rename from Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctDocumentQueryExecutionComponent.Client.cs rename to Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Distinct/DistinctQueryPipelineStage.Client.cs index 56e54d82fa..d46485d578 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctDocumentQueryExecutionComponent.Client.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Distinct/DistinctQueryPipelineStage.Client.cs @@ -1,36 +1,35 @@ //------------------------------------------------------------ // Copyright (c) Microsoft Corporation. All rights reserved. //------------------------------------------------------------ -namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.Distinct + +namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.Distinct { using System; using System.Collections.Generic; using System.Threading; using System.Threading.Tasks; using Microsoft.Azure.Cosmos.CosmosElements; - using Microsoft.Azure.Cosmos.Json; - using Microsoft.Azure.Cosmos.Query.Core.ContinuationTokens; using Microsoft.Azure.Cosmos.Query.Core.Exceptions; using Microsoft.Azure.Cosmos.Query.Core.Monads; - using Microsoft.Azure.Cosmos.Query.Core.QueryClient; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline; using Newtonsoft.Json; - internal abstract partial class DistinctDocumentQueryExecutionComponent : DocumentQueryExecutionComponentBase + internal abstract partial class DistinctQueryPipelineStage : QueryPipelineStageBase { /// /// Client implementaiton of Distinct. Here we only serialize the continuation token if there is a matching DISTINCT. /// - private sealed class ClientDistinctDocumentQueryExecutionComponent : DistinctDocumentQueryExecutionComponent + private sealed class ClientDistinctQueryPipelineStage : DistinctQueryPipelineStage { private static readonly string DisallowContinuationTokenMessage = "DISTINCT queries only return continuation tokens when there is a matching ORDER BY clause." + "For example if your query is 'SELECT DISTINCT VALUE c.name FROM c', then rewrite it as 'SELECT DISTINCT VALUE c.name FROM c ORDER BY c.name'."; private readonly DistinctQueryType distinctQueryType; - private ClientDistinctDocumentQueryExecutionComponent( + private ClientDistinctQueryPipelineStage( DistinctQueryType distinctQueryType, DistinctMap distinctMap, - IDocumentQueryExecutionComponent source) + IQueryPipelineStage source) : base(distinctMap, source) { if ((distinctQueryType != DistinctQueryType.Unordered) && (distinctQueryType != DistinctQueryType.Ordered)) @@ -41,9 +40,9 @@ private ClientDistinctDocumentQueryExecutionComponent( this.distinctQueryType = distinctQueryType; } - public static async Task> TryCreateAsync( + public static async Task> TryCreateAsync( CosmosElement requestContinuation, - Func>> tryCreateSourceAsync, + Func>> tryCreateSourceAsync, DistinctQueryType distinctQueryType) { if (tryCreateSourceAsync == null) @@ -56,7 +55,7 @@ public static async Task> TryCreateAs { if (!DistinctContinuationToken.TryParse(requestContinuation, out distinctContinuationToken)) { - return TryCatch.FromException( + return TryCatch.FromException( new MalformedContinuationTokenException( $"Invalid {nameof(DistinctContinuationToken)}: {requestContinuation}")); } @@ -83,7 +82,7 @@ public static async Task> TryCreateAs distinctMapToken); if (!tryCreateDistinctMap.Succeeded) { - return TryCatch.FromException(tryCreateDistinctMap.Exception); + return TryCatch.FromException(tryCreateDistinctMap.Exception); } CosmosElement sourceToken; @@ -92,7 +91,7 @@ public static async Task> TryCreateAs TryCatch tryParse = CosmosElement.Monadic.Parse(distinctContinuationToken.SourceToken); if (tryParse.Failed) { - return TryCatch.FromException( + return TryCatch.FromException( new MalformedContinuationTokenException( message: $"Invalid Source Token: {distinctContinuationToken.SourceToken}", innerException: tryParse.Exception)); @@ -105,38 +104,34 @@ public static async Task> TryCreateAs sourceToken = null; } - TryCatch tryCreateSource = await tryCreateSourceAsync(sourceToken); + TryCatch tryCreateSource = await tryCreateSourceAsync(sourceToken); if (!tryCreateSource.Succeeded) { - return TryCatch.FromException(tryCreateSource.Exception); + return TryCatch.FromException(tryCreateSource.Exception); } - return TryCatch.FromResult( - new ClientDistinctDocumentQueryExecutionComponent( + return TryCatch.FromResult( + new ClientDistinctQueryPipelineStage( distinctQueryType, tryCreateDistinctMap.Result, tryCreateSource.Result)); } - /// - /// Drains a page of results returning only distinct elements. - /// - /// The maximum number of items to drain. - /// The cancellation token. - /// A page of distinct results. - public override async Task DrainAsync(int maxElements, CancellationToken cancellationToken) + protected override async Task> GetNextPageAsync(CancellationToken cancellationToken) { cancellationToken.ThrowIfCancellationRequested(); - List distinctResults = new List(); - QueryResponseCore sourceResponse = await base.DrainAsync(maxElements, cancellationToken); - - if (!sourceResponse.IsSuccess) + await this.inputStage.MoveNextAsync(); + TryCatch tryGetSourcePage = this.inputStage.Current; + if (tryGetSourcePage.Failed) { - return sourceResponse; + return tryGetSourcePage; } - foreach (CosmosElement document in sourceResponse.CosmosElements) + QueryPage sourcePage = tryGetSourcePage.Result; + + List distinctResults = new List(); + foreach (CosmosElement document in sourcePage.Documents) { if (this.distinctMap.Add(document, out UInt128 hash)) { @@ -145,46 +140,44 @@ public override async Task DrainAsync(int maxElements, Cancel } // For clients we write out the continuation token if it's a streaming query. - QueryResponseCore queryResponseCore; + QueryPage queryPage; if (this.distinctQueryType == DistinctQueryType.Ordered) { - string updatedContinuationToken; - if (this.IsDone) + QueryState state; + if (sourcePage.State != null) { - updatedContinuationToken = null; + string updatedContinuationToken = new DistinctContinuationToken( + sourceToken: ((CosmosString)sourcePage.State.Value).Value, + distinctMapToken: this.distinctMap.GetContinuationToken()).ToString(); + state = new QueryState(CosmosString.Create(updatedContinuationToken)); } else { - updatedContinuationToken = new DistinctContinuationToken( - sourceToken: sourceResponse.ContinuationToken, - distinctMapToken: this.distinctMap.GetContinuationToken()).ToString(); + state = null; } - queryResponseCore = QueryResponseCore.CreateSuccess( - result: distinctResults, - continuationToken: updatedContinuationToken, - disallowContinuationTokenMessage: null, - activityId: sourceResponse.ActivityId, - requestCharge: sourceResponse.RequestCharge, - responseLengthBytes: sourceResponse.ResponseLengthBytes); + queryPage = new QueryPage( + documents: distinctResults, + requestCharge: sourcePage.RequestCharge, + activityId: sourcePage.ActivityId, + responseLengthInBytes: sourcePage.ResponseLengthInBytes, + cosmosQueryExecutionInfo: sourcePage.CosmosQueryExecutionInfo, + disallowContinuationTokenMessage: sourcePage.DisallowContinuationTokenMessage, + state: state); } else { - queryResponseCore = QueryResponseCore.CreateSuccess( - result: distinctResults, - continuationToken: null, - disallowContinuationTokenMessage: ClientDistinctDocumentQueryExecutionComponent.DisallowContinuationTokenMessage, - activityId: sourceResponse.ActivityId, - requestCharge: sourceResponse.RequestCharge, - responseLengthBytes: sourceResponse.ResponseLengthBytes); + queryPage = new QueryPage( + documents: distinctResults, + requestCharge: sourcePage.RequestCharge, + activityId: sourcePage.ActivityId, + responseLengthInBytes: sourcePage.ResponseLengthInBytes, + cosmosQueryExecutionInfo: sourcePage.CosmosQueryExecutionInfo, + disallowContinuationTokenMessage: ClientDistinctQueryPipelineStage.DisallowContinuationTokenMessage, + state: null); } - return queryResponseCore; - } - - public override CosmosElement GetCosmosElementContinuationToken() - { - throw new NotImplementedException(); + return TryCatch.FromResult(queryPage); } /// @@ -255,4 +248,4 @@ public override string ToString() } } } -} +} \ No newline at end of file diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctDocumentQueryExecutionComponent.Compute.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Distinct/DistinctQueryPipelineStage.Compute.cs similarity index 62% rename from Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctDocumentQueryExecutionComponent.Compute.cs rename to Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Distinct/DistinctQueryPipelineStage.Compute.cs index 89d931211c..b05fc86009 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctDocumentQueryExecutionComponent.Compute.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Distinct/DistinctQueryPipelineStage.Compute.cs @@ -1,40 +1,40 @@ //------------------------------------------------------------ // Copyright (c) Microsoft Corporation. All rights reserved. //------------------------------------------------------------ -namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.Distinct + +namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.Distinct { using System; using System.Collections.Generic; using System.Threading; using System.Threading.Tasks; using Microsoft.Azure.Cosmos.CosmosElements; - using Microsoft.Azure.Cosmos.Json; - using Microsoft.Azure.Cosmos.Query.Core.ContinuationTokens; using Microsoft.Azure.Cosmos.Query.Core.Exceptions; using Microsoft.Azure.Cosmos.Query.Core.Monads; - using Microsoft.Azure.Cosmos.Query.Core.QueryClient; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline; + using Newtonsoft.Json; - internal abstract partial class DistinctDocumentQueryExecutionComponent : DocumentQueryExecutionComponentBase + internal abstract partial class DistinctQueryPipelineStage : QueryPipelineStageBase { /// /// Compute implementation of DISTINCT. /// Here we never serialize the continuation token, but you can always retrieve it on demand with TryGetContinuationToken. /// - private sealed class ComputeDistinctDocumentQueryExecutionComponent : DistinctDocumentQueryExecutionComponent + private sealed class ComputeDistinctQueryPipelineStage : DistinctQueryPipelineStage { private static readonly string UseTryGetContinuationTokenMessage = $"Use TryGetContinuationToken"; - private ComputeDistinctDocumentQueryExecutionComponent( + private ComputeDistinctQueryPipelineStage( DistinctQueryType distinctQueryType, DistinctMap distinctMap, - IDocumentQueryExecutionComponent source) + IQueryPipelineStage source) : base(distinctMap, source) { } - public static async Task> TryCreateAsync( + public static async Task> TryCreateAsync( CosmosElement requestContinuation, - Func>> tryCreateSourceAsync, + Func>> tryCreateSourceAsync, DistinctQueryType distinctQueryType) { if (tryCreateSourceAsync == null) @@ -47,8 +47,9 @@ public static async Task> TryCreateAs { if (!DistinctContinuationToken.TryParse(requestContinuation, out distinctContinuationToken)) { - return TryCatch.FromException( - new MalformedContinuationTokenException($"Invalid {nameof(DistinctContinuationToken)}: {requestContinuation}")); + return TryCatch.FromException( + new MalformedContinuationTokenException( + $"Invalid {nameof(DistinctContinuationToken)}: {requestContinuation}")); } } else @@ -61,40 +62,36 @@ public static async Task> TryCreateAs distinctContinuationToken.DistinctMapToken); if (!tryCreateDistinctMap.Succeeded) { - return TryCatch.FromException(tryCreateDistinctMap.Exception); + return TryCatch.FromException(tryCreateDistinctMap.Exception); } - TryCatch tryCreateSource = await tryCreateSourceAsync( + TryCatch tryCreateSource = await tryCreateSourceAsync( distinctContinuationToken.SourceToken); if (!tryCreateSource.Succeeded) { - return TryCatch.FromException(tryCreateSource.Exception); + return TryCatch.FromException(tryCreateSource.Exception); } - return TryCatch.FromResult( - new ComputeDistinctDocumentQueryExecutionComponent( + return TryCatch.FromResult( + new ComputeDistinctQueryPipelineStage( distinctQueryType, tryCreateDistinctMap.Result, tryCreateSource.Result)); } - /// - /// Drains a page of results returning only distinct elements. - /// - /// The maximum number of items to drain. - /// The cancellation token. - /// A page of distinct results. - public override async Task DrainAsync(int maxElements, CancellationToken cancellationToken) + protected override async Task> GetNextPageAsync(CancellationToken cancellationToken) { - List distinctResults = new List(); - QueryResponseCore sourceResponse = await base.DrainAsync(maxElements, cancellationToken); - - if (!sourceResponse.IsSuccess) + await this.inputStage.MoveNextAsync(); + TryCatch tryGetSourcePage = this.inputStage.Current; + if (tryGetSourcePage.Failed) { - return sourceResponse; + return tryGetSourcePage; } - foreach (CosmosElement document in sourceResponse.CosmosElements) + QueryPage sourcePage = tryGetSourcePage.Result; + + List distinctResults = new List(); + foreach (CosmosElement document in sourcePage.Documents) { if (this.distinctMap.Add(document, out UInt128 hash)) { @@ -102,26 +99,29 @@ public override async Task DrainAsync(int maxElements, Cancel } } - return QueryResponseCore.CreateSuccess( - result: distinctResults, - continuationToken: null, - disallowContinuationTokenMessage: ComputeDistinctDocumentQueryExecutionComponent.UseTryGetContinuationTokenMessage, - activityId: sourceResponse.ActivityId, - requestCharge: sourceResponse.RequestCharge, - responseLengthBytes: sourceResponse.ResponseLengthBytes); - } - - public override CosmosElement GetCosmosElementContinuationToken() - { - if (this.IsDone) + QueryState queryState; + if (sourcePage.State != null) + { + DistinctContinuationToken distinctContinuationToken = new DistinctContinuationToken( + sourceToken: sourcePage.State.Value, + distinctMapToken: this.distinctMap.GetCosmosElementContinuationToken()); + queryState = new QueryState(DistinctContinuationToken.ToCosmosElement(distinctContinuationToken)); + } + else { - return default; + queryState = null; } - DistinctContinuationToken distinctContinuationToken = new DistinctContinuationToken( - sourceToken: this.Source.GetCosmosElementContinuationToken(), - distinctMapToken: this.distinctMap.GetCosmosElementContinuationToken()); - return DistinctContinuationToken.ToCosmosElement(distinctContinuationToken); + QueryPage queryPage = new QueryPage( + documents: distinctResults, + requestCharge: sourcePage.RequestCharge, + activityId: sourcePage.ActivityId, + responseLengthInBytes: sourcePage.ResponseLengthInBytes, + cosmosQueryExecutionInfo: sourcePage.CosmosQueryExecutionInfo, + disallowContinuationTokenMessage: ComputeDistinctQueryPipelineStage.UseTryGetContinuationTokenMessage, + state: queryState); + + return TryCatch.FromResult(queryPage); } private readonly struct DistinctContinuationToken @@ -190,4 +190,4 @@ public static bool TryParse( } } } -} +} \ No newline at end of file diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Distinct/DistinctQueryPipelineStage.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Distinct/DistinctQueryPipelineStage.cs new file mode 100644 index 0000000000..466f508ec6 --- /dev/null +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Distinct/DistinctQueryPipelineStage.cs @@ -0,0 +1,56 @@ +// ------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +// ------------------------------------------------------------ + +namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.Distinct +{ + using System; + using System.Threading.Tasks; + using Microsoft.Azure.Cosmos.CosmosElements; + using Microsoft.Azure.Cosmos.Query.Core.ExecutionContext; + using Microsoft.Azure.Cosmos.Query.Core.Monads; + + /// + /// Distinct queries return documents that are distinct with a page. + /// This means that documents are not guaranteed to be distinct across continuations and partitions. + /// The reasoning for this is because the backend treats each continuation of a query as a separate request + /// and partitions are not aware of each other. + /// The solution is that the client keeps a running hash set of all the documents it has already seen, + /// so that when it encounters a duplicate document from another continuation it will not be emitted to the user. + /// The only problem is that if the user chooses to go through the continuation token API for DocumentQuery instead + /// of while(HasMoreResults) ExecuteNextAsync, then will see duplicates across continuations. + /// There is no workaround for that use case, since the continuation token will have to include all the documents seen. + /// + internal abstract partial class DistinctQueryPipelineStage : QueryPipelineStageBase + { + /// + /// An DistinctMap that efficiently stores the documents that we have already seen. + /// + private readonly DistinctMap distinctMap; + + protected DistinctQueryPipelineStage( + DistinctMap distinctMap, + IQueryPipelineStage source) + : base(source) + { + this.distinctMap = distinctMap ?? throw new ArgumentNullException(nameof(distinctMap)); + } + + public static async Task> TryCreateAsync( + ExecutionEnvironment executionEnvironment, + CosmosElement requestContinuation, + Func>> tryCreateSourceAsync, + DistinctQueryType distinctQueryType) => executionEnvironment switch + { + ExecutionEnvironment.Client => await ClientDistinctQueryPipelineStage.TryCreateAsync( + requestContinuation, + tryCreateSourceAsync, + distinctQueryType), + ExecutionEnvironment.Compute => await ComputeDistinctQueryPipelineStage.TryCreateAsync( + requestContinuation, + tryCreateSourceAsync, + distinctQueryType), + _ => throw new ArgumentException($"Unknown {nameof(ExecutionEnvironment)}: {executionEnvironment}."), + }; + } +} diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctQueryType.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Distinct/DistinctQueryType.cs similarity index 93% rename from Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctQueryType.cs rename to Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Distinct/DistinctQueryType.cs index 02c9605092..e66d61d32b 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctQueryType.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Distinct/DistinctQueryType.cs @@ -1,7 +1,8 @@ //------------------------------------------------------------ // Copyright (c) Microsoft Corporation. All rights reserved. //------------------------------------------------------------ -namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.Distinct + +namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.Distinct { /// /// Enum of the type of distinct queries. diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/QueryPlan/QueryInfo.cs b/Microsoft.Azure.Cosmos/src/Query/Core/QueryPlan/QueryInfo.cs index 9bd52eb443..7ff6ca4d10 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/QueryPlan/QueryInfo.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/QueryPlan/QueryInfo.cs @@ -6,9 +6,9 @@ namespace Microsoft.Azure.Cosmos.Query.Core.QueryPlan { using System.Collections.Generic; using System.Linq; - using Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.Distinct; using Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.OrderBy; using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Aggregate; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Distinct; using Newtonsoft.Json; using Newtonsoft.Json.Converters; diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/CosmosQueryUnitTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/CosmosQueryUnitTests.cs index 4207a3da4d..10e813239e 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/CosmosQueryUnitTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/CosmosQueryUnitTests.cs @@ -24,6 +24,7 @@ namespace Microsoft.Azure.Cosmos.Tests using Microsoft.Azure.Cosmos.Query.Core.Metrics; using Microsoft.Azure.Cosmos.Query.Core.Monads; using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Aggregate; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Distinct; using Microsoft.Azure.Cosmos.Query.Core.QueryClient; using Microsoft.Azure.Cosmos.Query.Core.QueryPlan; using Microsoft.Azure.Cosmos.Resource.CosmosExceptions; diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/DistinctHashBaselineTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/DistinctHashBaselineTests.cs index 78ead5927d..cf8391e865 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/DistinctHashBaselineTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/DistinctHashBaselineTests.cs @@ -6,8 +6,7 @@ using System.Xml; using Microsoft.Azure.Cosmos.CosmosElements; using Microsoft.Azure.Cosmos.CosmosElements.Numbers; - using Microsoft.Azure.Cosmos.Query.Core; - using Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.Distinct; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Distinct; using Microsoft.Azure.Cosmos.Test.BaselineTest; using Microsoft.VisualStudio.TestTools.UnitTesting; diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/DistinctQueryPipelineStageTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/DistinctQueryPipelineStageTests.cs new file mode 100644 index 0000000000..6479b4f638 --- /dev/null +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/DistinctQueryPipelineStageTests.cs @@ -0,0 +1,65 @@ +namespace Microsoft.Azure.Cosmos.Tests.Query.Pipeline +{ + using System; + using System.Collections.Generic; + using System.Linq; + using System.Threading.Tasks; + using Microsoft.Azure.Cosmos.CosmosElements; + 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.Distinct; + using Microsoft.VisualStudio.TestTools.UnitTesting; + + [TestClass] + public sealed class DistinctQueryPipelineStageTests + { + [TestMethod] + public async Task SanityTests() + { + long[] values = new long[] { 42, 1337, 1337, 42 }; + IReadOnlyList> pages = values + .Select(value => new List() + { + CosmosElement.Parse($"{{\"item\": {value}}}") + }) + .ToList(); + + List elements = await DistinctQueryPipelineStageTests.CreateAndDrainAsync( + pages: pages, + executionEnvironment: ExecutionEnvironment.Compute, + continuationToken: null, + distinctQueryType: DistinctQueryType.Unordered); + + Assert.AreEqual(values.Distinct().Count(), elements.Count); + } + + private static async Task> CreateAndDrainAsync( + IReadOnlyList> pages, + ExecutionEnvironment executionEnvironment, + CosmosElement continuationToken, + DistinctQueryType distinctQueryType) + { + IQueryPipelineStage source = new MockQueryPipelineStage(pages); + + TryCatch tryCreateDistinctQueryPipelineStage = await DistinctQueryPipelineStage.TryCreateAsync( + executionEnvironment: executionEnvironment, + requestContinuation: continuationToken, + distinctQueryType: distinctQueryType, + tryCreateSourceAsync: (CosmosElement continuationToken) => Task.FromResult(TryCatch.FromResult(source))); + Assert.IsTrue(tryCreateDistinctQueryPipelineStage.Succeeded); + + IQueryPipelineStage distinctQueryPipelineStage = tryCreateDistinctQueryPipelineStage.Result; + + List elements = new List(); + await foreach (TryCatch page in new EnumerableStage(distinctQueryPipelineStage)) + { + page.ThrowIfFailed(); + + elements.AddRange(page.Result.Documents); + } + + return elements; + } + } +} From ca4b86fb00f70ffa2004a2d4eab298ef440fc70e Mon Sep 17 00:00:00 2001 From: Brandon Chong Date: Fri, 3 Jul 2020 01:33:46 -0700 Subject: [PATCH 26/85] got group by working --- ...yDocumentQueryExecutionComponent.Client.cs | 104 ------- .../GroupByDocumentQueryExecutionComponent.cs | 34 +- .../AggregateQueryPipelineStage.Client.cs | 3 +- .../AggregateQueryPipelineStage.Compute.cs | 74 ++--- .../Aggregate/AggregateQueryPipelineStage.cs | 18 +- .../Distinct/DistinctQueryPipelineStage.cs | 6 +- .../GroupByQueryPipelineStage.Client.cs | 108 +++++++ .../GroupByQueryPipelineStage.Compute.cs} | 174 ++++------- .../GroupBy/GroupByQueryPipelineStage.cs | 290 ++++++++++++++++++ .../Query/DistinctHashBenchmark.cs | 2 +- .../GroupByQueryPipelineStageTests.cs | 77 +++++ 11 files changed, 591 insertions(+), 299 deletions(-) delete mode 100644 Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/GroupBy/GroupByDocumentQueryExecutionComponent.Client.cs create mode 100644 Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/GroupBy/GroupByQueryPipelineStage.Client.cs rename Microsoft.Azure.Cosmos/src/Query/Core/{ExecutionComponent/GroupBy/GroupByDocumentQueryExecutionComponent.Compute.cs => Pipeline/GroupBy/GroupByQueryPipelineStage.Compute.cs} (55%) create mode 100644 Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/GroupBy/GroupByQueryPipelineStage.cs create mode 100644 Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/GroupByQueryPipelineStageTests.cs diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/GroupBy/GroupByDocumentQueryExecutionComponent.Client.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/GroupBy/GroupByDocumentQueryExecutionComponent.Client.cs deleted file mode 100644 index b83c90c257..0000000000 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/GroupBy/GroupByDocumentQueryExecutionComponent.Client.cs +++ /dev/null @@ -1,104 +0,0 @@ -//------------------------------------------------------------ -// Copyright (c) Microsoft Corporation. All rights reserved. -//------------------------------------------------------------ -namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.GroupBy -{ - using System; - using System.Collections.Generic; - using System.Threading; - using System.Threading.Tasks; - using Microsoft.Azure.Cosmos.CosmosElements; - using Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.Aggregate; - using Microsoft.Azure.Cosmos.Query.Core.Monads; - using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Aggregate; - using Microsoft.Azure.Cosmos.Query.Core.QueryClient; - - internal abstract partial class GroupByDocumentQueryExecutionComponent : DocumentQueryExecutionComponentBase - { - private sealed class ClientGroupByDocumentQueryExecutionComponent : GroupByDocumentQueryExecutionComponent - { - public const string ContinuationTokenNotSupportedWithGroupBy = "Continuation token is not supported for queries with GROUP BY. Do not use FeedResponse.ResponseContinuation or remove the GROUP BY from the query."; - - private ClientGroupByDocumentQueryExecutionComponent( - IDocumentQueryExecutionComponent source, - GroupingTable groupingTable) - : base( - source, - groupingTable) - { - } - - public static async Task> TryCreateAsync( - CosmosElement requestContinuation, - Func>> tryCreateSource, - IReadOnlyDictionary groupByAliasToAggregateType, - IReadOnlyList orderedAliases, - bool hasSelectValue) - { - TryCatch tryCreateGroupingTable = GroupingTable.TryCreateFromContinuationToken( - groupByAliasToAggregateType, - orderedAliases, - hasSelectValue, - continuationToken: null); - - if (!tryCreateGroupingTable.Succeeded) - { - return TryCatch.FromException(tryCreateGroupingTable.Exception); - } - - return (await tryCreateSource(requestContinuation)).Try(source => - { - return new ClientGroupByDocumentQueryExecutionComponent( - source, - tryCreateGroupingTable.Result); - }); - } - - public override async Task DrainAsync( - int maxElements, - CancellationToken cancellationToken) - { - cancellationToken.ThrowIfCancellationRequested(); - - // Draining GROUP BY is broken down into two stages: - - double requestCharge = 0.0; - long responseLengthBytes = 0; - while (!this.Source.IsDone) - { - // Stage 1: - // Drain the groupings fully from all continuation and all partitions - QueryResponseCore sourceResponse = await base.DrainAsync(int.MaxValue, cancellationToken); - if (!sourceResponse.IsSuccess) - { - return sourceResponse; - } - - requestCharge += sourceResponse.RequestCharge; - responseLengthBytes += sourceResponse.ResponseLengthBytes; - - this.AggregateGroupings(sourceResponse.CosmosElements); - } - - // Stage 2: - // Emit the results from the grouping table page by page - IReadOnlyList results = this.groupingTable.Drain(maxElements); - - QueryResponseCore response = QueryResponseCore.CreateSuccess( - result: results, - continuationToken: null, - disallowContinuationTokenMessage: ClientGroupByDocumentQueryExecutionComponent.ContinuationTokenNotSupportedWithGroupBy, - activityId: null, - requestCharge: requestCharge, - responseLengthBytes: responseLengthBytes); - - return response; - } - - public override CosmosElement GetCosmosElementContinuationToken() - { - throw new NotImplementedException(); - } - } - } -} \ No newline at end of file diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/GroupBy/GroupByDocumentQueryExecutionComponent.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/GroupBy/GroupByDocumentQueryExecutionComponent.cs index 36707da09a..8b92b6030a 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/GroupBy/GroupByDocumentQueryExecutionComponent.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/GroupBy/GroupByDocumentQueryExecutionComponent.cs @@ -57,7 +57,7 @@ protected GroupByDocumentQueryExecutionComponent( public override bool IsDone => this.groupingTable.IsDone; - public static async Task> TryCreateAsync( + public static Task> TryCreateAsync( ExecutionEnvironment executionEnvironment, CosmosElement continuationToken, Func>> tryCreateSourceAsync, @@ -65,37 +65,7 @@ public static async Task> TryCreateAs IReadOnlyList orderedAliases, bool hasSelectValue) { - if (tryCreateSourceAsync == null) - { - throw new ArgumentNullException(nameof(tryCreateSourceAsync)); - } - - TryCatch tryCreateGroupByComponent; - switch (executionEnvironment) - { - case ExecutionEnvironment.Client: - tryCreateGroupByComponent = await ClientGroupByDocumentQueryExecutionComponent.TryCreateAsync( - continuationToken, - tryCreateSourceAsync, - groupByAliasToAggregateType, - orderedAliases, - hasSelectValue); - break; - - case ExecutionEnvironment.Compute: - tryCreateGroupByComponent = await ComputeGroupByDocumentQueryExecutionComponent.TryCreateAsync( - continuationToken, - tryCreateSourceAsync, - groupByAliasToAggregateType, - orderedAliases, - hasSelectValue); - break; - - default: - throw new ArgumentException($"Unknown {nameof(ExecutionEnvironment)}: {executionEnvironment}"); - } - - return tryCreateGroupByComponent; + return default; } protected void AggregateGroupings(IReadOnlyList cosmosElements) diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Aggregate/AggregateQueryPipelineStage.Client.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Aggregate/AggregateQueryPipelineStage.Client.cs index fa8ff94db9..21a71e7070 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Aggregate/AggregateQueryPipelineStage.Client.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Aggregate/AggregateQueryPipelineStage.Client.cs @@ -73,9 +73,8 @@ protected override async Task> GetNextPageAsync(Cancellation double requestCharge = 0; long responseLengthBytes = 0; - while (this.inputStage.HasMoreResults) + while (await this.inputStage.MoveNextAsync()) { - await this.inputStage.MoveNextAsync(); TryCatch tryGetPageFromSource = this.inputStage.Current; if (tryGetPageFromSource.Failed) diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Aggregate/AggregateQueryPipelineStage.Compute.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Aggregate/AggregateQueryPipelineStage.Compute.cs index 1aa6ab0e67..8bdacfdd4d 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Aggregate/AggregateQueryPipelineStage.Compute.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Aggregate/AggregateQueryPipelineStage.Compute.cs @@ -94,7 +94,43 @@ protected override async Task> GetNextPageAsync(Cancellation cancellationToken.ThrowIfCancellationRequested(); // Draining aggregates is broken down into two stages - if (!this.inputStage.HasMoreResults) + QueryPage queryPage; + if (await this.inputStage.MoveNextAsync()) + { + // Stage 1: + // Drain the aggregates fully from all continuations and all partitions + // And return empty pages in the meantime. + TryCatch tryGetSourcePage = this.inputStage.Current; + if (tryGetSourcePage.Failed) + { + return tryGetSourcePage; + } + + QueryPage sourcePage = tryGetSourcePage.Result; + foreach (CosmosElement element in sourcePage.Documents) + { + RewrittenAggregateProjections rewrittenAggregateProjections = new RewrittenAggregateProjections( + this.isValueQuery, + element); + this.singleGroupAggregator.AddValues(rewrittenAggregateProjections.Payload); + } + + AggregateContinuationToken aggregateContinuationToken = new AggregateContinuationToken( + singleGroupAggregatorContinuationToken: this.singleGroupAggregator.GetCosmosElementContinuationToken(), + sourceContinuationToken: sourcePage.State != null ? sourcePage.State.Value : DoneSourceToken); + QueryState queryState = new QueryState(AggregateContinuationToken.ToCosmosElement(aggregateContinuationToken)); + QueryPage emptyPage = new QueryPage( + documents: EmptyResults, + requestCharge: sourcePage.RequestCharge, + activityId: sourcePage.ActivityId, + responseLengthInBytes: sourcePage.ResponseLengthInBytes, + cosmosQueryExecutionInfo: sourcePage.CosmosQueryExecutionInfo, + disallowContinuationTokenMessage: sourcePage.DisallowContinuationTokenMessage, + state: queryState); + + queryPage = emptyPage; + } + else { // Stage 2: // Return the final page after draining. @@ -114,42 +150,10 @@ protected override async Task> GetNextPageAsync(Cancellation disallowContinuationTokenMessage: default, state: default); - return TryCatch.FromResult(finalPage); - } - - // Stage 1: - // Drain the aggregates fully from all continuations and all partitions - // And return empty pages in the meantime. - await this.inputStage.MoveNextAsync(); - TryCatch tryGetSourcePage = this.inputStage.Current; - if (tryGetSourcePage.Failed) - { - return tryGetSourcePage; - } - - QueryPage sourcePage = tryGetSourcePage.Result; - foreach (CosmosElement element in sourcePage.Documents) - { - RewrittenAggregateProjections rewrittenAggregateProjections = new RewrittenAggregateProjections( - this.isValueQuery, - element); - this.singleGroupAggregator.AddValues(rewrittenAggregateProjections.Payload); + queryPage = finalPage; } - AggregateContinuationToken aggregateContinuationToken = new AggregateContinuationToken( - singleGroupAggregatorContinuationToken: this.singleGroupAggregator.GetCosmosElementContinuationToken(), - sourceContinuationToken: sourcePage.State != null ? sourcePage.State.Value : DoneSourceToken); - QueryState queryState = new QueryState(AggregateContinuationToken.ToCosmosElement(aggregateContinuationToken)); - QueryPage emptyPage = new QueryPage( - documents: EmptyResults, - requestCharge: sourcePage.RequestCharge, - activityId: sourcePage.ActivityId, - responseLengthInBytes: sourcePage.ResponseLengthInBytes, - cosmosQueryExecutionInfo: sourcePage.CosmosQueryExecutionInfo, - disallowContinuationTokenMessage: sourcePage.DisallowContinuationTokenMessage, - state: queryState); - - return TryCatch.FromResult(emptyPage); + return TryCatch.FromResult(queryPage); } private sealed class AggregateContinuationToken diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Aggregate/AggregateQueryPipelineStage.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Aggregate/AggregateQueryPipelineStage.cs index 273d44bcb0..c2216fb4ca 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Aggregate/AggregateQueryPipelineStage.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Aggregate/AggregateQueryPipelineStage.cs @@ -56,30 +56,23 @@ public AggregateQueryPipelineStage( this.isValueQuery = isValueQuery; } - public static async Task> TryCreateAsync( + public static Task> TryCreateAsync( ExecutionEnvironment executionEnvironment, IReadOnlyList aggregates, IReadOnlyDictionary aliasToAggregateType, IReadOnlyList orderedAliases, bool hasSelectValue, CosmosElement continuationToken, - Func>> tryCreateSourceAsync) - { - if (tryCreateSourceAsync == null) - { - throw new ArgumentNullException(nameof(tryCreateSourceAsync)); - } - - TryCatch tryCreateAggregate = executionEnvironment switch + Func>> tryCreateSourceAsync) => executionEnvironment switch { - ExecutionEnvironment.Client => await ClientAggregateQueryPipelineStage.TryCreateAsync( + ExecutionEnvironment.Client => ClientAggregateQueryPipelineStage.TryCreateAsync( aggregates, aliasToAggregateType, orderedAliases, hasSelectValue, continuationToken, tryCreateSourceAsync), - ExecutionEnvironment.Compute => await ComputeAggregateQueryPipelineStage.TryCreateAsync( + ExecutionEnvironment.Compute => ComputeAggregateQueryPipelineStage.TryCreateAsync( aggregates, aliasToAggregateType, orderedAliases, @@ -89,9 +82,6 @@ public static async Task> TryCreateAsync( _ => throw new ArgumentException($"Unknown {nameof(ExecutionEnvironment)}: {executionEnvironment}."), }; - return tryCreateAggregate; - } - /// /// Struct for getting the payload out of the rewritten projection. /// diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Distinct/DistinctQueryPipelineStage.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Distinct/DistinctQueryPipelineStage.cs index 466f508ec6..9b122ac281 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Distinct/DistinctQueryPipelineStage.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Distinct/DistinctQueryPipelineStage.cs @@ -36,17 +36,17 @@ protected DistinctQueryPipelineStage( this.distinctMap = distinctMap ?? throw new ArgumentNullException(nameof(distinctMap)); } - public static async Task> TryCreateAsync( + public static Task> TryCreateAsync( ExecutionEnvironment executionEnvironment, CosmosElement requestContinuation, Func>> tryCreateSourceAsync, DistinctQueryType distinctQueryType) => executionEnvironment switch { - ExecutionEnvironment.Client => await ClientDistinctQueryPipelineStage.TryCreateAsync( + ExecutionEnvironment.Client => ClientDistinctQueryPipelineStage.TryCreateAsync( requestContinuation, tryCreateSourceAsync, distinctQueryType), - ExecutionEnvironment.Compute => await ComputeDistinctQueryPipelineStage.TryCreateAsync( + ExecutionEnvironment.Compute => ComputeDistinctQueryPipelineStage.TryCreateAsync( requestContinuation, tryCreateSourceAsync, distinctQueryType), diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/GroupBy/GroupByQueryPipelineStage.Client.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/GroupBy/GroupByQueryPipelineStage.Client.cs new file mode 100644 index 0000000000..7a0ee1c0e8 --- /dev/null +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/GroupBy/GroupByQueryPipelineStage.Client.cs @@ -0,0 +1,108 @@ +// ------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +// ------------------------------------------------------------ + +namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.GroupBy +{ + using System; + using System.Collections.Generic; + using System.Threading; + using System.Threading.Tasks; + using Microsoft.Azure.Cosmos.CosmosElements; + using Microsoft.Azure.Cosmos.Query.Core.Monads; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Aggregate; + + internal abstract partial class GroupByQueryPipelineStage : QueryPipelineStageBase + { + private sealed class ClientGroupByQueryPipelineStage : GroupByQueryPipelineStage + { + public const string ContinuationTokenNotSupportedWithGroupBy = "Continuation token is not supported for queries with GROUP BY. Do not use FeedResponse.ResponseContinuation or remove the GROUP BY from the query."; + + private ClientGroupByQueryPipelineStage( + IQueryPipelineStage source, + GroupingTable groupingTable) + : base( + source, + groupingTable) + { + } + + public static async Task> TryCreateAsync( + CosmosElement requestContinuation, + Func>> tryCreateSourceAsync, + IReadOnlyDictionary groupByAliasToAggregateType, + IReadOnlyList orderedAliases, + bool hasSelectValue) + { + TryCatch tryCreateGroupingTable = GroupingTable.TryCreateFromContinuationToken( + groupByAliasToAggregateType, + orderedAliases, + hasSelectValue, + continuationToken: null); + + if (tryCreateGroupingTable.Failed) + { + return TryCatch.FromException(tryCreateGroupingTable.Exception); + } + + TryCatch tryCreateSource = await tryCreateSourceAsync(requestContinuation); + if (tryCreateSource.Failed) + { + return tryCreateSource; + } + + IQueryPipelineStage stage = new ClientGroupByQueryPipelineStage( + tryCreateSource.Result, + tryCreateGroupingTable.Result); + + return TryCatch.FromResult(stage); + } + + protected override async Task> GetNextPageAsync(CancellationToken cancellationToken) + { + cancellationToken.ThrowIfCancellationRequested(); + + // Draining GROUP BY is broken down into two stages: + + double requestCharge = 0.0; + long responseLengthInBytes = 0; + + int maxPageSize = 0; + + while (await this.inputStage.MoveNextAsync()) + { + // Stage 1: + // Drain the groupings fully from all continuation and all partitions + TryCatch tryGetSourcePage = this.inputStage.Current; + if (tryGetSourcePage.Failed) + { + return tryGetSourcePage; + } + + QueryPage sourcePage = tryGetSourcePage.Result; + + requestCharge += sourcePage.RequestCharge; + responseLengthInBytes += sourcePage.ResponseLengthInBytes; + maxPageSize = Math.Max(sourcePage.Documents.Count, maxPageSize); + + this.AggregateGroupings(sourcePage.Documents); + } + + // Stage 2: + // Emit the results from the grouping table page by page + IReadOnlyList results = this.groupingTable.Drain(maxPageSize); + + QueryPage queryPage = new QueryPage( + documents: results, + requestCharge: requestCharge, + activityId: null, + responseLengthInBytes: responseLengthInBytes, + cosmosQueryExecutionInfo: null, + disallowContinuationTokenMessage: ClientGroupByQueryPipelineStage.ContinuationTokenNotSupportedWithGroupBy, + state: null); + + return TryCatch.FromResult(queryPage); + } + } + } +} diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/GroupBy/GroupByDocumentQueryExecutionComponent.Compute.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/GroupBy/GroupByQueryPipelineStage.Compute.cs similarity index 55% rename from Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/GroupBy/GroupByDocumentQueryExecutionComponent.Compute.cs rename to Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/GroupBy/GroupByQueryPipelineStage.Compute.cs index b7b76112a4..45a523ed07 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/GroupBy/GroupByDocumentQueryExecutionComponent.Compute.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/GroupBy/GroupByQueryPipelineStage.Compute.cs @@ -1,26 +1,22 @@ -//------------------------------------------------------------ +// ------------------------------------------------------------ // Copyright (c) Microsoft Corporation. All rights reserved. -//------------------------------------------------------------ -namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.GroupBy +// ------------------------------------------------------------ + +namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.GroupBy { using System; using System.Collections.Generic; - using System.Runtime.CompilerServices; using System.Threading; using System.Threading.Tasks; using Microsoft.Azure.Cosmos.CosmosElements; - using Microsoft.Azure.Cosmos.Json; - using Microsoft.Azure.Cosmos.Query.Core.ContinuationTokens; using Microsoft.Azure.Cosmos.Query.Core.Exceptions; - using Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.Aggregate; using Microsoft.Azure.Cosmos.Query.Core.Metrics; using Microsoft.Azure.Cosmos.Query.Core.Monads; using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Aggregate; - using Microsoft.Azure.Cosmos.Query.Core.QueryClient; - internal abstract partial class GroupByDocumentQueryExecutionComponent : DocumentQueryExecutionComponentBase + internal abstract partial class GroupByQueryPipelineStage : QueryPipelineStageBase { - private sealed class ComputeGroupByDocumentQueryExecutionComponent : GroupByDocumentQueryExecutionComponent + private sealed class ComputeGroupByQueryPipelineStage : GroupByQueryPipelineStage { private const string DoneReadingGroupingsContinuationToken = "DONE"; private static readonly CosmosElement DoneCosmosElementToken = CosmosString.Create(DoneReadingGroupingsContinuationToken); @@ -28,8 +24,8 @@ private sealed class ComputeGroupByDocumentQueryExecutionComponent : GroupByDocu private static readonly IReadOnlyList EmptyResults = new List().AsReadOnly(); private static readonly IReadOnlyDictionary EmptyQueryMetrics = new Dictionary(); - private ComputeGroupByDocumentQueryExecutionComponent( - IDocumentQueryExecutionComponent source, + private ComputeGroupByQueryPipelineStage( + IQueryPipelineStage source, GroupingTable groupingTable) : base( source, @@ -37,9 +33,9 @@ private ComputeGroupByDocumentQueryExecutionComponent( { } - public static async Task> TryCreateAsync( + public static async Task> TryCreateAsync( CosmosElement requestContinuation, - Func>> tryCreateSourceAsync, + Func>> tryCreateSourceAsync, IReadOnlyDictionary groupByAliasToAggregateType, IReadOnlyList orderedAliases, bool hasSelectValue) @@ -49,8 +45,9 @@ public static async Task> TryCreateAs { if (!GroupByContinuationToken.TryParse(requestContinuation, out groupByContinuationToken)) { - return TryCatch.FromException( - new MalformedContinuationTokenException($"Invalid {nameof(GroupByContinuationToken)}: '{requestContinuation}'")); + return TryCatch.FromException( + new MalformedContinuationTokenException( + $"Invalid {nameof(GroupByContinuationToken)}: '{requestContinuation}'")); } } else @@ -60,11 +57,11 @@ public static async Task> TryCreateAs sourceContinuationToken: null); } - TryCatch tryCreateSource; + TryCatch tryCreateSource; if ((groupByContinuationToken.SourceContinuationToken is CosmosString sourceContinuationToken) - && (sourceContinuationToken.Value == ComputeGroupByDocumentQueryExecutionComponent.DoneReadingGroupingsContinuationToken)) + && (sourceContinuationToken.Value == ComputeGroupByQueryPipelineStage.DoneReadingGroupingsContinuationToken)) { - tryCreateSource = TryCatch.FromResult(DoneDocumentQueryExecutionComponent.Value); + tryCreateSource = TryCatch.FromResult(FinishedQueryPipelineStage.Value); } else { @@ -73,7 +70,7 @@ public static async Task> TryCreateAs if (!tryCreateSource.Succeeded) { - return TryCatch.FromException(tryCreateSource.Exception); + return TryCatch.FromException(tryCreateSource.Exception); } TryCatch tryCreateGroupingTable = GroupingTable.TryCreateFromContinuationToken( @@ -84,84 +81,81 @@ public static async Task> TryCreateAs if (!tryCreateGroupingTable.Succeeded) { - return TryCatch.FromException(tryCreateGroupingTable.Exception); + return TryCatch.FromException(tryCreateGroupingTable.Exception); } - return TryCatch.FromResult( - new ComputeGroupByDocumentQueryExecutionComponent( + return TryCatch.FromResult( + new ComputeGroupByQueryPipelineStage( tryCreateSource.Result, tryCreateGroupingTable.Result)); } - public override async Task DrainAsync( - int maxElements, - CancellationToken cancellationToken) + protected override async Task> GetNextPageAsync(CancellationToken cancellationToken) { cancellationToken.ThrowIfCancellationRequested(); // Draining GROUP BY is broken down into two stages: - QueryResponseCore response; - if (!this.Source.IsDone) + QueryPage queryPage; + if (await this.inputStage.MoveNextAsync()) { // Stage 1: // Drain the groupings fully from all continuation and all partitions - QueryResponseCore sourceResponse = await base.DrainAsync(int.MaxValue, cancellationToken); - if (!sourceResponse.IsSuccess) + TryCatch tryGetSourcePage = this.inputStage.Current; + if (tryGetSourcePage.Failed) { - return sourceResponse; + return tryGetSourcePage; } - this.AggregateGroupings(sourceResponse.CosmosElements); + QueryPage sourcePage = tryGetSourcePage.Result; + + this.AggregateGroupings(sourcePage.Documents); // We need to give empty pages until the results are fully drained. - response = QueryResponseCore.CreateSuccess( - result: EmptyResults, - continuationToken: null, - disallowContinuationTokenMessage: DocumentQueryExecutionComponentBase.UseCosmosElementContinuationTokenInstead, - activityId: sourceResponse.ActivityId, - requestCharge: sourceResponse.RequestCharge, - responseLengthBytes: sourceResponse.ResponseLengthBytes); + CosmosElement sourceContinuationToken = sourcePage.State == null ? DoneCosmosElementToken : sourcePage.State.Value; + GroupByContinuationToken groupByContinuationToken = new GroupByContinuationToken( + groupingTableContinuationToken: this.groupingTable.GetCosmosElementContinuationToken(), + sourceContinuationToken: sourceContinuationToken); + QueryState state = new QueryState(GroupByContinuationToken.ToCosmosElement(groupByContinuationToken)); + + queryPage = new QueryPage( + documents: EmptyResults, + requestCharge: sourcePage.RequestCharge, + activityId: sourcePage.ActivityId, + responseLengthInBytes: sourcePage.ResponseLengthInBytes, + cosmosQueryExecutionInfo: sourcePage.CosmosQueryExecutionInfo, + disallowContinuationTokenMessage: null, + state: state); } else { // Stage 2: // Emit the results from the grouping table page by page - IReadOnlyList results = this.groupingTable.Drain(maxElements); - - response = QueryResponseCore.CreateSuccess( - result: results, - continuationToken: null, - disallowContinuationTokenMessage: DocumentQueryExecutionComponentBase.UseCosmosElementContinuationTokenInstead, - activityId: null, - requestCharge: 0, - responseLengthBytes: 0); - } + IReadOnlyList results = this.groupingTable.Drain(10/*FIX THIS*/); - return response; - } - - public override CosmosElement GetCosmosElementContinuationToken() - { - if (this.IsDone) - { - return default; - } + QueryState state; + if (this.groupingTable.IsDone) + { + state = default; + } + else + { + GroupByContinuationToken groupByContinuationToken = new GroupByContinuationToken( + groupingTableContinuationToken: this.groupingTable.GetCosmosElementContinuationToken(), + sourceContinuationToken: DoneCosmosElementToken); + state = new QueryState(GroupByContinuationToken.ToCosmosElement(groupByContinuationToken)); + } - CosmosElement sourceContinuationToken; - if (this.Source.IsDone) - { - sourceContinuationToken = DoneCosmosElementToken; - } - else - { - sourceContinuationToken = this.Source.GetCosmosElementContinuationToken(); + queryPage = new QueryPage( + documents: results, + requestCharge: default, + activityId: default, + responseLengthInBytes: default, + cosmosQueryExecutionInfo: default, + disallowContinuationTokenMessage: default, + state: state); } - GroupByContinuationToken groupByContinuationToken = new GroupByContinuationToken( - groupingTableContinuationToken: this.groupingTable.GetCosmosElementContinuationToken(), - sourceContinuationToken: sourceContinuationToken); - - return GroupByContinuationToken.ToCosmosElement(groupByContinuationToken); + return TryCatch.FromResult(queryPage); } private readonly struct GroupByContinuationToken @@ -231,42 +225,6 @@ public static bool TryParse(CosmosElement value, out GroupByContinuationToken gr return true; } } - - private sealed class DoneDocumentQueryExecutionComponent : IDocumentQueryExecutionComponent - { - public static readonly DoneDocumentQueryExecutionComponent Value = new DoneDocumentQueryExecutionComponent(); - - private DoneDocumentQueryExecutionComponent() - { - } - - public bool IsDone => true; - - public void Dispose() - { - // Do Nothing - } - - public Task DrainAsync(int maxElements, CancellationToken token) - { - token.ThrowIfCancellationRequested(); - throw new NotImplementedException(); - } - - public CosmosElement GetCosmosElementContinuationToken() - { - throw new NotImplementedException(); - } - - public IReadOnlyDictionary GetQueryMetrics() - { - throw new NotImplementedException(); - } - - public void Stop() - { - } - } } } -} \ No newline at end of file +} diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/GroupBy/GroupByQueryPipelineStage.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/GroupBy/GroupByQueryPipelineStage.cs new file mode 100644 index 0000000000..e07b63372b --- /dev/null +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/GroupBy/GroupByQueryPipelineStage.cs @@ -0,0 +1,290 @@ +// ------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +// ------------------------------------------------------------ + +namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.GroupBy +{ + using System; + using System.Collections; + using System.Collections.Generic; + using System.Linq; + using System.Threading.Tasks; + using Microsoft.Azure.Cosmos.CosmosElements; + using Microsoft.Azure.Cosmos.Query.Core.Exceptions; + using Microsoft.Azure.Cosmos.Query.Core.ExecutionContext; + using Microsoft.Azure.Cosmos.Query.Core.Monads; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Aggregate; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Aggregate.Aggregators; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Distinct; + + /// + /// Query execution component that groups groupings across continuations and pages. + /// The general idea is a query gets rewritten from this: + /// + /// SELECT c.team, c.name, COUNT(1) AS count + /// FROM c + /// GROUP BY c.team, c.name + /// + /// To this: + /// + /// SELECT + /// [{"item": c.team}, {"item": c.name}] AS groupByItems, + /// {"team": c.team, "name": c.name, "count": {"item": COUNT(1)}} AS payload + /// FROM c + /// GROUP BY c.team, c.name + /// + /// With the following dictionary: + /// + /// { + /// "team": null, + /// "name": null, + /// "count" COUNT + /// } + /// + /// So we know how to aggregate each column. + /// At the end the columns are stitched together to make the grouped document. + /// + internal abstract partial class GroupByQueryPipelineStage : QueryPipelineStageBase + { + private readonly GroupingTable groupingTable; + + protected GroupByQueryPipelineStage( + IQueryPipelineStage source, + GroupingTable groupingTable) + : base(source) + { + this.groupingTable = groupingTable ?? throw new ArgumentNullException(nameof(groupingTable)); + } + + public static Task> TryCreateAsync( + ExecutionEnvironment executionEnvironment, + CosmosElement continuationToken, + Func>> tryCreateSourceAsync, + IReadOnlyDictionary groupByAliasToAggregateType, + IReadOnlyList orderedAliases, + bool hasSelectValue) => executionEnvironment switch + { + ExecutionEnvironment.Client => ClientGroupByQueryPipelineStage.TryCreateAsync( + continuationToken, + tryCreateSourceAsync, + groupByAliasToAggregateType, + orderedAliases, + hasSelectValue), + ExecutionEnvironment.Compute => ComputeGroupByQueryPipelineStage.TryCreateAsync( + continuationToken, + tryCreateSourceAsync, + groupByAliasToAggregateType, + orderedAliases, + hasSelectValue), + _ => throw new ArgumentException($"Unknown {nameof(ExecutionEnvironment)}: {executionEnvironment}"), + }; + + protected void AggregateGroupings(IReadOnlyList cosmosElements) + { + foreach (CosmosElement result in cosmosElements) + { + // Aggregate the values for all groupings across all continuations. + RewrittenGroupByProjection groupByItem = new RewrittenGroupByProjection(result); + this.groupingTable.AddPayload(groupByItem); + } + } + + /// + /// When a group by query gets rewritten the projection looks like: + /// + /// SELECT + /// [{"item": c.age}, {"item": c.name}] AS groupByItems, + /// {"age": c.age, "name": c.name} AS payload + /// + /// This struct just lets us easily access the "groupByItems" and "payload" property. + /// + protected readonly struct RewrittenGroupByProjection + { + private const string GroupByItemsPropertyName = "groupByItems"; + private const string PayloadPropertyName = "payload"; + + private readonly CosmosObject cosmosObject; + + public RewrittenGroupByProjection(CosmosElement cosmosElement) + { + if (cosmosElement == null) + { + throw new ArgumentNullException(nameof(cosmosElement)); + } + + if (!(cosmosElement is CosmosObject cosmosObject)) + { + throw new ArgumentException($"{nameof(cosmosElement)} must not be an object."); + } + + this.cosmosObject = cosmosObject; + } + + public CosmosArray GroupByItems + { + get + { + if (!this.cosmosObject.TryGetValue(GroupByItemsPropertyName, out CosmosElement cosmosElement)) + { + throw new InvalidOperationException($"Underlying object does not have an 'groupByItems' field."); + } + + if (!(cosmosElement is CosmosArray cosmosArray)) + { + throw new ArgumentException($"{nameof(RewrittenGroupByProjection)}['groupByItems'] was not an array."); + } + + return cosmosArray; + } + } + + public bool TryGetPayload(out CosmosElement payload) => this.cosmosObject.TryGetValue(PayloadPropertyName, out payload); + } + + protected sealed class GroupingTable : IEnumerable> + { + private static readonly IReadOnlyList EmptyAggregateOperators = new AggregateOperator[] { }; + + private readonly Dictionary table; + private readonly IReadOnlyDictionary groupByAliasToAggregateType; + private readonly IReadOnlyList orderedAliases; + private readonly bool hasSelectValue; + + private GroupingTable( + IReadOnlyDictionary groupByAliasToAggregateType, + IReadOnlyList orderedAliases, + bool hasSelectValue) + { + this.groupByAliasToAggregateType = groupByAliasToAggregateType ?? throw new ArgumentNullException(nameof(groupByAliasToAggregateType)); + this.orderedAliases = orderedAliases; + this.hasSelectValue = hasSelectValue; + this.table = new Dictionary(); + } + + public int Count => this.table.Count; + + public bool IsDone { get; private set; } + + public void AddPayload(RewrittenGroupByProjection rewrittenGroupByProjection) + { + // For VALUE queries the payload will be undefined if the field was undefined. + if (rewrittenGroupByProjection.TryGetPayload(out CosmosElement payload)) + { + UInt128 groupByKeysHash = DistinctHash.GetHash(rewrittenGroupByProjection.GroupByItems); + + if (!this.table.TryGetValue(groupByKeysHash, out SingleGroupAggregator singleGroupAggregator)) + { + singleGroupAggregator = SingleGroupAggregator.TryCreate( + EmptyAggregateOperators, + this.groupByAliasToAggregateType, + this.orderedAliases, + this.hasSelectValue, + continuationToken: null).Result; + this.table[groupByKeysHash] = singleGroupAggregator; + } + + singleGroupAggregator.AddValues(payload); + } + } + + public IReadOnlyList Drain(int maxItemCount) + { + List keys = this.table.Keys.Take(maxItemCount).ToList(); + List singleGroupAggregators = new List(keys.Count); + foreach (UInt128 key in keys) + { + SingleGroupAggregator singleGroupAggregator = this.table[key]; + singleGroupAggregators.Add(singleGroupAggregator); + } + + foreach (UInt128 key in keys) + { + this.table.Remove(key); + } + + List results = new List(); + foreach (SingleGroupAggregator singleGroupAggregator in singleGroupAggregators) + { + results.Add(singleGroupAggregator.GetResult()); + } + + if (this.Count == 0) + { + this.IsDone = true; + } + + return results; + } + + public CosmosElement GetCosmosElementContinuationToken() + { + Dictionary dictionary = new Dictionary(); + foreach (KeyValuePair kvp in this.table) + { + dictionary.Add(kvp.Key.ToString(), kvp.Value.GetCosmosElementContinuationToken()); + } + + return CosmosObject.Create(dictionary); + } + + public IEnumerator> GetEnumerator => this.table.GetEnumerator(); + + public static TryCatch TryCreateFromContinuationToken( + IReadOnlyDictionary groupByAliasToAggregateType, + IReadOnlyList orderedAliases, + bool hasSelectValue, + CosmosElement continuationToken) + { + GroupingTable groupingTable = new GroupingTable( + groupByAliasToAggregateType, + orderedAliases, + hasSelectValue); + + if (continuationToken != null) + { + if (!(continuationToken is CosmosObject groupingTableContinuationToken)) + { + return TryCatch.FromException( + new MalformedContinuationTokenException( + $"Invalid GroupingTableContinuationToken")); + } + + foreach (KeyValuePair kvp in groupingTableContinuationToken) + { + string key = kvp.Key; + CosmosElement value = kvp.Value; + + if (!UInt128.TryParse(key, out UInt128 groupByKey)) + { + return TryCatch.FromException( + new MalformedContinuationTokenException( + $"Invalid GroupingTableContinuationToken")); + } + + TryCatch tryCreateSingleGroupAggregator = SingleGroupAggregator.TryCreate( + EmptyAggregateOperators, + groupByAliasToAggregateType, + orderedAliases, + hasSelectValue, + value); + + if (tryCreateSingleGroupAggregator.Succeeded) + { + groupingTable.table[groupByKey] = tryCreateSingleGroupAggregator.Result; + } + else + { + return TryCatch.FromException(tryCreateSingleGroupAggregator.Exception); + } + } + } + + return TryCatch.FromResult(groupingTable); + } + + IEnumerator> IEnumerable>.GetEnumerator() => this.table.GetEnumerator(); + + IEnumerator IEnumerable.GetEnumerator() => this.table.GetEnumerator(); + } + } +} diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Performance.Tests/Query/DistinctHashBenchmark.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Performance.Tests/Query/DistinctHashBenchmark.cs index b909809d64..99e2943ba6 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Performance.Tests/Query/DistinctHashBenchmark.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Performance.Tests/Query/DistinctHashBenchmark.cs @@ -4,7 +4,7 @@ using BenchmarkDotNet.Attributes; using Microsoft.Azure.Cosmos.CosmosElements; using Microsoft.Azure.Cosmos.CosmosElements.Numbers; - using Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.Distinct; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Distinct; [MemoryDiagnoser] public class DistinctHashBenchmark diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/GroupByQueryPipelineStageTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/GroupByQueryPipelineStageTests.cs new file mode 100644 index 0000000000..e386129002 --- /dev/null +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/GroupByQueryPipelineStageTests.cs @@ -0,0 +1,77 @@ +//------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +//------------------------------------------------------------ + +namespace Microsoft.Azure.Cosmos.Tests.Query.Pipeline +{ + using System; + using System.Collections.Generic; + using System.Threading.Tasks; + using Microsoft.Azure.Cosmos.CosmosElements; + 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.Aggregate; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline.GroupBy; + using Microsoft.VisualStudio.TestTools.UnitTesting; + + [TestClass] + public class GroupByQueryPipelineStageTests + { + [TestMethod] + public async Task SinglePageAsync() + { + IReadOnlyList> pages = new List>() + { + new List() + { + CosmosElement.Parse("{\"groupByItems\": [{\"item\" : \"John\"}], \"payload\" : {\"name\": \"John\", \"count\": {\"item\": 42}}}") + } + }; + + List elements = await GroupByQueryPipelineStageTests.CreateAndDrainAsync( + pages: pages, + executionEnvironment: ExecutionEnvironment.Compute, + continuationToken: null, + groupByAliasToAggregateType: new Dictionary() { { "name", null }, { "count", AggregateOperator.Sum } }, + orderedAliases: new List() { "name", "count" }, + hasSelectValue: false); + + Assert.AreEqual(1, elements.Count); + Assert.AreEqual(42, Number64.ToLong(((elements[0] as CosmosObject)["count"] as CosmosNumber).Value)); + Assert.AreEqual("John", ((elements[0] as CosmosObject)["name"] as CosmosString).Value); + } + + private static async Task> CreateAndDrainAsync( + IReadOnlyList> pages, + ExecutionEnvironment executionEnvironment, + CosmosElement continuationToken, + IReadOnlyDictionary groupByAliasToAggregateType, + IReadOnlyList orderedAliases, + bool hasSelectValue) + { + IQueryPipelineStage source = new MockQueryPipelineStage(pages); + + TryCatch tryCreateGroupByStage = await GroupByQueryPipelineStage.TryCreateAsync( + executionEnvironment: executionEnvironment, + continuationToken: continuationToken, + tryCreateSourceAsync: (CosmosElement continuationToken) => Task.FromResult(TryCatch.FromResult(source)), + groupByAliasToAggregateType: groupByAliasToAggregateType, + orderedAliases: orderedAliases, + hasSelectValue: hasSelectValue); + Assert.IsTrue(tryCreateGroupByStage.Succeeded); + + IQueryPipelineStage groupByQueryPipelineStage = tryCreateGroupByStage.Result; + + List elements = new List(); + await foreach (TryCatch page in new EnumerableStage(groupByQueryPipelineStage)) + { + page.ThrowIfFailed(); + + elements.AddRange(page.Result.Documents); + } + + return elements; + } + } +} From 2e40aeeab86d6bf364361dc5a020cf5d2a586405 Mon Sep 17 00:00:00 2001 From: Brandon Chong Date: Sun, 5 Jul 2020 16:44:01 -0700 Subject: [PATCH 27/85] got the pipeline working --- ...ggregateDocumentQueryExecutionComponent.cs | 4 +- ...DistinctDocumentQueryExecutionComponent.cs | 4 +- .../GroupByDocumentQueryExecutionComponent.cs | 4 +- .../SkipDocumentQueryExecutionComponent.cs | 4 +- .../TakeDocumentQueryExecutionComponent.cs | 4 +- .../CosmosQueryExecutionContextFactory.cs | 2 +- ...OrderByItemQueryExecutionContext.Resume.cs | 2 +- ...arallelItemQueryExecutionContext.Resume.cs | 2 +- .../PipelinedDocumentQueryExecutionContext.cs | 44 ++--- .../AggregateQueryPipelineStage.Client.cs | 11 +- .../AggregateQueryPipelineStage.Compute.cs | 6 +- .../Aggregate/AggregateQueryPipelineStage.cs | 14 +- .../DistinctQueryPipelineStage.Client.cs | 10 +- .../DistinctQueryPipelineStage.Compute.cs | 11 +- .../Distinct/DistinctQueryPipelineStage.cs | 14 +- .../GroupByQueryPipelineStage.Client.cs | 6 +- .../GroupByQueryPipelineStage.Compute.cs | 6 +- .../GroupBy/GroupByQueryPipelineStage.cs | 13 +- .../Core/Pipeline/IQueryPipelineStage.cs | 3 - .../Pipeline/MonadicCreatePipelineStage.cs | 12 ++ ...rallelCrossPartitionQueryPageEnumerator.cs | 66 ------- .../Query/Core/Pipeline/PipelineFactory.cs | 116 ++++++++++++ .../Core/Pipeline/QueryPipelineStageBase.cs | 4 +- .../BackendQueryDataSource.cs | 2 +- .../IQueryDataSource.cs | 2 +- ...arallelCrossPartitionQueryPipelineStage.cs | 169 ++++++++++++++++++ .../QueryPartitionRangePageEnumerator.cs | 3 +- .../Skip/SkipQueryPipelineStage.Client.cs | 10 +- .../Skip/SkipQueryPipelineStage.Compute.cs | 10 +- .../Pipeline/Skip/SkipQueryPipelineStage.cs | 15 +- .../Take/TakeQueryPipelineStage.Client.cs | 20 +-- .../Take/TakeQueryPipelineStage.Compute.cs | 22 +-- .../Pipeline/Take/TakeQueryPipelineStage.cs | 25 +-- .../src/Query/Core/QueryPlan/QueryInfo.cs | 48 +---- .../CosmosQueryUnitTests.cs | 10 +- .../InMemoryCollectionQueryDataSource.cs | 2 +- .../AggregateQueryPipelineStageTests.cs | 5 +- .../DistinctQueryPipelineStageTests.cs | 4 +- .../Query/Pipeline/FactoryTests.cs | 40 +++++ .../GroupByQueryPipelineStageTests.cs | 4 +- .../Pipeline/SkipQueryPipelineStageTests.cs | 4 +- .../Pipeline/TakeQueryPipelineStageTests.cs | 4 +- .../QueryPartitionRangePageEnumeratorTests.cs | 3 +- 43 files changed, 500 insertions(+), 264 deletions(-) create mode 100644 Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/MonadicCreatePipelineStage.cs delete mode 100644 Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Partitions/ParallelCrossPartitionQueryPageEnumerator.cs create mode 100644 Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/PipelineFactory.cs rename Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/{Partitions => Remote}/BackendQueryDataSource.cs (96%) rename Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/{Partitions => Remote}/IQueryDataSource.cs (90%) create mode 100644 Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/ParallelCrossPartitionQueryPipelineStage.cs rename Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/{Partitions => Remote}/QueryPartitionRangePageEnumerator.cs (94%) create mode 100644 Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/FactoryTests.cs diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/AggregateDocumentQueryExecutionComponent.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/AggregateDocumentQueryExecutionComponent.cs index 1c00bd5f73..e26d329cb5 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/AggregateDocumentQueryExecutionComponent.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/AggregateDocumentQueryExecutionComponent.cs @@ -56,14 +56,14 @@ protected AggregateDocumentQueryExecutionComponent( this.isValueAggregateQuery = isValueAggregateQuery; } - public static Task> TryCreateAsync( + public static Task> MonadicCreateAsync( ExecutionEnvironment executionEnvironment, IReadOnlyList aggregates, IReadOnlyDictionary aliasToAggregateType, IReadOnlyList orderedAliases, bool hasSelectValue, CosmosElement continuationToken, - Func>> tryCreateSourceAsync) + Func>> monadicCreatePipelineStage) { return default; } diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctDocumentQueryExecutionComponent.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctDocumentQueryExecutionComponent.cs index 71c7ad185b..1e983601ff 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctDocumentQueryExecutionComponent.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctDocumentQueryExecutionComponent.cs @@ -25,10 +25,10 @@ protected DistinctDocumentQueryExecutionComponent( this.distinctMap = distinctMap ?? throw new ArgumentNullException(nameof(distinctMap)); } - public static Task> TryCreateAsync( + public static Task> MonadicCreateAsync( ExecutionEnvironment executionEnvironment, CosmosElement requestContinuation, - Func>> tryCreateSourceAsync, + Func>> monadicCreatePipelineStage, DistinctQueryType distinctQueryType) { return default; diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/GroupBy/GroupByDocumentQueryExecutionComponent.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/GroupBy/GroupByDocumentQueryExecutionComponent.cs index 8b92b6030a..7894065adc 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/GroupBy/GroupByDocumentQueryExecutionComponent.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/GroupBy/GroupByDocumentQueryExecutionComponent.cs @@ -57,10 +57,10 @@ protected GroupByDocumentQueryExecutionComponent( public override bool IsDone => this.groupingTable.IsDone; - public static Task> TryCreateAsync( + public static Task> MonadicCreateAsync( ExecutionEnvironment executionEnvironment, CosmosElement continuationToken, - Func>> tryCreateSourceAsync, + Func>> monadicCreatePipelineStage, IReadOnlyDictionary groupByAliasToAggregateType, IReadOnlyList orderedAliases, bool hasSelectValue) diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/SkipDocumentQueryExecutionComponent.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/SkipDocumentQueryExecutionComponent.cs index d07a69bf63..fc0539f3eb 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/SkipDocumentQueryExecutionComponent.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/SkipDocumentQueryExecutionComponent.cs @@ -26,11 +26,11 @@ protected SkipDocumentQueryExecutionComponent(IDocumentQueryExecutionComponent s this.skipCount = (int)skipCount; } - public static Task> TryCreateAsync( + public static Task> MonadicCreateAsync( ExecutionEnvironment executionEnvironment, int offsetCount, CosmosElement continuationToken, - Func>> tryCreateSourceAsync) + Func>> monadicCreatePipelineStage) { return default; } diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/TakeDocumentQueryExecutionComponent.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/TakeDocumentQueryExecutionComponent.cs index 7c9a56c714..07011fcb8d 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/TakeDocumentQueryExecutionComponent.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/TakeDocumentQueryExecutionComponent.cs @@ -26,11 +26,11 @@ protected TakeDocumentQueryExecutionComponent(IDocumentQueryExecutionComponent s this.skipCount = (int)skipCount; } - public static Task> TryCreateAsync( + public static Task> MonadicCreateAsync( ExecutionEnvironment executionEnvironment, int offsetCount, CosmosElement continuationToken, - Func>> tryCreateSourceAsync) + Func>> monadicCreatePipelineStage) { return default; } diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosQueryExecutionContextFactory.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosQueryExecutionContextFactory.cs index 16ca62bd30..2f35d372a4 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosQueryExecutionContextFactory.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosQueryExecutionContextFactory.cs @@ -401,7 +401,7 @@ private static async Task> TryCreateSpecia returnResultsInDeterministicOrder: inputParameters.ReturnResultsInDeterministicOrder, testSettings: inputParameters.TestInjections); - return await PipelinedDocumentQueryExecutionContext.TryCreateAsync( + return await PipelinedDocumentQueryExecutionContext.MonadicCreateAsync( inputParameters.ExecutionEnvironment, cosmosQueryContext, initParams, diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/OrderBy/CosmosOrderByItemQueryExecutionContext.Resume.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/OrderBy/CosmosOrderByItemQueryExecutionContext.Resume.cs index 5bda5c4671..8a0bad89c3 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/OrderBy/CosmosOrderByItemQueryExecutionContext.Resume.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/OrderBy/CosmosOrderByItemQueryExecutionContext.Resume.cs @@ -34,7 +34,7 @@ private static class Expressions public const string GreaterThanOrEqualTo = ">="; } - public static async Task> TryCreateAsync( + public static async Task> MonadicCreateAsync( CosmosQueryContext queryContext, CosmosCrossPartitionQueryExecutionContext.CrossPartitionInitParams initParams, CosmosElement requestContinuationToken, diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/Parallel/CosmosParallelItemQueryExecutionContext.Resume.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/Parallel/CosmosParallelItemQueryExecutionContext.Resume.cs index 0ad350287f..d5934bb618 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/Parallel/CosmosParallelItemQueryExecutionContext.Resume.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/Parallel/CosmosParallelItemQueryExecutionContext.Resume.cs @@ -20,7 +20,7 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.Parallel internal sealed partial class CosmosParallelItemQueryExecutionContext : CosmosCrossPartitionQueryExecutionContext { - public static async Task> TryCreateAsync( + public static async Task> MonadicCreateAsync( CosmosQueryContext queryContext, CosmosCrossPartitionQueryExecutionContext.CrossPartitionInitParams initParams, CosmosElement requestContinuationToken, diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/PipelinedDocumentQueryExecutionContext.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/PipelinedDocumentQueryExecutionContext.cs index 29c7bace17..d0dbfd5885 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/PipelinedDocumentQueryExecutionContext.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/PipelinedDocumentQueryExecutionContext.cs @@ -121,7 +121,7 @@ public override bool IsDone } } - public static async Task> TryCreateAsync( + public static async Task> MonadicCreateAsync( ExecutionEnvironment executionEnvironment, CosmosQueryContext queryContext, CosmosCrossPartitionQueryExecutionContext.CrossPartitionInitParams initParams, @@ -156,7 +156,7 @@ public static async Task> TryCreateAsync( Task> tryCreateOrderByComponentAsync(CosmosElement continuationToken) { - return CosmosOrderByItemQueryExecutionContext.TryCreateAsync( + return CosmosOrderByItemQueryExecutionContext.MonadicCreateAsync( queryContext, initParams, continuationToken, @@ -165,7 +165,7 @@ Task> tryCreateOrderByComponentAsync( Task> tryCreateParallelComponentAsync(CosmosElement continuationToken) { - return CosmosParallelItemQueryExecutionContext.TryCreateAsync( + return CosmosParallelItemQueryExecutionContext.MonadicCreateAsync( queryContext, initParams, continuationToken, @@ -184,42 +184,42 @@ Task> tryCreateParallelComponentAsync if (queryInfo.HasAggregates && !queryInfo.HasGroupBy) { - Func>> tryCreateSourceAsync = tryCreatePipelineAsync; + Func>> monadicCreatePipelineStage = tryCreatePipelineAsync; tryCreatePipelineAsync = async (continuationToken) => { - return await AggregateDocumentQueryExecutionComponent.TryCreateAsync( + return await AggregateDocumentQueryExecutionComponent.MonadicCreateAsync( executionEnvironment, queryInfo.Aggregates, queryInfo.GroupByAliasToAggregateType, queryInfo.GroupByAliases, queryInfo.HasSelectValue, continuationToken, - tryCreateSourceAsync); + monadicCreatePipelineStage); }; } if (queryInfo.HasDistinct) { - Func>> tryCreateSourceAsync = tryCreatePipelineAsync; + Func>> monadicCreatePipelineStage = tryCreatePipelineAsync; tryCreatePipelineAsync = async (continuationToken) => { - return await DistinctDocumentQueryExecutionComponent.TryCreateAsync( + return await DistinctDocumentQueryExecutionComponent.MonadicCreateAsync( executionEnvironment, continuationToken, - tryCreateSourceAsync, + monadicCreatePipelineStage, queryInfo.DistinctType); }; } if (queryInfo.HasGroupBy) { - Func>> tryCreateSourceAsync = tryCreatePipelineAsync; + Func>> monadicCreatePipelineStage = tryCreatePipelineAsync; tryCreatePipelineAsync = async (continuationToken) => { - return await GroupByDocumentQueryExecutionComponent.TryCreateAsync( + return await GroupByDocumentQueryExecutionComponent.MonadicCreateAsync( executionEnvironment, continuationToken, - tryCreateSourceAsync, + monadicCreatePipelineStage, queryInfo.GroupByAliasToAggregateType, queryInfo.GroupByAliases, queryInfo.HasSelectValue); @@ -228,40 +228,40 @@ Task> tryCreateParallelComponentAsync if (queryInfo.Offset.HasValue) { - Func>> tryCreateSourceAsync = tryCreatePipelineAsync; + Func>> monadicCreatePipelineStage = tryCreatePipelineAsync; tryCreatePipelineAsync = async (continuationToken) => { - return await SkipDocumentQueryExecutionComponent.TryCreateAsync( + return await SkipDocumentQueryExecutionComponent.MonadicCreateAsync( executionEnvironment, queryInfo.Offset.Value, continuationToken, - tryCreateSourceAsync); + monadicCreatePipelineStage); }; } if (queryInfo.Limit.HasValue) { - Func>> tryCreateSourceAsync = tryCreatePipelineAsync; + Func>> monadicCreatePipelineStage = tryCreatePipelineAsync; tryCreatePipelineAsync = async (continuationToken) => { - return await TakeDocumentQueryExecutionComponent.TryCreateAsync( + return await TakeDocumentQueryExecutionComponent.MonadicCreateAsync( executionEnvironment, queryInfo.Limit.Value, continuationToken, - tryCreateSourceAsync); + monadicCreatePipelineStage); }; } if (queryInfo.Top.HasValue) { - Func>> tryCreateSourceAsync = tryCreatePipelineAsync; + Func>> monadicCreatePipelineStage = tryCreatePipelineAsync; tryCreatePipelineAsync = async (continuationToken) => { - return await TakeDocumentQueryExecutionComponent.TryCreateAsync( + return await TakeDocumentQueryExecutionComponent.MonadicCreateAsync( executionEnvironment, queryInfo.Top.Value, continuationToken, - tryCreateSourceAsync); + monadicCreatePipelineStage); }; } @@ -317,7 +317,7 @@ public static async Task> TryCreatePassthr testSettings: initParams.TestSettings); // Return a parallel context, since we still want to be able to handle splits and concurrency / buffering. - return (await CosmosParallelItemQueryExecutionContext.TryCreateAsync( + return (await CosmosParallelItemQueryExecutionContext.MonadicCreateAsync( queryContext, initParams, requestContinuationToken, diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Aggregate/AggregateQueryPipelineStage.Client.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Aggregate/AggregateQueryPipelineStage.Client.cs index 21a71e7070..e27fe469dc 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Aggregate/AggregateQueryPipelineStage.Client.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Aggregate/AggregateQueryPipelineStage.Client.cs @@ -11,6 +11,7 @@ namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.Aggregate using Microsoft.Azure.Cosmos.CosmosElements; using Microsoft.Azure.Cosmos.Query.Core.Monads; using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Aggregate.Aggregators; + using static Microsoft.Azure.Cosmos.Query.Core.Pipeline.PipelineFactory; internal abstract partial class AggregateQueryPipelineStage : QueryPipelineStageBase { @@ -25,17 +26,17 @@ private ClientAggregateQueryPipelineStage( // all the work is done in the base constructor. } - public static async Task> TryCreateAsync( + public static TryCatch MonadicCreate( IReadOnlyList aggregates, IReadOnlyDictionary aliasToAggregateType, IReadOnlyList orderedAliases, bool hasSelectValue, CosmosElement continuationToken, - Func>> tryCreateSourceAsync) + MonadicCreatePipelineStage monadicCreatePipelineStage) { - if (tryCreateSourceAsync == null) + if (monadicCreatePipelineStage == null) { - throw new ArgumentNullException(nameof(tryCreateSourceAsync)); + throw new ArgumentNullException(nameof(monadicCreatePipelineStage)); } TryCatch tryCreateSingleGroupAggregator = SingleGroupAggregator.TryCreate( @@ -49,7 +50,7 @@ public static async Task> TryCreateAsync( return TryCatch.FromException(tryCreateSingleGroupAggregator.Exception); } - TryCatch tryCreateSource = await tryCreateSourceAsync(continuationToken); + TryCatch tryCreateSource = monadicCreatePipelineStage(continuationToken); if (tryCreateSource.Failed) { return tryCreateSource; diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Aggregate/AggregateQueryPipelineStage.Compute.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Aggregate/AggregateQueryPipelineStage.Compute.cs index 8bdacfdd4d..b6bdd1850c 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Aggregate/AggregateQueryPipelineStage.Compute.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Aggregate/AggregateQueryPipelineStage.Compute.cs @@ -30,13 +30,13 @@ private ComputeAggregateQueryPipelineStage( // all the work is done in the base constructor. } - public static async Task> TryCreateAsync( + public static TryCatch MonadicCreate( IReadOnlyList aggregates, IReadOnlyDictionary aliasToAggregateType, IReadOnlyList orderedAliases, bool hasSelectValue, CosmosElement continuationToken, - Func>> tryCreateSourceAsync) + MonadicCreatePipelineStage monadicCreatePipelineStage) { AggregateContinuationToken aggregateContinuationToken; if (continuationToken != null) @@ -73,7 +73,7 @@ public static async Task> TryCreateAsync( } else { - tryCreateSource = await tryCreateSourceAsync(aggregateContinuationToken.SourceContinuationToken); + tryCreateSource = monadicCreatePipelineStage(aggregateContinuationToken.SourceContinuationToken); } if (tryCreateSource.Failed) diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Aggregate/AggregateQueryPipelineStage.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Aggregate/AggregateQueryPipelineStage.cs index c2216fb4ca..cdcfb256b7 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Aggregate/AggregateQueryPipelineStage.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Aggregate/AggregateQueryPipelineStage.cs @@ -6,11 +6,13 @@ namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.Aggregate { using System; using System.Collections.Generic; + using System.Threading; using System.Threading.Tasks; using Microsoft.Azure.Cosmos.CosmosElements; using Microsoft.Azure.Cosmos.Query.Core.ExecutionContext; using Microsoft.Azure.Cosmos.Query.Core.Monads; using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Aggregate.Aggregators; + using static Microsoft.Azure.Cosmos.Query.Core.Pipeline.PipelineFactory; /// /// Stage that is able to aggregate local aggregates from multiple continuations and partitions. @@ -56,29 +58,29 @@ public AggregateQueryPipelineStage( this.isValueQuery = isValueQuery; } - public static Task> TryCreateAsync( + public static TryCatch MonadicCreate( ExecutionEnvironment executionEnvironment, IReadOnlyList aggregates, IReadOnlyDictionary aliasToAggregateType, IReadOnlyList orderedAliases, bool hasSelectValue, CosmosElement continuationToken, - Func>> tryCreateSourceAsync) => executionEnvironment switch + MonadicCreatePipelineStage monadicCreatePipelineStage) => executionEnvironment switch { - ExecutionEnvironment.Client => ClientAggregateQueryPipelineStage.TryCreateAsync( + ExecutionEnvironment.Client => ClientAggregateQueryPipelineStage.MonadicCreate( aggregates, aliasToAggregateType, orderedAliases, hasSelectValue, continuationToken, - tryCreateSourceAsync), - ExecutionEnvironment.Compute => ComputeAggregateQueryPipelineStage.TryCreateAsync( + monadicCreatePipelineStage), + ExecutionEnvironment.Compute => ComputeAggregateQueryPipelineStage.MonadicCreate( aggregates, aliasToAggregateType, orderedAliases, hasSelectValue, continuationToken, - tryCreateSourceAsync), + monadicCreatePipelineStage), _ => throw new ArgumentException($"Unknown {nameof(ExecutionEnvironment)}: {executionEnvironment}."), }; diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Distinct/DistinctQueryPipelineStage.Client.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Distinct/DistinctQueryPipelineStage.Client.cs index d46485d578..e2179679e5 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Distinct/DistinctQueryPipelineStage.Client.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Distinct/DistinctQueryPipelineStage.Client.cs @@ -40,14 +40,14 @@ private ClientDistinctQueryPipelineStage( this.distinctQueryType = distinctQueryType; } - public static async Task> TryCreateAsync( + public static TryCatch MonadicCreate( CosmosElement requestContinuation, - Func>> tryCreateSourceAsync, + MonadicCreatePipelineStage monadicCreatePipelineStage, DistinctQueryType distinctQueryType) { - if (tryCreateSourceAsync == null) + if (monadicCreatePipelineStage == null) { - throw new ArgumentNullException(nameof(tryCreateSourceAsync)); + throw new ArgumentNullException(nameof(monadicCreatePipelineStage)); } DistinctContinuationToken distinctContinuationToken; @@ -104,7 +104,7 @@ public static async Task> TryCreateAsync( sourceToken = null; } - TryCatch tryCreateSource = await tryCreateSourceAsync(sourceToken); + TryCatch tryCreateSource = monadicCreatePipelineStage(sourceToken); if (!tryCreateSource.Succeeded) { return TryCatch.FromException(tryCreateSource.Exception); diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Distinct/DistinctQueryPipelineStage.Compute.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Distinct/DistinctQueryPipelineStage.Compute.cs index b05fc86009..06992788f6 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Distinct/DistinctQueryPipelineStage.Compute.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Distinct/DistinctQueryPipelineStage.Compute.cs @@ -32,14 +32,14 @@ private ComputeDistinctQueryPipelineStage( { } - public static async Task> TryCreateAsync( + public static TryCatch MonadicCreate( CosmosElement requestContinuation, - Func>> tryCreateSourceAsync, + MonadicCreatePipelineStage monadicCreatePipelineStage, DistinctQueryType distinctQueryType) { - if (tryCreateSourceAsync == null) + if (monadicCreatePipelineStage == null) { - throw new ArgumentNullException(nameof(tryCreateSourceAsync)); + throw new ArgumentNullException(nameof(monadicCreatePipelineStage)); } DistinctContinuationToken distinctContinuationToken; @@ -65,8 +65,7 @@ public static async Task> TryCreateAsync( return TryCatch.FromException(tryCreateDistinctMap.Exception); } - TryCatch tryCreateSource = await tryCreateSourceAsync( - distinctContinuationToken.SourceToken); + TryCatch tryCreateSource = monadicCreatePipelineStage(distinctContinuationToken.SourceToken); if (!tryCreateSource.Succeeded) { return TryCatch.FromException(tryCreateSource.Exception); diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Distinct/DistinctQueryPipelineStage.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Distinct/DistinctQueryPipelineStage.cs index 9b122ac281..7f1b348324 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Distinct/DistinctQueryPipelineStage.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Distinct/DistinctQueryPipelineStage.cs @@ -5,7 +5,7 @@ namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.Distinct { using System; - using System.Threading.Tasks; + using System.Threading; using Microsoft.Azure.Cosmos.CosmosElements; using Microsoft.Azure.Cosmos.Query.Core.ExecutionContext; using Microsoft.Azure.Cosmos.Query.Core.Monads; @@ -36,19 +36,19 @@ protected DistinctQueryPipelineStage( this.distinctMap = distinctMap ?? throw new ArgumentNullException(nameof(distinctMap)); } - public static Task> TryCreateAsync( + public static TryCatch MonadicCreate( ExecutionEnvironment executionEnvironment, CosmosElement requestContinuation, - Func>> tryCreateSourceAsync, + MonadicCreatePipelineStage monadicCreatePipelineStage, DistinctQueryType distinctQueryType) => executionEnvironment switch { - ExecutionEnvironment.Client => ClientDistinctQueryPipelineStage.TryCreateAsync( + ExecutionEnvironment.Client => ClientDistinctQueryPipelineStage.MonadicCreate( requestContinuation, - tryCreateSourceAsync, + monadicCreatePipelineStage, distinctQueryType), - ExecutionEnvironment.Compute => ComputeDistinctQueryPipelineStage.TryCreateAsync( + ExecutionEnvironment.Compute => ComputeDistinctQueryPipelineStage.MonadicCreate( requestContinuation, - tryCreateSourceAsync, + monadicCreatePipelineStage, distinctQueryType), _ => throw new ArgumentException($"Unknown {nameof(ExecutionEnvironment)}: {executionEnvironment}."), }; diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/GroupBy/GroupByQueryPipelineStage.Client.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/GroupBy/GroupByQueryPipelineStage.Client.cs index 7a0ee1c0e8..37079609bd 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/GroupBy/GroupByQueryPipelineStage.Client.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/GroupBy/GroupByQueryPipelineStage.Client.cs @@ -27,9 +27,9 @@ private ClientGroupByQueryPipelineStage( { } - public static async Task> TryCreateAsync( + public static TryCatch MonadicCreate( CosmosElement requestContinuation, - Func>> tryCreateSourceAsync, + MonadicCreatePipelineStage monadicCreatePipelineStage, IReadOnlyDictionary groupByAliasToAggregateType, IReadOnlyList orderedAliases, bool hasSelectValue) @@ -45,7 +45,7 @@ public static async Task> TryCreateAsync( return TryCatch.FromException(tryCreateGroupingTable.Exception); } - TryCatch tryCreateSource = await tryCreateSourceAsync(requestContinuation); + TryCatch tryCreateSource = monadicCreatePipelineStage(requestContinuation); if (tryCreateSource.Failed) { return tryCreateSource; diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/GroupBy/GroupByQueryPipelineStage.Compute.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/GroupBy/GroupByQueryPipelineStage.Compute.cs index 45a523ed07..dd8016a7bb 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/GroupBy/GroupByQueryPipelineStage.Compute.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/GroupBy/GroupByQueryPipelineStage.Compute.cs @@ -33,9 +33,9 @@ private ComputeGroupByQueryPipelineStage( { } - public static async Task> TryCreateAsync( + public static TryCatch MonadicCreate( CosmosElement requestContinuation, - Func>> tryCreateSourceAsync, + MonadicCreatePipelineStage monadicCreatePipelineStage, IReadOnlyDictionary groupByAliasToAggregateType, IReadOnlyList orderedAliases, bool hasSelectValue) @@ -65,7 +65,7 @@ public static async Task> TryCreateAsync( } else { - tryCreateSource = await tryCreateSourceAsync(groupByContinuationToken.SourceContinuationToken); + tryCreateSource = monadicCreatePipelineStage(groupByContinuationToken.SourceContinuationToken); } if (!tryCreateSource.Succeeded) diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/GroupBy/GroupByQueryPipelineStage.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/GroupBy/GroupByQueryPipelineStage.cs index e07b63372b..6974eeacab 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/GroupBy/GroupByQueryPipelineStage.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/GroupBy/GroupByQueryPipelineStage.cs @@ -8,6 +8,7 @@ namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.GroupBy using System.Collections; using System.Collections.Generic; using System.Linq; + using System.Threading; using System.Threading.Tasks; using Microsoft.Azure.Cosmos.CosmosElements; using Microsoft.Azure.Cosmos.Query.Core.Exceptions; @@ -56,23 +57,23 @@ protected GroupByQueryPipelineStage( this.groupingTable = groupingTable ?? throw new ArgumentNullException(nameof(groupingTable)); } - public static Task> TryCreateAsync( + public static TryCatch MonadicCreate( ExecutionEnvironment executionEnvironment, CosmosElement continuationToken, - Func>> tryCreateSourceAsync, + MonadicCreatePipelineStage monadicCreatePipelineStage, IReadOnlyDictionary groupByAliasToAggregateType, IReadOnlyList orderedAliases, bool hasSelectValue) => executionEnvironment switch { - ExecutionEnvironment.Client => ClientGroupByQueryPipelineStage.TryCreateAsync( + ExecutionEnvironment.Client => ClientGroupByQueryPipelineStage.MonadicCreate( continuationToken, - tryCreateSourceAsync, + monadicCreatePipelineStage, groupByAliasToAggregateType, orderedAliases, hasSelectValue), - ExecutionEnvironment.Compute => ComputeGroupByQueryPipelineStage.TryCreateAsync( + ExecutionEnvironment.Compute => ComputeGroupByQueryPipelineStage.MonadicCreate( continuationToken, - tryCreateSourceAsync, + monadicCreatePipelineStage, groupByAliasToAggregateType, orderedAliases, hasSelectValue), diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/IQueryPipelineStage.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/IQueryPipelineStage.cs index f5ad3230e0..49c5832f5c 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/IQueryPipelineStage.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/IQueryPipelineStage.cs @@ -5,12 +5,9 @@ namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline { using System.Collections.Generic; - using System.Threading; - using System.Threading.Tasks; using Microsoft.Azure.Cosmos.Query.Core.Monads; internal interface IQueryPipelineStage : IAsyncEnumerator> { - bool HasMoreResults { get; } } } diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/MonadicCreatePipelineStage.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/MonadicCreatePipelineStage.cs new file mode 100644 index 0000000000..3f26975d14 --- /dev/null +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/MonadicCreatePipelineStage.cs @@ -0,0 +1,12 @@ +// ------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +// ------------------------------------------------------------ + +namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline +{ + using System.Threading.Tasks; + using Microsoft.Azure.Cosmos.CosmosElements; + using Microsoft.Azure.Cosmos.Query.Core.Monads; + + internal delegate TryCatch MonadicCreatePipelineStage(CosmosElement continuationToken); +} diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Partitions/ParallelCrossPartitionQueryPageEnumerator.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Partitions/ParallelCrossPartitionQueryPageEnumerator.cs deleted file mode 100644 index 37f9ad2c89..0000000000 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Partitions/ParallelCrossPartitionQueryPageEnumerator.cs +++ /dev/null @@ -1,66 +0,0 @@ -// ------------------------------------------------------------ -// Copyright (c) Microsoft Corporation. All rights reserved. -// ------------------------------------------------------------ - -namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.Partitions -{ - using System.Collections.Generic; - using Microsoft.Azure.Cosmos.Pagination; - - internal static class ParallelCrossPartitionQueryPageEnumerator - { - public static CrossPartitionRangePageEnumerator Create( - IFeedRangeProvider feedRangeProvider, - IQueryDataSource queryDataSource, - SqlQuerySpec sqlQuerySpec, - int pageSize, - CrossPartitionState state = default) - { - return new CrossPartitionRangePageEnumerator( - feedRangeProvider, - ParallelCrossPartitionQueryPageEnumerator.MakeCreateFunction(queryDataSource, sqlQuerySpec, pageSize), - Comparer.Singleton, - state: state); - } - - private static CreatePartitionRangePageEnumerator MakeCreateFunction( - IQueryDataSource queryDataSource, - SqlQuerySpec sqlQuerySpec, - int pageSize) => (FeedRange range, QueryState state) => new QueryPartitionRangePageEnumerator( - queryDataSource, - sqlQuerySpec, - range, - pageSize, - state); - - private sealed class Comparer : IComparer> - { - public static readonly Comparer Singleton = new Comparer(); - - public int Compare( - PartitionRangePageEnumerator partitionRangePageEnumerator1, - PartitionRangePageEnumerator partitionRangePageEnumerator2) - { - if (object.ReferenceEquals(partitionRangePageEnumerator1, partitionRangePageEnumerator2)) - { - return 0; - } - - if (partitionRangePageEnumerator1.HasMoreResults && !partitionRangePageEnumerator2.HasMoreResults) - { - return -1; - } - - if (!partitionRangePageEnumerator1.HasMoreResults && partitionRangePageEnumerator2.HasMoreResults) - { - return 1; - } - - // Either both don't have results or both do. - return string.CompareOrdinal( - ((FeedRangePartitionKeyRange)partitionRangePageEnumerator1.Range).PartitionKeyRangeId, - ((FeedRangePartitionKeyRange)partitionRangePageEnumerator2.Range).PartitionKeyRangeId); - } - } - } -} \ No newline at end of file diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/PipelineFactory.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/PipelineFactory.cs new file mode 100644 index 0000000000..3263f52720 --- /dev/null +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/PipelineFactory.cs @@ -0,0 +1,116 @@ +// ------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +// ------------------------------------------------------------ + +namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline +{ + using System; + using System.Threading; + using System.Threading.Tasks; + using Microsoft.Azure.Cosmos.CosmosElements; + using Microsoft.Azure.Cosmos.Pagination; + using Microsoft.Azure.Cosmos.Query.Core.ExecutionContext; + using Microsoft.Azure.Cosmos.Query.Core.Monads; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Aggregate; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Distinct; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline.GroupBy; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Skip; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Take; + using Microsoft.Azure.Cosmos.Query.Core.QueryPlan; + + internal static class PipelineFactory + { + public static TryCatch MonadicCreate( + ExecutionEnvironment executionEnvironment, + IFeedRangeProvider feedRangeProvider, + IQueryDataSource queryDataSource, + SqlQuerySpec sqlQuerySpec, + QueryInfo queryInfo, + int pageSize, + CosmosElement requestContinuationToken) + { + MonadicCreatePipelineStage monadicCreatePipelineStage; + if (queryInfo.HasOrderBy) + { + throw new NotImplementedException(); + } + else + { + monadicCreatePipelineStage = (continuationToken) => ParallelCrossPartitionQueryPipelineStage.MonadicCreate( + feedRangeProvider: feedRangeProvider, + queryDataSource: queryDataSource, + sqlQuerySpec: sqlQuerySpec, + pageSize: pageSize, + continuationToken: continuationToken); + } + + if (queryInfo.HasAggregates && !queryInfo.HasGroupBy) + { + MonadicCreatePipelineStage monadicCreateSourceStage = monadicCreatePipelineStage; + monadicCreatePipelineStage = (continuationToken) => AggregateQueryPipelineStage.MonadicCreate( + executionEnvironment, + queryInfo.Aggregates, + queryInfo.GroupByAliasToAggregateType, + queryInfo.GroupByAliases, + queryInfo.HasSelectValue, + continuationToken, + monadicCreateSourceStage); + } + + if (queryInfo.HasDistinct) + { + MonadicCreatePipelineStage monadicCreateSourceStage = monadicCreatePipelineStage; + monadicCreatePipelineStage = (continuationToken) => DistinctQueryPipelineStage.MonadicCreate( + executionEnvironment, + continuationToken, + monadicCreateSourceStage, + queryInfo.DistinctType); + } + + if (queryInfo.HasGroupBy) + { + MonadicCreatePipelineStage monadicCreateSourceStage = monadicCreatePipelineStage; + monadicCreatePipelineStage = (continuationToken) => GroupByQueryPipelineStage.MonadicCreate( + executionEnvironment, + continuationToken, + monadicCreateSourceStage, + queryInfo.GroupByAliasToAggregateType, + queryInfo.GroupByAliases, + queryInfo.HasSelectValue); + } + + if (queryInfo.HasOffset) + { + MonadicCreatePipelineStage monadicCreateSourceStage = monadicCreatePipelineStage; + monadicCreatePipelineStage = (continuationToken) => SkipQueryPipelineStage.MonadicCreate( + executionEnvironment, + queryInfo.Offset.Value, + continuationToken, + monadicCreateSourceStage); + } + + if (queryInfo.HasLimit) + { + MonadicCreatePipelineStage monadicCreateSourceStage = monadicCreatePipelineStage; + monadicCreatePipelineStage = (continuationToken) => TakeQueryPipelineStage.MonadicCreateLimitStage( + executionEnvironment, + queryInfo.Limit.Value, + continuationToken, + monadicCreateSourceStage); + } + + if (queryInfo.HasTop) + { + MonadicCreatePipelineStage monadicCreateSourceStage = monadicCreatePipelineStage; + monadicCreatePipelineStage = (continuationToken) => TakeQueryPipelineStage.MonadicCreateTopStage( + executionEnvironment, + queryInfo.Top.Value, + continuationToken, + monadicCreateSourceStage); + } + + return monadicCreatePipelineStage(requestContinuationToken); + } + } +} diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/QueryPipelineStageBase.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/QueryPipelineStageBase.cs index 28267a894d..1864ee7e99 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/QueryPipelineStageBase.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/QueryPipelineStageBase.cs @@ -36,14 +36,14 @@ public async ValueTask MoveNextAsync() return false; } + this.hasStarted = true; + this.Current = await this.GetNextPageAsync(default); if (this.Current.Succeeded) { this.State = this.Current.Result.State; } - this.hasStarted = true; - return true; } } diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Partitions/BackendQueryDataSource.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/BackendQueryDataSource.cs similarity index 96% rename from Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Partitions/BackendQueryDataSource.cs rename to Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/BackendQueryDataSource.cs index ced7eb235d..72f4af5fd6 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Partitions/BackendQueryDataSource.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/BackendQueryDataSource.cs @@ -2,7 +2,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // ------------------------------------------------------------ -namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.Partitions +namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote { using System; using System.Threading; diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Partitions/IQueryDataSource.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/IQueryDataSource.cs similarity index 90% rename from Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Partitions/IQueryDataSource.cs rename to Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/IQueryDataSource.cs index 5ac3ee6380..b2761f7c18 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Partitions/IQueryDataSource.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/IQueryDataSource.cs @@ -2,7 +2,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // ------------------------------------------------------------ -namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.Partitions +namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote { using System.Threading; using System.Threading.Tasks; diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/ParallelCrossPartitionQueryPipelineStage.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/ParallelCrossPartitionQueryPipelineStage.cs new file mode 100644 index 0000000000..b67d7aa306 --- /dev/null +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/ParallelCrossPartitionQueryPipelineStage.cs @@ -0,0 +1,169 @@ +// ------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +// ------------------------------------------------------------ + +namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote +{ + using System; + using System.Collections.Generic; + using System.Linq; + using System.Threading; + using System.Threading.Tasks; + using Microsoft.Azure.Cosmos.CosmosElements; + using Microsoft.Azure.Cosmos.Pagination; + using Microsoft.Azure.Cosmos.Query.Core.ContinuationTokens; + using Microsoft.Azure.Cosmos.Query.Core.Exceptions; + using Microsoft.Azure.Cosmos.Query.Core.Monads; + + internal sealed class ParallelCrossPartitionQueryPipelineStage : IQueryPipelineStage + { + private readonly CrossPartitionRangePageEnumerator crossPartitionRangePageEnumerator; + + private ParallelCrossPartitionQueryPipelineStage(CrossPartitionRangePageEnumerator crossPartitionRangePageEnumerator) + { + this.crossPartitionRangePageEnumerator = crossPartitionRangePageEnumerator ?? throw new ArgumentNullException(nameof(crossPartitionRangePageEnumerator)); + } + + public TryCatch Current + { + get + { + TryCatch> currentCrossPartitionPage = this.crossPartitionRangePageEnumerator.Current; + if (currentCrossPartitionPage.Failed) + { + return TryCatch.FromException(currentCrossPartitionPage.Exception); + } + + CrossPartitionPage crossPartitionPageResult = currentCrossPartitionPage.Result; + QueryPage backendQueryPage = crossPartitionPageResult.Page; + CrossPartitionState crossPartitionState = crossPartitionPageResult.State; + + List compositeContinuationTokens = new List(crossPartitionState.Value.Count); + foreach ((FeedRange range, QueryState state) in crossPartitionState.Value) + { + if (!(range is FeedRangeEpk epkRange)) + { + throw new InvalidOperationException(); + } + + CompositeContinuationToken compositeContinuationToken = new CompositeContinuationToken() + { + Range = epkRange.Range, + Token = ((CosmosString)state.Value).Value, + }; + + compositeContinuationTokens.Add(compositeContinuationToken); + } + + List cosmosElementContinuationTokens = compositeContinuationTokens.Select(token => CompositeContinuationToken.ToCosmosElement(token)).ToList(); + CosmosArray cosmosElementCompositeContinuationTokens = CosmosArray.Create(cosmosElementContinuationTokens); + + QueryPage crossPartitionQueryPage = new QueryPage( + backendQueryPage.Documents, + backendQueryPage.RequestCharge, + backendQueryPage.ActivityId, + backendQueryPage.ResponseLengthInBytes, + backendQueryPage.CosmosQueryExecutionInfo, + backendQueryPage.DisallowContinuationTokenMessage, + new QueryState(cosmosElementCompositeContinuationTokens)); + + return TryCatch.FromResult(crossPartitionQueryPage); + } + } + + public ValueTask DisposeAsync() => this.crossPartitionRangePageEnumerator.DisposeAsync(); + + public ValueTask MoveNextAsync() => this.crossPartitionRangePageEnumerator.MoveNextAsync(); + + public static TryCatch MonadicCreate( + IFeedRangeProvider feedRangeProvider, + IQueryDataSource queryDataSource, + SqlQuerySpec sqlQuerySpec, + int pageSize, + CosmosElement continuationToken) + { + CrossPartitionState state; + if (continuationToken == null) + { + state = default; + } + else + { + if (!(continuationToken is CosmosArray compositeContinuationTokenListRaw)) + { + return TryCatch.FromException( + new MalformedContinuationTokenException( + $"Invalid format for continuation token {continuationToken} for {nameof(ParallelCrossPartitionQueryPipelineStage)}")); + } + + List compositeContinuationTokens = new List(); + foreach (CosmosElement compositeContinuationTokenRaw in compositeContinuationTokenListRaw) + { + TryCatch tryCreateCompositeContinuationToken = CompositeContinuationToken.TryCreateFromCosmosElement(compositeContinuationTokenRaw); + if (tryCreateCompositeContinuationToken.Failed) + { + return TryCatch.FromException( + tryCreateCompositeContinuationToken.Exception); + } + + compositeContinuationTokens.Add(tryCreateCompositeContinuationToken.Result); + } + + List<(FeedRange, QueryState)> rangesAndStates = compositeContinuationTokens + .Select(token => ((FeedRange)new FeedRangeEpk(token.Range), new QueryState(CosmosString.Create(token.Token)))) + .ToList(); + + state = new CrossPartitionState(rangesAndStates); + } + + CrossPartitionRangePageEnumerator crossPartitionPageEnumerator = new CrossPartitionRangePageEnumerator( + feedRangeProvider, + ParallelCrossPartitionQueryPipelineStage.MakeCreateFunction(queryDataSource, sqlQuerySpec, pageSize), + Comparer.Singleton, + state: state); + + ParallelCrossPartitionQueryPipelineStage stage = new ParallelCrossPartitionQueryPipelineStage(crossPartitionPageEnumerator); + return TryCatch.FromResult(stage); + } + + private static CreatePartitionRangePageEnumerator MakeCreateFunction( + IQueryDataSource queryDataSource, + SqlQuerySpec sqlQuerySpec, + int pageSize) => (FeedRange range, QueryState state) => new QueryPartitionRangePageEnumerator( + queryDataSource, + sqlQuerySpec, + range, + pageSize, + state); + + private sealed class Comparer : IComparer> + { + public static readonly Comparer Singleton = new Comparer(); + + public int Compare( + PartitionRangePageEnumerator partitionRangePageEnumerator1, + PartitionRangePageEnumerator partitionRangePageEnumerator2) + { + if (object.ReferenceEquals(partitionRangePageEnumerator1, partitionRangePageEnumerator2)) + { + return 0; + } + + if (partitionRangePageEnumerator1.HasMoreResults && !partitionRangePageEnumerator2.HasMoreResults) + { + return -1; + } + + if (!partitionRangePageEnumerator1.HasMoreResults && partitionRangePageEnumerator2.HasMoreResults) + { + return 1; + } + + // Either both don't have results or both do. + return string.CompareOrdinal( + ((FeedRangePartitionKeyRange)partitionRangePageEnumerator1.Range).PartitionKeyRangeId, + ((FeedRangePartitionKeyRange)partitionRangePageEnumerator2.Range).PartitionKeyRangeId); + } + } + } +} \ No newline at end of file diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Partitions/QueryPartitionRangePageEnumerator.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/QueryPartitionRangePageEnumerator.cs similarity index 94% rename from Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Partitions/QueryPartitionRangePageEnumerator.cs rename to Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/QueryPartitionRangePageEnumerator.cs index 8025ef702c..c904ad16c3 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Partitions/QueryPartitionRangePageEnumerator.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/QueryPartitionRangePageEnumerator.cs @@ -2,7 +2,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // ------------------------------------------------------------ -namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.Partitions +namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote { using System; using System.Threading; @@ -10,7 +10,6 @@ namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.Partitions using Microsoft.Azure.Cosmos.CosmosElements; using Microsoft.Azure.Cosmos.Pagination; using Microsoft.Azure.Cosmos.Query.Core.Monads; - using Microsoft.Azure.Cosmos.SqlObjects; internal sealed class QueryPartitionRangePageEnumerator : PartitionRangePageEnumerator { diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Skip/SkipQueryPipelineStage.Client.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Skip/SkipQueryPipelineStage.Client.cs index b86ec80bc8..a8bf2b3676 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Skip/SkipQueryPipelineStage.Client.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Skip/SkipQueryPipelineStage.Client.cs @@ -25,14 +25,14 @@ private ClientSkipQueryPipelineStage(IQueryPipelineStage source, long skipCount) // Work is done in base constructor. } - public static async Task> TryCreateAsync( + public static TryCatch MonadicCreate( int offsetCount, CosmosElement continuationToken, - Func>> tryCreateSourceAsync) + MonadicCreatePipelineStage monadicCreatePipelineStage) { - if (tryCreateSourceAsync == null) + if (monadicCreatePipelineStage == null) { - throw new ArgumentNullException(nameof(tryCreateSourceAsync)); + throw new ArgumentNullException(nameof(monadicCreatePipelineStage)); } OffsetContinuationToken offsetContinuationToken; @@ -76,7 +76,7 @@ public static async Task> TryCreateAsync( sourceToken = null; } - TryCatch tryCreateSource = await tryCreateSourceAsync(sourceToken); + TryCatch tryCreateSource = monadicCreatePipelineStage(sourceToken); if (tryCreateSource.Failed) { return tryCreateSource; diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Skip/SkipQueryPipelineStage.Compute.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Skip/SkipQueryPipelineStage.Compute.cs index 0c5f56c4b1..fbfa845766 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Skip/SkipQueryPipelineStage.Compute.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Skip/SkipQueryPipelineStage.Compute.cs @@ -24,14 +24,14 @@ private ComputeSkipQueryPipelineStage(IQueryPipelineStage source, long skipCount // Work is done in base constructor. } - public static async Task> TryCreateAsync( + public static TryCatch MonadicCreate( int offsetCount, CosmosElement continuationToken, - Func>> tryCreateSourceAsync) + MonadicCreatePipelineStage monadicCreatePipelineStage) { - if (tryCreateSourceAsync == null) + if (monadicCreatePipelineStage == null) { - throw new ArgumentNullException(nameof(tryCreateSourceAsync)); + throw new ArgumentNullException(nameof(monadicCreatePipelineStage)); } OffsetContinuationToken offsetContinuationToken; @@ -58,7 +58,7 @@ public static async Task> TryCreateAsync( "offset count in continuation token can not be greater than the offsetcount in the query.")); } - TryCatch tryCreateSource = await tryCreateSourceAsync(offsetContinuationToken.SourceToken); + TryCatch tryCreateSource = monadicCreatePipelineStage(offsetContinuationToken.SourceToken); if (tryCreateSource.Failed) { return tryCreateSource; diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Skip/SkipQueryPipelineStage.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Skip/SkipQueryPipelineStage.cs index 3448218e5b..adf84a9fb6 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Skip/SkipQueryPipelineStage.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Skip/SkipQueryPipelineStage.cs @@ -7,6 +7,7 @@ namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.Skip using System; using System.Collections.Generic; using System.Text; + using System.Threading; using System.Threading.Tasks; using Microsoft.Azure.Cosmos.CosmosElements; using Microsoft.Azure.Cosmos.Query.Core.ExecutionContext; @@ -29,22 +30,22 @@ protected SkipQueryPipelineStage( this.skipCount = (int)skipCount; } - public static Task> TryCreateAsync( + public static TryCatch MonadicCreate( ExecutionEnvironment executionEnvironment, int offsetCount, CosmosElement continuationToken, - Func>> tryCreateSourceAsync) + MonadicCreatePipelineStage monadicCreatePipelineStage) { - Task> tryCreate = executionEnvironment switch + TryCatch tryCreate = executionEnvironment switch { - ExecutionEnvironment.Client => ClientSkipQueryPipelineStage.TryCreateAsync( + ExecutionEnvironment.Client => ClientSkipQueryPipelineStage.MonadicCreate( offsetCount, continuationToken, - tryCreateSourceAsync), - ExecutionEnvironment.Compute => ComputeSkipQueryPipelineStage.TryCreateAsync( + monadicCreatePipelineStage), + ExecutionEnvironment.Compute => ComputeSkipQueryPipelineStage.MonadicCreate( offsetCount, continuationToken, - tryCreateSourceAsync), + monadicCreatePipelineStage), _ => throw new ArgumentException($"Unknown {nameof(ExecutionEnvironment)}: {executionEnvironment}"), }; diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Take/TakeQueryPipelineStage.Client.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Take/TakeQueryPipelineStage.Client.cs index 9ffcb2c6ec..b760d14895 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Take/TakeQueryPipelineStage.Client.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Take/TakeQueryPipelineStage.Client.cs @@ -29,19 +29,19 @@ private ClientTakeQueryPipelineStage( this.takeEnum = takeEnum; } - public static async Task> TryCreateLimitStageAsync( + public static TryCatch MonadicCreateLimitStage( int limitCount, CosmosElement requestContinuationToken, - Func>> tryCreateSourceAsync) + MonadicCreatePipelineStage monadicCreatePipelineStage) { if (limitCount < 0) { throw new ArgumentException($"{nameof(limitCount)}: {limitCount} must be a non negative number."); } - if (tryCreateSourceAsync == null) + if (monadicCreatePipelineStage == null) { - throw new ArgumentNullException(nameof(tryCreateSourceAsync)); + throw new ArgumentNullException(nameof(monadicCreatePipelineStage)); } LimitContinuationToken limitContinuationToken; @@ -85,7 +85,7 @@ public static async Task> TryCreateLimitStageAsync sourceToken = null; } - TryCatch tryCreateSource = await tryCreateSourceAsync(sourceToken); + TryCatch tryCreateSource = monadicCreatePipelineStage(sourceToken); if (tryCreateSource.Failed) { return tryCreateSource; @@ -99,19 +99,19 @@ public static async Task> TryCreateLimitStageAsync return TryCatch.FromResult(stage); } - public static async Task> TryCreateTopStageAsync( + public static TryCatch MonadicCreateTopStage( int topCount, CosmosElement requestContinuationToken, - Func>> tryCreateSourceAsync) + MonadicCreatePipelineStage monadicCreatePipelineStage) { if (topCount < 0) { throw new ArgumentException($"{nameof(topCount)}: {topCount} must be a non negative number."); } - if (tryCreateSourceAsync == null) + if (monadicCreatePipelineStage == null) { - throw new ArgumentNullException(nameof(tryCreateSourceAsync)); + throw new ArgumentNullException(nameof(monadicCreatePipelineStage)); } TopContinuationToken topContinuationToken; @@ -155,7 +155,7 @@ public static async Task> TryCreateTopStageAsync( sourceToken = null; } - TryCatch tryCreateSource = await tryCreateSourceAsync(sourceToken); + TryCatch tryCreateSource = monadicCreatePipelineStage(sourceToken); if (tryCreateSource.Failed) { return tryCreateSource; diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Take/TakeQueryPipelineStage.Compute.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Take/TakeQueryPipelineStage.Compute.cs index 1d0693ae42..ca1057a815 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Take/TakeQueryPipelineStage.Compute.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Take/TakeQueryPipelineStage.Compute.cs @@ -27,35 +27,35 @@ private ComputeTakeQueryPipelineStage( // Work is done in the base class. } - public static Task> TryCreateLimitStageAsync( + public static TryCatch MonadicCreateLimitStage( int takeCount, CosmosElement requestContinuationToken, - Func>> tryCreateSourceAsync) => ComputeTakeQueryPipelineStage.TryCreateAsync( + MonadicCreatePipelineStage monadicCreatePipelineStage) => ComputeTakeQueryPipelineStage.MonadicCreate( takeCount, requestContinuationToken, - tryCreateSourceAsync); + monadicCreatePipelineStage); - public static Task> TryCreateTopStageAsync( + public static TryCatch MonadicCreateTopStage( int takeCount, CosmosElement requestContinuationToken, - Func>> tryCreateSourceAsync) => ComputeTakeQueryPipelineStage.TryCreateAsync( + MonadicCreatePipelineStage monadicCreatePipelineStage) => ComputeTakeQueryPipelineStage.MonadicCreate( takeCount, requestContinuationToken, - tryCreateSourceAsync); + monadicCreatePipelineStage); - private static async Task> TryCreateAsync( + private static TryCatch MonadicCreate( int takeCount, CosmosElement requestContinuationToken, - Func>> tryCreateSourceAsync) + MonadicCreatePipelineStage monadicCreatePipelineStage) { if (takeCount < 0) { throw new ArgumentException($"{nameof(takeCount)}: {takeCount} must be a non negative number."); } - if (tryCreateSourceAsync == null) + if (monadicCreatePipelineStage == null) { - throw new ArgumentNullException(nameof(tryCreateSourceAsync)); + throw new ArgumentNullException(nameof(monadicCreatePipelineStage)); } TakeContinuationToken takeContinuationToken; @@ -80,7 +80,7 @@ private static async Task> TryCreateAsync( $"{nameof(TakeContinuationToken.TakeCount)} in {nameof(TakeContinuationToken)}: {requestContinuationToken}: {takeContinuationToken.TakeCount} can not be greater than the limit count in the query: {takeCount}.")); } - TryCatch tryCreateSource = await tryCreateSourceAsync(takeContinuationToken.SourceToken); + TryCatch tryCreateSource = monadicCreatePipelineStage(takeContinuationToken.SourceToken); if (tryCreateSource.Failed) { return tryCreateSource; diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Take/TakeQueryPipelineStage.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Take/TakeQueryPipelineStage.cs index 401891eee8..2eea20bee8 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Take/TakeQueryPipelineStage.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Take/TakeQueryPipelineStage.cs @@ -5,6 +5,7 @@ namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.Take { using System; + using System.Threading; using System.Threading.Tasks; using Microsoft.Azure.Cosmos.CosmosElements; using Microsoft.Azure.Cosmos.Query.Core.ExecutionContext; @@ -22,37 +23,37 @@ protected TakeQueryPipelineStage( this.takeCount = takeCount; } - public static Task> TryCreateLimitStageAsync( + public static TryCatch MonadicCreateLimitStage( ExecutionEnvironment executionEnvironment, int limitCount, CosmosElement requestContinuationToken, - Func>> tryCreateSourceAsync) => executionEnvironment switch + MonadicCreatePipelineStage monadicCreatePipelineStage) => executionEnvironment switch { - ExecutionEnvironment.Client => ClientTakeQueryPipelineStage.TryCreateLimitStageAsync( + ExecutionEnvironment.Client => ClientTakeQueryPipelineStage.MonadicCreateLimitStage( limitCount, requestContinuationToken, - tryCreateSourceAsync), - ExecutionEnvironment.Compute => ComputeTakeQueryPipelineStage.TryCreateLimitStageAsync( + monadicCreatePipelineStage), + ExecutionEnvironment.Compute => ComputeTakeQueryPipelineStage.MonadicCreateLimitStage( limitCount, requestContinuationToken, - tryCreateSourceAsync), + monadicCreatePipelineStage), _ => throw new ArgumentOutOfRangeException($"Unknown {nameof(ExecutionEnvironment)}: {executionEnvironment}."), }; - public static Task> TryCreateTopStageAsync( + public static TryCatch MonadicCreateTopStage( ExecutionEnvironment executionEnvironment, int limitCount, CosmosElement requestContinuationToken, - Func>> tryCreateSourceAsync) => executionEnvironment switch + MonadicCreatePipelineStage monadicCreatePipelineStage) => executionEnvironment switch { - ExecutionEnvironment.Client => ClientTakeQueryPipelineStage.TryCreateTopStageAsync( + ExecutionEnvironment.Client => ClientTakeQueryPipelineStage.MonadicCreateTopStage( limitCount, requestContinuationToken, - tryCreateSourceAsync), - ExecutionEnvironment.Compute => ComputeTakeQueryPipelineStage.TryCreateTopStageAsync( + monadicCreatePipelineStage), + ExecutionEnvironment.Compute => ComputeTakeQueryPipelineStage.MonadicCreateTopStage( limitCount, requestContinuationToken, - tryCreateSourceAsync), + monadicCreatePipelineStage), _ => throw new ArgumentOutOfRangeException($"Unknown {nameof(ExecutionEnvironment)}: {executionEnvironment}."), }; } diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/QueryPlan/QueryInfo.cs b/Microsoft.Azure.Cosmos/src/Query/Core/QueryPlan/QueryInfo.cs index 7ff6ca4d10..9692b0f1d2 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/QueryPlan/QueryInfo.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/QueryPlan/QueryInfo.cs @@ -100,20 +100,8 @@ public bool HasSelectValue set; } - public bool HasDistinct - { - get - { - return this.DistinctType != DistinctQueryType.None; - } - } - public bool HasTop - { - get - { - return this.Top != null; - } - } + public bool HasDistinct => this.DistinctType != DistinctQueryType.None; + public bool HasTop => this.Top.HasValue; public bool HasAggregates { @@ -133,36 +121,12 @@ public bool HasAggregates } } - public bool HasGroupBy - { - get - { - return this.GroupByExpressions != null && this.GroupByExpressions.Count > 0; - } - } + public bool HasGroupBy => (this.GroupByExpressions != null) && (this.GroupByExpressions.Count > 0); - public bool HasOrderBy - { - get - { - return this.OrderBy != null && this.OrderBy.Count > 0; - } - } + public bool HasOrderBy => (this.OrderBy != null) && (this.OrderBy.Count > 0); - public bool HasOffset - { - get - { - return this.Offset != null; - } - } + public bool HasOffset => this.Offset.HasValue; - public bool HasLimit - { - get - { - return this.Limit != null; - } - } + public bool HasLimit => this.Limit.HasValue; } } \ No newline at end of file diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/CosmosQueryUnitTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/CosmosQueryUnitTests.cs index 10e813239e..facd287640 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/CosmosQueryUnitTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/CosmosQueryUnitTests.cs @@ -308,7 +308,7 @@ public async Task TestCosmosQueryPartitionKeyDefinition() AggregateOperator.Sum }; - components.Add((await AggregateDocumentQueryExecutionComponent.TryCreateAsync( + components.Add((await AggregateDocumentQueryExecutionComponent.MonadicCreateAsync( ExecutionEnvironment.Client, operators.ToArray(), new Dictionary() @@ -320,25 +320,25 @@ public async Task TestCosmosQueryPartitionKeyDefinition() null, func)).Result); - components.Add((await DistinctDocumentQueryExecutionComponent.TryCreateAsync( + components.Add((await DistinctDocumentQueryExecutionComponent.MonadicCreateAsync( ExecutionEnvironment.Client, null, func, DistinctQueryType.Ordered)).Result); - components.Add((await SkipDocumentQueryExecutionComponent.TryCreateAsync( + components.Add((await SkipDocumentQueryExecutionComponent.MonadicCreateAsync( ExecutionEnvironment.Client, 5, null, func)).Result); - components.Add((await TakeDocumentQueryExecutionComponent.TryCreateAsync( + components.Add((await TakeDocumentQueryExecutionComponent.MonadicCreateAsync( ExecutionEnvironment.Client, 5, null, func)).Result); - components.Add((await TakeDocumentQueryExecutionComponent.TryCreateAsync( + components.Add((await TakeDocumentQueryExecutionComponent.MonadicCreateAsync( ExecutionEnvironment.Client, 5, null, diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/InMemoryCollectionQueryDataSource.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/InMemoryCollectionQueryDataSource.cs index 2f934f9e68..ca9e1e2890 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/InMemoryCollectionQueryDataSource.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/InMemoryCollectionQueryDataSource.cs @@ -13,7 +13,7 @@ namespace Microsoft.Azure.Cosmos.Tests.Query using Microsoft.Azure.Cosmos.Query.Core; using Microsoft.Azure.Cosmos.Query.Core.Monads; using Microsoft.Azure.Cosmos.Query.Core.Pipeline; - using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Partitions; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote; internal sealed class InMemoryCollectionQueryDataSource : IQueryDataSource { diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/AggregateQueryPipelineStageTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/AggregateQueryPipelineStageTests.cs index 8da2f7b7c5..fbd9d6283a 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/AggregateQueryPipelineStageTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/AggregateQueryPipelineStageTests.cs @@ -7,6 +7,7 @@ namespace Microsoft.Azure.Cosmos.Tests.Query.Pipeline using System; using System.Collections.Generic; using System.Linq; + using System.Threading; using System.Threading.Tasks; using Microsoft.Azure.Cosmos.CosmosElements; using Microsoft.Azure.Cosmos.Query.Core.ExecutionContext; @@ -101,14 +102,14 @@ private static async Task> CreateAndDrain( { IQueryPipelineStage source = new MockQueryPipelineStage(pages); - TryCatch tryCreateAggregateQueryPipelineStage = await AggregateQueryPipelineStage.TryCreateAsync( + TryCatch tryCreateAggregateQueryPipelineStage = AggregateQueryPipelineStage.MonadicCreate( executionEnvironment: executionEnvironment, aggregates: aggregates, aliasToAggregateType: aliasToAggregateType, orderedAliases: orderedAliases, hasSelectValue: hasSelectValue, continuationToken: continuationToken, - tryCreateSourceAsync: (CosmosElement continuationToken) => Task.FromResult(TryCatch.FromResult(source))); + monadicCreatePipelineStage: (CosmosElement continuationToken) => TryCatch.FromResult(source)); Assert.IsTrue(tryCreateAggregateQueryPipelineStage.Succeeded); IQueryPipelineStage aggregateQueryPipelineStage = tryCreateAggregateQueryPipelineStage.Result; diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/DistinctQueryPipelineStageTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/DistinctQueryPipelineStageTests.cs index 6479b4f638..974669e1c0 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/DistinctQueryPipelineStageTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/DistinctQueryPipelineStageTests.cs @@ -42,11 +42,11 @@ private static async Task> CreateAndDrainAsync( { IQueryPipelineStage source = new MockQueryPipelineStage(pages); - TryCatch tryCreateDistinctQueryPipelineStage = await DistinctQueryPipelineStage.TryCreateAsync( + TryCatch tryCreateDistinctQueryPipelineStage = DistinctQueryPipelineStage.MonadicCreate( executionEnvironment: executionEnvironment, requestContinuation: continuationToken, distinctQueryType: distinctQueryType, - tryCreateSourceAsync: (CosmosElement continuationToken) => Task.FromResult(TryCatch.FromResult(source))); + monadicCreatePipelineStage: (CosmosElement continuationToken) => TryCatch.FromResult(source)); Assert.IsTrue(tryCreateDistinctQueryPipelineStage.Succeeded); IQueryPipelineStage distinctQueryPipelineStage = tryCreateDistinctQueryPipelineStage.Result; 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 new file mode 100644 index 0000000000..dd08f04005 --- /dev/null +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/FactoryTests.cs @@ -0,0 +1,40 @@ +//------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +//------------------------------------------------------------ + +namespace Microsoft.Azure.Cosmos.Tests.Query.Pipeline +{ + using Microsoft.Azure.Cosmos.Pagination; + 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.Remote; + using Microsoft.Azure.Cosmos.Query.Core.QueryPlan; + using Microsoft.VisualStudio.TestTools.UnitTesting; + using Moq; + + [TestClass] + public class FactoryTests + { + [TestMethod] + public void TestCreate() + { + Mock mockFeedRangeProvider = new Mock(); + Mock mockQueryDataSource = new Mock(); + + TryCatch monadicCreatePipeline = PipelineFactory.MonadicCreate( + ExecutionEnvironment.Compute, + feedRangeProvider: mockFeedRangeProvider.Object, + queryDataSource: mockQueryDataSource.Object, + sqlQuerySpec: new SqlQuerySpec("SELECT * FROM c"), + queryInfo: new QueryInfo() { }, + pageSize: 10, + requestContinuationToken: default); + Assert.IsTrue(monadicCreatePipeline.Succeeded); + + IQueryPipelineStage pipelineStage = monadicCreatePipeline.Result; + Assert.IsTrue(pipelineStage is ParallelCrossPartitionQueryPipelineStage); + } + } +} diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/GroupByQueryPipelineStageTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/GroupByQueryPipelineStageTests.cs index e386129002..63535acea0 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/GroupByQueryPipelineStageTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/GroupByQueryPipelineStageTests.cs @@ -52,10 +52,10 @@ private static async Task> CreateAndDrainAsync( { IQueryPipelineStage source = new MockQueryPipelineStage(pages); - TryCatch tryCreateGroupByStage = await GroupByQueryPipelineStage.TryCreateAsync( + TryCatch tryCreateGroupByStage = GroupByQueryPipelineStage.MonadicCreate( executionEnvironment: executionEnvironment, continuationToken: continuationToken, - tryCreateSourceAsync: (CosmosElement continuationToken) => Task.FromResult(TryCatch.FromResult(source)), + monadicCreatePipelineStage: (CosmosElement continuationToken) => TryCatch.FromResult(source), groupByAliasToAggregateType: groupByAliasToAggregateType, orderedAliases: orderedAliases, hasSelectValue: hasSelectValue); diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/SkipQueryPipelineStageTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/SkipQueryPipelineStageTests.cs index baa87ed9a3..80820a5de1 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/SkipQueryPipelineStageTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/SkipQueryPipelineStageTests.cs @@ -45,11 +45,11 @@ private static async Task> CreateAndDrainAsync( { IQueryPipelineStage source = new MockQueryPipelineStage(pages); - TryCatch tryCreateSkipQueryPipelineStage = await SkipQueryPipelineStage.TryCreateAsync( + TryCatch tryCreateSkipQueryPipelineStage = SkipQueryPipelineStage.MonadicCreate( executionEnvironment: executionEnvironment, offsetCount: offsetCount, continuationToken: continuationToken, - tryCreateSourceAsync: (CosmosElement continuationToken) => Task.FromResult(TryCatch.FromResult(source))); + monadicCreatePipelineStage: (CosmosElement continuationToken) => TryCatch.FromResult(source)); Assert.IsTrue(tryCreateSkipQueryPipelineStage.Succeeded); IQueryPipelineStage aggregateQueryPipelineStage = tryCreateSkipQueryPipelineStage.Result; diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/TakeQueryPipelineStageTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/TakeQueryPipelineStageTests.cs index fcbc04eda2..9fa35f087c 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/TakeQueryPipelineStageTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/TakeQueryPipelineStageTests.cs @@ -45,11 +45,11 @@ private static async Task> CreateAndDrainAsync( { IQueryPipelineStage source = new MockQueryPipelineStage(pages); - TryCatch tryCreateSkipQueryPipelineStage = await TakeQueryPipelineStage.TryCreateLimitStageAsync( + TryCatch tryCreateSkipQueryPipelineStage = TakeQueryPipelineStage.MonadicCreateLimitStage( executionEnvironment: executionEnvironment, limitCount: takeCount, requestContinuationToken: continuationToken, - tryCreateSourceAsync: (CosmosElement continuationToken) => Task.FromResult(TryCatch.FromResult(source))); + monadicCreatePipelineStage: (CosmosElement continuationToken) => TryCatch.FromResult(source)); Assert.IsTrue(tryCreateSkipQueryPipelineStage.Succeeded); IQueryPipelineStage takeQueryPipelineStage = tryCreateSkipQueryPipelineStage.Result; 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 bff9aac39b..48eb7fd0b9 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 @@ -5,10 +5,9 @@ using System.Threading.Tasks; using Microsoft.Azure.Cosmos.CosmosElements; using Microsoft.Azure.Cosmos.Pagination; - using Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.ItemProducers; using Microsoft.Azure.Cosmos.Query.Core.Monads; using Microsoft.Azure.Cosmos.Query.Core.Pipeline; - using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Partitions; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote; using Microsoft.Azure.Cosmos.Tests.Pagination; using Microsoft.Azure.Documents; using Microsoft.VisualStudio.TestTools.UnitTesting; From 12e8849bae00b8fffa145949ac8763acfa53d7c4 Mon Sep 17 00:00:00 2001 From: Brandon Chong Date: Mon, 6 Jul 2020 17:19:30 -0700 Subject: [PATCH 28/85] wired through query stages --- .../CatchAllCosmosQueryExecutionContext.cs | 70 --- .../CosmosQueryExecutionContextFactory.cs | 151 ++++--- ...ExecutionContextWithNameCacheStaleRetry.cs | 69 --- .../LazyCosmosQueryExecutionContext.cs | 98 ----- .../PipelinedDocumentQueryExecutionContext.cs | 410 ------------------ .../Query/Core/Monads/TryCatch{TResult}.cs | 14 + .../AggregateQueryPipelineStage.Compute.cs | 2 +- .../Pipeline/CatchAllQueryPipelineStage.cs | 40 ++ .../Core/Pipeline/EmptyQueryPipelineStage.cs | 28 ++ .../Pipeline/FaultedQueryPipelineStage.cs | 32 ++ .../Pipeline/FinishedQueryPipelineStage.cs | 30 -- .../GroupByQueryPipelineStage.Compute.cs | 2 +- .../Core/Pipeline/LazyQueryPipelineStage.cs | 55 +++ .../NameCacheStaleRetryQueryPipelineStage.cs | 66 +++ .../src/Query/v3Query/QueryIterator.cs | 74 ++-- .../src/Reactive/EmptyAsyncEnumerator.cs | 27 ++ .../src/Reactive/JustAsyncEnumerator.cs | 30 ++ .../CosmosQueryUnitTests.cs | 13 +- .../Query/Pipeline/MockQueryPipelineStage.cs | 2 +- 19 files changed, 429 insertions(+), 784 deletions(-) delete mode 100644 Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CatchAllCosmosQueryExecutionContext.cs delete mode 100644 Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosQueryExecutionContextWithNameCacheStaleRetry.cs delete mode 100644 Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/LazyCosmosQueryExecutionContext.cs delete mode 100644 Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/PipelinedDocumentQueryExecutionContext.cs create mode 100644 Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CatchAllQueryPipelineStage.cs create mode 100644 Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/EmptyQueryPipelineStage.cs create mode 100644 Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/FaultedQueryPipelineStage.cs delete mode 100644 Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/FinishedQueryPipelineStage.cs create mode 100644 Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/LazyQueryPipelineStage.cs create mode 100644 Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/NameCacheStaleRetryQueryPipelineStage.cs create mode 100644 Microsoft.Azure.Cosmos/src/Reactive/EmptyAsyncEnumerator.cs create mode 100644 Microsoft.Azure.Cosmos/src/Reactive/JustAsyncEnumerator.cs diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CatchAllCosmosQueryExecutionContext.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CatchAllCosmosQueryExecutionContext.cs deleted file mode 100644 index e37a658f46..0000000000 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CatchAllCosmosQueryExecutionContext.cs +++ /dev/null @@ -1,70 +0,0 @@ -// ------------------------------------------------------------ -// Copyright (c) Microsoft Corporation. All rights reserved. -// ------------------------------------------------------------ - -namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionContext -{ - using System; - using System.Threading; - using System.Threading.Tasks; - using Microsoft.Azure.Cosmos.CosmosElements; - using Microsoft.Azure.Cosmos.Json; - using Microsoft.Azure.Cosmos.Query.Core.QueryClient; - - internal sealed class CatchAllCosmosQueryExecutionContext : CosmosQueryExecutionContext - { - private readonly CosmosQueryExecutionContext cosmosQueryExecutionContext; - private bool hitException; - - public CatchAllCosmosQueryExecutionContext( - CosmosQueryExecutionContext cosmosQueryExecutionContext) - { - this.cosmosQueryExecutionContext = cosmosQueryExecutionContext ?? throw new ArgumentNullException(nameof(cosmosQueryExecutionContext)); - } - - public override bool IsDone => this.hitException || this.cosmosQueryExecutionContext.IsDone; - - public override void Dispose() - { - this.cosmosQueryExecutionContext.Dispose(); - } - - public override async Task ExecuteNextAsync(CancellationToken cancellationToken) - { - if (this.IsDone) - { - throw new InvalidOperationException( - $"Can not {nameof(ExecuteNextAsync)} from a {nameof(CosmosQueryExecutionContext)} where {nameof(this.IsDone)}."); - } - - cancellationToken.ThrowIfCancellationRequested(); - - QueryResponseCore queryResponseCore; - try - { - queryResponseCore = await this.cosmosQueryExecutionContext.ExecuteNextAsync(cancellationToken); - } - catch (OperationCanceledException) when (cancellationToken.IsCancellationRequested) - { - // Per cancellationToken.ThrowIfCancellationRequested(); line above, this function should still throw OperationCanceledException. - throw; - } - catch (Exception ex) - { - queryResponseCore = QueryResponseFactory.CreateFromException(ex); - } - - if (!queryResponseCore.IsSuccess) - { - this.hitException = true; - } - - return queryResponseCore; - } - - public override CosmosElement GetCosmosElementContinuationToken() - { - return this.cosmosQueryExecutionContext.GetCosmosElementContinuationToken(); - } - } -} diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosQueryExecutionContextFactory.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosQueryExecutionContextFactory.cs index 2f35d372a4..9d50f428c0 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosQueryExecutionContextFactory.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosQueryExecutionContextFactory.cs @@ -10,11 +10,14 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionContext using System.Threading.Tasks; using Microsoft.Azure.Cosmos; using Microsoft.Azure.Cosmos.CosmosElements; + using Microsoft.Azure.Cosmos.Pagination; using Microsoft.Azure.Cosmos.Query.Core; using Microsoft.Azure.Cosmos.Query.Core.ContinuationTokens; using Microsoft.Azure.Cosmos.Query.Core.Exceptions; using Microsoft.Azure.Cosmos.Query.Core.Monads; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline; using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Distinct; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote; using Microsoft.Azure.Cosmos.Query.Core.QueryClient; using Microsoft.Azure.Cosmos.Query.Core.QueryPlan; using Microsoft.Azure.Cosmos.SqlObjects; @@ -25,7 +28,7 @@ internal static class CosmosQueryExecutionContextFactory private const string InternalPartitionKeyDefinitionProperty = "x-ms-query-partitionkey-definition"; private const int PageSizeFactorForTop = 5; - public static CosmosQueryExecutionContext Create( + public static IQueryPipelineStage Create( CosmosQueryContext cosmosQueryContext, InputParameters inputParameters) { @@ -39,29 +42,26 @@ public static CosmosQueryExecutionContext Create( throw new ArgumentNullException(nameof(inputParameters)); } - CosmosQueryExecutionContextWithNameCacheStaleRetry cosmosQueryExecutionContextWithNameCacheStaleRetry = new CosmosQueryExecutionContextWithNameCacheStaleRetry( + NameCacheStaleRetryQueryPipelineStage nameCacheStaleRetryQueryPipelineStage = new NameCacheStaleRetryQueryPipelineStage( cosmosQueryContext: cosmosQueryContext, - cosmosQueryExecutionContextFactory: () => + queryPipelineStageFactory: () => { // Query Iterator requires that the creation of the query context is deferred until the user calls ReadNextAsync - AsyncLazy> lazyTryCreateCosmosQueryExecutionContext = new AsyncLazy>(valueFactory: (innerCancellationToken) => - { - innerCancellationToken.ThrowIfCancellationRequested(); - return CosmosQueryExecutionContextFactory.TryCreateCoreContextAsync( + AsyncLazy> lazyTryCreateStage = new AsyncLazy>( + valueFactory: (innerCancellationToken) => CosmosQueryExecutionContextFactory.TryCreateCoreContextAsync( cosmosQueryContext, inputParameters, - innerCancellationToken); - }); - LazyCosmosQueryExecutionContext lazyCosmosQueryExecutionContext = new LazyCosmosQueryExecutionContext(lazyTryCreateCosmosQueryExecutionContext); - return lazyCosmosQueryExecutionContext; - }); + innerCancellationToken)); - CatchAllCosmosQueryExecutionContext catchAllCosmosQueryExecutionContext = new CatchAllCosmosQueryExecutionContext(cosmosQueryExecutionContextWithNameCacheStaleRetry); + LazyQueryPipelineStage lazyQueryPipelineStage = new LazyQueryPipelineStage(lazyTryCreateStage: lazyTryCreateStage); + return lazyQueryPipelineStage; + }); - return catchAllCosmosQueryExecutionContext; + CatchAllQueryPipelineStage catchAllQueryPipelineStage = new CatchAllQueryPipelineStage(nameCacheStaleRetryQueryPipelineStage); + return catchAllQueryPipelineStage; } - private static async Task> TryCreateCoreContextAsync( + private static async Task> TryCreateCoreContextAsync( CosmosQueryContext cosmosQueryContext, InputParameters inputParameters, CancellationToken cancellationToken) @@ -78,14 +78,14 @@ private static async Task> TryCreateCoreCo continuationToken, out PipelineContinuationToken pipelineContinuationToken)) { - return TryCatch.FromException( + return TryCatch.FromException( new MalformedContinuationTokenException( $"Malformed {nameof(PipelineContinuationToken)}: {continuationToken}.")); } if (PipelineContinuationToken.IsTokenFromTheFuture(pipelineContinuationToken)) { - return TryCatch.FromException( + return TryCatch.FromException( new MalformedContinuationTokenException( $"{nameof(PipelineContinuationToken)} Continuation token is from a newer version of the SDK. " + $"Upgrade the SDK to avoid this issue." + @@ -96,7 +96,7 @@ private static async Task> TryCreateCoreCo pipelineContinuationToken, out PipelineContinuationTokenV1_1 latestVersionPipelineContinuationToken)) { - return TryCatch.FromException( + return TryCatch.FromException( new MalformedContinuationTokenException( $"{nameof(PipelineContinuationToken)}: '{continuationToken}' is no longer supported.")); } @@ -179,13 +179,12 @@ private static async Task> TryCreateCoreCo containerQueryProperties.ResourceId, inputParameters.PartitionKey.Value.InternalKey.GetEffectivePartitionKeyString(partitionKeyDefinition)); - return await CosmosQueryExecutionContextFactory.TryCreatePassthroughQueryExecutionContextAsync( + return CosmosQueryExecutionContextFactory.TryCreatePassthroughQueryExecutionContext( cosmosQueryContext, inputParameters, partitionedQueryExecutionInfo: new PartitionedQueryExecutionInfo(), targetRanges, - containerQueryProperties.ResourceId, - cancellationToken); + containerQueryProperties.ResourceId); } } } @@ -226,7 +225,7 @@ private static async Task> TryCreateCoreCo } } - public static async Task> TryCreateFromPartitionedQuerExecutionInfoAsync( + public static async Task> TryCreateFromPartitionedQuerExecutionInfoAsync( PartitionedQueryExecutionInfo partitionedQueryExecutionInfo, ContainerQueryProperties containerQueryProperties, CosmosQueryContext cosmosQueryContext, @@ -236,12 +235,12 @@ public static async Task> TryCreateFromPar cancellationToken.ThrowIfCancellationRequested(); List targetRanges = await CosmosQueryExecutionContextFactory.GetTargetPartitionKeyRangesAsync( - cosmosQueryContext.QueryClient, - cosmosQueryContext.ResourceLink.OriginalString, - partitionedQueryExecutionInfo, - containerQueryProperties, - inputParameters.Properties, - inputParameters.InitialFeedRange); + cosmosQueryContext.QueryClient, + cosmosQueryContext.ResourceLink.OriginalString, + partitionedQueryExecutionInfo, + containerQueryProperties, + inputParameters.Properties, + inputParameters.InitialFeedRange); bool singleLogicalPartitionKeyQuery = inputParameters.PartitionKey.HasValue || ((partitionedQueryExecutionInfo.QueryRanges.Count == 1) @@ -261,7 +260,7 @@ public static async Task> TryCreateFromPar bool createPassthoughQuery = streamingSinglePartitionQuery || streamingCrossContinuationQuery; - Task> tryCreateContextTask; + TryCatch tryCreatePipelineStage; if (createPassthoughQuery) { TestInjections.ResponseStats responseStats = inputParameters?.TestInjections?.Stats; @@ -270,13 +269,12 @@ public static async Task> TryCreateFromPar responseStats.PipelineType = TestInjections.PipelineType.Passthrough; } - tryCreateContextTask = CosmosQueryExecutionContextFactory.TryCreatePassthroughQueryExecutionContextAsync( + tryCreatePipelineStage = CosmosQueryExecutionContextFactory.TryCreatePassthroughQueryExecutionContext( cosmosQueryContext, inputParameters, partitionedQueryExecutionInfo, targetRanges, - containerQueryProperties.ResourceId, - cancellationToken); + containerQueryProperties.ResourceId); } else { @@ -311,25 +309,23 @@ public static async Task> TryCreateFromPar inputParameters.TestInjections); } - tryCreateContextTask = CosmosQueryExecutionContextFactory.TryCreateSpecializedDocumentQueryExecutionContextAsync( + tryCreatePipelineStage = CosmosQueryExecutionContextFactory.TryCreateSpecializedDocumentQueryExecutionContext( cosmosQueryContext, inputParameters, partitionedQueryExecutionInfo, targetRanges, - containerQueryProperties.ResourceId, - cancellationToken); + containerQueryProperties.ResourceId); } - return await tryCreateContextTask; + return tryCreatePipelineStage; } - private static Task> TryCreatePassthroughQueryExecutionContextAsync( + private static TryCatch TryCreatePassthroughQueryExecutionContext( CosmosQueryContext cosmosQueryContext, InputParameters inputParameters, PartitionedQueryExecutionInfo partitionedQueryExecutionInfo, List targetRanges, - string collectionRid, - CancellationToken cancellationToken) + string collectionRid) { CosmosCrossPartitionQueryExecutionContext.CrossPartitionInitParams initParams = new CosmosCrossPartitionQueryExecutionContext.CrossPartitionInitParams( sqlQuerySpec: inputParameters.SqlQuerySpec, @@ -343,26 +339,57 @@ private static Task> TryCreatePassthroughQ returnResultsInDeterministicOrder: inputParameters.ReturnResultsInDeterministicOrder, testSettings: inputParameters.TestInjections); - return PipelinedDocumentQueryExecutionContext.TryCreatePassthroughAsync( - inputParameters.ExecutionEnvironment, - cosmosQueryContext, - initParams, - inputParameters.InitialUserContinuationToken, - cancellationToken); + // Modify query plan + PartitionedQueryExecutionInfo passThroughQueryInfo = new PartitionedQueryExecutionInfo() + { + QueryInfo = new QueryInfo() + { + Aggregates = null, + DistinctType = DistinctQueryType.None, + GroupByAliases = null, + GroupByAliasToAggregateType = null, + GroupByExpressions = null, + HasSelectValue = false, + Limit = null, + Offset = null, + OrderBy = null, + OrderByExpressions = null, + RewrittenQuery = null, + Top = null, + }, + QueryRanges = initParams.PartitionedQueryExecutionInfo.QueryRanges, + }; + + initParams = new CosmosCrossPartitionQueryExecutionContext.CrossPartitionInitParams( + sqlQuerySpec: initParams.SqlQuerySpec, + collectionRid: initParams.CollectionRid, + partitionedQueryExecutionInfo: passThroughQueryInfo, + partitionKeyRanges: initParams.PartitionKeyRanges, + initialPageSize: initParams.MaxItemCount.GetValueOrDefault(1000), + maxConcurrency: initParams.MaxConcurrency, + maxItemCount: initParams.MaxItemCount, + maxBufferedItemCount: initParams.MaxBufferedItemCount, + returnResultsInDeterministicOrder: initParams.ReturnResultsInDeterministicOrder, + testSettings: initParams.TestSettings); + + // Return a parallel context, since we still want to be able to handle splits and concurrency / buffering. + return ParallelCrossPartitionQueryPipelineStage.MonadicCreate( + feedRangeProvider: new FeedRangeProvider(cosmosQueryContext.QueryClient, collectionRid), + queryDataSource: new BackendQueryDataSource(cosmosQueryContext), + sqlQuerySpec: inputParameters.SqlQuerySpec, + pageSize: initParams.InitialPageSize, + continuationToken: inputParameters.InitialUserContinuationToken); } - private static async Task> TryCreateSpecializedDocumentQueryExecutionContextAsync( + private static TryCatch TryCreateSpecializedDocumentQueryExecutionContext( CosmosQueryContext cosmosQueryContext, InputParameters inputParameters, PartitionedQueryExecutionInfo partitionedQueryExecutionInfo, List targetRanges, - string collectionRid, - CancellationToken cancellationToken) + string collectionRid) { QueryInfo queryInfo = partitionedQueryExecutionInfo.QueryInfo; - bool getLazyFeedResponse = queryInfo.HasTop; - // We need to compute the optimal initial page size for order-by queries long optimalPageSize = inputParameters.MaxItemCount; if (queryInfo.HasOrderBy) @@ -389,24 +416,14 @@ private static async Task> TryCreateSpecia (optimalPageSize > 0) && (optimalPageSize <= int.MaxValue), $"Invalid MaxItemCount {optimalPageSize}"); - CosmosCrossPartitionQueryExecutionContext.CrossPartitionInitParams initParams = new CosmosCrossPartitionQueryExecutionContext.CrossPartitionInitParams( + return PipelineFactory.MonadicCreate( + executionEnvironment: inputParameters.ExecutionEnvironment, + feedRangeProvider: new FeedRangeProvider(cosmosQueryContext.QueryClient, collectionRid), + queryDataSource: new BackendQueryDataSource(cosmosQueryContext), sqlQuerySpec: inputParameters.SqlQuerySpec, - collectionRid: collectionRid, - partitionedQueryExecutionInfo: partitionedQueryExecutionInfo, - partitionKeyRanges: targetRanges, - initialPageSize: (int)optimalPageSize, - maxConcurrency: inputParameters.MaxConcurrency, - maxItemCount: inputParameters.MaxItemCount, - maxBufferedItemCount: inputParameters.MaxBufferedItemCount, - returnResultsInDeterministicOrder: inputParameters.ReturnResultsInDeterministicOrder, - testSettings: inputParameters.TestInjections); - - return await PipelinedDocumentQueryExecutionContext.MonadicCreateAsync( - inputParameters.ExecutionEnvironment, - cosmosQueryContext, - initParams, - inputParameters.InitialUserContinuationToken, - cancellationToken); + queryInfo: partitionedQueryExecutionInfo.QueryInfo, + pageSize: (int)optimalPageSize, + requestContinuationToken: inputParameters.InitialUserContinuationToken); } /// diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosQueryExecutionContextWithNameCacheStaleRetry.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosQueryExecutionContextWithNameCacheStaleRetry.cs deleted file mode 100644 index 50be1cacd6..0000000000 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosQueryExecutionContextWithNameCacheStaleRetry.cs +++ /dev/null @@ -1,69 +0,0 @@ -// ------------------------------------------------------------ -// Copyright (c) Microsoft Corporation. All rights reserved. -// ------------------------------------------------------------ - -namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionContext -{ - using System; - using System.Threading; - using System.Threading.Tasks; - using Microsoft.Azure.Cosmos.CosmosElements; - using Microsoft.Azure.Cosmos.Json; - using Microsoft.Azure.Cosmos.Query.Core.QueryClient; - - internal sealed class CosmosQueryExecutionContextWithNameCacheStaleRetry : CosmosQueryExecutionContext - { - private readonly CosmosQueryContext cosmosQueryContext; - private readonly Func cosmosQueryExecutionContextFactory; - private CosmosQueryExecutionContext currentCosmosQueryExecutionContext; - private bool alreadyRetried; - - public CosmosQueryExecutionContextWithNameCacheStaleRetry( - CosmosQueryContext cosmosQueryContext, - Func cosmosQueryExecutionContextFactory) - { - this.cosmosQueryContext = cosmosQueryContext ?? throw new ArgumentNullException(nameof(cosmosQueryContext)); - this.cosmosQueryExecutionContextFactory = cosmosQueryExecutionContextFactory ?? throw new ArgumentNullException(nameof(cosmosQueryExecutionContextFactory)); - this.currentCosmosQueryExecutionContext = cosmosQueryExecutionContextFactory(); - } - - public override bool IsDone => this.currentCosmosQueryExecutionContext.IsDone; - - public override void Dispose() - { - this.currentCosmosQueryExecutionContext.Dispose(); - } - - public override async Task ExecuteNextAsync(CancellationToken cancellationToken) - { - cancellationToken.ThrowIfCancellationRequested(); - - // If the cache is stale the entire execute context has incorrect values and should be recreated. - // This should only be done for the first execution. - // If results have already been pulled, - // then an error should be returned to the user, - // since it's not possible to combine query results from multiple containers. - QueryResponseCore queryResponse = await this.currentCosmosQueryExecutionContext.ExecuteNextAsync(cancellationToken); - if ( - (queryResponse.StatusCode == System.Net.HttpStatusCode.Gone) && - (queryResponse.SubStatusCode == Documents.SubStatusCodes.NameCacheIsStale) && - !this.alreadyRetried) - { - await this.cosmosQueryContext.QueryClient.ForceRefreshCollectionCacheAsync( - this.cosmosQueryContext.ResourceLink.OriginalString, - cancellationToken); - this.alreadyRetried = true; - this.currentCosmosQueryExecutionContext.Dispose(); - this.currentCosmosQueryExecutionContext = this.cosmosQueryExecutionContextFactory(); - return await this.ExecuteNextAsync(cancellationToken); - } - - return queryResponse; - } - - public override CosmosElement GetCosmosElementContinuationToken() - { - return this.currentCosmosQueryExecutionContext.GetCosmosElementContinuationToken(); - } - } -} diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/LazyCosmosQueryExecutionContext.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/LazyCosmosQueryExecutionContext.cs deleted file mode 100644 index 1e7c61a02e..0000000000 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/LazyCosmosQueryExecutionContext.cs +++ /dev/null @@ -1,98 +0,0 @@ -// ------------------------------------------------------------ -// Copyright (c) Microsoft Corporation. All rights reserved. -// ------------------------------------------------------------ - -namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionContext -{ - using System; - using System.Collections.Generic; - using System.Threading; - using System.Threading.Tasks; - using Microsoft.Azure.Cosmos.CosmosElements; - using Microsoft.Azure.Cosmos.Diagnostics; - using Microsoft.Azure.Cosmos.Json; - using Microsoft.Azure.Cosmos.Query.Core.Monads; - using Microsoft.Azure.Cosmos.Query.Core.QueryClient; - - /// - /// Implementation of that composes another context and defers it's initialization until the first read. - /// - internal sealed class LazyCosmosQueryExecutionContext : CosmosQueryExecutionContext - { - private readonly AsyncLazy> lazyTryCreateCosmosQueryExecutionContext; - - public LazyCosmosQueryExecutionContext(AsyncLazy> lazyTryCreateCosmosQueryExecutionContext) - { - this.lazyTryCreateCosmosQueryExecutionContext = lazyTryCreateCosmosQueryExecutionContext ?? throw new ArgumentNullException(nameof(lazyTryCreateCosmosQueryExecutionContext)); - } - - public override bool IsDone - { - get - { - bool isDone; - if (this.lazyTryCreateCosmosQueryExecutionContext.ValueInitialized) - { - TryCatch tryCreateCosmosQueryExecutionContext = this.lazyTryCreateCosmosQueryExecutionContext.Result; - if (tryCreateCosmosQueryExecutionContext.Succeeded) - { - isDone = tryCreateCosmosQueryExecutionContext.Result.IsDone; - } - else - { - isDone = true; - } - } - else - { - isDone = false; - } - - return isDone; - } - } - - public override void Dispose() - { - if (this.lazyTryCreateCosmosQueryExecutionContext.ValueInitialized) - { - TryCatch tryCreateCosmosQueryExecutionContext = this.lazyTryCreateCosmosQueryExecutionContext.Result; - if (tryCreateCosmosQueryExecutionContext.Succeeded) - { - tryCreateCosmosQueryExecutionContext.Result.Dispose(); - } - } - } - - public override async Task ExecuteNextAsync(CancellationToken cancellationToken) - { - cancellationToken.ThrowIfCancellationRequested(); - - TryCatch tryCreateCosmosQueryExecutionContext = await this.lazyTryCreateCosmosQueryExecutionContext.GetValueAsync(cancellationToken); - if (!tryCreateCosmosQueryExecutionContext.Succeeded) - { - return QueryResponseFactory.CreateFromException(tryCreateCosmosQueryExecutionContext.Exception); - } - - CosmosQueryExecutionContext cosmosQueryExecutionContext = tryCreateCosmosQueryExecutionContext.Result; - QueryResponseCore queryResponseCore = await cosmosQueryExecutionContext.ExecuteNextAsync(cancellationToken); - return queryResponseCore; - } - - public override CosmosElement GetCosmosElementContinuationToken() - { - if (!this.lazyTryCreateCosmosQueryExecutionContext.ValueInitialized) - { - throw new InvalidOperationException(); - } - - TryCatch tryCreateCosmosQueryExecutionContext = this.lazyTryCreateCosmosQueryExecutionContext.Result; - if (!tryCreateCosmosQueryExecutionContext.Succeeded) - { - throw tryCreateCosmosQueryExecutionContext.Exception; - } - - return tryCreateCosmosQueryExecutionContext.Result.GetCosmosElementContinuationToken(); - } - } -} diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/PipelinedDocumentQueryExecutionContext.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/PipelinedDocumentQueryExecutionContext.cs deleted file mode 100644 index d0dbfd5885..0000000000 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/PipelinedDocumentQueryExecutionContext.cs +++ /dev/null @@ -1,410 +0,0 @@ -//------------------------------------------------------------ -// Copyright (c) Microsoft Corporation. All rights reserved. -//------------------------------------------------------------ - -namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionContext -{ - using System; - using System.Threading; - using System.Threading.Tasks; - using Microsoft.Azure.Cosmos; - using Microsoft.Azure.Cosmos.CosmosElements; - using Microsoft.Azure.Cosmos.Query.Core.ContinuationTokens; - using Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent; - using Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.Aggregate; - using Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.Distinct; - using Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.GroupBy; - using Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.SkipTake; - using Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.ItemProducers; - using Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.OrderBy; - using Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.Parallel; - using Microsoft.Azure.Cosmos.Query.Core.Monads; - using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Distinct; - using Microsoft.Azure.Cosmos.Query.Core.QueryClient; - using Microsoft.Azure.Cosmos.Query.Core.QueryPlan; - using Microsoft.Azure.Documents.Collections; - - /// - /// You can imagine the pipeline to be a directed acyclic graph where documents flow from multiple sources (the partitions) to a single sink (the client who calls on ExecuteNextAsync()). - /// The pipeline will consist of individual implementations of . - /// Every member of the pipeline has a source of documents (another member of the pipeline or an actual partition), - /// a method of draining documents (DrainAsync()) from said source, and a flag for whether that member of the pipeline is completely drained. - /// - /// The following is a diagram of the pipeline: - /// +--------------------------+ +--------------------------+ +--------------------------+ - /// | | | | | | - /// | Document Producer Tree 0 | | Document Producer Tree 1 | | Document Producer Tree N | - /// | | | | | | - /// +--------------------------+ +--------------------------+ +--------------------------+ - /// | | | - /// \ | / - /// \ | / - /// +---------------------------------------------------------+ - /// | | - /// | Parallel / Order By Document Query Execution Context | - /// | | - /// +---------------------------------------------------------+ - /// | - /// | - /// | - /// +---------------------------------------------------+ - /// | | - /// | Aggregate Document Query Execution Component | - /// | | - /// +---------------------------------------------------+ - /// | - /// | - /// | - /// +------------------------------------------+ - /// | | - /// | Top Document Query Execution Component | - /// | | - /// +------------------------------------------+ - /// | - /// | - /// | - /// +-----------------------------+ - /// | | - /// | Client | - /// | | - /// +-----------------------------+ - /// - /// - /// This class is responsible for constructing the pipelined described. - /// Note that the pipeline will always have one of or , - /// which both derive from as these are top level execution contexts. - /// These top level execution contexts have that are responsible for hitting the backend - /// and will optionally feed into and . - /// How these components are picked is based on , - /// which is a serialized form of this class and serves as a blueprint for construction. - /// - /// - /// Once the pipeline is constructed the client(sink of the graph) calls ExecuteNextAsync() which calls on DrainAsync(), - /// which by definition grabs documents from the parent component of the pipeline. - /// This bubbles down until you reach a component that has a DocumentProducer that fetches a document from the backend. - /// - /// -#nullable enable - internal sealed class PipelinedDocumentQueryExecutionContext : CosmosQueryExecutionContext - { - /// - /// The root level component that all calls will be forwarded to. - /// - private readonly IDocumentQueryExecutionComponent component; - - /// - /// The actual page size to drain. - /// - private readonly int actualPageSize; - - /// - /// Initializes a new instance of the PipelinedDocumentQueryExecutionContext class. - /// - /// The root level component that all calls will be forwarded to. - /// The actual page size to drain. - private PipelinedDocumentQueryExecutionContext( - IDocumentQueryExecutionComponent component, - int actualPageSize) - { - this.component = component ?? throw new ArgumentNullException($"{nameof(component)} can not be null."); - this.actualPageSize = (actualPageSize < 0) ? throw new ArgumentOutOfRangeException($"{nameof(actualPageSize)} can not be negative.") : actualPageSize; - } - - /// - /// Gets a value indicating whether this execution context is done draining documents. - /// - public override bool IsDone - { - get - { - return this.component.IsDone; - } - } - - public static async Task> MonadicCreateAsync( - ExecutionEnvironment executionEnvironment, - CosmosQueryContext queryContext, - CosmosCrossPartitionQueryExecutionContext.CrossPartitionInitParams initParams, - CosmosElement requestContinuationToken, - CancellationToken cancellationToken) - { - if (queryContext == null) - { - throw new ArgumentNullException(nameof(initParams)); - } - - cancellationToken.ThrowIfCancellationRequested(); - - QueryInfo queryInfo = initParams.PartitionedQueryExecutionInfo.QueryInfo; - - int initialPageSize = initParams.InitialPageSize; - if (queryInfo.HasGroupBy) - { - // The query will block until all groupings are gathered so we might as well speed up the process. - initParams = new CosmosCrossPartitionQueryExecutionContext.CrossPartitionInitParams( - sqlQuerySpec: initParams.SqlQuerySpec, - collectionRid: initParams.CollectionRid, - partitionedQueryExecutionInfo: initParams.PartitionedQueryExecutionInfo, - partitionKeyRanges: initParams.PartitionKeyRanges, - initialPageSize: int.MaxValue, - maxConcurrency: initParams.MaxConcurrency, - maxItemCount: int.MaxValue, - maxBufferedItemCount: initParams.MaxBufferedItemCount, - returnResultsInDeterministicOrder: true, - testSettings: initParams.TestSettings); - } - - Task> tryCreateOrderByComponentAsync(CosmosElement continuationToken) - { - return CosmosOrderByItemQueryExecutionContext.MonadicCreateAsync( - queryContext, - initParams, - continuationToken, - cancellationToken); - } - - Task> tryCreateParallelComponentAsync(CosmosElement continuationToken) - { - return CosmosParallelItemQueryExecutionContext.MonadicCreateAsync( - queryContext, - initParams, - continuationToken, - cancellationToken); - } - - Func>> tryCreatePipelineAsync; - if (queryInfo.HasOrderBy) - { - tryCreatePipelineAsync = tryCreateOrderByComponentAsync; - } - else - { - tryCreatePipelineAsync = tryCreateParallelComponentAsync; - } - - if (queryInfo.HasAggregates && !queryInfo.HasGroupBy) - { - Func>> monadicCreatePipelineStage = tryCreatePipelineAsync; - tryCreatePipelineAsync = async (continuationToken) => - { - return await AggregateDocumentQueryExecutionComponent.MonadicCreateAsync( - executionEnvironment, - queryInfo.Aggregates, - queryInfo.GroupByAliasToAggregateType, - queryInfo.GroupByAliases, - queryInfo.HasSelectValue, - continuationToken, - monadicCreatePipelineStage); - }; - } - - if (queryInfo.HasDistinct) - { - Func>> monadicCreatePipelineStage = tryCreatePipelineAsync; - tryCreatePipelineAsync = async (continuationToken) => - { - return await DistinctDocumentQueryExecutionComponent.MonadicCreateAsync( - executionEnvironment, - continuationToken, - monadicCreatePipelineStage, - queryInfo.DistinctType); - }; - } - - if (queryInfo.HasGroupBy) - { - Func>> monadicCreatePipelineStage = tryCreatePipelineAsync; - tryCreatePipelineAsync = async (continuationToken) => - { - return await GroupByDocumentQueryExecutionComponent.MonadicCreateAsync( - executionEnvironment, - continuationToken, - monadicCreatePipelineStage, - queryInfo.GroupByAliasToAggregateType, - queryInfo.GroupByAliases, - queryInfo.HasSelectValue); - }; - } - - if (queryInfo.Offset.HasValue) - { - Func>> monadicCreatePipelineStage = tryCreatePipelineAsync; - tryCreatePipelineAsync = async (continuationToken) => - { - return await SkipDocumentQueryExecutionComponent.MonadicCreateAsync( - executionEnvironment, - queryInfo.Offset.Value, - continuationToken, - monadicCreatePipelineStage); - }; - } - - if (queryInfo.Limit.HasValue) - { - Func>> monadicCreatePipelineStage = tryCreatePipelineAsync; - tryCreatePipelineAsync = async (continuationToken) => - { - return await TakeDocumentQueryExecutionComponent.MonadicCreateAsync( - executionEnvironment, - queryInfo.Limit.Value, - continuationToken, - monadicCreatePipelineStage); - }; - } - - if (queryInfo.Top.HasValue) - { - Func>> monadicCreatePipelineStage = tryCreatePipelineAsync; - tryCreatePipelineAsync = async (continuationToken) => - { - return await TakeDocumentQueryExecutionComponent.MonadicCreateAsync( - executionEnvironment, - queryInfo.Top.Value, - continuationToken, - monadicCreatePipelineStage); - }; - } - - return (await tryCreatePipelineAsync(requestContinuationToken)) - .Try((source) => new PipelinedDocumentQueryExecutionContext(source, initialPageSize)); - } - - public static async Task> TryCreatePassthroughAsync( - ExecutionEnvironment executionEnvironment, - CosmosQueryContext queryContext, - CosmosCrossPartitionQueryExecutionContext.CrossPartitionInitParams initParams, - CosmosElement requestContinuationToken, - CancellationToken cancellationToken) - { - if (queryContext == null) - { - throw new ArgumentNullException(nameof(queryContext)); - } - - cancellationToken.ThrowIfCancellationRequested(); - - // Modify query plan - PartitionedQueryExecutionInfo passThroughQueryInfo = new PartitionedQueryExecutionInfo() - { - QueryInfo = new QueryInfo() - { - Aggregates = null, - DistinctType = DistinctQueryType.None, - GroupByAliases = null, - GroupByAliasToAggregateType = null, - GroupByExpressions = null, - HasSelectValue = false, - Limit = null, - Offset = null, - OrderBy = null, - OrderByExpressions = null, - RewrittenQuery = null, - Top = null, - }, - QueryRanges = initParams.PartitionedQueryExecutionInfo.QueryRanges, - }; - - initParams = new CosmosCrossPartitionQueryExecutionContext.CrossPartitionInitParams( - sqlQuerySpec: initParams.SqlQuerySpec, - collectionRid: initParams.CollectionRid, - partitionedQueryExecutionInfo: passThroughQueryInfo, - partitionKeyRanges: initParams.PartitionKeyRanges, - initialPageSize: initParams.MaxItemCount.GetValueOrDefault(1000), - maxConcurrency: initParams.MaxConcurrency, - maxItemCount: initParams.MaxItemCount, - maxBufferedItemCount: initParams.MaxBufferedItemCount, - returnResultsInDeterministicOrder: initParams.ReturnResultsInDeterministicOrder, - testSettings: initParams.TestSettings); - - // Return a parallel context, since we still want to be able to handle splits and concurrency / buffering. - return (await CosmosParallelItemQueryExecutionContext.MonadicCreateAsync( - queryContext, - initParams, - requestContinuationToken, - cancellationToken)) - .Try((source) => new PipelinedDocumentQueryExecutionContext( - source, - initParams.InitialPageSize)); - } - - /// - /// Disposes of this context. - /// - public override void Dispose() - { - this.component.Dispose(); - } - - /// - /// Gets the next page of results from this context. - /// - /// The cancellation token. - /// A task to await on that in turn returns a DoucmentFeedResponse of results. - public async Task> ExecuteNextFeedResponseAsync(CancellationToken token) - { - QueryResponseCore feedResponse = await this.ExecuteNextAsync(token); - return new DocumentFeedResponse( - result: feedResponse.CosmosElements, - count: feedResponse.CosmosElements.Count, - responseHeaders: new DictionaryNameValueCollection(), - useETagAsContinuation: false, - queryMetrics: null, - requestStats: null, - disallowContinuationTokenMessage: feedResponse.DisallowContinuationTokenMessage, - responseLengthBytes: feedResponse.ResponseLengthBytes); - } - - /// - /// Gets the next page of results from this context. - /// - /// The cancellation token. - /// A task to await on that in turn returns a DoucmentFeedResponse of results. - public override async Task ExecuteNextAsync(CancellationToken token) - { - try - { - QueryResponseCore queryResponse = await this.component.DrainAsync(this.actualPageSize, token); - if (!queryResponse.IsSuccess) - { - this.component.Stop(); - return queryResponse; - } - - string? updatedContinuationToken; - if (queryResponse.DisallowContinuationTokenMessage == null) - { - if (queryResponse.ContinuationToken != null) - { - updatedContinuationToken = new PipelineContinuationTokenV0(CosmosElement.Parse(queryResponse.ContinuationToken)).ToString(); - } - else - { - updatedContinuationToken = null; - } - } - else - { - updatedContinuationToken = null; - } - - return QueryResponseCore.CreateSuccess( - result: queryResponse.CosmosElements, - continuationToken: updatedContinuationToken, - disallowContinuationTokenMessage: queryResponse.DisallowContinuationTokenMessage, - activityId: queryResponse.ActivityId, - requestCharge: queryResponse.RequestCharge, - responseLengthBytes: queryResponse.ResponseLengthBytes); - } - catch (Exception) - { - this.component.Stop(); - throw; - } - } - - public override CosmosElement GetCosmosElementContinuationToken() - { - return this.component.GetCosmosElementContinuationToken(); - } - } -} diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Monads/TryCatch{TResult}.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Monads/TryCatch{TResult}.cs index 07bf3f16f7..74588e157c 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Monads/TryCatch{TResult}.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Monads/TryCatch{TResult}.cs @@ -61,6 +61,20 @@ public Exception Exception } } + public Exception InnerMostException + { + get + { + Exception exception = this.Exception; + while (exception != null) + { + exception = exception.InnerException; + } + + return exception; + } + } + public void Match( Action onSuccess, Action onError) diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Aggregate/AggregateQueryPipelineStage.Compute.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Aggregate/AggregateQueryPipelineStage.Compute.cs index b6bdd1850c..66712e79c1 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Aggregate/AggregateQueryPipelineStage.Compute.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Aggregate/AggregateQueryPipelineStage.Compute.cs @@ -69,7 +69,7 @@ public static TryCatch MonadicCreate( TryCatch tryCreateSource; if (aggregateContinuationToken.SourceContinuationToken is CosmosString stringToken && (stringToken.Value == DoneSourceToken.Value)) { - tryCreateSource = TryCatch.FromResult(FinishedQueryPipelineStage.Value); + tryCreateSource = TryCatch.FromResult(EmptyQueryPipelineStage.Value); } else { diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CatchAllQueryPipelineStage.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CatchAllQueryPipelineStage.cs new file mode 100644 index 0000000000..db467fb529 --- /dev/null +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CatchAllQueryPipelineStage.cs @@ -0,0 +1,40 @@ +// ------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +// ------------------------------------------------------------ + +namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline +{ + using System; + using System.Threading; + using System.Threading.Tasks; + using Microsoft.Azure.Cosmos.Query.Core.Monads; + using Microsoft.Azure.Cosmos.Query.Core.QueryClient; + + internal sealed class CatchAllQueryPipelineStage : QueryPipelineStageBase + { + public CatchAllQueryPipelineStage(IQueryPipelineStage inputStage) + : base(inputStage) + { + } + + protected override async Task> GetNextPageAsync(CancellationToken cancellationToken) + { + try + { + await this.inputStage.MoveNextAsync(); + return this.inputStage.Current; + } + catch (OperationCanceledException) when (cancellationToken.IsCancellationRequested) + { + // Per cancellationToken.ThrowIfCancellationRequested(); line above, this function should still throw OperationCanceledException. + throw; + } + catch (Exception ex) + { + QueryResponseCore queryResponse = QueryResponseFactory.CreateFromException(ex); + CosmosException exception = queryResponse.CosmosException; + return TryCatch.FromException(exception); + } + } + } +} diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/EmptyQueryPipelineStage.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/EmptyQueryPipelineStage.cs new file mode 100644 index 0000000000..265eeb3440 --- /dev/null +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/EmptyQueryPipelineStage.cs @@ -0,0 +1,28 @@ +// ------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +// ------------------------------------------------------------ + +namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline +{ + using System.Threading.Tasks; + using Microsoft.Azure.Cosmos.Query.Core.Monads; + using Microsoft.Azure.Cosmos.Reactive; + + internal sealed class EmptyQueryPipelineStage : IQueryPipelineStage + { + public static readonly EmptyQueryPipelineStage Value = new EmptyQueryPipelineStage(); + + private readonly EmptyAsyncEnumerator> emptyAsyncEnumerator; + + public EmptyQueryPipelineStage() + { + this.emptyAsyncEnumerator = new EmptyAsyncEnumerator>(); + } + + public TryCatch Current => this.emptyAsyncEnumerator.Current; + + public ValueTask DisposeAsync() => this.emptyAsyncEnumerator.DisposeAsync(); + + public ValueTask MoveNextAsync() => this.emptyAsyncEnumerator.MoveNextAsync(); + } +} diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/FaultedQueryPipelineStage.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/FaultedQueryPipelineStage.cs new file mode 100644 index 0000000000..8761aea123 --- /dev/null +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/FaultedQueryPipelineStage.cs @@ -0,0 +1,32 @@ +// ------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +// ------------------------------------------------------------ + +namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline +{ + using System; + using System.Threading.Tasks; + using Microsoft.Azure.Cosmos.Query.Core.Monads; + using Microsoft.Azure.Cosmos.Reactive; + + internal sealed class FaultedQueryPipelineStage : IQueryPipelineStage + { + private readonly JustAsyncEnumerator> justAsyncEnumerator; + + public FaultedQueryPipelineStage(Exception exception) + { + if (exception == null) + { + throw new ArgumentNullException(nameof(exception)); + } + + this.justAsyncEnumerator = new JustAsyncEnumerator>(TryCatch.FromException(exception)); + } + + public TryCatch Current => this.justAsyncEnumerator.Current; + + public ValueTask DisposeAsync() => this.justAsyncEnumerator.DisposeAsync(); + + public ValueTask MoveNextAsync() => this.justAsyncEnumerator.MoveNextAsync(); + } +} diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/FinishedQueryPipelineStage.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/FinishedQueryPipelineStage.cs deleted file mode 100644 index 1728479894..0000000000 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/FinishedQueryPipelineStage.cs +++ /dev/null @@ -1,30 +0,0 @@ -// ------------------------------------------------------------ -// Copyright (c) Microsoft Corporation. All rights reserved. -// ------------------------------------------------------------ - -namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline -{ - using System; - using System.Threading; - using System.Threading.Tasks; - using Microsoft.Azure.Cosmos.Query.Core.Monads; - - internal sealed class FinishedQueryPipelineStage : IQueryPipelineStage - { - public static readonly FinishedQueryPipelineStage Value = new FinishedQueryPipelineStage(); - - private FinishedQueryPipelineStage() - { - } - - public bool HasMoreResults => false; - - public TryCatch Current => default; - - public ValueTask DisposeAsync() => default; - - public Task> GetNextPageAsync(CancellationToken cancellationToken) => throw new NotSupportedException(); - - public ValueTask MoveNextAsync() => new ValueTask(false); - } -} diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/GroupBy/GroupByQueryPipelineStage.Compute.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/GroupBy/GroupByQueryPipelineStage.Compute.cs index dd8016a7bb..55be638e72 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/GroupBy/GroupByQueryPipelineStage.Compute.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/GroupBy/GroupByQueryPipelineStage.Compute.cs @@ -61,7 +61,7 @@ public static TryCatch MonadicCreate( if ((groupByContinuationToken.SourceContinuationToken is CosmosString sourceContinuationToken) && (sourceContinuationToken.Value == ComputeGroupByQueryPipelineStage.DoneReadingGroupingsContinuationToken)) { - tryCreateSource = TryCatch.FromResult(FinishedQueryPipelineStage.Value); + tryCreateSource = TryCatch.FromResult(EmptyQueryPipelineStage.Value); } else { diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/LazyQueryPipelineStage.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/LazyQueryPipelineStage.cs new file mode 100644 index 0000000000..3757c21c6c --- /dev/null +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/LazyQueryPipelineStage.cs @@ -0,0 +1,55 @@ +// ------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +// ------------------------------------------------------------ + +namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline +{ + using System; + using System.Threading.Tasks; + using Microsoft.Azure.Cosmos.Query.Core.Monads; + + internal sealed class LazyQueryPipelineStage : IQueryPipelineStage + { + private readonly AsyncLazy> lazyTryCreateStage; + + public LazyQueryPipelineStage(AsyncLazy> lazyTryCreateStage) + { + this.lazyTryCreateStage = lazyTryCreateStage ?? throw new ArgumentNullException(nameof(lazyTryCreateStage)); + } + + public TryCatch Current { get; private set; } + + public ValueTask DisposeAsync() + { + if (this.lazyTryCreateStage.ValueInitialized) + { + TryCatch tryCreatePipelineStage = this.lazyTryCreateStage.Result; + if (tryCreatePipelineStage.Succeeded) + { + return tryCreatePipelineStage.Result.DisposeAsync(); + } + } + + return default; + } + + public async ValueTask MoveNextAsync() + { + TryCatch tryCreateStage = await this.lazyTryCreateStage.GetValueAsync(default); + if (tryCreateStage.Failed) + { + this.Current = TryCatch.FromException(tryCreateStage.Exception); + return true; + } + + IQueryPipelineStage stage = tryCreateStage.Result; + if (!await stage.MoveNextAsync()) + { + return false; + } + + this.Current = stage.Current; + return true; + } + } +} diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/NameCacheStaleRetryQueryPipelineStage.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/NameCacheStaleRetryQueryPipelineStage.cs new file mode 100644 index 0000000000..a54189a209 --- /dev/null +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/NameCacheStaleRetryQueryPipelineStage.cs @@ -0,0 +1,66 @@ +// ------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +// ------------------------------------------------------------ + +namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline +{ + using System; + using System.Net; + using System.Threading; + using System.Threading.Tasks; + using Microsoft.Azure.Cosmos.Query.Core.Monads; + using Microsoft.Azure.Cosmos.Query.Core.QueryClient; + + internal sealed class NameCacheStaleRetryQueryPipelineStage : IQueryPipelineStage + { + private readonly CosmosQueryContext cosmosQueryContext; + private readonly Func queryPipelineStageFactory; + private IQueryPipelineStage currentQueryPipelineStage; + private bool alreadyRetried; + + public NameCacheStaleRetryQueryPipelineStage( + CosmosQueryContext cosmosQueryContext, + Func queryPipelineStageFactory) + { + this.cosmosQueryContext = cosmosQueryContext ?? throw new ArgumentNullException(nameof(cosmosQueryContext)); + this.queryPipelineStageFactory = queryPipelineStageFactory ?? throw new ArgumentNullException(nameof(queryPipelineStageFactory)); + this.currentQueryPipelineStage = queryPipelineStageFactory(); + } + + public TryCatch Current { get; private set; } + + public ValueTask DisposeAsync() => this.currentQueryPipelineStage.DisposeAsync(); + + public async ValueTask MoveNextAsync() + { + if (!await this.currentQueryPipelineStage.MoveNextAsync()) + { + return false; + } + + TryCatch tryGetSourcePage = this.currentQueryPipelineStage.Current; + this.Current = tryGetSourcePage; + + if (tryGetSourcePage.Failed) + { + Exception exception = tryGetSourcePage.InnerMostException; + bool shouldRetry = (exception is CosmosException cosmosException) + && (cosmosException.StatusCode == HttpStatusCode.Gone) + && (cosmosException.SubStatusCode == (int)Documents.SubStatusCodes.NameCacheIsStale) + && !this.alreadyRetried; + if (shouldRetry) + { + await this.cosmosQueryContext.QueryClient.ForceRefreshCollectionCacheAsync( + this.cosmosQueryContext.ResourceLink.OriginalString, + default); + this.alreadyRetried = true; + await this.currentQueryPipelineStage.DisposeAsync(); + this.currentQueryPipelineStage = this.queryPipelineStageFactory(); + return await this.MoveNextAsync(); + } + } + + return true; + } + } +} diff --git a/Microsoft.Azure.Cosmos/src/Query/v3Query/QueryIterator.cs b/Microsoft.Azure.Cosmos/src/Query/v3Query/QueryIterator.cs index d79008c7a2..8c354325ce 100644 --- a/Microsoft.Azure.Cosmos/src/Query/v3Query/QueryIterator.cs +++ b/Microsoft.Azure.Cosmos/src/Query/v3Query/QueryIterator.cs @@ -5,38 +5,41 @@ namespace Microsoft.Azure.Cosmos.Query { using System; - using System.Collections.Generic; using System.Threading; using System.Threading.Tasks; using Microsoft.Azure.Cosmos.CosmosElements; - using Microsoft.Azure.Cosmos.Diagnostics; using Microsoft.Azure.Cosmos.Query.Core; using Microsoft.Azure.Cosmos.Query.Core.Exceptions; 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.QueryClient; using Microsoft.Azure.Cosmos.Query.Core.QueryPlan; + using Microsoft.Azure.Documents; internal sealed class QueryIterator : FeedIteratorInternal { private readonly CosmosQueryContextCore cosmosQueryContext; - private readonly CosmosQueryExecutionContext cosmosQueryExecutionContext; + private readonly IQueryPipelineStage queryPipelineStage; private readonly CosmosSerializationFormatOptions cosmosSerializationFormatOptions; private readonly RequestOptions requestOptions; private readonly CosmosClientContext clientContext; + private bool hasMoreResults; + private QueryIterator( CosmosQueryContextCore cosmosQueryContext, - CosmosQueryExecutionContext cosmosQueryExecutionContext, + IQueryPipelineStage cosmosQueryExecutionContext, CosmosSerializationFormatOptions cosmosSerializationFormatOptions, RequestOptions requestOptions, CosmosClientContext clientContext) { this.cosmosQueryContext = cosmosQueryContext ?? throw new ArgumentNullException(nameof(cosmosQueryContext)); - this.cosmosQueryExecutionContext = cosmosQueryExecutionContext ?? throw new ArgumentNullException(nameof(cosmosQueryExecutionContext)); + this.queryPipelineStage = cosmosQueryExecutionContext ?? throw new ArgumentNullException(nameof(cosmosQueryExecutionContext)); this.cosmosSerializationFormatOptions = cosmosSerializationFormatOptions; this.requestOptions = requestOptions; this.clientContext = clientContext ?? throw new ArgumentNullException(nameof(clientContext)); + this.hasMoreResults = true; } public static QueryIterator Create( @@ -82,7 +85,7 @@ public static QueryIterator Create( { return new QueryIterator( cosmosQueryContext, - new QueryExecutionContextWithException( + new FaultedQueryPipelineStage( new MalformedContinuationTokenException( message: $"Malformed Continuation Token: {continuationToken}", innerException: tryParse.Exception)), @@ -130,18 +133,19 @@ public static QueryIterator Create( clientContext); } - public override bool HasMoreResults => !this.cosmosQueryExecutionContext.IsDone; + public override bool HasMoreResults => this.hasMoreResults; public override async Task ReadNextAsync(CancellationToken cancellationToken = default) { CosmosDiagnosticsContext diagnostics = CosmosDiagnosticsContext.Create(this.requestOptions); using (diagnostics.GetOverallScope()) { - QueryResponseCore responseCore; + TryCatch tryGetQueryPage; try { // This catches exception thrown by the pipeline and converts it to QueryResponse - responseCore = await this.cosmosQueryExecutionContext.ExecuteNextAsync(cancellationToken); + await this.queryPipelineStage.MoveNextAsync(); + tryGetQueryPage = this.queryPipelineStage.Current; } catch (OperationCanceledException ex) when (!(ex is CosmosOperationCanceledException)) { @@ -149,61 +153,65 @@ public override async Task ReadNextAsync(CancellationToken canc } finally { - // This swaps the diagnostics in the context. This shows all the page reads between the previous ReadNextAsync and the current ReadNextAsync + // This swaps the diagnostics in the context. + // This shows all the page reads between the previous ReadNextAsync and the current ReadNextAsync diagnostics.AddDiagnosticsInternal(this.cosmosQueryContext.GetAndResetDiagnostics()); } - if (responseCore.IsSuccess) + if (tryGetQueryPage.Succeeded) { return QueryResponse.CreateSuccess( - result: responseCore.CosmosElements, - count: responseCore.CosmosElements.Count, - responseLengthBytes: responseCore.ResponseLengthBytes, + result: tryGetQueryPage.Result.Documents, + count: tryGetQueryPage.Result.Documents.Count, + responseLengthBytes: tryGetQueryPage.Result.ResponseLengthInBytes, diagnostics: diagnostics, serializationOptions: this.cosmosSerializationFormatOptions, responseHeaders: new CosmosQueryResponseMessageHeaders( - responseCore.ContinuationToken, - responseCore.DisallowContinuationTokenMessage, + (tryGetQueryPage.Result.State?.Value as CosmosString)?.Value, + tryGetQueryPage.Result.DisallowContinuationTokenMessage, this.cosmosQueryContext.ResourceTypeEnum, this.cosmosQueryContext.ContainerResourceId) { - RequestCharge = responseCore.RequestCharge, - ActivityId = responseCore.ActivityId, - SubStatusCode = responseCore.SubStatusCode ?? Documents.SubStatusCodes.Unknown + RequestCharge = tryGetQueryPage.Result.RequestCharge, + ActivityId = tryGetQueryPage.Result.ActivityId, + SubStatusCode = Documents.SubStatusCodes.Unknown }); } - if (responseCore.CosmosException != null) + CosmosException cosmosException = (CosmosException)tryGetQueryPage.InnerMostException; + SubStatusCodes subStatusCode; + if (Enum.IsDefined(typeof(SubStatusCodes), cosmosException.SubStatusCode)) + { + subStatusCode = (SubStatusCodes)cosmosException.SubStatusCode; + } + else { - return responseCore.CosmosException.ToCosmosResponseMessage(null); + subStatusCode = Documents.SubStatusCodes.Unknown; } return QueryResponse.CreateFailure( - statusCode: responseCore.StatusCode, - cosmosException: responseCore.CosmosException, + statusCode: cosmosException.StatusCode, + cosmosException: cosmosException, requestMessage: null, diagnostics: diagnostics, responseHeaders: new CosmosQueryResponseMessageHeaders( - responseCore.ContinuationToken, - responseCore.DisallowContinuationTokenMessage, + continauationToken: default, + disallowContinuationTokenMessage: default, this.cosmosQueryContext.ResourceTypeEnum, this.cosmosQueryContext.ContainerResourceId) { - RequestCharge = responseCore.RequestCharge, - ActivityId = responseCore.ActivityId, - SubStatusCode = responseCore.SubStatusCode ?? Documents.SubStatusCodes.Unknown, + RequestCharge = cosmosException.RequestCharge, + ActivityId = cosmosException.ActivityId, + SubStatusCode = subStatusCode, }); } } - public override CosmosElement GetCosmosElementContinuationToken() - { - return this.cosmosQueryExecutionContext.GetCosmosElementContinuationToken(); - } + public override CosmosElement GetCosmosElementContinuationToken() => this.queryPipelineStage.Current.Result.State.Value; protected override void Dispose(bool disposing) { - this.cosmosQueryExecutionContext.Dispose(); + this.queryPipelineStage.DisposeAsync(); base.Dispose(disposing); } } diff --git a/Microsoft.Azure.Cosmos/src/Reactive/EmptyAsyncEnumerator.cs b/Microsoft.Azure.Cosmos/src/Reactive/EmptyAsyncEnumerator.cs new file mode 100644 index 0000000000..c18e1f44b9 --- /dev/null +++ b/Microsoft.Azure.Cosmos/src/Reactive/EmptyAsyncEnumerator.cs @@ -0,0 +1,27 @@ +// ------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +// ------------------------------------------------------------ + +namespace Microsoft.Azure.Cosmos.Reactive +{ + using System.Collections.Generic; + using System.Threading.Tasks; + + /// + /// Emits no items but terminates normally. + /// + /// The type of the items. + /// + internal sealed class EmptyAsyncEnumerator : IAsyncEnumerator + { + public EmptyAsyncEnumerator() + { + } + + public T Current => default; + + public ValueTask DisposeAsync() => default; + + public ValueTask MoveNextAsync() => new ValueTask(false); + } +} diff --git a/Microsoft.Azure.Cosmos/src/Reactive/JustAsyncEnumerator.cs b/Microsoft.Azure.Cosmos/src/Reactive/JustAsyncEnumerator.cs new file mode 100644 index 0000000000..cd24a96e73 --- /dev/null +++ b/Microsoft.Azure.Cosmos/src/Reactive/JustAsyncEnumerator.cs @@ -0,0 +1,30 @@ +// ------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +// ------------------------------------------------------------ + +namespace Microsoft.Azure.Cosmos.Reactive +{ + using System.Collections.Generic; + using System.Threading.Tasks; + + /// + /// Emits a particular item (or series of items). + /// + /// The type of the item(s) + /// + internal sealed class JustAsyncEnumerator : IAsyncEnumerator + { + private readonly IEnumerator enumerator; + + public JustAsyncEnumerator(params T[] items) + { + this.enumerator = (IEnumerator)items.GetEnumerator(); + } + + public T Current => this.enumerator.Current; + + public ValueTask DisposeAsync() => default; + + public ValueTask MoveNextAsync() => new ValueTask(this.enumerator.MoveNext()); + } +} diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/CosmosQueryUnitTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/CosmosQueryUnitTests.cs index facd287640..fe0a221e2c 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/CosmosQueryUnitTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/CosmosQueryUnitTests.cs @@ -23,6 +23,7 @@ namespace Microsoft.Azure.Cosmos.Tests using Microsoft.Azure.Cosmos.Query.Core.ExecutionContext; using Microsoft.Azure.Cosmos.Query.Core.Metrics; using Microsoft.Azure.Cosmos.Query.Core.Monads; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline; using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Aggregate; using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Distinct; using Microsoft.Azure.Cosmos.Query.Core.QueryClient; @@ -285,13 +286,17 @@ public async Task TestCosmosQueryPartitionKeyDefinition() diagnosticsContext: new CosmosDiagnosticsContextCore(), correlatedActivityId: new Guid("221FC86C-1825-4284-B10E-A6029652CCA6")); - CosmosQueryExecutionContext context = CosmosQueryExecutionContextFactory.Create( + IQueryPipelineStage pipelineStage = CosmosQueryExecutionContextFactory.Create( cosmosQueryContext, inputParameters); - QueryResponseCore queryResponse = await context.ExecuteNextAsync(cancellationtoken); - Assert.AreEqual(HttpStatusCode.BadRequest, queryResponse.StatusCode); - Assert.IsTrue(queryResponse.CosmosException.ToString().Contains(exceptionMessage), "response error message did not contain the proper substring."); + Assert.IsTrue(await pipelineStage.MoveNextAsync()); + TryCatch tryGetPage = pipelineStage.Current; + Assert.IsTrue(tryGetPage.Failed); + Assert.AreEqual(HttpStatusCode.BadRequest, (tryGetPage.Exception as CosmosException).StatusCode); + Assert.IsTrue( + (tryGetPage.Exception as CosmosException).ToString().Contains(exceptionMessage), + "response error message did not contain the proper substring."); } private async Task<(IList components, QueryResponseCore response)> GetAllExecutionComponents() diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/MockQueryPipelineStage.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/MockQueryPipelineStage.cs index dbbea34dc3..67943111c1 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/MockQueryPipelineStage.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/MockQueryPipelineStage.cs @@ -18,7 +18,7 @@ internal sealed class MockQueryPipelineStage : QueryPipelineStageBase private int pageIndex; public MockQueryPipelineStage(IReadOnlyList> pages) - : base(FinishedQueryPipelineStage.Value) + : base(EmptyQueryPipelineStage.Value) { this.pages = pages ?? throw new ArgumentNullException(nameof(pages)); } From b050521c8feb70cd3e1dc76e7ed89e0a539e87e5 Mon Sep 17 00:00:00 2001 From: Brandon Chong Date: Wed, 8 Jul 2020 10:59:22 -0700 Subject: [PATCH 29/85] started parallel cpq debugging --- .../src/FeedRange/FeedRanges/FeedRangeEPK.cs | 5 + .../FeedRange/FeedRanges/FeedRangeInternal.cs | 5 + .../FeedRanges/FeedRangePartitionKey.cs | 5 + .../FeedRanges/FeedRangePartitionKeyRange.cs | 5 + .../IFeedRangeAsyncVisitor{TResult,TArg}.cs | 18 ++ .../CreatePartitionRangeEnumerator.cs | 2 +- .../CrossPartitionRangePageEnumerator.cs | 31 ++- .../src/Pagination/CrossPartitionState.cs | 4 +- .../src/Pagination/FeedRangeProvider.cs | 22 +- .../src/Pagination/IFeedRangeProvider.cs | 10 +- .../PartitionRangePageEnumerable.cs | 4 +- .../PartitionRangePageEnumerator.cs | 4 +- .../Query/Core/ExceptionToCosmosException.cs | 112 ++++++++++ .../Query/Core/Monads/TryCatch{TResult}.cs | 2 +- .../Pipeline/CatchAllQueryPipelineStage.cs | 5 +- .../Pipeline/Remote/BackendQueryDataSource.cs | 114 +++++++++- .../Core/Pipeline/Remote/IQueryDataSource.cs | 2 +- ...arallelCrossPartitionQueryPipelineStage.cs | 31 ++- .../QueryPartitionRangePageEnumerator.cs | 9 +- .../Query/v3Query/CosmosQueryClientCore.cs | 1 - .../Query/v3Query/CosmosQueryContextCore.cs | 2 - .../src/Query/v3Query/QueryIterator.cs | 2 +- .../src/Resource/CosmosClientContext.cs | 3 - .../CosmosExceptions/CosmosException.cs | 6 +- .../src/Routing/PartitionKeyHash.cs | 20 +- .../src/Routing/PartitionKeyHashRange.cs | 1 + .../Query/SanityQueryTests.cs | 29 ++- .../CosmosQueryUnitTests.cs | 2 +- .../InMemoryCollection.cs | 14 +- .../InMemoryCollectionFeedRangeProvider.cs | 133 ++++++++++-- .../InMemoryCollectionQueryDataSource.cs | 4 +- ...elCrossPartitionQueryPipelineStageTests.cs | 196 ++++++++++++++++++ 32 files changed, 705 insertions(+), 98 deletions(-) create mode 100644 Microsoft.Azure.Cosmos/src/FeedRange/FeedRanges/IFeedRangeAsyncVisitor{TResult,TArg}.cs create mode 100644 Microsoft.Azure.Cosmos/src/Query/Core/ExceptionToCosmosException.cs create mode 100644 Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/ParallelCrossPartitionQueryPipelineStageTests.cs diff --git a/Microsoft.Azure.Cosmos/src/FeedRange/FeedRanges/FeedRangeEPK.cs b/Microsoft.Azure.Cosmos/src/FeedRange/FeedRanges/FeedRangeEPK.cs index c6aab85d22..4043fcce91 100644 --- a/Microsoft.Azure.Cosmos/src/FeedRange/FeedRanges/FeedRangeEPK.cs +++ b/Microsoft.Azure.Cosmos/src/FeedRange/FeedRanges/FeedRangeEPK.cs @@ -50,6 +50,11 @@ public override Task AcceptAsync( IFeedRangeAsyncVisitor visitor, CancellationToken cancellationToken = default) => visitor.VisitAsync(this, cancellationToken); + public override Task AcceptAsync( + IFeedRangeAsyncVisitor visitor, + TArg argument, + CancellationToken cancellationToken) => visitor.VisitAsync(this, argument, cancellationToken); + public override string ToString() => this.Range.ToString(); } } diff --git a/Microsoft.Azure.Cosmos/src/FeedRange/FeedRanges/FeedRangeInternal.cs b/Microsoft.Azure.Cosmos/src/FeedRange/FeedRanges/FeedRangeInternal.cs index b38f3265e1..4008d4cd7a 100644 --- a/Microsoft.Azure.Cosmos/src/FeedRange/FeedRanges/FeedRangeInternal.cs +++ b/Microsoft.Azure.Cosmos/src/FeedRange/FeedRanges/FeedRangeInternal.cs @@ -29,6 +29,11 @@ public abstract Task> GetPartitionKeyRangesAsync( public abstract Task AcceptAsync(IFeedRangeAsyncVisitor visitor, CancellationToken cancellationToken = default); + public abstract Task AcceptAsync( + IFeedRangeAsyncVisitor visitor, + TArg argument, + CancellationToken cancellationToken); + public abstract override string ToString(); public override string ToJsonString() => JsonConvert.SerializeObject(this); diff --git a/Microsoft.Azure.Cosmos/src/FeedRange/FeedRanges/FeedRangePartitionKey.cs b/Microsoft.Azure.Cosmos/src/FeedRange/FeedRanges/FeedRangePartitionKey.cs index 000b343bcc..eb6fd39ebc 100644 --- a/Microsoft.Azure.Cosmos/src/FeedRange/FeedRanges/FeedRangePartitionKey.cs +++ b/Microsoft.Azure.Cosmos/src/FeedRange/FeedRanges/FeedRangePartitionKey.cs @@ -48,6 +48,11 @@ public override Task AcceptAsync( IFeedRangeAsyncVisitor visitor, CancellationToken cancellationToken = default) => visitor.VisitAsync(this, cancellationToken); + public override Task AcceptAsync( + IFeedRangeAsyncVisitor visitor, + TArg argument, + CancellationToken cancellationToken) => visitor.VisitAsync(this, argument, cancellationToken); + public override string ToString() => this.PartitionKey.InternalKey.ToJsonString(); } } diff --git a/Microsoft.Azure.Cosmos/src/FeedRange/FeedRanges/FeedRangePartitionKeyRange.cs b/Microsoft.Azure.Cosmos/src/FeedRange/FeedRanges/FeedRangePartitionKeyRange.cs index c6c0b50cd9..8b35392bec 100644 --- a/Microsoft.Azure.Cosmos/src/FeedRange/FeedRanges/FeedRangePartitionKeyRange.cs +++ b/Microsoft.Azure.Cosmos/src/FeedRange/FeedRanges/FeedRangePartitionKeyRange.cs @@ -80,6 +80,11 @@ public override Task AcceptAsync( IFeedRangeAsyncVisitor visitor, CancellationToken cancellationToken = default) => visitor.VisitAsync(this, cancellationToken); + public override Task AcceptAsync( + IFeedRangeAsyncVisitor visitor, + TArg argument, + CancellationToken cancellationToken) => visitor.VisitAsync(this, argument, cancellationToken); + public override string ToString() => this.PartitionKeyRangeId; } } diff --git a/Microsoft.Azure.Cosmos/src/FeedRange/FeedRanges/IFeedRangeAsyncVisitor{TResult,TArg}.cs b/Microsoft.Azure.Cosmos/src/FeedRange/FeedRanges/IFeedRangeAsyncVisitor{TResult,TArg}.cs new file mode 100644 index 0000000000..b940564e01 --- /dev/null +++ b/Microsoft.Azure.Cosmos/src/FeedRange/FeedRanges/IFeedRangeAsyncVisitor{TResult,TArg}.cs @@ -0,0 +1,18 @@ +// ------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +// ------------------------------------------------------------ + +namespace Microsoft.Azure.Cosmos +{ + using System.Threading; + using System.Threading.Tasks; + + internal interface IFeedRangeAsyncVisitor + { + public abstract Task VisitAsync(FeedRangePartitionKey feedRange, TArg argument, CancellationToken cancellationToken); + + public abstract Task VisitAsync(FeedRangePartitionKeyRange feedRange, TArg argument, CancellationToken cancellationToken); + + public abstract Task VisitAsync(FeedRangeEpk feedRange, TArg argument, CancellationToken cancellationToken); + } +} diff --git a/Microsoft.Azure.Cosmos/src/Pagination/CreatePartitionRangeEnumerator.cs b/Microsoft.Azure.Cosmos/src/Pagination/CreatePartitionRangeEnumerator.cs index 74d79fd5c5..4723204b2c 100644 --- a/Microsoft.Azure.Cosmos/src/Pagination/CreatePartitionRangeEnumerator.cs +++ b/Microsoft.Azure.Cosmos/src/Pagination/CreatePartitionRangeEnumerator.cs @@ -4,7 +4,7 @@ namespace Microsoft.Azure.Cosmos.Pagination { - internal delegate PartitionRangePageEnumerator CreatePartitionRangePageEnumerator(FeedRange feedRange, TState state) + internal delegate PartitionRangePageEnumerator CreatePartitionRangePageEnumerator(FeedRangeInternal feedRange, TState state) where TPage : Page where TState : State; } diff --git a/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionRangePageEnumerator.cs b/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionRangePageEnumerator.cs index eddeef7a35..957feb7ffd 100644 --- a/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionRangePageEnumerator.cs +++ b/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionRangePageEnumerator.cs @@ -40,7 +40,7 @@ public CrossPartitionRangePageEnumerator( this.lazyEnumerators = new AsyncLazy>>(async (CancellationToken token) => { - IReadOnlyList<(FeedRange, TState)> rangeAndStates; + IReadOnlyList<(FeedRangeInternal, TState)> rangeAndStates; if (state != default) { rangeAndStates = state.Value; @@ -48,10 +48,10 @@ public CrossPartitionRangePageEnumerator( else { // Fan out to all partitions with default state - IEnumerable ranges = await feedRangeProvider.GetFeedRangesAsync(token); + IEnumerable ranges = await feedRangeProvider.GetFeedRangesAsync(token); - List<(FeedRange, TState)> rangesAndStatesBuilder = new List<(FeedRange, TState)>(); - foreach (FeedRange range in ranges) + List<(FeedRangeInternal, TState)> rangesAndStatesBuilder = new List<(FeedRangeInternal, TState)>(); + foreach (FeedRangeInternal range in ranges) { rangesAndStatesBuilder.Add((range, default)); } @@ -60,7 +60,7 @@ public CrossPartitionRangePageEnumerator( } PriorityQueue> enumerators = new PriorityQueue>(comparer); - foreach ((FeedRange range, TState rangeState) in rangeAndStates) + foreach ((FeedRangeInternal range, TState rangeState) in rangeAndStates) { PartitionRangePageEnumerator enumerator = createPartitionRangeEnumerator(range, rangeState); enumerators.Enqueue(enumerator); @@ -94,10 +94,14 @@ public async ValueTask MoveNextAsync() if (IsSplitException(exception)) { // Handle split - IEnumerable childRanges = await this.feedRangeProvider.GetChildRangeAsync(currentPaginator.Range); - foreach (FeedRange childRange in childRanges) + IEnumerable childRanges = await this.feedRangeProvider.GetChildRangeAsync( + currentPaginator.Range, + cancellationToken: default); + foreach (FeedRangeInternal childRange in childRanges) { - PartitionRangePageEnumerator childPaginator = this.createPartitionRangeEnumerator(childRange, currentPaginator.State); + PartitionRangePageEnumerator childPaginator = this.createPartitionRangeEnumerator( + childRange, + currentPaginator.State); enumerators.Enqueue(childPaginator); } @@ -120,10 +124,17 @@ public async ValueTask MoveNextAsync() return true; } - List<(FeedRange, TState)> feedRangeAndStates = new List<(FeedRange, TState)>(enumerators.Count); + List<(FeedRangeInternal, TState)> feedRangeAndStates = new List<(FeedRangeInternal, TState)>(enumerators.Count); foreach (PartitionRangePageEnumerator enumerator in enumerators) { - feedRangeAndStates.Add((enumerator.Range, enumerator.State)); + // Converting to epk range just because some implementations prefer that. + // Eventually we should be range type agnostics, + // but we can't change the continution token format, + // since there are old SDKs that don't know how to resume from it. + FeedRangeEpk feedRangeEpk = await this.feedRangeProvider.ToEffectivePartitionKeyRangeAsync( + enumerator.Range, + cancellationToken: default); + feedRangeAndStates.Add((feedRangeEpk, enumerator.State)); } CrossPartitionState crossPartitionState = new CrossPartitionState(feedRangeAndStates); diff --git a/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionState.cs b/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionState.cs index 338652a341..4620145f59 100644 --- a/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionState.cs +++ b/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionState.cs @@ -10,11 +10,11 @@ namespace Microsoft.Azure.Cosmos.Pagination internal sealed class CrossPartitionState : State where TState : State { - public CrossPartitionState(IReadOnlyList<(FeedRange, TState)> value) + public CrossPartitionState(IReadOnlyList<(FeedRangeInternal, TState)> value) { this.Value = value ?? throw new ArgumentNullException(nameof(value)); } - public IReadOnlyList<(FeedRange, TState)> Value { get; } + public IReadOnlyList<(FeedRangeInternal, TState)> Value { get; } } } diff --git a/Microsoft.Azure.Cosmos/src/Pagination/FeedRangeProvider.cs b/Microsoft.Azure.Cosmos/src/Pagination/FeedRangeProvider.cs index 4b0cc74b4a..c972f5dccf 100644 --- a/Microsoft.Azure.Cosmos/src/Pagination/FeedRangeProvider.cs +++ b/Microsoft.Azure.Cosmos/src/Pagination/FeedRangeProvider.cs @@ -21,9 +21,9 @@ public FeedRangeProvider(CosmosQueryClient cosmosQueryClient, string collectionR this.collectionRid = collectionRid ?? throw new ArgumentNullException(nameof(collectionRid)); } - public async Task> GetChildRangeAsync( - FeedRange feedRange, - CancellationToken cancellationToken = default) + public async Task> GetChildRangeAsync( + FeedRangeInternal feedRange, + CancellationToken cancellationToken) { cancellationToken.ThrowIfCancellationRequested(); @@ -42,7 +42,7 @@ public async Task> GetChildRangeAsync( new Documents.Routing.Range(feedRangeEpk.Range.Min, feedRangeEpk.Range.Max, isMaxInclusive: true, isMinInclusive: false), forceRefresh: true); - List childFeedRanges = new List(replacementRanges.Count); + List childFeedRanges = new List(replacementRanges.Count); foreach (Documents.PartitionKeyRange replacementRange in replacementRanges) { @@ -52,9 +52,19 @@ public async Task> GetChildRangeAsync( return childFeedRanges; } - public Task> GetFeedRangesAsync( - CancellationToken cancellationToken = default) => this.GetChildRangeAsync( + public Task> GetFeedRangesAsync( + CancellationToken cancellationToken) => this.GetChildRangeAsync( FeedRangeEpk.FullRange, cancellationToken); + + public Task ToEffectivePartitionKeyRangeAsync(FeedRangeInternal feedRange, CancellationToken cancellationToken) + { + throw new NotImplementedException(); + } + + public Task ToPhysicalPartitionKeyRangeAsync(FeedRangeInternal feedRange, CancellationToken cancellationToken) + { + throw new NotImplementedException(); + } } } diff --git a/Microsoft.Azure.Cosmos/src/Pagination/IFeedRangeProvider.cs b/Microsoft.Azure.Cosmos/src/Pagination/IFeedRangeProvider.cs index 5a5cf416ce..a8c9122e18 100644 --- a/Microsoft.Azure.Cosmos/src/Pagination/IFeedRangeProvider.cs +++ b/Microsoft.Azure.Cosmos/src/Pagination/IFeedRangeProvider.cs @@ -10,10 +10,12 @@ namespace Microsoft.Azure.Cosmos.Pagination internal interface IFeedRangeProvider { - public Task> GetChildRangeAsync( - FeedRange feedRange, - CancellationToken cancellationToken = default); + public Task> GetChildRangeAsync(FeedRangeInternal feedRange, CancellationToken cancellationToken); - public Task> GetFeedRangesAsync(CancellationToken cancellationToken = default); + public Task> GetFeedRangesAsync(CancellationToken cancellationToken); + + public Task ToPhysicalPartitionKeyRangeAsync(FeedRangeInternal feedRange, CancellationToken cancellationToken); + + public Task ToEffectivePartitionKeyRangeAsync(FeedRangeInternal feedRange, CancellationToken cancellationToken); } } diff --git a/Microsoft.Azure.Cosmos/src/Pagination/PartitionRangePageEnumerable.cs b/Microsoft.Azure.Cosmos/src/Pagination/PartitionRangePageEnumerable.cs index f9ca6d40d9..f009d75cda 100644 --- a/Microsoft.Azure.Cosmos/src/Pagination/PartitionRangePageEnumerable.cs +++ b/Microsoft.Azure.Cosmos/src/Pagination/PartitionRangePageEnumerable.cs @@ -13,12 +13,12 @@ internal sealed class PartitionRangePageEnumerable : IAsyncEnumer where TPage : Page where TState : State { - private readonly FeedRange range; + private readonly FeedRangeInternal range; private readonly TState state; private readonly CreatePartitionRangePageEnumerator createPartitionRangeEnumerator; public PartitionRangePageEnumerable( - FeedRange range, + FeedRangeInternal range, TState state, CreatePartitionRangePageEnumerator createPartitionRangeEnumerator) { diff --git a/Microsoft.Azure.Cosmos/src/Pagination/PartitionRangePageEnumerator.cs b/Microsoft.Azure.Cosmos/src/Pagination/PartitionRangePageEnumerator.cs index aa35740874..b66bc920b0 100644 --- a/Microsoft.Azure.Cosmos/src/Pagination/PartitionRangePageEnumerator.cs +++ b/Microsoft.Azure.Cosmos/src/Pagination/PartitionRangePageEnumerator.cs @@ -18,13 +18,13 @@ internal abstract class PartitionRangePageEnumerator : IAsyncEnum { private bool hasStarted; - protected PartitionRangePageEnumerator(FeedRange range, TState state = null) + protected PartitionRangePageEnumerator(FeedRangeInternal range, TState state = null) { this.Range = range; this.State = state; } - public FeedRange Range { get; } + public FeedRangeInternal Range { get; } public TryCatch Current { get; private set; } diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExceptionToCosmosException.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExceptionToCosmosException.cs new file mode 100644 index 0000000000..e0ca44e49f --- /dev/null +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExceptionToCosmosException.cs @@ -0,0 +1,112 @@ +// ------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +// ------------------------------------------------------------ + +namespace Microsoft.Azure.Cosmos.Query.Core +{ + using System; + using Microsoft.Azure.Cosmos.Query.Core.Exceptions; + using Microsoft.Azure.Cosmos.Query.Core.Monads; + using Microsoft.Azure.Cosmos.Resource.CosmosExceptions; + + internal sealed class ExceptionToCosmosException + { + private static readonly string EmptyGuidString = Guid.Empty.ToString(); + + public static CosmosException CreateFromException(Exception exception) + { + if (exception is CosmosException cosmosException) + { + return cosmosException; + } + + if (exception is Microsoft.Azure.Documents.DocumentClientException documentClientException) + { + return CreateFromDocumentClientException(documentClientException); + } + + if (exception is QueryException queryException) + { + return queryException.Accept(QueryExceptionConverter.Singleton); + } + + if (exception is ExceptionWithStackTraceException exceptionWithStackTrace) + { + return CreateFromExceptionWithStackTrace(exceptionWithStackTrace); + } + + if (exception.InnerException != null) + { + // retry with the inner exception + return ExceptionToCosmosException.CreateFromException(exception.InnerException); + } + + return CosmosExceptionFactory.CreateInternalServerErrorException( + subStatusCode: default, + message: exception.Message, + stackTrace: exception.StackTrace, + activityId: EmptyGuidString, + requestCharge: 0, + retryAfter: null, + headers: null, + diagnosticsContext: null, + innerException: exception); + } + + private static CosmosException CreateFromDocumentClientException(Microsoft.Azure.Documents.DocumentClientException documentClientException) + { + CosmosException cosmosException = CosmosExceptionFactory.Create( + documentClientException, + null); + + return cosmosException; + } + + private static CosmosException CreateFromExceptionWithStackTrace(ExceptionWithStackTraceException exceptionWithStackTrace) + { + // Use the original stack trace from the inner exception. + if (exceptionWithStackTrace.InnerException is Microsoft.Azure.Documents.DocumentClientException + || exceptionWithStackTrace.InnerException is CosmosException) + { + return ExceptionToCosmosException.CreateFromException(exceptionWithStackTrace.InnerException); + } + + CosmosException cosmosException = ExceptionToCosmosException.CreateFromException(exceptionWithStackTrace.InnerException); + return CosmosExceptionFactory.Create( + cosmosException.StatusCode, + cosmosException.SubStatusCode, + cosmosException.Message, + exceptionWithStackTrace.StackTrace, + cosmosException.ActivityId, + cosmosException.RequestCharge, + cosmosException.RetryAfter, + cosmosException.Headers, + cosmosException.DiagnosticsContext, + cosmosException.Error, + cosmosException.InnerException); + } + + private sealed class QueryExceptionConverter : QueryExceptionVisitor + { + public static readonly QueryExceptionConverter Singleton = new QueryExceptionConverter(); + + private QueryExceptionConverter() + { + } + + public override CosmosException Visit(MalformedContinuationTokenException malformedContinuationTokenException) => CosmosExceptionFactory.CreateBadRequestException( + message: malformedContinuationTokenException.Message, + stackTrace: malformedContinuationTokenException.StackTrace, + innerException: malformedContinuationTokenException); + + public override CosmosException Visit(UnexpectedQueryPartitionProviderException unexpectedQueryPartitionProviderException) => CosmosExceptionFactory.CreateInternalServerErrorException( + message: $"{nameof(CosmosException)} due to {nameof(UnexpectedQueryPartitionProviderException)}", + innerException: unexpectedQueryPartitionProviderException); + + public override CosmosException Visit(ExpectedQueryPartitionProviderException expectedQueryPartitionProviderException) => CosmosExceptionFactory.CreateBadRequestException( + message: expectedQueryPartitionProviderException.Message, + stackTrace: expectedQueryPartitionProviderException.StackTrace, + innerException: expectedQueryPartitionProviderException); + } + } +} diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Monads/TryCatch{TResult}.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Monads/TryCatch{TResult}.cs index 74588e157c..d7bf48d9ec 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Monads/TryCatch{TResult}.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Monads/TryCatch{TResult}.cs @@ -66,7 +66,7 @@ public Exception InnerMostException get { Exception exception = this.Exception; - while (exception != null) + while (exception.InnerException != null) { exception = exception.InnerException; } diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CatchAllQueryPipelineStage.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CatchAllQueryPipelineStage.cs index db467fb529..0a5af1767f 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CatchAllQueryPipelineStage.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CatchAllQueryPipelineStage.cs @@ -31,9 +31,8 @@ protected override async Task> GetNextPageAsync(Cancellation } catch (Exception ex) { - QueryResponseCore queryResponse = QueryResponseFactory.CreateFromException(ex); - CosmosException exception = queryResponse.CosmosException; - return TryCatch.FromException(exception); + CosmosException cosmosException = ExceptionToCosmosException.CreateFromException(ex); + return TryCatch.FromException(cosmosException); } } } diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/BackendQueryDataSource.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/BackendQueryDataSource.cs index 72f4af5fd6..ee33cacb6f 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/BackendQueryDataSource.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/BackendQueryDataSource.cs @@ -5,6 +5,7 @@ namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote { using System; + using System.Collections.Generic; using System.Threading; using System.Threading.Tasks; using Microsoft.Azure.Cosmos.Query.Core.Monads; @@ -13,26 +14,117 @@ namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote internal sealed class BackendQueryDataSource : IQueryDataSource { - private readonly CosmosQueryContext cosmosQueryContext; + private readonly ExecuteQueryBasedOnFeedRangeVisitor visitor; public BackendQueryDataSource(CosmosQueryContext cosmosQueryContext) { - this.cosmosQueryContext = cosmosQueryContext ?? throw new ArgumentNullException(nameof(cosmosQueryContext)); + this.visitor = new ExecuteQueryBasedOnFeedRangeVisitor(cosmosQueryContext); } public Task> ExecuteQueryAsync( SqlQuerySpec sqlQuerySpec, string continuationToken, - int partitionKeyRangeId, + FeedRangeInternal feedRange, int pageSize, - CancellationToken cancellationToken) => this.cosmosQueryContext.ExecuteQueryAsync( - querySpecForInit: sqlQuerySpec, - continuationToken: continuationToken, - partitionKeyRange: new PartitionKeyRangeIdentity( + CancellationToken cancellationToken) => feedRange.AcceptAsync( + this.visitor, + new VisitorArguments(sqlQuerySpec, continuationToken, pageSize), + cancellationToken); + + private readonly struct VisitorArguments + { + public VisitorArguments(SqlQuerySpec sqlQuerySpec, string continuationToken, int pageSize) + { + this.SqlQuerySpec = sqlQuerySpec; + this.ContinuationToken = continuationToken; + this.PageSize = pageSize; + } + + public SqlQuerySpec SqlQuerySpec { get; } + + public string ContinuationToken { get; } + + public int PageSize { get; } + } + + private sealed class ExecuteQueryBasedOnFeedRangeVisitor : IFeedRangeAsyncVisitor, VisitorArguments> + { + private readonly CosmosQueryContext cosmosQueryContext; + + public ExecuteQueryBasedOnFeedRangeVisitor(CosmosQueryContext cosmosQueryContext) + { + this.cosmosQueryContext = cosmosQueryContext ?? throw new ArgumentNullException(nameof(cosmosQueryContext)); + } + + public Task> VisitAsync( + FeedRangePartitionKey feedRange, + VisitorArguments argument, + CancellationToken cancellationToken) => this.cosmosQueryContext.QueryClient.ExecuteItemQueryAsync( + this.cosmosQueryContext.ResourceLink, + this.cosmosQueryContext.ResourceTypeEnum, + this.cosmosQueryContext.OperationTypeEnum, + this.cosmosQueryContext.CorrelatedActivityId, + new QueryRequestOptions() + { + PartitionKey = feedRange.PartitionKey, + }, + queryPageDiagnostics: default, + argument.SqlQuerySpec, + argument.ContinuationToken, + partitionKeyRange: default, + isContinuationExpected: true, + argument.PageSize, + cancellationToken); + + public Task> VisitAsync( + FeedRangePartitionKeyRange feedRange, + VisitorArguments argument, + CancellationToken cancellationToken) => this.cosmosQueryContext.ExecuteQueryAsync( + querySpecForInit: argument.SqlQuerySpec, + continuationToken: argument.ContinuationToken, + partitionKeyRange: new PartitionKeyRangeIdentity( + this.cosmosQueryContext.ContainerResourceId, + feedRange.PartitionKeyRangeId), + isContinuationExpected: this.cosmosQueryContext.IsContinuationExpected, + pageSize: argument.PageSize, + cancellationToken: cancellationToken); + + public async Task> VisitAsync( + FeedRangeEpk feedRange, + VisitorArguments argument, + CancellationToken cancellationToken) + { + // Check to see if it lines up exactly with one physical partition + IReadOnlyList feedRanges = await this.cosmosQueryContext.QueryClient.TryGetOverlappingRangesAsync( this.cosmosQueryContext.ContainerResourceId, - partitionKeyRangeId.ToString()), - isContinuationExpected: this.cosmosQueryContext.IsContinuationExpected, - pageSize: pageSize, - cancellationToken: cancellationToken); + feedRange.Range, + forceRefresh: true); + + if (feedRanges.Count != 1) + { + // Simulate a split exception, since we don't have a partition key range id to route to. + CosmosException goneException = new CosmosException( + message: $"Epk Range: {feedRange.Range} is gone.", + statusCode: System.Net.HttpStatusCode.Gone, + subStatusCode: (int)SubStatusCodes.PartitionKeyRangeGone, + activityId: Guid.NewGuid().ToString(), + requestCharge: default); + + return TryCatch.FromException(goneException); + } + + // If the epk range aligns exactly to a physical partition, then continue as if PKRangeId was given + PartitionKeyRange singleRange = feedRanges[0]; + Documents.Routing.Range range = singleRange.ToRange(); + if (!feedRange.Range.Equals(range)) + { + // User want's use to query on a sub partition, which is currently not possible. + throw new NotImplementedException("Can not query on a sub partition."); + } + + FeedRangeInternal partitionKeyRangeFeedRange = new FeedRangePartitionKeyRange(singleRange.Id); + return await partitionKeyRangeFeedRange.AcceptAsync(this, argument, cancellationToken); + } + } } } diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/IQueryDataSource.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/IQueryDataSource.cs index b2761f7c18..6d90772469 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/IQueryDataSource.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/IQueryDataSource.cs @@ -13,7 +13,7 @@ internal interface IQueryDataSource public Task> ExecuteQueryAsync( SqlQuerySpec sqlQuerySpec, string continuationToken, - int partitionKeyRangeId, + FeedRangeInternal feedRange, int pageSize, CancellationToken cancellationToken); } diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/ParallelCrossPartitionQueryPipelineStage.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/ParallelCrossPartitionQueryPipelineStage.cs index b67d7aa306..870a3d1ece 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/ParallelCrossPartitionQueryPipelineStage.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/ParallelCrossPartitionQueryPipelineStage.cs @@ -4,6 +4,7 @@ namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote { +#nullable enable using System; using System.Collections.Generic; using System.Linq; @@ -39,7 +40,7 @@ public TryCatch Current CrossPartitionState crossPartitionState = crossPartitionPageResult.State; List compositeContinuationTokens = new List(crossPartitionState.Value.Count); - foreach ((FeedRange range, QueryState state) in crossPartitionState.Value) + foreach ((FeedRangeInternal range, QueryState state) in crossPartitionState.Value) { if (!(range is FeedRangeEpk epkRange)) { @@ -49,13 +50,15 @@ public TryCatch Current CompositeContinuationToken compositeContinuationToken = new CompositeContinuationToken() { Range = epkRange.Range, - Token = ((CosmosString)state.Value).Value, + Token = state != null ? ((CosmosString)state.Value).Value : null, }; compositeContinuationTokens.Add(compositeContinuationToken); } - List cosmosElementContinuationTokens = compositeContinuationTokens.Select(token => CompositeContinuationToken.ToCosmosElement(token)).ToList(); + List cosmosElementContinuationTokens = compositeContinuationTokens + .Select(token => CompositeContinuationToken.ToCosmosElement(token)) + .ToList(); CosmosArray cosmosElementCompositeContinuationTokens = CosmosArray.Create(cosmosElementContinuationTokens); QueryPage crossPartitionQueryPage = new QueryPage( @@ -80,9 +83,14 @@ public static TryCatch MonadicCreate( IQueryDataSource queryDataSource, SqlQuerySpec sqlQuerySpec, int pageSize, - CosmosElement continuationToken) + CosmosElement? continuationToken) { - CrossPartitionState state; + if (pageSize <= 0) + { + throw new ArgumentOutOfRangeException(nameof(pageSize)); + } + + CrossPartitionState? state; if (continuationToken == null) { state = default; @@ -96,6 +104,13 @@ public static TryCatch MonadicCreate( $"Invalid format for continuation token {continuationToken} for {nameof(ParallelCrossPartitionQueryPipelineStage)}")); } + if (compositeContinuationTokenListRaw.Count == 0) + { + return TryCatch.FromException( + new MalformedContinuationTokenException( + $"Invalid format for continuation token {continuationToken} for {nameof(ParallelCrossPartitionQueryPipelineStage)}")); + } + List compositeContinuationTokens = new List(); foreach (CosmosElement compositeContinuationTokenRaw in compositeContinuationTokenListRaw) { @@ -109,8 +124,8 @@ public static TryCatch MonadicCreate( compositeContinuationTokens.Add(tryCreateCompositeContinuationToken.Result); } - List<(FeedRange, QueryState)> rangesAndStates = compositeContinuationTokens - .Select(token => ((FeedRange)new FeedRangeEpk(token.Range), new QueryState(CosmosString.Create(token.Token)))) + List<(FeedRangeInternal, QueryState)> rangesAndStates = compositeContinuationTokens + .Select(token => ((FeedRangeInternal)new FeedRangeEpk(token.Range), new QueryState(CosmosString.Create(token.Token)))) .ToList(); state = new CrossPartitionState(rangesAndStates); @@ -129,7 +144,7 @@ public static TryCatch MonadicCreate( private static CreatePartitionRangePageEnumerator MakeCreateFunction( IQueryDataSource queryDataSource, SqlQuerySpec sqlQuerySpec, - int pageSize) => (FeedRange range, QueryState state) => new QueryPartitionRangePageEnumerator( + int pageSize) => (FeedRangeInternal range, QueryState state) => new QueryPartitionRangePageEnumerator( queryDataSource, sqlQuerySpec, range, diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/QueryPartitionRangePageEnumerator.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/QueryPartitionRangePageEnumerator.cs index c904ad16c3..050c4c738b 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/QueryPartitionRangePageEnumerator.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/QueryPartitionRangePageEnumerator.cs @@ -20,7 +20,7 @@ internal sealed class QueryPartitionRangePageEnumerator : PartitionRangePageEnum public QueryPartitionRangePageEnumerator( IQueryDataSource queryDataSource, SqlQuerySpec sqlQuerySpec, - FeedRange feedRange, + FeedRangeInternal feedRange, int pageSize, QueryState state = default) : base(feedRange, state) @@ -28,17 +28,12 @@ public QueryPartitionRangePageEnumerator( this.queryDataSource = queryDataSource ?? throw new ArgumentNullException(nameof(queryDataSource)); this.sqlQuerySpec = sqlQuerySpec ?? throw new ArgumentNullException(nameof(sqlQuerySpec)); this.pageSize = pageSize; - - if (!(feedRange is FeedRangePartitionKeyRange)) - { - throw new ArgumentOutOfRangeException(nameof(feedRange)); - } } public override Task> GetNextPageAsync(CancellationToken cancellationToken) => this.queryDataSource.ExecuteQueryAsync( sqlQuerySpec: this.sqlQuerySpec, continuationToken: this.State == null ? null : ((CosmosString)this.State.Value).Value, - partitionKeyRangeId: int.Parse(((FeedRangePartitionKeyRange)this.Range).PartitionKeyRangeId), + feedRange: this.Range, pageSize: this.pageSize, cancellationToken); diff --git a/Microsoft.Azure.Cosmos/src/Query/v3Query/CosmosQueryClientCore.cs b/Microsoft.Azure.Cosmos/src/Query/v3Query/CosmosQueryClientCore.cs index 1038deea1e..1a45e9cdd0 100644 --- a/Microsoft.Azure.Cosmos/src/Query/v3Query/CosmosQueryClientCore.cs +++ b/Microsoft.Azure.Cosmos/src/Query/v3Query/CosmosQueryClientCore.cs @@ -16,7 +16,6 @@ namespace Microsoft.Azure.Cosmos using Microsoft.Azure.Cosmos.Diagnostics; using Microsoft.Azure.Cosmos.Json; using Microsoft.Azure.Cosmos.Query.Core; - using Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.ItemProducers; using Microsoft.Azure.Cosmos.Query.Core.Monads; using Microsoft.Azure.Cosmos.Query.Core.Pipeline; using Microsoft.Azure.Cosmos.Query.Core.QueryClient; diff --git a/Microsoft.Azure.Cosmos/src/Query/v3Query/CosmosQueryContextCore.cs b/Microsoft.Azure.Cosmos/src/Query/v3Query/CosmosQueryContextCore.cs index 7522abc00f..70f22fc035 100644 --- a/Microsoft.Azure.Cosmos/src/Query/v3Query/CosmosQueryContextCore.cs +++ b/Microsoft.Azure.Cosmos/src/Query/v3Query/CosmosQueryContextCore.cs @@ -8,8 +8,6 @@ namespace Microsoft.Azure.Cosmos.Query using System.Threading.Tasks; using Microsoft.Azure.Cosmos.Diagnostics; using Microsoft.Azure.Cosmos.Query.Core; - using Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.ItemProducers; - using Microsoft.Azure.Cosmos.Query.Core.Metrics; using Microsoft.Azure.Cosmos.Query.Core.Monads; using Microsoft.Azure.Cosmos.Query.Core.Pipeline; using Microsoft.Azure.Cosmos.Query.Core.QueryClient; diff --git a/Microsoft.Azure.Cosmos/src/Query/v3Query/QueryIterator.cs b/Microsoft.Azure.Cosmos/src/Query/v3Query/QueryIterator.cs index 8c354325ce..bcb37ac781 100644 --- a/Microsoft.Azure.Cosmos/src/Query/v3Query/QueryIterator.cs +++ b/Microsoft.Azure.Cosmos/src/Query/v3Query/QueryIterator.cs @@ -178,7 +178,7 @@ public override async Task ReadNextAsync(CancellationToken canc }); } - CosmosException cosmosException = (CosmosException)tryGetQueryPage.InnerMostException; + CosmosException cosmosException = ExceptionToCosmosException.CreateFromException(tryGetQueryPage.Exception); SubStatusCodes subStatusCode; if (Enum.IsDefined(typeof(SubStatusCodes), cosmosException.SubStatusCode)) { diff --git a/Microsoft.Azure.Cosmos/src/Resource/CosmosClientContext.cs b/Microsoft.Azure.Cosmos/src/Resource/CosmosClientContext.cs index 53d6d8009e..815a2eb1eb 100644 --- a/Microsoft.Azure.Cosmos/src/Resource/CosmosClientContext.cs +++ b/Microsoft.Azure.Cosmos/src/Resource/CosmosClientContext.cs @@ -5,13 +5,10 @@ namespace Microsoft.Azure.Cosmos { using System; - using System.Globalization; using System.IO; - using System.Text; using System.Threading; using System.Threading.Tasks; using Microsoft.Azure.Cosmos.Handlers; - using Microsoft.Azure.Cosmos.Query; using Microsoft.Azure.Documents; /// diff --git a/Microsoft.Azure.Cosmos/src/Resource/CosmosExceptions/CosmosException.cs b/Microsoft.Azure.Cosmos/src/Resource/CosmosExceptions/CosmosException.cs index 95a18e8993..0d319db99e 100644 --- a/Microsoft.Azure.Cosmos/src/Resource/CosmosExceptions/CosmosException.cs +++ b/Microsoft.Azure.Cosmos/src/Resource/CosmosExceptions/CosmosException.cs @@ -17,7 +17,7 @@ public class CosmosException : Exception private readonly string stackTrace; internal CosmosException( - HttpStatusCode statusCodes, + HttpStatusCode statusCode, string message, int subStatusCode, string stackTrace, @@ -29,7 +29,7 @@ internal CosmosException( Error error, Exception innerException) : base(CosmosException.GetMessageHelper( - statusCodes, + statusCode, subStatusCode, message, activityId), innerException) @@ -37,7 +37,7 @@ internal CosmosException( this.ResponseBody = message; this.stackTrace = stackTrace; this.ActivityId = activityId; - this.StatusCode = statusCodes; + this.StatusCode = statusCode; this.SubStatusCode = subStatusCode; this.RetryAfter = retryAfter; this.RequestCharge = requestCharge; diff --git a/Microsoft.Azure.Cosmos/src/Routing/PartitionKeyHash.cs b/Microsoft.Azure.Cosmos/src/Routing/PartitionKeyHash.cs index 2633236978..fe35c9e832 100644 --- a/Microsoft.Azure.Cosmos/src/Routing/PartitionKeyHash.cs +++ b/Microsoft.Azure.Cosmos/src/Routing/PartitionKeyHash.cs @@ -61,14 +61,22 @@ public override bool Equals(object obj) return this.Equals(effectivePartitionKey); } - public bool Equals(PartitionKeyHash other) - { - return this.Value.Equals(other.Value); - } + public bool Equals(PartitionKeyHash other) => this.Value.Equals(other.Value); + + public override int GetHashCode() => this.Value.GetHashCode(); + + public override string ToString() => this.Value.ToString(); - public override int GetHashCode() + public static bool TryParse(string value, out PartitionKeyHash parsedValue) { - return this.Value.GetHashCode(); + if (!UInt128.TryParse(value, out UInt128 uInt128)) + { + parsedValue = default; + return false; + } + + parsedValue = new PartitionKeyHash(uInt128); + return true; } public static class V1 diff --git a/Microsoft.Azure.Cosmos/src/Routing/PartitionKeyHashRange.cs b/Microsoft.Azure.Cosmos/src/Routing/PartitionKeyHashRange.cs index a0aecb5676..34d16bee88 100644 --- a/Microsoft.Azure.Cosmos/src/Routing/PartitionKeyHashRange.cs +++ b/Microsoft.Azure.Cosmos/src/Routing/PartitionKeyHashRange.cs @@ -6,6 +6,7 @@ namespace Microsoft.Azure.Cosmos.Routing { using System; using System.Text; + using Microsoft.Azure.Documents; internal readonly struct PartitionKeyHashRange : IComparable, IEquatable { diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/Query/SanityQueryTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/Query/SanityQueryTests.cs index 228db1a6d6..c1fd74e5f4 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/Query/SanityQueryTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/Query/SanityQueryTests.cs @@ -20,6 +20,33 @@ namespace Microsoft.Azure.Cosmos.EmulatorTests.Query [TestClass] public sealed class SanityQueryTests : QueryTestsBase { + + [TestMethod] + public async Task Sanity() + { + int seed = (int)(DateTime.UtcNow - new DateTime(1970, 1, 1)).TotalSeconds; + uint numberOfDocuments = 100; + QueryOracleUtil util = new QueryOracle2(seed); + IEnumerable inputDocuments = util.GetDocuments(numberOfDocuments); + + await this.CreateIngestQueryDeleteAsync( + ConnectionModes.Direct, + CollectionTypes.SinglePartition | CollectionTypes.MultiPartition, + inputDocuments, + ImplementationAsync); + + async Task ImplementationAsync(Container container, IReadOnlyList documents) + { + List queryResults = await QueryTestsBase.RunQueryAsync( + container, + "SELECT * FROM c"); + + Assert.AreEqual( + documents.Count(), + queryResults.Count); + } + } + [TestMethod] public async Task TestBasicCrossPartitionQueryAsync() { @@ -86,7 +113,7 @@ async Task ImplementationAsync(Container container, IReadOnlyList GC.WaitForPendingFinalizers(); GC.Collect(); - foreach(WeakReference weakReference in weakReferences) + foreach (WeakReference weakReference in weakReferences) { Assert.IsFalse(weakReference.IsAlive); } diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/CosmosQueryUnitTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/CosmosQueryUnitTests.cs index fe0a221e2c..74cb6eaa20 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/CosmosQueryUnitTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/CosmosQueryUnitTests.cs @@ -380,7 +380,7 @@ public async Task TestCosmosQueryPartitionKeyDefinition() System.Net.HttpStatusCode.Unauthorized, SubStatusCodes.PartitionKeyMismatch, new CosmosException( - statusCodes: HttpStatusCode.Unauthorized, + statusCode: HttpStatusCode.Unauthorized, message: "Random error message", subStatusCode: default, stackTrace: default, diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/InMemoryCollection.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/InMemoryCollection.cs index 6f36c5457e..eb8f28bdbc 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/InMemoryCollection.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/InMemoryCollection.cs @@ -133,7 +133,11 @@ record = default; return TryCatch<(List, long?)>.FromResult((page, page.Last().ResourceIdentifier)); } - public TryCatch<(List, long?)> Query(SqlQuerySpec sqlQuerySpec, int partitionKeyRangeId, long resourceIdentifier, int pageSize) + public TryCatch<(List, long?)> Query( + SqlQuerySpec sqlQuerySpec, + int partitionKeyRangeId, + long resourceIdentifier, + int pageSize) { if (sqlQuerySpec == null) { @@ -221,9 +225,13 @@ public void Split(int partitionKeyRangeId) public (int, int) GetChildRanges(int partitionKeyRangeId) => this.parentToChildMapping[partitionKeyRangeId]; - public bool Return429() => (this.failureConfigs != null) && this.failureConfigs.Inject429s && ((this.random.Next() % 2) == 0); + public PartitionKeyHashRange GetHashRange(int partitionKeyRangeId) => this.partitionKeyRangeIdToHashRange[partitionKeyRangeId]; - public bool ReturnEmptyPage() => (this.failureConfigs != null) && this.failureConfigs.InjectEmptyPages && ((this.random.Next() % 2) == 0); + public int GetPartitionKeyRangeId(PartitionKeyHashRange hashRange) => this.partitionKeyRangeIdToHashRange.Where(kvp => kvp.Value.Equals(hashRange)).First().Key; + + private bool Return429() => (this.failureConfigs != null) && this.failureConfigs.Inject429s && ((this.random.Next() % 2) == 0); + + private bool ReturnEmptyPage() => (this.failureConfigs != null) && this.failureConfigs.InjectEmptyPages && ((this.random.Next() % 2) == 0); private static PartitionKeyHash GetHashFromPayload(CosmosObject payload, PartitionKeyDefinition partitionKeyDefinition) { diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/InMemoryCollectionFeedRangeProvider.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/InMemoryCollectionFeedRangeProvider.cs index c7134016db..6761121d84 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/InMemoryCollectionFeedRangeProvider.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/InMemoryCollectionFeedRangeProvider.cs @@ -6,50 +6,149 @@ using System.Threading.Tasks; using Microsoft.Azure.Cosmos; using Microsoft.Azure.Cosmos.Pagination; + using Microsoft.Azure.Cosmos.Routing; internal sealed class InMemoryCollectionFeedRangeProvider : IFeedRangeProvider { private readonly InMemoryCollection inMemoryCollection; + private readonly FeedRangeToEffectivePartitionKeyRangeVisitor feedRangeToEffectivePartitionKeyRangeVisitor; + private readonly FeedRangeToPhysicalPartitionKeyRange feedRangeToPhysicalPartitionKeyRange; public InMemoryCollectionFeedRangeProvider(InMemoryCollection inMemoryCollection) { this.inMemoryCollection = inMemoryCollection ?? throw new ArgumentNullException(nameof(inMemoryCollection)); + this.feedRangeToEffectivePartitionKeyRangeVisitor = new FeedRangeToEffectivePartitionKeyRangeVisitor(this.inMemoryCollection); + this.feedRangeToPhysicalPartitionKeyRange = new FeedRangeToPhysicalPartitionKeyRange(this.inMemoryCollection); } - public Task> GetChildRangeAsync( - FeedRange feedRange, - CancellationToken cancellationToken = default) + public async Task> GetChildRangeAsync( + FeedRangeInternal feedRange, + CancellationToken cancellationToken) { cancellationToken.ThrowIfCancellationRequested(); - if (!(feedRange is FeedRangePartitionKeyRange feedRangePartitionKeyRange)) - { - throw new ArgumentOutOfRangeException(nameof(feedRange)); - } + FeedRangePartitionKeyRange feedRangePartitionKeyRange = await this.ToPhysicalPartitionKeyRangeAsync(feedRange, cancellationToken); int partitionKeyRangeId = int.Parse(feedRangePartitionKeyRange.PartitionKeyRangeId); (int leftChild, int rightChild) = this.inMemoryCollection.GetChildRanges(partitionKeyRangeId); - return Task.FromResult( - (IEnumerable)new List() - { - new FeedRangePartitionKeyRange(leftChild.ToString()), - new FeedRangePartitionKeyRange(rightChild.ToString()), - }); + return new List() + { + new FeedRangePartitionKeyRange(leftChild.ToString()), + new FeedRangePartitionKeyRange(rightChild.ToString()), + }; } - public Task> GetFeedRangesAsync(CancellationToken cancellationToken = default) + public Task> GetFeedRangesAsync(CancellationToken cancellationToken) { cancellationToken.ThrowIfCancellationRequested(); - List ranges = new List(); + List ranges = new List(); foreach (int partitionKeyRangeId in this.inMemoryCollection.PartitionKeyRangeFeedReed().Keys) { - FeedRange range = new FeedRangePartitionKeyRange(partitionKeyRangeId.ToString()); + FeedRangeInternal range = new FeedRangePartitionKeyRange(partitionKeyRangeId.ToString()); ranges.Add(range); } - return Task.FromResult((IEnumerable)ranges); + return Task.FromResult((IEnumerable)ranges); + } + + public Task ToEffectivePartitionKeyRangeAsync( + FeedRangeInternal feedRange, + CancellationToken cancellationToken) => feedRange.AcceptAsync(this.feedRangeToEffectivePartitionKeyRangeVisitor); + + public Task ToPhysicalPartitionKeyRangeAsync( + FeedRangeInternal feedRange, + CancellationToken cancellationToken) => feedRange.AcceptAsync(this.feedRangeToPhysicalPartitionKeyRange); + + private sealed class FeedRangeToEffectivePartitionKeyRangeVisitor : IFeedRangeAsyncVisitor + { + private readonly InMemoryCollection inMemoryCollection; + + public FeedRangeToEffectivePartitionKeyRangeVisitor(InMemoryCollection inMemoryCollection) + { + this.inMemoryCollection = inMemoryCollection ?? throw new ArgumentNullException(nameof(inMemoryCollection)); + } + + public Task VisitAsync(FeedRangePartitionKey feedRange, CancellationToken cancellationToken = default) + { + throw new NotImplementedException(); + } + + public Task VisitAsync(FeedRangePartitionKeyRange feedRange, CancellationToken cancellationToken = default) + { + PartitionKeyHashRange hashRange = this.inMemoryCollection.GetHashRange(int.Parse(feedRange.PartitionKeyRangeId)); + Documents.Routing.Range range = new Documents.Routing.Range( + min: hashRange.StartInclusive.HasValue ? hashRange.StartInclusive.Value.ToString() : string.Empty, + max: hashRange.EndExclusive.HasValue ? hashRange.EndExclusive.Value.ToString() : string.Empty, + isMinInclusive: true, + isMaxInclusive: false); + + return Task.FromResult(new FeedRangeEpk(range)); + } + + public Task VisitAsync(FeedRangeEpk feedRange, CancellationToken cancellationToken = default) + { + return Task.FromResult(feedRange); + } + } + + private sealed class FeedRangeToPhysicalPartitionKeyRange : IFeedRangeAsyncVisitor + { + private readonly InMemoryCollection inMemoryCollection; + + public FeedRangeToPhysicalPartitionKeyRange(InMemoryCollection inMemoryCollection) + { + this.inMemoryCollection = inMemoryCollection ?? throw new ArgumentNullException(nameof(inMemoryCollection)); + } + + public Task VisitAsync(FeedRangePartitionKey feedRange, CancellationToken cancellationToken = default) + { + throw new NotImplementedException(); + } + + public Task VisitAsync(FeedRangePartitionKeyRange feedRange, CancellationToken cancellationToken = default) + { + return Task.FromResult(feedRange); + } + + public Task VisitAsync(FeedRangeEpk feedRange, CancellationToken cancellationToken = default) + { + Documents.Routing.Range range = feedRange.Range; + + PartitionKeyHash? startInclusive; + if (range.Min == string.Empty) + { + startInclusive = default; + } + else + { + if (!PartitionKeyHash.TryParse(range.Min, out PartitionKeyHash parsedValue)) + { + throw new InvalidOperationException($"Failed to parse: {range.Min}."); + } + + startInclusive = parsedValue; + } + + PartitionKeyHash? endExclusive; + if (range.Max == string.Empty) + { + endExclusive = default; + } + else + { + if (!PartitionKeyHash.TryParse(range.Max, out PartitionKeyHash parsedValue)) + { + throw new InvalidOperationException($"Failed to parse: {range.Max}."); + } + + endExclusive = parsedValue; + } + + int pkRangeId = this.inMemoryCollection.GetPartitionKeyRangeId(new PartitionKeyHashRange(startInclusive, endExclusive)); + return Task.FromResult(new FeedRangePartitionKeyRange(pkRangeId.ToString())); + } } } } diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/InMemoryCollectionQueryDataSource.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/InMemoryCollectionQueryDataSource.cs index ca9e1e2890..b9937a77e6 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/InMemoryCollectionQueryDataSource.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/InMemoryCollectionQueryDataSource.cs @@ -27,13 +27,13 @@ public InMemoryCollectionQueryDataSource(InMemoryCollection inMemoryCollection) public Task> ExecuteQueryAsync( SqlQuerySpec sqlQuerySpec, string continuationToken, - int partitionKeyRangeId, + FeedRangeInternal feedRange, int pageSize, CancellationToken cancellationToken) { TryCatch<(List records, long? resourceIdentifer)> tryExecuteQuery = this.inMemoryCollection.Query( sqlQuerySpec, - partitionKeyRangeId, + int.Parse(((FeedRangePartitionKeyRange)feedRange).PartitionKeyRangeId), continuationToken != null ? long.Parse(continuationToken) : 0, pageSize); if (tryExecuteQuery.Failed) 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 new file mode 100644 index 0000000000..8fbef7357a --- /dev/null +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/ParallelCrossPartitionQueryPipelineStageTests.cs @@ -0,0 +1,196 @@ +//------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +//------------------------------------------------------------ + +namespace Microsoft.Azure.Cosmos.Tests.Query.Pipeline +{ + using System.Collections.Generic; + using System.Threading.Tasks; + using Microsoft.Azure.Cosmos.CosmosElements; + using Microsoft.Azure.Cosmos.Pagination; + using Microsoft.Azure.Cosmos.Query.Core; + using Microsoft.Azure.Cosmos.Query.Core.ContinuationTokens; + using Microsoft.Azure.Cosmos.Query.Core.Exceptions; + using Microsoft.Azure.Cosmos.Query.Core.Monads; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote; + using Microsoft.Azure.Cosmos.SqlObjects; + using Microsoft.Azure.Cosmos.Tests.Pagination; + using Microsoft.Azure.Documents; + using Microsoft.VisualStudio.TestTools.UnitTesting; + using Moq; + + [TestClass] + public class ParallelCrossPartitionQueryPipelineStageTests + { + [TestMethod] + public void MonadicCreate_NullContinuationToken() + { + Mock mockFeedRangeProvider = new Mock(); + Mock mockQueryDataSource = new Mock(); + TryCatch monadicCreate = ParallelCrossPartitionQueryPipelineStage.MonadicCreate( + feedRangeProvider: mockFeedRangeProvider.Object, + queryDataSource: mockQueryDataSource.Object, + sqlQuerySpec: new SqlQuerySpec("SELECT * FROM c"), + pageSize: 10, + continuationToken: null); + Assert.IsTrue(monadicCreate.Succeeded); + } + + [TestMethod] + public void MonadicCreate_NonCosmosArrayContinuationToken() + { + Mock mockFeedRangeProvider = new Mock(); + Mock mockQueryDataSource = new Mock(); + TryCatch monadicCreate = ParallelCrossPartitionQueryPipelineStage.MonadicCreate( + feedRangeProvider: mockFeedRangeProvider.Object, + queryDataSource: mockQueryDataSource.Object, + sqlQuerySpec: new SqlQuerySpec("SELECT * FROM c"), + pageSize: 10, + continuationToken: CosmosObject.Create(new Dictionary())); + Assert.IsTrue(monadicCreate.Failed); + Assert.IsTrue(monadicCreate.InnerMostException is MalformedContinuationTokenException); + } + + [TestMethod] + public void MonadicCreate_EmptyArrayContinuationToken() + { + Mock mockFeedRangeProvider = new Mock(); + Mock mockQueryDataSource = new Mock(); + TryCatch monadicCreate = ParallelCrossPartitionQueryPipelineStage.MonadicCreate( + feedRangeProvider: mockFeedRangeProvider.Object, + queryDataSource: mockQueryDataSource.Object, + sqlQuerySpec: new SqlQuerySpec("SELECT * FROM c"), + pageSize: 10, + continuationToken: CosmosArray.Create(new List())); + Assert.IsTrue(monadicCreate.Failed); + Assert.IsTrue(monadicCreate.InnerMostException is MalformedContinuationTokenException); + } + + [TestMethod] + public void MonadicCreate_NonCompositeContinuationToken() + { + Mock mockFeedRangeProvider = new Mock(); + Mock mockQueryDataSource = new Mock(); + TryCatch monadicCreate = ParallelCrossPartitionQueryPipelineStage.MonadicCreate( + feedRangeProvider: mockFeedRangeProvider.Object, + queryDataSource: mockQueryDataSource.Object, + sqlQuerySpec: new SqlQuerySpec("SELECT * FROM c"), + pageSize: 10, + continuationToken: CosmosArray.Create(new List() { CosmosString.Create("asdf") })); + Assert.IsTrue(monadicCreate.Failed); + Assert.IsTrue(monadicCreate.InnerMostException is MalformedContinuationTokenException); + } + + [TestMethod] + public void MonadicCreate_SingleCompositeContinuationToken() + { + Mock mockFeedRangeProvider = new Mock(); + Mock mockQueryDataSource = new Mock(); + + CompositeContinuationToken token = new CompositeContinuationToken() + { + Range = new Documents.Routing.Range("A", "B", true, false), + Token = "asdf", + }; + + TryCatch monadicCreate = ParallelCrossPartitionQueryPipelineStage.MonadicCreate( + feedRangeProvider: mockFeedRangeProvider.Object, + queryDataSource: mockQueryDataSource.Object, + sqlQuerySpec: new SqlQuerySpec("SELECT * FROM c"), + pageSize: 10, + continuationToken: CosmosArray.Create(new List() { CompositeContinuationToken.ToCosmosElement(token) })); + Assert.IsTrue(monadicCreate.Succeeded); + } + + [TestMethod] + public void MonadicCreate_MultipleCompositeContinuationToken() + { + Mock mockFeedRangeProvider = new Mock(); + Mock mockQueryDataSource = new Mock(); + + CompositeContinuationToken token = new CompositeContinuationToken() + { + Range = new Documents.Routing.Range("A", "B", true, false), + Token = "asdf", + }; + + TryCatch monadicCreate = ParallelCrossPartitionQueryPipelineStage.MonadicCreate( + feedRangeProvider: mockFeedRangeProvider.Object, + queryDataSource: mockQueryDataSource.Object, + sqlQuerySpec: new SqlQuerySpec("SELECT * FROM c"), + pageSize: 10, + continuationToken: CosmosArray.Create( + new List() + { + CompositeContinuationToken.ToCosmosElement(token), + CompositeContinuationToken.ToCosmosElement(token) + })); + Assert.IsTrue(monadicCreate.Succeeded); + } + + [TestMethod] + public async Task TestDrainFully_StartFromBeginingAsync() + { + int numItems = 1000; + InMemoryCollection inMemoryCollection = CreateInMemoryCollection(numItems); + IFeedRangeProvider feedRangeProvider = new InMemoryCollectionFeedRangeProvider(inMemoryCollection); + IQueryDataSource queryDataSource = new InMemoryCollectionQueryDataSource(inMemoryCollection); + + TryCatch monadicCreate = ParallelCrossPartitionQueryPipelineStage.MonadicCreate( + feedRangeProvider: feedRangeProvider, + queryDataSource: queryDataSource, + sqlQuerySpec: new SqlQuerySpec("SELECT * FROM c"), + pageSize: 10, + continuationToken: default); + Assert.IsTrue(monadicCreate.Succeeded); + IQueryPipelineStage queryPipelineStage = monadicCreate.Result; + + List documents = new List(); + while (await queryPipelineStage.MoveNextAsync()) + { + TryCatch tryGetQueryPage = queryPipelineStage.Current; + Assert.IsTrue(tryGetQueryPage.Succeeded); + + QueryPage queryPage = tryGetQueryPage.Result; + documents.AddRange(queryPage.Documents); + } + + Assert.AreEqual(numItems, documents.Count); + } + + private static InMemoryCollection CreateInMemoryCollection(int numItems, InMemoryCollection.FailureConfigs failureConfigs = null) + { + PartitionKeyDefinition partitionKeyDefinition = new PartitionKeyDefinition() + { + Paths = new System.Collections.ObjectModel.Collection() + { + "/pk" + }, + Kind = PartitionKind.Hash, + Version = PartitionKeyDefinitionVersion.V2, + }; + + InMemoryCollection inMemoryCollection = new InMemoryCollection(partitionKeyDefinition, failureConfigs); + + inMemoryCollection.Split(partitionKeyRangeId: 0); + + inMemoryCollection.Split(partitionKeyRangeId: 1); + inMemoryCollection.Split(partitionKeyRangeId: 2); + + inMemoryCollection.Split(partitionKeyRangeId: 3); + inMemoryCollection.Split(partitionKeyRangeId: 4); + inMemoryCollection.Split(partitionKeyRangeId: 5); + inMemoryCollection.Split(partitionKeyRangeId: 6); + + for (int i = 0; i < numItems; i++) + { + // Insert an item + CosmosObject item = CosmosObject.Parse($"{{\"pk\" : {i} }}"); + inMemoryCollection.CreateItem(item); + } + + return inMemoryCollection; + } + } +} From 84e907b06f4c8c41250de175be926318ecd5d465 Mon Sep 17 00:00:00 2001 From: Brandon Chong Date: Wed, 8 Jul 2020 16:48:02 -0700 Subject: [PATCH 30/85] not putting back finished enumerators --- .../CrossPartitionRangePageEnumerator.cs | 21 +++++++++---------- ...arallelCrossPartitionQueryPipelineStage.cs | 1 - ...CollectionPartitionRangeEnumeratorTests.cs | 14 ++----------- 3 files changed, 12 insertions(+), 24 deletions(-) diff --git a/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionRangePageEnumerator.cs b/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionRangePageEnumerator.cs index 957feb7ffd..fb14c93e7a 100644 --- a/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionRangePageEnumerator.cs +++ b/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionRangePageEnumerator.cs @@ -75,13 +75,19 @@ public CrossPartitionRangePageEnumerator( public async ValueTask MoveNextAsync() { PriorityQueue> enumerators = await this.lazyEnumerators.GetValueAsync(cancellationToken: default); - PartitionRangePageEnumerator currentPaginator = enumerators.Dequeue(); - bool movedNext = await currentPaginator.MoveNextAsync(); - if (!movedNext) + if (enumerators.Count == 0) { return false; } + PartitionRangePageEnumerator currentPaginator = enumerators.Dequeue(); + if (!await currentPaginator.MoveNextAsync()) + { + // Current enumerator is empty, + // so recursively retry on the next enumerator. + return await this.MoveNextAsync(); + } + if (currentPaginator.Current.Failed) { // Check if it's a retryable exception. @@ -127,14 +133,7 @@ public async ValueTask MoveNextAsync() List<(FeedRangeInternal, TState)> feedRangeAndStates = new List<(FeedRangeInternal, TState)>(enumerators.Count); foreach (PartitionRangePageEnumerator enumerator in enumerators) { - // Converting to epk range just because some implementations prefer that. - // Eventually we should be range type agnostics, - // but we can't change the continution token format, - // since there are old SDKs that don't know how to resume from it. - FeedRangeEpk feedRangeEpk = await this.feedRangeProvider.ToEffectivePartitionKeyRangeAsync( - enumerator.Range, - cancellationToken: default); - feedRangeAndStates.Add((feedRangeEpk, enumerator.State)); + feedRangeAndStates.Add((enumerator.Range, enumerator.State)); } CrossPartitionState crossPartitionState = new CrossPartitionState(feedRangeAndStates); diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/ParallelCrossPartitionQueryPipelineStage.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/ParallelCrossPartitionQueryPipelineStage.cs index 870a3d1ece..377a825295 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/ParallelCrossPartitionQueryPipelineStage.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/ParallelCrossPartitionQueryPipelineStage.cs @@ -8,7 +8,6 @@ namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote using System; using System.Collections.Generic; using System.Linq; - using System.Threading; using System.Threading.Tasks; using Microsoft.Azure.Cosmos.CosmosElements; using Microsoft.Azure.Cosmos.Pagination; diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/CrossPartitionInMemoryCollectionPartitionRangeEnumeratorTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/CrossPartitionInMemoryCollectionPartitionRangeEnumeratorTests.cs index d90b900a55..5bf3255660 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/CrossPartitionInMemoryCollectionPartitionRangeEnumeratorTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/CrossPartitionInMemoryCollectionPartitionRangeEnumeratorTests.cs @@ -163,7 +163,7 @@ internal override IAsyncEnumerable createEnumerator(Cosmos.FeedRange range, InMemoryCollectionState state) => new InMemoryCollectionPartitionRangeEnumerator( inMemoryCollection, - partitionKeyRangeId: int.Parse(((FeedRangePartitionKeyRange)range).PartitionKeyRangeId), + partitionKeyRangeId: int.Parse(feedRangeProvider.ToPhysicalPartitionKeyRangeAsync((FeedRangeInternal)range, cancellationToken: default).Result.PartitionKeyRangeId), pageSize: 10, state: state); @@ -181,7 +181,7 @@ internal override IAsyncEnumerator createEnumerator(Cosmos.FeedRange range, InMemoryCollectionState state) => new InMemoryCollectionPartitionRangeEnumerator( inMemoryCollection, - partitionKeyRangeId: int.Parse(((FeedRangePartitionKeyRange)range).PartitionKeyRangeId), + partitionKeyRangeId: int.Parse(feedRangeProvider.ToPhysicalPartitionKeyRangeAsync((FeedRangeInternal)range, cancellationToken: default).Result.PartitionKeyRangeId), pageSize: 10, state: state); @@ -212,16 +212,6 @@ public int Compare( return 0; } - if (partitionRangePageEnumerator1.HasMoreResults && !partitionRangePageEnumerator2.HasMoreResults) - { - return -1; - } - - if (!partitionRangePageEnumerator1.HasMoreResults && partitionRangePageEnumerator2.HasMoreResults) - { - return 1; - } - // Either both don't have results or both do. return string.CompareOrdinal( ((FeedRangePartitionKeyRange)partitionRangePageEnumerator1.Range).PartitionKeyRangeId, From 51dc1298c2e854a0d746b983867f14b3d94c0888 Mon Sep 17 00:00:00 2001 From: Brandon Chong Date: Wed, 8 Jul 2020 17:50:20 -0700 Subject: [PATCH 31/85] need to stop using feedrange --- .../CrossPartitionRangePageEnumerable.cs | 1 + .../CrossPartitionRangePageEnumerator.cs | 15 +++++++- ...arallelCrossPartitionQueryPipelineStage.cs | 24 ++++++------ Microsoft.Azure.Cosmos/src/TaskHelper.cs | 1 + ...elCrossPartitionQueryPipelineStageTests.cs | 38 ++++++++++++++++++- 5 files changed, 65 insertions(+), 14 deletions(-) diff --git a/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionRangePageEnumerable.cs b/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionRangePageEnumerable.cs index 01f991397e..d71e67c3d4 100644 --- a/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionRangePageEnumerable.cs +++ b/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionRangePageEnumerable.cs @@ -38,6 +38,7 @@ public IAsyncEnumerator>> GetAsyncEnu this.feedRangeProvider, this.createPartitionRangeEnumerator, this.comparer, + forceEpkRange: false, this.state); } } diff --git a/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionRangePageEnumerator.cs b/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionRangePageEnumerator.cs index fb14c93e7a..6818cf1679 100644 --- a/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionRangePageEnumerator.cs +++ b/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionRangePageEnumerator.cs @@ -23,15 +23,18 @@ internal sealed class CrossPartitionRangePageEnumerator : IAsyncE private readonly IFeedRangeProvider feedRangeProvider; private readonly CreatePartitionRangePageEnumerator createPartitionRangeEnumerator; private readonly AsyncLazy>> lazyEnumerators; + private readonly bool forceEpkRange; public CrossPartitionRangePageEnumerator( IFeedRangeProvider feedRangeProvider, CreatePartitionRangePageEnumerator createPartitionRangeEnumerator, IComparer> comparer, + bool forceEpkRange = false, CrossPartitionState state = default) { this.feedRangeProvider = feedRangeProvider ?? throw new ArgumentNullException(nameof(feedRangeProvider)); this.createPartitionRangeEnumerator = createPartitionRangeEnumerator ?? throw new ArgumentNullException(nameof(createPartitionRangeEnumerator)); + this.forceEpkRange = forceEpkRange; if (comparer == null) { @@ -133,7 +136,17 @@ public async ValueTask MoveNextAsync() List<(FeedRangeInternal, TState)> feedRangeAndStates = new List<(FeedRangeInternal, TState)>(enumerators.Count); foreach (PartitionRangePageEnumerator enumerator in enumerators) { - feedRangeAndStates.Add((enumerator.Range, enumerator.State)); + FeedRangeInternal feedRangeInternal; + if (this.forceEpkRange) + { + feedRangeInternal = await this.feedRangeProvider.ToEffectivePartitionKeyRangeAsync(enumerator.Range, cancellationToken: default); + } + else + { + feedRangeInternal = enumerator.Range; + } + + feedRangeAndStates.Add((feedRangeInternal, enumerator.State)); } CrossPartitionState crossPartitionState = new CrossPartitionState(feedRangeAndStates); diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/ParallelCrossPartitionQueryPipelineStage.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/ParallelCrossPartitionQueryPipelineStage.cs index 377a825295..19816bdee1 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/ParallelCrossPartitionQueryPipelineStage.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/ParallelCrossPartitionQueryPipelineStage.cs @@ -4,7 +4,6 @@ namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote { -#nullable enable using System; using System.Collections.Generic; using System.Linq; @@ -19,7 +18,8 @@ internal sealed class ParallelCrossPartitionQueryPipelineStage : IQueryPipelineS { private readonly CrossPartitionRangePageEnumerator crossPartitionRangePageEnumerator; - private ParallelCrossPartitionQueryPipelineStage(CrossPartitionRangePageEnumerator crossPartitionRangePageEnumerator) + private ParallelCrossPartitionQueryPipelineStage( + CrossPartitionRangePageEnumerator crossPartitionRangePageEnumerator) { this.crossPartitionRangePageEnumerator = crossPartitionRangePageEnumerator ?? throw new ArgumentNullException(nameof(crossPartitionRangePageEnumerator)); } @@ -41,14 +41,11 @@ public TryCatch Current List compositeContinuationTokens = new List(crossPartitionState.Value.Count); foreach ((FeedRangeInternal range, QueryState state) in crossPartitionState.Value) { - if (!(range is FeedRangeEpk epkRange)) - { - throw new InvalidOperationException(); - } + FeedRangeEpk feedRangeEpk = (FeedRangeEpk)range; CompositeContinuationToken compositeContinuationToken = new CompositeContinuationToken() { - Range = epkRange.Range, + Range = feedRangeEpk.Range, Token = state != null ? ((CosmosString)state.Value).Value : null, }; @@ -82,14 +79,14 @@ public static TryCatch MonadicCreate( IQueryDataSource queryDataSource, SqlQuerySpec sqlQuerySpec, int pageSize, - CosmosElement? continuationToken) + CosmosElement continuationToken) { if (pageSize <= 0) { throw new ArgumentOutOfRangeException(nameof(pageSize)); } - CrossPartitionState? state; + CrossPartitionState state; if (continuationToken == null) { state = default; @@ -124,7 +121,9 @@ public static TryCatch MonadicCreate( } List<(FeedRangeInternal, QueryState)> rangesAndStates = compositeContinuationTokens - .Select(token => ((FeedRangeInternal)new FeedRangeEpk(token.Range), new QueryState(CosmosString.Create(token.Token)))) + .Select(token => ( + (FeedRangeInternal)new FeedRangeEpk(token.Range), + token.Token != null ? new QueryState(CosmosString.Create(token.Token)) : null)) .ToList(); state = new CrossPartitionState(rangesAndStates); @@ -134,6 +133,7 @@ public static TryCatch MonadicCreate( feedRangeProvider, ParallelCrossPartitionQueryPipelineStage.MakeCreateFunction(queryDataSource, sqlQuerySpec, pageSize), Comparer.Singleton, + forceEpkRange: true, state: state); ParallelCrossPartitionQueryPipelineStage stage = new ParallelCrossPartitionQueryPipelineStage(crossPartitionPageEnumerator); @@ -175,8 +175,8 @@ public int Compare( // Either both don't have results or both do. return string.CompareOrdinal( - ((FeedRangePartitionKeyRange)partitionRangePageEnumerator1.Range).PartitionKeyRangeId, - ((FeedRangePartitionKeyRange)partitionRangePageEnumerator2.Range).PartitionKeyRangeId); + ((FeedRangeEpk)partitionRangePageEnumerator1.Range).Range.Min, + ((FeedRangeEpk)partitionRangePageEnumerator2.Range).Range.Min); } } } diff --git a/Microsoft.Azure.Cosmos/src/TaskHelper.cs b/Microsoft.Azure.Cosmos/src/TaskHelper.cs index a82e8f091d..73709f8ef5 100644 --- a/Microsoft.Azure.Cosmos/src/TaskHelper.cs +++ b/Microsoft.Azure.Cosmos/src/TaskHelper.cs @@ -4,6 +4,7 @@ namespace Microsoft.Azure.Cosmos { using System; + using System.Collections.Generic; using System.Threading; using System.Threading.Tasks; using Microsoft.Azure.Documents; 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 8fbef7357a..3b5a3cab6d 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 @@ -14,7 +14,6 @@ namespace Microsoft.Azure.Cosmos.Tests.Query.Pipeline using Microsoft.Azure.Cosmos.Query.Core.Monads; using Microsoft.Azure.Cosmos.Query.Core.Pipeline; using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote; - using Microsoft.Azure.Cosmos.SqlObjects; using Microsoft.Azure.Cosmos.Tests.Pagination; using Microsoft.Azure.Documents; using Microsoft.VisualStudio.TestTools.UnitTesting; @@ -159,6 +158,43 @@ public async Task TestDrainFully_StartFromBeginingAsync() Assert.AreEqual(numItems, documents.Count); } + [TestMethod] + public async Task TestDrainFully_WithStateResume() + { + int numItems = 1000; + InMemoryCollection inMemoryCollection = CreateInMemoryCollection(numItems); + IFeedRangeProvider feedRangeProvider = new InMemoryCollectionFeedRangeProvider(inMemoryCollection); + IQueryDataSource queryDataSource = new InMemoryCollectionQueryDataSource(inMemoryCollection); + + TryCatch monadicCreate = ParallelCrossPartitionQueryPipelineStage.MonadicCreate( + feedRangeProvider: feedRangeProvider, + queryDataSource: queryDataSource, + sqlQuerySpec: new SqlQuerySpec("SELECT * FROM c"), + pageSize: 10, + continuationToken: default); + Assert.IsTrue(monadicCreate.Succeeded); + IQueryPipelineStage queryPipelineStage = monadicCreate.Result; + + List documents = new List(); + while (await queryPipelineStage.MoveNextAsync()) + { + TryCatch tryGetQueryPage = queryPipelineStage.Current; + Assert.IsTrue(tryGetQueryPage.Succeeded); + + QueryPage queryPage = tryGetQueryPage.Result; + documents.AddRange(queryPage.Documents); + + queryPipelineStage = ParallelCrossPartitionQueryPipelineStage.MonadicCreate( + feedRangeProvider: feedRangeProvider, + queryDataSource: queryDataSource, + sqlQuerySpec: new SqlQuerySpec("SELECT * FROM c"), + pageSize: 10, + continuationToken: queryPage.State.Value).Result; + } + + Assert.AreEqual(numItems, documents.Count); + } + private static InMemoryCollection CreateInMemoryCollection(int numItems, InMemoryCollection.FailureConfigs failureConfigs = null) { PartitionKeyDefinition partitionKeyDefinition = new PartitionKeyDefinition() From c62e1f30904dae703849d0229b9a6e696621ddab Mon Sep 17 00:00:00 2001 From: Brandon Chong Date: Thu, 9 Jul 2020 14:31:42 -0700 Subject: [PATCH 32/85] moved to pkrange instead of feed range --- .../CreatePartitionRangeEnumerator.cs | 6 ++- .../CrossPartitionRangePageEnumerator.cs | 23 +++++---- .../src/Pagination/CrossPartitionState.cs | 5 +- .../src/Pagination/FeedRangeProvider.cs | 48 +++++++------------ .../src/Pagination/IFeedRangeProvider.cs | 9 ++-- .../PartitionRangePageEnumerable.cs | 5 +- .../PartitionRangePageEnumerator.cs | 5 +- ...CollectionPartitionRangeEnumeratorTests.cs | 28 ++++++----- .../InMemoryCollectionFeedRangeProvider.cs | 45 ++++++++++------- ...emoryCollectionPartitionRangeEnumerator.cs | 16 ++++--- ...CollectionPartitionRangeEnumeratorTests.cs | 22 ++++----- 11 files changed, 111 insertions(+), 101 deletions(-) diff --git a/Microsoft.Azure.Cosmos/src/Pagination/CreatePartitionRangeEnumerator.cs b/Microsoft.Azure.Cosmos/src/Pagination/CreatePartitionRangeEnumerator.cs index 74d79fd5c5..43210b0bc8 100644 --- a/Microsoft.Azure.Cosmos/src/Pagination/CreatePartitionRangeEnumerator.cs +++ b/Microsoft.Azure.Cosmos/src/Pagination/CreatePartitionRangeEnumerator.cs @@ -4,7 +4,11 @@ namespace Microsoft.Azure.Cosmos.Pagination { - internal delegate PartitionRangePageEnumerator CreatePartitionRangePageEnumerator(FeedRange feedRange, TState state) + using Microsoft.Azure.Documents; + + internal delegate PartitionRangePageEnumerator CreatePartitionRangePageEnumerator( + PartitionKeyRange feedRange, + TState state) where TPage : Page where TState : State; } diff --git a/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionRangePageEnumerator.cs b/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionRangePageEnumerator.cs index eddeef7a35..ca6b1aa241 100644 --- a/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionRangePageEnumerator.cs +++ b/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionRangePageEnumerator.cs @@ -12,6 +12,7 @@ namespace Microsoft.Azure.Cosmos.Pagination using Microsoft.Azure.Cosmos.Query.Core; using Microsoft.Azure.Cosmos.Query.Core.Collections; using Microsoft.Azure.Cosmos.Query.Core.Monads; + using Microsoft.Azure.Documents; /// /// Coordinates draining pages from multiple , while maintaining a global sort order and handling repartitioning (splits, merge). @@ -40,7 +41,7 @@ public CrossPartitionRangePageEnumerator( this.lazyEnumerators = new AsyncLazy>>(async (CancellationToken token) => { - IReadOnlyList<(FeedRange, TState)> rangeAndStates; + IReadOnlyList<(PartitionKeyRange, TState)> rangeAndStates; if (state != default) { rangeAndStates = state.Value; @@ -48,10 +49,10 @@ public CrossPartitionRangePageEnumerator( else { // Fan out to all partitions with default state - IEnumerable ranges = await feedRangeProvider.GetFeedRangesAsync(token); + IEnumerable ranges = await feedRangeProvider.GetFeedRangesAsync(token); - List<(FeedRange, TState)> rangesAndStatesBuilder = new List<(FeedRange, TState)>(); - foreach (FeedRange range in ranges) + List<(PartitionKeyRange, TState)> rangesAndStatesBuilder = new List<(PartitionKeyRange, TState)>(); + foreach (PartitionKeyRange range in ranges) { rangesAndStatesBuilder.Add((range, default)); } @@ -60,7 +61,7 @@ public CrossPartitionRangePageEnumerator( } PriorityQueue> enumerators = new PriorityQueue>(comparer); - foreach ((FeedRange range, TState rangeState) in rangeAndStates) + foreach ((PartitionKeyRange range, TState rangeState) in rangeAndStates) { PartitionRangePageEnumerator enumerator = createPartitionRangeEnumerator(range, rangeState); enumerators.Enqueue(enumerator); @@ -94,10 +95,14 @@ public async ValueTask MoveNextAsync() if (IsSplitException(exception)) { // Handle split - IEnumerable childRanges = await this.feedRangeProvider.GetChildRangeAsync(currentPaginator.Range); - foreach (FeedRange childRange in childRanges) + IEnumerable childRanges = await this.feedRangeProvider.GetChildRangeAsync( + currentPaginator.Range, + cancellationToken: default); + foreach (PartitionKeyRange childRange in childRanges) { - PartitionRangePageEnumerator childPaginator = this.createPartitionRangeEnumerator(childRange, currentPaginator.State); + PartitionRangePageEnumerator childPaginator = this.createPartitionRangeEnumerator( + childRange, + currentPaginator.State); enumerators.Enqueue(childPaginator); } @@ -120,7 +125,7 @@ public async ValueTask MoveNextAsync() return true; } - List<(FeedRange, TState)> feedRangeAndStates = new List<(FeedRange, TState)>(enumerators.Count); + List<(PartitionKeyRange, TState)> feedRangeAndStates = new List<(PartitionKeyRange, TState)>(enumerators.Count); foreach (PartitionRangePageEnumerator enumerator in enumerators) { feedRangeAndStates.Add((enumerator.Range, enumerator.State)); diff --git a/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionState.cs b/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionState.cs index 338652a341..c244e8f5c5 100644 --- a/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionState.cs +++ b/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionState.cs @@ -6,15 +6,16 @@ namespace Microsoft.Azure.Cosmos.Pagination { using System; using System.Collections.Generic; + using Microsoft.Azure.Documents; internal sealed class CrossPartitionState : State where TState : State { - public CrossPartitionState(IReadOnlyList<(FeedRange, TState)> value) + public CrossPartitionState(IReadOnlyList<(PartitionKeyRange, TState)> value) { this.Value = value ?? throw new ArgumentNullException(nameof(value)); } - public IReadOnlyList<(FeedRange, TState)> Value { get; } + public IReadOnlyList<(PartitionKeyRange, TState)> Value { get; } } } diff --git a/Microsoft.Azure.Cosmos/src/Pagination/FeedRangeProvider.cs b/Microsoft.Azure.Cosmos/src/Pagination/FeedRangeProvider.cs index 4b0cc74b4a..391fbe92e5 100644 --- a/Microsoft.Azure.Cosmos/src/Pagination/FeedRangeProvider.cs +++ b/Microsoft.Azure.Cosmos/src/Pagination/FeedRangeProvider.cs @@ -9,9 +9,16 @@ namespace Microsoft.Azure.Cosmos.Pagination using System.Threading; using System.Threading.Tasks; using Microsoft.Azure.Cosmos.Query.Core.QueryClient; + using Microsoft.Azure.Documents; internal sealed class FeedRangeProvider : IFeedRangeProvider { + private static readonly PartitionKeyRange FullRange = new PartitionKeyRange() + { + MinInclusive = Documents.Routing.PartitionKeyInternal.MinimumInclusiveEffectivePartitionKey, + MaxExclusive = Documents.Routing.PartitionKeyInternal.MaximumExclusiveEffectivePartitionKey, + }; + private readonly CosmosQueryClient cosmosQueryClient; private readonly string collectionRid; @@ -21,40 +28,17 @@ public FeedRangeProvider(CosmosQueryClient cosmosQueryClient, string collectionR this.collectionRid = collectionRid ?? throw new ArgumentNullException(nameof(collectionRid)); } - public async Task> GetChildRangeAsync( - FeedRange feedRange, - CancellationToken cancellationToken = default) - { - cancellationToken.ThrowIfCancellationRequested(); - - if (feedRange == null) - { - throw new ArgumentNullException(nameof(feedRange)); - } - - if (!(feedRange is FeedRangeEpk feedRangeEpk)) - { - throw new ArgumentOutOfRangeException(nameof(feedRange)); - } - - IReadOnlyList replacementRanges = await this.cosmosQueryClient.TryGetOverlappingRangesAsync( + public Task> GetChildRangeAsync( + PartitionKeyRange partitionKeyRange, + CancellationToken cancellationToken) => this.cosmosQueryClient.TryGetOverlappingRangesAsync( this.collectionRid, - new Documents.Routing.Range(feedRangeEpk.Range.Min, feedRangeEpk.Range.Max, isMaxInclusive: true, isMinInclusive: false), - forceRefresh: true); - - List childFeedRanges = new List(replacementRanges.Count); - - foreach (Documents.PartitionKeyRange replacementRange in replacementRanges) - { - childFeedRanges.Add(new FeedRangeEpk(replacementRange.ToRange())); - } - - return childFeedRanges; - } + partitionKeyRange.ToRange(), + forceRefresh: true) + .ContinueWith((task) => (IEnumerable)task.Result); - public Task> GetFeedRangesAsync( - CancellationToken cancellationToken = default) => this.GetChildRangeAsync( - FeedRangeEpk.FullRange, + public Task> GetFeedRangesAsync( + CancellationToken cancellationToken) => this.GetChildRangeAsync( + FeedRangeProvider.FullRange, cancellationToken); } } diff --git a/Microsoft.Azure.Cosmos/src/Pagination/IFeedRangeProvider.cs b/Microsoft.Azure.Cosmos/src/Pagination/IFeedRangeProvider.cs index 5a5cf416ce..301a03c57d 100644 --- a/Microsoft.Azure.Cosmos/src/Pagination/IFeedRangeProvider.cs +++ b/Microsoft.Azure.Cosmos/src/Pagination/IFeedRangeProvider.cs @@ -7,13 +7,14 @@ namespace Microsoft.Azure.Cosmos.Pagination using System.Collections.Generic; using System.Threading; using System.Threading.Tasks; + using Microsoft.Azure.Documents; internal interface IFeedRangeProvider { - public Task> GetChildRangeAsync( - FeedRange feedRange, - CancellationToken cancellationToken = default); + public Task> GetChildRangeAsync( + PartitionKeyRange partitionKeyRange, + CancellationToken cancellationToken); - public Task> GetFeedRangesAsync(CancellationToken cancellationToken = default); + public Task> GetFeedRangesAsync(CancellationToken cancellationToken); } } diff --git a/Microsoft.Azure.Cosmos/src/Pagination/PartitionRangePageEnumerable.cs b/Microsoft.Azure.Cosmos/src/Pagination/PartitionRangePageEnumerable.cs index f9ca6d40d9..295bc8af6c 100644 --- a/Microsoft.Azure.Cosmos/src/Pagination/PartitionRangePageEnumerable.cs +++ b/Microsoft.Azure.Cosmos/src/Pagination/PartitionRangePageEnumerable.cs @@ -8,17 +8,18 @@ namespace Microsoft.Azure.Cosmos.Pagination using System.Collections.Generic; using System.Threading; using Microsoft.Azure.Cosmos.Query.Core.Monads; + using Microsoft.Azure.Documents; internal sealed class PartitionRangePageEnumerable : IAsyncEnumerable> where TPage : Page where TState : State { - private readonly FeedRange range; + private readonly PartitionKeyRange range; private readonly TState state; private readonly CreatePartitionRangePageEnumerator createPartitionRangeEnumerator; public PartitionRangePageEnumerable( - FeedRange range, + PartitionKeyRange range, TState state, CreatePartitionRangePageEnumerator createPartitionRangeEnumerator) { diff --git a/Microsoft.Azure.Cosmos/src/Pagination/PartitionRangePageEnumerator.cs b/Microsoft.Azure.Cosmos/src/Pagination/PartitionRangePageEnumerator.cs index d7e22984d9..81bf67da22 100644 --- a/Microsoft.Azure.Cosmos/src/Pagination/PartitionRangePageEnumerator.cs +++ b/Microsoft.Azure.Cosmos/src/Pagination/PartitionRangePageEnumerator.cs @@ -8,6 +8,7 @@ namespace Microsoft.Azure.Cosmos.Pagination using System.Threading; using System.Threading.Tasks; using Microsoft.Azure.Cosmos.Query.Core.Monads; + using Microsoft.Azure.Documents; /// /// Has the ability to page through a partition range. @@ -18,13 +19,13 @@ internal abstract class PartitionRangePageEnumerator : IAsyncEnum { private bool hasStarted; - protected PartitionRangePageEnumerator(FeedRange range, TState state = null) + protected PartitionRangePageEnumerator(PartitionKeyRange range, TState state = null) { this.Range = range; this.State = state; } - public FeedRange Range { get; } + public PartitionKeyRange Range { get; } public TryCatch Current { get; private set; } diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/CrossPartitionInMemoryCollectionPartitionRangeEnumeratorTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/CrossPartitionInMemoryCollectionPartitionRangeEnumeratorTests.cs index d90b900a55..a10eac6656 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/CrossPartitionInMemoryCollectionPartitionRangeEnumeratorTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/CrossPartitionInMemoryCollectionPartitionRangeEnumeratorTests.cs @@ -161,11 +161,13 @@ internal override IAsyncEnumerable state = null) { IFeedRangeProvider feedRangeProvider = new InMemoryCollectionFeedRangeProvider(inMemoryCollection); - PartitionRangePageEnumerator createEnumerator(Cosmos.FeedRange range, InMemoryCollectionState state) => new InMemoryCollectionPartitionRangeEnumerator( - inMemoryCollection, - partitionKeyRangeId: int.Parse(((FeedRangePartitionKeyRange)range).PartitionKeyRangeId), - pageSize: 10, - state: state); + PartitionRangePageEnumerator createEnumerator( + PartitionKeyRange range, + InMemoryCollectionState state) => new InMemoryCollectionPartitionRangeEnumerator( + inMemoryCollection, + partitionKeyRangeId: int.Parse(range.Id), + pageSize: 10, + state: state); return new CrossPartitionRangePageEnumerable( feedRangeProvider: feedRangeProvider, @@ -179,11 +181,13 @@ internal override IAsyncEnumerator state = null) { IFeedRangeProvider feedRangeProvider = new InMemoryCollectionFeedRangeProvider(inMemoryCollection); - PartitionRangePageEnumerator createEnumerator(Cosmos.FeedRange range, InMemoryCollectionState state) => new InMemoryCollectionPartitionRangeEnumerator( - inMemoryCollection, - partitionKeyRangeId: int.Parse(((FeedRangePartitionKeyRange)range).PartitionKeyRangeId), - pageSize: 10, - state: state); + PartitionRangePageEnumerator createEnumerator( + PartitionKeyRange range, + InMemoryCollectionState state) => new InMemoryCollectionPartitionRangeEnumerator( + inMemoryCollection, + partitionKeyRangeId: int.Parse(range.Id), + pageSize: 10, + state: state); CrossPartitionRangePageEnumerator enumerator = new CrossPartitionRangePageEnumerator( feedRangeProvider: feedRangeProvider, @@ -224,8 +228,8 @@ public int Compare( // Either both don't have results or both do. return string.CompareOrdinal( - ((FeedRangePartitionKeyRange)partitionRangePageEnumerator1.Range).PartitionKeyRangeId, - ((FeedRangePartitionKeyRange)partitionRangePageEnumerator2.Range).PartitionKeyRangeId); + partitionRangePageEnumerator1.Range.MinInclusive, + partitionRangePageEnumerator2.Range.MinInclusive); } } } diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/InMemoryCollectionFeedRangeProvider.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/InMemoryCollectionFeedRangeProvider.cs index c7134016db..c4135de738 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/InMemoryCollectionFeedRangeProvider.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/InMemoryCollectionFeedRangeProvider.cs @@ -4,8 +4,8 @@ using System.Collections.Generic; using System.Threading; using System.Threading.Tasks; - using Microsoft.Azure.Cosmos; using Microsoft.Azure.Cosmos.Pagination; + using Microsoft.Azure.Documents; internal sealed class InMemoryCollectionFeedRangeProvider : IFeedRangeProvider { @@ -16,40 +16,51 @@ public InMemoryCollectionFeedRangeProvider(InMemoryCollection inMemoryCollection this.inMemoryCollection = inMemoryCollection ?? throw new ArgumentNullException(nameof(inMemoryCollection)); } - public Task> GetChildRangeAsync( - FeedRange feedRange, - CancellationToken cancellationToken = default) + public Task> GetChildRangeAsync( + PartitionKeyRange feedRange, + CancellationToken cancellationToken) { cancellationToken.ThrowIfCancellationRequested(); - if (!(feedRange is FeedRangePartitionKeyRange feedRangePartitionKeyRange)) - { - throw new ArgumentOutOfRangeException(nameof(feedRange)); - } - - int partitionKeyRangeId = int.Parse(feedRangePartitionKeyRange.PartitionKeyRangeId); + int partitionKeyRangeId = int.Parse(feedRange.Id); (int leftChild, int rightChild) = this.inMemoryCollection.GetChildRanges(partitionKeyRangeId); return Task.FromResult( - (IEnumerable)new List() + (IEnumerable)new List() { - new FeedRangePartitionKeyRange(leftChild.ToString()), - new FeedRangePartitionKeyRange(rightChild.ToString()), + new PartitionKeyRange() + { + Id = leftChild.ToString(), + MinInclusive = leftChild.ToString(), + MaxExclusive = leftChild.ToString(), + }, + new PartitionKeyRange() + { + Id = rightChild.ToString(), + MinInclusive = rightChild.ToString(), + MaxExclusive = rightChild.ToString(), + } }); } - public Task> GetFeedRangesAsync(CancellationToken cancellationToken = default) + public Task> GetFeedRangesAsync(CancellationToken cancellationToken) { cancellationToken.ThrowIfCancellationRequested(); - List ranges = new List(); + List ranges = new List(); foreach (int partitionKeyRangeId in this.inMemoryCollection.PartitionKeyRangeFeedReed().Keys) { - FeedRange range = new FeedRangePartitionKeyRange(partitionKeyRangeId.ToString()); + PartitionKeyRange range = new PartitionKeyRange() + { + Id = partitionKeyRangeId.ToString(), + MinInclusive = partitionKeyRangeId.ToString(), + MaxExclusive = partitionKeyRangeId.ToString(), + }; + ranges.Add(range); } - return Task.FromResult((IEnumerable)ranges); + return Task.FromResult((IEnumerable)ranges); } } } diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/InMemoryCollectionPartitionRangeEnumerator.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/InMemoryCollectionPartitionRangeEnumerator.cs index de09150abd..a24d3b7c7c 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/InMemoryCollectionPartitionRangeEnumerator.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/InMemoryCollectionPartitionRangeEnumerator.cs @@ -10,6 +10,7 @@ namespace Microsoft.Azure.Cosmos.Tests.Pagination using System.Threading.Tasks; using Microsoft.Azure.Cosmos.Pagination; using Microsoft.Azure.Cosmos.Query.Core.Monads; + using Microsoft.Azure.Documents; internal sealed class InMemoryCollectionPartitionRangeEnumerator : PartitionRangePageEnumerator { @@ -22,18 +23,21 @@ public InMemoryCollectionPartitionRangeEnumerator( int partitionKeyRangeId, int pageSize, InMemoryCollectionState state = null) - : base(new FeedRangePartitionKeyRange(partitionKeyRangeId.ToString()), state ?? new InMemoryCollectionState(resourceIdentifier: 0)) + : base( + new PartitionKeyRange() + { + Id = partitionKeyRangeId.ToString(), + MinInclusive = partitionKeyRangeId.ToString(), + MaxExclusive = partitionKeyRangeId.ToString() + }, + state ?? new InMemoryCollectionState(resourceIdentifier: 0)) { this.inMemoryCollection = inMemoryCollection ?? throw new ArgumentNullException(nameof(inMemoryCollection)); this.partitionKeyRangeId = partitionKeyRangeId; this.pageSize = pageSize; } - public override ValueTask DisposeAsync() - { - // Do Nothing - return default; - } + public override ValueTask DisposeAsync() => default; public override Task> GetNextPageAsync(CancellationToken cancellationToken = default) { diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/SinglePartitionInMemoryCollectionPartitionRangeEnumeratorTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/SinglePartitionInMemoryCollectionPartitionRangeEnumeratorTests.cs index c803c4eff6..af4c3d5795 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/SinglePartitionInMemoryCollectionPartitionRangeEnumeratorTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/SinglePartitionInMemoryCollectionPartitionRangeEnumeratorTests.cs @@ -85,11 +85,11 @@ public async Task TestSplitAsync() foreach (int partitionKeyRangeId in new int[] { 1, 2 }) { PartitionRangePageEnumerable enumerable = new PartitionRangePageEnumerable( - range: new FeedRangePartitionKeyRange(partitionKeyRangeId.ToString()), - state: state, - (range, state) => new InMemoryCollectionPartitionRangeEnumerator( + range: new PartitionKeyRange() { Id = partitionKeyRangeId.ToString() }, + state: state, + (range, state) => new InMemoryCollectionPartitionRangeEnumerator( inMemoryCollection, - partitionKeyRangeId: int.Parse(((FeedRangePartitionKeyRange)range).PartitionKeyRangeId), + partitionKeyRangeId: int.Parse(range.Id), pageSize: 10, state: state)); HashSet resourceIdentifiers = await this.DrainFullyAsync(enumerable); @@ -131,28 +131,22 @@ internal override InMemoryCollection CreateInMemoryCollection(int numItems, InMe internal override IAsyncEnumerable> CreateEnumerable( InMemoryCollection inMemoryCollection, - InMemoryCollectionState state = null) - { - return new PartitionRangePageEnumerable( - range: new FeedRangePartitionKeyRange("0"), + InMemoryCollectionState state = null) => new PartitionRangePageEnumerable( + range: new PartitionKeyRange() { Id = "0" }, state: state, (range, state) => new InMemoryCollectionPartitionRangeEnumerator( inMemoryCollection, - partitionKeyRangeId: int.Parse(((FeedRangePartitionKeyRange)range).PartitionKeyRangeId), + partitionKeyRangeId: int.Parse(range.Id), pageSize: 10, state: state)); - } internal override IAsyncEnumerator> CreateEnumerator( InMemoryCollection inMemoryCollection, - InMemoryCollectionState state = null) - { - return new InMemoryCollectionPartitionRangeEnumerator( + InMemoryCollectionState state = null) => new InMemoryCollectionPartitionRangeEnumerator( inMemoryCollection, partitionKeyRangeId: 0, pageSize: 10, state: state); - } } } } From 797a761b29f75eeb942f143ab0ac0ba92e67692b Mon Sep 17 00:00:00 2001 From: Brandon Chong Date: Thu, 9 Jul 2020 15:31:01 -0700 Subject: [PATCH 33/85] fixed baseline --- .../src/Microsoft.Azure.Cosmos.csproj | 6 +++--- .../DirectContractTests.cs | 12 ++++++++++-- 2 files changed, 13 insertions(+), 5 deletions(-) diff --git a/Microsoft.Azure.Cosmos/src/Microsoft.Azure.Cosmos.csproj b/Microsoft.Azure.Cosmos/src/Microsoft.Azure.Cosmos.csproj index 0b28f17738..fbbfc74f70 100644 --- a/Microsoft.Azure.Cosmos/src/Microsoft.Azure.Cosmos.csproj +++ b/Microsoft.Azure.Cosmos/src/Microsoft.Azure.Cosmos.csproj @@ -86,15 +86,15 @@ - + - - + + diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/DirectContractTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/DirectContractTests.cs index 8e5b2246c1..96a915b399 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/DirectContractTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/DirectContractTests.cs @@ -96,7 +96,16 @@ public void ProjectPackageDependenciesTest() string csprojFile = "Microsoft.Azure.Cosmos.csproj"; Dictionary projDependencies = DirectContractTests.GetPackageReferencies(csprojFile); Dictionary baselineDependencies = JsonConvert.DeserializeObject< Dictionary>( - "{\"System.Numerics.Vectors\":\"4.5.0\",\"Newtonsoft.Json\":\"10.0.2\",\"System.Configuration.ConfigurationManager\":\"4.5.0\",\"System.Memory\":\"4.5.1\",\"System.Runtime.CompilerServices.Unsafe\":\"4.5.1\",\"System.Threading.Tasks.Extensions\":\"4.5.1\",\"System.ValueTuple\":\"4.5.0\"}"); + "{" + + "\"System.Numerics.Vectors\":\"4.5.0\"," + + "\"Newtonsoft.Json\":\"10.0.2\"," + + "\"Microsoft.Bcl.AsyncInterfaces\":\"1.0.0\"," + + "\"System.Configuration.ConfigurationManager\":\"4.5.0\"," + + "\"System.Memory\":\"4.5.1\"," + + "\"System.Runtime.CompilerServices.Unsafe\":\"4.5.2\"," + + "\"System.Threading.Tasks.Extensions\":\"4.5.2\"," + + "\"System.ValueTuple\":\"4.5.0\"" + + "}"); Assert.AreEqual(projDependencies.Count, baselineDependencies.Count); foreach(KeyValuePair projectDependency in projDependencies) @@ -104,7 +113,6 @@ public void ProjectPackageDependenciesTest() string baselineVersion = baselineDependencies[projectDependency.Key]; Assert.AreEqual(baselineVersion, projectDependency.Value); } - } [TestMethod] From bb855001689260b34fa698a9e38176108223c131 Mon Sep 17 00:00:00 2001 From: Brandon Chong Date: Fri, 10 Jul 2020 01:41:25 -0700 Subject: [PATCH 34/85] got basic cases working --- .../src/Routing/PartitionKeyHash.cs | 11 ++ .../InMemoryCollection.cs | 47 ++++++- .../InMemoryCollectionFeedRangeProvider.cs | 117 ++---------------- .../InMemoryCollectionQueryDataSource.cs | 3 +- 4 files changed, 62 insertions(+), 116 deletions(-) diff --git a/Microsoft.Azure.Cosmos/src/Routing/PartitionKeyHash.cs b/Microsoft.Azure.Cosmos/src/Routing/PartitionKeyHash.cs index fe35c9e832..42ef970915 100644 --- a/Microsoft.Azure.Cosmos/src/Routing/PartitionKeyHash.cs +++ b/Microsoft.Azure.Cosmos/src/Routing/PartitionKeyHash.cs @@ -5,6 +5,7 @@ namespace Microsoft.Azure.Cosmos.Routing { using System; + using System.Runtime.CompilerServices; using System.Runtime.InteropServices; using System.Text; using Microsoft.Azure.Documents.Routing; @@ -79,6 +80,16 @@ public static bool TryParse(string value, out PartitionKeyHash parsedValue) return true; } + public static PartitionKeyHash Parse(string value) + { + if (!PartitionKeyHash.TryParse(value, out PartitionKeyHash parsedValue)) + { + throw new FormatException(); + } + + return parsedValue; + } + public static class V1 { private const int MaxStringLength = 100; diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/InMemoryCollection.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/InMemoryCollection.cs index eb8f28bdbc..c4067ad339 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/InMemoryCollection.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/InMemoryCollection.cs @@ -8,6 +8,9 @@ namespace Microsoft.Azure.Cosmos.Tests using System.Collections; using System.Collections.Generic; using System.Linq; + using System.Threading; + using System.Threading.Tasks; + using Microsoft.Azure.Cosmos; using Microsoft.Azure.Cosmos.CosmosElements; using Microsoft.Azure.Cosmos.Query.Core; using Microsoft.Azure.Cosmos.Query.Core.Monads; @@ -135,7 +138,7 @@ record = default; public TryCatch<(List, long?)> Query( SqlQuerySpec sqlQuerySpec, - int partitionKeyRangeId, + FeedRangeInternal feedRange, long resourceIdentifier, int pageSize) { @@ -149,7 +152,41 @@ record = default; throw new NotSupportedException("InMemoryCollection only supports SELECT * FROM c queries"); } - // for now just always do a "SELECT * FROM c" query + int partitionKeyRangeId; + if (feedRange is FeedRangePartitionKeyRange feedRangePartitionKeyRange) + { + partitionKeyRangeId = int.Parse(feedRangePartitionKeyRange.PartitionKeyRangeId); + } + else if (feedRange is FeedRangeEpk feedRangeEpk) + { + PartitionKeyHash? startInclusive = (feedRangeEpk.Range.Min != string.Empty) ? PartitionKeyHash.Parse(feedRangeEpk.Range.Min) : (PartitionKeyHash?)null; + PartitionKeyHash? endExclusive = (feedRangeEpk.Range.Max != string.Empty) ? PartitionKeyHash.Parse(feedRangeEpk.Range.Max) : (PartitionKeyHash?)null; + + PartitionKeyHashRange partitionKeyHashRange = new PartitionKeyHashRange(startInclusive, endExclusive); + IEnumerable pkrangeId = this.partitionKeyRangeIdToHashRange + .Where(kvp => kvp.Value.Equals(partitionKeyHashRange)) + .Select(kvp => kvp.Key); + if (pkrangeId.Count() != 1) + { + // Epk range no longer maps to a physical partition + // so return a gone exception + return TryCatch<(List, long?)>.FromException( + new CosmosException( + message: $"Epk Range {feedRangeEpk.Range} is gone", + statusCode: System.Net.HttpStatusCode.Gone, + subStatusCode: (int)SubStatusCodes.PartitionKeyRangeGone, + activityId: Guid.NewGuid().ToString(), + requestCharge: default)); + } + + partitionKeyRangeId = pkrangeId.First(); + } + else + { + throw new NotImplementedException(); + } + + // For now just do a read feed return this.ReadFeed(partitionKeyRangeId, resourceIdentifier, pageSize); } @@ -227,8 +264,6 @@ public void Split(int partitionKeyRangeId) public PartitionKeyHashRange GetHashRange(int partitionKeyRangeId) => this.partitionKeyRangeIdToHashRange[partitionKeyRangeId]; - public int GetPartitionKeyRangeId(PartitionKeyHashRange hashRange) => this.partitionKeyRangeIdToHashRange.Where(kvp => kvp.Value.Equals(hashRange)).First().Key; - private bool Return429() => (this.failureConfigs != null) && this.failureConfigs.Inject429s && ((this.random.Next() % 2) == 0); private bool ReturnEmptyPage() => (this.failureConfigs != null) && this.failureConfigs.InjectEmptyPages && ((this.random.Next() % 2) == 0); @@ -247,7 +282,7 @@ private static CosmosElement GetPartitionKeyFromPayload(CosmosObject payload, Pa throw new ArgumentOutOfRangeException("Can only support hash partitioning"); } - if (partitionKeyDefinition.Version != PartitionKeyDefinitionVersion.V2) + if (partitionKeyDefinition.Version != Documents.PartitionKeyDefinitionVersion.V2) { throw new ArgumentOutOfRangeException("Can only support hash v2"); } @@ -282,7 +317,7 @@ private static PartitionKeyHash GetHashFromPartitionKey(CosmosElement partitionK throw new ArgumentOutOfRangeException("Can only support hash partitioning"); } - if (partitionKeyDefinition.Version != PartitionKeyDefinitionVersion.V2) + if (partitionKeyDefinition.Version != Documents.PartitionKeyDefinitionVersion.V2) { throw new ArgumentOutOfRangeException("Can only support hash v2"); } diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/InMemoryCollectionFeedRangeProvider.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/InMemoryCollectionFeedRangeProvider.cs index 6251673385..f568a2cc2b 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/InMemoryCollectionFeedRangeProvider.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/InMemoryCollectionFeedRangeProvider.cs @@ -11,14 +11,10 @@ internal sealed class InMemoryCollectionFeedRangeProvider : IFeedRangeProvider { private readonly InMemoryCollection inMemoryCollection; - private readonly FeedRangeToEffectivePartitionKeyRangeVisitor feedRangeToEffectivePartitionKeyRangeVisitor; - private readonly FeedRangeToPhysicalPartitionKeyRange feedRangeToPhysicalPartitionKeyRange; public InMemoryCollectionFeedRangeProvider(InMemoryCollection inMemoryCollection) { this.inMemoryCollection = inMemoryCollection ?? throw new ArgumentNullException(nameof(inMemoryCollection)); - this.feedRangeToEffectivePartitionKeyRangeVisitor = new FeedRangeToEffectivePartitionKeyRangeVisitor(this.inMemoryCollection); - this.feedRangeToPhysicalPartitionKeyRange = new FeedRangeToPhysicalPartitionKeyRange(this.inMemoryCollection); } public Task> GetChildRangeAsync( @@ -29,6 +25,8 @@ public Task> GetChildRangeAsync( int partitionKeyRangeId = int.Parse(feedRange.Id); (int leftChild, int rightChild) = this.inMemoryCollection.GetChildRanges(partitionKeyRangeId); + PartitionKeyHashRange leftChildHashRange = this.inMemoryCollection.GetHashRange(leftChild); + PartitionKeyHashRange rightChildHashRange = this.inMemoryCollection.GetHashRange(rightChild); return Task.FromResult( (IEnumerable)new List() @@ -36,14 +34,14 @@ public Task> GetChildRangeAsync( new PartitionKeyRange() { Id = leftChild.ToString(), - MinInclusive = leftChild.ToString(), - MaxExclusive = leftChild.ToString(), + MinInclusive = leftChildHashRange.StartInclusive.HasValue ? leftChildHashRange.StartInclusive.Value.ToString() : string.Empty, + MaxExclusive = leftChildHashRange.EndExclusive.HasValue ? leftChildHashRange.EndExclusive.Value.ToString() : string.Empty, }, new PartitionKeyRange() { Id = rightChild.ToString(), - MinInclusive = rightChild.ToString(), - MaxExclusive = rightChild.ToString(), + MinInclusive = rightChildHashRange.StartInclusive.HasValue ? rightChildHashRange.StartInclusive.Value.ToString() : string.Empty, + MaxExclusive = rightChildHashRange.EndExclusive.HasValue ? rightChildHashRange.EndExclusive.Value.ToString() : string.Empty, } }); } @@ -55,11 +53,12 @@ public Task> GetFeedRangesAsync(CancellationToken List ranges = new List(); foreach (int partitionKeyRangeId in this.inMemoryCollection.PartitionKeyRangeFeedReed().Keys) { + PartitionKeyHashRange hashRange = this.inMemoryCollection.GetHashRange(partitionKeyRangeId); PartitionKeyRange range = new PartitionKeyRange() { Id = partitionKeyRangeId.ToString(), - MinInclusive = partitionKeyRangeId.ToString(), - MaxExclusive = partitionKeyRangeId.ToString(), + MinInclusive = hashRange.StartInclusive.HasValue ? hashRange.StartInclusive.Value.ToString() : string.Empty, + MaxExclusive = hashRange.EndExclusive.HasValue ? hashRange.EndExclusive.Value.ToString() : string.Empty, }; ranges.Add(range); @@ -67,103 +66,5 @@ public Task> GetFeedRangesAsync(CancellationToken return Task.FromResult((IEnumerable)ranges); } - - public Task ToEffectivePartitionKeyRangeAsync( - FeedRangeInternal feedRange, - CancellationToken cancellationToken) => feedRange.AcceptAsync(this.feedRangeToEffectivePartitionKeyRangeVisitor); - - public Task ToPhysicalPartitionKeyRangeAsync( - FeedRangeInternal feedRange, - CancellationToken cancellationToken) => feedRange.AcceptAsync(this.feedRangeToPhysicalPartitionKeyRange); - - private sealed class FeedRangeToEffectivePartitionKeyRangeVisitor : IFeedRangeAsyncVisitor - { - private readonly InMemoryCollection inMemoryCollection; - - public FeedRangeToEffectivePartitionKeyRangeVisitor(InMemoryCollection inMemoryCollection) - { - this.inMemoryCollection = inMemoryCollection ?? throw new ArgumentNullException(nameof(inMemoryCollection)); - } - - public Task VisitAsync(FeedRangePartitionKey feedRange, CancellationToken cancellationToken = default) - { - throw new NotImplementedException(); - } - - public Task VisitAsync(FeedRangePartitionKeyRange feedRange, CancellationToken cancellationToken = default) - { - PartitionKeyHashRange hashRange = this.inMemoryCollection.GetHashRange(int.Parse(feedRange.PartitionKeyRangeId)); - Documents.Routing.Range range = new Documents.Routing.Range( - min: hashRange.StartInclusive.HasValue ? hashRange.StartInclusive.Value.ToString() : string.Empty, - max: hashRange.EndExclusive.HasValue ? hashRange.EndExclusive.Value.ToString() : string.Empty, - isMinInclusive: true, - isMaxInclusive: false); - - return Task.FromResult(new FeedRangeEpk(range)); - } - - public Task VisitAsync(FeedRangeEpk feedRange, CancellationToken cancellationToken = default) - { - return Task.FromResult(feedRange); - } - } - - private sealed class FeedRangeToPhysicalPartitionKeyRange : IFeedRangeAsyncVisitor - { - private readonly InMemoryCollection inMemoryCollection; - - public FeedRangeToPhysicalPartitionKeyRange(InMemoryCollection inMemoryCollection) - { - this.inMemoryCollection = inMemoryCollection ?? throw new ArgumentNullException(nameof(inMemoryCollection)); - } - - public Task VisitAsync(FeedRangePartitionKey feedRange, CancellationToken cancellationToken = default) - { - throw new NotImplementedException(); - } - - public Task VisitAsync(FeedRangePartitionKeyRange feedRange, CancellationToken cancellationToken = default) - { - return Task.FromResult(feedRange); - } - - public Task VisitAsync(FeedRangeEpk feedRange, CancellationToken cancellationToken = default) - { - Documents.Routing.Range range = feedRange.Range; - - PartitionKeyHash? startInclusive; - if (range.Min == string.Empty) - { - startInclusive = default; - } - else - { - if (!PartitionKeyHash.TryParse(range.Min, out PartitionKeyHash parsedValue)) - { - throw new InvalidOperationException($"Failed to parse: {range.Min}."); - } - - startInclusive = parsedValue; - } - - PartitionKeyHash? endExclusive; - if (range.Max == string.Empty) - { - endExclusive = default; - } - else - { - if (!PartitionKeyHash.TryParse(range.Max, out PartitionKeyHash parsedValue)) - { - throw new InvalidOperationException($"Failed to parse: {range.Max}."); - } - - endExclusive = parsedValue; - } - - int pkRangeId = this.inMemoryCollection.GetPartitionKeyRangeId(new PartitionKeyHashRange(startInclusive, endExclusive)); - return Task.FromResult(new FeedRangePartitionKeyRange(pkRangeId.ToString())); - } - } } } diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/InMemoryCollectionQueryDataSource.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/InMemoryCollectionQueryDataSource.cs index 0a19672fed..9cb9a2b45d 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/InMemoryCollectionQueryDataSource.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/InMemoryCollectionQueryDataSource.cs @@ -8,6 +8,7 @@ namespace Microsoft.Azure.Cosmos.Tests.Query using System.Collections.Generic; using System.Threading; using System.Threading.Tasks; + using Microsoft.Azure.Cosmos; using Microsoft.Azure.Cosmos.CosmosElements; using Microsoft.Azure.Cosmos.CosmosElements.Numbers; using Microsoft.Azure.Cosmos.Query.Core; @@ -75,7 +76,5 @@ private readonly struct VisitorArguments { } - - private sealed class ExecuteQueryBasedOnFeedRangeVisitor : IFeedRangeAsyncVisitor, VisitorArguments> } } From 27afdbc92c2b67c2121c8d59612daae233e46072 Mon Sep 17 00:00:00 2001 From: Brandon Chong Date: Fri, 10 Jul 2020 01:55:41 -0700 Subject: [PATCH 35/85] got continuation token support working --- .../CrossPartitionRangePageEnumerator.cs | 22 +++++++++--- ...arallelCrossPartitionQueryPipelineStage.cs | 36 ++++++++++++------- ...elCrossPartitionQueryPipelineStageTests.cs | 5 ++- 3 files changed, 44 insertions(+), 19 deletions(-) diff --git a/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionRangePageEnumerator.cs b/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionRangePageEnumerator.cs index 639ba1d71d..55d22482bf 100644 --- a/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionRangePageEnumerator.cs +++ b/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionRangePageEnumerator.cs @@ -122,7 +122,10 @@ public async ValueTask MoveNextAsync() } } - enumerators.Enqueue(currentPaginator); + if (currentPaginator.State != null) + { + enumerators.Enqueue(currentPaginator); + } TryCatch backendPage = currentPaginator.Current; if (backendPage.Failed) @@ -131,13 +134,22 @@ public async ValueTask MoveNextAsync() return true; } - List<(PartitionKeyRange, TState)> feedRangeAndStates = new List<(PartitionKeyRange, TState)>(enumerators.Count); - foreach (PartitionRangePageEnumerator enumerator in enumerators) + CrossPartitionState crossPartitionState; + if (enumerators.Count == 0) + { + crossPartitionState = null; + } + else { - feedRangeAndStates.Add((enumerator.Range, enumerator.State)); + List<(PartitionKeyRange, TState)> feedRangeAndStates = new List<(PartitionKeyRange, TState)>(enumerators.Count); + foreach (PartitionRangePageEnumerator enumerator in enumerators) + { + feedRangeAndStates.Add((enumerator.Range, enumerator.State)); + } + + crossPartitionState = new CrossPartitionState(feedRangeAndStates); } - CrossPartitionState crossPartitionState = new CrossPartitionState(feedRangeAndStates); this.Current = TryCatch>.FromResult( new CrossPartitionPage(backendPage.Result, crossPartitionState)); return true; diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/ParallelCrossPartitionQueryPipelineStage.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/ParallelCrossPartitionQueryPipelineStage.cs index ed9c555ad7..b7c4cdc007 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/ParallelCrossPartitionQueryPipelineStage.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/ParallelCrossPartitionQueryPipelineStage.cs @@ -39,22 +39,32 @@ public TryCatch Current QueryPage backendQueryPage = crossPartitionPageResult.Page; CrossPartitionState crossPartitionState = crossPartitionPageResult.State; - List compositeContinuationTokens = new List(crossPartitionState.Value.Count); - foreach ((PartitionKeyRange range, QueryState state) in crossPartitionState.Value) + QueryState queryState; + if (crossPartitionState == null) { - CompositeContinuationToken compositeContinuationToken = new CompositeContinuationToken() + queryState = null; + } + else + { + List compositeContinuationTokens = new List(crossPartitionState.Value.Count); + foreach ((PartitionKeyRange range, QueryState state) in crossPartitionState.Value) { - Range = range.ToRange(), - Token = state != null ? ((CosmosString)state.Value).Value : null, - }; + CompositeContinuationToken compositeContinuationToken = new CompositeContinuationToken() + { + Range = range.ToRange(), + Token = state != null ? ((CosmosString)state.Value).Value : null, + }; - compositeContinuationTokens.Add(compositeContinuationToken); - } + compositeContinuationTokens.Add(compositeContinuationToken); + } - List cosmosElementContinuationTokens = compositeContinuationTokens - .Select(token => CompositeContinuationToken.ToCosmosElement(token)) - .ToList(); - CosmosArray cosmosElementCompositeContinuationTokens = CosmosArray.Create(cosmosElementContinuationTokens); + List cosmosElementContinuationTokens = compositeContinuationTokens + .Select(token => CompositeContinuationToken.ToCosmosElement(token)) + .ToList(); + CosmosArray cosmosElementCompositeContinuationTokens = CosmosArray.Create(cosmosElementContinuationTokens); + + queryState = new QueryState(cosmosElementCompositeContinuationTokens); + } QueryPage crossPartitionQueryPage = new QueryPage( backendQueryPage.Documents, @@ -63,7 +73,7 @@ public TryCatch Current backendQueryPage.ResponseLengthInBytes, backendQueryPage.CosmosQueryExecutionInfo, backendQueryPage.DisallowContinuationTokenMessage, - new QueryState(cosmosElementCompositeContinuationTokens)); + queryState); return TryCatch.FromResult(crossPartitionQueryPage); } 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 3b5a3cab6d..90a1d7c794 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 @@ -184,12 +184,15 @@ public async Task TestDrainFully_WithStateResume() QueryPage queryPage = tryGetQueryPage.Result; documents.AddRange(queryPage.Documents); - queryPipelineStage = ParallelCrossPartitionQueryPipelineStage.MonadicCreate( + if (queryPage.State != null) + { + queryPipelineStage = ParallelCrossPartitionQueryPipelineStage.MonadicCreate( feedRangeProvider: feedRangeProvider, queryDataSource: queryDataSource, sqlQuerySpec: new SqlQuerySpec("SELECT * FROM c"), pageSize: 10, continuationToken: queryPage.State.Value).Result; + } } Assert.AreEqual(numItems, documents.Count); From 33080a0257fe5307e4f8d972b51a0d8913e5ba19 Mon Sep 17 00:00:00 2001 From: Brandon Chong Date: Fri, 10 Jul 2020 02:00:40 -0700 Subject: [PATCH 36/85] simplified query draining with continuation token path --- ...elCrossPartitionQueryPipelineStageTests.cs | 35 ++++++++----------- 1 file changed, 15 insertions(+), 20 deletions(-) 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 90a1d7c794..77bffee270 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 @@ -166,34 +166,29 @@ public async Task TestDrainFully_WithStateResume() IFeedRangeProvider feedRangeProvider = new InMemoryCollectionFeedRangeProvider(inMemoryCollection); IQueryDataSource queryDataSource = new InMemoryCollectionQueryDataSource(inMemoryCollection); - TryCatch monadicCreate = ParallelCrossPartitionQueryPipelineStage.MonadicCreate( - feedRangeProvider: feedRangeProvider, - queryDataSource: queryDataSource, - sqlQuerySpec: new SqlQuerySpec("SELECT * FROM c"), - pageSize: 10, - continuationToken: default); - Assert.IsTrue(monadicCreate.Succeeded); - IQueryPipelineStage queryPipelineStage = monadicCreate.Result; - List documents = new List(); - while (await queryPipelineStage.MoveNextAsync()) + + QueryState queryState = null; + do { + TryCatch monadicCreate = ParallelCrossPartitionQueryPipelineStage.MonadicCreate( + feedRangeProvider: feedRangeProvider, + queryDataSource: queryDataSource, + sqlQuerySpec: new SqlQuerySpec("SELECT * FROM c"), + pageSize: 10, + continuationToken: queryState?.Value); + Assert.IsTrue(monadicCreate.Succeeded); + IQueryPipelineStage queryPipelineStage = monadicCreate.Result; + + Assert.IsTrue(await queryPipelineStage.MoveNextAsync()); TryCatch tryGetQueryPage = queryPipelineStage.Current; Assert.IsTrue(tryGetQueryPage.Succeeded); QueryPage queryPage = tryGetQueryPage.Result; documents.AddRange(queryPage.Documents); - if (queryPage.State != null) - { - queryPipelineStage = ParallelCrossPartitionQueryPipelineStage.MonadicCreate( - feedRangeProvider: feedRangeProvider, - queryDataSource: queryDataSource, - sqlQuerySpec: new SqlQuerySpec("SELECT * FROM c"), - pageSize: 10, - continuationToken: queryPage.State.Value).Result; - } - } + queryState = queryPage.State; + } while (queryState != null); Assert.AreEqual(numItems, documents.Count); } From 7fbb306ff7fe61d01206f0e6d44d7110565d40a1 Mon Sep 17 00:00:00 2001 From: Brandon Chong Date: Mon, 13 Jul 2020 15:23:24 -0700 Subject: [PATCH 37/85] resolved iteration comment by adding DocumentContainer as an abstraction layer --- .../src/Microsoft.Azure.Cosmos.csproj | 2 +- ...=> CreatePartitionRangeAsyncEnumerator.cs} | 2 +- ...CrossPartitionRangePageAsyncEnumerable.cs} | 14 +- ...CrossPartitionRangePageAsyncEnumerator.cs} | 32 +-- .../src/Pagination/DocumentContainer.cs | 218 +++++++++++++++ .../src/Pagination/DocumentContainerPage.cs | 22 ++ .../src/Pagination/DocumentContainerState.cs | 16 ++ .../src/Pagination/FeedRangeProvider.cs | 44 --- .../src/Pagination/IFeedRangeProvider.cs | 13 +- ...s => PartitionRangePageAsyncEnumerable.cs} | 8 +- ...s => PartitionRangePageAsyncEnumerator.cs} | 10 +- .../src/Pagination/Record.cs | 33 +++ .../src/Query/Core/Monads/TryCatch.cs | 65 ++--- .../Query/Core/Monads/TryCatch{TResult}.cs | 17 ++ .../InMemoryCollection.cs | 252 +++++++++++------- .../InMemoryCollectionTests.cs | 168 +++++++----- ...CollectionPartitionRangeEnumeratorTests.cs | 237 ---------------- ...sPartitionPartitionRangeEnumeratorTests.cs | 199 ++++++++++++++ ...cumentContainerPartitionRangeEnumerator.cs | 48 ++++ .../InMemoryCollectionFeedRangeProvider.cs | 66 ----- .../Pagination/InMemoryCollectionPage.cs | 20 -- ...emoryCollectionPartitionRangeEnumerator.cs | 62 ----- .../Pagination/InMemoryCollectionState.cs | 18 -- ...ts.cs => PartitionRangeEnumeratorTests.cs} | 98 +++++-- ...PartitionPartitionRangeEnumeratorTests.cs} | 63 ++--- 25 files changed, 977 insertions(+), 750 deletions(-) rename Microsoft.Azure.Cosmos/src/Pagination/{CreatePartitionRangeEnumerator.cs => CreatePartitionRangeAsyncEnumerator.cs} (76%) rename Microsoft.Azure.Cosmos/src/Pagination/{CrossPartitionRangePageEnumerable.cs => CrossPartitionRangePageAsyncEnumerable.cs} (68%) rename Microsoft.Azure.Cosmos/src/Pagination/{CrossPartitionRangePageEnumerator.cs => CrossPartitionRangePageAsyncEnumerator.cs} (75%) create mode 100644 Microsoft.Azure.Cosmos/src/Pagination/DocumentContainer.cs create mode 100644 Microsoft.Azure.Cosmos/src/Pagination/DocumentContainerPage.cs create mode 100644 Microsoft.Azure.Cosmos/src/Pagination/DocumentContainerState.cs delete mode 100644 Microsoft.Azure.Cosmos/src/Pagination/FeedRangeProvider.cs rename Microsoft.Azure.Cosmos/src/Pagination/{PartitionRangePageEnumerable.cs => PartitionRangePageAsyncEnumerable.cs} (76%) rename Microsoft.Azure.Cosmos/src/Pagination/{PartitionRangePageEnumerator.cs => PartitionRangePageAsyncEnumerator.cs} (72%) create mode 100644 Microsoft.Azure.Cosmos/src/Pagination/Record.cs delete mode 100644 Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/CrossPartitionInMemoryCollectionPartitionRangeEnumeratorTests.cs create mode 100644 Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/CrossPartitionPartitionRangeEnumeratorTests.cs create mode 100644 Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/DocumentContainerPartitionRangeEnumerator.cs delete mode 100644 Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/InMemoryCollectionFeedRangeProvider.cs delete mode 100644 Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/InMemoryCollectionPage.cs delete mode 100644 Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/InMemoryCollectionPartitionRangeEnumerator.cs delete mode 100644 Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/InMemoryCollectionState.cs rename Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/{InMemoryCollectionPartitionRangeEnumeratorTests.cs => PartitionRangeEnumeratorTests.cs} (58%) rename Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/{SinglePartitionInMemoryCollectionPartitionRangeEnumeratorTests.cs => SinglePartitionPartitionRangeEnumeratorTests.cs} (58%) diff --git a/Microsoft.Azure.Cosmos/src/Microsoft.Azure.Cosmos.csproj b/Microsoft.Azure.Cosmos/src/Microsoft.Azure.Cosmos.csproj index fbbfc74f70..3710c70007 100644 --- a/Microsoft.Azure.Cosmos/src/Microsoft.Azure.Cosmos.csproj +++ b/Microsoft.Azure.Cosmos/src/Microsoft.Azure.Cosmos.csproj @@ -92,7 +92,7 @@ - + diff --git a/Microsoft.Azure.Cosmos/src/Pagination/CreatePartitionRangeEnumerator.cs b/Microsoft.Azure.Cosmos/src/Pagination/CreatePartitionRangeAsyncEnumerator.cs similarity index 76% rename from Microsoft.Azure.Cosmos/src/Pagination/CreatePartitionRangeEnumerator.cs rename to Microsoft.Azure.Cosmos/src/Pagination/CreatePartitionRangeAsyncEnumerator.cs index 43210b0bc8..2d1696b8aa 100644 --- a/Microsoft.Azure.Cosmos/src/Pagination/CreatePartitionRangeEnumerator.cs +++ b/Microsoft.Azure.Cosmos/src/Pagination/CreatePartitionRangeAsyncEnumerator.cs @@ -6,7 +6,7 @@ namespace Microsoft.Azure.Cosmos.Pagination { using Microsoft.Azure.Documents; - internal delegate PartitionRangePageEnumerator CreatePartitionRangePageEnumerator( + internal delegate PartitionRangePageAsyncEnumerator CreatePartitionRangePageAsyncEnumerator( PartitionKeyRange feedRange, TState state) where TPage : Page diff --git a/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionRangePageEnumerable.cs b/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionRangePageAsyncEnumerable.cs similarity index 68% rename from Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionRangePageEnumerable.cs rename to Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionRangePageAsyncEnumerable.cs index 01f991397e..c4cb0cb669 100644 --- a/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionRangePageEnumerable.cs +++ b/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionRangePageAsyncEnumerable.cs @@ -9,19 +9,19 @@ namespace Microsoft.Azure.Cosmos.Pagination using System.Threading; using Microsoft.Azure.Cosmos.Query.Core.Monads; - internal sealed class CrossPartitionRangePageEnumerable : IAsyncEnumerable>> + internal sealed class CrossPartitionRangePageAsyncEnumerable : IAsyncEnumerable>> where TPage : Page where TState : State { private readonly CrossPartitionState state; - private readonly CreatePartitionRangePageEnumerator createPartitionRangeEnumerator; - private readonly IComparer> comparer; + private readonly CreatePartitionRangePageAsyncEnumerator createPartitionRangeEnumerator; + private readonly IComparer> comparer; private readonly IFeedRangeProvider feedRangeProvider; - public CrossPartitionRangePageEnumerable( + public CrossPartitionRangePageAsyncEnumerable( IFeedRangeProvider feedRangeProvider, - CreatePartitionRangePageEnumerator createPartitionRangeEnumerator, - IComparer> comparer, + CreatePartitionRangePageAsyncEnumerator createPartitionRangeEnumerator, + IComparer> comparer, CrossPartitionState state = default) { this.feedRangeProvider = feedRangeProvider ?? throw new ArgumentNullException(nameof(comparer)); @@ -34,7 +34,7 @@ public IAsyncEnumerator>> GetAsyncEnu { cancellationToken.ThrowIfCancellationRequested(); - return new CrossPartitionRangePageEnumerator( + return new CrossPartitionRangePageAsyncEnumerator( this.feedRangeProvider, this.createPartitionRangeEnumerator, this.comparer, diff --git a/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionRangePageEnumerator.cs b/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionRangePageAsyncEnumerator.cs similarity index 75% rename from Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionRangePageEnumerator.cs rename to Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionRangePageAsyncEnumerator.cs index ca6b1aa241..47ac66f0fc 100644 --- a/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionRangePageEnumerator.cs +++ b/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionRangePageAsyncEnumerator.cs @@ -15,20 +15,20 @@ namespace Microsoft.Azure.Cosmos.Pagination using Microsoft.Azure.Documents; /// - /// Coordinates draining pages from multiple , while maintaining a global sort order and handling repartitioning (splits, merge). + /// Coordinates draining pages from multiple , while maintaining a global sort order and handling repartitioning (splits, merge). /// - internal sealed class CrossPartitionRangePageEnumerator : IAsyncEnumerator>> + internal sealed class CrossPartitionRangePageAsyncEnumerator : IAsyncEnumerator>> where TPage : Page where TState : State { private readonly IFeedRangeProvider feedRangeProvider; - private readonly CreatePartitionRangePageEnumerator createPartitionRangeEnumerator; - private readonly AsyncLazy>> lazyEnumerators; + private readonly CreatePartitionRangePageAsyncEnumerator createPartitionRangeEnumerator; + private readonly AsyncLazy>> lazyEnumerators; - public CrossPartitionRangePageEnumerator( + public CrossPartitionRangePageAsyncEnumerator( IFeedRangeProvider feedRangeProvider, - CreatePartitionRangePageEnumerator createPartitionRangeEnumerator, - IComparer> comparer, + CreatePartitionRangePageAsyncEnumerator createPartitionRangeEnumerator, + IComparer> comparer, CrossPartitionState state = default) { this.feedRangeProvider = feedRangeProvider ?? throw new ArgumentNullException(nameof(feedRangeProvider)); @@ -39,7 +39,7 @@ public CrossPartitionRangePageEnumerator( throw new ArgumentNullException(nameof(comparer)); } - this.lazyEnumerators = new AsyncLazy>>(async (CancellationToken token) => + this.lazyEnumerators = new AsyncLazy>>(async (CancellationToken token) => { IReadOnlyList<(PartitionKeyRange, TState)> rangeAndStates; if (state != default) @@ -60,10 +60,10 @@ public CrossPartitionRangePageEnumerator( rangeAndStates = rangesAndStatesBuilder; } - PriorityQueue> enumerators = new PriorityQueue>(comparer); + PriorityQueue> enumerators = new PriorityQueue>(comparer); foreach ((PartitionKeyRange range, TState rangeState) in rangeAndStates) { - PartitionRangePageEnumerator enumerator = createPartitionRangeEnumerator(range, rangeState); + PartitionRangePageAsyncEnumerator enumerator = createPartitionRangeEnumerator(range, rangeState); enumerators.Enqueue(enumerator); } @@ -75,8 +75,8 @@ public CrossPartitionRangePageEnumerator( public async ValueTask MoveNextAsync() { - PriorityQueue> enumerators = await this.lazyEnumerators.GetValueAsync(cancellationToken: default); - PartitionRangePageEnumerator currentPaginator = enumerators.Dequeue(); + PriorityQueue> enumerators = await this.lazyEnumerators.GetValueAsync(cancellationToken: default); + PartitionRangePageAsyncEnumerator currentPaginator = enumerators.Dequeue(); bool movedNext = await currentPaginator.MoveNextAsync(); if (!movedNext) { @@ -100,7 +100,7 @@ public async ValueTask MoveNextAsync() cancellationToken: default); foreach (PartitionKeyRange childRange in childRanges) { - PartitionRangePageEnumerator childPaginator = this.createPartitionRangeEnumerator( + PartitionRangePageAsyncEnumerator childPaginator = this.createPartitionRangeEnumerator( childRange, currentPaginator.State); enumerators.Enqueue(childPaginator); @@ -126,7 +126,7 @@ public async ValueTask MoveNextAsync() } List<(PartitionKeyRange, TState)> feedRangeAndStates = new List<(PartitionKeyRange, TState)>(enumerators.Count); - foreach (PartitionRangePageEnumerator enumerator in enumerators) + foreach (PartitionRangePageAsyncEnumerator enumerator in enumerators) { feedRangeAndStates.Add((enumerator.Range, enumerator.State)); } @@ -146,8 +146,8 @@ public ValueTask DisposeAsync() private static bool IsSplitException(Exception exeception) { return exeception is CosmosException cosmosException - && cosmosException.StatusCode == HttpStatusCode.Gone - && cosmosException.SubStatusCode == (int)Documents.SubStatusCodes.PartitionKeyRangeGone; + && (cosmosException.StatusCode == HttpStatusCode.Gone) + && (cosmosException.SubStatusCode == (int)Documents.SubStatusCodes.PartitionKeyRangeGone); } private static bool IsMergeException(Exception exception) diff --git a/Microsoft.Azure.Cosmos/src/Pagination/DocumentContainer.cs b/Microsoft.Azure.Cosmos/src/Pagination/DocumentContainer.cs new file mode 100644 index 0000000000..afb1004436 --- /dev/null +++ b/Microsoft.Azure.Cosmos/src/Pagination/DocumentContainer.cs @@ -0,0 +1,218 @@ +// ------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +// ------------------------------------------------------------ + +namespace Microsoft.Azure.Cosmos.Pagination +{ + using System; + using System.Collections.Generic; + using System.Threading; + using System.Threading.Tasks; + using Microsoft.Azure.Cosmos.CosmosElements; + using Microsoft.Azure.Cosmos.Query.Core.Monads; + using Microsoft.Azure.Documents; + + internal abstract class DocumentContainer : IFeedRangeProvider + { + private static readonly CosmosException RequestRateTooLargeException = new CosmosException( + message: "Request Rate Too Large", + statusCode: (System.Net.HttpStatusCode)429, + subStatusCode: default, + activityId: Guid.NewGuid().ToString(), + requestCharge: default); + + private static readonly Task>> ThrottleForGetRanges = Task.FromResult( + TryCatch>.FromException( + RequestRateTooLargeException)); + + private static readonly Task> ThrottleForCreateItem = Task.FromResult( + TryCatch.FromException( + RequestRateTooLargeException)); + + private static readonly Task> ThrottleForFeedOperation = Task.FromResult( + TryCatch.FromException( + RequestRateTooLargeException)); + + private static readonly PartitionKeyRange FullRange = new PartitionKeyRange() + { + MinInclusive = Documents.Routing.PartitionKeyInternal.MinimumInclusiveEffectivePartitionKey, + MaxExclusive = Documents.Routing.PartitionKeyInternal.MaximumExclusiveEffectivePartitionKey, + }; + + private readonly FailureConfigs failureConfigs; + private readonly Random random; + + protected DocumentContainer(FailureConfigs failureConfigs = null) + { + this.failureConfigs = failureConfigs; + this.random = new Random(); + } + + protected abstract Task>> MonadicGetChildRangeImplementationAsync( + PartitionKeyRange partitionKeyRange, + CancellationToken cancellationToken); + + public Task>> MonadicGetChildRangeAsync( + PartitionKeyRange partitionKeyRange, + CancellationToken cancellationToken) + { + if (this.ShouldReturn429()) + { + return ThrottleForGetRanges; + } + + return this.MonadicGetChildRangeImplementationAsync( + partitionKeyRange, + cancellationToken); + } + + public Task> GetChildRangeAsync( + PartitionKeyRange partitionKeyRange, + CancellationToken cancellationToken) => TryCatch>.UnsafeGetResultAsync( + this.MonadicGetChildRangeAsync( + partitionKeyRange, + cancellationToken), + cancellationToken); + + public Task>> MonadicGetFeedRangesAsync( + CancellationToken cancellationToken) => this.MonadicGetChildRangeAsync( + DocumentContainer.FullRange, + cancellationToken); + + public Task> GetFeedRangesAsync( + CancellationToken cancellationToken) => TryCatch>.UnsafeGetResultAsync( + this.MonadicGetFeedRangesAsync( + cancellationToken), + cancellationToken); + + protected abstract Task> MonadicCreateItemImplementationAsync( + CosmosObject payload, + CancellationToken cancellationToken); + + public Task> MonadicCreateItemAsync( + CosmosObject payload, + CancellationToken cancellationToken) + { + if (this.ShouldReturn429()) + { + return ThrottleForCreateItem; + } + + return this.MonadicCreateItemImplementationAsync( + payload, + cancellationToken); + } + + public Task CreateItemAsync( + CosmosObject payload, + CancellationToken cancellationToken) => TryCatch>.UnsafeGetResultAsync( + this.MonadicCreateItemAsync( + payload, + cancellationToken), + cancellationToken); + + protected abstract Task> MonadicReadItemImplementationAsync( + CosmosElement partitionKey, + Guid identifer, + CancellationToken cancellationToken); + + public Task> MonadicReadItemAsync( + CosmosElement partitionKey, + Guid identifer, + CancellationToken cancellationToken) + { + if (this.ShouldReturn429()) + { + return ThrottleForCreateItem; + } + + return this.MonadicReadItemImplementationAsync( + partitionKey, + identifer, + cancellationToken); + } + + public Task ReadItemAsync( + CosmosElement partitionKey, + Guid identifier, + CancellationToken cancellationToken) => TryCatch.UnsafeGetResultAsync( + this.MonadicReadItemAsync( + partitionKey, + identifier, + cancellationToken), + cancellationToken); + + protected abstract Task> MonadicReadFeedImplementationAsync( + int partitionKeyRangeId, + long resourceIdentifer, + int pageSize, + CancellationToken cancellationToken); + + public Task> MonadicReadFeedAsync( + int partitionKeyRangeId, + long resourceIdentifer, + int pageSize, + CancellationToken cancellationToken) + { + if (this.ShouldReturn429()) + { + return ThrottleForFeedOperation; + } + + if (this.ShouldReturnEmptyPage()) + { + return Task.FromResult( + TryCatch.FromResult( + new DocumentContainerPage( + new List(), + new DocumentContainerState(resourceIdentifer)))); + } + + return this.MonadicReadFeedImplementationAsync( + partitionKeyRangeId, + resourceIdentifer, + pageSize, + cancellationToken); + } + + public Task ReadFeedAsync( + int partitionKeyRangeId, + long resourceIdentifier, + int pageSize, + CancellationToken cancellationToken) => TryCatch.UnsafeGetResultAsync( + this.MonadicReadFeedAsync( + partitionKeyRangeId, + resourceIdentifier, + pageSize, + cancellationToken), + cancellationToken); + + public abstract Task MonadicSplitAsync( + int partitionKeyRangeId, + CancellationToken cancellationToken); + + public Task SplitAsync( + int partitionKeyRangeId, + CancellationToken cancellationToken) => TryCatch.UnsafeWaitAsync( + this.MonadicSplitAsync( + partitionKeyRangeId, + cancellationToken), + cancellationToken); + private bool ShouldReturn429() => (this.failureConfigs != null) && this.failureConfigs.Inject429s && ((this.random.Next() % 2) == 0); + + private bool ShouldReturnEmptyPage() => (this.failureConfigs != null) && this.failureConfigs.InjectEmptyPages && ((this.random.Next() % 2) == 0); + + public sealed class FailureConfigs + { + public FailureConfigs(bool inject429s, bool injectEmptyPages) + { + this.Inject429s = inject429s; + this.InjectEmptyPages = injectEmptyPages; + } + + public bool Inject429s { get; } + + public bool InjectEmptyPages { get; } + } + } +} diff --git a/Microsoft.Azure.Cosmos/src/Pagination/DocumentContainerPage.cs b/Microsoft.Azure.Cosmos/src/Pagination/DocumentContainerPage.cs new file mode 100644 index 0000000000..9cb6fb5104 --- /dev/null +++ b/Microsoft.Azure.Cosmos/src/Pagination/DocumentContainerPage.cs @@ -0,0 +1,22 @@ +// ------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +// ------------------------------------------------------------ + +namespace Microsoft.Azure.Cosmos.Pagination +{ + using System; + using System.Collections.Generic; + + internal sealed class DocumentContainerPage : Page + { + public DocumentContainerPage( + IReadOnlyList records, + DocumentContainerState state = null) + : base(state) + { + this.Records = records ?? throw new ArgumentNullException(nameof(records)); + } + + public IReadOnlyList Records { get; } + } +} diff --git a/Microsoft.Azure.Cosmos/src/Pagination/DocumentContainerState.cs b/Microsoft.Azure.Cosmos/src/Pagination/DocumentContainerState.cs new file mode 100644 index 0000000000..cbc29a3542 --- /dev/null +++ b/Microsoft.Azure.Cosmos/src/Pagination/DocumentContainerState.cs @@ -0,0 +1,16 @@ +// ------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +// ------------------------------------------------------------ + +namespace Microsoft.Azure.Cosmos.Pagination +{ + internal sealed class DocumentContainerState : State + { + public DocumentContainerState(long resourceIdentifier) + { + this.ResourceIdentifer = resourceIdentifier; + } + + public long ResourceIdentifer { get; } + } +} diff --git a/Microsoft.Azure.Cosmos/src/Pagination/FeedRangeProvider.cs b/Microsoft.Azure.Cosmos/src/Pagination/FeedRangeProvider.cs deleted file mode 100644 index 391fbe92e5..0000000000 --- a/Microsoft.Azure.Cosmos/src/Pagination/FeedRangeProvider.cs +++ /dev/null @@ -1,44 +0,0 @@ -// ------------------------------------------------------------ -// Copyright (c) Microsoft Corporation. All rights reserved. -// ------------------------------------------------------------ - -namespace Microsoft.Azure.Cosmos.Pagination -{ - using System; - using System.Collections.Generic; - using System.Threading; - using System.Threading.Tasks; - using Microsoft.Azure.Cosmos.Query.Core.QueryClient; - using Microsoft.Azure.Documents; - - internal sealed class FeedRangeProvider : IFeedRangeProvider - { - private static readonly PartitionKeyRange FullRange = new PartitionKeyRange() - { - MinInclusive = Documents.Routing.PartitionKeyInternal.MinimumInclusiveEffectivePartitionKey, - MaxExclusive = Documents.Routing.PartitionKeyInternal.MaximumExclusiveEffectivePartitionKey, - }; - - private readonly CosmosQueryClient cosmosQueryClient; - private readonly string collectionRid; - - public FeedRangeProvider(CosmosQueryClient cosmosQueryClient, string collectionRid) - { - this.cosmosQueryClient = cosmosQueryClient ?? throw new ArgumentNullException(nameof(cosmosQueryClient)); - this.collectionRid = collectionRid ?? throw new ArgumentNullException(nameof(collectionRid)); - } - - public Task> GetChildRangeAsync( - PartitionKeyRange partitionKeyRange, - CancellationToken cancellationToken) => this.cosmosQueryClient.TryGetOverlappingRangesAsync( - this.collectionRid, - partitionKeyRange.ToRange(), - forceRefresh: true) - .ContinueWith((task) => (IEnumerable)task.Result); - - public Task> GetFeedRangesAsync( - CancellationToken cancellationToken) => this.GetChildRangeAsync( - FeedRangeProvider.FullRange, - cancellationToken); - } -} diff --git a/Microsoft.Azure.Cosmos/src/Pagination/IFeedRangeProvider.cs b/Microsoft.Azure.Cosmos/src/Pagination/IFeedRangeProvider.cs index 301a03c57d..6ec3401c51 100644 --- a/Microsoft.Azure.Cosmos/src/Pagination/IFeedRangeProvider.cs +++ b/Microsoft.Azure.Cosmos/src/Pagination/IFeedRangeProvider.cs @@ -7,14 +7,23 @@ namespace Microsoft.Azure.Cosmos.Pagination using System.Collections.Generic; using System.Threading; using System.Threading.Tasks; + using Microsoft.Azure.Cosmos.Query.Core.Monads; using Microsoft.Azure.Documents; internal interface IFeedRangeProvider { - public Task> GetChildRangeAsync( + Task>> MonadicGetChildRangeAsync( PartitionKeyRange partitionKeyRange, CancellationToken cancellationToken); - public Task> GetFeedRangesAsync(CancellationToken cancellationToken); + Task> GetChildRangeAsync( + PartitionKeyRange partitionKeyRange, + CancellationToken cancellationToken); + + Task>> MonadicGetFeedRangesAsync( + CancellationToken cancellationToken); + + Task> GetFeedRangesAsync( + CancellationToken cancellationToken); } } diff --git a/Microsoft.Azure.Cosmos/src/Pagination/PartitionRangePageEnumerable.cs b/Microsoft.Azure.Cosmos/src/Pagination/PartitionRangePageAsyncEnumerable.cs similarity index 76% rename from Microsoft.Azure.Cosmos/src/Pagination/PartitionRangePageEnumerable.cs rename to Microsoft.Azure.Cosmos/src/Pagination/PartitionRangePageAsyncEnumerable.cs index 295bc8af6c..7fe6f697c7 100644 --- a/Microsoft.Azure.Cosmos/src/Pagination/PartitionRangePageEnumerable.cs +++ b/Microsoft.Azure.Cosmos/src/Pagination/PartitionRangePageAsyncEnumerable.cs @@ -10,18 +10,18 @@ namespace Microsoft.Azure.Cosmos.Pagination using Microsoft.Azure.Cosmos.Query.Core.Monads; using Microsoft.Azure.Documents; - internal sealed class PartitionRangePageEnumerable : IAsyncEnumerable> + internal sealed class PartitionRangePageAsyncEnumerable : IAsyncEnumerable> where TPage : Page where TState : State { private readonly PartitionKeyRange range; private readonly TState state; - private readonly CreatePartitionRangePageEnumerator createPartitionRangeEnumerator; + private readonly CreatePartitionRangePageAsyncEnumerator createPartitionRangeEnumerator; - public PartitionRangePageEnumerable( + public PartitionRangePageAsyncEnumerable( PartitionKeyRange range, TState state, - CreatePartitionRangePageEnumerator createPartitionRangeEnumerator) + CreatePartitionRangePageAsyncEnumerator createPartitionRangeEnumerator) { this.range = range ?? throw new ArgumentNullException(nameof(range)); this.state = state; diff --git a/Microsoft.Azure.Cosmos/src/Pagination/PartitionRangePageEnumerator.cs b/Microsoft.Azure.Cosmos/src/Pagination/PartitionRangePageAsyncEnumerator.cs similarity index 72% rename from Microsoft.Azure.Cosmos/src/Pagination/PartitionRangePageEnumerator.cs rename to Microsoft.Azure.Cosmos/src/Pagination/PartitionRangePageAsyncEnumerator.cs index 81bf67da22..df6bd63d60 100644 --- a/Microsoft.Azure.Cosmos/src/Pagination/PartitionRangePageEnumerator.cs +++ b/Microsoft.Azure.Cosmos/src/Pagination/PartitionRangePageAsyncEnumerator.cs @@ -13,13 +13,13 @@ namespace Microsoft.Azure.Cosmos.Pagination /// /// Has the ability to page through a partition range. /// - internal abstract class PartitionRangePageEnumerator : IAsyncEnumerator> + internal abstract class PartitionRangePageAsyncEnumerator : IAsyncEnumerator> where TPage : Page where TState : State { private bool hasStarted; - protected PartitionRangePageEnumerator(PartitionKeyRange range, TState state = null) + protected PartitionRangePageAsyncEnumerator(PartitionKeyRange range, TState state = default) { this.Range = range; this.State = state; @@ -31,7 +31,7 @@ protected PartitionRangePageEnumerator(PartitionKeyRange range, TState state = n public TState State { get; private set; } - public bool HasMoreResults => !this.hasStarted || (this.State != default); + private bool HasMoreResults => !this.hasStarted || (this.State != default); public async ValueTask MoveNextAsync() { @@ -42,7 +42,7 @@ public async ValueTask MoveNextAsync() this.hasStarted = true; - this.Current = await this.GetNextPageAsync(); + this.Current = await this.GetNextPageAsync(cancellationToken: default); if (this.Current.Succeeded) { this.State = this.Current.Result.State; @@ -51,7 +51,7 @@ public async ValueTask MoveNextAsync() return true; } - public abstract Task> GetNextPageAsync(CancellationToken cancellationToken = default); + protected abstract Task> GetNextPageAsync(CancellationToken cancellationToken); public abstract ValueTask DisposeAsync(); } diff --git a/Microsoft.Azure.Cosmos/src/Pagination/Record.cs b/Microsoft.Azure.Cosmos/src/Pagination/Record.cs new file mode 100644 index 0000000000..99f34d0360 --- /dev/null +++ b/Microsoft.Azure.Cosmos/src/Pagination/Record.cs @@ -0,0 +1,33 @@ +// ------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +// ------------------------------------------------------------ + +namespace Microsoft.Azure.Cosmos.Pagination +{ + using System; + using Microsoft.Azure.Cosmos.CosmosElements; + + internal sealed class Record + { + private Record(long resourceIdentifier, long timestamp, Guid identifier, CosmosObject payload) + { + this.ResourceIdentifier = resourceIdentifier < 0 ? throw new ArgumentOutOfRangeException(nameof(resourceIdentifier)) : resourceIdentifier; + this.Timestamp = timestamp < 0 ? throw new ArgumentOutOfRangeException(nameof(timestamp)) : timestamp; + this.Identifier = identifier; + this.Payload = payload ?? throw new ArgumentNullException(nameof(payload)); + } + + public long ResourceIdentifier { get; } + + public long Timestamp { get; } + + public Guid Identifier { get; } + + public CosmosObject Payload { get; } + + public static Record Create(long previousResourceIdentifier, CosmosObject payload) + { + return new Record(previousResourceIdentifier + 1, DateTime.UtcNow.Ticks, Guid.NewGuid(), payload); + } + } +} diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Monads/TryCatch.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Monads/TryCatch.cs index 7e6a14ed76..d1a2e0e0bb 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Monads/TryCatch.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Monads/TryCatch.cs @@ -5,6 +5,8 @@ namespace Microsoft.Azure.Cosmos.Query.Core.Monads { using System; + using System.Runtime.CompilerServices; + using System.Threading; using System.Threading.Tasks; #if INTERNAL @@ -31,42 +33,28 @@ private TryCatch(TryCatch voidTryCatch) public void Match( Action onSuccess, - Action onError) - { - this.voidTryCatch.Match( + Action onError) => this.voidTryCatch.Match( onSuccess: (dummy) => { onSuccess(); }, onError: onError); - } public TryCatch Try( - Action onSuccess) - { - return new TryCatch(this.voidTryCatch.Try(onSuccess: (dummy) => { onSuccess(); })); - } + Action onSuccess) => new TryCatch(this.voidTryCatch.Try(onSuccess: (dummy) => { onSuccess(); })); public TryCatch Try( - Func onSuccess) - { - return this.voidTryCatch.Try(onSuccess: (dummy) => { return onSuccess(); }); - } + Func onSuccess) => this.voidTryCatch.Try( + onSuccess: (dummy) => { return onSuccess(); }); public Task> TryAsync( - Func> onSuccess) - { - return this.voidTryCatch.TryAsync(onSuccess: (dummy) => { return onSuccess(); }); - } + Func> onSuccess) => this.voidTryCatch.TryAsync( + onSuccess: (dummy) => { return onSuccess(); }); public TryCatch Catch( - Action onError) - { - return new TryCatch(this.voidTryCatch.Catch(onError)); - } + Action onError) => new TryCatch(this.voidTryCatch.Catch(onError)); public async Task CatchAsync( - Func onError) - { - return new TryCatch(await this.voidTryCatch.CatchAsync(onError)); - } + Func onError) => new TryCatch(await this.voidTryCatch.CatchAsync(onError)); + + public void ThrowIfFailed() => this.voidTryCatch.ThrowIfFailed(); public override bool Equals(object obj) { @@ -83,25 +71,22 @@ public override bool Equals(object obj) return false; } - public bool Equals(TryCatch other) - { - return this.voidTryCatch.Equals(other.voidTryCatch); - } + public bool Equals(TryCatch other) => this.voidTryCatch.Equals(other.voidTryCatch); - public override int GetHashCode() - { - return this.voidTryCatch.GetHashCode(); - } + public override int GetHashCode() => this.voidTryCatch.GetHashCode(); - public static TryCatch FromResult() - { - return new TryCatch(TryCatch.FromResult(default)); - } + public static TryCatch FromResult() => new TryCatch(TryCatch.FromResult(default)); - public static TryCatch FromException(Exception exception) - { - return new TryCatch(TryCatch.FromException(exception)); - } + public static TryCatch FromException(Exception exception) => new TryCatch(TryCatch.FromException(exception)); + + public static Task UnsafeWaitAsync( + Task tryCatchTask, + CancellationToken cancellationToken) => tryCatchTask.ContinueWith( + antecedent => + { + antecedent.Result.ThrowIfFailed(); + }, + cancellationToken); /// /// Represents a void return type. diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Monads/TryCatch{TResult}.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Monads/TryCatch{TResult}.cs index 07bf3f16f7..22b4b763c6 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Monads/TryCatch{TResult}.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Monads/TryCatch{TResult}.cs @@ -7,6 +7,7 @@ namespace Microsoft.Azure.Cosmos.Query.Core.Monads using System; using System.Diagnostics; using System.Runtime.ExceptionServices; + using System.Threading; using System.Threading.Tasks; #if INTERNAL @@ -215,5 +216,21 @@ public static bool ConvertToTryGet(TryCatch tryCatch, out T result) result = tryCatch.Result; return true; } + + public static T UnsafeGetResult(TryCatch tryCatch) + { + tryCatch.ThrowIfFailed(); + return tryCatch.Result; + } + + public static Task UnsafeGetResultAsync(Task> tryCatch, CancellationToken cancellationToken) + { + return tryCatch + .ContinueWith(antecedent => + { + antecedent.Result.ThrowIfFailed(); + return antecedent.Result.Result; + }, cancellationToken); + } } } diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/InMemoryCollection.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/InMemoryCollection.cs index 023426bc2d..435f1c138d 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/InMemoryCollection.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/InMemoryCollection.cs @@ -8,34 +8,29 @@ namespace Microsoft.Azure.Cosmos.Tests using System.Collections; using System.Collections.Generic; using System.Linq; + using System.Threading; + using System.Threading.Tasks; using Microsoft.Azure.Cosmos.CosmosElements; + using Microsoft.Azure.Cosmos.Pagination; using Microsoft.Azure.Cosmos.Query.Core.Monads; using Microsoft.Azure.Cosmos.Routing; using Microsoft.Azure.Documents; // Collection useful for mocking requests and repartitioning (splits / merge). - internal sealed class InMemoryCollection + internal sealed class InMemoryCollection : DocumentContainer { - private static readonly CosmosException RequestRateTooLargeException = new CosmosException( - message: "Request Rate Too Large", - statusCode: (System.Net.HttpStatusCode)429, - subStatusCode: default, - activityId: Guid.NewGuid().ToString(), - requestCharge: default); private readonly PartitionKeyDefinition partitionKeyDefinition; private readonly Dictionary parentToChildMapping; - private readonly FailureConfigs failureConfigs; - private readonly Random random; private PartitionKeyHashRangeDictionary partitionedRecords; private Dictionary partitionKeyRangeIdToHashRange; - public InMemoryCollection(PartitionKeyDefinition partitionKeyDefinition, FailureConfigs failureConfigs = default) + public InMemoryCollection( + PartitionKeyDefinition partitionKeyDefinition, + FailureConfigs failureConfigs = default) + : base(failureConfigs) { this.partitionKeyDefinition = partitionKeyDefinition ?? throw new ArgumentNullException(nameof(partitionKeyDefinition)); - this.failureConfigs = failureConfigs; - this.random = new Random(); - PartitionKeyHashRange fullRange = new PartitionKeyHashRange(startInclusive: null, endExclusive: null); PartitionKeyHashRanges partitionKeyHashRanges = PartitionKeyHashRanges.Create(new PartitionKeyHashRange[] { fullRange }); this.partitionedRecords = new PartitionKeyHashRangeDictionary(partitionKeyHashRanges); @@ -47,8 +42,77 @@ public InMemoryCollection(PartitionKeyDefinition partitionKeyDefinition, Failure this.parentToChildMapping = new Dictionary(); } - public Record CreateItem(CosmosObject payload) + protected override Task>> MonadicGetChildRangeImplementationAsync( + PartitionKeyRange partitionKeyRange, + CancellationToken cancellationToken) + { + cancellationToken.ThrowIfCancellationRequested(); + + PartitionKeyRange CreateRangeFromId(int id) + { + PartitionKeyHashRange hashRange = this.partitionKeyRangeIdToHashRange[id]; + return new PartitionKeyRange() + { + Id = id.ToString(), + MinInclusive = hashRange.StartInclusive.HasValue ? hashRange.StartInclusive.Value.ToString() : string.Empty, + MaxExclusive = hashRange.EndExclusive.HasValue ? hashRange.EndExclusive.Value.ToString() : string.Empty, + }; + } + + bool isFullRange = (partitionKeyRange.MinInclusive == Documents.Routing.PartitionKeyInternal.MinimumInclusiveEffectivePartitionKey) && + (partitionKeyRange.MaxExclusive == Documents.Routing.PartitionKeyInternal.MaximumExclusiveEffectivePartitionKey); + if (isFullRange) + { + List ranges = new List(); + foreach (int id in this.partitionKeyRangeIdToHashRange.Keys) + { + ranges.Add(CreateRangeFromId(id)); + } + + return Task.FromResult(TryCatch>.FromResult(ranges)); + } + + if (!int.TryParse(partitionKeyRange.Id, out int partitionKeyRangeId)) + { + return Task.FromResult( + TryCatch>.FromException( + new FormatException( + $"PartitionKeyRangeId: {partitionKeyRange.Id} is not an integer."))); + } + + if (!this.parentToChildMapping.TryGetValue(partitionKeyRangeId, out (int left, int right) children)) + { + return Task.FromResult( + TryCatch>.FromException( + new KeyNotFoundException( + $"PartitionKeyRangeId: {partitionKeyRangeId} does not exist."))); + } + + List childRanges = new List() + { + new PartitionKeyRange() + { + Id = children.left.ToString(), + MinInclusive = children.left.ToString(), + MaxExclusive = children.left.ToString(), + }, + new PartitionKeyRange() + { + Id = children.right.ToString(), + MinInclusive = children.right.ToString(), + MaxExclusive = children.right.ToString(), + } + }; + + return Task.FromResult(TryCatch>.FromResult(childRanges)); + } + + protected override Task> MonadicCreateItemImplementationAsync( + CosmosObject payload, + CancellationToken cancellationToken) { + cancellationToken.ThrowIfCancellationRequested(); + if (payload == null) { throw new ArgumentNullException(nameof(payload)); @@ -61,57 +125,79 @@ public Record CreateItem(CosmosObject payload) this.partitionedRecords[partitionKeyHash] = records; } - return records.Add(payload); + Record recordAdded = records.Add(payload); + + return Task.FromResult(TryCatch.FromResult(recordAdded)); } - public bool TryReadItem(CosmosElement partitionKey, Guid identifier, out Record record) + protected override Task> MonadicReadItemImplementationAsync( + CosmosElement partitionKey, + Guid identifier, + CancellationToken cancellationToken) { - PartitionKeyHash partitionKeyHash = GetHashFromPartitionKey(partitionKey, this.partitionKeyDefinition); + cancellationToken.ThrowIfCancellationRequested(); + + static Task> CreateNotFoundException(CosmosElement partitionKey, Guid identifer) + { + return Task.FromResult( + TryCatch.FromException( + new CosmosException( + message: $"Document with partitionKey: {partitionKey?.ToString() ?? "UNDEFINED"} not found.", + statusCode: System.Net.HttpStatusCode.NotFound, + subStatusCode: default, + activityId: Guid.NewGuid().ToString(), + requestCharge: 42))); + } + + PartitionKeyHash partitionKeyHash = GetHashFromPartitionKey( + partitionKey, + this.partitionKeyDefinition); + if (!this.partitionedRecords.TryGetValue(partitionKeyHash, out Records records)) { - record = default; - return false; + return CreateNotFoundException(partitionKey, identifier); } foreach (Record candidate in records) { bool identifierMatches = candidate.Identifier == identifier; - CosmosElement candidatePartitionKey = GetPartitionKeyFromPayload(candidate.Payload, this.partitionKeyDefinition); - bool partitionKeyMatches = CosmosElementEqualityComparer.Value.Equals(candidatePartitionKey, partitionKey); + CosmosElement candidatePartitionKey = GetPartitionKeyFromPayload( + candidate.Payload, + this.partitionKeyDefinition); + bool partitionKeyMatches = CosmosElementEqualityComparer.Value.Equals( + candidatePartitionKey, + partitionKey); if (identifierMatches && partitionKeyMatches) { - record = candidate; - return true; + return Task.FromResult(TryCatch.FromResult(candidate)); } } - record = default; - return false; + return CreateNotFoundException(partitionKey, identifier); } - public TryCatch<(List, long?)> ReadFeed(int partitionKeyRangeId, long resourceIndentifer, int pageSize) + protected override Task> MonadicReadFeedImplementationAsync( + int partitionKeyRangeId, + long resourceIdentifer, + int pageSize, + CancellationToken cancellationToken) { - if (this.Return429()) - { - return TryCatch<(List, long?)>.FromException(RequestRateTooLargeException); - } - - if (this.ReturnEmptyPage()) - { - return TryCatch<(List, long?)>.FromResult((new List(), resourceIndentifer)); - } + cancellationToken.ThrowIfCancellationRequested(); - if (!this.partitionKeyRangeIdToHashRange.TryGetValue(partitionKeyRangeId, out PartitionKeyHashRange range)) + if (!this.partitionKeyRangeIdToHashRange.TryGetValue( + partitionKeyRangeId, + out PartitionKeyHashRange range)) { - return TryCatch<(List, long?)>.FromException( - new CosmosException( + return Task.FromResult( + TryCatch.FromException( + new CosmosException( message: $"PartitionKeyRangeId {partitionKeyRangeId} is gone", statusCode: System.Net.HttpStatusCode.Gone, subStatusCode: (int)SubStatusCodes.PartitionKeyRangeGone, activityId: Guid.NewGuid().ToString(), - requestCharge: default)); + requestCharge: 42))); } if (!this.partitionedRecords.TryGetValue(range, out Records records)) @@ -120,26 +206,45 @@ record = default; } List page = records - .Where(record => record.ResourceIdentifier > resourceIndentifer) + .Where(record => record.ResourceIdentifier > resourceIdentifer) .Take(pageSize) .ToList(); if (page.Count == 0) { - return TryCatch<(List, long?)>.FromResult((page, null)); + return Task.FromResult( + TryCatch.FromResult( + new DocumentContainerPage( + records: page, + state: default))); } - return TryCatch<(List, long?)>.FromResult((page, page.Last().ResourceIdentifier)); + return Task.FromResult( + TryCatch.FromResult( + new DocumentContainerPage( + records: page, + state: new DocumentContainerState(page.Last().ResourceIdentifier)))); } - public IReadOnlyDictionary PartitionKeyRangeFeedReed() => this.partitionKeyRangeIdToHashRange; - - public void Split(int partitionKeyRangeId) + public override Task MonadicSplitAsync( + int partitionKeyRangeId, + CancellationToken cancellationToken) { + cancellationToken.ThrowIfCancellationRequested(); + // Get the current range and records - if (!this.partitionKeyRangeIdToHashRange.TryGetValue(partitionKeyRangeId, out PartitionKeyHashRange parentRange)) + if (!this.partitionKeyRangeIdToHashRange.TryGetValue( + partitionKeyRangeId, + out PartitionKeyHashRange parentRange)) { - throw new InvalidOperationException("Failed to find the range."); + return Task.FromResult( + TryCatch.FromException( + new CosmosException( + message: $"PartitionKeyRangeId {partitionKeyRangeId} is gone", + statusCode: System.Net.HttpStatusCode.Gone, + subStatusCode: (int)SubStatusCodes.PartitionKeyRangeGone, + activityId: Guid.NewGuid().ToString(), + requestCharge: 42))); } if (!this.partitionedRecords.TryGetValue(parentRange, out Records parentRecords)) @@ -150,7 +255,9 @@ public void Split(int partitionKeyRangeId) int maxPartitionKeyRangeId = this.partitionKeyRangeIdToHashRange.Keys.Max(); // Split the range space - PartitionKeyHashRanges partitionKeyHashRanges = PartitionKeyHashRangeSplitterAndMerger.SplitRange(parentRange, 2); + PartitionKeyHashRanges partitionKeyHashRanges = PartitionKeyHashRangeSplitterAndMerger.SplitRange( + parentRange, + rangeCount: 2); // Update the partition routing map this.parentToChildMapping[partitionKeyRangeId] = (maxPartitionKeyRangeId + 1, maxPartitionKeyRangeId + 2); @@ -200,15 +307,15 @@ public void Split(int partitionKeyRangeId) records.Add(record); } - } - - public (int, int) GetChildRanges(int partitionKeyRangeId) => this.parentToChildMapping[partitionKeyRangeId]; - public bool Return429() => (this.failureConfigs != null) && this.failureConfigs.Inject429s && ((this.random.Next() % 2) == 0); + return Task.FromResult(TryCatch.FromResult()); + } - public bool ReturnEmptyPage() => (this.failureConfigs != null) && this.failureConfigs.InjectEmptyPages && ((this.random.Next() % 2) == 0); + public IEnumerable PartitionKeyRangeIds => this.partitionKeyRangeIdToHashRange.Keys; - private static PartitionKeyHash GetHashFromPayload(CosmosObject payload, PartitionKeyDefinition partitionKeyDefinition) + private static PartitionKeyHash GetHashFromPayload( + CosmosObject payload, + PartitionKeyDefinition partitionKeyDefinition) { CosmosElement partitionKey = GetPartitionKeyFromPayload(payload, partitionKeyDefinition); return GetHashFromPartitionKey(partitionKey, partitionKeyDefinition); @@ -279,43 +386,6 @@ private static PartitionKeyHash GetHashFromPartitionKey(CosmosElement partitionK return partitionKeyHash; } - public sealed class Record - { - private Record(long resourceIdentifier, long timestamp, Guid identifier, CosmosObject payload) - { - this.ResourceIdentifier = resourceIdentifier < 0 ? throw new ArgumentOutOfRangeException(nameof(resourceIdentifier)) : resourceIdentifier; - this.Timestamp = timestamp < 0 ? throw new ArgumentOutOfRangeException(nameof(timestamp)) : timestamp; - this.Identifier = identifier; - this.Payload = payload ?? throw new ArgumentNullException(nameof(payload)); - } - - public long ResourceIdentifier { get; } - - public long Timestamp { get; } - - public Guid Identifier { get; } - - public CosmosObject Payload { get; } - - public static Record Create(long previousResourceIdentifier, CosmosObject payload) - { - return new Record(previousResourceIdentifier + 1, DateTime.UtcNow.Ticks, Guid.NewGuid(), payload); - } - } - - public sealed class FailureConfigs - { - public FailureConfigs(bool inject429s, bool injectEmptyPages) - { - this.Inject429s = inject429s; - this.InjectEmptyPages = injectEmptyPages; - } - - public bool Inject429s { get; } - - public bool InjectEmptyPages { get; } - } - private sealed class Records : IReadOnlyList { private readonly List storage; diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/InMemoryCollectionTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/InMemoryCollectionTests.cs index 14c62fecdc..0da6880823 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/InMemoryCollectionTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/InMemoryCollectionTests.cs @@ -7,10 +7,11 @@ namespace Microsoft.Azure.Cosmos.Tests using System; using System.Collections.Generic; using System.Linq; + using System.Threading.Tasks; using Microsoft.Azure.Cosmos.CosmosElements; using Microsoft.Azure.Cosmos.CosmosElements.Numbers; + using Microsoft.Azure.Cosmos.Pagination; using Microsoft.Azure.Cosmos.Query.Core.Monads; - using Microsoft.Azure.Cosmos.Serialization.HybridRow.RecordIO; using Microsoft.Azure.Documents; using Microsoft.VisualStudio.TestTools.UnitTesting; @@ -18,7 +19,7 @@ namespace Microsoft.Azure.Cosmos.Tests public class InMemoryCollectionTests { [TestMethod] - public void TestCrud() + public async Task TestCrudAsync() { PartitionKeyDefinition partitionKeyDefinition = new PartitionKeyDefinition() { @@ -34,23 +35,22 @@ public void TestCrud() // Insert an item CosmosObject item = CosmosObject.Parse("{\"pk\" : 42 }"); - InMemoryCollection.Record record = inMemoryCollection.CreateItem(item); + Record record = await inMemoryCollection.CreateItemAsync(item, cancellationToken: default); Assert.IsNotNull(record); Assert.AreNotEqual(Guid.Empty, record.Identifier); Assert.AreEqual(1, record.ResourceIdentifier); // Try to read it back - Assert.IsTrue( - inMemoryCollection.TryReadItem( - partitionKey: CosmosNumber64.Create(42), - record.Identifier, - out InMemoryCollection.Record readRecord)); + Record readRecord = await inMemoryCollection.ReadItemAsync( + partitionKey: CosmosNumber64.Create(42), + record.Identifier, + cancellationToken: default); Assert.AreEqual(item.ToString(), readRecord.Payload.ToString()); } [TestMethod] - public void TestPartitionKey() + public async Task TestPartitionKeyAsync() { PartitionKeyDefinition partitionKeyDefinition = new PartitionKeyDefinition() { @@ -66,22 +66,22 @@ public void TestPartitionKey() // Insert an item CosmosObject item1 = CosmosObject.Parse("{\"pk\" : 42 }"); - InMemoryCollection.Record record1 = inMemoryCollection.CreateItem(item1); + Record record1 = await inMemoryCollection.CreateItemAsync(item1, cancellationToken: default); // Insert into another partition key CosmosObject item2 = CosmosObject.Parse("{\"pk\" : 1337 }"); - InMemoryCollection.Record record2 = inMemoryCollection.CreateItem(item2); + Record record2 = await inMemoryCollection.CreateItemAsync(item2, cancellationToken: default); // Try to read back an id with wrong pk - Assert.IsFalse( - inMemoryCollection.TryReadItem( - partitionKey: item1["pk"], - record2.Identifier, - out InMemoryCollection.Record _)); + TryCatch monadicReadItem = await inMemoryCollection.MonadicReadItemAsync( + partitionKey: item1["pk"], + record2.Identifier, + cancellationToken: default); + Assert.IsFalse(monadicReadItem.Succeeded); } [TestMethod] - public void UndefinedPartitionKey() + public async Task TestUndefinedPartitionKeyAsync() { PartitionKeyDefinition partitionKeyDefinition = new PartitionKeyDefinition() { @@ -97,18 +97,18 @@ public void UndefinedPartitionKey() // Insert an item CosmosObject item = CosmosObject.Parse("{}"); - InMemoryCollection.Record record = inMemoryCollection.CreateItem(item); + Record record = await inMemoryCollection.CreateItemAsync(item, cancellationToken: default); // Try to read back an id with wrong pk - Assert.IsTrue( - inMemoryCollection.TryReadItem( - partitionKey: null, - record.Identifier, - out InMemoryCollection.Record _)); + TryCatch monadicReadItem = await inMemoryCollection.MonadicReadItemAsync( + partitionKey: null, + record.Identifier, + cancellationToken: default); + Assert.IsFalse(monadicReadItem.Succeeded); } [TestMethod] - public void Split() + public async Task TestSplitAsync() { PartitionKeyDefinition partitionKeyDefinition = new PartitionKeyDefinition() { @@ -122,33 +122,39 @@ public void Split() InMemoryCollection inMemoryCollection = new InMemoryCollection(partitionKeyDefinition); - Assert.AreEqual(1, inMemoryCollection.PartitionKeyRangeFeedReed().Count); + Assert.AreEqual(1, (await inMemoryCollection.GetFeedRangesAsync(cancellationToken: default)).Count); int numItemsToInsert = 10; for (int i = 0; i < numItemsToInsert; i++) { // Insert an item CosmosObject item = CosmosObject.Parse($"{{\"pk\" : {i} }}"); - inMemoryCollection.CreateItem(item); + await inMemoryCollection.CreateItemAsync(item, cancellationToken: default); } - inMemoryCollection.Split(partitionKeyRangeId: 0); + await inMemoryCollection.SplitAsync(partitionKeyRangeId: 0, cancellationToken: default); - Assert.AreEqual(2, inMemoryCollection.PartitionKeyRangeFeedReed().Count); - (int leftChild, int rightChild) = inMemoryCollection.GetChildRanges(partitionKeyRangeId: 0); - Assert.AreEqual(1, leftChild); - Assert.AreEqual(2, rightChild); + Assert.AreEqual(2, (await inMemoryCollection.GetFeedRangesAsync(cancellationToken: default)).Count); + List ranges = await inMemoryCollection.GetChildRangeAsync( + new PartitionKeyRange() + { + Id = "0" + }, + cancellationToken: default); + Assert.AreEqual(2, ranges.Count); + Assert.AreEqual(1, int.Parse(ranges[0].Id)); + Assert.AreEqual(2, int.Parse(ranges[1].Id)); - int AssertChildPartition(int partitionKeyRangeId) + async Task AssertChildPartitionAsync(int partitionKeyRangeId) { - TryCatch<(List records, long? continuation)> tryGetPartitionRecords = inMemoryCollection.ReadFeed( + DocumentContainerPage readFeedPage = await inMemoryCollection.ReadFeedAsync( partitionKeyRangeId: partitionKeyRangeId, - resourceIndentifer: 0, - pageSize: 100); - tryGetPartitionRecords.ThrowIfFailed(); + resourceIdentifier: 0, + pageSize: 100, + cancellationToken: default); List values = new List(); - foreach (InMemoryCollection.Record record in tryGetPartitionRecords.Result.records) + foreach (Record record in readFeedPage.Records) { values.Add(Number64.ToLong((record.Payload["pk"] as CosmosNumber).Value)); } @@ -159,12 +165,12 @@ int AssertChildPartition(int partitionKeyRangeId) return values.Count; } - int count = AssertChildPartition(partitionKeyRangeId: 1) + AssertChildPartition(partitionKeyRangeId: 2); + int count = await AssertChildPartitionAsync(partitionKeyRangeId: 1) + await AssertChildPartitionAsync(partitionKeyRangeId: 2); Assert.AreEqual(numItemsToInsert, count); } [TestMethod] - public void MultiSplit() + public async Task TestMultiSplitAsync() { PartitionKeyDefinition partitionKeyDefinition = new PartitionKeyDefinition() { @@ -183,41 +189,59 @@ public void MultiSplit() { // Insert an item CosmosObject item = CosmosObject.Parse($"{{\"pk\" : {i} }}"); - inMemoryCollection.CreateItem(item); + await inMemoryCollection.CreateItemAsync(item, cancellationToken: default); } - Assert.AreEqual(1, inMemoryCollection.PartitionKeyRangeFeedReed().Count); - - inMemoryCollection.Split(partitionKeyRangeId: 0); - - Assert.AreEqual(2, inMemoryCollection.PartitionKeyRangeFeedReed().Count); - (int leftChild, int rightChild) = inMemoryCollection.GetChildRanges(partitionKeyRangeId: 0); - Assert.AreEqual(1, leftChild); - Assert.AreEqual(2, rightChild); - - inMemoryCollection.Split(partitionKeyRangeId: 1); - inMemoryCollection.Split(partitionKeyRangeId: 2); - - - Assert.AreEqual(4, inMemoryCollection.PartitionKeyRangeFeedReed().Count); - (int leftChild1, int rightChild1) = inMemoryCollection.GetChildRanges(partitionKeyRangeId: 1); - Assert.AreEqual(3, leftChild1); - Assert.AreEqual(4, rightChild1); - - (int leftChild2, int rightChild2) = inMemoryCollection.GetChildRanges(partitionKeyRangeId: 2); - Assert.AreEqual(5, leftChild2); - Assert.AreEqual(6, rightChild2); - - int AssertChildPartition(int partitionKeyRangeId) + Assert.AreEqual(1, (await inMemoryCollection.GetFeedRangesAsync(cancellationToken: default)).Count); + + await inMemoryCollection.SplitAsync(partitionKeyRangeId: 0, cancellationToken: default); + + Assert.AreEqual(2, (await inMemoryCollection.GetFeedRangesAsync(cancellationToken: default)).Count); + List ranges = await inMemoryCollection.GetChildRangeAsync( + new PartitionKeyRange() + { + Id = "0" + }, + cancellationToken: default); + Assert.AreEqual(2, ranges.Count); + Assert.AreEqual(1, int.Parse(ranges[0].Id)); + Assert.AreEqual(2, int.Parse(ranges[1].Id)); + + await inMemoryCollection.SplitAsync(partitionKeyRangeId: 1, cancellationToken: default); + await inMemoryCollection.SplitAsync(partitionKeyRangeId: 2, cancellationToken: default); + + + Assert.AreEqual(4, (await inMemoryCollection.GetFeedRangesAsync(cancellationToken: default)).Count); + List ranges1 = await inMemoryCollection.GetChildRangeAsync( + new PartitionKeyRange() + { + Id = "1" + }, + cancellationToken: default); + Assert.AreEqual(2, ranges.Count); + Assert.AreEqual(3, int.Parse(ranges[0].Id)); + Assert.AreEqual(4, int.Parse(ranges[1].Id)); + + List ranges2 = await inMemoryCollection.GetChildRangeAsync( + new PartitionKeyRange() + { + Id = "2" + }, + cancellationToken: default); + Assert.AreEqual(2, ranges.Count); + Assert.AreEqual(5, int.Parse(ranges[0].Id)); + Assert.AreEqual(6, int.Parse(ranges[1].Id)); + + async Task AssertChildPartitionAsync(int partitionKeyRangeId) { - TryCatch<(List records, long? continuation)> tryGetPartitionRecords = inMemoryCollection.ReadFeed( + DocumentContainerPage page = await inMemoryCollection.ReadFeedAsync( partitionKeyRangeId: partitionKeyRangeId, - resourceIndentifer: 0, - pageSize: 100); - tryGetPartitionRecords.ThrowIfFailed(); + resourceIdentifier: 0, + pageSize: 100, + cancellationToken: default); List values = new List(); - foreach (InMemoryCollection.Record record in tryGetPartitionRecords.Result.records) + foreach (Record record in page.Records) { values.Add(Number64.ToLong((record.Payload["pk"] as CosmosNumber).Value)); } @@ -228,10 +252,10 @@ int AssertChildPartition(int partitionKeyRangeId) return values.Count; } - int count = AssertChildPartition(partitionKeyRangeId: 3) - + AssertChildPartition(partitionKeyRangeId: 4) - + AssertChildPartition(partitionKeyRangeId: 5) - + AssertChildPartition(partitionKeyRangeId: 6); + int count = await AssertChildPartitionAsync(partitionKeyRangeId: 3) + + await AssertChildPartitionAsync(partitionKeyRangeId: 4) + + await AssertChildPartitionAsync(partitionKeyRangeId: 5) + + await AssertChildPartitionAsync(partitionKeyRangeId: 6); Assert.AreEqual(numItemsToInsert, count); } } diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/CrossPartitionInMemoryCollectionPartitionRangeEnumeratorTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/CrossPartitionInMemoryCollectionPartitionRangeEnumeratorTests.cs deleted file mode 100644 index a10eac6656..0000000000 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/CrossPartitionInMemoryCollectionPartitionRangeEnumeratorTests.cs +++ /dev/null @@ -1,237 +0,0 @@ -//------------------------------------------------------------ -// Copyright (c) Microsoft Corporation. All rights reserved. -//------------------------------------------------------------ - -namespace Microsoft.Azure.Cosmos.Tests.Pagination -{ - using System; - using System.Collections.Generic; - using System.Linq; - using System.Threading.Tasks; - using Microsoft.Azure.Cosmos.CosmosElements; - using Microsoft.Azure.Cosmos.Pagination; - using Microsoft.Azure.Cosmos.Query.Core.Monads; - using Microsoft.Azure.Documents; - using Microsoft.VisualStudio.TestTools.UnitTesting; - - [TestClass] - public sealed class CrossPartitionInMemoryCollectionPartitionRangeEnumeratorTests - { - [TestMethod] - public async Task Test429sAsync() - { - Implementation implementation = new Implementation(); - await implementation.Test429sAsync(); - } - - [TestMethod] - public async Task Test429sWithContinuationsAsync() - { - Implementation implementation = new Implementation(); - await implementation.Test429sWithContinuationsAsync(); - } - - [TestMethod] - public async Task TestDrainFullyAsync() - { - Implementation implementation = new Implementation(); - await implementation.TestDrainFullyAsync(); - } - - [TestMethod] - public async Task TestEmptyPages() - { - Implementation implementation = new Implementation(); - await implementation.TestEmptyPages(); - } - - [TestMethod] - public async Task TestResumingFromStateAsync() - { - Implementation implementation = new Implementation(); - await implementation.TestResumingFromStateAsync(); - } - - [TestMethod] - public async Task TestSplitWithDuringDrainAsync() - { - Implementation implementation = new Implementation(); - await implementation.TestSplitWithDuringDrainAsync(); - } - - [TestMethod] - public async Task TestSplitWithResumeContinuationAsync() - { - Implementation implementation = new Implementation(); - await implementation.TestSplitWithResumeContinuationAsync(); - } - - private sealed class Implementation : InMemoryCollectionPartitionRangeEnumeratorTests, CrossPartitionState> - { - [TestMethod] - public async Task TestSplitWithResumeContinuationAsync() - { - int numItems = 1000; - InMemoryCollection inMemoryCollection = this.CreateInMemoryCollection(numItems); - IAsyncEnumerator>> enumerator = this.CreateEnumerator(inMemoryCollection); - - (HashSet firstDrainResults, CrossPartitionState state) = await this.PartialDrainAsync(enumerator, numIterations: 3); - - int minPartitionKeyRangeId = inMemoryCollection.PartitionKeyRangeFeedReed().Keys.Min(); - int maxPartitionKeyRangeId = inMemoryCollection.PartitionKeyRangeFeedReed().Keys.Max(); - // Split the partition we were reading from - inMemoryCollection.Split(minPartitionKeyRangeId); - - // And a partition we have let to read from - inMemoryCollection.Split((minPartitionKeyRangeId + maxPartitionKeyRangeId) / 2); - - // Resume from state - IAsyncEnumerable>> enumerable = this.CreateEnumerable(inMemoryCollection, state); - - HashSet secondDrainResults = await this.DrainFullyAsync(enumerable); - Assert.AreEqual(numItems, firstDrainResults.Count + secondDrainResults.Count); - } - - [TestMethod] - public async Task TestSplitWithDuringDrainAsync() - { - int numItems = 1000; - InMemoryCollection inMemoryCollection = this.CreateInMemoryCollection(numItems); - IAsyncEnumerable>> enumerable = this.CreateEnumerable(inMemoryCollection); - - HashSet identifiers = new HashSet(); - Random random = new Random(); - await foreach (TryCatch> tryGetPage in enumerable) - { - if (random.Next() % 2 == 0) - { - List partitionKeyRangeIds = inMemoryCollection.PartitionKeyRangeFeedReed().Keys.ToList(); - int randomIdToSplit = partitionKeyRangeIds[random.Next(0, partitionKeyRangeIds.Count)]; - inMemoryCollection.Split(randomIdToSplit); - } - - tryGetPage.ThrowIfFailed(); - - List records = this.GetRecordsFromPage(tryGetPage.Result); - foreach (InMemoryCollection.Record record in records) - { - identifiers.Add(record.Identifier); - } - } - - Assert.AreEqual(numItems, identifiers.Count); - } - - internal override InMemoryCollection CreateInMemoryCollection(int numItems, InMemoryCollection.FailureConfigs failureConfigs = null) - { - PartitionKeyDefinition partitionKeyDefinition = new PartitionKeyDefinition() - { - Paths = new System.Collections.ObjectModel.Collection() - { - "/pk" - }, - Kind = PartitionKind.Hash, - Version = PartitionKeyDefinitionVersion.V2, - }; - - InMemoryCollection inMemoryCollection = new InMemoryCollection(partitionKeyDefinition, failureConfigs); - - inMemoryCollection.Split(partitionKeyRangeId: 0); - - inMemoryCollection.Split(partitionKeyRangeId: 1); - inMemoryCollection.Split(partitionKeyRangeId: 2); - - inMemoryCollection.Split(partitionKeyRangeId: 3); - inMemoryCollection.Split(partitionKeyRangeId: 4); - inMemoryCollection.Split(partitionKeyRangeId: 5); - inMemoryCollection.Split(partitionKeyRangeId: 6); - - for (int i = 0; i < numItems; i++) - { - // Insert an item - CosmosObject item = CosmosObject.Parse($"{{\"pk\" : {i} }}"); - inMemoryCollection.CreateItem(item); - } - - return inMemoryCollection; - } - - internal override IAsyncEnumerable>> CreateEnumerable( - InMemoryCollection inMemoryCollection, - CrossPartitionState state = null) - { - IFeedRangeProvider feedRangeProvider = new InMemoryCollectionFeedRangeProvider(inMemoryCollection); - PartitionRangePageEnumerator createEnumerator( - PartitionKeyRange range, - InMemoryCollectionState state) => new InMemoryCollectionPartitionRangeEnumerator( - inMemoryCollection, - partitionKeyRangeId: int.Parse(range.Id), - pageSize: 10, - state: state); - - return new CrossPartitionRangePageEnumerable( - feedRangeProvider: feedRangeProvider, - createPartitionRangeEnumerator: createEnumerator, - comparer: PartitionRangePageEnumeratorComparer.Singleton, - state: state); - } - - internal override IAsyncEnumerator>> CreateEnumerator( - InMemoryCollection inMemoryCollection, - CrossPartitionState state = null) - { - IFeedRangeProvider feedRangeProvider = new InMemoryCollectionFeedRangeProvider(inMemoryCollection); - PartitionRangePageEnumerator createEnumerator( - PartitionKeyRange range, - InMemoryCollectionState state) => new InMemoryCollectionPartitionRangeEnumerator( - inMemoryCollection, - partitionKeyRangeId: int.Parse(range.Id), - pageSize: 10, - state: state); - - CrossPartitionRangePageEnumerator enumerator = new CrossPartitionRangePageEnumerator( - feedRangeProvider: feedRangeProvider, - createPartitionRangeEnumerator: createEnumerator, - comparer: PartitionRangePageEnumeratorComparer.Singleton, - state: state); - - return enumerator; - } - - internal override List GetRecordsFromPage(CrossPartitionPage page) - { - return page.Page.Records; - } - - private sealed class PartitionRangePageEnumeratorComparer : IComparer> - { - public static readonly PartitionRangePageEnumeratorComparer Singleton = new PartitionRangePageEnumeratorComparer(); - - public int Compare( - PartitionRangePageEnumerator partitionRangePageEnumerator1, - PartitionRangePageEnumerator partitionRangePageEnumerator2) - { - if (object.ReferenceEquals(partitionRangePageEnumerator1, partitionRangePageEnumerator2)) - { - return 0; - } - - if (partitionRangePageEnumerator1.HasMoreResults && !partitionRangePageEnumerator2.HasMoreResults) - { - return -1; - } - - if (!partitionRangePageEnumerator1.HasMoreResults && partitionRangePageEnumerator2.HasMoreResults) - { - return 1; - } - - // Either both don't have results or both do. - return string.CompareOrdinal( - partitionRangePageEnumerator1.Range.MinInclusive, - partitionRangePageEnumerator2.Range.MinInclusive); - } - } - } - } -} diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/CrossPartitionPartitionRangeEnumeratorTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/CrossPartitionPartitionRangeEnumeratorTests.cs new file mode 100644 index 0000000000..e9a68b8a2a --- /dev/null +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/CrossPartitionPartitionRangeEnumeratorTests.cs @@ -0,0 +1,199 @@ +//------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +//------------------------------------------------------------ + +namespace Microsoft.Azure.Cosmos.Tests.Pagination +{ + using System; + using System.Collections.Generic; + using System.Linq; + using System.Threading.Tasks; + using Microsoft.Azure.Cosmos.CosmosElements; + using Microsoft.Azure.Cosmos.Pagination; + using Microsoft.Azure.Cosmos.Query.Core.Monads; + using Microsoft.Azure.Documents; + using Microsoft.VisualStudio.TestTools.UnitTesting; + + [TestClass] + public sealed class CrossPartitionPartitionRangeEnumeratorTests + { + [TestMethod] + public async Task Test429sAsync() + { + Implementation implementation = new Implementation(); + await implementation.Test429sAsync(); + } + + [TestMethod] + public async Task Test429sWithContinuationsAsync() + { + Implementation implementation = new Implementation(); + await implementation.Test429sWithContinuationsAsync(); + } + + [TestMethod] + public async Task TestDrainFullyAsync() + { + Implementation implementation = new Implementation(); + await implementation.TestDrainFullyAsync(); + } + + [TestMethod] + public async Task TestEmptyPages() + { + Implementation implementation = new Implementation(); + await implementation.TestEmptyPages(); + } + + [TestMethod] + public async Task TestResumingFromStateAsync() + { + Implementation implementation = new Implementation(); + await implementation.TestResumingFromStateAsync(); + } + + [TestMethod] + public async Task TestSplitWithDuringDrainAsync() + { + Implementation implementation = new Implementation(); + await implementation.TestSplitWithDuringDrainAsync(); + } + + [TestMethod] + public async Task TestSplitWithResumeContinuationAsync() + { + Implementation implementation = new Implementation(); + await implementation.TestSplitWithResumeContinuationAsync(); + } + + private sealed class Implementation : PartitionRangeEnumeratorTests, CrossPartitionState> + { + public Implementation() + : base(singlePartition: false) + { + } + [TestMethod] + public async Task TestSplitWithResumeContinuationAsync() + { + int numItems = 1000; + DocumentContainer inMemoryCollection = await this.CreateDocumentContainerAsync(numItems); + IAsyncEnumerator>> enumerator = this.CreateEnumerator(inMemoryCollection); + + (HashSet firstDrainResults, CrossPartitionState state) = await this.PartialDrainAsync(enumerator, numIterations: 3); + + int minPartitionKeyRangeId = (await inMemoryCollection.GetFeedRangesAsync(cancellationToken: default)) + .Select(range => int.Parse(range.Id)).Min(); + int maxPartitionKeyRangeId = (await inMemoryCollection.GetFeedRangesAsync(cancellationToken: default)) + .Select(range => int.Parse(range.Id)).Max(); + // Split the partition we were reading from + await inMemoryCollection.SplitAsync(minPartitionKeyRangeId, cancellationToken: default); + + // And a partition we have let to read from + await inMemoryCollection.SplitAsync((minPartitionKeyRangeId + maxPartitionKeyRangeId) / 2, cancellationToken: default); + + // Resume from state + IAsyncEnumerable>> enumerable = this.CreateEnumerable(inMemoryCollection, state); + + HashSet secondDrainResults = await this.DrainFullyAsync(enumerable); + Assert.AreEqual(numItems, firstDrainResults.Count + secondDrainResults.Count); + } + + [TestMethod] + public async Task TestSplitWithDuringDrainAsync() + { + int numItems = 1000; + DocumentContainer inMemoryCollection = await this.CreateDocumentContainerAsync(numItems); + IAsyncEnumerable>> enumerable = this.CreateEnumerable(inMemoryCollection); + + HashSet identifiers = new HashSet(); + Random random = new Random(); + await foreach (TryCatch> tryGetPage in enumerable) + { + if (random.Next() % 2 == 0) + { + List partitionKeyRangeIds = (await inMemoryCollection.GetFeedRangesAsync(cancellationToken: default)) + .Select(range => int.Parse(range.Id)) + .ToList(); + int randomIdToSplit = partitionKeyRangeIds[random.Next(0, partitionKeyRangeIds.Count)]; + await inMemoryCollection.SplitAsync(randomIdToSplit, cancellationToken: default); + } + + tryGetPage.ThrowIfFailed(); + + IReadOnlyList records = this.GetRecordsFromPage(tryGetPage.Result); + foreach (Record record in records) + { + identifiers.Add(record.Identifier); + } + } + + Assert.AreEqual(numItems, identifiers.Count); + } + + public override IAsyncEnumerable>> CreateEnumerable( + DocumentContainer inMemoryCollection, + CrossPartitionState state = null) + { + PartitionRangePageAsyncEnumerator createEnumerator( + PartitionKeyRange range, + DocumentContainerState state) => new DocumentContainerPartitionRangeEnumerator( + inMemoryCollection, + partitionKeyRangeId: int.Parse(range.Id), + pageSize: 10, + state: state); + + return new CrossPartitionRangePageAsyncEnumerable( + feedRangeProvider: inMemoryCollection, + createPartitionRangeEnumerator: createEnumerator, + comparer: PartitionRangePageEnumeratorComparer.Singleton, + state: state); + } + + public override IAsyncEnumerator>> CreateEnumerator( + DocumentContainer inMemoryCollection, + CrossPartitionState state = null) + { + PartitionRangePageAsyncEnumerator createEnumerator( + PartitionKeyRange range, + DocumentContainerState state) => new DocumentContainerPartitionRangeEnumerator( + inMemoryCollection, + partitionKeyRangeId: int.Parse(range.Id), + pageSize: 10, + state: state); + + CrossPartitionRangePageAsyncEnumerator enumerator = new CrossPartitionRangePageAsyncEnumerator( + feedRangeProvider: inMemoryCollection, + createPartitionRangeEnumerator: createEnumerator, + comparer: PartitionRangePageEnumeratorComparer.Singleton, + state: state); + + return enumerator; + } + + public override IReadOnlyList GetRecordsFromPage(CrossPartitionPage page) + { + return page.Page.Records; + } + + private sealed class PartitionRangePageEnumeratorComparer : IComparer> + { + public static readonly PartitionRangePageEnumeratorComparer Singleton = new PartitionRangePageEnumeratorComparer(); + + public int Compare( + PartitionRangePageAsyncEnumerator partitionRangePageEnumerator1, + PartitionRangePageAsyncEnumerator partitionRangePageEnumerator2) + { + if (object.ReferenceEquals(partitionRangePageEnumerator1, partitionRangePageEnumerator2)) + { + return 0; + } + + // Either both don't have results or both do. + return string.CompareOrdinal( + partitionRangePageEnumerator1.Range.MinInclusive, + partitionRangePageEnumerator2.Range.MinInclusive); + } + } + } + } +} diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/DocumentContainerPartitionRangeEnumerator.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/DocumentContainerPartitionRangeEnumerator.cs new file mode 100644 index 0000000000..758a52d179 --- /dev/null +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/DocumentContainerPartitionRangeEnumerator.cs @@ -0,0 +1,48 @@ +//------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +//------------------------------------------------------------ + +namespace Microsoft.Azure.Cosmos.Tests.Pagination +{ + using System; + using System.Collections.Generic; + using System.Threading; + using System.Threading.Tasks; + using Microsoft.Azure.Cosmos.Pagination; + using Microsoft.Azure.Cosmos.Query.Core.Monads; + using Microsoft.Azure.Documents; + + internal sealed class DocumentContainerPartitionRangeEnumerator : PartitionRangePageAsyncEnumerator + { + private readonly DocumentContainer documentContainer; + private readonly int pageSize; + private readonly int partitionKeyRangeId; + + public DocumentContainerPartitionRangeEnumerator( + DocumentContainer documentContainer, + int partitionKeyRangeId, + int pageSize, + DocumentContainerState state = null) + : base( + new PartitionKeyRange() + { + Id = partitionKeyRangeId.ToString(), + MinInclusive = partitionKeyRangeId.ToString(), + MaxExclusive = partitionKeyRangeId.ToString() + }, + state ?? new DocumentContainerState(resourceIdentifier: 0)) + { + this.documentContainer = documentContainer ?? throw new ArgumentNullException(nameof(documentContainer)); + this.partitionKeyRangeId = partitionKeyRangeId; + this.pageSize = pageSize; + } + + public override ValueTask DisposeAsync() => default; + + protected override Task> GetNextPageAsync(CancellationToken cancellationToken = default) => this.documentContainer.MonadicReadFeedAsync( + partitionKeyRangeId: this.partitionKeyRangeId, + resourceIdentifer: this.State.ResourceIdentifer, + pageSize: this.pageSize, + cancellationToken: default); + } +} diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/InMemoryCollectionFeedRangeProvider.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/InMemoryCollectionFeedRangeProvider.cs deleted file mode 100644 index c4135de738..0000000000 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/InMemoryCollectionFeedRangeProvider.cs +++ /dev/null @@ -1,66 +0,0 @@ -namespace Microsoft.Azure.Cosmos.Tests.Pagination -{ - using System; - using System.Collections.Generic; - using System.Threading; - using System.Threading.Tasks; - using Microsoft.Azure.Cosmos.Pagination; - using Microsoft.Azure.Documents; - - internal sealed class InMemoryCollectionFeedRangeProvider : IFeedRangeProvider - { - private readonly InMemoryCollection inMemoryCollection; - - public InMemoryCollectionFeedRangeProvider(InMemoryCollection inMemoryCollection) - { - this.inMemoryCollection = inMemoryCollection ?? throw new ArgumentNullException(nameof(inMemoryCollection)); - } - - public Task> GetChildRangeAsync( - PartitionKeyRange feedRange, - CancellationToken cancellationToken) - { - cancellationToken.ThrowIfCancellationRequested(); - - int partitionKeyRangeId = int.Parse(feedRange.Id); - (int leftChild, int rightChild) = this.inMemoryCollection.GetChildRanges(partitionKeyRangeId); - - return Task.FromResult( - (IEnumerable)new List() - { - new PartitionKeyRange() - { - Id = leftChild.ToString(), - MinInclusive = leftChild.ToString(), - MaxExclusive = leftChild.ToString(), - }, - new PartitionKeyRange() - { - Id = rightChild.ToString(), - MinInclusive = rightChild.ToString(), - MaxExclusive = rightChild.ToString(), - } - }); - } - - public Task> GetFeedRangesAsync(CancellationToken cancellationToken) - { - cancellationToken.ThrowIfCancellationRequested(); - - List ranges = new List(); - foreach (int partitionKeyRangeId in this.inMemoryCollection.PartitionKeyRangeFeedReed().Keys) - { - PartitionKeyRange range = new PartitionKeyRange() - { - Id = partitionKeyRangeId.ToString(), - MinInclusive = partitionKeyRangeId.ToString(), - MaxExclusive = partitionKeyRangeId.ToString(), - }; - - ranges.Add(range); - } - - return Task.FromResult((IEnumerable)ranges); - } - } -} diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/InMemoryCollectionPage.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/InMemoryCollectionPage.cs deleted file mode 100644 index aea233388c..0000000000 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/InMemoryCollectionPage.cs +++ /dev/null @@ -1,20 +0,0 @@ -//------------------------------------------------------------ -// Copyright (c) Microsoft Corporation. All rights reserved. -//------------------------------------------------------------ - -namespace Microsoft.Azure.Cosmos.Tests.Pagination -{ - using System.Collections.Generic; - using Microsoft.Azure.Cosmos.Pagination; - - internal sealed class InMemoryCollectionPage : Page - { - public InMemoryCollectionPage(List records, InMemoryCollectionState state) - : base(state) - { - this.Records = records; - } - - public List Records { get; } - } -} diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/InMemoryCollectionPartitionRangeEnumerator.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/InMemoryCollectionPartitionRangeEnumerator.cs deleted file mode 100644 index a24d3b7c7c..0000000000 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/InMemoryCollectionPartitionRangeEnumerator.cs +++ /dev/null @@ -1,62 +0,0 @@ -//------------------------------------------------------------ -// Copyright (c) Microsoft Corporation. All rights reserved. -//------------------------------------------------------------ - -namespace Microsoft.Azure.Cosmos.Tests.Pagination -{ - using System; - using System.Collections.Generic; - using System.Threading; - using System.Threading.Tasks; - using Microsoft.Azure.Cosmos.Pagination; - using Microsoft.Azure.Cosmos.Query.Core.Monads; - using Microsoft.Azure.Documents; - - internal sealed class InMemoryCollectionPartitionRangeEnumerator : PartitionRangePageEnumerator - { - private readonly InMemoryCollection inMemoryCollection; - private readonly int pageSize; - private readonly int partitionKeyRangeId; - - public InMemoryCollectionPartitionRangeEnumerator( - InMemoryCollection inMemoryCollection, - int partitionKeyRangeId, - int pageSize, - InMemoryCollectionState state = null) - : base( - new PartitionKeyRange() - { - Id = partitionKeyRangeId.ToString(), - MinInclusive = partitionKeyRangeId.ToString(), - MaxExclusive = partitionKeyRangeId.ToString() - }, - state ?? new InMemoryCollectionState(resourceIdentifier: 0)) - { - this.inMemoryCollection = inMemoryCollection ?? throw new ArgumentNullException(nameof(inMemoryCollection)); - this.partitionKeyRangeId = partitionKeyRangeId; - this.pageSize = pageSize; - } - - public override ValueTask DisposeAsync() => default; - - public override Task> GetNextPageAsync(CancellationToken cancellationToken = default) - { - cancellationToken.ThrowIfCancellationRequested(); - - TryCatch<(List records, long? continuation)> tryReadPage = this.inMemoryCollection.ReadFeed( - partitionKeyRangeId: this.partitionKeyRangeId, - resourceIndentifer: ((InMemoryCollectionState)this.State).ResourceIdentifier, - pageSize: this.pageSize); - - if (tryReadPage.Failed) - { - return Task.FromResult(TryCatch.FromException(tryReadPage.Exception)); - } - - InMemoryCollectionState inMemoryCollectionState = tryReadPage.Result.continuation.HasValue ? new InMemoryCollectionState(tryReadPage.Result.continuation.Value) : default; - InMemoryCollectionPage page = new InMemoryCollectionPage(tryReadPage.Result.records, inMemoryCollectionState); - - return Task.FromResult(TryCatch.FromResult(page)); - } - } -} diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/InMemoryCollectionState.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/InMemoryCollectionState.cs deleted file mode 100644 index 45c0d9f402..0000000000 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/InMemoryCollectionState.cs +++ /dev/null @@ -1,18 +0,0 @@ -//------------------------------------------------------------ -// Copyright (c) Microsoft Corporation. All rights reserved. -//------------------------------------------------------------ - -namespace Microsoft.Azure.Cosmos.Tests.Pagination -{ - using Microsoft.Azure.Cosmos.Pagination; - - internal sealed class InMemoryCollectionState : State - { - public InMemoryCollectionState(long resourceIdentifier) - { - this.ResourceIdentifier = resourceIdentifier; - } - - public long ResourceIdentifier { get; } - } -} diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/InMemoryCollectionPartitionRangeEnumeratorTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/PartitionRangeEnumeratorTests.cs similarity index 58% rename from Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/InMemoryCollectionPartitionRangeEnumeratorTests.cs rename to Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/PartitionRangeEnumeratorTests.cs index 189c8ee6dd..e365daca8e 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/InMemoryCollectionPartitionRangeEnumeratorTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/PartitionRangeEnumeratorTests.cs @@ -3,19 +3,27 @@ using System; using System.Collections.Generic; using System.Threading.Tasks; + using Microsoft.Azure.Cosmos.CosmosElements; using Microsoft.Azure.Cosmos.Pagination; using Microsoft.Azure.Cosmos.Query.Core.Monads; + using Microsoft.Azure.Documents; using Microsoft.VisualStudio.TestTools.UnitTesting; - internal abstract class InMemoryCollectionPartitionRangeEnumeratorTests + internal abstract class PartitionRangeEnumeratorTests where TPage : Page where TState : State { + private readonly bool singlePartition; + protected PartitionRangeEnumeratorTests(bool singlePartition) + { + this.singlePartition = singlePartition; + } + [TestMethod] public async Task TestDrainFullyAsync() { int numItems = 1000; - InMemoryCollection inMemoryCollection = this.CreateInMemoryCollection(numItems); + DocumentContainer inMemoryCollection = await this.CreateDocumentContainerAsync(numItems); IAsyncEnumerable> enumerable = this.CreateEnumerable(inMemoryCollection); HashSet identifiers = await this.DrainFullyAsync(enumerable); Assert.AreEqual(numItems, identifiers.Count); @@ -25,7 +33,7 @@ public async Task TestDrainFullyAsync() public async Task TestResumingFromStateAsync() { int numItems = 1000; - InMemoryCollection inMemoryCollection = this.CreateInMemoryCollection(numItems); + DocumentContainer inMemoryCollection = await this.CreateDocumentContainerAsync(numItems); IAsyncEnumerator> enumerator = this.CreateEnumerator(inMemoryCollection); (HashSet firstDrainResults, TState state) = await this.PartialDrainAsync(enumerator, numIterations: 3); @@ -40,9 +48,9 @@ public async Task TestResumingFromStateAsync() public async Task Test429sAsync() { int numItems = 100; - InMemoryCollection inMemoryCollection = this.CreateInMemoryCollection( + DocumentContainer inMemoryCollection = await this.CreateDocumentContainerAsync( numItems, - new InMemoryCollection.FailureConfigs( + new DocumentContainer.FailureConfigs( inject429s: true, injectEmptyPages: false)); @@ -66,8 +74,8 @@ public async Task Test429sAsync() } else { - List records = this.GetRecordsFromPage(tryGetPage.Result); - foreach (InMemoryCollection.Record record in records) + IReadOnlyList records = this.GetRecordsFromPage(tryGetPage.Result); + foreach (Record record in records) { identifiers.Add(record.Identifier); } @@ -81,9 +89,9 @@ public async Task Test429sAsync() public async Task Test429sWithContinuationsAsync() { int numItems = 100; - InMemoryCollection inMemoryCollection = this.CreateInMemoryCollection( + DocumentContainer inMemoryCollection = await this.CreateDocumentContainerAsync( numItems, - new InMemoryCollection.FailureConfigs( + new DocumentContainer.FailureConfigs( inject429s: true, injectEmptyPages: false)); @@ -113,8 +121,8 @@ public async Task Test429sWithContinuationsAsync() } else { - List records = this.GetRecordsFromPage(tryGetPage.Result); - foreach (InMemoryCollection.Record record in records) + IReadOnlyList records = this.GetRecordsFromPage(tryGetPage.Result); + foreach (Record record in records) { identifiers.Add(record.Identifier); } @@ -130,9 +138,9 @@ public async Task Test429sWithContinuationsAsync() public async Task TestEmptyPages() { int numItems = 100; - InMemoryCollection inMemoryCollection = this.CreateInMemoryCollection( + DocumentContainer inMemoryCollection = await this.CreateDocumentContainerAsync( numItems, - new InMemoryCollection.FailureConfigs( + new DocumentContainer.FailureConfigs( inject429s: false, injectEmptyPages: true)); IAsyncEnumerable> enumerable = this.CreateEnumerable(inMemoryCollection); @@ -140,24 +148,68 @@ public async Task TestEmptyPages() Assert.AreEqual(numItems, identifiers.Count); } - internal abstract List GetRecordsFromPage(TPage page); + public abstract IReadOnlyList GetRecordsFromPage(TPage page); - internal abstract InMemoryCollection CreateInMemoryCollection(int numItems, InMemoryCollection.FailureConfigs failureConfigs = default); + public abstract IAsyncEnumerable> CreateEnumerable(DocumentContainer documentContainer, TState state = null); + + public abstract IAsyncEnumerator> CreateEnumerator(DocumentContainer documentContainer, TState state = null); + + public async Task CreateDocumentContainerAsync( + int numItems, + DocumentContainer.FailureConfigs failureConfigs = default) + { + PartitionKeyDefinition partitionKeyDefinition = new PartitionKeyDefinition() + { + Paths = new System.Collections.ObjectModel.Collection() + { + "/pk" + }, + Kind = PartitionKind.Hash, + Version = PartitionKeyDefinitionVersion.V2, + }; + + InMemoryCollection inMemoryCollection = new InMemoryCollection(partitionKeyDefinition, failureConfigs); + + if (!this.singlePartition) + { + await inMemoryCollection.SplitAsync(partitionKeyRangeId: 0, cancellationToken: default); - internal abstract IAsyncEnumerable> CreateEnumerable(InMemoryCollection inMemoryCollection, TState state = null); + await inMemoryCollection.SplitAsync(partitionKeyRangeId: 1, cancellationToken: default); + await inMemoryCollection.SplitAsync(partitionKeyRangeId: 2, cancellationToken: default); - internal abstract IAsyncEnumerator> CreateEnumerator(InMemoryCollection inMemoryCollection, TState state = null); + await inMemoryCollection.SplitAsync(partitionKeyRangeId: 3, cancellationToken: default); + await inMemoryCollection.SplitAsync(partitionKeyRangeId: 4, cancellationToken: default); + await inMemoryCollection.SplitAsync(partitionKeyRangeId: 5, cancellationToken: default); + await inMemoryCollection.SplitAsync(partitionKeyRangeId: 6, cancellationToken: default); + } + + for (int i = 0; i < numItems; i++) + { + // Insert an item + CosmosObject item = CosmosObject.Parse($"{{\"pk\" : {i} }}"); + while (true) + { + TryCatch monadicCreateRecord = await inMemoryCollection.MonadicCreateItemAsync(item, cancellationToken: default); + if (monadicCreateRecord.Succeeded) + { + break; + } + } + } + + return inMemoryCollection; + } - internal async Task> DrainFullyAsync(IAsyncEnumerable> enumerable) + public async Task> DrainFullyAsync(IAsyncEnumerable> enumerable) { HashSet identifiers = new HashSet(); await foreach (TryCatch tryGetPage in enumerable) { tryGetPage.ThrowIfFailed(); - List records = this.GetRecordsFromPage(tryGetPage.Result); + IReadOnlyList records = this.GetRecordsFromPage(tryGetPage.Result); - foreach (InMemoryCollection.Record record in records) + foreach (Record record in records) { identifiers.Add(record.Identifier); } @@ -166,7 +218,7 @@ internal async Task> DrainFullyAsync(IAsyncEnumerable, TState)> PartialDrainAsync( + public async Task<(HashSet, TState)> PartialDrainAsync( IAsyncEnumerator> enumerator, int numIterations) { @@ -181,9 +233,9 @@ internal async Task> DrainFullyAsync(IAsyncEnumerable tryGetPage = enumerator.Current; tryGetPage.ThrowIfFailed(); - List records = this.GetRecordsFromPage(tryGetPage.Result); + IReadOnlyList records = this.GetRecordsFromPage(tryGetPage.Result); - foreach (InMemoryCollection.Record record in records) + foreach (Record record in records) { identifiers.Add(record.Identifier); } diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/SinglePartitionInMemoryCollectionPartitionRangeEnumeratorTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/SinglePartitionPartitionRangeEnumeratorTests.cs similarity index 58% rename from Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/SinglePartitionInMemoryCollectionPartitionRangeEnumeratorTests.cs rename to Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/SinglePartitionPartitionRangeEnumeratorTests.cs index af4c3d5795..1a098b7980 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/SinglePartitionInMemoryCollectionPartitionRangeEnumeratorTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/SinglePartitionPartitionRangeEnumeratorTests.cs @@ -14,7 +14,7 @@ namespace Microsoft.Azure.Cosmos.Tests.Pagination using Microsoft.Azure.Cosmos.Query.Core.Monads; [TestClass] - public sealed class SinglePartitionInMemoryCollectionPartitionRangeEnumeratorTests + public sealed class SinglePartitionPartitionRangeEnumeratorTests { [TestMethod] public async Task Test429sAsync() @@ -59,22 +59,27 @@ public async Task TestSplitAsync() } [TestClass] - private sealed class Implementation : InMemoryCollectionPartitionRangeEnumeratorTests + private sealed class Implementation : PartitionRangeEnumeratorTests { + public Implementation() + : base(singlePartition: true) + { + } + [TestMethod] public async Task TestSplitAsync() { int numItems = 100; - InMemoryCollection inMemoryCollection = this.CreateInMemoryCollection(numItems); - InMemoryCollectionPartitionRangeEnumerator enumerator = new InMemoryCollectionPartitionRangeEnumerator( + DocumentContainer inMemoryCollection = await this.CreateDocumentContainerAsync(numItems); + DocumentContainerPartitionRangeEnumerator enumerator = new DocumentContainerPartitionRangeEnumerator( inMemoryCollection, partitionKeyRangeId: 0, pageSize: 10); - (HashSet parentIdentifiers, InMemoryCollectionState state) = await this.PartialDrainAsync(enumerator, numIterations: 3); + (HashSet parentIdentifiers, DocumentContainerState state) = await this.PartialDrainAsync(enumerator, numIterations: 3); // Split the partition - inMemoryCollection.Split(partitionKeyRangeId: 0); + await inMemoryCollection.SplitAsync(partitionKeyRangeId: 0, cancellationToken: default); // Try To read from the partition that is gone. await enumerator.MoveNextAsync(); @@ -84,10 +89,10 @@ public async Task TestSplitAsync() HashSet childIdentifiers = new HashSet(); foreach (int partitionKeyRangeId in new int[] { 1, 2 }) { - PartitionRangePageEnumerable enumerable = new PartitionRangePageEnumerable( + PartitionRangePageAsyncEnumerable enumerable = new PartitionRangePageAsyncEnumerable( range: new PartitionKeyRange() { Id = partitionKeyRangeId.ToString() }, state: state, - (range, state) => new InMemoryCollectionPartitionRangeEnumerator( + (range, state) => new DocumentContainerPartitionRangeEnumerator( inMemoryCollection, partitionKeyRangeId: int.Parse(range.Id), pageSize: 10, @@ -100,49 +105,25 @@ public async Task TestSplitAsync() Assert.AreEqual(numItems, parentIdentifiers.Count + childIdentifiers.Count); } - internal override List GetRecordsFromPage(InMemoryCollectionPage page) + public override IReadOnlyList GetRecordsFromPage(DocumentContainerPage page) { return page.Records; } - internal override InMemoryCollection CreateInMemoryCollection(int numItems, InMemoryCollection.FailureConfigs failureConfigs = null) - { - PartitionKeyDefinition partitionKeyDefinition = new PartitionKeyDefinition() - { - Paths = new System.Collections.ObjectModel.Collection() - { - "/pk" - }, - Kind = PartitionKind.Hash, - Version = PartitionKeyDefinitionVersion.V2, - }; - - InMemoryCollection inMemoryCollection = new InMemoryCollection(partitionKeyDefinition, failureConfigs); - - for (int i = 0; i < numItems; i++) - { - // Insert an item - CosmosObject item = CosmosObject.Parse($"{{\"pk\" : {i} }}"); - inMemoryCollection.CreateItem(item); - } - - return inMemoryCollection; - } - - internal override IAsyncEnumerable> CreateEnumerable( - InMemoryCollection inMemoryCollection, - InMemoryCollectionState state = null) => new PartitionRangePageEnumerable( + public override IAsyncEnumerable> CreateEnumerable( + DocumentContainer documentContainer, + DocumentContainerState state = null) => new PartitionRangePageAsyncEnumerable( range: new PartitionKeyRange() { Id = "0" }, state: state, - (range, state) => new InMemoryCollectionPartitionRangeEnumerator( - inMemoryCollection, + (range, state) => new DocumentContainerPartitionRangeEnumerator( + documentContainer, partitionKeyRangeId: int.Parse(range.Id), pageSize: 10, state: state)); - internal override IAsyncEnumerator> CreateEnumerator( - InMemoryCollection inMemoryCollection, - InMemoryCollectionState state = null) => new InMemoryCollectionPartitionRangeEnumerator( + public override IAsyncEnumerator> CreateEnumerator( + DocumentContainer inMemoryCollection, + DocumentContainerState state = null) => new DocumentContainerPartitionRangeEnumerator( inMemoryCollection, partitionKeyRangeId: 0, pageSize: 10, From 34a913991bcbdacd3300623df106962fe703ccb0 Mon Sep 17 00:00:00 2001 From: Brandon Chong Date: Mon, 13 Jul 2020 19:13:49 -0700 Subject: [PATCH 38/85] stopped returning 429 for get partition key ranges, since there isn't anything for us to do --- .../CrossPartitionRangePageAsyncEnumerator.cs | 34 ++++++++++++++----- .../src/Pagination/DocumentContainer.cs | 14 +------- .../PartitionRangeEnumeratorTests.cs | 20 +++++++---- 3 files changed, 41 insertions(+), 27 deletions(-) diff --git a/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionRangePageAsyncEnumerator.cs b/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionRangePageAsyncEnumerator.cs index 47ac66f0fc..459e6a1f14 100644 --- a/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionRangePageAsyncEnumerator.cs +++ b/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionRangePageAsyncEnumerator.cs @@ -76,13 +76,19 @@ public CrossPartitionRangePageAsyncEnumerator( public async ValueTask MoveNextAsync() { PriorityQueue> enumerators = await this.lazyEnumerators.GetValueAsync(cancellationToken: default); - PartitionRangePageAsyncEnumerator currentPaginator = enumerators.Dequeue(); - bool movedNext = await currentPaginator.MoveNextAsync(); - if (!movedNext) + if (enumerators.Count == 0) { return false; } + PartitionRangePageAsyncEnumerator currentPaginator = enumerators.Dequeue(); + if (!await currentPaginator.MoveNextAsync()) + { + // Current enumerator is empty, + // so recursively retry on the next enumerator. + return await this.MoveNextAsync(); + } + if (currentPaginator.Current.Failed) { // Check if it's a retryable exception. @@ -116,7 +122,10 @@ public async ValueTask MoveNextAsync() } } - enumerators.Enqueue(currentPaginator); + if (currentPaginator.State != null) + { + enumerators.Enqueue(currentPaginator); + } TryCatch backendPage = currentPaginator.Current; if (backendPage.Failed) @@ -125,13 +134,22 @@ public async ValueTask MoveNextAsync() return true; } - List<(PartitionKeyRange, TState)> feedRangeAndStates = new List<(PartitionKeyRange, TState)>(enumerators.Count); - foreach (PartitionRangePageAsyncEnumerator enumerator in enumerators) + CrossPartitionState crossPartitionState; + if (enumerators.Count == 0) + { + crossPartitionState = null; + } + else { - feedRangeAndStates.Add((enumerator.Range, enumerator.State)); + List<(PartitionKeyRange, TState)> feedRangeAndStates = new List<(PartitionKeyRange, TState)>(enumerators.Count); + foreach (PartitionRangePageAsyncEnumerator enumerator in enumerators) + { + feedRangeAndStates.Add((enumerator.Range, enumerator.State)); + } + + crossPartitionState = new CrossPartitionState(feedRangeAndStates); } - CrossPartitionState crossPartitionState = new CrossPartitionState(feedRangeAndStates); this.Current = TryCatch>.FromResult( new CrossPartitionPage(backendPage.Result, crossPartitionState)); return true; diff --git a/Microsoft.Azure.Cosmos/src/Pagination/DocumentContainer.cs b/Microsoft.Azure.Cosmos/src/Pagination/DocumentContainer.cs index afb1004436..e102e926f2 100644 --- a/Microsoft.Azure.Cosmos/src/Pagination/DocumentContainer.cs +++ b/Microsoft.Azure.Cosmos/src/Pagination/DocumentContainer.cs @@ -21,10 +21,6 @@ internal abstract class DocumentContainer : IFeedRangeProvider activityId: Guid.NewGuid().ToString(), requestCharge: default); - private static readonly Task>> ThrottleForGetRanges = Task.FromResult( - TryCatch>.FromException( - RequestRateTooLargeException)); - private static readonly Task> ThrottleForCreateItem = Task.FromResult( TryCatch.FromException( RequestRateTooLargeException)); @@ -54,17 +50,9 @@ protected abstract Task>> MonadicGetChildRangeI public Task>> MonadicGetChildRangeAsync( PartitionKeyRange partitionKeyRange, - CancellationToken cancellationToken) - { - if (this.ShouldReturn429()) - { - return ThrottleForGetRanges; - } - - return this.MonadicGetChildRangeImplementationAsync( + CancellationToken cancellationToken) => this.MonadicGetChildRangeImplementationAsync( partitionKeyRange, cancellationToken); - } public Task> GetChildRangeAsync( PartitionKeyRange partitionKeyRange, diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/PartitionRangeEnumeratorTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/PartitionRangeEnumeratorTests.cs index e365daca8e..2f2e005468 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/PartitionRangeEnumeratorTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/PartitionRangeEnumeratorTests.cs @@ -185,16 +185,24 @@ public async Task CreateDocumentContainerAsync( for (int i = 0; i < numItems; i++) { - // Insert an item - CosmosObject item = CosmosObject.Parse($"{{\"pk\" : {i} }}"); - while (true) + try { - TryCatch monadicCreateRecord = await inMemoryCollection.MonadicCreateItemAsync(item, cancellationToken: default); - if (monadicCreateRecord.Succeeded) + // Insert an item + CosmosObject item = CosmosObject.Parse($"{{\"pk\" : {i} }}"); + //await inMemoryCollection.CreateItemAsync(item, cancellationToken: default); + while (true) { - break; + TryCatch monadicCreateRecord = await inMemoryCollection.MonadicCreateItemAsync(item, cancellationToken: default); + if (monadicCreateRecord.Succeeded) + { + break; + } } } + catch(Exception ex) + { + throw ex; + } } return inMemoryCollection; From 6365de5d879159cf6b00d1fe2839514d7b379353 Mon Sep 17 00:00:00 2001 From: Brandon Chong Date: Mon, 13 Jul 2020 21:43:48 -0700 Subject: [PATCH 39/85] updated tests --- .../DocumentContainerTests.cs} | 185 +++++++++--------- .../{ => Pagination}/InMemoryCollection.cs | 16 +- .../Pagination/InMemoryCollectionTests.cs | 18 ++ 3 files changed, 112 insertions(+), 107 deletions(-) rename Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/{InMemoryCollectionTests.cs => Pagination/DocumentContainerTests.cs} (51%) rename Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/{ => Pagination}/InMemoryCollection.cs (97%) create mode 100644 Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/InMemoryCollectionTests.cs diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/InMemoryCollectionTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/DocumentContainerTests.cs similarity index 51% rename from Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/InMemoryCollectionTests.cs rename to Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/DocumentContainerTests.cs index 0da6880823..e29b3a6e50 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/InMemoryCollectionTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/DocumentContainerTests.cs @@ -2,7 +2,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. //------------------------------------------------------------ -namespace Microsoft.Azure.Cosmos.Tests +namespace Microsoft.Azure.Cosmos.Tests.Pagination { using System; using System.Collections.Generic; @@ -16,32 +16,54 @@ namespace Microsoft.Azure.Cosmos.Tests using Microsoft.VisualStudio.TestTools.UnitTesting; [TestClass] - public class InMemoryCollectionTests + public abstract class DocumentContainerTests { + private static readonly PartitionKeyDefinition PartitionKeyDefinition = new PartitionKeyDefinition() + { + Paths = new System.Collections.ObjectModel.Collection() + { + "/pk" + }, + Kind = PartitionKind.Hash, + Version = PartitionKeyDefinitionVersion.V2, + }; + + internal abstract DocumentContainer CreateDocumentContainer( + PartitionKeyDefinition partitionKeyDefinition, + DocumentContainer.FailureConfigs failureConfigs = default); + [TestMethod] - public async Task TestCrudAsync() + public async Task TestGetFeedRanges() { - PartitionKeyDefinition partitionKeyDefinition = new PartitionKeyDefinition() + DocumentContainer documentContainer = this.CreateDocumentContainer(PartitionKeyDefinition); + { - Paths = new System.Collections.ObjectModel.Collection() - { - "/pk" - }, - Kind = PartitionKind.Hash, - Version = PartitionKeyDefinitionVersion.V2, - }; + List ranges = await documentContainer.GetFeedRangesAsync(cancellationToken: default); + Assert.AreEqual(expected: 1, ranges.Count); + } - InMemoryCollection inMemoryCollection = new InMemoryCollection(partitionKeyDefinition); + await documentContainer.SplitAsync(partitionKeyRangeId: 0, cancellationToken: default); + + { + List ranges = await documentContainer.GetFeedRangesAsync(cancellationToken: default); + Assert.AreEqual(expected: 2, ranges.Count); + } + } + + [TestMethod] + public async Task TestCrudAsync() + { + DocumentContainer documentContainer = this.CreateDocumentContainer(PartitionKeyDefinition); // Insert an item CosmosObject item = CosmosObject.Parse("{\"pk\" : 42 }"); - Record record = await inMemoryCollection.CreateItemAsync(item, cancellationToken: default); + Record record = await documentContainer.CreateItemAsync(item, cancellationToken: default); Assert.IsNotNull(record); Assert.AreNotEqual(Guid.Empty, record.Identifier); Assert.AreEqual(1, record.ResourceIdentifier); // Try to read it back - Record readRecord = await inMemoryCollection.ReadItemAsync( + Record readRecord = await documentContainer.ReadItemAsync( partitionKey: CosmosNumber64.Create(42), record.Identifier, cancellationToken: default); @@ -52,28 +74,18 @@ public async Task TestCrudAsync() [TestMethod] public async Task TestPartitionKeyAsync() { - PartitionKeyDefinition partitionKeyDefinition = new PartitionKeyDefinition() - { - Paths = new System.Collections.ObjectModel.Collection() - { - "/pk" - }, - Kind = PartitionKind.Hash, - Version = PartitionKeyDefinitionVersion.V2, - }; - - InMemoryCollection inMemoryCollection = new InMemoryCollection(partitionKeyDefinition); + DocumentContainer documentContainer = this.CreateDocumentContainer(PartitionKeyDefinition); // Insert an item CosmosObject item1 = CosmosObject.Parse("{\"pk\" : 42 }"); - Record record1 = await inMemoryCollection.CreateItemAsync(item1, cancellationToken: default); + Record record1 = await documentContainer.CreateItemAsync(item1, cancellationToken: default); // Insert into another partition key CosmosObject item2 = CosmosObject.Parse("{\"pk\" : 1337 }"); - Record record2 = await inMemoryCollection.CreateItemAsync(item2, cancellationToken: default); + Record record2 = await documentContainer.CreateItemAsync(item2, cancellationToken: default); // Try to read back an id with wrong pk - TryCatch monadicReadItem = await inMemoryCollection.MonadicReadItemAsync( + TryCatch monadicReadItem = await documentContainer.MonadicReadItemAsync( partitionKey: item1["pk"], record2.Identifier, cancellationToken: default); @@ -83,59 +95,39 @@ public async Task TestPartitionKeyAsync() [TestMethod] public async Task TestUndefinedPartitionKeyAsync() { - PartitionKeyDefinition partitionKeyDefinition = new PartitionKeyDefinition() - { - Paths = new System.Collections.ObjectModel.Collection() - { - "/pk" - }, - Kind = PartitionKind.Hash, - Version = PartitionKeyDefinitionVersion.V2, - }; - - InMemoryCollection inMemoryCollection = new InMemoryCollection(partitionKeyDefinition); + DocumentContainer documentContainer = this.CreateDocumentContainer(PartitionKeyDefinition); // Insert an item CosmosObject item = CosmosObject.Parse("{}"); - Record record = await inMemoryCollection.CreateItemAsync(item, cancellationToken: default); + Record record = await documentContainer.CreateItemAsync(item, cancellationToken: default); - // Try to read back an id with wrong pk - TryCatch monadicReadItem = await inMemoryCollection.MonadicReadItemAsync( + // Try to read back an id with null undefined partition key + TryCatch monadicReadItem = await documentContainer.MonadicReadItemAsync( partitionKey: null, record.Identifier, cancellationToken: default); - Assert.IsFalse(monadicReadItem.Succeeded); + Assert.IsTrue(monadicReadItem.Succeeded); } [TestMethod] public async Task TestSplitAsync() { - PartitionKeyDefinition partitionKeyDefinition = new PartitionKeyDefinition() - { - Paths = new System.Collections.ObjectModel.Collection() - { - "/pk" - }, - Kind = PartitionKind.Hash, - Version = PartitionKeyDefinitionVersion.V2, - }; + DocumentContainer documentContainer = this.CreateDocumentContainer(PartitionKeyDefinition); - InMemoryCollection inMemoryCollection = new InMemoryCollection(partitionKeyDefinition); - - Assert.AreEqual(1, (await inMemoryCollection.GetFeedRangesAsync(cancellationToken: default)).Count); + Assert.AreEqual(1, (await documentContainer.GetFeedRangesAsync(cancellationToken: default)).Count); int numItemsToInsert = 10; for (int i = 0; i < numItemsToInsert; i++) { // Insert an item CosmosObject item = CosmosObject.Parse($"{{\"pk\" : {i} }}"); - await inMemoryCollection.CreateItemAsync(item, cancellationToken: default); + await documentContainer.CreateItemAsync(item, cancellationToken: default); } - await inMemoryCollection.SplitAsync(partitionKeyRangeId: 0, cancellationToken: default); + await documentContainer.SplitAsync(partitionKeyRangeId: 0, cancellationToken: default); - Assert.AreEqual(2, (await inMemoryCollection.GetFeedRangesAsync(cancellationToken: default)).Count); - List ranges = await inMemoryCollection.GetChildRangeAsync( + Assert.AreEqual(2, (await documentContainer.GetFeedRangesAsync(cancellationToken: default)).Count); + List ranges = await documentContainer.GetChildRangeAsync( new PartitionKeyRange() { Id = "0" @@ -147,7 +139,7 @@ public async Task TestSplitAsync() async Task AssertChildPartitionAsync(int partitionKeyRangeId) { - DocumentContainerPage readFeedPage = await inMemoryCollection.ReadFeedAsync( + DocumentContainerPage readFeedPage = await documentContainer.ReadFeedAsync( partitionKeyRangeId: partitionKeyRangeId, resourceIdentifier: 0, pageSize: 100, @@ -182,59 +174,64 @@ public async Task TestMultiSplitAsync() Version = PartitionKeyDefinitionVersion.V2, }; - InMemoryCollection inMemoryCollection = new InMemoryCollection(partitionKeyDefinition); + DocumentContainer documentContainer = this.CreateDocumentContainer(partitionKeyDefinition); int numItemsToInsert = 10; for (int i = 0; i < numItemsToInsert; i++) { // Insert an item CosmosObject item = CosmosObject.Parse($"{{\"pk\" : {i} }}"); - await inMemoryCollection.CreateItemAsync(item, cancellationToken: default); + await documentContainer.CreateItemAsync(item, cancellationToken: default); } - Assert.AreEqual(1, (await inMemoryCollection.GetFeedRangesAsync(cancellationToken: default)).Count); - - await inMemoryCollection.SplitAsync(partitionKeyRangeId: 0, cancellationToken: default); - - Assert.AreEqual(2, (await inMemoryCollection.GetFeedRangesAsync(cancellationToken: default)).Count); - List ranges = await inMemoryCollection.GetChildRangeAsync( - new PartitionKeyRange() - { - Id = "0" - }, - cancellationToken: default); - Assert.AreEqual(2, ranges.Count); - Assert.AreEqual(1, int.Parse(ranges[0].Id)); - Assert.AreEqual(2, int.Parse(ranges[1].Id)); + Assert.AreEqual(1, (await documentContainer.GetFeedRangesAsync(cancellationToken: default)).Count); - await inMemoryCollection.SplitAsync(partitionKeyRangeId: 1, cancellationToken: default); - await inMemoryCollection.SplitAsync(partitionKeyRangeId: 2, cancellationToken: default); + await documentContainer.SplitAsync(partitionKeyRangeId: 0, cancellationToken: default); + { + Assert.AreEqual(2, (await documentContainer.GetFeedRangesAsync(cancellationToken: default)).Count); + List ranges = await documentContainer.GetChildRangeAsync( + new PartitionKeyRange() + { + Id = "0" + }, + cancellationToken: default); + Assert.AreEqual(2, ranges.Count); + Assert.AreEqual(1, int.Parse(ranges[0].Id)); + Assert.AreEqual(2, int.Parse(ranges[1].Id)); + } + + await documentContainer.SplitAsync(partitionKeyRangeId: 1, cancellationToken: default); + await documentContainer.SplitAsync(partitionKeyRangeId: 2, cancellationToken: default); - Assert.AreEqual(4, (await inMemoryCollection.GetFeedRangesAsync(cancellationToken: default)).Count); - List ranges1 = await inMemoryCollection.GetChildRangeAsync( - new PartitionKeyRange() - { - Id = "1" - }, - cancellationToken: default); - Assert.AreEqual(2, ranges.Count); - Assert.AreEqual(3, int.Parse(ranges[0].Id)); - Assert.AreEqual(4, int.Parse(ranges[1].Id)); + { + Assert.AreEqual(4, (await documentContainer.GetFeedRangesAsync(cancellationToken: default)).Count); + List ranges = await documentContainer.GetChildRangeAsync( + new PartitionKeyRange() + { + Id = "1" + }, + cancellationToken: default); + Assert.AreEqual(2, ranges.Count); + Assert.AreEqual(3, int.Parse(ranges[0].Id)); + Assert.AreEqual(4, int.Parse(ranges[1].Id)); + } - List ranges2 = await inMemoryCollection.GetChildRangeAsync( + { + List ranges = await documentContainer.GetChildRangeAsync( new PartitionKeyRange() { Id = "2" }, cancellationToken: default); - Assert.AreEqual(2, ranges.Count); - Assert.AreEqual(5, int.Parse(ranges[0].Id)); - Assert.AreEqual(6, int.Parse(ranges[1].Id)); - + Assert.AreEqual(2, ranges.Count); + Assert.AreEqual(5, int.Parse(ranges[0].Id)); + Assert.AreEqual(6, int.Parse(ranges[1].Id)); + } + async Task AssertChildPartitionAsync(int partitionKeyRangeId) { - DocumentContainerPage page = await inMemoryCollection.ReadFeedAsync( + DocumentContainerPage page = await documentContainer.ReadFeedAsync( partitionKeyRangeId: partitionKeyRangeId, resourceIdentifier: 0, pageSize: 100, diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/InMemoryCollection.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/InMemoryCollection.cs similarity index 97% rename from Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/InMemoryCollection.cs rename to Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/InMemoryCollection.cs index 435f1c138d..0e2dde9e0f 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/InMemoryCollection.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/InMemoryCollection.cs @@ -2,7 +2,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. //------------------------------------------------------------ -namespace Microsoft.Azure.Cosmos.Tests +namespace Microsoft.Azure.Cosmos.Tests.Pagination { using System; using System.Collections; @@ -90,18 +90,8 @@ PartitionKeyRange CreateRangeFromId(int id) List childRanges = new List() { - new PartitionKeyRange() - { - Id = children.left.ToString(), - MinInclusive = children.left.ToString(), - MaxExclusive = children.left.ToString(), - }, - new PartitionKeyRange() - { - Id = children.right.ToString(), - MinInclusive = children.right.ToString(), - MaxExclusive = children.right.ToString(), - } + CreateRangeFromId(children.left), + CreateRangeFromId(children.right), }; return Task.FromResult(TryCatch>.FromResult(childRanges)); diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/InMemoryCollectionTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/InMemoryCollectionTests.cs new file mode 100644 index 0000000000..543e73f888 --- /dev/null +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/InMemoryCollectionTests.cs @@ -0,0 +1,18 @@ +//------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +//------------------------------------------------------------ + +namespace Microsoft.Azure.Cosmos.Tests.Pagination +{ + using Microsoft.Azure.Cosmos.Pagination; + using Microsoft.Azure.Documents; + using Microsoft.VisualStudio.TestTools.UnitTesting; + + [TestClass] + public sealed class InMemoryCollectionTests : DocumentContainerTests + { + internal override DocumentContainer CreateDocumentContainer( + PartitionKeyDefinition partitionKeyDefinition, + DocumentContainer.FailureConfigs failureConfigs = null) => new InMemoryCollection(partitionKeyDefinition, failureConfigs); + } +} From 6737a04725cce82221f8e523f58262ba514fd9ab Mon Sep 17 00:00:00 2001 From: Brandon Chong Date: Tue, 14 Jul 2020 15:23:25 -0700 Subject: [PATCH 40/85] added document container interface --- .../src/Pagination/DocumentContainer.cs | 2 +- .../src/Pagination/IDocumentContainer.cs | 53 +++++++++++++++++++ ...sPartitionPartitionRangeEnumeratorTests.cs | 8 +-- ...cumentContainerPartitionRangeEnumerator.cs | 4 +- .../Pagination/DocumentContainerTests.cs | 14 ++--- .../Pagination/InMemoryCollectionTests.cs | 2 +- .../PartitionRangeEnumeratorTests.cs | 16 +++--- ...ePartitionPartitionRangeEnumeratorTests.cs | 6 +-- 8 files changed, 79 insertions(+), 26 deletions(-) create mode 100644 Microsoft.Azure.Cosmos/src/Pagination/IDocumentContainer.cs diff --git a/Microsoft.Azure.Cosmos/src/Pagination/DocumentContainer.cs b/Microsoft.Azure.Cosmos/src/Pagination/DocumentContainer.cs index e102e926f2..7910489848 100644 --- a/Microsoft.Azure.Cosmos/src/Pagination/DocumentContainer.cs +++ b/Microsoft.Azure.Cosmos/src/Pagination/DocumentContainer.cs @@ -12,7 +12,7 @@ namespace Microsoft.Azure.Cosmos.Pagination using Microsoft.Azure.Cosmos.Query.Core.Monads; using Microsoft.Azure.Documents; - internal abstract class DocumentContainer : IFeedRangeProvider + internal abstract class DocumentContainer : IDocumentContainer { private static readonly CosmosException RequestRateTooLargeException = new CosmosException( message: "Request Rate Too Large", diff --git a/Microsoft.Azure.Cosmos/src/Pagination/IDocumentContainer.cs b/Microsoft.Azure.Cosmos/src/Pagination/IDocumentContainer.cs new file mode 100644 index 0000000000..72d10fde1f --- /dev/null +++ b/Microsoft.Azure.Cosmos/src/Pagination/IDocumentContainer.cs @@ -0,0 +1,53 @@ +// ------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +// ------------------------------------------------------------ + +namespace Microsoft.Azure.Cosmos.Pagination +{ + using System; + using System.Threading; + using System.Threading.Tasks; + using Microsoft.Azure.Cosmos.CosmosElements; + using Microsoft.Azure.Cosmos.Query.Core.Monads; + + internal interface IDocumentContainer : IFeedRangeProvider + { + Task> MonadicCreateItemAsync( + CosmosObject payload, + CancellationToken cancellationToken); + + Task CreateItemAsync( + CosmosObject payload, + CancellationToken cancellationToken); + + Task> MonadicReadItemAsync( + CosmosElement partitionKey, + Guid identifer, + CancellationToken cancellationToken); + + Task ReadItemAsync( + CosmosElement partitionKey, + Guid identifier, + CancellationToken cancellationToken); + + Task> MonadicReadFeedAsync( + int partitionKeyRangeId, + long resourceIdentifer, + int pageSize, + CancellationToken cancellationToken); + + Task ReadFeedAsync( + int partitionKeyRangeId, + long resourceIdentifier, + int pageSize, + CancellationToken cancellationToken); + + Task MonadicSplitAsync( + int partitionKeyRangeId, + CancellationToken cancellationToken); + + Task SplitAsync( + int partitionKeyRangeId, + CancellationToken cancellationToken); + } +} diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/CrossPartitionPartitionRangeEnumeratorTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/CrossPartitionPartitionRangeEnumeratorTests.cs index e9a68b8a2a..aeb838cf02 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/CrossPartitionPartitionRangeEnumeratorTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/CrossPartitionPartitionRangeEnumeratorTests.cs @@ -76,7 +76,7 @@ public Implementation() public async Task TestSplitWithResumeContinuationAsync() { int numItems = 1000; - DocumentContainer inMemoryCollection = await this.CreateDocumentContainerAsync(numItems); + IDocumentContainer inMemoryCollection = await this.CreateDocumentContainerAsync(numItems); IAsyncEnumerator>> enumerator = this.CreateEnumerator(inMemoryCollection); (HashSet firstDrainResults, CrossPartitionState state) = await this.PartialDrainAsync(enumerator, numIterations: 3); @@ -102,7 +102,7 @@ public async Task TestSplitWithResumeContinuationAsync() public async Task TestSplitWithDuringDrainAsync() { int numItems = 1000; - DocumentContainer inMemoryCollection = await this.CreateDocumentContainerAsync(numItems); + IDocumentContainer inMemoryCollection = await this.CreateDocumentContainerAsync(numItems); IAsyncEnumerable>> enumerable = this.CreateEnumerable(inMemoryCollection); HashSet identifiers = new HashSet(); @@ -131,7 +131,7 @@ public async Task TestSplitWithDuringDrainAsync() } public override IAsyncEnumerable>> CreateEnumerable( - DocumentContainer inMemoryCollection, + IDocumentContainer inMemoryCollection, CrossPartitionState state = null) { PartitionRangePageAsyncEnumerator createEnumerator( @@ -150,7 +150,7 @@ PartitionRangePageAsyncEnumerator } public override IAsyncEnumerator>> CreateEnumerator( - DocumentContainer inMemoryCollection, + IDocumentContainer inMemoryCollection, CrossPartitionState state = null) { PartitionRangePageAsyncEnumerator createEnumerator( diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/DocumentContainerPartitionRangeEnumerator.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/DocumentContainerPartitionRangeEnumerator.cs index 758a52d179..09a0d18a85 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/DocumentContainerPartitionRangeEnumerator.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/DocumentContainerPartitionRangeEnumerator.cs @@ -14,12 +14,12 @@ namespace Microsoft.Azure.Cosmos.Tests.Pagination internal sealed class DocumentContainerPartitionRangeEnumerator : PartitionRangePageAsyncEnumerator { - private readonly DocumentContainer documentContainer; + private readonly IDocumentContainer documentContainer; private readonly int pageSize; private readonly int partitionKeyRangeId; public DocumentContainerPartitionRangeEnumerator( - DocumentContainer documentContainer, + IDocumentContainer documentContainer, int partitionKeyRangeId, int pageSize, DocumentContainerState state = null) diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/DocumentContainerTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/DocumentContainerTests.cs index e29b3a6e50..5d9e48f918 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/DocumentContainerTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/DocumentContainerTests.cs @@ -28,14 +28,14 @@ public abstract class DocumentContainerTests Version = PartitionKeyDefinitionVersion.V2, }; - internal abstract DocumentContainer CreateDocumentContainer( + internal abstract IDocumentContainer CreateDocumentContainer( PartitionKeyDefinition partitionKeyDefinition, DocumentContainer.FailureConfigs failureConfigs = default); [TestMethod] public async Task TestGetFeedRanges() { - DocumentContainer documentContainer = this.CreateDocumentContainer(PartitionKeyDefinition); + IDocumentContainer documentContainer = this.CreateDocumentContainer(PartitionKeyDefinition); { List ranges = await documentContainer.GetFeedRangesAsync(cancellationToken: default); @@ -53,7 +53,7 @@ public async Task TestGetFeedRanges() [TestMethod] public async Task TestCrudAsync() { - DocumentContainer documentContainer = this.CreateDocumentContainer(PartitionKeyDefinition); + IDocumentContainer documentContainer = this.CreateDocumentContainer(PartitionKeyDefinition); // Insert an item CosmosObject item = CosmosObject.Parse("{\"pk\" : 42 }"); @@ -74,7 +74,7 @@ public async Task TestCrudAsync() [TestMethod] public async Task TestPartitionKeyAsync() { - DocumentContainer documentContainer = this.CreateDocumentContainer(PartitionKeyDefinition); + IDocumentContainer documentContainer = this.CreateDocumentContainer(PartitionKeyDefinition); // Insert an item CosmosObject item1 = CosmosObject.Parse("{\"pk\" : 42 }"); @@ -95,7 +95,7 @@ public async Task TestPartitionKeyAsync() [TestMethod] public async Task TestUndefinedPartitionKeyAsync() { - DocumentContainer documentContainer = this.CreateDocumentContainer(PartitionKeyDefinition); + IDocumentContainer documentContainer = this.CreateDocumentContainer(PartitionKeyDefinition); // Insert an item CosmosObject item = CosmosObject.Parse("{}"); @@ -112,7 +112,7 @@ public async Task TestUndefinedPartitionKeyAsync() [TestMethod] public async Task TestSplitAsync() { - DocumentContainer documentContainer = this.CreateDocumentContainer(PartitionKeyDefinition); + IDocumentContainer documentContainer = this.CreateDocumentContainer(PartitionKeyDefinition); Assert.AreEqual(1, (await documentContainer.GetFeedRangesAsync(cancellationToken: default)).Count); @@ -174,7 +174,7 @@ public async Task TestMultiSplitAsync() Version = PartitionKeyDefinitionVersion.V2, }; - DocumentContainer documentContainer = this.CreateDocumentContainer(partitionKeyDefinition); + IDocumentContainer documentContainer = this.CreateDocumentContainer(partitionKeyDefinition); int numItemsToInsert = 10; for (int i = 0; i < numItemsToInsert; i++) diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/InMemoryCollectionTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/InMemoryCollectionTests.cs index 543e73f888..46099a3090 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/InMemoryCollectionTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/InMemoryCollectionTests.cs @@ -11,7 +11,7 @@ namespace Microsoft.Azure.Cosmos.Tests.Pagination [TestClass] public sealed class InMemoryCollectionTests : DocumentContainerTests { - internal override DocumentContainer CreateDocumentContainer( + internal override IDocumentContainer CreateDocumentContainer( PartitionKeyDefinition partitionKeyDefinition, DocumentContainer.FailureConfigs failureConfigs = null) => new InMemoryCollection(partitionKeyDefinition, failureConfigs); } diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/PartitionRangeEnumeratorTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/PartitionRangeEnumeratorTests.cs index 2f2e005468..e119ed7860 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/PartitionRangeEnumeratorTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/PartitionRangeEnumeratorTests.cs @@ -23,7 +23,7 @@ protected PartitionRangeEnumeratorTests(bool singlePartition) public async Task TestDrainFullyAsync() { int numItems = 1000; - DocumentContainer inMemoryCollection = await this.CreateDocumentContainerAsync(numItems); + IDocumentContainer inMemoryCollection = await this.CreateDocumentContainerAsync(numItems); IAsyncEnumerable> enumerable = this.CreateEnumerable(inMemoryCollection); HashSet identifiers = await this.DrainFullyAsync(enumerable); Assert.AreEqual(numItems, identifiers.Count); @@ -33,7 +33,7 @@ public async Task TestDrainFullyAsync() public async Task TestResumingFromStateAsync() { int numItems = 1000; - DocumentContainer inMemoryCollection = await this.CreateDocumentContainerAsync(numItems); + IDocumentContainer inMemoryCollection = await this.CreateDocumentContainerAsync(numItems); IAsyncEnumerator> enumerator = this.CreateEnumerator(inMemoryCollection); (HashSet firstDrainResults, TState state) = await this.PartialDrainAsync(enumerator, numIterations: 3); @@ -48,7 +48,7 @@ public async Task TestResumingFromStateAsync() public async Task Test429sAsync() { int numItems = 100; - DocumentContainer inMemoryCollection = await this.CreateDocumentContainerAsync( + IDocumentContainer inMemoryCollection = await this.CreateDocumentContainerAsync( numItems, new DocumentContainer.FailureConfigs( inject429s: true, @@ -89,7 +89,7 @@ public async Task Test429sAsync() public async Task Test429sWithContinuationsAsync() { int numItems = 100; - DocumentContainer inMemoryCollection = await this.CreateDocumentContainerAsync( + IDocumentContainer inMemoryCollection = await this.CreateDocumentContainerAsync( numItems, new DocumentContainer.FailureConfigs( inject429s: true, @@ -138,7 +138,7 @@ public async Task Test429sWithContinuationsAsync() public async Task TestEmptyPages() { int numItems = 100; - DocumentContainer inMemoryCollection = await this.CreateDocumentContainerAsync( + IDocumentContainer inMemoryCollection = await this.CreateDocumentContainerAsync( numItems, new DocumentContainer.FailureConfigs( inject429s: false, @@ -150,11 +150,11 @@ public async Task TestEmptyPages() public abstract IReadOnlyList GetRecordsFromPage(TPage page); - public abstract IAsyncEnumerable> CreateEnumerable(DocumentContainer documentContainer, TState state = null); + public abstract IAsyncEnumerable> CreateEnumerable(IDocumentContainer documentContainer, TState state = null); - public abstract IAsyncEnumerator> CreateEnumerator(DocumentContainer documentContainer, TState state = null); + public abstract IAsyncEnumerator> CreateEnumerator(IDocumentContainer documentContainer, TState state = null); - public async Task CreateDocumentContainerAsync( + public async Task CreateDocumentContainerAsync( int numItems, DocumentContainer.FailureConfigs failureConfigs = default) { diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/SinglePartitionPartitionRangeEnumeratorTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/SinglePartitionPartitionRangeEnumeratorTests.cs index 1a098b7980..ff2fefb271 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/SinglePartitionPartitionRangeEnumeratorTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/SinglePartitionPartitionRangeEnumeratorTests.cs @@ -70,7 +70,7 @@ public Implementation() public async Task TestSplitAsync() { int numItems = 100; - DocumentContainer inMemoryCollection = await this.CreateDocumentContainerAsync(numItems); + IDocumentContainer inMemoryCollection = await this.CreateDocumentContainerAsync(numItems); DocumentContainerPartitionRangeEnumerator enumerator = new DocumentContainerPartitionRangeEnumerator( inMemoryCollection, partitionKeyRangeId: 0, @@ -111,7 +111,7 @@ public override IReadOnlyList GetRecordsFromPage(DocumentContainerPage p } public override IAsyncEnumerable> CreateEnumerable( - DocumentContainer documentContainer, + IDocumentContainer documentContainer, DocumentContainerState state = null) => new PartitionRangePageAsyncEnumerable( range: new PartitionKeyRange() { Id = "0" }, state: state, @@ -122,7 +122,7 @@ public override IAsyncEnumerable> CreateEnumerab state: state)); public override IAsyncEnumerator> CreateEnumerator( - DocumentContainer inMemoryCollection, + IDocumentContainer inMemoryCollection, DocumentContainerState state = null) => new DocumentContainerPartitionRangeEnumerator( inMemoryCollection, partitionKeyRangeId: 0, From da3fb94af0b7093c5a4651d6daca892c7d600997 Mon Sep 17 00:00:00 2001 From: Brandon Chong Date: Tue, 14 Jul 2020 17:12:18 -0700 Subject: [PATCH 41/85] added end to end tests --- .../Query/Pipeline/FullPipelineTests.cs | 190 ++++++++++++++++++ 1 file changed, 190 insertions(+) create mode 100644 Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/FullPipelineTests.cs 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 new file mode 100644 index 0000000000..00a866dc9d --- /dev/null +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/FullPipelineTests.cs @@ -0,0 +1,190 @@ +//------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +//------------------------------------------------------------ + +namespace Microsoft.Azure.Cosmos.Tests.Query.Pipeline +{ + using System.Collections.Generic; + using System.Collections.ObjectModel; + using System.Linq; + using System.Threading.Tasks; + using Microsoft.Azure.Cosmos.CosmosElements; + using Microsoft.Azure.Cosmos.Pagination; + 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.QueryPlan; + using Microsoft.Azure.Cosmos.Tests.Pagination; + using Microsoft.Azure.Documents; + using Microsoft.VisualStudio.TestTools.UnitTesting; + + [TestClass] + public class FullPipelineTests + { + private static readonly Dictionary DefaultQueryEngineConfiguration = new Dictionary() + { + {"maxSqlQueryInputLength", 30720}, + {"maxJoinsPerSqlQuery", 5}, + {"maxLogicalAndPerSqlQuery", 200}, + {"maxLogicalOrPerSqlQuery", 200}, + {"maxUdfRefPerSqlQuery", 2}, + {"maxInExpressionItemsCount", 8000}, + {"queryMaxInMemorySortDocumentCount", 500}, + {"maxQueryRequestTimeoutFraction", 0.90}, + {"sqlAllowNonFiniteNumbers", false}, + {"sqlAllowAggregateFunctions", true}, + {"sqlAllowSubQuery", true}, + {"sqlAllowScalarSubQuery", false}, + {"allowNewKeywords", true}, + {"sqlAllowLike", false}, + {"sqlAllowGroupByClause", false}, + {"maxSpatialQueryCells", 12}, + {"spatialMaxGeometryPointCount", 256}, + {"sqlDisableQueryILOptimization", false}, + {"sqlDisableFilterPlanOptimization", false} + }; + + private static readonly QueryPartitionProvider queryPartitionProvider = new QueryPartitionProvider(DefaultQueryEngineConfiguration); + private static readonly PartitionKeyDefinition partitionKeyDefinition = new PartitionKeyDefinition() + { + Paths = new Collection() + { + "/pk" + }, + Kind = PartitionKind.Hash, + Version = PartitionKeyDefinitionVersion.V2, + }; + + [TestMethod] + public async Task SelectStar() + { + List documents = new List(); + for (int i = 0; i < 250; i++) + { + documents.Add(CosmosObject.Parse($"{{\"pk\" : {i} }}")); + } + + List documentsQueried = await ExecuteQueryAsync( + query: "SELECT * FROM c", + documents: documents); + + Assert.AreEqual(expected: documents.Count, actual: documentsQueried.Count); + } + + private static async Task> ExecuteQueryAsync( + string query, + IReadOnlyList documents) + { + IDocumentContainer documentContainer = await CreateDocumentContainerAsync(documents); + + List resultsFromDrainWithoutState = await DrainWithoutStateAsync(query, documentContainer); + List resultsFromDrainWithState = await DrainWithStateAsync(query, documentContainer); + + Assert.IsTrue(resultsFromDrainWithoutState.SequenceEqual(resultsFromDrainWithState)); + + return resultsFromDrainWithoutState; + } + + private static async Task> DrainWithoutStateAsync(string query, IDocumentContainer documentContainer) + { + IQueryPipelineStage pipelineStage = CreatePipeline(documentContainer, query); + + List elements = new List(); + while (await pipelineStage.MoveNextAsync()) + { + TryCatch tryGetQueryPage = pipelineStage.Current; + tryGetQueryPage.ThrowIfFailed(); + + elements.AddRange(tryGetQueryPage.Result.Documents); + } + + return elements; + } + + private static async Task> DrainWithStateAsync(string query, IDocumentContainer documentContainer) + { + IQueryPipelineStage pipelineStage; + CosmosElement state = null; + + List elements = new List(); + do + { + pipelineStage = CreatePipeline(documentContainer, query, state); + + await pipelineStage.MoveNextAsync(); + TryCatch tryGetQueryPage = pipelineStage.Current; + + tryGetQueryPage.ThrowIfFailed(); + + elements.AddRange(tryGetQueryPage.Result.Documents); + state = tryGetQueryPage.Result.State?.Value; + } + while (state != null); + + return elements; + } + + private static async Task CreateDocumentContainerAsync( + IReadOnlyList documents, + DocumentContainer.FailureConfigs failureConfigs = null) + { + InMemoryCollection inMemoryCollection = new InMemoryCollection(partitionKeyDefinition, failureConfigs); + + await inMemoryCollection.SplitAsync(partitionKeyRangeId: 0, cancellationToken: default); + + await inMemoryCollection.SplitAsync(partitionKeyRangeId: 1, cancellationToken: default); + await inMemoryCollection.SplitAsync(partitionKeyRangeId: 2, cancellationToken: default); + + await inMemoryCollection.SplitAsync(partitionKeyRangeId: 3, cancellationToken: default); + await inMemoryCollection.SplitAsync(partitionKeyRangeId: 4, cancellationToken: default); + await inMemoryCollection.SplitAsync(partitionKeyRangeId: 5, cancellationToken: default); + await inMemoryCollection.SplitAsync(partitionKeyRangeId: 6, cancellationToken: default); + + foreach (CosmosObject document in documents) + { + while (true) + { + TryCatch monadicCreateRecord = await inMemoryCollection.MonadicCreateItemAsync( + document, + cancellationToken: default); + if (monadicCreateRecord.Succeeded) + { + break; + } + } + } + + return inMemoryCollection; + } + + private static IQueryPipelineStage CreatePipeline(IDocumentContainer documentContainer, string query, CosmosElement state = null) + { + TryCatch tryCreatePipeline = PipelineFactory.MonadicCreate( + ExecutionEnvironment.Compute, + documentContainer, + new SqlQuerySpec(query), + GetQueryPlan(query), + pageSize: 10, + requestContinuationToken: state); + + tryCreatePipeline.ThrowIfFailed(); + + return tryCreatePipeline.Result; + } + + private static QueryInfo GetQueryPlan(string query) + { + TryCatch info = queryPartitionProvider.TryGetPartitionedQueryExecutionInfoInternal( + new SqlQuerySpec(query), + partitionKeyDefinition, + requireFormattableOrderByQuery: true, + isContinuationExpected: false, + allowNonValueAggregateQuery: true, + hasLogicalPartitionKey: false); + + info.ThrowIfFailed(); + return info.Result.QueryInfo; + } + } +} From c8a5cfa6d944b461369c5f7d7dee45617c411333 Mon Sep 17 00:00:00 2001 From: Brandon Chong Date: Tue, 14 Jul 2020 18:05:18 -0700 Subject: [PATCH 42/85] got basic query working end to end --- .../src/Query/v3Query/QueryIterator.cs | 9 +++++++-- .../Query/SanityQueryTests.cs | 1 - 2 files changed, 7 insertions(+), 3 deletions(-) diff --git a/Microsoft.Azure.Cosmos/src/Query/v3Query/QueryIterator.cs b/Microsoft.Azure.Cosmos/src/Query/v3Query/QueryIterator.cs index 7bf4b2a312..7d6103685e 100644 --- a/Microsoft.Azure.Cosmos/src/Query/v3Query/QueryIterator.cs +++ b/Microsoft.Azure.Cosmos/src/Query/v3Query/QueryIterator.cs @@ -164,6 +164,11 @@ public override async Task ReadNextAsync(CancellationToken canc if (tryGetQueryPage.Succeeded) { + if (tryGetQueryPage.Result.State == null) + { + this.hasMoreResults = false; + } + return QueryResponse.CreateSuccess( result: tryGetQueryPage.Result.Documents, count: tryGetQueryPage.Result.Documents.Count, @@ -171,7 +176,7 @@ public override async Task ReadNextAsync(CancellationToken canc diagnostics: diagnostics, serializationOptions: this.cosmosSerializationFormatOptions, responseHeaders: new CosmosQueryResponseMessageHeaders( - (tryGetQueryPage.Result.State?.Value as CosmosString)?.Value, + tryGetQueryPage.Result.State?.Value.ToString(), tryGetQueryPage.Result.DisallowContinuationTokenMessage, this.cosmosQueryContext.ResourceTypeEnum, this.cosmosQueryContext.ContainerResourceId) @@ -211,7 +216,7 @@ public override async Task ReadNextAsync(CancellationToken canc } } - public override CosmosElement GetCosmosElementContinuationToken() => this.queryPipelineStage.Current.Result.State.Value; + public override CosmosElement GetCosmosElementContinuationToken() => this.queryPipelineStage.Current.Result.State?.Value; protected override void Dispose(bool disposing) { diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/Query/SanityQueryTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/Query/SanityQueryTests.cs index c1fd74e5f4..d96fddd75a 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/Query/SanityQueryTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/Query/SanityQueryTests.cs @@ -20,7 +20,6 @@ namespace Microsoft.Azure.Cosmos.EmulatorTests.Query [TestClass] public sealed class SanityQueryTests : QueryTestsBase { - [TestMethod] public async Task Sanity() { From 61ed00592a37d4a01612dd2f97822149986ac3b2 Mon Sep 17 00:00:00 2001 From: Brandon Chong Date: Fri, 17 Jul 2020 18:10:52 -0700 Subject: [PATCH 43/85] composition over inheritance --- .../src/Pagination/DocumentContainer.cs | 119 ++------------ .../src/Pagination/IDocumentContainer.cs | 22 +-- .../src/Pagination/IFeedRangeProvider.cs | 10 +- .../Pagination/IMonadicDocumentContainer.cs | 34 ++++ .../Pagination/IMonadicFeedRangeProvider.cs | 22 +++ .../Pagination/DocumentContainerTests.cs | 2 +- .../Pagination/FlakyDocumentContainer.cs | 145 ++++++++++++++++++ .../Pagination/InMemoryCollectionTests.cs | 5 +- ...moryCollection.cs => InMemoryContainer.cs} | 29 ++-- .../PartitionRangeEnumeratorTests.cs | 49 +++--- 10 files changed, 265 insertions(+), 172 deletions(-) create mode 100644 Microsoft.Azure.Cosmos/src/Pagination/IMonadicDocumentContainer.cs create mode 100644 Microsoft.Azure.Cosmos/src/Pagination/IMonadicFeedRangeProvider.cs create mode 100644 Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/FlakyDocumentContainer.cs rename Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/{InMemoryCollection.cs => InMemoryContainer.cs} (94%) diff --git a/Microsoft.Azure.Cosmos/src/Pagination/DocumentContainer.cs b/Microsoft.Azure.Cosmos/src/Pagination/DocumentContainer.cs index 7910489848..ae95d9b9be 100644 --- a/Microsoft.Azure.Cosmos/src/Pagination/DocumentContainer.cs +++ b/Microsoft.Azure.Cosmos/src/Pagination/DocumentContainer.cs @@ -12,45 +12,21 @@ namespace Microsoft.Azure.Cosmos.Pagination using Microsoft.Azure.Cosmos.Query.Core.Monads; using Microsoft.Azure.Documents; - internal abstract class DocumentContainer : IDocumentContainer + /// + /// Composes a and creates an . + /// + internal sealed class DocumentContainer : IDocumentContainer { - private static readonly CosmosException RequestRateTooLargeException = new CosmosException( - message: "Request Rate Too Large", - statusCode: (System.Net.HttpStatusCode)429, - subStatusCode: default, - activityId: Guid.NewGuid().ToString(), - requestCharge: default); + private readonly IMonadicDocumentContainer monadicDocumentContainer; - private static readonly Task> ThrottleForCreateItem = Task.FromResult( - TryCatch.FromException( - RequestRateTooLargeException)); - - private static readonly Task> ThrottleForFeedOperation = Task.FromResult( - TryCatch.FromException( - RequestRateTooLargeException)); - - private static readonly PartitionKeyRange FullRange = new PartitionKeyRange() - { - MinInclusive = Documents.Routing.PartitionKeyInternal.MinimumInclusiveEffectivePartitionKey, - MaxExclusive = Documents.Routing.PartitionKeyInternal.MaximumExclusiveEffectivePartitionKey, - }; - - private readonly FailureConfigs failureConfigs; - private readonly Random random; - - protected DocumentContainer(FailureConfigs failureConfigs = null) + public DocumentContainer(IMonadicDocumentContainer monadicDocumentContainer) { - this.failureConfigs = failureConfigs; - this.random = new Random(); + this.monadicDocumentContainer = monadicDocumentContainer ?? throw new ArgumentNullException(nameof(monadicDocumentContainer)); } - protected abstract Task>> MonadicGetChildRangeImplementationAsync( - PartitionKeyRange partitionKeyRange, - CancellationToken cancellationToken); - public Task>> MonadicGetChildRangeAsync( PartitionKeyRange partitionKeyRange, - CancellationToken cancellationToken) => this.MonadicGetChildRangeImplementationAsync( + CancellationToken cancellationToken) => this.monadicDocumentContainer.MonadicGetChildRangeAsync( partitionKeyRange, cancellationToken); @@ -63,8 +39,7 @@ public Task> GetChildRangeAsync( cancellationToken); public Task>> MonadicGetFeedRangesAsync( - CancellationToken cancellationToken) => this.MonadicGetChildRangeAsync( - DocumentContainer.FullRange, + CancellationToken cancellationToken) => this.monadicDocumentContainer.MonadicGetFeedRangesAsync( cancellationToken); public Task> GetFeedRangesAsync( @@ -73,23 +48,11 @@ public Task> GetFeedRangesAsync( cancellationToken), cancellationToken); - protected abstract Task> MonadicCreateItemImplementationAsync( - CosmosObject payload, - CancellationToken cancellationToken); - public Task> MonadicCreateItemAsync( CosmosObject payload, - CancellationToken cancellationToken) - { - if (this.ShouldReturn429()) - { - return ThrottleForCreateItem; - } - - return this.MonadicCreateItemImplementationAsync( + CancellationToken cancellationToken) => this.monadicDocumentContainer.MonadicCreateItemAsync( payload, cancellationToken); - } public Task CreateItemAsync( CosmosObject payload, @@ -99,26 +62,13 @@ public Task CreateItemAsync( cancellationToken), cancellationToken); - protected abstract Task> MonadicReadItemImplementationAsync( - CosmosElement partitionKey, - Guid identifer, - CancellationToken cancellationToken); - public Task> MonadicReadItemAsync( CosmosElement partitionKey, Guid identifer, - CancellationToken cancellationToken) - { - if (this.ShouldReturn429()) - { - return ThrottleForCreateItem; - } - - return this.MonadicReadItemImplementationAsync( + CancellationToken cancellationToken) => this.monadicDocumentContainer.MonadicReadItemAsync( partitionKey, identifer, cancellationToken); - } public Task ReadItemAsync( CosmosElement partitionKey, @@ -130,38 +80,15 @@ public Task ReadItemAsync( cancellationToken), cancellationToken); - protected abstract Task> MonadicReadFeedImplementationAsync( - int partitionKeyRangeId, - long resourceIdentifer, - int pageSize, - CancellationToken cancellationToken); - public Task> MonadicReadFeedAsync( int partitionKeyRangeId, long resourceIdentifer, int pageSize, - CancellationToken cancellationToken) - { - if (this.ShouldReturn429()) - { - return ThrottleForFeedOperation; - } - - if (this.ShouldReturnEmptyPage()) - { - return Task.FromResult( - TryCatch.FromResult( - new DocumentContainerPage( - new List(), - new DocumentContainerState(resourceIdentifer)))); - } - - return this.MonadicReadFeedImplementationAsync( + CancellationToken cancellationToken) => this.monadicDocumentContainer.MonadicReadFeedAsync( partitionKeyRangeId, resourceIdentifer, pageSize, cancellationToken); - } public Task ReadFeedAsync( int partitionKeyRangeId, @@ -175,9 +102,11 @@ public Task ReadFeedAsync( cancellationToken), cancellationToken); - public abstract Task MonadicSplitAsync( + public Task MonadicSplitAsync( int partitionKeyRangeId, - CancellationToken cancellationToken); + CancellationToken cancellationToken) => this.monadicDocumentContainer.MonadicSplitAsync( + partitionKeyRangeId, + cancellationToken); public Task SplitAsync( int partitionKeyRangeId, @@ -186,21 +115,5 @@ public Task SplitAsync( partitionKeyRangeId, cancellationToken), cancellationToken); - private bool ShouldReturn429() => (this.failureConfigs != null) && this.failureConfigs.Inject429s && ((this.random.Next() % 2) == 0); - - private bool ShouldReturnEmptyPage() => (this.failureConfigs != null) && this.failureConfigs.InjectEmptyPages && ((this.random.Next() % 2) == 0); - - public sealed class FailureConfigs - { - public FailureConfigs(bool inject429s, bool injectEmptyPages) - { - this.Inject429s = inject429s; - this.InjectEmptyPages = injectEmptyPages; - } - - public bool Inject429s { get; } - - public bool InjectEmptyPages { get; } - } } } diff --git a/Microsoft.Azure.Cosmos/src/Pagination/IDocumentContainer.cs b/Microsoft.Azure.Cosmos/src/Pagination/IDocumentContainer.cs index 72d10fde1f..504ce137df 100644 --- a/Microsoft.Azure.Cosmos/src/Pagination/IDocumentContainer.cs +++ b/Microsoft.Azure.Cosmos/src/Pagination/IDocumentContainer.cs @@ -8,44 +8,24 @@ namespace Microsoft.Azure.Cosmos.Pagination using System.Threading; using System.Threading.Tasks; using Microsoft.Azure.Cosmos.CosmosElements; - using Microsoft.Azure.Cosmos.Query.Core.Monads; - internal interface IDocumentContainer : IFeedRangeProvider + internal interface IDocumentContainer : IMonadicDocumentContainer, IFeedRangeProvider { - Task> MonadicCreateItemAsync( - CosmosObject payload, - CancellationToken cancellationToken); - Task CreateItemAsync( CosmosObject payload, CancellationToken cancellationToken); - Task> MonadicReadItemAsync( - CosmosElement partitionKey, - Guid identifer, - CancellationToken cancellationToken); - Task ReadItemAsync( CosmosElement partitionKey, Guid identifier, CancellationToken cancellationToken); - Task> MonadicReadFeedAsync( - int partitionKeyRangeId, - long resourceIdentifer, - int pageSize, - CancellationToken cancellationToken); - Task ReadFeedAsync( int partitionKeyRangeId, long resourceIdentifier, int pageSize, CancellationToken cancellationToken); - Task MonadicSplitAsync( - int partitionKeyRangeId, - CancellationToken cancellationToken); - Task SplitAsync( int partitionKeyRangeId, CancellationToken cancellationToken); diff --git a/Microsoft.Azure.Cosmos/src/Pagination/IFeedRangeProvider.cs b/Microsoft.Azure.Cosmos/src/Pagination/IFeedRangeProvider.cs index 6ec3401c51..5ee0da6b06 100644 --- a/Microsoft.Azure.Cosmos/src/Pagination/IFeedRangeProvider.cs +++ b/Microsoft.Azure.Cosmos/src/Pagination/IFeedRangeProvider.cs @@ -7,22 +7,14 @@ namespace Microsoft.Azure.Cosmos.Pagination using System.Collections.Generic; using System.Threading; using System.Threading.Tasks; - using Microsoft.Azure.Cosmos.Query.Core.Monads; using Microsoft.Azure.Documents; - internal interface IFeedRangeProvider + internal interface IFeedRangeProvider : IMonadicFeedRangeProvider { - Task>> MonadicGetChildRangeAsync( - PartitionKeyRange partitionKeyRange, - CancellationToken cancellationToken); - Task> GetChildRangeAsync( PartitionKeyRange partitionKeyRange, CancellationToken cancellationToken); - Task>> MonadicGetFeedRangesAsync( - CancellationToken cancellationToken); - Task> GetFeedRangesAsync( CancellationToken cancellationToken); } diff --git a/Microsoft.Azure.Cosmos/src/Pagination/IMonadicDocumentContainer.cs b/Microsoft.Azure.Cosmos/src/Pagination/IMonadicDocumentContainer.cs new file mode 100644 index 0000000000..f7694cdbae --- /dev/null +++ b/Microsoft.Azure.Cosmos/src/Pagination/IMonadicDocumentContainer.cs @@ -0,0 +1,34 @@ +// ------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +// ------------------------------------------------------------ + +namespace Microsoft.Azure.Cosmos.Pagination +{ + using System; + using System.Threading; + using System.Threading.Tasks; + using Microsoft.Azure.Cosmos.CosmosElements; + using Microsoft.Azure.Cosmos.Query.Core.Monads; + + internal interface IMonadicDocumentContainer : IMonadicFeedRangeProvider + { + Task> MonadicCreateItemAsync( + CosmosObject payload, + CancellationToken cancellationToken); + + Task> MonadicReadItemAsync( + CosmosElement partitionKey, + Guid identifer, + CancellationToken cancellationToken); + + Task> MonadicReadFeedAsync( + int partitionKeyRangeId, + long resourceIdentifer, + int pageSize, + CancellationToken cancellationToken); + + Task MonadicSplitAsync( + int partitionKeyRangeId, + CancellationToken cancellationToken); + } +} diff --git a/Microsoft.Azure.Cosmos/src/Pagination/IMonadicFeedRangeProvider.cs b/Microsoft.Azure.Cosmos/src/Pagination/IMonadicFeedRangeProvider.cs new file mode 100644 index 0000000000..587121b5f1 --- /dev/null +++ b/Microsoft.Azure.Cosmos/src/Pagination/IMonadicFeedRangeProvider.cs @@ -0,0 +1,22 @@ +// ------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +// ------------------------------------------------------------ + +namespace Microsoft.Azure.Cosmos.Pagination +{ + using System.Collections.Generic; + using System.Threading; + using System.Threading.Tasks; + using Microsoft.Azure.Cosmos.Query.Core.Monads; + using Microsoft.Azure.Documents; + + internal interface IMonadicFeedRangeProvider + { + Task>> MonadicGetChildRangeAsync( + PartitionKeyRange partitionKeyRange, + CancellationToken cancellationToken); + + Task>> MonadicGetFeedRangesAsync( + CancellationToken cancellationToken); + } +} diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/DocumentContainerTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/DocumentContainerTests.cs index 5d9e48f918..85f368152c 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/DocumentContainerTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/DocumentContainerTests.cs @@ -30,7 +30,7 @@ public abstract class DocumentContainerTests internal abstract IDocumentContainer CreateDocumentContainer( PartitionKeyDefinition partitionKeyDefinition, - DocumentContainer.FailureConfigs failureConfigs = default); + FlakyDocumentContainer.FailureConfigs failureConfigs = default); [TestMethod] public async Task TestGetFeedRanges() diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/FlakyDocumentContainer.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/FlakyDocumentContainer.cs new file mode 100644 index 0000000000..edd2944a2a --- /dev/null +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/FlakyDocumentContainer.cs @@ -0,0 +1,145 @@ +//------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +//------------------------------------------------------------ + +namespace Microsoft.Azure.Cosmos.Tests.Pagination +{ + using System; + using System.Collections.Generic; + using System.Threading; + using System.Threading.Tasks; + using Microsoft.Azure.Cosmos.CosmosElements; + using Microsoft.Azure.Cosmos.Pagination; + using Microsoft.Azure.Cosmos.Query.Core.Monads; + using Microsoft.Azure.Documents; + + /// + /// Implementation of that composes another and randomly adds in exceptions. + /// This is useful for mocking throttles and other edge cases like empty pages. + /// + internal sealed class FlakyDocumentContainer : IMonadicDocumentContainer + { + private readonly FailureConfigs failureConfigs; + private readonly Random random; + + private static readonly CosmosException RequestRateTooLargeException = new CosmosException( + message: "Request Rate Too Large", + statusCode: (System.Net.HttpStatusCode)429, + subStatusCode: default, + activityId: Guid.NewGuid().ToString(), + requestCharge: default); + + private static readonly Task> ThrottleForCreateItem = Task.FromResult( + TryCatch.FromException( + RequestRateTooLargeException)); + + private static readonly Task> ThrottleForFeedOperation = Task.FromResult( + TryCatch.FromException( + RequestRateTooLargeException)); + + private readonly IMonadicDocumentContainer documentContainer; + + public FlakyDocumentContainer( + IMonadicDocumentContainer documentContainer, + FailureConfigs failureConfigs) + { + this.documentContainer = documentContainer ?? throw new ArgumentNullException(nameof(documentContainer)); + this.failureConfigs = failureConfigs ?? throw new ArgumentNullException(nameof(failureConfigs)); + this.random = new Random(); + } + + public Task> MonadicCreateItemAsync( + CosmosObject payload, + CancellationToken cancellationToken) + { + if (this.ShouldReturn429()) + { + return ThrottleForCreateItem; + } + + return this.documentContainer.MonadicCreateItemAsync( + payload, + cancellationToken); + } + + public Task> MonadicReadItemAsync( + CosmosElement partitionKey, + Guid identifer, + CancellationToken cancellationToken) + { + if (this.ShouldReturn429()) + { + return ThrottleForCreateItem; + } + + return this.documentContainer.MonadicReadItemAsync( + partitionKey, + identifer, + cancellationToken); + } + + public Task> MonadicReadFeedAsync( + int partitionKeyRangeId, + long resourceIdentifer, + int pageSize, + CancellationToken cancellationToken) + { + if (this.ShouldReturn429()) + { + return ThrottleForFeedOperation; + } + + if (this.ShouldReturnEmptyPage()) + { + return Task.FromResult( + TryCatch.FromResult( + new DocumentContainerPage( + new List(), + new DocumentContainerState(resourceIdentifer)))); + } + + return this.documentContainer.MonadicReadFeedAsync( + partitionKeyRangeId, + resourceIdentifer, + pageSize, + cancellationToken); + } + + public Task MonadicSplitAsync( + int partitionKeyRangeId, + CancellationToken cancellationToken) => this.documentContainer.MonadicSplitAsync( + partitionKeyRangeId, + cancellationToken); + + public Task>> MonadicGetChildRangeAsync( + PartitionKeyRange partitionKeyRange, + CancellationToken cancellationToken) => this.documentContainer.MonadicGetChildRangeAsync( + partitionKeyRange, + cancellationToken); + + public Task>> MonadicGetFeedRangesAsync( + CancellationToken cancellationToken) => this.documentContainer.MonadicGetFeedRangesAsync( + cancellationToken); + + private bool ShouldReturn429() => (this.failureConfigs != null) + && this.failureConfigs.Inject429s + && ((this.random.Next() % 2) == 0); + + private bool ShouldReturnEmptyPage() => (this.failureConfigs != null) + && this.failureConfigs.InjectEmptyPages + && ((this.random.Next() % 2) == 0); + + public sealed class FailureConfigs + { + public FailureConfigs(bool inject429s, bool injectEmptyPages) + { + this.Inject429s = inject429s; + this.InjectEmptyPages = injectEmptyPages; + } + + public bool Inject429s { get; } + + public bool InjectEmptyPages { get; } + } + } +} diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/InMemoryCollectionTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/InMemoryCollectionTests.cs index 46099a3090..2f6b6c5d21 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/InMemoryCollectionTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/InMemoryCollectionTests.cs @@ -13,6 +13,9 @@ public sealed class InMemoryCollectionTests : DocumentContainerTests { internal override IDocumentContainer CreateDocumentContainer( PartitionKeyDefinition partitionKeyDefinition, - DocumentContainer.FailureConfigs failureConfigs = null) => new InMemoryCollection(partitionKeyDefinition, failureConfigs); + FlakyDocumentContainer.FailureConfigs failureConfigs = null) => new DocumentContainer( + new FlakyDocumentContainer( + new InMemoryContainer(partitionKeyDefinition), + failureConfigs)); } } diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/InMemoryCollection.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/InMemoryContainer.cs similarity index 94% rename from Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/InMemoryCollection.cs rename to Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/InMemoryContainer.cs index 0e2dde9e0f..bf028faefc 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/InMemoryCollection.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/InMemoryContainer.cs @@ -17,18 +17,22 @@ namespace Microsoft.Azure.Cosmos.Tests.Pagination using Microsoft.Azure.Documents; // Collection useful for mocking requests and repartitioning (splits / merge). - internal sealed class InMemoryCollection : DocumentContainer + internal sealed class InMemoryContainer : IMonadicDocumentContainer { + private static readonly PartitionKeyRange FullRange = new PartitionKeyRange() + { + MinInclusive = Documents.Routing.PartitionKeyInternal.MinimumInclusiveEffectivePartitionKey, + MaxExclusive = Documents.Routing.PartitionKeyInternal.MaximumExclusiveEffectivePartitionKey, + }; + private readonly PartitionKeyDefinition partitionKeyDefinition; private readonly Dictionary parentToChildMapping; private PartitionKeyHashRangeDictionary partitionedRecords; private Dictionary partitionKeyRangeIdToHashRange; - public InMemoryCollection( - PartitionKeyDefinition partitionKeyDefinition, - FailureConfigs failureConfigs = default) - : base(failureConfigs) + public InMemoryContainer( + PartitionKeyDefinition partitionKeyDefinition) { this.partitionKeyDefinition = partitionKeyDefinition ?? throw new ArgumentNullException(nameof(partitionKeyDefinition)); PartitionKeyHashRange fullRange = new PartitionKeyHashRange(startInclusive: null, endExclusive: null); @@ -42,7 +46,12 @@ public InMemoryCollection( this.parentToChildMapping = new Dictionary(); } - protected override Task>> MonadicGetChildRangeImplementationAsync( + public Task>> MonadicGetFeedRangesAsync( + CancellationToken cancellationToken) => this.MonadicGetChildRangeAsync( + FullRange, + cancellationToken); + + public Task>> MonadicGetChildRangeAsync( PartitionKeyRange partitionKeyRange, CancellationToken cancellationToken) { @@ -97,7 +106,7 @@ PartitionKeyRange CreateRangeFromId(int id) return Task.FromResult(TryCatch>.FromResult(childRanges)); } - protected override Task> MonadicCreateItemImplementationAsync( + public Task> MonadicCreateItemAsync( CosmosObject payload, CancellationToken cancellationToken) { @@ -120,7 +129,7 @@ protected override Task> MonadicCreateItemImplementationAsync( return Task.FromResult(TryCatch.FromResult(recordAdded)); } - protected override Task> MonadicReadItemImplementationAsync( + public Task> MonadicReadItemAsync( CosmosElement partitionKey, Guid identifier, CancellationToken cancellationToken) @@ -168,7 +177,7 @@ static Task> CreateNotFoundException(CosmosElement partitionKey return CreateNotFoundException(partitionKey, identifier); } - protected override Task> MonadicReadFeedImplementationAsync( + public Task> MonadicReadFeedAsync( int partitionKeyRangeId, long resourceIdentifer, int pageSize, @@ -216,7 +225,7 @@ protected override Task> MonadicReadFeedImplemen state: new DocumentContainerState(page.Last().ResourceIdentifier)))); } - public override Task MonadicSplitAsync( + public Task MonadicSplitAsync( int partitionKeyRangeId, CancellationToken cancellationToken) { diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/PartitionRangeEnumeratorTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/PartitionRangeEnumeratorTests.cs index e119ed7860..a3774ef5e1 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/PartitionRangeEnumeratorTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/PartitionRangeEnumeratorTests.cs @@ -50,7 +50,7 @@ public async Task Test429sAsync() int numItems = 100; IDocumentContainer inMemoryCollection = await this.CreateDocumentContainerAsync( numItems, - new DocumentContainer.FailureConfigs( + new FlakyDocumentContainer.FailureConfigs( inject429s: true, injectEmptyPages: false)); @@ -91,7 +91,7 @@ public async Task Test429sWithContinuationsAsync() int numItems = 100; IDocumentContainer inMemoryCollection = await this.CreateDocumentContainerAsync( numItems, - new DocumentContainer.FailureConfigs( + new FlakyDocumentContainer.FailureConfigs( inject429s: true, injectEmptyPages: false)); @@ -140,7 +140,7 @@ public async Task TestEmptyPages() int numItems = 100; IDocumentContainer inMemoryCollection = await this.CreateDocumentContainerAsync( numItems, - new DocumentContainer.FailureConfigs( + new FlakyDocumentContainer.FailureConfigs( inject429s: false, injectEmptyPages: true)); IAsyncEnumerable> enumerable = this.CreateEnumerable(inMemoryCollection); @@ -156,7 +156,7 @@ public async Task TestEmptyPages() public async Task CreateDocumentContainerAsync( int numItems, - DocumentContainer.FailureConfigs failureConfigs = default) + FlakyDocumentContainer.FailureConfigs failureConfigs = default) { PartitionKeyDefinition partitionKeyDefinition = new PartitionKeyDefinition() { @@ -168,44 +168,39 @@ public async Task CreateDocumentContainerAsync( Version = PartitionKeyDefinitionVersion.V2, }; - InMemoryCollection inMemoryCollection = new InMemoryCollection(partitionKeyDefinition, failureConfigs); + InMemoryContainer inMemoryCollection = new InMemoryContainer(partitionKeyDefinition); + FlakyDocumentContainer flakyDocumentContainer = new FlakyDocumentContainer(inMemoryCollection, failureConfigs); + DocumentContainer documentContainer = new DocumentContainer(flakyDocumentContainer); if (!this.singlePartition) { - await inMemoryCollection.SplitAsync(partitionKeyRangeId: 0, cancellationToken: default); + await documentContainer.SplitAsync(partitionKeyRangeId: 0, cancellationToken: default); - await inMemoryCollection.SplitAsync(partitionKeyRangeId: 1, cancellationToken: default); - await inMemoryCollection.SplitAsync(partitionKeyRangeId: 2, cancellationToken: default); + await documentContainer.SplitAsync(partitionKeyRangeId: 1, cancellationToken: default); + await documentContainer.SplitAsync(partitionKeyRangeId: 2, cancellationToken: default); - await inMemoryCollection.SplitAsync(partitionKeyRangeId: 3, cancellationToken: default); - await inMemoryCollection.SplitAsync(partitionKeyRangeId: 4, cancellationToken: default); - await inMemoryCollection.SplitAsync(partitionKeyRangeId: 5, cancellationToken: default); - await inMemoryCollection.SplitAsync(partitionKeyRangeId: 6, cancellationToken: default); + await documentContainer.SplitAsync(partitionKeyRangeId: 3, cancellationToken: default); + await documentContainer.SplitAsync(partitionKeyRangeId: 4, cancellationToken: default); + await documentContainer.SplitAsync(partitionKeyRangeId: 5, cancellationToken: default); + await documentContainer.SplitAsync(partitionKeyRangeId: 6, cancellationToken: default); } for (int i = 0; i < numItems; i++) { - try + // Insert an item + CosmosObject item = CosmosObject.Parse($"{{\"pk\" : {i} }}"); + //await inMemoryCollection.CreateItemAsync(item, cancellationToken: default); + while (true) { - // Insert an item - CosmosObject item = CosmosObject.Parse($"{{\"pk\" : {i} }}"); - //await inMemoryCollection.CreateItemAsync(item, cancellationToken: default); - while (true) + TryCatch monadicCreateRecord = await inMemoryCollection.MonadicCreateItemAsync(item, cancellationToken: default); + if (monadicCreateRecord.Succeeded) { - TryCatch monadicCreateRecord = await inMemoryCollection.MonadicCreateItemAsync(item, cancellationToken: default); - if (monadicCreateRecord.Succeeded) - { - break; - } + break; } } - catch(Exception ex) - { - throw ex; - } } - return inMemoryCollection; + return documentContainer; } public async Task> DrainFullyAsync(IAsyncEnumerable> enumerable) From 464d5cca259f8cb066beb91ca7b8553ed9b2717c Mon Sep 17 00:00:00 2001 From: Brandon Chong Date: Mon, 20 Jul 2020 14:32:52 -0700 Subject: [PATCH 44/85] started order by, but need the first PR to go through. To many merge conflicts --- ...OrderByCrossPartitionQueryPipelineStage.cs | 116 ++++++++++++++++++ ...arallelCrossPartitionQueryPipelineStage.cs | 97 ++++++++------- 2 files changed, 170 insertions(+), 43 deletions(-) create mode 100644 Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderByCrossPartitionQueryPipelineStage.cs diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderByCrossPartitionQueryPipelineStage.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderByCrossPartitionQueryPipelineStage.cs new file mode 100644 index 0000000000..0c3300d8ca --- /dev/null +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderByCrossPartitionQueryPipelineStage.cs @@ -0,0 +1,116 @@ +// ------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +// ------------------------------------------------------------ + +namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote +{ + using System; + using System.Collections.Generic; + using System.Threading.Tasks; + using Microsoft.Azure.Cosmos.CosmosElements; + using Microsoft.Azure.Cosmos.Pagination; + using Microsoft.Azure.Cosmos.Query.Core.ContinuationTokens; + using Microsoft.Azure.Cosmos.Query.Core.Exceptions; + using Microsoft.Azure.Cosmos.Query.Core.Monads; + + internal sealed class OrderByCrossPartitionQueryPipelineStage : IQueryPipelineStage + { + private readonly CrossPartitionRangePageAsyncEnumerator crossPartitionRangePageAsyncEnumerator; + + private OrderByCrossPartitionQueryPipelineStage( + CrossPartitionRangePageAsyncEnumerator crossPartitionRangePageAsyncEnumerator) + { + this.crossPartitionRangePageAsyncEnumerator = crossPartitionRangePageAsyncEnumerator ?? throw new ArgumentNullException(nameof(crossPartitionRangePageAsyncEnumerator)); + } + + public TryCatch Current => throw new NotImplementedException(); + + public ValueTask DisposeAsync() => this.crossPartitionRangePageAsyncEnumerator.DisposeAsync(); + + public ValueTask MoveNextAsync() + { + throw new NotImplementedException(); + } + + public static TryCatch MonadicCreate( + IDocumentContainer documentContainer, + SqlQuerySpec sqlQuerySpec, + int pageSize, + CosmosElement continuationToken) + { + if (pageSize <= 0) + { + throw new ArgumentOutOfRangeException(nameof(pageSize)); + } + + TryCatch> monadicExtractState = MonadicExtractState(continuationToken); + if (monadicExtractState.Failed) + { + return TryCatch.FromException(monadicExtractState.Exception); + } + + CrossPartitionState state = monadicExtractState.Result; + + + } + + private static TryCatch> MonadicExtractState( + CosmosElement continuationToken, + int numOrderByColumns) + { + if (continuationToken == null) + { + return TryCatch>.FromResult(default); + } + + if (!(continuationToken is CosmosArray cosmosArray)) + { + return TryCatch>.FromException( + new MalformedContinuationTokenException( + $"Order by continuation token must be an array: {continuationToken}.")); + } + + List orderByContinuationTokens = new List(); + foreach (CosmosElement arrayItem in cosmosArray) + { + TryCatch tryCreateOrderByContinuationToken = OrderByContinuationToken.TryCreateFromCosmosElement(arrayItem); + if (!tryCreateOrderByContinuationToken.Succeeded) + { + return TryCatch>.FromException(tryCreateOrderByContinuationToken.Exception); + } + + orderByContinuationTokens.Add(tryCreateOrderByContinuationToken.Result); + } + + if (orderByContinuationTokens.Count == 0) + { + return TryCatch>.FromException( + new MalformedContinuationTokenException( + $"Order by continuation token cannot be empty: {continuationToken}.")); + } + + foreach (OrderByContinuationToken suppliedOrderByContinuationToken in orderByContinuationTokens) + { + if (suppliedOrderByContinuationToken.OrderByItems.Count != numOrderByColumns) + { + return TryCatch>.FromException( + new MalformedContinuationTokenException( + $"Invalid order-by items in continuation token {continuationToken} for OrderBy~Context.")); + } + } + } + + private sealed class OrderByQueryPage : Page + { + public OrderByQueryPage(QueryPage queryPage) + : base(queryPage.State) + { + this.Enumerator = queryPage.Documents.GetEnumerator(); + } + + public QueryPage Page { get; } + + public IEnumerator Enumerator { get; } + } + } +} diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/ParallelCrossPartitionQueryPipelineStage.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/ParallelCrossPartitionQueryPipelineStage.cs index 556a89443b..ea7e1bcbe5 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/ParallelCrossPartitionQueryPipelineStage.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/ParallelCrossPartitionQueryPipelineStage.cs @@ -94,52 +94,13 @@ public static TryCatch MonadicCreate( throw new ArgumentOutOfRangeException(nameof(pageSize)); } - CrossPartitionState state; - if (continuationToken == null) + TryCatch> monadicExtractState = MonadicExtractState(continuationToken); + if (monadicExtractState.Failed) { - state = default; + return TryCatch.FromException(monadicExtractState.Exception); } - else - { - if (!(continuationToken is CosmosArray compositeContinuationTokenListRaw)) - { - return TryCatch.FromException( - new MalformedContinuationTokenException( - $"Invalid format for continuation token {continuationToken} for {nameof(ParallelCrossPartitionQueryPipelineStage)}")); - } - - if (compositeContinuationTokenListRaw.Count == 0) - { - return TryCatch.FromException( - new MalformedContinuationTokenException( - $"Invalid format for continuation token {continuationToken} for {nameof(ParallelCrossPartitionQueryPipelineStage)}")); - } - - List compositeContinuationTokens = new List(); - foreach (CosmosElement compositeContinuationTokenRaw in compositeContinuationTokenListRaw) - { - TryCatch tryCreateCompositeContinuationToken = CompositeContinuationToken.TryCreateFromCosmosElement(compositeContinuationTokenRaw); - if (tryCreateCompositeContinuationToken.Failed) - { - return TryCatch.FromException( - tryCreateCompositeContinuationToken.Exception); - } - - compositeContinuationTokens.Add(tryCreateCompositeContinuationToken.Result); - } - - List<(PartitionKeyRange, QueryState)> rangesAndStates = compositeContinuationTokens - .Select(token => ( - new PartitionKeyRange() - { - MinInclusive = token.Range.Min, - MaxExclusive = token.Range.Max, - }, - token.Token != null ? new QueryState(CosmosString.Create(token.Token)) : null)) - .ToList(); - state = new CrossPartitionState(rangesAndStates); - } + CrossPartitionState state = monadicExtractState.Result; CrossPartitionRangePageAsyncEnumerator crossPartitionPageEnumerator = new CrossPartitionRangePageAsyncEnumerator( documentContainer, @@ -151,6 +112,56 @@ public static TryCatch MonadicCreate( return TryCatch.FromResult(stage); } + private static TryCatch> MonadicExtractState( + CosmosElement continuationToken) + { + if (continuationToken == null) + { + return TryCatch>.FromResult(default); + } + + if (!(continuationToken is CosmosArray compositeContinuationTokenListRaw)) + { + return TryCatch>.FromException( + new MalformedContinuationTokenException( + $"Invalid format for continuation token {continuationToken} for {nameof(ParallelCrossPartitionQueryPipelineStage)}")); + } + + if (compositeContinuationTokenListRaw.Count == 0) + { + return TryCatch>.FromException( + new MalformedContinuationTokenException( + $"Invalid format for continuation token {continuationToken} for {nameof(ParallelCrossPartitionQueryPipelineStage)}")); + } + + List compositeContinuationTokens = new List(); + foreach (CosmosElement compositeContinuationTokenRaw in compositeContinuationTokenListRaw) + { + TryCatch tryCreateCompositeContinuationToken = CompositeContinuationToken.TryCreateFromCosmosElement(compositeContinuationTokenRaw); + if (tryCreateCompositeContinuationToken.Failed) + { + return TryCatch>.FromException( + tryCreateCompositeContinuationToken.Exception); + } + + compositeContinuationTokens.Add(tryCreateCompositeContinuationToken.Result); + } + + List<(PartitionKeyRange, QueryState)> rangesAndStates = compositeContinuationTokens + .Select(token => ( + new PartitionKeyRange() + { + MinInclusive = token.Range.Min, + MaxExclusive = token.Range.Max, + }, + token.Token != null ? new QueryState(CosmosString.Create(token.Token)) : null)) + .ToList(); + + CrossPartitionState state = new CrossPartitionState(rangesAndStates); + + return TryCatch>.FromResult(state); + } + private static CreatePartitionRangePageAsyncEnumerator MakeCreateFunction( IQueryDataSource queryDataSource, SqlQuerySpec sqlQuerySpec, From 965bdf37fda28d7fac0ff2d6c6a3366a0a7faab7 Mon Sep 17 00:00:00 2001 From: Brandon Chong Date: Thu, 23 Jul 2020 19:08:51 -0700 Subject: [PATCH 45/85] fix build errors --- .../Contracts/ContractTests.cs | 2 +- .../Pagination/InMemoryContainer.cs | 4 +--- 2 files changed, 2 insertions(+), 4 deletions(-) diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/Contracts/ContractTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/Contracts/ContractTests.cs index 0b9422c348..6ba70c45bc 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/Contracts/ContractTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/Contracts/ContractTests.cs @@ -208,7 +208,7 @@ public async Task ChangeFeed_FeedRange_FromV2SDK() count += response.Count; string migratedContinuation = firstResponse.ContinuationToken; Assert.IsTrue(FeedRangeContinuation.TryParse(migratedContinuation, out FeedRangeContinuation feedRangeContinuation)); - Assert.IsTrue(feedRangeContinuation.FeedRange is FeedRangeEPK); + Assert.IsTrue(feedRangeContinuation.FeedRange is FeedRangeEpk); } } 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 bf028faefc..85a2098c41 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 @@ -164,9 +164,7 @@ static Task> CreateNotFoundException(CosmosElement partitionKey CosmosElement candidatePartitionKey = GetPartitionKeyFromPayload( candidate.Payload, this.partitionKeyDefinition); - bool partitionKeyMatches = CosmosElementEqualityComparer.Value.Equals( - candidatePartitionKey, - partitionKey); + bool partitionKeyMatches = candidatePartitionKey.Equals(partitionKey); if (identifierMatches && partitionKeyMatches) { From bfed6196d4bba6740a042e14b7980d2d2240032d Mon Sep 17 00:00:00 2001 From: Brandon Chong Date: Tue, 28 Jul 2020 23:00:52 -0700 Subject: [PATCH 46/85] need to make serious changes for order by --- .../CrossPartitionRangeItemAsyncEnumerator.cs | 199 ++++++ .../src/Pagination/ItemEnumeratorPage.cs | 22 + .../PartitionRangePageAsyncEnumerator.cs | 8 +- ...smosCrossPartitionQueryExecutionContext.cs | 153 ----- ...OrderByItemQueryExecutionContext.Resume.cs | 572 ----------------- .../CosmosOrderByItemQueryExecutionContext.cs | 20 - ...CosmosParallelItemQueryExecutionContext.cs | 7 - .../CrossPartitionQueryPipelineStage.cs | 157 +++++ ...OrderByCrossPartitionQueryPipelineStage.cs | 578 +++++++++++++++++- .../Core/Pipeline/Remote/OrderByQueryPage.cs | 23 + ...yQueryPartitionRangePageAsyncEnumerator.cs | 55 ++ ...arallelCrossPartitionQueryPipelineStage.cs | 7 + 12 files changed, 1020 insertions(+), 781 deletions(-) create mode 100644 Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionRangeItemAsyncEnumerator.cs create mode 100644 Microsoft.Azure.Cosmos/src/Pagination/ItemEnumeratorPage.cs create mode 100644 Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/CrossPartitionQueryPipelineStage.cs create mode 100644 Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderByQueryPage.cs create mode 100644 Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderByQueryPartitionRangePageAsyncEnumerator.cs diff --git a/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionRangeItemAsyncEnumerator.cs b/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionRangeItemAsyncEnumerator.cs new file mode 100644 index 0000000000..52eca81425 --- /dev/null +++ b/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionRangeItemAsyncEnumerator.cs @@ -0,0 +1,199 @@ +// ------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +// ------------------------------------------------------------ + +namespace Microsoft.Azure.Cosmos.Pagination +{ + using System; + using System.Collections.Generic; + using System.Net; + using System.Threading.Tasks; + using Microsoft.Azure.Cosmos.CosmosElements; + using Microsoft.Azure.Cosmos.Query.Core.Collections; + using Microsoft.Azure.Cosmos.Query.Core.Monads; + using Microsoft.Azure.Documents; + + /// + /// Coordinates draining pages from multiple , while maintaining a global sort order and handling repartitioning (splits, merge). + /// Also does things on a per item basis. + /// + internal sealed class CrossPartitionRangeItemAsyncEnumerator : IAsyncEnumerator> + where TState : State + { + private readonly IFeedRangeProvider feedRangeProvider; + private readonly CreatePartitionRangePageAsyncEnumerator, TState> createPartitionRangeEnumerator; + private readonly Func, TState>, Task>>> intializeAsync; + private readonly PriorityQueue, TState>> enumerators; + private readonly Queue, TState>> uninitializedEnumerators; + + public CrossPartitionRangeItemAsyncEnumerator( + IFeedRangeProvider feedRangeProvider, + CreatePartitionRangePageAsyncEnumerator, TState> createPartitionRangeEnumerator, + IEnumerable, TState>> uninitializedEnumerators, + Func, TState>, Task>>> initializeAsync, + IComparer itemComparer) + { + this.feedRangeProvider = feedRangeProvider ?? throw new ArgumentNullException(nameof(feedRangeProvider)); + this.createPartitionRangeEnumerator = createPartitionRangeEnumerator ?? throw new ArgumentNullException(nameof(createPartitionRangeEnumerator)); + this.intializeAsync = initializeAsync ?? throw new ArgumentNullException(nameof(initializeAsync)); + this.uninitializedEnumerators = new Queue, TState>>(uninitializedEnumerators); + this.enumerators = new PriorityQueue, TState>>(new EnumeratorComparer()); + } + + public TryCatch, TState>> Current { get; private set; } + + public async ValueTask MoveNextAsync() + { + if (this.uninitializedEnumerators.Count != 0) + { + PartitionRangePageAsyncEnumerator, TState> uninitializedEnumerator = this.uninitializedEnumerators.Dequeue(); + TryCatch> initializeMonad = await this.intializeAsync(uninitializedEnumerator); + if (initializeMonad.Failed) + { + if (!await this.TryHandleExceptionAsync(uninitializedEnumerator)) + { + this.uninitializedEnumerators.Enqueue(uninitializedEnumerator); + } + + this.Current = TryCatch, TState>>.FromException(uninitializedEnumerator.Current.Exception); + return true; + } + + // Once the enumerator has been initialized we can add it back to the priority queue + this.enumerators.Enqueue(uninitializedEnumerator); + + // We want to report back the metrics from initialization, so that the user has accurate metrics, + // But we need to make up a fake continuation token, since we aren't in a valid state to continue from. + this.Current = TryCatch, TState>>.FromResult( + new CrossPartitionPage, TState>( + initializeMonad.Result, + new CrossPartitionState(new List<(PartitionKeyRange, TState)>()))); + + // and recursively retry + return await this.MoveNextAsync(); + } + + if (this.enumerators.Count == 0) + { + return false; + } + + PartitionRangePageAsyncEnumerator, TState> currentEnumerator = this.enumerators.Dequeue(); + if (currentEnumerator.Current.Result.Items.Count == 0) + { + + } + + if (!await currentEnumerator.MoveNextAsync()) + { + // Current enumerator is empty, + // so recursively retry on the next enumerator. + return await this.MoveNextAsync(); + } + + if (currentEnumerator.Current.Failed) + { + + } + + if (currentEnumerator.State != null) + { + currentEnumerator.Enqueue(currentPaginator); + } + + TryCatch> backendPage = currentEnumerator.Current; + if (backendPage.Failed) + { + this.Current = TryCatch, TState>>.FromException(backendPage.Exception); + return true; + } + + CrossPartitionState crossPartitionState; + if (enumerators.Count == 0) + { + crossPartitionState = null; + } + else + { + List<(PartitionKeyRange, TState)> feedRangeAndStates = new List<(PartitionKeyRange, TState)>(enumerators.Count); + foreach (PartitionRangePageAsyncEnumerator, TState> enumerator in enumerators) + { + feedRangeAndStates.Add((enumerator.Range, enumerator.State)); + } + + crossPartitionState = new CrossPartitionState(feedRangeAndStates); + } + + this.Current = TryCatch, TState>>.FromResult( + new CrossPartitionPage, TState>(backendPage.Result, crossPartitionState)); + return true; + } + + private async Task TryHandleExceptionAsync(PartitionRangePageAsyncEnumerator, TState> enumerator) + { + if (!enumerator.Current.Failed) + { + throw new InvalidOperationException("Enumerator was not in a faulted state."); + } + + Exception exception = enumerator.Current.Exception; + + // Check if it's a retryable exception. + while (exception.InnerException != null) + { + exception = exception.InnerException; + } + + if (IsSplitException(exception)) + { + // Handle split + IEnumerable childRanges = await this.feedRangeProvider.GetChildRangeAsync( + enumerator.Range, + cancellationToken: default); + foreach (PartitionKeyRange childRange in childRanges) + { + PartitionRangePageAsyncEnumerator, TState> childPaginator = this.createPartitionRangeEnumerator( + childRange, + enumerator.State); + this.uninitializedEnumerators.Enqueue(childPaginator); + } + + return true; + } + + if (IsMergeException(exception)) + { + throw new NotImplementedException(); + } + + return false; + } + + public ValueTask DisposeAsync() + { + // Do Nothing. + return default; + } + + private static bool IsSplitException(Exception exeception) + { + return exeception is CosmosException cosmosException + && (cosmosException.StatusCode == HttpStatusCode.Gone) + && (cosmosException.SubStatusCode == (int)Documents.SubStatusCodes.PartitionKeyRangeGone); + } + + private static bool IsMergeException(Exception exception) + { + // TODO: code this out + return false; + } + + private sealed class EnumeratorComparer : IComparer, TState>> + { + public int Compare(PartitionRangePageAsyncEnumerator, TState> x, PartitionRangePageAsyncEnumerator, TState> y) + { + throw new NotImplementedException(); + } + } + } +} diff --git a/Microsoft.Azure.Cosmos/src/Pagination/ItemEnumeratorPage.cs b/Microsoft.Azure.Cosmos/src/Pagination/ItemEnumeratorPage.cs new file mode 100644 index 0000000000..fe98eb0a8e --- /dev/null +++ b/Microsoft.Azure.Cosmos/src/Pagination/ItemEnumeratorPage.cs @@ -0,0 +1,22 @@ +// ------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +// ------------------------------------------------------------ + +namespace Microsoft.Azure.Cosmos.Pagination +{ + using System; + using System.Collections.Generic; + using System.Text; + using Microsoft.Azure.Cosmos.CosmosElements; + + internal sealed class ItemEnumeratorPage : Page + where TState : State + { + public ItemEnumeratorPage(IReadOnlyList cosmosElements) + { + this.Items = new Queue(cosmosElements); + } + + public Queue Items { get; } + } +} diff --git a/Microsoft.Azure.Cosmos/src/Pagination/PartitionRangePageAsyncEnumerator.cs b/Microsoft.Azure.Cosmos/src/Pagination/PartitionRangePageAsyncEnumerator.cs index 7340d5e9bb..1be4829fb7 100644 --- a/Microsoft.Azure.Cosmos/src/Pagination/PartitionRangePageAsyncEnumerator.cs +++ b/Microsoft.Azure.Cosmos/src/Pagination/PartitionRangePageAsyncEnumerator.cs @@ -17,8 +17,6 @@ internal abstract class PartitionRangePageAsyncEnumerator : IAsyn where TPage : Page where TState : State { - private bool hasStarted; - protected PartitionRangePageAsyncEnumerator(PartitionKeyRange range, TState state = default) { this.Range = range; @@ -31,7 +29,9 @@ protected PartitionRangePageAsyncEnumerator(PartitionKeyRange range, TState stat public TState State { get; private set; } - private bool HasMoreResults => !this.hasStarted || (this.State != default); + public bool HasStarted { get; private set; } + + private bool HasMoreResults => !this.HasStarted || (this.State != default); public async ValueTask MoveNextAsync() { @@ -46,7 +46,7 @@ public async ValueTask MoveNextAsync() this.State = this.Current.Result.State; } - this.hasStarted = true; + this.HasStarted = true; return true; } diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosCrossPartitionQueryExecutionContext.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosCrossPartitionQueryExecutionContext.cs index f295f7d456..7e60f27d3e 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosCrossPartitionQueryExecutionContext.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosCrossPartitionQueryExecutionContext.cs @@ -431,129 +431,6 @@ protected async Task TryInitializeAsync( return TryCatch.FromResult(); } - public static TryCatch> TryGetInitializationInfo( - IReadOnlyList partitionKeyRanges, - IReadOnlyList partitionedContinuationTokens) - where PartitionedToken : IPartitionedToken - { - if (partitionKeyRanges == null) - { - throw new ArgumentNullException(nameof(partitionKeyRanges)); - } - - if (partitionedContinuationTokens == null) - { - throw new ArgumentNullException(nameof(partitionedContinuationTokens)); - } - - if (partitionKeyRanges.Count < 1) - { - throw new ArgumentException(nameof(partitionKeyRanges)); - } - - if (partitionedContinuationTokens.Count < 1) - { - throw new ArgumentException(nameof(partitionKeyRanges)); - } - - if (partitionedContinuationTokens.Count > partitionKeyRanges.Count) - { - throw new ArgumentException($"{nameof(partitionedContinuationTokens)} can not have more elements than {nameof(partitionKeyRanges)}."); - } - - // Find the continuation token for the partition we left off on: - PartitionedToken firstContinuationToken = partitionedContinuationTokens - .OrderBy((partitionedToken) => partitionedToken.PartitionRange.Min) - .First(); - - // Segment the ranges based off that: - ReadOnlyMemory sortedRanges = partitionKeyRanges - .OrderBy((partitionKeyRange) => partitionKeyRange.MinInclusive) - .ToArray(); - - PartitionKeyRange firstContinuationRange = new PartitionKeyRange - { - MinInclusive = firstContinuationToken.PartitionRange.Min, - MaxExclusive = firstContinuationToken.PartitionRange.Max - }; - - int matchedIndex = sortedRanges.Span.BinarySearch( - firstContinuationRange, - Comparer.Create((range1, range2) => string.CompareOrdinal(range1.MinInclusive, range2.MinInclusive))); - if (matchedIndex < 0) - { - return TryCatch>.FromException( - new MalformedContinuationTokenException( - $"{RMResources.InvalidContinuationToken} - Could not find continuation token: {firstContinuationToken}")); - } - - ReadOnlyMemory partitionsLeftOfTarget = matchedIndex == 0 ? ReadOnlyMemory.Empty : sortedRanges.Slice(start: 0, length: matchedIndex); - ReadOnlyMemory targetPartition = sortedRanges.Slice(start: matchedIndex, length: 1); - ReadOnlyMemory partitionsRightOfTarget = matchedIndex == sortedRanges.Length - 1 ? ReadOnlyMemory.Empty : sortedRanges.Slice(start: matchedIndex + 1); - - // Create the continuation token mapping for each region. - IReadOnlyDictionary mappingForPartitionsLeftOfTarget = MatchRangesToContinuationTokens( - partitionsLeftOfTarget, - partitionedContinuationTokens); - IReadOnlyDictionary mappingForTargetPartition = MatchRangesToContinuationTokens( - targetPartition, - partitionedContinuationTokens); - IReadOnlyDictionary mappingForPartitionsRightOfTarget = MatchRangesToContinuationTokens( - partitionsRightOfTarget, - partitionedContinuationTokens); - - return TryCatch>.FromResult( - new PartitionMapping( - partitionsLeftOfTarget: mappingForPartitionsLeftOfTarget, - targetPartition: mappingForTargetPartition, - partitionsRightOfTarget: mappingForPartitionsRightOfTarget)); - } - - /// - /// Matches ranges to their corresponding continuation token. - /// Note that most ranges don't have a corresponding continuation token, so their value will be set to null. - /// Also note that in the event of a split two or more ranges will match to the same continuation token. - /// - /// The type of token we are matching with. - /// The partition key ranges to match. - /// The continuation tokens to match with. - /// A dictionary of ranges matched with their continuation tokens. - public static IReadOnlyDictionary MatchRangesToContinuationTokens( - ReadOnlyMemory partitionKeyRanges, - IReadOnlyList partitionedContinuationTokens) - where PartitionedToken : IPartitionedToken - { - if (partitionedContinuationTokens == null) - { - throw new ArgumentNullException(nameof(partitionedContinuationTokens)); - } - - Dictionary partitionKeyRangeToToken = new Dictionary(); - ReadOnlySpan partitionKeyRangeSpan = partitionKeyRanges.Span; - for (int i = 0; i < partitionKeyRangeSpan.Length; i++) - { - PartitionKeyRange partitionKeyRange = partitionKeyRangeSpan[i]; - foreach (PartitionedToken partitionedToken in partitionedContinuationTokens) - { - // See if continuation token includes the range - if ((partitionKeyRange.MinInclusive.CompareTo(partitionedToken.PartitionRange.Min) >= 0) - && (partitionKeyRange.MaxExclusive.CompareTo(partitionedToken.PartitionRange.Max) <= 0)) - { - partitionKeyRangeToToken[partitionKeyRange] = partitionedToken; - break; - } - } - - if (!partitionKeyRangeToToken.ContainsKey(partitionKeyRange)) - { - // Could not find a matching token so just set it to null - partitionKeyRangeToToken[partitionKeyRange] = default; - } - } - - return partitionKeyRangeToToken; - } - protected virtual long GetAndResetResponseLengthBytes() { return Interlocked.Exchange(ref this.totalResponseLengthBytes, 0); @@ -623,36 +500,6 @@ private void OnItemProducerTreeCompleteFetching( public abstract CosmosElement GetCosmosElementContinuationToken(); - public readonly struct InitInfo - { - public InitInfo(int targetIndex, IReadOnlyDictionary continuationTokens) - { - this.TargetIndex = targetIndex; - this.ContinuationTokens = continuationTokens; - } - - public int TargetIndex { get; } - - public IReadOnlyDictionary ContinuationTokens { get; } - } - - public readonly struct PartitionMapping - { - public PartitionMapping( - IReadOnlyDictionary partitionsLeftOfTarget, - IReadOnlyDictionary targetPartition, - IReadOnlyDictionary partitionsRightOfTarget) - { - this.PartitionsLeftOfTarget = partitionsLeftOfTarget ?? throw new ArgumentNullException(nameof(partitionsLeftOfTarget)); - this.TargetPartition = targetPartition ?? throw new ArgumentNullException(nameof(targetPartition)); - this.PartitionsRightOfTarget = partitionsRightOfTarget ?? throw new ArgumentNullException(nameof(partitionsRightOfTarget)); - } - - public IReadOnlyDictionary PartitionsLeftOfTarget { get; } - public IReadOnlyDictionary TargetPartition { get; } - public IReadOnlyDictionary PartitionsRightOfTarget { get; } - } - /// /// All CrossPartitionQueries need this information on top of the parameter for DocumentQueryExecutionContextBase. /// I moved it out into it's own type, so that we don't have to keep passing around all the individual parameters in the factory pattern. diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/OrderBy/CosmosOrderByItemQueryExecutionContext.Resume.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/OrderBy/CosmosOrderByItemQueryExecutionContext.Resume.cs index 8a0bad89c3..21d64177ea 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/OrderBy/CosmosOrderByItemQueryExecutionContext.Resume.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/OrderBy/CosmosOrderByItemQueryExecutionContext.Resume.cs @@ -23,17 +23,6 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.OrderBy internal sealed partial class CosmosOrderByItemQueryExecutionContext : CosmosCrossPartitionQueryExecutionContext { - private const string TrueFilter = "true"; - - private static class Expressions - { - public const string LessThan = "<"; - public const string LessThanOrEqualTo = "<="; - public const string EqualTo = "="; - public const string GreaterThan = ">"; - public const string GreaterThanOrEqualTo = ">="; - } - public static async Task> MonadicCreateAsync( CosmosQueryContext queryContext, CosmosCrossPartitionQueryExecutionContext.CrossPartitionInitParams initParams, @@ -51,9 +40,6 @@ public static async Task> MonadicCrea cancellationToken.ThrowIfCancellationRequested(); - // TODO (brchon): For now we are not honoring non deterministic ORDER BY queries, since there is a bug in the continuation logic. - // We can turn it back on once the bug is fixed. - // This shouldn't hurt any query results. OrderByItemProducerTreeComparer orderByItemProducerTreeComparer = new OrderByItemProducerTreeComparer(initParams.PartitionedQueryExecutionInfo.QueryInfo.OrderBy.ToArray()); CosmosOrderByItemQueryExecutionContext context = new CosmosOrderByItemQueryExecutionContext( initPararms: queryContext, @@ -84,563 +70,5 @@ public static async Task> MonadicCrea cancellationToken: cancellationToken)) .Try(() => context); } - - private async Task TryInitializeAsync( - SqlQuerySpec sqlQuerySpec, - CosmosElement requestContinuation, - string collectionRid, - IReadOnlyList partitionKeyRanges, - int initialPageSize, - IReadOnlyList orderByColumns, - CancellationToken cancellationToken) - { - if (sqlQuerySpec == null) - { - throw new ArgumentNullException(nameof(sqlQuerySpec)); - } - - if (collectionRid == null) - { - throw new ArgumentNullException(nameof(collectionRid)); - } - - if (partitionKeyRanges == null) - { - throw new ArgumentNullException(nameof(partitionKeyRanges)); - } - - if (orderByColumns == null) - { - throw new ArgumentNullException(nameof(orderByColumns)); - } - - cancellationToken.ThrowIfCancellationRequested(); - - if (requestContinuation == null) - { - // Start off all the partition key ranges with null continuation - SqlQuerySpec rewrittenQueryForOrderBy = new SqlQuerySpec( - sqlQuerySpec.QueryText.Replace(oldValue: FormatPlaceHolder, newValue: True), - sqlQuerySpec.Parameters); - Dictionary partitionKeyRangeToContinuationToken = new Dictionary(); - foreach (PartitionKeyRange partitionKeyRange in partitionKeyRanges) - { - partitionKeyRangeToContinuationToken.Add(key: partitionKeyRange, value: null); - } - - return await base.TryInitializeAsync( - collectionRid, - initialPageSize, - rewrittenQueryForOrderBy, - partitionKeyRangeToContinuationToken, - deferFirstPage: false, - filter: null, - tryFilterAsync: null, - cancellationToken); - } - - TryCatch> tryGetOrderByContinuationTokenMapping = TryGetOrderByContinuationTokenMapping( - partitionKeyRanges, - requestContinuation, - orderByColumns.Count); - if (!tryGetOrderByContinuationTokenMapping.Succeeded) - { - return TryCatch.FromException(tryGetOrderByContinuationTokenMapping.Exception); - } - - IReadOnlyList orderByItems = tryGetOrderByContinuationTokenMapping - .Result - .TargetPartition - .Values - .First() - .OrderByItems - .Select(x => x.Item) - .ToList(); - if (orderByItems.Count != orderByColumns.Count) - { - return TryCatch.FromException( - new MalformedContinuationTokenException($"Order By Items from continuation token did not match the query text. Order by item count: {orderByItems.Count()} did not match column count {orderByColumns.Count()}. Continuation token: {requestContinuation}")); - } - - ReadOnlyMemory<(OrderByColumn, CosmosElement)> columnAndItems = orderByColumns.Zip(orderByItems, (column, item) => (column, item)).ToArray(); - - // For ascending order-by, left of target partition has filter expression > value, - // right of target partition has filter expression >= value, - // and target partition takes the previous filter from continuation (or true if no continuation) - (string leftFilter, string targetFilter, string rightFilter) = CosmosOrderByItemQueryExecutionContext.GetFormattedFilters(columnAndItems); - List<(IReadOnlyDictionary, string)> tokenMappingAndFilters = new List<(IReadOnlyDictionary, string)>() - { - { (tryGetOrderByContinuationTokenMapping.Result.PartitionsLeftOfTarget, leftFilter) }, - { (tryGetOrderByContinuationTokenMapping.Result.TargetPartition, targetFilter) }, - { (tryGetOrderByContinuationTokenMapping.Result.PartitionsRightOfTarget, rightFilter) }, - }; - - IReadOnlyList sortOrders = orderByColumns.Select(column => column.SortOrder).ToList(); - foreach ((IReadOnlyDictionary tokenMapping, string filter) in tokenMappingAndFilters) - { - SqlQuerySpec rewrittenQueryForOrderBy = new SqlQuerySpec( - sqlQuerySpec.QueryText.Replace(oldValue: FormatPlaceHolder, newValue: filter), - sqlQuerySpec.Parameters); - - TryCatch tryInitialize = await base.TryInitializeAsync( - collectionRid, - initialPageSize, - rewrittenQueryForOrderBy, - tokenMapping.ToDictionary(kvp => kvp.Key, kvp => kvp.Value?.CompositeContinuationToken.Token), - deferFirstPage: false, - filter, - tryFilterAsync: async (itemProducerTree) => - { - if (!tokenMapping.TryGetValue( - itemProducerTree.Root.PartitionKeyRange, - out OrderByContinuationToken continuationToken)) - { - throw new InvalidOperationException($"Failed to retrieve {nameof(OrderByContinuationToken)}."); - } - - if (continuationToken == null) - { - return TryCatch.FromResult(); - } - - return await this.TryFilterAsync( - itemProducerTree, - sortOrders, - continuationToken, - cancellationToken); - }, - cancellationToken); - if (!tryInitialize.Succeeded) - { - return tryInitialize; - } - } - - return TryCatch.FromResult(); - } - - private static TryCatch> TryGetOrderByContinuationTokenMapping( - IReadOnlyList partitionKeyRanges, - CosmosElement continuationToken, - int numOrderByItems) - { - if (partitionKeyRanges == null) - { - throw new ArgumentOutOfRangeException(nameof(partitionKeyRanges)); - } - - if (numOrderByItems < 0) - { - throw new ArgumentOutOfRangeException(nameof(numOrderByItems)); - } - - if (continuationToken == null) - { - throw new ArgumentNullException(nameof(continuationToken)); - } - - TryCatch> tryExtractContinuationTokens = TryExtractContinuationTokens(continuationToken, numOrderByItems); - if (!tryExtractContinuationTokens.Succeeded) - { - return TryCatch>.FromException(tryExtractContinuationTokens.Exception); - } - - return TryGetInitializationInfo( - partitionKeyRanges, - tryExtractContinuationTokens.Result); - } - - private static TryCatch> TryExtractContinuationTokens( - CosmosElement requestContinuation, - int numOrderByItems) - { - if (requestContinuation == null) - { - throw new ArgumentNullException("continuation can not be null or empty."); - } - - if (numOrderByItems < 0) - { - throw new ArgumentOutOfRangeException(nameof(numOrderByItems)); - } - - if (!(requestContinuation is CosmosArray cosmosArray)) - { - return TryCatch>.FromException( - new MalformedContinuationTokenException( - $"Order by continuation token must be an array: {requestContinuation}.")); - } - - List orderByContinuationTokens = new List(); - foreach (CosmosElement arrayItem in cosmosArray) - { - TryCatch tryCreateOrderByContinuationToken = OrderByContinuationToken.TryCreateFromCosmosElement(arrayItem); - if (!tryCreateOrderByContinuationToken.Succeeded) - { - return TryCatch>.FromException(tryCreateOrderByContinuationToken.Exception); - } - - orderByContinuationTokens.Add(tryCreateOrderByContinuationToken.Result); - } - - if (orderByContinuationTokens.Count == 0) - { - return TryCatch>.FromException( - new MalformedContinuationTokenException( - $"Order by continuation token cannot be empty: {requestContinuation}.")); - } - - foreach (OrderByContinuationToken suppliedOrderByContinuationToken in orderByContinuationTokens) - { - if (suppliedOrderByContinuationToken.OrderByItems.Count != numOrderByItems) - { - return TryCatch>.FromException( - new MalformedContinuationTokenException( - $"Invalid order-by items in continuation token {requestContinuation} for OrderBy~Context.")); - } - } - - return TryCatch>.FromResult(orderByContinuationTokens); - } - - /// - /// When resuming an order by query we need to filter the document producers. - /// - /// The producer to filter down. - /// The sort orders. - /// The continuation token. - /// The cancellation token. - /// A task to await on. - private async Task TryFilterAsync( - ItemProducerTree producer, - IReadOnlyList sortOrders, - OrderByContinuationToken continuationToken, - CancellationToken cancellationToken) - { - cancellationToken.ThrowIfCancellationRequested(); - // When we resume a query on a partition there is a possibility that we only read a partial page from the backend - // meaning that will we repeat some documents if we didn't do anything about it. - // The solution is to filter all the documents that come before in the sort order, since we have already emitted them to the client. - // The key is to seek until we get an order by value that matches the order by value we left off on. - // Once we do that we need to seek to the correct _rid within the term, - // since there might be many documents with the same order by value we left off on. - - foreach (ItemProducerTree tree in producer) - { - if (!ResourceId.TryParse(continuationToken.Rid, out ResourceId continuationRid)) - { - return TryCatch.FromException( - new MalformedContinuationTokenException( - $"Invalid Rid in the continuation token {continuationToken.CompositeContinuationToken.Token} for OrderBy~Context.")); - } - - Dictionary resourceIds = new Dictionary(); - int itemToSkip = continuationToken.SkipCount; - bool continuationRidVerified = false; - - while (true) - { - if (tree.Current == null) - { - // This document producer doesn't have anymore items. - break; - } - - OrderByQueryResult orderByResult = new OrderByQueryResult(tree.Current); - // Throw away documents until it matches the item from the continuation token. - int cmp = 0; - for (int i = 0; i < sortOrders.Count; ++i) - { - cmp = ItemComparer.Instance.Compare( - continuationToken.OrderByItems[i].Item, - orderByResult.OrderByItems[i].Item); - - if (cmp != 0) - { - cmp = sortOrders[i] == SortOrder.Ascending ? cmp : -cmp; - break; - } - } - - if (cmp < 0) - { - // We might have passed the item due to deletions and filters. - break; - } - - if (cmp == 0) - { - if (!resourceIds.TryGetValue(orderByResult.Rid, out ResourceId rid)) - { - if (!ResourceId.TryParse(orderByResult.Rid, out rid)) - { - return TryCatch.FromException( - new MalformedContinuationTokenException( - $"Invalid Rid in the continuation token {continuationToken.CompositeContinuationToken.Token} for OrderBy~Context~TryParse.")); - } - - resourceIds.Add(orderByResult.Rid, rid); - } - - if (!continuationRidVerified) - { - if (continuationRid.Database != rid.Database || continuationRid.DocumentCollection != rid.DocumentCollection) - { - return TryCatch.FromException( - new MalformedContinuationTokenException( - $"Invalid Rid in the continuation token {continuationToken.CompositeContinuationToken.Token} for OrderBy~Context.")); - } - - continuationRidVerified = true; - } - - // Once the item matches the order by items from the continuation tokens - // We still need to remove all the documents that have a lower rid in the rid sort order. - // If there is a tie in the sort order the documents should be in _rid order in the same direction as the index (given by the backend) - cmp = continuationRid.Document.CompareTo(rid.Document); - if ((producer.CosmosQueryExecutionInfo == null) || producer.CosmosQueryExecutionInfo.ReverseRidEnabled) - { - // If reverse rid is enabled on the backend then fallback to the old way of doing it. - if (sortOrders[0] == SortOrder.Descending) - { - cmp = -cmp; - } - } - else - { - // Go by the whatever order the index wants - if (producer.CosmosQueryExecutionInfo.ReverseIndexScan) - { - cmp = -cmp; - } - } - - // We might have passed the item due to deletions and filters. - // We also have a skip count for JOINs - if (cmp < 0 || (cmp == 0 && itemToSkip-- <= 0)) - { - break; - } - } - - if (!tree.TryMoveNextDocumentWithinPage()) - { - while (true) - { - (bool successfullyMovedNext, QueryResponseCore? failureResponse) = await tree.TryMoveNextPageAsync(cancellationToken); - if (!successfullyMovedNext) - { - if (failureResponse.HasValue) - { - return TryCatch.FromException( - failureResponse.Value.CosmosException); - } - - break; - } - - if (tree.IsAtBeginningOfPage) - { - break; - } - - if (tree.TryMoveNextDocumentWithinPage()) - { - break; - } - } - } - } - } - - return TryCatch.FromResult(); - } - - private static void AppendToBuilders((StringBuilder leftFilter, StringBuilder targetFilter, StringBuilder rightFilter) builders, object str) - { - CosmosOrderByItemQueryExecutionContext.AppendToBuilders(builders, str, str, str); - } - - private static void AppendToBuilders((StringBuilder leftFilter, StringBuilder targetFilter, StringBuilder rightFilter) builders, object left, object target, object right) - { - builders.leftFilter.Append(left); - builders.targetFilter.Append(target); - builders.rightFilter.Append(right); - } - - private static (string leftFilter, string targetFilter, string rightFilter) GetFormattedFilters( - ReadOnlyMemory<(OrderByColumn orderByColumn, CosmosElement orderByItem)> columnAndItems) - { - // When we run cross partition queries, - // we only serialize the continuation token for the partition that we left off on. - // The only problem is that when we resume the order by query, - // we don't have continuation tokens for all other partition. - // The saving grace is that the data has a composite sort order(query sort order, partition key range id) - // so we can generate range filters which in turn the backend will turn into rid based continuation tokens, - // which is enough to get the streams of data flowing from all partitions. - // The details of how this is done is described below: - int numOrderByItems = columnAndItems.Length; - bool isSingleOrderBy = numOrderByItems == 1; - StringBuilder left = new StringBuilder(); - StringBuilder target = new StringBuilder(); - StringBuilder right = new StringBuilder(); - - (StringBuilder, StringBuilder, StringBuilder) builders = (left, target, right); - - if (isSingleOrderBy) - { - //For a single order by query we resume the continuations in this manner - // Suppose the query is SELECT* FROM c ORDER BY c.string ASC - // And we left off on partition N with the value "B" - // Then - // All the partitions to the left will have finished reading "B" - // Partition N is still reading "B" - // All the partitions to the right have let to read a "B - // Therefore the filters should be - // > "B" , >= "B", and >= "B" respectively - // Repeat the same logic for DESC and you will get - // < "B", <= "B", and <= "B" respectively - // The general rule becomes - // For ASC - // > for partitions to the left - // >= for the partition we left off on - // >= for the partitions to the right - // For DESC - // < for partitions to the left - // <= for the partition we left off on - // <= for the partitions to the right - (OrderByColumn orderByColumn, CosmosElement orderByItem) = columnAndItems.Span[0]; - (string expression, SortOrder sortOrder) = (orderByColumn.Expression, orderByColumn.SortOrder); - - StringBuilder sb = new StringBuilder(); - CosmosElementToQueryLiteral cosmosElementToQueryLiteral = new CosmosElementToQueryLiteral(sb); - orderByItem.Accept(cosmosElementToQueryLiteral); - - string orderByItemToString = sb.ToString(); - - left.Append($"{expression} {(sortOrder == SortOrder.Descending ? Expressions.LessThan : Expressions.GreaterThan)} {orderByItemToString}"); - target.Append($"{expression} {(sortOrder == SortOrder.Descending ? Expressions.LessThanOrEqualTo : Expressions.GreaterThanOrEqualTo)} {orderByItemToString}"); - right.Append($"{expression} {(sortOrder == SortOrder.Descending ? Expressions.LessThanOrEqualTo : Expressions.GreaterThanOrEqualTo)} {orderByItemToString}"); - } - else - { - //For a multi order by query - // Suppose the query is SELECT* FROM c ORDER BY c.string ASC, c.number ASC - // And we left off on partition N with the value("A", 1) - // Then - // All the partitions to the left will have finished reading("A", 1) - // Partition N is still reading("A", 1) - // All the partitions to the right have let to read a "(A", 1) - // The filters are harder to derive since their are multiple columns - // But the problem reduces to "How do you know one document comes after another in a multi order by query" - // The answer is to just look at it one column at a time. - // For this particular scenario: - // If a first column is greater ex. ("B", blah), then the document comes later in the sort order - // Therefore we want all documents where the first column is greater than "A" which means > "A" - // Or if the first column is a tie, then you look at the second column ex. ("A", blah). - // Therefore we also want all documents where the first column was a tie but the second column is greater which means = "A" AND > 1 - // Therefore the filters should be - // (> "A") OR (= "A" AND > 1), (> "A") OR (= "A" AND >= 1), (> "A") OR (= "A" AND >= 1) - // Notice that if we repeated the same logic we for single order by we would have gotten - // > "A" AND > 1, >= "A" AND >= 1, >= "A" AND >= 1 - // which is wrong since we missed some documents - // Repeat the same logic for ASC, DESC - // (> "A") OR (= "A" AND < 1), (> "A") OR (= "A" AND <= 1), (> "A") OR (= "A" AND <= 1) - // Again for DESC, ASC - // (< "A") OR (= "A" AND > 1), (< "A") OR (= "A" AND >= 1), (< "A") OR (= "A" AND >= 1) - // And again for DESC DESC - // (< "A") OR (= "A" AND < 1), (< "A") OR (= "A" AND <= 1), (< "A") OR (= "A" AND <= 1) - // The general we look at all prefixes of the order by columns to look for tie breakers. - // Except for the full prefix whose last column follows the rules for single item order by - // And then you just OR all the possibilities together - for (int prefixLength = 1; prefixLength <= numOrderByItems; prefixLength++) - { - ReadOnlySpan<(OrderByColumn orderByColumn, CosmosElement orderByItem)> columnAndItemPrefix = columnAndItems.Span.Slice(start: 0, length: prefixLength); - - bool lastPrefix = prefixLength == numOrderByItems; - bool firstPrefix = prefixLength == 1; - - CosmosOrderByItemQueryExecutionContext.AppendToBuilders(builders, "("); - - for (int index = 0; index < prefixLength; index++) - { - string expression = columnAndItemPrefix[index].orderByColumn.Expression; - SortOrder sortOrder = columnAndItemPrefix[index].orderByColumn.SortOrder; - CosmosElement orderByItem = columnAndItemPrefix[index].orderByItem; - bool lastItem = index == prefixLength - 1; - - // Append Expression - CosmosOrderByItemQueryExecutionContext.AppendToBuilders(builders, expression); - CosmosOrderByItemQueryExecutionContext.AppendToBuilders(builders, " "); - - // Append binary operator - if (lastItem) - { - string inequality = sortOrder == SortOrder.Descending ? Expressions.LessThan : Expressions.GreaterThan; - CosmosOrderByItemQueryExecutionContext.AppendToBuilders(builders, inequality); - if (lastPrefix) - { - CosmosOrderByItemQueryExecutionContext.AppendToBuilders(builders, string.Empty, Expressions.EqualTo, Expressions.EqualTo); - } - } - else - { - CosmosOrderByItemQueryExecutionContext.AppendToBuilders(builders, Expressions.EqualTo); - } - - // Append SortOrder - StringBuilder sb = new StringBuilder(); - CosmosElementToQueryLiteral cosmosElementToQueryLiteral = new CosmosElementToQueryLiteral(sb); - orderByItem.Accept(cosmosElementToQueryLiteral); - string orderByItemToString = sb.ToString(); - CosmosOrderByItemQueryExecutionContext.AppendToBuilders(builders, " "); - CosmosOrderByItemQueryExecutionContext.AppendToBuilders(builders, orderByItemToString); - CosmosOrderByItemQueryExecutionContext.AppendToBuilders(builders, " "); - - if (!lastItem) - { - CosmosOrderByItemQueryExecutionContext.AppendToBuilders(builders, "AND "); - } - } - - CosmosOrderByItemQueryExecutionContext.AppendToBuilders(builders, ")"); - if (!lastPrefix) - { - CosmosOrderByItemQueryExecutionContext.AppendToBuilders(builders, " OR "); - } - } - } - - // For the target filter we can make an optimization to just return "true", - // since we already have the backend continuation token to resume with. - return (left.ToString(), TrueFilter, right.ToString()); - } - - private readonly struct OrderByInitInfo - { - public OrderByInitInfo( - RangeFilterInitializationInfo[] filters, - IReadOnlyDictionary continuationTokens) - { - this.Filters = filters; - this.ContinuationTokens = continuationTokens; - } - - public RangeFilterInitializationInfo[] Filters { get; } - - public IReadOnlyDictionary ContinuationTokens { get; } - } - - private readonly struct OrderByColumn - { - public OrderByColumn(string expression, SortOrder sortOrder) - { - this.Expression = expression ?? throw new ArgumentNullException(nameof(expression)); - this.SortOrder = sortOrder; - } - - public string Expression { get; } - public SortOrder SortOrder { get; } - } } } diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/OrderBy/CosmosOrderByItemQueryExecutionContext.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/OrderBy/CosmosOrderByItemQueryExecutionContext.cs index 8debd89d5a..aa542a1e19 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/OrderBy/CosmosOrderByItemQueryExecutionContext.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/OrderBy/CosmosOrderByItemQueryExecutionContext.cs @@ -10,28 +10,8 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.OrderBy using Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.ItemProducers; using Microsoft.Azure.Cosmos.Query.Core.QueryClient; - /// - /// CosmosOrderByItemQueryExecutionContext is a concrete implementation for CrossPartitionQueryExecutionContext. - /// This class is responsible for draining cross partition queries that have order by conditions. - /// The way order by queries work is that they are doing a k-way merge of sorted lists from each partition with an added condition. - /// The added condition is that if 2 or more top documents from different partitions are equivalent then we drain from the left most partition first. - /// This way we can generate a single continuation token for all n partitions. - /// This class is able to stop and resume execution by generating continuation tokens and reconstructing an execution context from said token. - /// internal sealed partial class CosmosOrderByItemQueryExecutionContext { - /// - /// Order by queries are rewritten to allow us to inject a filter. - /// This placeholder is so that we can just string replace it with the filter we want without having to understand the structure of the query. - /// - private const string FormatPlaceHolder = "{documentdb-formattableorderbyquery-filter}"; - - /// - /// If query does not need a filter then we replace the FormatPlaceHolder with "true", since - /// "SELECT * FROM c WHERE blah and true" is the same as "SELECT * FROM c where blah" - /// - private const string True = "true"; - /// /// Function to determine the priority of fetches. /// Basically we are fetching from the partition with the least number of buffered documents first. diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/Parallel/CosmosParallelItemQueryExecutionContext.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/Parallel/CosmosParallelItemQueryExecutionContext.cs index 56a2f6225e..bca7f6fe39 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/Parallel/CosmosParallelItemQueryExecutionContext.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/Parallel/CosmosParallelItemQueryExecutionContext.cs @@ -9,13 +9,6 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.Parallel using Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.ItemProducers; using Microsoft.Azure.Cosmos.Query.Core.QueryClient; - /// - /// CosmosParallelItemQueryExecutionContext is a concrete implementation for CrossPartitionQueryExecutionContext. - /// This class is responsible for draining cross partition queries that do not have order by conditions. - /// The way parallel queries work is that it drains from the left most partition first. - /// This class handles draining in the correct order and can also stop and resume the query - /// by generating a continuation token and resuming from said continuation token. - /// internal sealed partial class CosmosParallelItemQueryExecutionContext : CosmosCrossPartitionQueryExecutionContext { /// diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/CrossPartitionQueryPipelineStage.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/CrossPartitionQueryPipelineStage.cs new file mode 100644 index 0000000000..8eaa0d089a --- /dev/null +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/CrossPartitionQueryPipelineStage.cs @@ -0,0 +1,157 @@ +// ------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +// ------------------------------------------------------------ + +namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote +{ + using System; + using System.Collections.Generic; + using System.Linq; + using Microsoft.Azure.Cosmos.Query.Core.ContinuationTokens; + using Microsoft.Azure.Cosmos.Query.Core.Exceptions; + using Microsoft.Azure.Cosmos.Query.Core.Monads; + using Microsoft.Azure.Documents; + + internal abstract class CrossPartitionQueryPipelineStage : IQueryPipelineStage + { + protected static TryCatch> MonadicGetPartitionMapping( + IReadOnlyList partitionKeyRanges, + IReadOnlyList partitionedContinuationTokens) + where PartitionedToken : IPartitionedToken + { + if (partitionKeyRanges == null) + { + throw new ArgumentNullException(nameof(partitionKeyRanges)); + } + + if (partitionedContinuationTokens == null) + { + throw new ArgumentNullException(nameof(partitionedContinuationTokens)); + } + + if (partitionKeyRanges.Count < 1) + { + throw new ArgumentException(nameof(partitionKeyRanges)); + } + + if (partitionedContinuationTokens.Count < 1) + { + throw new ArgumentException(nameof(partitionKeyRanges)); + } + + if (partitionedContinuationTokens.Count > partitionKeyRanges.Count) + { + throw new ArgumentException($"{nameof(partitionedContinuationTokens)} can not have more elements than {nameof(partitionKeyRanges)}."); + } + + // Find the continuation token for the partition we left off on: + PartitionedToken firstContinuationToken = partitionedContinuationTokens + .OrderBy((partitionedToken) => partitionedToken.PartitionRange.Min) + .First(); + + // Segment the ranges based off that: + ReadOnlyMemory sortedRanges = partitionKeyRanges + .OrderBy((partitionKeyRange) => partitionKeyRange.MinInclusive) + .ToArray(); + + PartitionKeyRange firstContinuationRange = new PartitionKeyRange + { + MinInclusive = firstContinuationToken.PartitionRange.Min, + MaxExclusive = firstContinuationToken.PartitionRange.Max + }; + + int matchedIndex = sortedRanges.Span.BinarySearch( + firstContinuationRange, + Comparer.Create((range1, range2) => string.CompareOrdinal(range1.MinInclusive, range2.MinInclusive))); + if (matchedIndex < 0) + { + return TryCatch>.FromException( + new MalformedContinuationTokenException( + $"{RMResources.InvalidContinuationToken} - Could not find continuation token: {firstContinuationToken}")); + } + + ReadOnlyMemory partitionsLeftOfTarget = matchedIndex == 0 ? ReadOnlyMemory.Empty : sortedRanges.Slice(start: 0, length: matchedIndex); + ReadOnlyMemory targetPartition = sortedRanges.Slice(start: matchedIndex, length: 1); + ReadOnlyMemory partitionsRightOfTarget = matchedIndex == sortedRanges.Length - 1 ? ReadOnlyMemory.Empty : sortedRanges.Slice(start: matchedIndex + 1); + + // Create the continuation token mapping for each region. + IReadOnlyDictionary mappingForPartitionsLeftOfTarget = MatchRangesToContinuationTokens( + partitionsLeftOfTarget, + partitionedContinuationTokens); + IReadOnlyDictionary mappingForTargetPartition = MatchRangesToContinuationTokens( + targetPartition, + partitionedContinuationTokens); + IReadOnlyDictionary mappingForPartitionsRightOfTarget = MatchRangesToContinuationTokens( + partitionsRightOfTarget, + partitionedContinuationTokens); + + return TryCatch>.FromResult( + new PartitionMapping( + partitionsLeftOfTarget: mappingForPartitionsLeftOfTarget, + targetPartition: mappingForTargetPartition, + partitionsRightOfTarget: mappingForPartitionsRightOfTarget)); + } + + /// + /// Matches ranges to their corresponding continuation token. + /// Note that most ranges don't have a corresponding continuation token, so their value will be set to null. + /// Also note that in the event of a split two or more ranges will match to the same continuation token. + /// + /// The type of token we are matching with. + /// The partition key ranges to match. + /// The continuation tokens to match with. + /// A dictionary of ranges matched with their continuation tokens. + protected static IReadOnlyDictionary MatchRangesToContinuationTokens( + ReadOnlyMemory partitionKeyRanges, + IReadOnlyList partitionedContinuationTokens) + where PartitionedToken : IPartitionedToken + { + if (partitionedContinuationTokens == null) + { + throw new ArgumentNullException(nameof(partitionedContinuationTokens)); + } + + Dictionary partitionKeyRangeToToken = new Dictionary(); + ReadOnlySpan partitionKeyRangeSpan = partitionKeyRanges.Span; + for (int i = 0; i < partitionKeyRangeSpan.Length; i++) + { + PartitionKeyRange partitionKeyRange = partitionKeyRangeSpan[i]; + foreach (PartitionedToken partitionedToken in partitionedContinuationTokens) + { + // See if continuation token includes the range + if ((partitionKeyRange.MinInclusive.CompareTo(partitionedToken.PartitionRange.Min) >= 0) + && (partitionKeyRange.MaxExclusive.CompareTo(partitionedToken.PartitionRange.Max) <= 0)) + { + partitionKeyRangeToToken[partitionKeyRange] = partitionedToken; + break; + } + } + + if (!partitionKeyRangeToToken.ContainsKey(partitionKeyRange)) + { + // Could not find a matching token so just set it to null + partitionKeyRangeToToken[partitionKeyRange] = default; + } + } + + return partitionKeyRangeToToken; + } + + protected readonly struct PartitionMapping + { + public PartitionMapping( + IReadOnlyDictionary partitionsLeftOfTarget, + IReadOnlyDictionary targetPartition, + IReadOnlyDictionary partitionsRightOfTarget) + { + this.PartitionsLeftOfTarget = partitionsLeftOfTarget ?? throw new ArgumentNullException(nameof(partitionsLeftOfTarget)); + this.TargetPartition = targetPartition ?? throw new ArgumentNullException(nameof(targetPartition)); + this.PartitionsRightOfTarget = partitionsRightOfTarget ?? throw new ArgumentNullException(nameof(partitionsRightOfTarget)); + } + + public IReadOnlyDictionary PartitionsLeftOfTarget { get; } + public IReadOnlyDictionary TargetPartition { get; } + public IReadOnlyDictionary PartitionsRightOfTarget { get; } + } + } +} diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderByCrossPartitionQueryPipelineStage.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderByCrossPartitionQueryPipelineStage.cs index 0c3300d8ca..a110e6aa2b 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderByCrossPartitionQueryPipelineStage.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderByCrossPartitionQueryPipelineStage.cs @@ -6,17 +6,65 @@ namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote { using System; using System.Collections.Generic; + using System.Linq; + using System.Text; + using System.Threading; using System.Threading.Tasks; using Microsoft.Azure.Cosmos.CosmosElements; using Microsoft.Azure.Cosmos.Pagination; + using Microsoft.Azure.Cosmos.Query.Core.Collections; using Microsoft.Azure.Cosmos.Query.Core.ContinuationTokens; using Microsoft.Azure.Cosmos.Query.Core.Exceptions; + using Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.ItemProducers; + using Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.OrderBy; using Microsoft.Azure.Cosmos.Query.Core.Monads; + using Microsoft.Azure.Documents; - internal sealed class OrderByCrossPartitionQueryPipelineStage : IQueryPipelineStage + /// + /// CosmosOrderByItemQueryExecutionContext is a concrete implementation for CrossPartitionQueryExecutionContext. + /// This class is responsible for draining cross partition queries that have order by conditions. + /// The way order by queries work is that they are doing a k-way merge of sorted lists from each partition with an added condition. + /// The added condition is that if 2 or more top documents from different partitions are equivalent then we drain from the left most partition first. + /// This way we can generate a single continuation token for all n partitions. + /// This class is able to stop and resume execution by generating continuation tokens and reconstructing an execution context from said token. + /// + internal sealed class OrderByCrossPartitionQueryPipelineStage : CrossPartitionQueryPipelineStage { + private static class Expressions + { + public const string LessThan = "<"; + public const string LessThanOrEqualTo = "<="; + public const string EqualTo = "="; + public const string GreaterThan = ">"; + public const string GreaterThanOrEqualTo = ">="; + } + + /// + /// Order by queries are rewritten to allow us to inject a filter. + /// This placeholder is so that we can just string replace it with the filter we want without having to understand the structure of the query. + /// + private const string FormatPlaceHolder = "{documentdb-formattableorderbyquery-filter}"; + + /// + /// If query does not need a filter then we replace the FormatPlaceHolder with "true", since + /// "SELECT * FROM c WHERE blah and true" is the same as "SELECT * FROM c where blah" + /// + private const string TrueFilter = "true"; + + /// + /// For ORDER BY queries we need to drain the first page of every enumerator before adding it to the prioirty queue for k-way merge. + /// Since this requires async work, we defer it until the user calls MoveNextAsync(). + /// + private readonly Queue<(OrderByQueryPartitionRangePageAsyncEnumerator, OrderByContinuationToken)> uninitializedEnumerators; + + private readonly PriorityQueue initializedEnumerators; + private readonly CrossPartitionRangePageAsyncEnumerator crossPartitionRangePageAsyncEnumerator; + private readonly List sortOrders; + + private readonly IDocumentContainer documentContainer; + private OrderByCrossPartitionQueryPipelineStage( CrossPartitionRangePageAsyncEnumerator crossPartitionRangePageAsyncEnumerator) { @@ -27,90 +75,570 @@ private OrderByCrossPartitionQueryPipelineStage( public ValueTask DisposeAsync() => this.crossPartitionRangePageAsyncEnumerator.DisposeAsync(); - public ValueTask MoveNextAsync() + public async ValueTask MoveNextAsync() { - throw new NotImplementedException(); + if (this.uninitializedEnumerators.Count != 0) + { + (OrderByQueryPartitionRangePageAsyncEnumerator enumerator, OrderByContinuationToken token) = this.uninitializedEnumerators.Dequeue(); + if (token != null) + { + TryCatch monadicFilterAsync = await OrderByCrossPartitionQueryPipelineStage.MonadicFilterAsync( + enumerator, + this.sortOrders, + token, + cancellationToken: default); + if (monadicFilterAsync.Failed) + { + // Check if it's a retryable exception. + Exception exception = monadicFilterAsync.Exception; + while (exception.InnerException != null) + { + exception = exception.InnerException; + } + + if (IsSplitException(exception)) + { + // Handle split + IEnumerable childRanges = await this.documentContainer.GetChildRangeAsync( + enumerator.Range, + cancellationToken: default); + foreach (PartitionKeyRange childRange in childRanges) + { + OrderByQueryPartitionRangePageAsyncEnumerator childPaginator = new OrderByQueryPartitionRangePageAsyncEnumerator( + documentContainer, + sqlQuerySpec, + range, + pageSize, + state: default) + enumerators.Enqueue(childPaginator); + } + + // Recursively retry + return await this.MoveNextAsync(); + } + } + } + else + { + + } + } } public static TryCatch MonadicCreate( IDocumentContainer documentContainer, SqlQuerySpec sqlQuerySpec, + IReadOnlyList targetRanges, + IReadOnlyList orderByColumns, int pageSize, CosmosElement continuationToken) { + // TODO (brchon): For now we are not honoring non deterministic ORDER BY queries, since there is a bug in the continuation logic. + // We can turn it back on once the bug is fixed. + // This shouldn't hurt any query results. + + if (documentContainer == null) + { + throw new ArgumentNullException(nameof(documentContainer)); + } + + if (sqlQuerySpec == null) + { + throw new ArgumentNullException(nameof(sqlQuerySpec)); + } + + if (targetRanges == null) + { + throw new ArgumentNullException(nameof(targetRanges)); + } + + if (targetRanges.Count == 0) + { + throw new ArgumentException($"{nameof(targetRanges)} must not be empty."); + } + + if (orderByColumns == null) + { + throw new ArgumentNullException(nameof(orderByColumns)); + } + + if (orderByColumns.Count == 0) + { + throw new ArgumentException($"{nameof(orderByColumns)} must not be empty."); + } + if (pageSize <= 0) { throw new ArgumentOutOfRangeException(nameof(pageSize)); } - TryCatch> monadicExtractState = MonadicExtractState(continuationToken); - if (monadicExtractState.Failed) + List remoteEnumerators; + if (continuationToken == null) + { + // Start off all the partition key ranges with null continuation + SqlQuerySpec rewrittenQueryForOrderBy = new SqlQuerySpec( + sqlQuerySpec.QueryText.Replace(oldValue: FormatPlaceHolder, newValue: TrueFilter), + sqlQuerySpec.Parameters); + + remoteEnumerators = targetRanges + .Select(range => new OrderByQueryPartitionRangePageAsyncEnumerator( + documentContainer, + sqlQuerySpec, + range, + pageSize, + state: default)) + .ToList(); + } + else + { + TryCatch> monadicGetOrderByContinuationTokenMapping = MonadicGetOrderByContinuationTokenMapping( + targetRanges, + continuationToken, + orderByColumns.Count); + if (monadicGetOrderByContinuationTokenMapping.Failed) + { + return TryCatch.FromException(monadicGetOrderByContinuationTokenMapping.Exception); + } + + PartitionMapping partitionMapping = monadicGetOrderByContinuationTokenMapping.Result; + + IReadOnlyList orderByItems = partitionMapping + .TargetPartition + .Values + .First() + .OrderByItems + .Select(x => x.Item) + .ToList(); + if (orderByItems.Count != orderByColumns.Count) + { + return TryCatch.FromException( + new MalformedContinuationTokenException( + $"Order By Items from continuation token did not match the query text. " + + $"Order by item count: {orderByItems.Count()} did not match column count {orderByColumns.Count()}. " + + $"Continuation token: {continuationToken}")); + } + + ReadOnlyMemory<(OrderByColumn, CosmosElement)> columnAndItems = orderByColumns.Zip(orderByItems, (column, item) => (column, item)).ToArray(); + + // For ascending order-by, left of target partition has filter expression > value, + // right of target partition has filter expression >= value, + // and target partition takes the previous filter from continuation (or true if no continuation) + (string leftFilter, string targetFilter, string rightFilter) = OrderByCrossPartitionQueryPipelineStage.GetFormattedFilters(columnAndItems); + List<(IReadOnlyDictionary, string)> tokenMappingAndFilters = new List<(IReadOnlyDictionary, string)>() + { + { (partitionMapping.PartitionsLeftOfTarget, leftFilter) }, + { (partitionMapping.TargetPartition, targetFilter) }, + { (partitionMapping.PartitionsRightOfTarget, rightFilter) }, + }; + + IReadOnlyList sortOrders = orderByColumns.Select(column => column.SortOrder).ToList(); + remoteEnumerators = new List(); + foreach ((IReadOnlyDictionary tokenMapping, string filter) in tokenMappingAndFilters) + { + SqlQuerySpec rewrittenQueryForOrderBy = new SqlQuerySpec( + sqlQuerySpec.QueryText.Replace(oldValue: FormatPlaceHolder, newValue: filter), + sqlQuerySpec.Parameters); + + foreach (KeyValuePair kvp in tokenMapping) + { + PartitionKeyRange range = kvp.Key; + OrderByContinuationToken token = kvp.Value; + OrderByQueryPartitionRangePageAsyncEnumerator remoteEnumerator = new OrderByQueryPartitionRangePageAsyncEnumerator( + documentContainer, + sqlQuerySpec, + range, + pageSize, + state: token != null ? new QueryState(CosmosString.Create(token.CompositeContinuationToken.Token)) : null); + + remoteEnumerators.Add(remoteEnumerator); + } + } + } + } + + private static TryCatch> MonadicGetOrderByContinuationTokenMapping( + IReadOnlyList partitionKeyRanges, + CosmosElement continuationToken, + int numOrderByItems) + { + if (partitionKeyRanges == null) + { + throw new ArgumentOutOfRangeException(nameof(partitionKeyRanges)); + } + + if (numOrderByItems < 0) { - return TryCatch.FromException(monadicExtractState.Exception); + throw new ArgumentOutOfRangeException(nameof(numOrderByItems)); } - CrossPartitionState state = monadicExtractState.Result; + if (continuationToken == null) + { + throw new ArgumentNullException(nameof(continuationToken)); + } + TryCatch> monadicExtractContinuationTokens = MonadicExtractOrderByTokens(continuationToken, numOrderByItems); + if (monadicExtractContinuationTokens.Failed) + { + return TryCatch>.FromException(monadicExtractContinuationTokens.Exception); + } + return MonadicGetPartitionMapping( + partitionKeyRanges, + monadicExtractContinuationTokens.Result); } - private static TryCatch> MonadicExtractState( + private static TryCatch> MonadicExtractOrderByTokens( CosmosElement continuationToken, int numOrderByColumns) { if (continuationToken == null) { - return TryCatch>.FromResult(default); + return TryCatch>.FromResult(default); } if (!(continuationToken is CosmosArray cosmosArray)) { - return TryCatch>.FromException( + return TryCatch>.FromException( new MalformedContinuationTokenException( $"Order by continuation token must be an array: {continuationToken}.")); } + if (cosmosArray.Count == 0) + { + return TryCatch>.FromException( + new MalformedContinuationTokenException( + $"Order by continuation token cannot be empty: {continuationToken}.")); + } + List orderByContinuationTokens = new List(); foreach (CosmosElement arrayItem in cosmosArray) { TryCatch tryCreateOrderByContinuationToken = OrderByContinuationToken.TryCreateFromCosmosElement(arrayItem); if (!tryCreateOrderByContinuationToken.Succeeded) { - return TryCatch>.FromException(tryCreateOrderByContinuationToken.Exception); + return TryCatch>.FromException(tryCreateOrderByContinuationToken.Exception); } orderByContinuationTokens.Add(tryCreateOrderByContinuationToken.Result); } - if (orderByContinuationTokens.Count == 0) - { - return TryCatch>.FromException( - new MalformedContinuationTokenException( - $"Order by continuation token cannot be empty: {continuationToken}.")); - } - foreach (OrderByContinuationToken suppliedOrderByContinuationToken in orderByContinuationTokens) { if (suppliedOrderByContinuationToken.OrderByItems.Count != numOrderByColumns) { - return TryCatch>.FromException( + return TryCatch>.FromException( new MalformedContinuationTokenException( $"Invalid order-by items in continuation token {continuationToken} for OrderBy~Context.")); } } + + return TryCatch>.FromResult(orderByContinuationTokens); + } + + private static void AppendToBuilders((StringBuilder leftFilter, StringBuilder targetFilter, StringBuilder rightFilter) builders, object str) + { + OrderByCrossPartitionQueryPipelineStage.AppendToBuilders(builders, str, str, str); + } + + private static void AppendToBuilders((StringBuilder leftFilter, StringBuilder targetFilter, StringBuilder rightFilter) builders, object left, object target, object right) + { + builders.leftFilter.Append(left); + builders.targetFilter.Append(target); + builders.rightFilter.Append(right); } - private sealed class OrderByQueryPage : Page + private static (string leftFilter, string targetFilter, string rightFilter) GetFormattedFilters( + ReadOnlyMemory<(OrderByColumn orderByColumn, CosmosElement orderByItem)> columnAndItems) { - public OrderByQueryPage(QueryPage queryPage) - : base(queryPage.State) + // When we run cross partition queries, + // we only serialize the continuation token for the partition that we left off on. + // The only problem is that when we resume the order by query, + // we don't have continuation tokens for all other partition. + // The saving grace is that the data has a composite sort order(query sort order, partition key range id) + // so we can generate range filters which in turn the backend will turn into rid based continuation tokens, + // which is enough to get the streams of data flowing from all partitions. + // The details of how this is done is described below: + int numOrderByItems = columnAndItems.Length; + bool isSingleOrderBy = numOrderByItems == 1; + StringBuilder left = new StringBuilder(); + StringBuilder target = new StringBuilder(); + StringBuilder right = new StringBuilder(); + + (StringBuilder, StringBuilder, StringBuilder) builders = (left, target, right); + + if (isSingleOrderBy) { - this.Enumerator = queryPage.Documents.GetEnumerator(); + //For a single order by query we resume the continuations in this manner + // Suppose the query is SELECT* FROM c ORDER BY c.string ASC + // And we left off on partition N with the value "B" + // Then + // All the partitions to the left will have finished reading "B" + // Partition N is still reading "B" + // All the partitions to the right have let to read a "B + // Therefore the filters should be + // > "B" , >= "B", and >= "B" respectively + // Repeat the same logic for DESC and you will get + // < "B", <= "B", and <= "B" respectively + // The general rule becomes + // For ASC + // > for partitions to the left + // >= for the partition we left off on + // >= for the partitions to the right + // For DESC + // < for partitions to the left + // <= for the partition we left off on + // <= for the partitions to the right + (OrderByColumn orderByColumn, CosmosElement orderByItem) = columnAndItems.Span[0]; + (string expression, SortOrder sortOrder) = (orderByColumn.Expression, orderByColumn.SortOrder); + + StringBuilder sb = new StringBuilder(); + CosmosElementToQueryLiteral cosmosElementToQueryLiteral = new CosmosElementToQueryLiteral(sb); + orderByItem.Accept(cosmosElementToQueryLiteral); + + string orderByItemToString = sb.ToString(); + + left.Append($"{expression} {(sortOrder == SortOrder.Descending ? Expressions.LessThan : Expressions.GreaterThan)} {orderByItemToString}"); + target.Append($"{expression} {(sortOrder == SortOrder.Descending ? Expressions.LessThanOrEqualTo : Expressions.GreaterThanOrEqualTo)} {orderByItemToString}"); + right.Append($"{expression} {(sortOrder == SortOrder.Descending ? Expressions.LessThanOrEqualTo : Expressions.GreaterThanOrEqualTo)} {orderByItemToString}"); + } + else + { + //For a multi order by query + // Suppose the query is SELECT* FROM c ORDER BY c.string ASC, c.number ASC + // And we left off on partition N with the value("A", 1) + // Then + // All the partitions to the left will have finished reading("A", 1) + // Partition N is still reading("A", 1) + // All the partitions to the right have let to read a "(A", 1) + // The filters are harder to derive since their are multiple columns + // But the problem reduces to "How do you know one document comes after another in a multi order by query" + // The answer is to just look at it one column at a time. + // For this particular scenario: + // If a first column is greater ex. ("B", blah), then the document comes later in the sort order + // Therefore we want all documents where the first column is greater than "A" which means > "A" + // Or if the first column is a tie, then you look at the second column ex. ("A", blah). + // Therefore we also want all documents where the first column was a tie but the second column is greater which means = "A" AND > 1 + // Therefore the filters should be + // (> "A") OR (= "A" AND > 1), (> "A") OR (= "A" AND >= 1), (> "A") OR (= "A" AND >= 1) + // Notice that if we repeated the same logic we for single order by we would have gotten + // > "A" AND > 1, >= "A" AND >= 1, >= "A" AND >= 1 + // which is wrong since we missed some documents + // Repeat the same logic for ASC, DESC + // (> "A") OR (= "A" AND < 1), (> "A") OR (= "A" AND <= 1), (> "A") OR (= "A" AND <= 1) + // Again for DESC, ASC + // (< "A") OR (= "A" AND > 1), (< "A") OR (= "A" AND >= 1), (< "A") OR (= "A" AND >= 1) + // And again for DESC DESC + // (< "A") OR (= "A" AND < 1), (< "A") OR (= "A" AND <= 1), (< "A") OR (= "A" AND <= 1) + // The general we look at all prefixes of the order by columns to look for tie breakers. + // Except for the full prefix whose last column follows the rules for single item order by + // And then you just OR all the possibilities together + for (int prefixLength = 1; prefixLength <= numOrderByItems; prefixLength++) + { + ReadOnlySpan<(OrderByColumn orderByColumn, CosmosElement orderByItem)> columnAndItemPrefix = columnAndItems.Span.Slice(start: 0, length: prefixLength); + + bool lastPrefix = prefixLength == numOrderByItems; + bool firstPrefix = prefixLength == 1; + + OrderByCrossPartitionQueryPipelineStage.AppendToBuilders(builders, "("); + + for (int index = 0; index < prefixLength; index++) + { + string expression = columnAndItemPrefix[index].orderByColumn.Expression; + SortOrder sortOrder = columnAndItemPrefix[index].orderByColumn.SortOrder; + CosmosElement orderByItem = columnAndItemPrefix[index].orderByItem; + bool lastItem = index == prefixLength - 1; + + // Append Expression + OrderByCrossPartitionQueryPipelineStage.AppendToBuilders(builders, expression); + OrderByCrossPartitionQueryPipelineStage.AppendToBuilders(builders, " "); + + // Append binary operator + if (lastItem) + { + string inequality = sortOrder == SortOrder.Descending ? Expressions.LessThan : Expressions.GreaterThan; + OrderByCrossPartitionQueryPipelineStage.AppendToBuilders(builders, inequality); + if (lastPrefix) + { + OrderByCrossPartitionQueryPipelineStage.AppendToBuilders(builders, string.Empty, Expressions.EqualTo, Expressions.EqualTo); + } + } + else + { + OrderByCrossPartitionQueryPipelineStage.AppendToBuilders(builders, Expressions.EqualTo); + } + + // Append SortOrder + StringBuilder sb = new StringBuilder(); + CosmosElementToQueryLiteral cosmosElementToQueryLiteral = new CosmosElementToQueryLiteral(sb); + orderByItem.Accept(cosmosElementToQueryLiteral); + string orderByItemToString = sb.ToString(); + OrderByCrossPartitionQueryPipelineStage.AppendToBuilders(builders, " "); + OrderByCrossPartitionQueryPipelineStage.AppendToBuilders(builders, orderByItemToString); + OrderByCrossPartitionQueryPipelineStage.AppendToBuilders(builders, " "); + + if (!lastItem) + { + OrderByCrossPartitionQueryPipelineStage.AppendToBuilders(builders, "AND "); + } + } + + OrderByCrossPartitionQueryPipelineStage.AppendToBuilders(builders, ")"); + if (!lastPrefix) + { + OrderByCrossPartitionQueryPipelineStage.AppendToBuilders(builders, " OR "); + } + } } - public QueryPage Page { get; } + // For the target filter we can make an optimization to just return "true", + // since we already have the backend continuation token to resume with. + return (left.ToString(), TrueFilter, right.ToString()); + } + + /// + /// When resuming an order by query we need to filter the document producers. + /// + /// The producer to filter down. + /// The sort orders. + /// The continuation token. + /// The cancellation token. + /// A task to await on. + private static async Task MonadicFilterAsync( + OrderByQueryPartitionRangePageAsyncEnumerator enumerator, + IReadOnlyList sortOrders, + OrderByContinuationToken continuationToken, + CancellationToken cancellationToken) + { + cancellationToken.ThrowIfCancellationRequested(); + // When we resume a query on a partition there is a possibility that we only read a partial page from the backend + // meaning that will we repeat some documents if we didn't do anything about it. + // The solution is to filter all the documents that come before in the sort order, since we have already emitted them to the client. + // The key is to seek until we get an order by value that matches the order by value we left off on. + // Once we do that we need to seek to the correct _rid within the term, + // since there might be many documents with the same order by value we left off on. + + if (!ResourceId.TryParse(continuationToken.Rid, out ResourceId continuationRid)) + { + return TryCatch.FromException( + new MalformedContinuationTokenException( + $"Invalid Rid in the continuation token {continuationToken.CompositeContinuationToken.Token} for OrderBy~Context.")); + } + + Dictionary resourceIds = new Dictionary(); + int itemToSkip = continuationToken.SkipCount; + bool continuationRidVerified = false; + + // Throw away documents until it matches the item from the continuation token. + while (await enumerator.MoveNextAsync()) + { + TryCatch monadicOrderByQueryPage = enumerator.Current; + if (monadicOrderByQueryPage.Failed) + { + return TryCatch.FromException(monadicOrderByQueryPage.Exception); + } + + OrderByQueryPage orderByQueryPage = monadicOrderByQueryPage.Result; + IEnumerator documents = orderByQueryPage.Enumerator; + while (documents.MoveNext()) + { + OrderByQueryResult orderByResult = new OrderByQueryResult(documents.Current); + + int cmp = 0; + for (int i = 0; (i < sortOrders.Count) && (cmp == 0); ++i) + { + cmp = ItemComparer.Instance.Compare( + continuationToken.OrderByItems[i].Item, + orderByResult.OrderByItems[i].Item); + + if (cmp != 0) + { + cmp = sortOrders[i] == SortOrder.Ascending ? cmp : -cmp; + } + } + + if (cmp < 0) + { + // We might have passed the item due to deletions and filters. + return TryCatch.FromResult(); + } + + if (cmp == 0) + { + if (!resourceIds.TryGetValue(orderByResult.Rid, out ResourceId rid)) + { + if (!ResourceId.TryParse(orderByResult.Rid, out rid)) + { + return TryCatch.FromException( + new MalformedContinuationTokenException( + $"Invalid Rid in the continuation token {continuationToken.CompositeContinuationToken.Token} for OrderBy~Context~TryParse.")); + } + + resourceIds.Add(orderByResult.Rid, rid); + } + + if (!continuationRidVerified) + { + if (continuationRid.Database != rid.Database || continuationRid.DocumentCollection != rid.DocumentCollection) + { + return TryCatch.FromException( + new MalformedContinuationTokenException( + $"Invalid Rid in the continuation token {continuationToken.CompositeContinuationToken.Token} for OrderBy~Context.")); + } + + continuationRidVerified = true; + } + + // Once the item matches the order by items from the continuation tokens + // We still need to remove all the documents that have a lower rid in the rid sort order. + // If there is a tie in the sort order the documents should be in _rid order in the same direction as the index (given by the backend) + cmp = continuationRid.Document.CompareTo(rid.Document); + if ((orderByQueryPage.Page.CosmosQueryExecutionInfo == null) || orderByQueryPage.Page.CosmosQueryExecutionInfo.ReverseRidEnabled) + { + // If reverse rid is enabled on the backend then fallback to the old way of doing it. + if (sortOrders[0] == SortOrder.Descending) + { + cmp = -cmp; + } + } + else + { + // Go by the whatever order the index wants + if (orderByQueryPage.Page.CosmosQueryExecutionInfo.ReverseIndexScan) + { + cmp = -cmp; + } + } + + // We might have passed the item due to deletions and filters. + // We also have a skip count for JOINs + if (cmp < 0 || (cmp == 0 && itemToSkip-- <= 0)) + { + return TryCatch.FromResult(); + } + } + } + } + + return TryCatch.FromResult(); + } + + private static bool IsSplitException(Exception exeception) + { + return exeception is CosmosException cosmosException + && (cosmosException.StatusCode == HttpStatusCode.Gone) + && (cosmosException.SubStatusCode == (int)Documents.SubStatusCodes.PartitionKeyRangeGone); + } + + public readonly struct OrderByColumn + { + public OrderByColumn(string expression, SortOrder sortOrder) + { + this.Expression = expression ?? throw new ArgumentNullException(nameof(expression)); + this.SortOrder = sortOrder; + } - public IEnumerator Enumerator { get; } + public string Expression { get; } + public SortOrder SortOrder { get; } } } } diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderByQueryPage.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderByQueryPage.cs new file mode 100644 index 0000000000..5467d32485 --- /dev/null +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderByQueryPage.cs @@ -0,0 +1,23 @@ +// ------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +// ------------------------------------------------------------ + +namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote +{ + using System.Collections.Generic; + using Microsoft.Azure.Cosmos.CosmosElements; + using Microsoft.Azure.Cosmos.Pagination; + + internal sealed class OrderByQueryPage : Page + { + public OrderByQueryPage(QueryPage queryPage) + : base(queryPage.State) + { + this.Enumerator = queryPage.Documents.GetEnumerator(); + } + + public QueryPage Page { get; } + + public IEnumerator Enumerator { get; } + } +} diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderByQueryPartitionRangePageAsyncEnumerator.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderByQueryPartitionRangePageAsyncEnumerator.cs new file mode 100644 index 0000000000..1d9b8f269b --- /dev/null +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderByQueryPartitionRangePageAsyncEnumerator.cs @@ -0,0 +1,55 @@ +// ------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +// ------------------------------------------------------------ + +namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote +{ + using System; + using System.Threading; + using System.Threading.Tasks; + using Microsoft.Azure.Cosmos.CosmosElements; + using Microsoft.Azure.Cosmos.Pagination; + using Microsoft.Azure.Cosmos.Query.Core.Monads; + using Microsoft.Azure.Documents; + + internal sealed class OrderByQueryPartitionRangePageAsyncEnumerator : PartitionRangePageAsyncEnumerator + { + private readonly IQueryDataSource queryDataSource; + private readonly SqlQuerySpec sqlQuerySpec; + private readonly int pageSize; + + public OrderByQueryPartitionRangePageAsyncEnumerator( + IQueryDataSource queryDataSource, + SqlQuerySpec sqlQuerySpec, + PartitionKeyRange feedRange, + int pageSize, + QueryState state = default) + : base(feedRange, state) + { + this.queryDataSource = queryDataSource ?? throw new ArgumentNullException(nameof(queryDataSource)); + this.sqlQuerySpec = sqlQuerySpec ?? throw new ArgumentNullException(nameof(sqlQuerySpec)); + this.pageSize = pageSize; + } + + public override ValueTask DisposeAsync() => default; + + protected override Task> GetNextPageAsync(CancellationToken cancellationToken) => this.queryDataSource + .MonadicQueryAsync( + sqlQuerySpec: this.sqlQuerySpec, + continuationToken: this.State == null ? null : ((CosmosString)this.State.Value).Value, + feedRange: new FeedRangeEpk(this.Range.ToRange()), + pageSize: this.pageSize, + cancellationToken) + .ContinueWith>(antecedent => + { + TryCatch monadicQueryPage = antecedent.Result; + if (monadicQueryPage.Failed) + { + return TryCatch.FromException(monadicQueryPage.Exception); + } + + QueryPage queryPage = monadicQueryPage.Result; + return TryCatch.FromResult(new OrderByQueryPage(queryPage)); + }); + } +} diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/ParallelCrossPartitionQueryPipelineStage.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/ParallelCrossPartitionQueryPipelineStage.cs index ea7e1bcbe5..04d498d827 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/ParallelCrossPartitionQueryPipelineStage.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/ParallelCrossPartitionQueryPipelineStage.cs @@ -15,6 +15,13 @@ namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote using Microsoft.Azure.Cosmos.Query.Core.Monads; using Microsoft.Azure.Documents; + /// + /// is an implementation of that drain results from multiple remote nodes. + /// This class is responsible for draining cross partition queries that do not have order by conditions. + /// The way parallel queries work is that it drains from the left most partition first. + /// This class handles draining in the correct order and can also stop and resume the query + /// by generating a continuation token and resuming from said continuation token. + /// internal sealed class ParallelCrossPartitionQueryPipelineStage : IQueryPipelineStage { private readonly CrossPartitionRangePageAsyncEnumerator crossPartitionRangePageAsyncEnumerator; From da5d2d066f5afa273664705d28e42ff3dcd94f4e Mon Sep 17 00:00:00 2001 From: Brandon Chong Date: Thu, 30 Jul 2020 14:21:34 -0700 Subject: [PATCH 47/85] Revert "Internal Query : Fix Remove Antlr dependency for now (#1626)" This reverts commit ed7faaef5a9940f6b4a3a332020c74f0ea0aabc3. --- .../src/Microsoft.Azure.Cosmos.csproj | 1 + .../src/Query/Core/Parser/CstToAstVisitor.cs | 34 +++++++++---------- .../src/Query/Core/Parser/ErrorListener.cs | 2 -- .../src/Query/Core/Parser/IsqlListener.cs | 2 -- .../src/Query/Core/Parser/IsqlVisitor.cs | 2 -- .../src/Query/Core/Parser/LASets.cs | 3 -- .../src/Query/Core/Parser/sqlBaseVisitor.cs | 2 -- .../src/Query/Core/Parser/sqlLexer.cs | 2 -- .../src/Query/Core/Parser/sqlParser.cs | 4 +-- .../src/SqlObjects/SqlQuery.cs | 9 ----- .../Query/SanityQueryTests.cs | 2 -- 11 files changed, 18 insertions(+), 45 deletions(-) diff --git a/Microsoft.Azure.Cosmos/src/Microsoft.Azure.Cosmos.csproj b/Microsoft.Azure.Cosmos/src/Microsoft.Azure.Cosmos.csproj index 52e330d396..068bda3a63 100644 --- a/Microsoft.Azure.Cosmos/src/Microsoft.Azure.Cosmos.csproj +++ b/Microsoft.Azure.Cosmos/src/Microsoft.Azure.Cosmos.csproj @@ -82,6 +82,7 @@ + diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Parser/CstToAstVisitor.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Parser/CstToAstVisitor.cs index aa794658f3..2cc04391f8 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Parser/CstToAstVisitor.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Parser/CstToAstVisitor.cs @@ -4,7 +4,6 @@ namespace Microsoft.Azure.Cosmos.Query.Core.Parser { -#if false using System; using System.Collections.Generic; using Antlr4.Runtime.Tree; @@ -127,7 +126,7 @@ public override SqlObject VisitSql_query(sqlParser.Sql_queryContext context) sqlOffsetLimitClause); } - #region SELECT + #region SELECT public override SqlObject VisitSelect_clause(sqlParser.Select_clauseContext context) { SqlSelectSpec sqlSelectSpec = (SqlSelectSpec)this.Visit(context.selection()); @@ -201,8 +200,8 @@ public override SqlObject VisitTop_spec(sqlParser.Top_specContext context) Number64 topCount = CstToAstVisitor.GetNumber64ValueFromNode(context.NUMERIC_LITERAL()); return SqlTopSpec.Create(SqlNumberLiteral.Create(topCount)); } - #endregion - #region FROM + #endregion + #region FROM public override SqlObject VisitFrom_clause(sqlParser.From_clauseContext context) { Contract.Requires(context != null); @@ -311,16 +310,16 @@ public override SqlObject VisitStringPathExpression(sqlParser.StringPathExpressi return SqlStringPathExpression.Create(pathExpression, stringIndex); } - #endregion - #region WHERE + #endregion + #region WHERE public override SqlObject VisitWhere_clause(sqlParser.Where_clauseContext context) { Contract.Requires(context != null); SqlScalarExpression sqlScalarExpression = (SqlScalarExpression)this.Visit(context.scalar_expression()); return SqlWhereClause.Create(sqlScalarExpression); } - #endregion - #region GROUP BY + #endregion + #region GROUP BY public override SqlObject VisitGroup_by_clause(sqlParser.Group_by_clauseContext context) { Contract.Requires(context != null); @@ -333,8 +332,8 @@ public override SqlObject VisitGroup_by_clause(sqlParser.Group_by_clauseContext return SqlGroupByClause.Create(groupByColumns); } - #endregion - #region ORDER BY + #endregion + #region ORDER BY public override SqlObject VisitOrder_by_clause(sqlParser.Order_by_clauseContext context) { Contract.Requires(context != null); @@ -366,8 +365,8 @@ public override SqlObject VisitOrder_by_clause(sqlParser.Order_by_clauseContext return SqlOrderbyClause.Create(orderByItems); } - #endregion - #region OFFSET LIMIT + #endregion + #region OFFSET LIMIT public override SqlObject VisitOffset_limit_clause(sqlParser.Offset_limit_clauseContext context) { Contract.Requires(context != null); @@ -384,8 +383,8 @@ public override SqlObject VisitOffset_limit_clause(sqlParser.Offset_limit_clause return SqlOffsetLimitClause.Create(sqlOffsetSpec, sqlLimitSpec); } - #endregion - #region ScalarExpressions + #endregion + #region ScalarExpressions public override SqlObject VisitArrayCreateScalarExpression(sqlParser.ArrayCreateScalarExpressionContext context) { Contract.Requires(context != null); @@ -638,9 +637,9 @@ public override SqlObject VisitUnaryScalarExpression(sqlParser.UnaryScalarExpres return SqlUnaryScalarExpression.Create(unaryOperator, expression); } - #endregion + #endregion - #region NOT IMPLEMENTED ON PURPOSE + #region NOT IMPLEMENTED ON PURPOSE public override SqlObject VisitBinary_operator(sqlParser.Binary_operatorContext context) { throw new NotSupportedException(); @@ -690,7 +689,7 @@ public override SqlObject VisitScalar_expression_list(sqlParser.Scalar_expressio { throw new NotSupportedException(); } - #endregion + #endregion private sealed class UnknownSqlObjectException : ArgumentOutOfRangeException { @@ -725,5 +724,4 @@ private static Number64 GetNumber64ValueFromNode(IParseTree parseTree) return number64; } } -#endif } diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Parser/ErrorListener.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Parser/ErrorListener.cs index 2f04662eec..ee201c73e5 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Parser/ErrorListener.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Parser/ErrorListener.cs @@ -6,7 +6,6 @@ namespace Microsoft.Azure.Cosmos.Query.Core.Parser { - #if false using System; using System.Collections.Generic; using System.IO; @@ -67,5 +66,4 @@ public override void SyntaxError(TextWriter output, IRecognizer recognizer, S of base.SyntaxError(output, recognizer, offendingSymbol, line, col, msg, e); } } -#endif } \ No newline at end of file diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Parser/IsqlListener.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Parser/IsqlListener.cs index 0b7d1cd50f..b77c525ef3 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Parser/IsqlListener.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Parser/IsqlListener.cs @@ -21,7 +21,6 @@ namespace Microsoft.Azure.Cosmos.Query.Core.Parser { - #if false using Antlr4.Runtime.Misc; using IParseTreeListener = Antlr4.Runtime.Tree.IParseTreeListener; @@ -583,5 +582,4 @@ internal interface IsqlListener : IParseTreeListener /// The parse tree. void ExitLiteral([NotNull] sqlParser.LiteralContext context); } -#endif } diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Parser/IsqlVisitor.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Parser/IsqlVisitor.cs index 830df59460..802ad222ca 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Parser/IsqlVisitor.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Parser/IsqlVisitor.cs @@ -21,7 +21,6 @@ namespace Microsoft.Azure.Cosmos.Query.Core.Parser { - #if false using Antlr4.Runtime.Misc; using Antlr4.Runtime.Tree; using IToken = Antlr4.Runtime.IToken; @@ -360,5 +359,4 @@ internal interface IsqlVisitor : IParseTreeVisitor /// The visitor result. Result VisitLiteral([NotNull] sqlParser.LiteralContext context); } -#endif } \ No newline at end of file diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Parser/LASets.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Parser/LASets.cs index eb5e79db54..3b3dccb941 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Parser/LASets.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Parser/LASets.cs @@ -6,7 +6,6 @@ namespace Microsoft.Azure.Cosmos.Query.Core.Parser { -#if false using System; using System.Collections.Generic; using System.Diagnostics; @@ -18,7 +17,6 @@ namespace Microsoft.Azure.Cosmos.Query.Core.Parser internal class LASets { - private readonly Dictionary, bool> visited = new Dictionary, bool>(); private readonly bool logParse = false; private readonly bool logClosure = false; @@ -797,5 +795,4 @@ private string PrintResult(List> all_parses) return sb.ToString(); } } -#endif } diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Parser/sqlBaseVisitor.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Parser/sqlBaseVisitor.cs index f2189c00cd..11a2fb31b1 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Parser/sqlBaseVisitor.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Parser/sqlBaseVisitor.cs @@ -21,7 +21,6 @@ namespace Microsoft.Azure.Cosmos.Query.Core.Parser { - #if false using Antlr4.Runtime.Misc; using Antlr4.Runtime.Tree; @@ -563,5 +562,4 @@ internal partial class sqlBaseVisitor : AbstractParseTreeVisitor /// The visitor result. public virtual Result VisitLiteral([NotNull] sqlParser.LiteralContext context) { return VisitChildren(context); } } -#endif } \ No newline at end of file diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Parser/sqlLexer.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Parser/sqlLexer.cs index cd6dd56397..6d10bc8dc0 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Parser/sqlLexer.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Parser/sqlLexer.cs @@ -21,7 +21,6 @@ namespace Microsoft.Azure.Cosmos.Query.Core.Parser { - #if false using System; using System.IO; using Antlr4.Runtime; @@ -563,5 +562,4 @@ static sqlLexer() public static readonly ATN _ATN = new ATNDeserializer().Deserialize(_serializedATN); } -#endif } diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Parser/sqlParser.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Parser/sqlParser.cs index cc93367f5e..0811b9af4e 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Parser/sqlParser.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Parser/sqlParser.cs @@ -21,7 +21,6 @@ namespace Microsoft.Azure.Cosmos.Query.Core.Parser { - #if false using System; using System.IO; using Antlr4.Runtime; @@ -3326,5 +3325,4 @@ private bool scalar_expression_sempred(Scalar_expressionContext _localctx, int p public static readonly ATN _ATN = new ATNDeserializer().Deserialize(_serializedATN); } -#endif -} \ No newline at end of file +} diff --git a/Microsoft.Azure.Cosmos/src/SqlObjects/SqlQuery.cs b/Microsoft.Azure.Cosmos/src/SqlObjects/SqlQuery.cs index be01e2802f..601949a24e 100644 --- a/Microsoft.Azure.Cosmos/src/SqlObjects/SqlQuery.cs +++ b/Microsoft.Azure.Cosmos/src/SqlObjects/SqlQuery.cs @@ -5,10 +5,8 @@ namespace Microsoft.Azure.Cosmos.SqlObjects { using System; using System.Runtime.ExceptionServices; -#if false using Antlr4.Runtime; using Antlr4.Runtime.Misc; -#endif using Microsoft.Azure.Cosmos.Query.Core.Parser; using Microsoft.Azure.Cosmos.SqlObjects.Visitors; @@ -76,7 +74,6 @@ public static bool TryParse(string text, out SqlQuery sqlQuery) throw new ArgumentNullException(nameof(text)); } -#if false AntlrInputStream str = new AntlrInputStream(text); sqlLexer lexer = new sqlLexer(str); CommonTokenStream tokens = new CommonTokenStream(lexer); @@ -106,13 +103,8 @@ public static bool TryParse(string text, out SqlQuery sqlQuery) sqlQuery = (SqlQuery)CstToAstVisitor.Singleton.Visit(programContext); return true; -#else - sqlQuery = default; - return false; -#endif } -#if false private sealed class ThrowExceptionOnErrors : IAntlrErrorStrategy { public static readonly ThrowExceptionOnErrors Singleton = new ThrowExceptionOnErrors(); @@ -153,6 +145,5 @@ public void Sync(Parser recognizer) // Do nothing } } -#endif } } diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/Query/SanityQueryTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/Query/SanityQueryTests.cs index 228db1a6d6..5da1ac134b 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/Query/SanityQueryTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/Query/SanityQueryTests.cs @@ -675,12 +675,10 @@ async Task> AssertPassthroughAsync(string query, Cosmos.Part Assert.IsTrue(feedOptions.TestSettings.Stats.PipelineType.HasValue); Assert.AreEqual(TestInjections.PipelineType.Passthrough, feedOptions.TestSettings.Stats.PipelineType.Value); -#if false if (pk.HasValue) { Assert.AreEqual(0, cosmosQueryClientCore.QueryPlanCalls); } -#endif return queryResults; } From 48a9c7c9ea3f77d914ac9986d2ca13c65b82a50f Mon Sep 17 00:00:00 2001 From: Brandon Chong Date: Thu, 30 Jul 2020 15:29:26 -0700 Subject: [PATCH 48/85] added test for antlr dependancy --- .../Contracts/DirectContractTests.cs | 55 +++++++++++-------- 1 file changed, 32 insertions(+), 23 deletions(-) diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Contracts/DirectContractTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Contracts/DirectContractTests.cs index 1b53d2c861..cc214509e8 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Contracts/DirectContractTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Contracts/DirectContractTests.cs @@ -27,7 +27,7 @@ public void TestInteropTest() CosmosClient client = new CosmosClient(connectionString: null); Assert.Fail(); } - catch(ArgumentNullException) + catch (ArgumentNullException) { } @@ -62,10 +62,10 @@ public void MappedRegionsTest() .Select(e => e.Name) .ToArray(); - if(locationNames.Length > cosmosRegions.Length) + if (locationNames.Length > cosmosRegions.Length) { HashSet missingLocationNames = new HashSet(locationNames); - foreach(string region in cosmosRegions) + foreach (string region in cosmosRegions) { missingLocationNames.Remove(region); } @@ -94,31 +94,40 @@ public void CustomJsonReaderTest() public void ProjectPackageDependenciesTest() { string csprojFile = "Microsoft.Azure.Cosmos.csproj"; - Dictionary projDependencies = DirectContractTests.GetPackageReferencies(csprojFile); - Dictionary baselineDependencies = JsonConvert.DeserializeObject< Dictionary>( - "{\"System.Numerics.Vectors\":\"4.5.0\",\"Newtonsoft.Json\":\"10.0.2\",\"System.Configuration.ConfigurationManager\":\"4.5.0\",\"System.Memory\":\"4.5.1\",\"System.Buffers\":\"4.5.1\",\"System.Runtime.CompilerServices.Unsafe\":\"4.5.1\",\"System.Threading.Tasks.Extensions\":\"4.5.1\",\"System.ValueTuple\":\"4.5.0\"}"); - - Assert.AreEqual(projDependencies.Count, baselineDependencies.Count); - foreach(KeyValuePair projectDependency in projDependencies) + Dictionary projectDependencies = DirectContractTests.GetPackageReferencies(csprojFile); + Dictionary baselineDependencies = new Dictionary() + { + { "System.Numerics.Vectors", new Version(4, 5, 0) }, + { "Newtonsoft.Json", new Version(10, 0, 2) }, + { "System.Configuration.ConfigurationManager", new Version(4, 5, 0) }, + { "System.Memory", new Version(4, 5, 1) }, + { "System.Buffers", new Version(4, 5, 1) }, + { "System.Runtime.CompilerServices.Unsafe", new Version(4, 5, 1) }, + { "System.Threading.Tasks.Extensions", new Version(4, 5, 1) }, + { "System.ValueTuple", new Version(4, 5, 0) }, + { "Antlr4.Runtime.Standard", new Version(4, 8, 0) }, + }; + + Assert.AreEqual(projectDependencies.Count, baselineDependencies.Count); + foreach (KeyValuePair projectDependency in projectDependencies) { - string baselineVersion = baselineDependencies[projectDependency.Key]; + Version baselineVersion = baselineDependencies[projectDependency.Key]; Assert.AreEqual(baselineVersion, projectDependency.Value); } - } [TestMethod] public void PackageDependenciesTest() { string csprojFile = "Microsoft.Azure.Cosmos.csproj"; - Dictionary projDependencies = DirectContractTests.GetPackageReferencies(csprojFile); + Dictionary projDependencies = DirectContractTests.GetPackageReferencies(csprojFile); string[] files = Directory.GetFiles(Directory.GetCurrentDirectory(), "*.nuspec"); - Dictionary allDependencies = new Dictionary(); + Dictionary allDependencies = new Dictionary(); foreach (string nuspecFile in files) { - Dictionary nuspecDependencies = DirectContractTests.GetNuspecDependencies(nuspecFile); - foreach(KeyValuePair e in nuspecDependencies) + Dictionary nuspecDependencies = DirectContractTests.GetNuspecDependencies(nuspecFile); + foreach (KeyValuePair e in nuspecDependencies) { if (!allDependencies.ContainsKey(e.Key)) { @@ -126,7 +135,7 @@ public void PackageDependenciesTest() } else { - string existingValue = allDependencies[e.Key]; + Version existingValue = allDependencies[e.Key]; if (existingValue.CompareTo(e.Value) > 0) { allDependencies[e.Key] = e.Value; @@ -136,7 +145,7 @@ public void PackageDependenciesTest() } // Dependency version should match - foreach(KeyValuePair e in allDependencies) + foreach (KeyValuePair e in allDependencies) { Assert.AreEqual(e.Value, projDependencies[e.Key]); } @@ -144,7 +153,7 @@ public void PackageDependenciesTest() CollectionAssert.IsSubsetOf(allDependencies.Keys, projDependencies.Keys); } - private static Dictionary GetPackageReferencies(string csprojName) + private static Dictionary GetPackageReferencies(string csprojName) { string fullCsprojName = Path.Combine(Directory.GetCurrentDirectory(), csprojName); Trace.TraceInformation($"Testing dependencies for csporj file {fullCsprojName}"); @@ -156,7 +165,7 @@ private static Dictionary GetPackageReferencies(string csprojNam int prjRefCount = new Regex(" projReferences = new Dictionary(); + Dictionary projReferences = new Dictionary(); foreach (Match m in matches) { if (m.Groups["PrivateAssets"].Captures.Count != 0) @@ -165,14 +174,14 @@ private static Dictionary GetPackageReferencies(string csprojNam } else { - projReferences[m.Groups["Include"].Value] = m.Groups["Version"].Value; + projReferences[m.Groups["Include"].Value] = Version.Parse(m.Groups["Version"].Value); } } return projReferences; } - private static Dictionary GetNuspecDependencies(string nuspecFile) + private static Dictionary GetNuspecDependencies(string nuspecFile) { Trace.TraceInformation($"Testing dependencies for nuspec file {nuspecFile}"); string nuspecContent = File.ReadAllText(nuspecFile); @@ -183,10 +192,10 @@ private static Dictionary GetNuspecDependencies(string nuspecFil int dependencyCount = new Regex(" dependencies = new Dictionary(); + Dictionary dependencies = new Dictionary(); foreach (Match m in matches) { - dependencies[m.Groups["id"].Value] = m.Groups["version"].Value; + dependencies[m.Groups["id"].Value] = Version.Parse(m.Groups["version"].Value); } return dependencies; From 5006be310e23e7668bbec0efbce3550e09a2a078 Mon Sep 17 00:00:00 2001 From: Brandon Chong Date: Fri, 31 Jul 2020 16:52:52 -0700 Subject: [PATCH 49/85] need to wire up the item enumerator to order by --- .../CrossPartitionRangeItemAsyncEnumerator.cs | 199 ------------------ .../src/Pagination/ItemEnumeratorPage.cs | 22 -- ...OrderByCrossPartitionQueryPipelineStage.cs | 80 ++----- .../Pipeline/Remote/OrderByItemEnumerator.cs | 191 +++++++++++++++++ ...yQueryPartitionRangePageAsyncEnumerator.cs | 4 + ...eryPipelineStage.cs => PartitionMapper.cs} | 9 +- 6 files changed, 218 insertions(+), 287 deletions(-) delete mode 100644 Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionRangeItemAsyncEnumerator.cs delete mode 100644 Microsoft.Azure.Cosmos/src/Pagination/ItemEnumeratorPage.cs create mode 100644 Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderByItemEnumerator.cs rename Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/{CrossPartitionQueryPipelineStage.cs => PartitionMapper.cs} (95%) diff --git a/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionRangeItemAsyncEnumerator.cs b/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionRangeItemAsyncEnumerator.cs deleted file mode 100644 index 52eca81425..0000000000 --- a/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionRangeItemAsyncEnumerator.cs +++ /dev/null @@ -1,199 +0,0 @@ -// ------------------------------------------------------------ -// Copyright (c) Microsoft Corporation. All rights reserved. -// ------------------------------------------------------------ - -namespace Microsoft.Azure.Cosmos.Pagination -{ - using System; - using System.Collections.Generic; - using System.Net; - using System.Threading.Tasks; - using Microsoft.Azure.Cosmos.CosmosElements; - using Microsoft.Azure.Cosmos.Query.Core.Collections; - using Microsoft.Azure.Cosmos.Query.Core.Monads; - using Microsoft.Azure.Documents; - - /// - /// Coordinates draining pages from multiple , while maintaining a global sort order and handling repartitioning (splits, merge). - /// Also does things on a per item basis. - /// - internal sealed class CrossPartitionRangeItemAsyncEnumerator : IAsyncEnumerator> - where TState : State - { - private readonly IFeedRangeProvider feedRangeProvider; - private readonly CreatePartitionRangePageAsyncEnumerator, TState> createPartitionRangeEnumerator; - private readonly Func, TState>, Task>>> intializeAsync; - private readonly PriorityQueue, TState>> enumerators; - private readonly Queue, TState>> uninitializedEnumerators; - - public CrossPartitionRangeItemAsyncEnumerator( - IFeedRangeProvider feedRangeProvider, - CreatePartitionRangePageAsyncEnumerator, TState> createPartitionRangeEnumerator, - IEnumerable, TState>> uninitializedEnumerators, - Func, TState>, Task>>> initializeAsync, - IComparer itemComparer) - { - this.feedRangeProvider = feedRangeProvider ?? throw new ArgumentNullException(nameof(feedRangeProvider)); - this.createPartitionRangeEnumerator = createPartitionRangeEnumerator ?? throw new ArgumentNullException(nameof(createPartitionRangeEnumerator)); - this.intializeAsync = initializeAsync ?? throw new ArgumentNullException(nameof(initializeAsync)); - this.uninitializedEnumerators = new Queue, TState>>(uninitializedEnumerators); - this.enumerators = new PriorityQueue, TState>>(new EnumeratorComparer()); - } - - public TryCatch, TState>> Current { get; private set; } - - public async ValueTask MoveNextAsync() - { - if (this.uninitializedEnumerators.Count != 0) - { - PartitionRangePageAsyncEnumerator, TState> uninitializedEnumerator = this.uninitializedEnumerators.Dequeue(); - TryCatch> initializeMonad = await this.intializeAsync(uninitializedEnumerator); - if (initializeMonad.Failed) - { - if (!await this.TryHandleExceptionAsync(uninitializedEnumerator)) - { - this.uninitializedEnumerators.Enqueue(uninitializedEnumerator); - } - - this.Current = TryCatch, TState>>.FromException(uninitializedEnumerator.Current.Exception); - return true; - } - - // Once the enumerator has been initialized we can add it back to the priority queue - this.enumerators.Enqueue(uninitializedEnumerator); - - // We want to report back the metrics from initialization, so that the user has accurate metrics, - // But we need to make up a fake continuation token, since we aren't in a valid state to continue from. - this.Current = TryCatch, TState>>.FromResult( - new CrossPartitionPage, TState>( - initializeMonad.Result, - new CrossPartitionState(new List<(PartitionKeyRange, TState)>()))); - - // and recursively retry - return await this.MoveNextAsync(); - } - - if (this.enumerators.Count == 0) - { - return false; - } - - PartitionRangePageAsyncEnumerator, TState> currentEnumerator = this.enumerators.Dequeue(); - if (currentEnumerator.Current.Result.Items.Count == 0) - { - - } - - if (!await currentEnumerator.MoveNextAsync()) - { - // Current enumerator is empty, - // so recursively retry on the next enumerator. - return await this.MoveNextAsync(); - } - - if (currentEnumerator.Current.Failed) - { - - } - - if (currentEnumerator.State != null) - { - currentEnumerator.Enqueue(currentPaginator); - } - - TryCatch> backendPage = currentEnumerator.Current; - if (backendPage.Failed) - { - this.Current = TryCatch, TState>>.FromException(backendPage.Exception); - return true; - } - - CrossPartitionState crossPartitionState; - if (enumerators.Count == 0) - { - crossPartitionState = null; - } - else - { - List<(PartitionKeyRange, TState)> feedRangeAndStates = new List<(PartitionKeyRange, TState)>(enumerators.Count); - foreach (PartitionRangePageAsyncEnumerator, TState> enumerator in enumerators) - { - feedRangeAndStates.Add((enumerator.Range, enumerator.State)); - } - - crossPartitionState = new CrossPartitionState(feedRangeAndStates); - } - - this.Current = TryCatch, TState>>.FromResult( - new CrossPartitionPage, TState>(backendPage.Result, crossPartitionState)); - return true; - } - - private async Task TryHandleExceptionAsync(PartitionRangePageAsyncEnumerator, TState> enumerator) - { - if (!enumerator.Current.Failed) - { - throw new InvalidOperationException("Enumerator was not in a faulted state."); - } - - Exception exception = enumerator.Current.Exception; - - // Check if it's a retryable exception. - while (exception.InnerException != null) - { - exception = exception.InnerException; - } - - if (IsSplitException(exception)) - { - // Handle split - IEnumerable childRanges = await this.feedRangeProvider.GetChildRangeAsync( - enumerator.Range, - cancellationToken: default); - foreach (PartitionKeyRange childRange in childRanges) - { - PartitionRangePageAsyncEnumerator, TState> childPaginator = this.createPartitionRangeEnumerator( - childRange, - enumerator.State); - this.uninitializedEnumerators.Enqueue(childPaginator); - } - - return true; - } - - if (IsMergeException(exception)) - { - throw new NotImplementedException(); - } - - return false; - } - - public ValueTask DisposeAsync() - { - // Do Nothing. - return default; - } - - private static bool IsSplitException(Exception exeception) - { - return exeception is CosmosException cosmosException - && (cosmosException.StatusCode == HttpStatusCode.Gone) - && (cosmosException.SubStatusCode == (int)Documents.SubStatusCodes.PartitionKeyRangeGone); - } - - private static bool IsMergeException(Exception exception) - { - // TODO: code this out - return false; - } - - private sealed class EnumeratorComparer : IComparer, TState>> - { - public int Compare(PartitionRangePageAsyncEnumerator, TState> x, PartitionRangePageAsyncEnumerator, TState> y) - { - throw new NotImplementedException(); - } - } - } -} diff --git a/Microsoft.Azure.Cosmos/src/Pagination/ItemEnumeratorPage.cs b/Microsoft.Azure.Cosmos/src/Pagination/ItemEnumeratorPage.cs deleted file mode 100644 index fe98eb0a8e..0000000000 --- a/Microsoft.Azure.Cosmos/src/Pagination/ItemEnumeratorPage.cs +++ /dev/null @@ -1,22 +0,0 @@ -// ------------------------------------------------------------ -// Copyright (c) Microsoft Corporation. All rights reserved. -// ------------------------------------------------------------ - -namespace Microsoft.Azure.Cosmos.Pagination -{ - using System; - using System.Collections.Generic; - using System.Text; - using Microsoft.Azure.Cosmos.CosmosElements; - - internal sealed class ItemEnumeratorPage : Page - where TState : State - { - public ItemEnumeratorPage(IReadOnlyList cosmosElements) - { - this.Items = new Queue(cosmosElements); - } - - public Queue Items { get; } - } -} diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderByCrossPartitionQueryPipelineStage.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderByCrossPartitionQueryPipelineStage.cs index a110e6aa2b..a57183b43a 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderByCrossPartitionQueryPipelineStage.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderByCrossPartitionQueryPipelineStage.cs @@ -7,6 +7,7 @@ namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote using System; using System.Collections.Generic; using System.Linq; + using System.Net; using System.Text; using System.Threading; using System.Threading.Tasks; @@ -19,6 +20,7 @@ namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote using Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.OrderBy; using Microsoft.Azure.Cosmos.Query.Core.Monads; using Microsoft.Azure.Documents; + using static Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote.PartitionMapper; /// /// CosmosOrderByItemQueryExecutionContext is a concrete implementation for CrossPartitionQueryExecutionContext. @@ -51,78 +53,28 @@ private static class Expressions /// private const string TrueFilter = "true"; - /// - /// For ORDER BY queries we need to drain the first page of every enumerator before adding it to the prioirty queue for k-way merge. - /// Since this requires async work, we defer it until the user calls MoveNextAsync(). - /// - private readonly Queue<(OrderByQueryPartitionRangePageAsyncEnumerator, OrderByContinuationToken)> uninitializedEnumerators; - - private readonly PriorityQueue initializedEnumerators; - - private readonly CrossPartitionRangePageAsyncEnumerator crossPartitionRangePageAsyncEnumerator; - - private readonly List sortOrders; - - private readonly IDocumentContainer documentContainer; + private readonly OrderByItemEnumerator orderByItemEnumerator; + private readonly int pageSize; - private OrderByCrossPartitionQueryPipelineStage( - CrossPartitionRangePageAsyncEnumerator crossPartitionRangePageAsyncEnumerator) + private OrderByCrossPartitionQueryPipelineStage(OrderByItemEnumerator orderByItemEnumerator) { - this.crossPartitionRangePageAsyncEnumerator = crossPartitionRangePageAsyncEnumerator ?? throw new ArgumentNullException(nameof(crossPartitionRangePageAsyncEnumerator)); + this.orderByItemEnumerator = orderByItemEnumerator ?? throw new ArgumentNullException(nameof(orderByItemEnumerator)); } - public TryCatch Current => throw new NotImplementedException(); + public TryCatch Current { get; private set; } - public ValueTask DisposeAsync() => this.crossPartitionRangePageAsyncEnumerator.DisposeAsync(); + public ValueTask DisposeAsync() => this.orderByItemEnumerator.DisposeAsync(); public async ValueTask MoveNextAsync() { - if (this.uninitializedEnumerators.Count != 0) + if (this.orderByItemEnumerator.Current.Failed) { - (OrderByQueryPartitionRangePageAsyncEnumerator enumerator, OrderByContinuationToken token) = this.uninitializedEnumerators.Dequeue(); - if (token != null) - { - TryCatch monadicFilterAsync = await OrderByCrossPartitionQueryPipelineStage.MonadicFilterAsync( - enumerator, - this.sortOrders, - token, - cancellationToken: default); - if (monadicFilterAsync.Failed) - { - // Check if it's a retryable exception. - Exception exception = monadicFilterAsync.Exception; - while (exception.InnerException != null) - { - exception = exception.InnerException; - } - - if (IsSplitException(exception)) - { - // Handle split - IEnumerable childRanges = await this.documentContainer.GetChildRangeAsync( - enumerator.Range, - cancellationToken: default); - foreach (PartitionKeyRange childRange in childRanges) - { - OrderByQueryPartitionRangePageAsyncEnumerator childPaginator = new OrderByQueryPartitionRangePageAsyncEnumerator( - documentContainer, - sqlQuerySpec, - range, - pageSize, - state: default) - enumerators.Enqueue(childPaginator); - } - - // Recursively retry - return await this.MoveNextAsync(); - } - } - } - else - { - - } + // If we have a buffered exception, + // then return that first. + this.Current = this.orderByItemEnumerator.Current; + return true; } + } public static TryCatch MonadicCreate( @@ -254,6 +206,10 @@ public static TryCatch MonadicCreate( } } } + + OrderByItemEnumerator orderByItemEnumerator = new OrderByItemEnumerator( + documentContainer, + ) } private static TryCatch> MonadicGetOrderByContinuationTokenMapping( diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderByItemEnumerator.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderByItemEnumerator.cs new file mode 100644 index 0000000000..fb0f65e66d --- /dev/null +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderByItemEnumerator.cs @@ -0,0 +1,191 @@ +// ------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +// ------------------------------------------------------------ + +namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote +{ + using System; + using System.Collections.Generic; + using System.Net; + using System.Threading.Tasks; + using Microsoft.Azure.Cosmos.CosmosElements; + using Microsoft.Azure.Cosmos.Pagination; + using Microsoft.Azure.Cosmos.Query.Core.Collections; + using Microsoft.Azure.Cosmos.Query.Core.Monads; + using Microsoft.Azure.Documents; + + internal sealed class OrderByItemEnumerator : IAsyncEnumerator> + { + private static readonly QueryState UninitializedQueryState = new QueryState(CosmosString.Create("ORDER BY NOT INITIALIZED YET!")); + private static readonly IReadOnlyList EmptyPage = new List(); + private readonly IDocumentContainer documentContainer; + private readonly Func createEnumerator; + private readonly Func>> intializeAsync; + private readonly PriorityQueue enumerators; + private readonly Queue uninitializedEnumerators; + + public OrderByItemEnumerator( + IDocumentContainer documentContainer, + Func createEnumerator, + IEnumerable uninitializedEnumerators, + Func>> initializeAsync, + IComparer itemComparer) + { + this.documentContainer = documentContainer ?? throw new ArgumentNullException(nameof(documentContainer)); + this.createEnumerator = createEnumerator ?? throw new ArgumentNullException(nameof(createEnumerator)); + this.intializeAsync = initializeAsync ?? throw new ArgumentNullException(nameof(initializeAsync)); + this.uninitializedEnumerators = new Queue(uninitializedEnumerators); + this.enumerators = new PriorityQueue(itemComparer); + } + + public TryCatch Current { get; private set; } + + public async ValueTask MoveNextAsync() + { + if (this.uninitializedEnumerators.Count != 0) + { + OrderByQueryPartitionRangePageAsyncEnumerator uninitializedEnumerator = this.uninitializedEnumerators.Dequeue(); + TryCatch initializeMonad = await this.intializeAsync(uninitializedEnumerator); + if (initializeMonad.Failed) + { + if (!await this.TryHandleExceptionAsync(uninitializedEnumerator)) + { + this.uninitializedEnumerators.Enqueue(uninitializedEnumerator); + } + + this.Current = TryCatch.FromException(uninitializedEnumerator.Current.Exception); + return true; + } + + // Once the enumerator has been initialized we can add it back to the priority queue + this.enumerators.Enqueue(uninitializedEnumerator); + + // We want to report back the metrics from initialization, so that the user has accurate metrics, + // But we need to make up a fake continuation token, since we aren't in a valid state to continue from. + QueryPage initialiazationPage = initializeMonad.Result; + this.Current = TryCatch.FromResult( + new QueryPage( + documents: EmptyPage, + requestCharge: initialiazationPage.RequestCharge, + activityId: initialiazationPage.ActivityId, + responseLengthInBytes: initialiazationPage.ResponseLengthInBytes, + cosmosQueryExecutionInfo: initialiazationPage.CosmosQueryExecutionInfo, + disallowContinuationTokenMessage: initialiazationPage.DisallowContinuationTokenMessage, + state: UninitializedQueryState)); + return true; + } + + if (this.enumerators.Count == 0) + { + return false; + } + + OrderByQueryPartitionRangePageAsyncEnumerator currentEnumerator = this.enumerators.Dequeue(); + if (!currentEnumerator.Current.Result.Enumerator.MoveNext()) + { + // The order by page ran out of results + if (await currentEnumerator.MoveNextAsync()) + { + this.enumerators.Enqueue(currentEnumerator); + + TryCatch monadicOrderByQueryPage = currentEnumerator.Current; + if (monadicOrderByQueryPage.Failed) + { + this.Current = TryCatch.FromException(monadicOrderByQueryPage.Exception); + return true; + } + + // Return an empty page with the query stats + QueryPage page = monadicOrderByQueryPage.Result.Page; + this.Current = TryCatch.FromResult( + new QueryPage( + documents: EmptyPage, + requestCharge: page.RequestCharge, + activityId: page.ActivityId, + responseLengthInBytes: page.ResponseLengthInBytes, + cosmosQueryExecutionInfo: page.CosmosQueryExecutionInfo, + disallowContinuationTokenMessage: page.DisallowContinuationTokenMessage, + state: page.State)); + return true; + } + + // recursively retry + return await this.MoveNextAsync(); + } + + // Create a query page with just the one document we consumed + // No stats to report, since we already reported it when we moved to this page. + CosmosElement orderByItem = currentEnumerator.Current.Result.Enumerator.Current; + this.Current = TryCatch.FromResult( + new QueryPage( + documents: new List() { orderByItem }, + requestCharge: 0, + activityId: default, + responseLengthInBytes: 0, + cosmosQueryExecutionInfo: default, + disallowContinuationTokenMessage: default, + state: currentEnumerator.Current.Result.State)); + return true; + } + + private async Task TryHandleExceptionAsync(OrderByQueryPartitionRangePageAsyncEnumerator enumerator) + { + if (!enumerator.Current.Failed) + { + throw new InvalidOperationException("Enumerator was not in a faulted state."); + } + + Exception exception = enumerator.Current.Exception; + + // Check if it's a retryable exception. + while (exception.InnerException != null) + { + exception = exception.InnerException; + } + + if (IsSplitException(exception)) + { + // Handle split + IEnumerable childRanges = await this.documentContainer.GetChildRangeAsync( + enumerator.Range, + cancellationToken: default); + foreach (PartitionKeyRange childRange in childRanges) + { + OrderByQueryPartitionRangePageAsyncEnumerator childPaginator = this.createEnumerator( + childRange, + enumerator.State, + enumerator.Filter); + this.uninitializedEnumerators.Enqueue(childPaginator); + } + + return true; + } + + if (IsMergeException(exception)) + { + throw new NotImplementedException(); + } + + return false; + } + + public ValueTask DisposeAsync() + { + // Do Nothing. + return default; + } + + private static bool IsSplitException(Exception exeception) + { + return exeception is CosmosException cosmosException + && (cosmosException.StatusCode == HttpStatusCode.Gone) + && (cosmosException.SubStatusCode == (int)Documents.SubStatusCodes.PartitionKeyRangeGone); + } + + private static bool IsMergeException(Exception exception) + { + // TODO: code this out + return false; + } + } +} diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderByQueryPartitionRangePageAsyncEnumerator.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderByQueryPartitionRangePageAsyncEnumerator.cs index 1d9b8f269b..d50f79b6e0 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderByQueryPartitionRangePageAsyncEnumerator.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderByQueryPartitionRangePageAsyncEnumerator.cs @@ -23,14 +23,18 @@ public OrderByQueryPartitionRangePageAsyncEnumerator( SqlQuerySpec sqlQuerySpec, PartitionKeyRange feedRange, int pageSize, + string filter, QueryState state = default) : base(feedRange, state) { this.queryDataSource = queryDataSource ?? throw new ArgumentNullException(nameof(queryDataSource)); this.sqlQuerySpec = sqlQuerySpec ?? throw new ArgumentNullException(nameof(sqlQuerySpec)); this.pageSize = pageSize; + this.Filter = filter; } + public string Filter { get; } + public override ValueTask DisposeAsync() => default; protected override Task> GetNextPageAsync(CancellationToken cancellationToken) => this.queryDataSource diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/CrossPartitionQueryPipelineStage.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/PartitionMapper.cs similarity index 95% rename from Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/CrossPartitionQueryPipelineStage.cs rename to Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/PartitionMapper.cs index 8eaa0d089a..7d4e7b52d1 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/CrossPartitionQueryPipelineStage.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/PartitionMapper.cs @@ -7,14 +7,15 @@ namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote using System; using System.Collections.Generic; using System.Linq; + using System.Threading.Tasks; using Microsoft.Azure.Cosmos.Query.Core.ContinuationTokens; using Microsoft.Azure.Cosmos.Query.Core.Exceptions; using Microsoft.Azure.Cosmos.Query.Core.Monads; using Microsoft.Azure.Documents; - internal abstract class CrossPartitionQueryPipelineStage : IQueryPipelineStage + internal static class PartitionMapper { - protected static TryCatch> MonadicGetPartitionMapping( + public static TryCatch> MonadicGetPartitionMapping( IReadOnlyList partitionKeyRanges, IReadOnlyList partitionedContinuationTokens) where PartitionedToken : IPartitionedToken @@ -101,7 +102,7 @@ protected static TryCatch> MonadicGetPartitio /// The partition key ranges to match. /// The continuation tokens to match with. /// A dictionary of ranges matched with their continuation tokens. - protected static IReadOnlyDictionary MatchRangesToContinuationTokens( + public static IReadOnlyDictionary MatchRangesToContinuationTokens( ReadOnlyMemory partitionKeyRanges, IReadOnlyList partitionedContinuationTokens) where PartitionedToken : IPartitionedToken @@ -137,7 +138,7 @@ protected static IReadOnlyDictionary MatchR return partitionKeyRangeToToken; } - protected readonly struct PartitionMapping + public readonly struct PartitionMapping { public PartitionMapping( IReadOnlyDictionary partitionsLeftOfTarget, From e0f8a07289e98fe4f25d34b59bf97d2bef9a317b Mon Sep 17 00:00:00 2001 From: Brandon Chong Date: Mon, 3 Aug 2020 10:36:45 -0700 Subject: [PATCH 50/85] need to refactor for failures inside of the filter logic --- ...OrderByCrossPartitionQueryPipelineStage.cs | 194 +++++++++--------- .../Remote/OrderByEnumeratorComparer.cs | 148 +++++++++++++ .../Pipeline/Remote/OrderByItemEnumerator.cs | 29 +-- 3 files changed, 262 insertions(+), 109 deletions(-) create mode 100644 Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderByEnumeratorComparer.cs diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderByCrossPartitionQueryPipelineStage.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderByCrossPartitionQueryPipelineStage.cs index a57183b43a..5869fdff69 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderByCrossPartitionQueryPipelineStage.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderByCrossPartitionQueryPipelineStage.cs @@ -30,8 +30,11 @@ namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote /// This way we can generate a single continuation token for all n partitions. /// This class is able to stop and resume execution by generating continuation tokens and reconstructing an execution context from said token. /// - internal sealed class OrderByCrossPartitionQueryPipelineStage : CrossPartitionQueryPipelineStage + internal sealed class OrderByCrossPartitionQueryPipelineStage : IQueryPipelineStage { + private static readonly QueryState UninitializedQueryState = new QueryState(CosmosString.Create("ORDER BY NOT INITIALIZED YET!")); + private static readonly IReadOnlyList EmptyPage = new List(); + private static class Expressions { public const string LessThan = "<"; @@ -56,7 +59,7 @@ private static class Expressions private readonly OrderByItemEnumerator orderByItemEnumerator; private readonly int pageSize; - private OrderByCrossPartitionQueryPipelineStage(OrderByItemEnumerator orderByItemEnumerator) + private OrderByCrossPartitionQueryPipelineStage() { this.orderByItemEnumerator = orderByItemEnumerator ?? throw new ArgumentNullException(nameof(orderByItemEnumerator)); } @@ -67,14 +70,9 @@ private OrderByCrossPartitionQueryPipelineStage(OrderByItemEnumerator orderByIte public async ValueTask MoveNextAsync() { - if (this.orderByItemEnumerator.Current.Failed) - { - // If we have a buffered exception, - // then return that first. - this.Current = this.orderByItemEnumerator.Current; - return true; - } - + bool movedNext = await this.orderByItemEnumerator.MoveNextAsync(); + this.Current = this.orderByItemEnumerator.Current; + return movedNext; } public static TryCatch MonadicCreate( @@ -124,7 +122,9 @@ public static TryCatch MonadicCreate( throw new ArgumentOutOfRangeException(nameof(pageSize)); } - List remoteEnumerators; + IReadOnlyList sortOrders = orderByColumns.Select(column => column.SortOrder).ToList(); + + List<(OrderByQueryPartitionRangePageAsyncEnumerator, OrderByContinuationToken)> enumeratorsAndTokens; if (continuationToken == null) { // Start off all the partition key ranges with null continuation @@ -132,13 +132,14 @@ public static TryCatch MonadicCreate( sqlQuerySpec.QueryText.Replace(oldValue: FormatPlaceHolder, newValue: TrueFilter), sqlQuerySpec.Parameters); - remoteEnumerators = targetRanges - .Select(range => new OrderByQueryPartitionRangePageAsyncEnumerator( + enumeratorsAndTokens = targetRanges + .Select(range => (new OrderByQueryPartitionRangePageAsyncEnumerator( documentContainer, sqlQuerySpec, range, pageSize, - state: default)) + TrueFilter, + state: default), (OrderByContinuationToken)null)) .ToList(); } else @@ -183,8 +184,7 @@ public static TryCatch MonadicCreate( { (partitionMapping.PartitionsRightOfTarget, rightFilter) }, }; - IReadOnlyList sortOrders = orderByColumns.Select(column => column.SortOrder).ToList(); - remoteEnumerators = new List(); + enumeratorsAndTokens = new List<(OrderByQueryPartitionRangePageAsyncEnumerator, OrderByContinuationToken)>(); foreach ((IReadOnlyDictionary tokenMapping, string filter) in tokenMappingAndFilters) { SqlQuerySpec rewrittenQueryForOrderBy = new SqlQuerySpec( @@ -200,16 +200,43 @@ public static TryCatch MonadicCreate( sqlQuerySpec, range, pageSize, + filter, state: token != null ? new QueryState(CosmosString.Create(token.CompositeContinuationToken.Token)) : null); - remoteEnumerators.Add(remoteEnumerator); + enumeratorsAndTokens.Add((remoteEnumerator, token)); } } } OrderByItemEnumerator orderByItemEnumerator = new OrderByItemEnumerator( documentContainer, - ) + createEnumerator: (range, state, filter) => new OrderByQueryPartitionRangePageAsyncEnumerator( + documentContainer, + sqlQuerySpec, + range, + pageSize, + filter, + state: state), + enumeratorsAndTokens, + initializeAsync: async (enumerator, token) => + { + TryCatch queryPage; + if (token is null) + { + await enumerator.MoveNextAsync(); + queryPage = enumerator.Current; + } + else + { + queryPage = await MonadicFilterAsync(enumerator, sortOrders, token, cancellationToken: default); + } + + return queryPage; + }, + new OrderByEnumeratorComparer(sortOrders)); + + OrderByCrossPartitionQueryPipelineStage stage = new OrderByCrossPartitionQueryPipelineStage(orderByItemEnumerator); + return TryCatch.FromResult(stage); } private static TryCatch> MonadicGetOrderByContinuationTokenMapping( @@ -393,7 +420,6 @@ private static (string leftFilter, string targetFilter, string rightFilter) GetF ReadOnlySpan<(OrderByColumn orderByColumn, CosmosElement orderByItem)> columnAndItemPrefix = columnAndItems.Span.Slice(start: 0, length: prefixLength); bool lastPrefix = prefixLength == numOrderByItems; - bool firstPrefix = prefixLength == 1; OrderByCrossPartitionQueryPipelineStage.AppendToBuilders(builders, "("); @@ -459,7 +485,7 @@ private static (string leftFilter, string targetFilter, string rightFilter) GetF /// The continuation token. /// The cancellation token. /// A task to await on. - private static async Task MonadicFilterAsync( + private static async Task)>> MonadicFilterNextAsync( OrderByQueryPartitionRangePageAsyncEnumerator enumerator, IReadOnlyList sortOrders, OrderByContinuationToken continuationToken, @@ -475,107 +501,83 @@ private static async Task MonadicFilterAsync( if (!ResourceId.TryParse(continuationToken.Rid, out ResourceId continuationRid)) { - return TryCatch.FromException( + return TryCatch<(bool, TryCatch)>.FromException( new MalformedContinuationTokenException( $"Invalid Rid in the continuation token {continuationToken.CompositeContinuationToken.Token} for OrderBy~Context.")); } - Dictionary resourceIds = new Dictionary(); int itemToSkip = continuationToken.SkipCount; - bool continuationRidVerified = false; // Throw away documents until it matches the item from the continuation token. - while (await enumerator.MoveNextAsync()) + if (!await enumerator.MoveNextAsync()) { - TryCatch monadicOrderByQueryPage = enumerator.Current; - if (monadicOrderByQueryPage.Failed) - { - return TryCatch.FromException(monadicOrderByQueryPage.Exception); - } + return TryCatch<(bool, TryCatch)>.FromResult((true, enumerator.Current)); + } - OrderByQueryPage orderByQueryPage = monadicOrderByQueryPage.Result; - IEnumerator documents = orderByQueryPage.Enumerator; - while (documents.MoveNext()) + TryCatch monadicOrderByQueryPage = enumerator.Current; + if (monadicOrderByQueryPage.Failed) + { + return TryCatch<(bool, TryCatch)>.FromException(monadicOrderByQueryPage.Exception); + } + + OrderByQueryPage orderByQueryPage = monadicOrderByQueryPage.Result; + IEnumerator documents = orderByQueryPage.Enumerator; + while (documents.MoveNext()) + { + OrderByQueryResult orderByResult = new OrderByQueryResult(documents.Current); + ResourceId rid = ResourceId.Parse(orderByResult.Rid); + int cmp = 0; + for (int i = 0; (i < sortOrders.Count) && (cmp == 0); ++i) { - OrderByQueryResult orderByResult = new OrderByQueryResult(documents.Current); + cmp = ItemComparer.Instance.Compare( + continuationToken.OrderByItems[i].Item, + orderByResult.OrderByItems[i].Item); - int cmp = 0; - for (int i = 0; (i < sortOrders.Count) && (cmp == 0); ++i) + if (cmp != 0) { - cmp = ItemComparer.Instance.Compare( - continuationToken.OrderByItems[i].Item, - orderByResult.OrderByItems[i].Item); - - if (cmp != 0) - { - cmp = sortOrders[i] == SortOrder.Ascending ? cmp : -cmp; - } + cmp = sortOrders[i] == SortOrder.Ascending ? cmp : -cmp; } + } - if (cmp < 0) - { - // We might have passed the item due to deletions and filters. - return TryCatch.FromResult(); - } + if (cmp < 0) + { + // We might have passed the item due to deletions and filters. + return TryCatch<(bool, TryCatch)>.FromResult((true, enumerator.Current)); + } - if (cmp == 0) + if (cmp == 0) + { + // Once the item matches the order by items from the continuation tokens + // We still need to remove all the documents that have a lower rid in the rid sort order. + // If there is a tie in the sort order the documents should be in _rid order in the same direction as the index (given by the backend) + cmp = continuationRid.Document.CompareTo(rid.Document); + if ((orderByQueryPage.Page.CosmosQueryExecutionInfo == null) || orderByQueryPage.Page.CosmosQueryExecutionInfo.ReverseRidEnabled) { - if (!resourceIds.TryGetValue(orderByResult.Rid, out ResourceId rid)) + // If reverse rid is enabled on the backend then fallback to the old way of doing it. + if (sortOrders[0] == SortOrder.Descending) { - if (!ResourceId.TryParse(orderByResult.Rid, out rid)) - { - return TryCatch.FromException( - new MalformedContinuationTokenException( - $"Invalid Rid in the continuation token {continuationToken.CompositeContinuationToken.Token} for OrderBy~Context~TryParse.")); - } - - resourceIds.Add(orderByResult.Rid, rid); + cmp = -cmp; } - - if (!continuationRidVerified) - { - if (continuationRid.Database != rid.Database || continuationRid.DocumentCollection != rid.DocumentCollection) - { - return TryCatch.FromException( - new MalformedContinuationTokenException( - $"Invalid Rid in the continuation token {continuationToken.CompositeContinuationToken.Token} for OrderBy~Context.")); - } - - continuationRidVerified = true; - } - - // Once the item matches the order by items from the continuation tokens - // We still need to remove all the documents that have a lower rid in the rid sort order. - // If there is a tie in the sort order the documents should be in _rid order in the same direction as the index (given by the backend) - cmp = continuationRid.Document.CompareTo(rid.Document); - if ((orderByQueryPage.Page.CosmosQueryExecutionInfo == null) || orderByQueryPage.Page.CosmosQueryExecutionInfo.ReverseRidEnabled) - { - // If reverse rid is enabled on the backend then fallback to the old way of doing it. - if (sortOrders[0] == SortOrder.Descending) - { - cmp = -cmp; - } - } - else + } + else + { + // Go by the whatever order the index wants + if (orderByQueryPage.Page.CosmosQueryExecutionInfo.ReverseIndexScan) { - // Go by the whatever order the index wants - if (orderByQueryPage.Page.CosmosQueryExecutionInfo.ReverseIndexScan) - { - cmp = -cmp; - } + cmp = -cmp; } + } - // We might have passed the item due to deletions and filters. - // We also have a skip count for JOINs - if (cmp < 0 || (cmp == 0 && itemToSkip-- <= 0)) - { - return TryCatch.FromResult(); - } + // We might have passed the item due to deletions and filters. + // We also have a skip count for JOINs + if (cmp < 0 || (cmp == 0 && itemToSkip-- <= 0)) + { + return TryCatch<(bool, TryCatch)>.FromResult((true, enumerator.Current)); } } } - return TryCatch.FromResult(); + return TryCatch<(bool, TryCatch)>.FromResult((false, enumerator.Current)); } private static bool IsSplitException(Exception exeception) diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderByEnumeratorComparer.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderByEnumeratorComparer.cs new file mode 100644 index 0000000000..03db9cef4e --- /dev/null +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderByEnumeratorComparer.cs @@ -0,0 +1,148 @@ +// ------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +// ------------------------------------------------------------ + +namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote +{ + using System; + using System.Collections.Generic; + using System.Diagnostics; + using System.Linq; + using System.Text; + using Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.ItemProducers; + using Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.OrderBy; + + /// + /// For cross partition order by queries we serve documents from the partition + /// that has the next document in the sort order of the query. + /// If there is a tie, then we break the tie by picking the leftmost partition. + /// + internal sealed class OrderByEnumeratorComparer : IComparer + { + /// + /// The sort orders for the query (1 for each order by in the query). + /// Until composite indexing is released this will just be an array of length 1. + /// + private readonly IReadOnlyList sortOrders; + + /// + /// Initializes a new instance of the OrderByConsumeComparer class. + /// + /// The sort orders for the query. + public OrderByEnumeratorComparer(IReadOnlyList sortOrders) + { + if (sortOrders == null) + { + throw new ArgumentNullException("Sort Orders array can not be null for an order by comparer."); + } + + if (sortOrders.Count == 0) + { + throw new ArgumentException("Sort Orders array can not be empty for an order by comparerer."); + } + + this.sortOrders = new List(sortOrders); + } + + /// + /// Compares two document producer trees and returns an integer with the relation of which has the document that comes first in the sort order. + /// + /// The first document producer tree. + /// The second document producer tree. + /// + /// Less than zero if the document in the first document producer comes first. + /// Zero if the documents are equivalent. + /// Greater than zero if the document in the second document producer comes first. + /// + public int Compare( + OrderByQueryPartitionRangePageAsyncEnumerator enumerator1, + OrderByQueryPartitionRangePageAsyncEnumerator enumerator2) + { + if (object.ReferenceEquals(enumerator1, enumerator2)) + { + return 0; + } + + if (enumerator1.Current.Failed && !enumerator2.Current.Failed) + { + return -1; + } + + if (!enumerator1.Current.Failed && enumerator2.Current.Failed) + { + return 1; + } + + if (enumerator1.Current.Failed && enumerator2.Current.Failed) + { + return string.CompareOrdinal(enumerator1.Range.MinInclusive, enumerator2.Range.MinInclusive); + } + + OrderByQueryResult result1 = new OrderByQueryResult(enumerator1.Current.Result.Page.Documents.First()); + OrderByQueryResult result2 = new OrderByQueryResult(enumerator2.Current.Result.Page.Documents.First()); + + // First compare the documents based on the sort order of the query. + int cmp = this.CompareOrderByItems(result1.OrderByItems, result2.OrderByItems); + if (cmp != 0) + { + // If there is no tie just return that. + return cmp; + } + + // If there is a tie, then break the tie by picking the one from the left most partition. + return string.CompareOrdinal(enumerator1.Range.MinInclusive, enumerator2.Range.MinInclusive); + + } + + /// + /// Takes the items relevant to the sort and return an integer defining the relationship. + /// + /// The items relevant to the sort from the first partition. + /// The items relevant to the sort from the second partition. + /// The sort relationship. + /// + /// Suppose the query was "SELECT * FROM c ORDER BY c.name asc, c.age desc", + /// then items1 could be ["Brandon", 22] and items2 could be ["Felix", 28] + /// Then we would first compare "Brandon" to "Felix" and say that "Brandon" comes first in an ascending lex order (we don't even have to look at age). + /// If items1 was ["Brandon", 22] and items2 was ["Brandon", 23] then we would say have to look at the age to break the tie and in this case 23 comes first in a descending order. + /// Some examples of composite order by: http://www.dofactory.com/sql/order-by + /// + public int CompareOrderByItems(IReadOnlyList items1, IReadOnlyList items2) + { + if (object.ReferenceEquals(items1, items2)) + { + return 0; + } + + Debug.Assert( + items1 != null && items2 != null, + "Order-by items must be present."); + + Debug.Assert( + items1.Count == items2.Count, + "OrderByResult instances should have the same number of order-by items."); + + Debug.Assert( + items1.Count > 0, + "OrderByResult instances should have at least 1 order-by item."); + + Debug.Assert( + this.sortOrders.Count == items1.Count, + "SortOrders must match size of order-by items."); + + for (int i = 0; i < this.sortOrders.Count; ++i) + { + int cmp = ItemComparer.Instance.Compare( + items1[i].Item, + items2[i].Item); + + if (cmp != 0) + { + return this.sortOrders[i] != SortOrder.Descending ? cmp : -cmp; + } + } + + return 0; + } + } +} diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderByItemEnumerator.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderByItemEnumerator.cs index fb0f65e66d..1887c2f5bd 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderByItemEnumerator.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderByItemEnumerator.cs @@ -11,6 +11,7 @@ namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote using Microsoft.Azure.Cosmos.CosmosElements; using Microsoft.Azure.Cosmos.Pagination; using Microsoft.Azure.Cosmos.Query.Core.Collections; + using Microsoft.Azure.Cosmos.Query.Core.ContinuationTokens; using Microsoft.Azure.Cosmos.Query.Core.Monads; using Microsoft.Azure.Documents; @@ -20,21 +21,21 @@ internal sealed class OrderByItemEnumerator : IAsyncEnumerator EmptyPage = new List(); private readonly IDocumentContainer documentContainer; private readonly Func createEnumerator; - private readonly Func>> intializeAsync; + private readonly Func>> intializeAsync; private readonly PriorityQueue enumerators; - private readonly Queue uninitializedEnumerators; + private readonly Queue<(OrderByQueryPartitionRangePageAsyncEnumerator, OrderByContinuationToken)> uninitializedEnumeratorsAndTokens; public OrderByItemEnumerator( IDocumentContainer documentContainer, Func createEnumerator, - IEnumerable uninitializedEnumerators, - Func>> initializeAsync, + IEnumerable<(OrderByQueryPartitionRangePageAsyncEnumerator, OrderByContinuationToken)> uninitializedEnumeratorsAndTokens, + Func>> initializeAsync, IComparer itemComparer) { this.documentContainer = documentContainer ?? throw new ArgumentNullException(nameof(documentContainer)); this.createEnumerator = createEnumerator ?? throw new ArgumentNullException(nameof(createEnumerator)); this.intializeAsync = initializeAsync ?? throw new ArgumentNullException(nameof(initializeAsync)); - this.uninitializedEnumerators = new Queue(uninitializedEnumerators); + this.uninitializedEnumeratorsAndTokens = new Queue<(OrderByQueryPartitionRangePageAsyncEnumerator, OrderByContinuationToken)>(uninitializedEnumeratorsAndTokens ?? throw new ArgumentNullException(nameof(uninitializedEnumeratorsAndTokens))); this.enumerators = new PriorityQueue(itemComparer); } @@ -42,15 +43,15 @@ public OrderByItemEnumerator( public async ValueTask MoveNextAsync() { - if (this.uninitializedEnumerators.Count != 0) + if (this.uninitializedEnumeratorsAndTokens.Count != 0) { - OrderByQueryPartitionRangePageAsyncEnumerator uninitializedEnumerator = this.uninitializedEnumerators.Dequeue(); - TryCatch initializeMonad = await this.intializeAsync(uninitializedEnumerator); + (OrderByQueryPartitionRangePageAsyncEnumerator uninitializedEnumerator, OrderByContinuationToken token) = this.uninitializedEnumeratorsAndTokens.Dequeue(); + TryCatch initializeMonad = await this.intializeAsync(uninitializedEnumerator, token); if (initializeMonad.Failed) { - if (!await this.TryHandleExceptionAsync(uninitializedEnumerator)) + if (!await this.TryHandleExceptionAsync(uninitializedEnumerator, token)) { - this.uninitializedEnumerators.Enqueue(uninitializedEnumerator); + this.uninitializedEnumeratorsAndTokens.Enqueue((uninitializedEnumerator, token)); } this.Current = TryCatch.FromException(uninitializedEnumerator.Current.Exception); @@ -62,7 +63,7 @@ public async ValueTask MoveNextAsync() // We want to report back the metrics from initialization, so that the user has accurate metrics, // But we need to make up a fake continuation token, since we aren't in a valid state to continue from. - QueryPage initialiazationPage = initializeMonad.Result; + QueryPage initialiazationPage = initializeMonad.Result.Page; this.Current = TryCatch.FromResult( new QueryPage( documents: EmptyPage, @@ -128,7 +129,9 @@ public async ValueTask MoveNextAsync() return true; } - private async Task TryHandleExceptionAsync(OrderByQueryPartitionRangePageAsyncEnumerator enumerator) + private async Task TryHandleExceptionAsync( + OrderByQueryPartitionRangePageAsyncEnumerator enumerator, + OrderByContinuationToken token) { if (!enumerator.Current.Failed) { @@ -155,7 +158,7 @@ private async Task TryHandleExceptionAsync(OrderByQueryPartitionRangePageA childRange, enumerator.State, enumerator.Filter); - this.uninitializedEnumerators.Enqueue(childPaginator); + this.uninitializedEnumeratorsAndTokens.Enqueue((childPaginator, token)); } return true; From c699205a72ac9b23edd4d2e8048f03a621e41101 Mon Sep 17 00:00:00 2001 From: Brandon Chong Date: Mon, 3 Aug 2020 11:36:59 -0700 Subject: [PATCH 51/85] drafted out a solution without splits and continuation token --- .../src/CosmosElements/CosmosArray.cs | 1 + .../src/CosmosElements/CosmosObject.cs | 1 + ...OrderByItemQueryExecutionContext.Resume.cs | 13 +- ...arallelItemQueryExecutionContext.Resume.cs | 5 +- .../Aggregate/AggregateQueryPipelineStage.cs | 2 +- .../NameCacheStaleRetryQueryPipelineStage.cs | 2 +- ...OrderByCrossPartitionQueryPipelineStage.cs | 199 +++++++++++++----- .../Pipeline/Remote/OrderByItemEnumerator.cs | 194 ----------------- .../Core/Pipeline/Remote/PartitionMapper.cs | 1 - .../Core/QueryClient/CosmosQueryClient.cs | 2 +- .../Query/v3Query/CosmosQueryClientCore.cs | 2 +- .../Query/MockCosmosQueryClient.cs | 2 +- .../FeedRange/FeedRangeTests.cs | 4 +- .../Query/ContinuationResumeLogicTests.cs | 7 +- 14 files changed, 170 insertions(+), 265 deletions(-) delete mode 100644 Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderByItemEnumerator.cs diff --git a/Microsoft.Azure.Cosmos/src/CosmosElements/CosmosArray.cs b/Microsoft.Azure.Cosmos/src/CosmosElements/CosmosArray.cs index 669c789163..041502ffca 100644 --- a/Microsoft.Azure.Cosmos/src/CosmosElements/CosmosArray.cs +++ b/Microsoft.Azure.Cosmos/src/CosmosElements/CosmosArray.cs @@ -12,6 +12,7 @@ namespace Microsoft.Azure.Cosmos.CosmosElements using Microsoft.Azure.Cosmos.Json; using Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.Distinct; using Microsoft.Azure.Cosmos.Query.Core.Monads; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Distinct; #if INTERNAL #pragma warning disable CS1591 // Missing XML comment for publicly visible type or member diff --git a/Microsoft.Azure.Cosmos/src/CosmosElements/CosmosObject.cs b/Microsoft.Azure.Cosmos/src/CosmosElements/CosmosObject.cs index a63393e201..9f4bcf2e48 100644 --- a/Microsoft.Azure.Cosmos/src/CosmosElements/CosmosObject.cs +++ b/Microsoft.Azure.Cosmos/src/CosmosElements/CosmosObject.cs @@ -12,6 +12,7 @@ namespace Microsoft.Azure.Cosmos.CosmosElements using Microsoft.Azure.Cosmos.Json; using Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.Distinct; using Microsoft.Azure.Cosmos.Query.Core.Monads; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Distinct; #if INTERNAL #pragma warning disable CS1591 // Missing XML comment for publicly visible type or member diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/OrderBy/CosmosOrderByItemQueryExecutionContext.Resume.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/OrderBy/CosmosOrderByItemQueryExecutionContext.Resume.cs index 21d64177ea..f099b268da 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/OrderBy/CosmosOrderByItemQueryExecutionContext.Resume.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/OrderBy/CosmosOrderByItemQueryExecutionContext.Resume.cs @@ -18,6 +18,7 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.OrderBy using Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.ItemProducers; using Microsoft.Azure.Cosmos.Query.Core.Monads; using Microsoft.Azure.Cosmos.Query.Core.QueryClient; + using static Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote.OrderByCrossPartitionQueryPipelineStage; using PartitionKeyRange = Documents.PartitionKeyRange; using ResourceId = Documents.ResourceId; @@ -60,15 +61,9 @@ public static async Task> MonadicCrea .Zip(sortOrders, (expression, order) => new OrderByColumn(expression, order)) .ToList(); - return (await context.TryInitializeAsync( - sqlQuerySpec: initParams.SqlQuerySpec, - requestContinuation: requestContinuationToken, - collectionRid: initParams.CollectionRid, - partitionKeyRanges: initParams.PartitionKeyRanges, - initialPageSize: initParams.InitialPageSize, - orderByColumns: columns, - cancellationToken: cancellationToken)) - .Try(() => context); + await Task.Delay(0); + + return default; } } } diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/Parallel/CosmosParallelItemQueryExecutionContext.Resume.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/Parallel/CosmosParallelItemQueryExecutionContext.Resume.cs index d5934bb618..19670c41b7 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/Parallel/CosmosParallelItemQueryExecutionContext.Resume.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/Parallel/CosmosParallelItemQueryExecutionContext.Resume.cs @@ -16,6 +16,7 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.Parallel using Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.ItemProducers; using Microsoft.Azure.Cosmos.Query.Core.Monads; using Microsoft.Azure.Cosmos.Query.Core.QueryClient; + using static Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote.PartitionMapper; using PartitionKeyRange = Documents.PartitionKeyRange; internal sealed partial class CosmosParallelItemQueryExecutionContext : CosmosCrossPartitionQueryExecutionContext @@ -143,9 +144,7 @@ private static TryCatch> TryGetPart return TryCatch>.FromException(tryParseCompositeContinuationTokens.Exception); } - return CosmosCrossPartitionQueryExecutionContext.TryGetInitializationInfo( - partitionKeyRanges, - tryParseCompositeContinuationTokens.Result); + return default; } private static TryCatch> TryParseCompositeContinuationList( diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Aggregate/AggregateQueryPipelineStage.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Aggregate/AggregateQueryPipelineStage.cs index cdcfb256b7..03ac6e4aa4 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Aggregate/AggregateQueryPipelineStage.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Aggregate/AggregateQueryPipelineStage.cs @@ -101,7 +101,7 @@ public RewrittenAggregateProjections(bool isValueAggregateQuery, CosmosElement r // SELECT VALUE [{"item": {"sum": SUM(c.blah), "count": COUNT(c.blah)}}] if (!(raw is CosmosArray aggregates)) { - throw new ArgumentException($"{nameof(RewrittenAggregateProjections)} was not an array for a value aggregate query. Type is: {raw.Type}"); + throw new ArgumentException($"{nameof(RewrittenAggregateProjections)} was not an array for a value aggregate query. Type is: {raw.GetType()}"); } this.Payload = aggregates[0]; diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/NameCacheStaleRetryQueryPipelineStage.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/NameCacheStaleRetryQueryPipelineStage.cs index a54189a209..20b30e7fc7 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/NameCacheStaleRetryQueryPipelineStage.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/NameCacheStaleRetryQueryPipelineStage.cs @@ -51,7 +51,7 @@ public async ValueTask MoveNextAsync() if (shouldRetry) { await this.cosmosQueryContext.QueryClient.ForceRefreshCollectionCacheAsync( - this.cosmosQueryContext.ResourceLink.OriginalString, + this.cosmosQueryContext.ResourceLink, default); this.alreadyRetried = true; await this.currentQueryPipelineStage.DisposeAsync(); diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderByCrossPartitionQueryPipelineStage.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderByCrossPartitionQueryPipelineStage.cs index 5869fdff69..f174031a57 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderByCrossPartitionQueryPipelineStage.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderByCrossPartitionQueryPipelineStage.cs @@ -32,18 +32,6 @@ namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote /// internal sealed class OrderByCrossPartitionQueryPipelineStage : IQueryPipelineStage { - private static readonly QueryState UninitializedQueryState = new QueryState(CosmosString.Create("ORDER BY NOT INITIALIZED YET!")); - private static readonly IReadOnlyList EmptyPage = new List(); - - private static class Expressions - { - public const string LessThan = "<"; - public const string LessThanOrEqualTo = "<="; - public const string EqualTo = "="; - public const string GreaterThan = ">"; - public const string GreaterThanOrEqualTo = ">="; - } - /// /// Order by queries are rewritten to allow us to inject a filter. /// This placeholder is so that we can just string replace it with the filter we want without having to understand the structure of the query. @@ -56,23 +44,161 @@ private static class Expressions /// private const string TrueFilter = "true"; - private readonly OrderByItemEnumerator orderByItemEnumerator; + private static readonly QueryState UninitializedQueryState = new QueryState(CosmosString.Create("ORDER BY NOT INITIALIZED YET!")); + private static readonly IReadOnlyList EmptyPage = new List(); + + private readonly IDocumentContainer documentContainer; + private readonly SqlQuerySpec sqlQuerySpec; + private readonly IReadOnlyList sortOrders; + private readonly PriorityQueue enumerators; + private readonly Queue<(OrderByQueryPartitionRangePageAsyncEnumerator, OrderByContinuationToken)> uninitializedEnumeratorsAndTokens; private readonly int pageSize; - private OrderByCrossPartitionQueryPipelineStage() + private QueryState state; + + private static class Expressions + { + public const string LessThan = "<"; + public const string LessThanOrEqualTo = "<="; + public const string EqualTo = "="; + public const string GreaterThan = ">"; + public const string GreaterThanOrEqualTo = ">="; + } + + private OrderByCrossPartitionQueryPipelineStage( + IDocumentContainer documentContainer, + SqlQuerySpec sqlQuerySpec, + IReadOnlyList sortOrders, + int pageSize, + IEnumerable<(OrderByQueryPartitionRangePageAsyncEnumerator, OrderByContinuationToken)> uninitializedEnumeratorsAndTokens, + QueryState state) { - this.orderByItemEnumerator = orderByItemEnumerator ?? throw new ArgumentNullException(nameof(orderByItemEnumerator)); + this.documentContainer = documentContainer ?? throw new ArgumentNullException(nameof(documentContainer)); + this.sqlQuerySpec = sqlQuerySpec ?? throw new ArgumentNullException(nameof(sqlQuerySpec)); + this.sortOrders = sortOrders ?? throw new ArgumentNullException(nameof(sortOrders)); + this.enumerators = new PriorityQueue(new OrderByEnumeratorComparer(this.sortOrders)); + this.pageSize = pageSize < 0 ? throw new ArgumentOutOfRangeException($"{nameof(pageSize)} must be a non negative number.") : pageSize; + this.uninitializedEnumeratorsAndTokens = new Queue<(OrderByQueryPartitionRangePageAsyncEnumerator, OrderByContinuationToken)>(uninitializedEnumeratorsAndTokens ?? throw new ArgumentNullException(nameof(uninitializedEnumeratorsAndTokens))); + this.state = state; } public TryCatch Current { get; private set; } - public ValueTask DisposeAsync() => this.orderByItemEnumerator.DisposeAsync(); + public ValueTask DisposeAsync() => default; public async ValueTask MoveNextAsync() { - bool movedNext = await this.orderByItemEnumerator.MoveNextAsync(); - this.Current = this.orderByItemEnumerator.Current; - return movedNext; + this.state = null; + if (this.uninitializedEnumeratorsAndTokens.Count != 0) + { + (OrderByQueryPartitionRangePageAsyncEnumerator uninitializedEnumerator, OrderByContinuationToken token) = this.uninitializedEnumeratorsAndTokens.Dequeue(); + if (token is null) + { + // We need to prime the page + if (!await uninitializedEnumerator.MoveNextAsync()) + { + throw new InvalidOperationException("Failed to prime the enumerator."); + } + + if (uninitializedEnumerator.Current.Failed) + { + //TODO HANDLE SPLIT + this.uninitializedEnumeratorsAndTokens.Enqueue((uninitializedEnumerator, token)); + this.Current = TryCatch.FromException(uninitializedEnumerator.Current.Exception); + } + else + { + this.enumerators.Enqueue(uninitializedEnumerator); + QueryPage page = uninitializedEnumerator.Current.Result.Page; + // Just return an empty page with the stats + this.Current = TryCatch.FromResult( + new QueryPage( + documents: EmptyPage, + requestCharge: page.RequestCharge, + activityId: page.ActivityId, + responseLengthInBytes: page.ResponseLengthInBytes, + cosmosQueryExecutionInfo: page.CosmosQueryExecutionInfo, + disallowContinuationTokenMessage: page.DisallowContinuationTokenMessage, + state: this.state)); + } + + return true; + } + else + { + // We need to actually filter the enumerator + TryCatch<(bool, TryCatch)> filterMonad = await FilterNextAsync( + uninitializedEnumerator, + this.sortOrders, + token, + cancellationToken: default); + + if (filterMonad.Failed) + { + //TODO HANDLE SPLIT + this.Current = TryCatch.FromException(filterMonad.Exception); + return true; + } + + (bool doneFiltering, TryCatch monadicQueryByPage) = filterMonad.Result; + if (doneFiltering) + { + this.enumerators.Enqueue(uninitializedEnumerator); + } + else + { + this.uninitializedEnumeratorsAndTokens.Enqueue((uninitializedEnumerator, token)); + } + + QueryPage page = uninitializedEnumerator.Current.Result.Page; + // Just return an empty page with the stats + this.Current = TryCatch.FromResult( + new QueryPage( + documents: EmptyPage, + requestCharge: page.RequestCharge, + activityId: page.ActivityId, + responseLengthInBytes: page.ResponseLengthInBytes, + cosmosQueryExecutionInfo: page.CosmosQueryExecutionInfo, + disallowContinuationTokenMessage: page.DisallowContinuationTokenMessage, + state: this.state)); + + return true; + } + } + + if (this.enumerators.Count == 0) + { + return false; + } + + // Try to form a page with as many items in the sorted order without having to do async work. + List documents = new List(); + while (documents.Count < this.pageSize) + { + OrderByQueryPartitionRangePageAsyncEnumerator currentEnumerator = this.enumerators.Dequeue(); + if (!currentEnumerator.Current.Result.Enumerator.MoveNext()) + { + // The order by page ran out of results + // mark the enumerator as unitialized and it will get requeueed on the next iteration with a fresh page. + this.uninitializedEnumeratorsAndTokens.Enqueue((currentEnumerator, (OrderByContinuationToken)null)); + break; + } + + documents.Add(currentEnumerator.Current.Result.Enumerator.Current); + } + + // Return a page of results + // No stats to report, since we already reported it when we moved to this page. + this.Current = TryCatch.FromResult( + new QueryPage( + documents: documents, + requestCharge: 0, + activityId: default, + responseLengthInBytes: 0, + cosmosQueryExecutionInfo: default, + disallowContinuationTokenMessage: default, + state: default /*TODO figure out continuation logic*/)); + return true; } public static TryCatch MonadicCreate( @@ -122,8 +248,6 @@ public static TryCatch MonadicCreate( throw new ArgumentOutOfRangeException(nameof(pageSize)); } - IReadOnlyList sortOrders = orderByColumns.Select(column => column.SortOrder).ToList(); - List<(OrderByQueryPartitionRangePageAsyncEnumerator, OrderByContinuationToken)> enumeratorsAndTokens; if (continuationToken == null) { @@ -208,34 +332,13 @@ public static TryCatch MonadicCreate( } } - OrderByItemEnumerator orderByItemEnumerator = new OrderByItemEnumerator( + OrderByCrossPartitionQueryPipelineStage stage = new OrderByCrossPartitionQueryPipelineStage( documentContainer, - createEnumerator: (range, state, filter) => new OrderByQueryPartitionRangePageAsyncEnumerator( - documentContainer, - sqlQuerySpec, - range, - pageSize, - filter, - state: state), + sqlQuerySpec, + orderByColumns.Select(column => column.SortOrder).ToList(), + pageSize, enumeratorsAndTokens, - initializeAsync: async (enumerator, token) => - { - TryCatch queryPage; - if (token is null) - { - await enumerator.MoveNextAsync(); - queryPage = enumerator.Current; - } - else - { - queryPage = await MonadicFilterAsync(enumerator, sortOrders, token, cancellationToken: default); - } - - return queryPage; - }, - new OrderByEnumeratorComparer(sortOrders)); - - OrderByCrossPartitionQueryPipelineStage stage = new OrderByCrossPartitionQueryPipelineStage(orderByItemEnumerator); + continuationToken == null ? null : new QueryState(continuationToken)); return TryCatch.FromResult(stage); } @@ -485,7 +588,7 @@ private static (string leftFilter, string targetFilter, string rightFilter) GetF /// The continuation token. /// The cancellation token. /// A task to await on. - private static async Task)>> MonadicFilterNextAsync( + private static async Task)>> FilterNextAsync( OrderByQueryPartitionRangePageAsyncEnumerator enumerator, IReadOnlyList sortOrders, OrderByContinuationToken continuationToken, diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderByItemEnumerator.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderByItemEnumerator.cs deleted file mode 100644 index 1887c2f5bd..0000000000 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderByItemEnumerator.cs +++ /dev/null @@ -1,194 +0,0 @@ -// ------------------------------------------------------------ -// Copyright (c) Microsoft Corporation. All rights reserved. -// ------------------------------------------------------------ - -namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote -{ - using System; - using System.Collections.Generic; - using System.Net; - using System.Threading.Tasks; - using Microsoft.Azure.Cosmos.CosmosElements; - using Microsoft.Azure.Cosmos.Pagination; - using Microsoft.Azure.Cosmos.Query.Core.Collections; - using Microsoft.Azure.Cosmos.Query.Core.ContinuationTokens; - using Microsoft.Azure.Cosmos.Query.Core.Monads; - using Microsoft.Azure.Documents; - - internal sealed class OrderByItemEnumerator : IAsyncEnumerator> - { - private static readonly QueryState UninitializedQueryState = new QueryState(CosmosString.Create("ORDER BY NOT INITIALIZED YET!")); - private static readonly IReadOnlyList EmptyPage = new List(); - private readonly IDocumentContainer documentContainer; - private readonly Func createEnumerator; - private readonly Func>> intializeAsync; - private readonly PriorityQueue enumerators; - private readonly Queue<(OrderByQueryPartitionRangePageAsyncEnumerator, OrderByContinuationToken)> uninitializedEnumeratorsAndTokens; - - public OrderByItemEnumerator( - IDocumentContainer documentContainer, - Func createEnumerator, - IEnumerable<(OrderByQueryPartitionRangePageAsyncEnumerator, OrderByContinuationToken)> uninitializedEnumeratorsAndTokens, - Func>> initializeAsync, - IComparer itemComparer) - { - this.documentContainer = documentContainer ?? throw new ArgumentNullException(nameof(documentContainer)); - this.createEnumerator = createEnumerator ?? throw new ArgumentNullException(nameof(createEnumerator)); - this.intializeAsync = initializeAsync ?? throw new ArgumentNullException(nameof(initializeAsync)); - this.uninitializedEnumeratorsAndTokens = new Queue<(OrderByQueryPartitionRangePageAsyncEnumerator, OrderByContinuationToken)>(uninitializedEnumeratorsAndTokens ?? throw new ArgumentNullException(nameof(uninitializedEnumeratorsAndTokens))); - this.enumerators = new PriorityQueue(itemComparer); - } - - public TryCatch Current { get; private set; } - - public async ValueTask MoveNextAsync() - { - if (this.uninitializedEnumeratorsAndTokens.Count != 0) - { - (OrderByQueryPartitionRangePageAsyncEnumerator uninitializedEnumerator, OrderByContinuationToken token) = this.uninitializedEnumeratorsAndTokens.Dequeue(); - TryCatch initializeMonad = await this.intializeAsync(uninitializedEnumerator, token); - if (initializeMonad.Failed) - { - if (!await this.TryHandleExceptionAsync(uninitializedEnumerator, token)) - { - this.uninitializedEnumeratorsAndTokens.Enqueue((uninitializedEnumerator, token)); - } - - this.Current = TryCatch.FromException(uninitializedEnumerator.Current.Exception); - return true; - } - - // Once the enumerator has been initialized we can add it back to the priority queue - this.enumerators.Enqueue(uninitializedEnumerator); - - // We want to report back the metrics from initialization, so that the user has accurate metrics, - // But we need to make up a fake continuation token, since we aren't in a valid state to continue from. - QueryPage initialiazationPage = initializeMonad.Result.Page; - this.Current = TryCatch.FromResult( - new QueryPage( - documents: EmptyPage, - requestCharge: initialiazationPage.RequestCharge, - activityId: initialiazationPage.ActivityId, - responseLengthInBytes: initialiazationPage.ResponseLengthInBytes, - cosmosQueryExecutionInfo: initialiazationPage.CosmosQueryExecutionInfo, - disallowContinuationTokenMessage: initialiazationPage.DisallowContinuationTokenMessage, - state: UninitializedQueryState)); - return true; - } - - if (this.enumerators.Count == 0) - { - return false; - } - - OrderByQueryPartitionRangePageAsyncEnumerator currentEnumerator = this.enumerators.Dequeue(); - if (!currentEnumerator.Current.Result.Enumerator.MoveNext()) - { - // The order by page ran out of results - if (await currentEnumerator.MoveNextAsync()) - { - this.enumerators.Enqueue(currentEnumerator); - - TryCatch monadicOrderByQueryPage = currentEnumerator.Current; - if (monadicOrderByQueryPage.Failed) - { - this.Current = TryCatch.FromException(monadicOrderByQueryPage.Exception); - return true; - } - - // Return an empty page with the query stats - QueryPage page = monadicOrderByQueryPage.Result.Page; - this.Current = TryCatch.FromResult( - new QueryPage( - documents: EmptyPage, - requestCharge: page.RequestCharge, - activityId: page.ActivityId, - responseLengthInBytes: page.ResponseLengthInBytes, - cosmosQueryExecutionInfo: page.CosmosQueryExecutionInfo, - disallowContinuationTokenMessage: page.DisallowContinuationTokenMessage, - state: page.State)); - return true; - } - - // recursively retry - return await this.MoveNextAsync(); - } - - // Create a query page with just the one document we consumed - // No stats to report, since we already reported it when we moved to this page. - CosmosElement orderByItem = currentEnumerator.Current.Result.Enumerator.Current; - this.Current = TryCatch.FromResult( - new QueryPage( - documents: new List() { orderByItem }, - requestCharge: 0, - activityId: default, - responseLengthInBytes: 0, - cosmosQueryExecutionInfo: default, - disallowContinuationTokenMessage: default, - state: currentEnumerator.Current.Result.State)); - return true; - } - - private async Task TryHandleExceptionAsync( - OrderByQueryPartitionRangePageAsyncEnumerator enumerator, - OrderByContinuationToken token) - { - if (!enumerator.Current.Failed) - { - throw new InvalidOperationException("Enumerator was not in a faulted state."); - } - - Exception exception = enumerator.Current.Exception; - - // Check if it's a retryable exception. - while (exception.InnerException != null) - { - exception = exception.InnerException; - } - - if (IsSplitException(exception)) - { - // Handle split - IEnumerable childRanges = await this.documentContainer.GetChildRangeAsync( - enumerator.Range, - cancellationToken: default); - foreach (PartitionKeyRange childRange in childRanges) - { - OrderByQueryPartitionRangePageAsyncEnumerator childPaginator = this.createEnumerator( - childRange, - enumerator.State, - enumerator.Filter); - this.uninitializedEnumeratorsAndTokens.Enqueue((childPaginator, token)); - } - - return true; - } - - if (IsMergeException(exception)) - { - throw new NotImplementedException(); - } - - return false; - } - - public ValueTask DisposeAsync() - { - // Do Nothing. - return default; - } - - private static bool IsSplitException(Exception exeception) - { - return exeception is CosmosException cosmosException - && (cosmosException.StatusCode == HttpStatusCode.Gone) - && (cosmosException.SubStatusCode == (int)Documents.SubStatusCodes.PartitionKeyRangeGone); - } - - private static bool IsMergeException(Exception exception) - { - // TODO: code this out - return false; - } - } -} diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/PartitionMapper.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/PartitionMapper.cs index 7d4e7b52d1..5b5a9d8b2a 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/PartitionMapper.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/PartitionMapper.cs @@ -7,7 +7,6 @@ namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote using System; using System.Collections.Generic; using System.Linq; - using System.Threading.Tasks; using Microsoft.Azure.Cosmos.Query.Core.ContinuationTokens; using Microsoft.Azure.Cosmos.Query.Core.Exceptions; using Microsoft.Azure.Cosmos.Query.Core.Monads; diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/QueryClient/CosmosQueryClient.cs b/Microsoft.Azure.Cosmos/src/Query/Core/QueryClient/CosmosQueryClient.cs index 21c2e6adee..66bb2639d9 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/QueryClient/CosmosQueryClient.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/QueryClient/CosmosQueryClient.cs @@ -46,7 +46,7 @@ public abstract Task> TryGetPartitionedQ CancellationToken cancellationToken); public abstract Task> ExecuteItemQueryAsync( - Uri resourceUri, + string resourceUri, Documents.ResourceType resourceType, Documents.OperationType operationType, Guid clientQueryCorrelationId, diff --git a/Microsoft.Azure.Cosmos/src/Query/v3Query/CosmosQueryClientCore.cs b/Microsoft.Azure.Cosmos/src/Query/v3Query/CosmosQueryClientCore.cs index d2801f9d28..c262d9801e 100644 --- a/Microsoft.Azure.Cosmos/src/Query/v3Query/CosmosQueryClientCore.cs +++ b/Microsoft.Azure.Cosmos/src/Query/v3Query/CosmosQueryClientCore.cs @@ -118,7 +118,7 @@ public override async Task> TryGetPartit } public override async Task> ExecuteItemQueryAsync( - Uri resourceUri, + string resourceUri, ResourceType resourceType, OperationType operationType, Guid clientQueryCorrelationId, diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/Query/MockCosmosQueryClient.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/Query/MockCosmosQueryClient.cs index 48ad198fce..dd9b120122 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/Query/MockCosmosQueryClient.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/Query/MockCosmosQueryClient.cs @@ -61,7 +61,7 @@ public override Task ExecuteQueryPlanRequestAsync } public override Task> ExecuteItemQueryAsync( - Uri resourceUri, + string resourceUri, ResourceType resourceType, OperationType operationType, Guid clientQueryCorrelationId, diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/FeedRange/FeedRangeTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/FeedRange/FeedRangeTests.cs index cd71bacb83..96686e2374 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/FeedRange/FeedRangeTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/FeedRange/FeedRangeTests.cs @@ -205,8 +205,8 @@ public void FeedRangeEPK_ToJsonFromJson() string representation = feedRangeEPK.ToJsonString(); FeedRangeEpk feedRangeEPKDeserialized = Cosmos.FeedRange.FromJsonString(representation) as FeedRangeEpk; Assert.IsNotNull(feedRangeEPKDeserialized); - Assert.AreEqual(FeedRangeEpk.Range.Min, feedRangeEPKDeserialized.Range.Min); - Assert.AreEqual(FeedRangeEpk.Range.Max, feedRangeEPKDeserialized.Range.Max); + Assert.AreEqual(feedRangeEPK.Range.Min, feedRangeEPKDeserialized.Range.Min); + Assert.AreEqual(feedRangeEPK.Range.Max, feedRangeEPKDeserialized.Range.Max); } [TestMethod] diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/ContinuationResumeLogicTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/ContinuationResumeLogicTests.cs index 2bba535810..5a589533e0 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/ContinuationResumeLogicTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/ContinuationResumeLogicTests.cs @@ -6,10 +6,11 @@ using Microsoft.Azure.Cosmos.Query.Core.ContinuationTokens; using Microsoft.Azure.Cosmos.Query.Core.ExecutionContext; using Microsoft.Azure.Cosmos.Query.Core.Monads; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote; using Microsoft.Azure.Documents; using Microsoft.VisualStudio.TestTools.UnitTesting; using Newtonsoft.Json; - using static Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.CosmosCrossPartitionQueryExecutionContext; + using static Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote.PartitionMapper; [TestClass] public class ContinuationResumeLogicTests @@ -301,7 +302,7 @@ private static void RunMatchRangesToContinuationTokens( IEnumerable partitionKeyRanges, IEnumerable partitionedTokens) { - IReadOnlyDictionary actualMapping = CosmosCrossPartitionQueryExecutionContext.MatchRangesToContinuationTokens( + IReadOnlyDictionary actualMapping = PartitionMapper.MatchRangesToContinuationTokens( partitionKeyRanges.OrderBy(x => Guid.NewGuid()).ToArray(), partitionedTokens.OrderBy(x => Guid.NewGuid()).ToList()); @@ -317,7 +318,7 @@ private static void RunTryGetInitializationInfo( IEnumerable partitionKeyRanges, IEnumerable partitionedTokens) { - TryCatch> tryGetInitializationInfo = CosmosCrossPartitionQueryExecutionContext.TryGetInitializationInfo( + TryCatch> tryGetInitializationInfo = PartitionMapper.MonadicGetPartitionMapping( partitionKeyRanges.OrderBy(x => Guid.NewGuid()).ToArray(), partitionedTokens.OrderBy(x => Guid.NewGuid()).ToList()); Assert.IsTrue(tryGetInitializationInfo.Succeeded); From c299f7d13544ba02b9c227ba01877e49b363d51e Mon Sep 17 00:00:00 2001 From: Brandon Chong Date: Mon, 3 Aug 2020 15:38:44 -0700 Subject: [PATCH 52/85] got basic order by working --- ...OrderByCrossPartitionQueryPipelineStage.cs | 43 ++- .../Remote/OrderByEnumeratorComparer.cs | 4 +- .../Core/Pipeline/Remote/OrderByQueryPage.cs | 2 + .../Pagination/InMemoryContainer.cs | 172 ++++++++-- ...ByCrossPartitionQueryPipelineStageTests.cs | 310 ++++++++++++++++++ 5 files changed, 489 insertions(+), 42 deletions(-) create mode 100644 Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/OrderByCrossPartitionQueryPipelineStageTests.cs diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderByCrossPartitionQueryPipelineStage.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderByCrossPartitionQueryPipelineStage.cs index f174031a57..73897e2ee4 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderByCrossPartitionQueryPipelineStage.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderByCrossPartitionQueryPipelineStage.cs @@ -97,7 +97,17 @@ public async ValueTask MoveNextAsync() // We need to prime the page if (!await uninitializedEnumerator.MoveNextAsync()) { - throw new InvalidOperationException("Failed to prime the enumerator."); + // No more documents, so just return an empty page + this.Current = TryCatch.FromResult( + new QueryPage( + documents: EmptyPage, + requestCharge: 0, + activityId: string.Empty, + responseLengthInBytes: 0, + cosmosQueryExecutionInfo: default, + disallowContinuationTokenMessage: default, + state: this.state)); + return true; } if (uninitializedEnumerator.Current.Failed) @@ -108,7 +118,16 @@ public async ValueTask MoveNextAsync() } else { - this.enumerators.Enqueue(uninitializedEnumerator); + if (!uninitializedEnumerator.Current.Result.Enumerator.MoveNext()) + { + // Page was empty + this.uninitializedEnumeratorsAndTokens.Enqueue((uninitializedEnumerator, token)); + } + else + { + this.enumerators.Enqueue(uninitializedEnumerator); + } + QueryPage page = uninitializedEnumerator.Current.Result.Page; // Just return an empty page with the stats this.Current = TryCatch.FromResult( @@ -143,7 +162,11 @@ public async ValueTask MoveNextAsync() (bool doneFiltering, TryCatch monadicQueryByPage) = filterMonad.Result; if (doneFiltering) { - this.enumerators.Enqueue(uninitializedEnumerator); + // TODO test empty pages + if (uninitializedEnumerator.Current.Result.Enumerator.Current != null) + { + this.enumerators.Enqueue(uninitializedEnumerator); + } } else { @@ -176,15 +199,23 @@ public async ValueTask MoveNextAsync() while (documents.Count < this.pageSize) { OrderByQueryPartitionRangePageAsyncEnumerator currentEnumerator = this.enumerators.Dequeue(); + OrderByQueryResult orderByQueryResult = new OrderByQueryResult(currentEnumerator.Current.Result.Enumerator.Current); + documents.Add(orderByQueryResult.Payload); + if (!currentEnumerator.Current.Result.Enumerator.MoveNext()) { // The order by page ran out of results - // mark the enumerator as unitialized and it will get requeueed on the next iteration with a fresh page. - this.uninitializedEnumeratorsAndTokens.Enqueue((currentEnumerator, (OrderByContinuationToken)null)); + if (currentEnumerator.State != null) + { + // If the continuation isn't null + // then mark the enumerator as unitialized and it will get requeueed on the next iteration with a fresh page. + this.uninitializedEnumeratorsAndTokens.Enqueue((currentEnumerator, (OrderByContinuationToken)null)); + } + break; } - documents.Add(currentEnumerator.Current.Result.Enumerator.Current); + this.enumerators.Enqueue(currentEnumerator); } // Return a page of results diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderByEnumeratorComparer.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderByEnumeratorComparer.cs index 03db9cef4e..8e6125891b 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderByEnumeratorComparer.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderByEnumeratorComparer.cs @@ -78,8 +78,8 @@ public int Compare( return string.CompareOrdinal(enumerator1.Range.MinInclusive, enumerator2.Range.MinInclusive); } - OrderByQueryResult result1 = new OrderByQueryResult(enumerator1.Current.Result.Page.Documents.First()); - OrderByQueryResult result2 = new OrderByQueryResult(enumerator2.Current.Result.Page.Documents.First()); + OrderByQueryResult result1 = new OrderByQueryResult(enumerator1.Current.Result.Enumerator.Current); + OrderByQueryResult result2 = new OrderByQueryResult(enumerator2.Current.Result.Enumerator.Current); // First compare the documents based on the sort order of the query. int cmp = this.CompareOrderByItems(result1.OrderByItems, result2.OrderByItems); diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderByQueryPage.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderByQueryPage.cs index 5467d32485..11b6218a87 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderByQueryPage.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderByQueryPage.cs @@ -4,6 +4,7 @@ namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote { + using System; using System.Collections.Generic; using Microsoft.Azure.Cosmos.CosmosElements; using Microsoft.Azure.Cosmos.Pagination; @@ -13,6 +14,7 @@ internal sealed class OrderByQueryPage : Page public OrderByQueryPage(QueryPage queryPage) : base(queryPage.State) { + this.Page = queryPage ?? throw new ArgumentNullException(nameof(queryPage)); this.Enumerator = queryPage.Documents.GetEnumerator(); } 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 780e955370..a1a2b6bf07 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 @@ -298,54 +298,158 @@ public async Task> MonadicQueryAsync( throw new ArgumentNullException(nameof(sqlQuerySpec)); } - if (!sqlQuerySpec.QueryText.Equals("SELECT * FROM c", StringComparison.OrdinalIgnoreCase)) + const string selectStar = "SELECT * FROM c"; + const string orderBy = "SELECT r._rid, [{\"item\": c._ts}] AS orderByItems, c AS payload FROM Root AS c ORDER BY c._ts"; + + if (!sqlQuerySpec.QueryText.Equals(selectStar, StringComparison.OrdinalIgnoreCase) + && !sqlQuerySpec.QueryText.Equals(orderBy, StringComparison.OrdinalIgnoreCase)) { - throw new NotSupportedException("InMemoryCollection only supports SELECT * FROM c queries"); + throw new NotSupportedException("InMemoryCollection only supports SELECT * FROM c and order by _ts queries"); } - long resourceIdentifier = continuationToken != null ? long.Parse(continuationToken) : 0; + if (sqlQuerySpec.QueryText.Equals(selectStar, StringComparison.OrdinalIgnoreCase)) + { + // For now just do a read feed + long resourceIdentifier = continuationToken != null ? long.Parse(continuationToken) : 0; + TryCatch tryGetPage = await this.MonadicReadFeedAsync( + partitionKeyRangeId, + resourceIdentifier, + pageSize, + cancellationToken); - // For now just do a read feed - TryCatch tryGetPage = await this.MonadicReadFeedAsync( - partitionKeyRangeId, - resourceIdentifier, - pageSize, - cancellationToken); + if (tryGetPage.Failed) + { + return TryCatch.FromException(tryGetPage.Exception); + } - if (tryGetPage.Failed) - { - return TryCatch.FromException(tryGetPage.Exception); - } + DocumentContainerPage page = tryGetPage.Result; + List documents = new List(page.Records.Count); + foreach (Record record in page.Records) + { + Dictionary keyValuePairs = new Dictionary + { + ["_rid"] = CosmosNumber64.Create(record.ResourceIdentifier), + ["_ts"] = CosmosNumber64.Create(record.Timestamp), + ["id"] = CosmosString.Create(record.Identifier) + }; - DocumentContainerPage page = tryGetPage.Result; - List documents = new List(page.Records.Count); - foreach (Record record in page.Records) + foreach (KeyValuePair property in record.Payload) + { + keyValuePairs[property.Key] = property.Value; + } + + documents.Add(CosmosObject.Create(keyValuePairs)); + } + + QueryPage queryPage = new QueryPage( + documents: documents, + requestCharge: 42, + activityId: Guid.NewGuid().ToString(), + responseLengthInBytes: 1337, + cosmosQueryExecutionInfo: default, + disallowContinuationTokenMessage: default, + state: page.State != null ? new QueryState(CosmosString.Create(page.State.ResourceIdentifer.ToString())) : null); + + return TryCatch.FromResult(queryPage); + } + else { - Dictionary keyValuePairs = new Dictionary + // We need to simulate an order by + (long timestamp, long resourceIdentifer) timestampAndIdentifier; + if (continuationToken == null) { - ["_rid"] = CosmosNumber64.Create(record.ResourceIdentifier), - ["_ts"] = CosmosNumber64.Create(record.Timestamp), - ["id"] = CosmosString.Create(record.Identifier) - }; + timestampAndIdentifier = (-1, -1); + } + else + { + string[] tokens = continuationToken.Split(","); + timestampAndIdentifier = (long.Parse(tokens[0]), long.Parse(tokens[1])); + } - foreach (KeyValuePair property in record.Payload) + if (!this.partitionKeyRangeIdToHashRange.TryGetValue( + partitionKeyRangeId, + out PartitionKeyHashRange range)) { - keyValuePairs[property.Key] = property.Value; + return TryCatch.FromException( + new CosmosException( + message: $"PartitionKeyRangeId {partitionKeyRangeId} is gone", + statusCode: System.Net.HttpStatusCode.Gone, + subStatusCode: (int)SubStatusCodes.PartitionKeyRangeGone, + activityId: Guid.NewGuid().ToString(), + requestCharge: 42)); } - documents.Add(CosmosObject.Create(keyValuePairs)); - } + if (!this.partitionedRecords.TryGetValue(range, out Records records)) + { + throw new InvalidOperationException("failed to find the range."); + } + + List results = records + .OrderBy(record => record.Timestamp) + .ThenBy(record => record.ResourceIdentifier) + .Where(record => + (record.Timestamp > timestampAndIdentifier.timestamp) + || ((record.Timestamp == timestampAndIdentifier.timestamp) && (record.ResourceIdentifier > timestampAndIdentifier.resourceIdentifer))) + .Take(pageSize) + .ToList(); - QueryPage queryPage = new QueryPage( - documents: documents, - requestCharge: 42, - activityId: Guid.NewGuid().ToString(), - responseLengthInBytes: 1337, - cosmosQueryExecutionInfo: default, - disallowContinuationTokenMessage: default, - state: page.State != null ? new QueryState(CosmosString.Create(page.State.ResourceIdentifer.ToString())) : null); + if (results.Count == 0) + { + return TryCatch.FromResult( + new QueryPage( + documents: new List(), + requestCharge: 42, + activityId: Guid.NewGuid().ToString(), + responseLengthInBytes: 1337, + cosmosQueryExecutionInfo: default, + disallowContinuationTokenMessage: default, + state: null)); + } - return TryCatch.FromResult(queryPage); + List documents = new List(results.Count); + foreach (Record record in results) + { + Dictionary payloadProperties = new Dictionary() + { + ["_rid"] = CosmosNumber64.Create(record.ResourceIdentifier), + ["_ts"] = CosmosNumber64.Create(record.Timestamp), + ["id"] = CosmosString.Create(record.Identifier) + }; + + foreach (KeyValuePair property in record.Payload) + { + payloadProperties[property.Key] = property.Value; + } + + Dictionary keyValuePairs = new Dictionary + { + ["_rid"] = CosmosNumber64.Create(record.ResourceIdentifier), + ["orderByItems"] = CosmosArray.Create( + new List() + { + CosmosObject.Create( + new Dictionary() + { + ["item"] = CosmosNumber64.Create(record.Timestamp) + }) + }), + ["payload"] = CosmosObject.Create(payloadProperties), + }; + + documents.Add(CosmosObject.Create(keyValuePairs)); + } + + QueryPage queryPage = new QueryPage( + documents: documents, + requestCharge: 42, + activityId: Guid.NewGuid().ToString(), + responseLengthInBytes: 1337, + cosmosQueryExecutionInfo: default, + disallowContinuationTokenMessage: default, + state: new QueryState(CosmosString.Create($"{results.Last().Timestamp},{results.Last().ResourceIdentifier}"))); + + return TryCatch.FromResult(queryPage); + } } public Task> MonadicQueryAsync( diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/OrderByCrossPartitionQueryPipelineStageTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/OrderByCrossPartitionQueryPipelineStageTests.cs new file mode 100644 index 0000000000..83796b418f --- /dev/null +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/OrderByCrossPartitionQueryPipelineStageTests.cs @@ -0,0 +1,310 @@ +//------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +//------------------------------------------------------------ + +namespace Microsoft.Azure.Cosmos.Tests.Query.Pipeline +{ + using System.Collections.Generic; + using System.Linq; + using System.Threading.Tasks; + using Microsoft.Azure.Cosmos.CosmosElements; + using Microsoft.Azure.Cosmos.Pagination; + using Microsoft.Azure.Cosmos.Query.Core; + using Microsoft.Azure.Cosmos.Query.Core.ContinuationTokens; + using Microsoft.Azure.Cosmos.Query.Core.Exceptions; + using Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.OrderBy; + using Microsoft.Azure.Cosmos.Query.Core.Monads; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote; + using Microsoft.Azure.Cosmos.Tests.Pagination; + using Microsoft.Azure.Documents; + using Microsoft.VisualStudio.TestTools.UnitTesting; + using Moq; + + [TestClass] + public class OrderByCrossPartitionQueryPipelineStageTests + { + [TestMethod] + public void MonadicCreate_NullContinuationToken() + { + Mock mockDocumentContainer = new Mock(); + + TryCatch monadicCreate = OrderByCrossPartitionQueryPipelineStage.MonadicCreate( + documentContainer: mockDocumentContainer.Object, + sqlQuerySpec: new SqlQuerySpec("SELECT * FROM c ORDER BY c._ts"), + targetRanges: new List() { new PartitionKeyRange() }, + orderByColumns: new List() + { + new OrderByCrossPartitionQueryPipelineStage.OrderByColumn("_ts", Cosmos.Query.Core.ExecutionContext.OrderBy.SortOrder.Ascending) + }, + pageSize: 10, + continuationToken: null); + Assert.IsTrue(monadicCreate.Succeeded); + } + + [TestMethod] + public void MonadicCreate_NonCosmosArrayContinuationToken() + { + Mock mockDocumentContainer = new Mock(); + + TryCatch monadicCreate = OrderByCrossPartitionQueryPipelineStage.MonadicCreate( + documentContainer: mockDocumentContainer.Object, + sqlQuerySpec: new SqlQuerySpec("SELECT * FROM c ORDER BY c._ts"), + targetRanges: new List() { new PartitionKeyRange() }, + orderByColumns: new List() + { + new OrderByCrossPartitionQueryPipelineStage.OrderByColumn("_ts", Cosmos.Query.Core.ExecutionContext.OrderBy.SortOrder.Ascending) + }, + pageSize: 10, + continuationToken: CosmosObject.Create(new Dictionary())); + Assert.IsTrue(monadicCreate.Failed); + Assert.IsTrue(monadicCreate.InnerMostException is MalformedContinuationTokenException); + } + + [TestMethod] + public void MonadicCreate_EmptyArrayContinuationToken() + { + Mock mockDocumentContainer = new Mock(); + + TryCatch monadicCreate = OrderByCrossPartitionQueryPipelineStage.MonadicCreate( + documentContainer: mockDocumentContainer.Object, + sqlQuerySpec: new SqlQuerySpec("SELECT * FROM c ORDER BY c._ts"), + targetRanges: new List() { new PartitionKeyRange() }, + orderByColumns: new List() + { + new OrderByCrossPartitionQueryPipelineStage.OrderByColumn("_ts", Cosmos.Query.Core.ExecutionContext.OrderBy.SortOrder.Ascending) + }, + pageSize: 10, + continuationToken: CosmosArray.Create(new List())); + Assert.IsTrue(monadicCreate.Failed); + Assert.IsTrue(monadicCreate.InnerMostException is MalformedContinuationTokenException); + } + + [TestMethod] + public void MonadicCreate_NonCompositeContinuationToken() + { + Mock mockDocumentContainer = new Mock(); + + TryCatch monadicCreate = OrderByCrossPartitionQueryPipelineStage.MonadicCreate( + documentContainer: mockDocumentContainer.Object, + sqlQuerySpec: new SqlQuerySpec("SELECT * FROM c ORDER BY c._ts"), + targetRanges: new List() { new PartitionKeyRange() }, + orderByColumns: new List() + { + new OrderByCrossPartitionQueryPipelineStage.OrderByColumn("_ts", Cosmos.Query.Core.ExecutionContext.OrderBy.SortOrder.Ascending) + }, + pageSize: 10, + continuationToken: CosmosArray.Create(new List() { CosmosString.Create("asdf") })); + Assert.IsTrue(monadicCreate.Failed); + Assert.IsTrue(monadicCreate.InnerMostException is MalformedContinuationTokenException); + } + + [TestMethod] + public void MonadicCreate_SingleOrderByContinuationToken() + { + Mock mockDocumentContainer = new Mock(); + + CompositeContinuationToken compositeContinuationToken = new CompositeContinuationToken() + { + Range = new Documents.Routing.Range("A", "B", true, false), + Token = "asdf", + }; + + OrderByContinuationToken orderByContinuationToken = new OrderByContinuationToken( + compositeContinuationToken, + new List() { new OrderByItem(CosmosObject.Create(new Dictionary() { { "item", CosmosString.Create("asdf") } })) }, + rid: "rid", + skipCount: 42, + filter: "filter"); + + TryCatch monadicCreate = OrderByCrossPartitionQueryPipelineStage.MonadicCreate( + documentContainer: mockDocumentContainer.Object, + sqlQuerySpec: new SqlQuerySpec("SELECT * FROM c ORDER BY c._ts"), + targetRanges: new List() { new PartitionKeyRange() { Id = "0", MinInclusive = "A", MaxExclusive = "B" } }, + orderByColumns: new List() + { + new OrderByCrossPartitionQueryPipelineStage.OrderByColumn("_ts", Cosmos.Query.Core.ExecutionContext.OrderBy.SortOrder.Ascending) + }, + pageSize: 10, + continuationToken: CosmosArray.Create( + new List() + { + OrderByContinuationToken.ToCosmosElement(orderByContinuationToken) + })); + Assert.IsTrue(monadicCreate.Succeeded); + } + + [TestMethod] + public void MonadicCreate_MultipleOrderByContinuationToken() + { + Mock mockDocumentContainer = new Mock(); + + CompositeContinuationToken token = new CompositeContinuationToken() + { + Range = new Documents.Routing.Range("A", "B", true, false), + Token = "asdf", + }; + + CompositeContinuationToken compositeContinuationToken1 = new CompositeContinuationToken() + { + Range = new Documents.Routing.Range("A", "B", true, false), + Token = "asdf", + }; + + OrderByContinuationToken orderByContinuationToken1 = new OrderByContinuationToken( + compositeContinuationToken1, + new List() { new OrderByItem(CosmosObject.Create(new Dictionary() { { "item", CosmosString.Create("asdf") } })) }, + rid: "rid", + skipCount: 42, + filter: "filter"); + + CompositeContinuationToken compositeContinuationToken2 = new CompositeContinuationToken() + { + Range = new Documents.Routing.Range("B", "C", true, false), + Token = "asdf", + }; + + OrderByContinuationToken orderByContinuationToken2 = new OrderByContinuationToken( + compositeContinuationToken2, + new List() { new OrderByItem(CosmosObject.Create(new Dictionary() { { "item", CosmosString.Create("asdf") } })) }, + rid: "rid", + skipCount: 42, + filter: "filter"); + + TryCatch monadicCreate = OrderByCrossPartitionQueryPipelineStage.MonadicCreate( + documentContainer: mockDocumentContainer.Object, + sqlQuerySpec: new SqlQuerySpec("SELECT * FROM c ORDER BY c._ts"), + targetRanges: new List() + { + new PartitionKeyRange() { Id = "0", MinInclusive = "A", MaxExclusive = "B" }, + new PartitionKeyRange() { Id = "1", MinInclusive = "B", MaxExclusive = "C" } + }, + orderByColumns: new List() + { + new OrderByCrossPartitionQueryPipelineStage.OrderByColumn("_ts", Cosmos.Query.Core.ExecutionContext.OrderBy.SortOrder.Ascending) + }, + pageSize: 10, + continuationToken: CosmosArray.Create( + new List() + { + OrderByContinuationToken.ToCosmosElement(orderByContinuationToken1), + OrderByContinuationToken.ToCosmosElement(orderByContinuationToken2) + })); + Assert.IsTrue(monadicCreate.Succeeded); + } + + [TestMethod] + public async Task TestDrainFully_StartFromBeginingAsync() + { + int numItems = 100; + IDocumentContainer documentContainer = await CreateDocumentContainerAsync(numItems); + + TryCatch monadicCreate = OrderByCrossPartitionQueryPipelineStage.MonadicCreate( + documentContainer: documentContainer, + sqlQuerySpec: new SqlQuerySpec("SELECT r._rid, [{\"item\": c._ts}] AS orderByItems, c AS payload FROM Root AS c ORDER BY c._ts"), + targetRanges: await documentContainer.GetFeedRangesAsync(cancellationToken: default), + orderByColumns: new List() + { + new OrderByCrossPartitionQueryPipelineStage.OrderByColumn("_ts", SortOrder.Ascending) + }, + pageSize: 10, + continuationToken: null); + Assert.IsTrue(monadicCreate.Succeeded); + IQueryPipelineStage queryPipelineStage = monadicCreate.Result; + + List documents = new List(); + while (await queryPipelineStage.MoveNextAsync()) + { + TryCatch tryGetQueryPage = queryPipelineStage.Current; + Assert.IsTrue(tryGetQueryPage.Succeeded); + + QueryPage queryPage = tryGetQueryPage.Result; + documents.AddRange(queryPage.Documents); + } + + Assert.AreEqual(numItems, documents.Count); + Assert.IsTrue(documents.OrderBy(document => ((CosmosObject)document)["_ts"]).ToList().SequenceEqual(documents)); + } + + [TestMethod] + public async Task TestDrainFully_WithStateResume() + { + int numItems = 1000; + IDocumentContainer documentContainer = await CreateDocumentContainerAsync(numItems); + + List documents = new List(); + + QueryState queryState = null; + do + { + TryCatch monadicCreate = ParallelCrossPartitionQueryPipelineStage.MonadicCreate( + documentContainer: documentContainer, + sqlQuerySpec: new SqlQuerySpec("SELECT * FROM c"), + pageSize: 10, + continuationToken: queryState?.Value); + Assert.IsTrue(monadicCreate.Succeeded); + IQueryPipelineStage queryPipelineStage = monadicCreate.Result; + + Assert.IsTrue(await queryPipelineStage.MoveNextAsync()); + TryCatch tryGetQueryPage = queryPipelineStage.Current; + Assert.IsTrue(tryGetQueryPage.Succeeded); + + QueryPage queryPage = tryGetQueryPage.Result; + documents.AddRange(queryPage.Documents); + + queryState = queryPage.State; + } while (queryState != null); + + Assert.AreEqual(numItems, documents.Count); + } + + private static async Task CreateDocumentContainerAsync( + int numItems, + FlakyDocumentContainer.FailureConfigs failureConfigs = null) + { + PartitionKeyDefinition partitionKeyDefinition = new PartitionKeyDefinition() + { + Paths = new System.Collections.ObjectModel.Collection() + { + "/pk" + }, + Kind = PartitionKind.Hash, + Version = PartitionKeyDefinitionVersion.V2, + }; + + IMonadicDocumentContainer monadicDocumentContainer = new InMemoryContainer(partitionKeyDefinition); + if (failureConfigs != null) + { + monadicDocumentContainer = new FlakyDocumentContainer(monadicDocumentContainer, failureConfigs); + } + + DocumentContainer documentContainer = new DocumentContainer(monadicDocumentContainer); + + await documentContainer.SplitAsync(partitionKeyRangeId: 0, cancellationToken: default); + + await documentContainer.SplitAsync(partitionKeyRangeId: 1, cancellationToken: default); + await documentContainer.SplitAsync(partitionKeyRangeId: 2, cancellationToken: default); + + await documentContainer.SplitAsync(partitionKeyRangeId: 3, cancellationToken: default); + await documentContainer.SplitAsync(partitionKeyRangeId: 4, cancellationToken: default); + await documentContainer.SplitAsync(partitionKeyRangeId: 5, cancellationToken: default); + await documentContainer.SplitAsync(partitionKeyRangeId: 6, cancellationToken: default); + + for (int i = 0; i < numItems; i++) + { + // Insert an item + CosmosObject item = CosmosObject.Parse($"{{\"pk\" : {i} }}"); + while (true) + { + TryCatch monadicCreateRecord = await documentContainer.MonadicCreateItemAsync(item, cancellationToken: default); + if (monadicCreateRecord.Succeeded) + { + break; + } + } + } + + return documentContainer; + } + } +} From 9c37663aa22ec24f882510ddde2730b4179639ae Mon Sep 17 00:00:00 2001 From: Brandon Chong Date: Wed, 5 Aug 2020 15:53:29 -0700 Subject: [PATCH 53/85] need to get the parser involved for this one --- ...OrderByCrossPartitionQueryPipelineStage.cs | 53 ++++++++++++++--- ...yQueryPartitionRangePageAsyncEnumerator.cs | 43 ++++++++------ .../Pagination/InMemoryContainer.cs | 4 +- ...ByCrossPartitionQueryPipelineStageTests.cs | 57 ++++++++++++------- 4 files changed, 109 insertions(+), 48 deletions(-) diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderByCrossPartitionQueryPipelineStage.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderByCrossPartitionQueryPipelineStage.cs index 73897e2ee4..7cd15b7bcd 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderByCrossPartitionQueryPipelineStage.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderByCrossPartitionQueryPipelineStage.cs @@ -194,12 +194,15 @@ public async ValueTask MoveNextAsync() return false; } + OrderByQueryPartitionRangePageAsyncEnumerator currentEnumerator = default; + OrderByQueryResult orderByQueryResult = default; + // Try to form a page with as many items in the sorted order without having to do async work. List documents = new List(); while (documents.Count < this.pageSize) { - OrderByQueryPartitionRangePageAsyncEnumerator currentEnumerator = this.enumerators.Dequeue(); - OrderByQueryResult orderByQueryResult = new OrderByQueryResult(currentEnumerator.Current.Result.Enumerator.Current); + currentEnumerator = this.enumerators.Dequeue(); + orderByQueryResult = new OrderByQueryResult(currentEnumerator.Current.Result.Enumerator.Current); documents.Add(orderByQueryResult.Payload); if (!currentEnumerator.Current.Result.Enumerator.MoveNext()) @@ -211,13 +214,37 @@ public async ValueTask MoveNextAsync() // then mark the enumerator as unitialized and it will get requeueed on the next iteration with a fresh page. this.uninitializedEnumeratorsAndTokens.Enqueue((currentEnumerator, (OrderByContinuationToken)null)); } - + break; } this.enumerators.Enqueue(currentEnumerator); } + string continuationTokenString; + if ((this.enumerators.Count == 0) && (this.uninitializedEnumeratorsAndTokens.Count == 0)) + { + continuationTokenString = null; + } + else + { + OrderByContinuationToken orderByContinuationToken = new OrderByContinuationToken( + new CompositeContinuationToken() + { + Token = currentEnumerator.StartOfPageState != null ? ((CosmosString)currentEnumerator.StartOfPageState.Value).Value : null, + Range = currentEnumerator.Range.ToRange(), + }, + orderByQueryResult.OrderByItems, + orderByQueryResult.Rid, + skipCount: 0, + filter: currentEnumerator.Filter); + + CosmosElement cosmosElementOrderByContinuationToken = OrderByContinuationToken.ToCosmosElement(orderByContinuationToken); + CosmosArray continuationTokenList = CosmosArray.Create(new List() { cosmosElementOrderByContinuationToken }); + + continuationTokenString = continuationTokenList.ToString(); + } + // Return a page of results // No stats to report, since we already reported it when we moved to this page. this.Current = TryCatch.FromResult( @@ -228,7 +255,7 @@ public async ValueTask MoveNextAsync() responseLengthInBytes: 0, cosmosQueryExecutionInfo: default, disallowContinuationTokenMessage: default, - state: default /*TODO figure out continuation logic*/)); + state: continuationTokenString != null ? new QueryState(CosmosString.Create(continuationTokenString)) : null)); return true; } @@ -356,7 +383,7 @@ public static TryCatch MonadicCreate( range, pageSize, filter, - state: token != null ? new QueryState(CosmosString.Create(token.CompositeContinuationToken.Token)) : null); + state: token?.CompositeContinuationToken?.Token != null ? new QueryState(CosmosString.Create(token.CompositeContinuationToken.Token)) : null); enumeratorsAndTokens.Add((remoteEnumerator, token)); } @@ -413,13 +440,25 @@ private static TryCatch> MonadicExtractOrderByTok return TryCatch>.FromResult(default); } - if (!(continuationToken is CosmosArray cosmosArray)) + if (!(continuationToken is CosmosString continuationTokenString)) + { + return TryCatch>.FromException( + new MalformedContinuationTokenException( + $"Order by continuation token must be a string: {continuationToken}.")); + } + + string rawJson = continuationTokenString.Value; + + TryCatch monadicCosmosArray = CosmosArray.Monadic.Parse(rawJson); + if (monadicCosmosArray.Failed) { return TryCatch>.FromException( new MalformedContinuationTokenException( - $"Order by continuation token must be an array: {continuationToken}.")); + $"Order by continuation token must be an array: {continuationToken}.", + monadicCosmosArray.Exception)); } + CosmosArray cosmosArray = monadicCosmosArray.Result; if (cosmosArray.Count == 0) { return TryCatch>.FromException( diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderByQueryPartitionRangePageAsyncEnumerator.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderByQueryPartitionRangePageAsyncEnumerator.cs index d50f79b6e0..b646b97797 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderByQueryPartitionRangePageAsyncEnumerator.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderByQueryPartitionRangePageAsyncEnumerator.cs @@ -31,29 +31,36 @@ public OrderByQueryPartitionRangePageAsyncEnumerator( this.sqlQuerySpec = sqlQuerySpec ?? throw new ArgumentNullException(nameof(sqlQuerySpec)); this.pageSize = pageSize; this.Filter = filter; + this.StartOfPageState = state; } public string Filter { get; } + public QueryState StartOfPageState { get; private set; } + public override ValueTask DisposeAsync() => default; - protected override Task> GetNextPageAsync(CancellationToken cancellationToken) => this.queryDataSource - .MonadicQueryAsync( - sqlQuerySpec: this.sqlQuerySpec, - continuationToken: this.State == null ? null : ((CosmosString)this.State.Value).Value, - feedRange: new FeedRangeEpk(this.Range.ToRange()), - pageSize: this.pageSize, - cancellationToken) - .ContinueWith>(antecedent => - { - TryCatch monadicQueryPage = antecedent.Result; - if (monadicQueryPage.Failed) - { - return TryCatch.FromException(monadicQueryPage.Exception); - } - - QueryPage queryPage = monadicQueryPage.Result; - return TryCatch.FromResult(new OrderByQueryPage(queryPage)); - }); + protected override Task> GetNextPageAsync(CancellationToken cancellationToken) + { + this.StartOfPageState = this.State; + return this.queryDataSource + .MonadicQueryAsync( + sqlQuerySpec: this.sqlQuerySpec, + continuationToken: this.State == null ? null : ((CosmosString)this.State.Value).Value, + feedRange: new FeedRangeEpk(this.Range.ToRange()), + pageSize: this.pageSize, + cancellationToken) + .ContinueWith>(antecedent => + { + TryCatch monadicQueryPage = antecedent.Result; + if (monadicQueryPage.Failed) + { + return TryCatch.FromException(monadicQueryPage.Exception); + } + + QueryPage queryPage = monadicQueryPage.Result; + return TryCatch.FromResult(new OrderByQueryPage(queryPage)); + }); + } } } 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 a1a2b6bf07..24dd0bc93d 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 @@ -299,7 +299,7 @@ public async Task> MonadicQueryAsync( } const string selectStar = "SELECT * FROM c"; - const string orderBy = "SELECT r._rid, [{\"item\": c._ts}] AS orderByItems, c AS payload FROM Root AS c ORDER BY c._ts"; + const string orderBy = @"SELECT r._rid, [{""item"": c._ts}] AS orderByItems, c AS payload FROM c WHERE ({documentdb-formattableorderbyquery-filter}) ORDER BY c._ts"; if (!sqlQuerySpec.QueryText.Equals(selectStar, StringComparison.OrdinalIgnoreCase) && !sqlQuerySpec.QueryText.Equals(orderBy, StringComparison.OrdinalIgnoreCase)) @@ -423,7 +423,7 @@ public async Task> MonadicQueryAsync( Dictionary keyValuePairs = new Dictionary { - ["_rid"] = CosmosNumber64.Create(record.ResourceIdentifier), + ["_rid"] = CosmosString.Create(record.ResourceIdentifier.ToString()), ["orderByItems"] = CosmosArray.Create( new List() { diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/OrderByCrossPartitionQueryPipelineStageTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/OrderByCrossPartitionQueryPipelineStageTests.cs index 83796b418f..56f3c6a8fe 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/OrderByCrossPartitionQueryPipelineStageTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/OrderByCrossPartitionQueryPipelineStageTests.cs @@ -126,11 +126,12 @@ public void MonadicCreate_SingleOrderByContinuationToken() new OrderByCrossPartitionQueryPipelineStage.OrderByColumn("_ts", Cosmos.Query.Core.ExecutionContext.OrderBy.SortOrder.Ascending) }, pageSize: 10, - continuationToken: CosmosArray.Create( - new List() - { - OrderByContinuationToken.ToCosmosElement(orderByContinuationToken) - })); + continuationToken: CosmosString.Create( + CosmosArray.Create( + new List() + { + OrderByContinuationToken.ToCosmosElement(orderByContinuationToken) + }).ToString())); Assert.IsTrue(monadicCreate.Succeeded); } @@ -184,12 +185,13 @@ public void MonadicCreate_MultipleOrderByContinuationToken() new OrderByCrossPartitionQueryPipelineStage.OrderByColumn("_ts", Cosmos.Query.Core.ExecutionContext.OrderBy.SortOrder.Ascending) }, pageSize: 10, - continuationToken: CosmosArray.Create( - new List() - { - OrderByContinuationToken.ToCosmosElement(orderByContinuationToken1), - OrderByContinuationToken.ToCosmosElement(orderByContinuationToken2) - })); + continuationToken: CosmosString.Create( + CosmosArray.Create( + new List() + { + OrderByContinuationToken.ToCosmosElement(orderByContinuationToken1), + OrderByContinuationToken.ToCosmosElement(orderByContinuationToken2) + }).ToString())); Assert.IsTrue(monadicCreate.Succeeded); } @@ -201,7 +203,7 @@ public async Task TestDrainFully_StartFromBeginingAsync() TryCatch monadicCreate = OrderByCrossPartitionQueryPipelineStage.MonadicCreate( documentContainer: documentContainer, - sqlQuerySpec: new SqlQuerySpec("SELECT r._rid, [{\"item\": c._ts}] AS orderByItems, c AS payload FROM Root AS c ORDER BY c._ts"), + sqlQuerySpec: new SqlQuerySpec(@"SELECT r._rid, [{""item"": c._ts}] AS orderByItems, c AS payload FROM c WHERE ({documentdb-formattableorderbyquery-filter}) ORDER BY c._ts"), targetRanges: await documentContainer.GetFeedRangesAsync(cancellationToken: default), orderByColumns: new List() { @@ -237,22 +239,35 @@ public async Task TestDrainFully_WithStateResume() QueryState queryState = null; do { - TryCatch monadicCreate = ParallelCrossPartitionQueryPipelineStage.MonadicCreate( + TryCatch monadicCreate = OrderByCrossPartitionQueryPipelineStage.MonadicCreate( documentContainer: documentContainer, - sqlQuerySpec: new SqlQuerySpec("SELECT * FROM c"), + sqlQuerySpec: new SqlQuerySpec(@"SELECT r._rid, [{""item"": c._ts}] AS orderByItems, c AS payload FROM c WHERE ({documentdb-formattableorderbyquery-filter}) ORDER BY c._ts"), + targetRanges: await documentContainer.GetFeedRangesAsync(cancellationToken: default), + orderByColumns: new List() + { + new OrderByCrossPartitionQueryPipelineStage.OrderByColumn("_ts", SortOrder.Ascending) + }, pageSize: 10, continuationToken: queryState?.Value); Assert.IsTrue(monadicCreate.Succeeded); IQueryPipelineStage queryPipelineStage = monadicCreate.Result; - Assert.IsTrue(await queryPipelineStage.MoveNextAsync()); - TryCatch tryGetQueryPage = queryPipelineStage.Current; - Assert.IsTrue(tryGetQueryPage.Succeeded); - - QueryPage queryPage = tryGetQueryPage.Result; - documents.AddRange(queryPage.Documents); + QueryPage queryPage; + do + { + // We need to drain out all the initial empty pages, + // since they are non resumable state. + Assert.IsTrue(await queryPipelineStage.MoveNextAsync()); + TryCatch tryGetQueryPage = queryPipelineStage.Current; + if (tryGetQueryPage.Failed) + { + Assert.Fail(tryGetQueryPage.Exception.ToString()); + } - queryState = queryPage.State; + queryPage = tryGetQueryPage.Result; + documents.AddRange(queryPage.Documents); + queryState = queryPage.State; + } while (queryPage.Documents.Count == 0); } while (queryState != null); Assert.AreEqual(numItems, documents.Count); From 8ec4ce2811ee25650aefd84f523a44b7ee293f20 Mon Sep 17 00:00:00 2001 From: Brandon Chong Date: Wed, 5 Aug 2020 16:46:13 -0700 Subject: [PATCH 54/85] added monadic parser --- .../CosmosQueryExecutionContextFactory.cs | 3 +- .../src/Query/Core/Parser/ErrorListener.cs | 31 +++-- .../src/Query/Core/Parser/ParseException.cs | 16 +++ .../src/Query/Core/Parser/QueryParser.cs | 116 ++++++++++++++++++ .../src/SqlObjects/SqlQuery.cs | 83 ------------- .../Query/ParsingBenchmark.cs | 3 +- 6 files changed, 155 insertions(+), 97 deletions(-) create mode 100644 Microsoft.Azure.Cosmos/src/Query/Core/Parser/ParseException.cs create mode 100644 Microsoft.Azure.Cosmos/src/Query/Core/Parser/QueryParser.cs diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosQueryExecutionContextFactory.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosQueryExecutionContextFactory.cs index b42497b11b..c505ff73fa 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosQueryExecutionContextFactory.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosQueryExecutionContextFactory.cs @@ -15,6 +15,7 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionContext using Microsoft.Azure.Cosmos.Query.Core.Exceptions; using Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.Distinct; using Microsoft.Azure.Cosmos.Query.Core.Monads; + using Microsoft.Azure.Cosmos.Query.Core.Parser; using Microsoft.Azure.Cosmos.Query.Core.QueryClient; using Microsoft.Azure.Cosmos.Query.Core.QueryPlan; using Microsoft.Azure.Cosmos.SqlObjects; @@ -154,7 +155,7 @@ private static async Task> TryCreateCoreCo SqlQuery sqlQuery; using (cosmosQueryContext.CreateDiagnosticScope("QueryParsing")) { - parsed = SqlQuery.TryParse(inputParameters.SqlQuerySpec.QueryText, out sqlQuery); + parsed = QueryParser.TryParse(inputParameters.SqlQuerySpec.QueryText, out sqlQuery); } if (parsed) diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Parser/ErrorListener.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Parser/ErrorListener.cs index ee201c73e5..1d1ab9df9c 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Parser/ErrorListener.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Parser/ErrorListener.cs @@ -13,9 +13,9 @@ namespace Microsoft.Azure.Cosmos.Query.Core.Parser using Antlr4.Runtime; using Antlr4.Runtime.Misc; - internal sealed class ErrorListener : ConsoleErrorListener + internal sealed class ErrorListener : IAntlrErrorListener { - public bool hadError; + public ParseException parseException; private readonly Parser parser; private readonly Lexer lexer; private readonly CommonTokenStream tokenStream; @@ -29,10 +29,15 @@ public ErrorListener(Parser parser, Lexer lexer, CommonTokenStream token_stream) this.firstTime = true; } - public override void SyntaxError(TextWriter output, IRecognizer recognizer, S offendingSymbol, int line, - int col, string msg, RecognitionException e) + public void SyntaxError( + TextWriter output, + IRecognizer recognizer, + S offendingSymbol, + int line, + int col, + string msg, + RecognitionException recognitionException) { - this.hadError = true; if (this.firstTime) { try @@ -41,29 +46,31 @@ public override void SyntaxError(TextWriter output, IRecognizer recognizer, S of LASets la_sets = new LASets(); IntervalSet set = la_sets.Compute(this.parser, this.tokenStream, line, col); List result = new List(); + foreach (int r in set.ToList()) { string rule_name = this.parser.Vocabulary.GetSymbolicName(r); result.Add(rule_name); } + + string message; if (result.Any()) { - System.Console.Error.WriteLine("Parse error line/col " + line + "/" + col - + ", expecting " - + string.Join(", ", result)); + message = $"Parse error line:{line}, col:{col}, expecting: {string.Join(", ", result)}"; + } else { - base.SyntaxError(output, recognizer, offendingSymbol, line, col, msg, e); + message = $"Parse error: message:{msg} offendingSymbol: {offendingSymbol}, line:{line}, col:{col}"; } - return; + this.parseException = new ParseException(message, recognitionException); } - catch (Exception) + catch (Exception ex) { + this.parseException = new ParseException($"Unknown parse exception", ex); } } - base.SyntaxError(output, recognizer, offendingSymbol, line, col, msg, e); } } } \ No newline at end of file diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Parser/ParseException.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Parser/ParseException.cs new file mode 100644 index 0000000000..495b41ce18 --- /dev/null +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Parser/ParseException.cs @@ -0,0 +1,16 @@ +// ------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +// ------------------------------------------------------------ + +namespace Microsoft.Azure.Cosmos.Query.Core.Parser +{ + using System; + + internal sealed class ParseException : Exception + { + public ParseException(string message = null, Exception innerException = null) + : base(message, innerException) + { + } + } +} diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Parser/QueryParser.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Parser/QueryParser.cs new file mode 100644 index 0000000000..ff9804e806 --- /dev/null +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Parser/QueryParser.cs @@ -0,0 +1,116 @@ +// ------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +// ------------------------------------------------------------ + +namespace Microsoft.Azure.Cosmos.Query.Core.Parser +{ + using System; + using System.Runtime.ExceptionServices; + using Antlr4.Runtime; + using Antlr4.Runtime.Misc; + using Microsoft.Azure.Cosmos.Query.Core.Monads; + using Microsoft.Azure.Cosmos.SqlObjects; + + internal static class QueryParser + { + public static class Monadic + { + public static TryCatch Parse(string text) + { + if (text == null) + { + throw new ArgumentNullException(nameof(text)); + } + + AntlrInputStream str = new AntlrInputStream(text); + sqlLexer lexer = new sqlLexer(str); + CommonTokenStream tokens = new CommonTokenStream(lexer); + sqlParser parser = new sqlParser(tokens) + { + ErrorHandler = ThrowExceptionOnErrors.Singleton, + }; + ErrorListener listener = new ErrorListener(parser, lexer, tokens); + parser.AddErrorListener(listener); + + sqlParser.ProgramContext programContext; + try + { + programContext = parser.program(); + } + catch (Exception ex) + { + return TryCatch.FromException(ex); + } + + if (listener.parseException != null) + { + return TryCatch.FromException(listener.parseException); + } + + SqlQuery sqlQuery = (SqlQuery)CstToAstVisitor.Singleton.Visit(programContext); + return TryCatch.FromResult(sqlQuery); + } + + private sealed class ThrowExceptionOnErrors : IAntlrErrorStrategy + { + public static readonly ThrowExceptionOnErrors Singleton = new ThrowExceptionOnErrors(); + + public bool InErrorRecoveryMode(Parser recognizer) + { + return false; + } + + public void Recover(Parser recognizer, RecognitionException e) + { + ExceptionDispatchInfo.Capture(e).Throw(); + } + + [return: NotNull] + public IToken RecoverInline(Parser recognizer) + { + throw new NotSupportedException("can not recover."); + } + + public void ReportError(Parser recognizer, RecognitionException e) + { + ExceptionDispatchInfo.Capture(e).Throw(); + } + + public void ReportMatch(Parser recognizer) + { + // Do nothing + } + + public void Reset(Parser recognizer) + { + // Do nothing + } + + public void Sync(Parser recognizer) + { + // Do nothing + } + } + } + + public static bool TryParse(string text, out SqlQuery sqlQuery) + { + TryCatch monadicParse = Monadic.Parse(text); + if (monadicParse.Failed) + { + sqlQuery = default; + return false; + } + + sqlQuery = monadicParse.Result; + return false; + } + + public static SqlQuery Parse(string text) + { + TryCatch monadicParse = Monadic.Parse(text); + monadicParse.ThrowIfFailed(); + return monadicParse.Result; + } + } +} diff --git a/Microsoft.Azure.Cosmos/src/SqlObjects/SqlQuery.cs b/Microsoft.Azure.Cosmos/src/SqlObjects/SqlQuery.cs index 601949a24e..2cb28f4593 100644 --- a/Microsoft.Azure.Cosmos/src/SqlObjects/SqlQuery.cs +++ b/Microsoft.Azure.Cosmos/src/SqlObjects/SqlQuery.cs @@ -4,10 +4,6 @@ namespace Microsoft.Azure.Cosmos.SqlObjects { using System; - using System.Runtime.ExceptionServices; - using Antlr4.Runtime; - using Antlr4.Runtime.Misc; - using Microsoft.Azure.Cosmos.Query.Core.Parser; using Microsoft.Azure.Cosmos.SqlObjects.Visitors; #if INTERNAL @@ -66,84 +62,5 @@ public static SqlQuery Create( groupByClause, orderByClause, offsetLimitClause); - - public static bool TryParse(string text, out SqlQuery sqlQuery) - { - if (text == null) - { - throw new ArgumentNullException(nameof(text)); - } - - AntlrInputStream str = new AntlrInputStream(text); - sqlLexer lexer = new sqlLexer(str); - CommonTokenStream tokens = new CommonTokenStream(lexer); - sqlParser parser = new sqlParser(tokens) - { - ErrorHandler = ThrowExceptionOnErrors.Singleton, - }; - ErrorListener listener = new ErrorListener(parser, lexer, tokens); - parser.AddErrorListener(listener); - - sqlParser.ProgramContext programContext; - try - { - programContext = parser.program(); - } - catch (Exception) - { - sqlQuery = default; - return false; - } - - if (listener.hadError) - { - sqlQuery = default; - return false; - } - - sqlQuery = (SqlQuery)CstToAstVisitor.Singleton.Visit(programContext); - return true; - } - - private sealed class ThrowExceptionOnErrors : IAntlrErrorStrategy - { - public static readonly ThrowExceptionOnErrors Singleton = new ThrowExceptionOnErrors(); - - public bool InErrorRecoveryMode(Parser recognizer) - { - return false; - } - - public void Recover(Parser recognizer, RecognitionException e) - { - ExceptionDispatchInfo.Capture(e).Throw(); - } - - [return: NotNull] - public IToken RecoverInline(Parser recognizer) - { - throw new NotSupportedException("can not recover."); - } - - public void ReportError(Parser recognizer, RecognitionException e) - { - ExceptionDispatchInfo.Capture(e).Throw(); - } - - public void ReportMatch(Parser recognizer) - { - // Do nothing - } - - public void Reset(Parser recognizer) - { - // Do nothing - } - - public void Sync(Parser recognizer) - { - // Do nothing - } - } } } diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Performance.Tests/Query/ParsingBenchmark.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Performance.Tests/Query/ParsingBenchmark.cs index d7a696835d..775131a7c0 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Performance.Tests/Query/ParsingBenchmark.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Performance.Tests/Query/ParsingBenchmark.cs @@ -10,6 +10,7 @@ namespace Microsoft.Azure.Cosmos.Query using BenchmarkDotNet.Attributes; using Microsoft.Azure.Cosmos.Query.Core; using Microsoft.Azure.Cosmos.Query.Core.Monads; + using Microsoft.Azure.Cosmos.Query.Core.Parser; using Microsoft.Azure.Cosmos.Query.Core.QueryPlan; using Microsoft.Azure.Cosmos.SqlObjects; using Microsoft.Azure.Documents; @@ -108,7 +109,7 @@ public void ParseBenchmark(QueryLength queryLength, ParserType parserType) private static void ParseUsingMangedParser(SqlQuerySpec sqlQuerySpec) { - if (!SqlQuery.TryParse(sqlQuerySpec.QueryText, out SqlQuery sqlQuery)) + if (!QueryParser.TryParse(sqlQuerySpec.QueryText, out SqlQuery sqlQuery)) { throw new InvalidOperationException("FAILED TO PARSE QUERY."); } From 18dc83c9c9ead87f7bc98d44becc4279311491c0 Mon Sep 17 00:00:00 2001 From: Brandon Chong Date: Wed, 5 Aug 2020 19:42:05 -0700 Subject: [PATCH 55/85] whoops --- Microsoft.Azure.Cosmos/src/Query/Core/Parser/QueryParser.cs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Parser/QueryParser.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Parser/QueryParser.cs index ff9804e806..05c0ee0ee4 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Parser/QueryParser.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Parser/QueryParser.cs @@ -103,7 +103,7 @@ public static bool TryParse(string text, out SqlQuery sqlQuery) } sqlQuery = monadicParse.Result; - return false; + return true; } public static SqlQuery Parse(string text) From eaeab623bd7c58f16668616924cd2b7f1c30eb7b Mon Sep 17 00:00:00 2001 From: Brandon Chong Date: Thu, 6 Aug 2020 10:51:12 -0700 Subject: [PATCH 56/85] hooked up parser and offline engine to the inmemorycontainer --- .../src/Pagination/DocumentContainer.cs | 4 +- .../src/Pagination/DocumentContainerState.cs | 6 +- .../src/Pagination/IDocumentContainer.cs | 5 +- .../Pagination/IMonadicDocumentContainer.cs | 2 +- .../NetworkAttachedDocumentContainer.cs | 4 +- .../src/Pagination/Record.cs | 19 +- ...OrderByCrossPartitionQueryPipelineStage.cs | 4 +- ...cumentContainerPartitionRangeEnumerator.cs | 2 +- .../Pagination/DocumentContainerTests.cs | 4 +- .../Pagination/FlakyDocumentContainer.cs | 2 +- .../Pagination/InMemoryContainer.cs | 206 ++++++------------ ...ByCrossPartitionQueryPipelineStageTests.cs | 21 +- .../QueryPartitionRangePageEnumeratorTests.cs | 2 +- 13 files changed, 118 insertions(+), 163 deletions(-) diff --git a/Microsoft.Azure.Cosmos/src/Pagination/DocumentContainer.cs b/Microsoft.Azure.Cosmos/src/Pagination/DocumentContainer.cs index c9d4d2e385..fc3ee75fe2 100644 --- a/Microsoft.Azure.Cosmos/src/Pagination/DocumentContainer.cs +++ b/Microsoft.Azure.Cosmos/src/Pagination/DocumentContainer.cs @@ -84,7 +84,7 @@ public Task ReadItemAsync( public Task> MonadicReadFeedAsync( int partitionKeyRangeId, - long resourceIdentifer, + ResourceId resourceIdentifer, int pageSize, CancellationToken cancellationToken) => this.monadicDocumentContainer.MonadicReadFeedAsync( partitionKeyRangeId, @@ -94,7 +94,7 @@ public Task> MonadicReadFeedAsync( public Task ReadFeedAsync( int partitionKeyRangeId, - long resourceIdentifier, + ResourceId resourceIdentifier, int pageSize, CancellationToken cancellationToken) => TryCatch.UnsafeGetResultAsync( this.MonadicReadFeedAsync( diff --git a/Microsoft.Azure.Cosmos/src/Pagination/DocumentContainerState.cs b/Microsoft.Azure.Cosmos/src/Pagination/DocumentContainerState.cs index cbc29a3542..dd187d638d 100644 --- a/Microsoft.Azure.Cosmos/src/Pagination/DocumentContainerState.cs +++ b/Microsoft.Azure.Cosmos/src/Pagination/DocumentContainerState.cs @@ -4,13 +4,15 @@ namespace Microsoft.Azure.Cosmos.Pagination { + using Microsoft.Azure.Documents; + internal sealed class DocumentContainerState : State { - public DocumentContainerState(long resourceIdentifier) + public DocumentContainerState(ResourceId resourceIdentifier) { this.ResourceIdentifer = resourceIdentifier; } - public long ResourceIdentifer { get; } + public ResourceId ResourceIdentifer { get; } } } diff --git a/Microsoft.Azure.Cosmos/src/Pagination/IDocumentContainer.cs b/Microsoft.Azure.Cosmos/src/Pagination/IDocumentContainer.cs index 2dd017dc63..42d223beb3 100644 --- a/Microsoft.Azure.Cosmos/src/Pagination/IDocumentContainer.cs +++ b/Microsoft.Azure.Cosmos/src/Pagination/IDocumentContainer.cs @@ -4,12 +4,11 @@ namespace Microsoft.Azure.Cosmos.Pagination { - using System; using System.Threading; using System.Threading.Tasks; using Microsoft.Azure.Cosmos.CosmosElements; - using Microsoft.Azure.Cosmos.Query.Core.Monads; using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote; + using Microsoft.Azure.Documents; internal interface IDocumentContainer : IMonadicDocumentContainer, IFeedRangeProvider, IQueryDataSource { @@ -24,7 +23,7 @@ Task ReadItemAsync( Task ReadFeedAsync( int partitionKeyRangeId, - long resourceIdentifier, + ResourceId resourceIdentifier, int pageSize, CancellationToken cancellationToken); diff --git a/Microsoft.Azure.Cosmos/src/Pagination/IMonadicDocumentContainer.cs b/Microsoft.Azure.Cosmos/src/Pagination/IMonadicDocumentContainer.cs index 8695784b0e..fc63f3760f 100644 --- a/Microsoft.Azure.Cosmos/src/Pagination/IMonadicDocumentContainer.cs +++ b/Microsoft.Azure.Cosmos/src/Pagination/IMonadicDocumentContainer.cs @@ -28,7 +28,7 @@ Task> MonadicReadItemAsync( Task> MonadicReadFeedAsync( int partitionKeyRangeId, - long resourceIdentifer, + ResourceId resourceIdentifer, int pageSize, CancellationToken cancellationToken); diff --git a/Microsoft.Azure.Cosmos/src/Pagination/NetworkAttachedDocumentContainer.cs b/Microsoft.Azure.Cosmos/src/Pagination/NetworkAttachedDocumentContainer.cs index 9d91f87443..1428941d60 100644 --- a/Microsoft.Azure.Cosmos/src/Pagination/NetworkAttachedDocumentContainer.cs +++ b/Microsoft.Azure.Cosmos/src/Pagination/NetworkAttachedDocumentContainer.cs @@ -63,7 +63,7 @@ public async Task> MonadicCreateItemAsync( CosmosObject insertedDocument = tryInsertDocument.Resource; string identifier = ((CosmosString)insertedDocument["id"]).Value; - long resourceIdentifier = BitConverter.ToInt64(ResourceId.Parse(((CosmosString)insertedDocument["_rid"]).Value).Value, startIndex: 0); + ResourceId resourceIdentifier = ResourceId.Parse(((CosmosString)insertedDocument["_rid"]).Value); long timestamp = Number64.ToLong(((CosmosNumber)insertedDocument["_ts"]).Value); Record record = new Record(resourceIdentifier, timestamp, identifier, insertedDocument); @@ -92,7 +92,7 @@ public Task>> MonadicGetChildRangeAsync( public Task> MonadicReadFeedAsync( int partitionKeyRangeId, - long resourceIdentifer, + ResourceId resourceIdentifer, int pageSize, CancellationToken cancellationToken) { diff --git a/Microsoft.Azure.Cosmos/src/Pagination/Record.cs b/Microsoft.Azure.Cosmos/src/Pagination/Record.cs index 980a206cb1..5496256e3b 100644 --- a/Microsoft.Azure.Cosmos/src/Pagination/Record.cs +++ b/Microsoft.Azure.Cosmos/src/Pagination/Record.cs @@ -5,23 +5,25 @@ namespace Microsoft.Azure.Cosmos.Pagination { using System; + using System.Reflection; using Microsoft.Azure.Cosmos.CosmosElements; + using Microsoft.Azure.Documents; internal sealed class Record { public Record( - long resourceIdentifier, + ResourceId resourceIdentifier, long timestamp, string identifier, CosmosObject payload) { - this.ResourceIdentifier = resourceIdentifier < 0 ? throw new ArgumentOutOfRangeException(nameof(resourceIdentifier)) : resourceIdentifier; + this.ResourceIdentifier = resourceIdentifier; this.Timestamp = timestamp < 0 ? throw new ArgumentOutOfRangeException(nameof(timestamp)) : timestamp; this.Identifier = identifier ?? throw new ArgumentNullException(nameof(identifier)); this.Payload = payload ?? throw new ArgumentNullException(nameof(payload)); } - public long ResourceIdentifier { get; } + public ResourceId ResourceIdentifier { get; } public long Timestamp { get; } @@ -29,10 +31,17 @@ public Record( public CosmosObject Payload { get; } - public static Record Create(long previousResourceIdentifier, CosmosObject payload) + public static Record Create(ResourceId previousResourceIdentifier, CosmosObject payload) { + const string dummyRidString = "AYIMAMmFOw8YAAAAAAAAAA=="; + ResourceId resourceId = ResourceId.Parse(dummyRidString); + PropertyInfo prop = resourceId + .GetType() + .GetProperty("Document", BindingFlags.NonPublic | BindingFlags.Public | BindingFlags.Instance); + prop.SetValue(resourceId, previousResourceIdentifier.Document + 1); + return new Record( - previousResourceIdentifier + 1, + resourceId, DateTime.UtcNow.Ticks, Guid.NewGuid().ToString(), payload); diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderByCrossPartitionQueryPipelineStage.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderByCrossPartitionQueryPipelineStage.cs index 7cd15b7bcd..9cb9893461 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderByCrossPartitionQueryPipelineStage.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderByCrossPartitionQueryPipelineStage.cs @@ -317,7 +317,7 @@ public static TryCatch MonadicCreate( enumeratorsAndTokens = targetRanges .Select(range => (new OrderByQueryPartitionRangePageAsyncEnumerator( documentContainer, - sqlQuerySpec, + rewrittenQueryForOrderBy, range, pageSize, TrueFilter, @@ -379,7 +379,7 @@ public static TryCatch MonadicCreate( OrderByContinuationToken token = kvp.Value; OrderByQueryPartitionRangePageAsyncEnumerator remoteEnumerator = new OrderByQueryPartitionRangePageAsyncEnumerator( documentContainer, - sqlQuerySpec, + rewrittenQueryForOrderBy, range, pageSize, filter, diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/DocumentContainerPartitionRangeEnumerator.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/DocumentContainerPartitionRangeEnumerator.cs index 09a0d18a85..41910827e0 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/DocumentContainerPartitionRangeEnumerator.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/DocumentContainerPartitionRangeEnumerator.cs @@ -30,7 +30,7 @@ public DocumentContainerPartitionRangeEnumerator( MinInclusive = partitionKeyRangeId.ToString(), MaxExclusive = partitionKeyRangeId.ToString() }, - state ?? new DocumentContainerState(resourceIdentifier: 0)) + state ?? new DocumentContainerState(resourceIdentifier: ResourceId.Empty)) { this.documentContainer = documentContainer ?? throw new ArgumentNullException(nameof(documentContainer)); this.partitionKeyRangeId = partitionKeyRangeId; diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/DocumentContainerTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/DocumentContainerTests.cs index 50c0dbca3a..a2b0d23994 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/DocumentContainerTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/DocumentContainerTests.cs @@ -205,7 +205,7 @@ async Task AssertChildPartitionAsync(int partitionKeyRangeId) { DocumentContainerPage readFeedPage = await documentContainer.ReadFeedAsync( partitionKeyRangeId: partitionKeyRangeId, - resourceIdentifier: 0, + resourceIdentifier: ResourceId.Empty, pageSize: 100, cancellationToken: default); @@ -297,7 +297,7 @@ async Task AssertChildPartitionAsync(int partitionKeyRangeId) { DocumentContainerPage page = await documentContainer.ReadFeedAsync( partitionKeyRangeId: partitionKeyRangeId, - resourceIdentifier: 0, + resourceIdentifier: ResourceId.Empty, pageSize: 100, cancellationToken: default); diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/FlakyDocumentContainer.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/FlakyDocumentContainer.cs index 891b1a1bc6..c1d31583dd 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/FlakyDocumentContainer.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/FlakyDocumentContainer.cs @@ -89,7 +89,7 @@ public Task> MonadicReadItemAsync( public Task> MonadicReadFeedAsync( int partitionKeyRangeId, - long resourceIdentifer, + ResourceId resourceIdentifer, int pageSize, CancellationToken cancellationToken) { 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 24dd0bc93d..9d38bc4663 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 @@ -16,9 +16,12 @@ namespace Microsoft.Azure.Cosmos.Tests.Pagination using Microsoft.Azure.Cosmos.Pagination; using Microsoft.Azure.Cosmos.Query.Core; using Microsoft.Azure.Cosmos.Query.Core.Monads; + using Microsoft.Azure.Cosmos.Query.Core.Parser; using Microsoft.Azure.Cosmos.Query.Core.Pipeline; using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote; using Microsoft.Azure.Cosmos.Routing; + using Microsoft.Azure.Cosmos.SqlObjects; + using Microsoft.Azure.Cosmos.Tests.Query.OfflineEngine; using Microsoft.Azure.Documents; // Collection useful for mocking requests and repartitioning (splits / merge). @@ -230,7 +233,7 @@ static Task> CreateNotFoundException(CosmosElement partitionKey public Task> MonadicReadFeedAsync( int partitionKeyRangeId, - long resourceIdentifer, + ResourceId resourceIdentifer, int pageSize, CancellationToken cancellationToken) { @@ -256,7 +259,7 @@ public Task> MonadicReadFeedAsync( } List page = records - .Where(record => record.ResourceIdentifier > resourceIdentifer) + .Where(record => record.ResourceIdentifier.Document > resourceIdentifer.Document) .Take(pageSize) .ToList(); @@ -286,7 +289,7 @@ public Task> MonadicQueryAsync( throw new NotImplementedException(); } - public async Task> MonadicQueryAsync( + public Task> MonadicQueryAsync( SqlQuerySpec sqlQuerySpec, string continuationToken, int partitionKeyRangeId, @@ -298,158 +301,89 @@ public async Task> MonadicQueryAsync( throw new ArgumentNullException(nameof(sqlQuerySpec)); } - const string selectStar = "SELECT * FROM c"; - const string orderBy = @"SELECT r._rid, [{""item"": c._ts}] AS orderByItems, c AS payload FROM c WHERE ({documentdb-formattableorderbyquery-filter}) ORDER BY c._ts"; - - if (!sqlQuerySpec.QueryText.Equals(selectStar, StringComparison.OrdinalIgnoreCase) - && !sqlQuerySpec.QueryText.Equals(orderBy, StringComparison.OrdinalIgnoreCase)) + TryCatch monadicParse = QueryParser.Monadic.Parse(sqlQuerySpec.QueryText); + if (monadicParse.Failed) { - throw new NotSupportedException("InMemoryCollection only supports SELECT * FROM c and order by _ts queries"); + return Task.FromResult(TryCatch.FromException(monadicParse.Exception)); } - if (sqlQuerySpec.QueryText.Equals(selectStar, StringComparison.OrdinalIgnoreCase)) - { - // For now just do a read feed - long resourceIdentifier = continuationToken != null ? long.Parse(continuationToken) : 0; - TryCatch tryGetPage = await this.MonadicReadFeedAsync( - partitionKeyRangeId, - resourceIdentifier, - pageSize, - cancellationToken); - - if (tryGetPage.Failed) - { - return TryCatch.FromException(tryGetPage.Exception); - } - - DocumentContainerPage page = tryGetPage.Result; - List documents = new List(page.Records.Count); - foreach (Record record in page.Records) - { - Dictionary keyValuePairs = new Dictionary - { - ["_rid"] = CosmosNumber64.Create(record.ResourceIdentifier), - ["_ts"] = CosmosNumber64.Create(record.Timestamp), - ["id"] = CosmosString.Create(record.Identifier) - }; - - foreach (KeyValuePair property in record.Payload) - { - keyValuePairs[property.Key] = property.Value; - } - - documents.Add(CosmosObject.Create(keyValuePairs)); - } - - QueryPage queryPage = new QueryPage( - documents: documents, - requestCharge: 42, - activityId: Guid.NewGuid().ToString(), - responseLengthInBytes: 1337, - cosmosQueryExecutionInfo: default, - disallowContinuationTokenMessage: default, - state: page.State != null ? new QueryState(CosmosString.Create(page.State.ResourceIdentifer.ToString())) : null); + SqlQuery sqlQuery = monadicParse.Result; - return TryCatch.FromResult(queryPage); - } - else + if (!this.partitionKeyRangeIdToHashRange.TryGetValue( + partitionKeyRangeId, + out PartitionKeyHashRange range)) { - // We need to simulate an order by - (long timestamp, long resourceIdentifer) timestampAndIdentifier; - if (continuationToken == null) - { - timestampAndIdentifier = (-1, -1); - } - else - { - string[] tokens = continuationToken.Split(","); - timestampAndIdentifier = (long.Parse(tokens[0]), long.Parse(tokens[1])); - } - - if (!this.partitionKeyRangeIdToHashRange.TryGetValue( - partitionKeyRangeId, - out PartitionKeyHashRange range)) - { - return TryCatch.FromException( - new CosmosException( + return Task.FromResult(TryCatch.FromException( + new CosmosException( message: $"PartitionKeyRangeId {partitionKeyRangeId} is gone", statusCode: System.Net.HttpStatusCode.Gone, subStatusCode: (int)SubStatusCodes.PartitionKeyRangeGone, activityId: Guid.NewGuid().ToString(), - requestCharge: 42)); - } + requestCharge: 42))); + } - if (!this.partitionedRecords.TryGetValue(range, out Records records)) - { - throw new InvalidOperationException("failed to find the range."); - } + if (!this.partitionedRecords.TryGetValue(range, out Records records)) + { + throw new InvalidOperationException("failed to find the range."); + } - List results = records - .OrderBy(record => record.Timestamp) - .ThenBy(record => record.ResourceIdentifier) - .Where(record => - (record.Timestamp > timestampAndIdentifier.timestamp) - || ((record.Timestamp == timestampAndIdentifier.timestamp) && (record.ResourceIdentifier > timestampAndIdentifier.resourceIdentifer))) - .Take(pageSize) - .ToList(); + List documents = new List(); + foreach (Record record in records) + { + Dictionary keyValuePairs = new Dictionary + { + ["_rid"] = CosmosString.Create(record.ResourceIdentifier.ToString()), + ["_ts"] = CosmosNumber64.Create(record.Timestamp), + ["id"] = CosmosString.Create(record.Identifier) + }; - if (results.Count == 0) + foreach (KeyValuePair property in record.Payload) { - return TryCatch.FromResult( - new QueryPage( - documents: new List(), - requestCharge: 42, - activityId: Guid.NewGuid().ToString(), - responseLengthInBytes: 1337, - cosmosQueryExecutionInfo: default, - disallowContinuationTokenMessage: default, - state: null)); + keyValuePairs[property.Key] = property.Value; } - List documents = new List(results.Count); - foreach (Record record in results) - { - Dictionary payloadProperties = new Dictionary() - { - ["_rid"] = CosmosNumber64.Create(record.ResourceIdentifier), - ["_ts"] = CosmosNumber64.Create(record.Timestamp), - ["id"] = CosmosString.Create(record.Identifier) - }; + documents.Add(CosmosObject.Create(keyValuePairs)); + } - foreach (KeyValuePair property in record.Payload) - { - payloadProperties[property.Key] = property.Value; - } + IEnumerable queryResults = SqlInterpreter.ExecuteQuery(documents, sqlQuery); - Dictionary keyValuePairs = new Dictionary - { - ["_rid"] = CosmosString.Create(record.ResourceIdentifier.ToString()), - ["orderByItems"] = CosmosArray.Create( - new List() - { - CosmosObject.Create( - new Dictionary() - { - ["item"] = CosmosNumber64.Create(record.Timestamp) - }) - }), - ["payload"] = CosmosObject.Create(payloadProperties), - }; - - documents.Add(CosmosObject.Create(keyValuePairs)); + int index; + if (continuationToken != null) + { + if (!int.TryParse(continuationToken, out index)) + { + return Task.FromResult( + TryCatch.FromException( + new Exception("FAILED TO PARSE CONTINUATION TOKEN"))); } + } + else + { + index = 0; + } - QueryPage queryPage = new QueryPage( - documents: documents, - requestCharge: 42, - activityId: Guid.NewGuid().ToString(), - responseLengthInBytes: 1337, - cosmosQueryExecutionInfo: default, - disallowContinuationTokenMessage: default, - state: new QueryState(CosmosString.Create($"{results.Last().Timestamp},{results.Last().ResourceIdentifier}"))); + List queryPageResults = queryResults.Skip(index).Take(pageSize).ToList(); - return TryCatch.FromResult(queryPage); + QueryState queryState; + if (queryPageResults.Count == 0) + { + queryState = null; + } + else + { + queryState = new QueryState(CosmosString.Create((index + pageSize).ToString())); } + + return Task.FromResult( + TryCatch.FromResult( + new QueryPage( + queryPageResults, + requestCharge: 42, + activityId: Guid.NewGuid().ToString(), + responseLengthInBytes: 1337, + cosmosQueryExecutionInfo: default, + disallowContinuationTokenMessage: default, + state: queryState))); } public Task> MonadicQueryAsync( @@ -644,10 +578,10 @@ public Records() public Record Add(CosmosObject payload) { - long previousResourceId; + ResourceId previousResourceId; if (this.Count == 0) { - previousResourceId = 0; + previousResourceId = ResourceId.Empty; } else { diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/OrderByCrossPartitionQueryPipelineStageTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/OrderByCrossPartitionQueryPipelineStageTests.cs index 56f3c6a8fe..b96f2907e7 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/OrderByCrossPartitionQueryPipelineStageTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/OrderByCrossPartitionQueryPipelineStageTests.cs @@ -203,11 +203,15 @@ public async Task TestDrainFully_StartFromBeginingAsync() TryCatch monadicCreate = OrderByCrossPartitionQueryPipelineStage.MonadicCreate( documentContainer: documentContainer, - sqlQuerySpec: new SqlQuerySpec(@"SELECT r._rid, [{""item"": c._ts}] AS orderByItems, c AS payload FROM c WHERE ({documentdb-formattableorderbyquery-filter}) ORDER BY c._ts"), + sqlQuerySpec: new SqlQuerySpec(@" + SELECT c._rid AS _rid, [{""item"": c._ts}] AS orderByItems, c AS payload + FROM c + WHERE {documentdb-formattableorderbyquery-filter} + ORDER BY c._ts"), targetRanges: await documentContainer.GetFeedRangesAsync(cancellationToken: default), orderByColumns: new List() { - new OrderByCrossPartitionQueryPipelineStage.OrderByColumn("_ts", SortOrder.Ascending) + new OrderByCrossPartitionQueryPipelineStage.OrderByColumn("c._ts", SortOrder.Ascending) }, pageSize: 10, continuationToken: null); @@ -218,7 +222,10 @@ public async Task TestDrainFully_StartFromBeginingAsync() while (await queryPipelineStage.MoveNextAsync()) { TryCatch tryGetQueryPage = queryPipelineStage.Current; - Assert.IsTrue(tryGetQueryPage.Succeeded); + if (tryGetQueryPage.Failed) + { + Assert.Fail(tryGetQueryPage.Exception.ToString()); + } QueryPage queryPage = tryGetQueryPage.Result; documents.AddRange(queryPage.Documents); @@ -241,11 +248,15 @@ public async Task TestDrainFully_WithStateResume() { TryCatch monadicCreate = OrderByCrossPartitionQueryPipelineStage.MonadicCreate( documentContainer: documentContainer, - sqlQuerySpec: new SqlQuerySpec(@"SELECT r._rid, [{""item"": c._ts}] AS orderByItems, c AS payload FROM c WHERE ({documentdb-formattableorderbyquery-filter}) ORDER BY c._ts"), + sqlQuerySpec: new SqlQuerySpec(@" + SELECT c._rid AS _rid, [{""item"": c._ts}] AS orderByItems, c AS payload + FROM c + WHERE {documentdb-formattableorderbyquery-filter} + ORDER BY c._ts"), targetRanges: await documentContainer.GetFeedRangesAsync(cancellationToken: default), orderByColumns: new List() { - new OrderByCrossPartitionQueryPipelineStage.OrderByColumn("_ts", SortOrder.Ascending) + new OrderByCrossPartitionQueryPipelineStage.OrderByColumn("c._ts", SortOrder.Ascending) }, pageSize: 10, continuationToken: queryState?.Value); 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 fcc32ad8c8..c99dcd5e0e 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 @@ -110,7 +110,7 @@ public override IReadOnlyList GetRecordsFromPage(QueryPage page) foreach (CosmosElement element in page.Documents) { CosmosObject document = (CosmosObject)element; - long resourceIdentifier = Number64.ToLong(((CosmosNumber)document["_rid"]).Value); + ResourceId resourceIdentifier = ResourceId.Parse(((CosmosString)document["_rid"]).Value); long timestamp = Number64.ToLong(((CosmosNumber)document["_ts"]).Value); string identifer = ((CosmosString)document["id"]).Value; From 555c9d69bbb3c5e4f089c173743e58c345faf327 Mon Sep 17 00:00:00 2001 From: Brandon Chong Date: Fri, 7 Aug 2020 17:13:02 -0700 Subject: [PATCH 57/85] got continuation token support working --- ...OrderByCrossPartitionQueryPipelineStage.cs | 48 +++++++++++++++---- ...ByCrossPartitionQueryPipelineStageTests.cs | 4 +- 2 files changed, 41 insertions(+), 11 deletions(-) diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderByCrossPartitionQueryPipelineStage.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderByCrossPartitionQueryPipelineStage.cs index 9cb9893461..0c7b2ca31e 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderByCrossPartitionQueryPipelineStage.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderByCrossPartitionQueryPipelineStage.cs @@ -44,7 +44,7 @@ internal sealed class OrderByCrossPartitionQueryPipelineStage : IQueryPipelineSt /// private const string TrueFilter = "true"; - private static readonly QueryState UninitializedQueryState = new QueryState(CosmosString.Create("ORDER BY NOT INITIALIZED YET!")); + private static readonly QueryState InitializingQueryState = new QueryState(CosmosString.Create("ORDER BY NOT INITIALIZED YET!")); private static readonly IReadOnlyList EmptyPage = new List(); private readonly IDocumentContainer documentContainer; @@ -79,7 +79,7 @@ private OrderByCrossPartitionQueryPipelineStage( this.enumerators = new PriorityQueue(new OrderByEnumeratorComparer(this.sortOrders)); this.pageSize = pageSize < 0 ? throw new ArgumentOutOfRangeException($"{nameof(pageSize)} must be a non negative number.") : pageSize; this.uninitializedEnumeratorsAndTokens = new Queue<(OrderByQueryPartitionRangePageAsyncEnumerator, OrderByContinuationToken)>(uninitializedEnumeratorsAndTokens ?? throw new ArgumentNullException(nameof(uninitializedEnumeratorsAndTokens))); - this.state = state; + this.state = state ?? InitializingQueryState; } public TryCatch Current { get; private set; } @@ -88,7 +88,6 @@ private OrderByCrossPartitionQueryPipelineStage( public async ValueTask MoveNextAsync() { - this.state = null; if (this.uninitializedEnumeratorsAndTokens.Count != 0) { (OrderByQueryPartitionRangePageAsyncEnumerator uninitializedEnumerator, OrderByContinuationToken token) = this.uninitializedEnumeratorsAndTokens.Dequeue(); @@ -121,7 +120,25 @@ public async ValueTask MoveNextAsync() if (!uninitializedEnumerator.Current.Result.Enumerator.MoveNext()) { // Page was empty - this.uninitializedEnumeratorsAndTokens.Enqueue((uninitializedEnumerator, token)); + if (uninitializedEnumerator.State != null) + { + this.uninitializedEnumeratorsAndTokens.Enqueue((uninitializedEnumerator, token)); + } + + if ((this.uninitializedEnumeratorsAndTokens.Count == 0) && (this.enumerators.Count == 0)) + { + // Query did not match any results. We need to emit a fake empty page with null continuation + this.Current = TryCatch.FromResult( + new QueryPage( + documents: EmptyPage, + requestCharge: 0, + activityId: Guid.NewGuid().ToString(), + responseLengthInBytes: 0, + cosmosQueryExecutionInfo: default, + disallowContinuationTokenMessage: default, + state: null)); + return true; + } } else { @@ -162,11 +179,25 @@ public async ValueTask MoveNextAsync() (bool doneFiltering, TryCatch monadicQueryByPage) = filterMonad.Result; if (doneFiltering) { - // TODO test empty pages if (uninitializedEnumerator.Current.Result.Enumerator.Current != null) { this.enumerators.Enqueue(uninitializedEnumerator); } + else if ((this.uninitializedEnumeratorsAndTokens.Count == 0) && (this.enumerators.Count == 0)) + { + // Query did not match any results. + // We need to emit a fake empty page with null continuation + this.Current = TryCatch.FromResult( + new QueryPage( + documents: EmptyPage, + requestCharge: 0, + activityId: Guid.NewGuid().ToString(), + responseLengthInBytes: 0, + cosmosQueryExecutionInfo: default, + disallowContinuationTokenMessage: default, + state: null)); + return true; + } } else { @@ -191,6 +222,7 @@ public async ValueTask MoveNextAsync() if (this.enumerators.Count == 0) { + // Finished draining. return false; } @@ -721,7 +753,7 @@ private static (string leftFilter, string targetFilter, string rightFilter) GetF if (cmp == 0) { // Once the item matches the order by items from the continuation tokens - // We still need to remove all the documents that have a lower rid in the rid sort order. + // We still need to remove all the documents that have a lower or same rid in the rid sort order. // If there is a tie in the sort order the documents should be in _rid order in the same direction as the index (given by the backend) cmp = continuationRid.Document.CompareTo(rid.Document); if ((orderByQueryPage.Page.CosmosQueryExecutionInfo == null) || orderByQueryPage.Page.CosmosQueryExecutionInfo.ReverseRidEnabled) @@ -741,9 +773,7 @@ private static (string leftFilter, string targetFilter, string rightFilter) GetF } } - // We might have passed the item due to deletions and filters. - // We also have a skip count for JOINs - if (cmp < 0 || (cmp == 0 && itemToSkip-- <= 0)) + if (cmp < 0) { return TryCatch<(bool, TryCatch)>.FromResult((true, enumerator.Current)); } diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/OrderByCrossPartitionQueryPipelineStageTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/OrderByCrossPartitionQueryPipelineStageTests.cs index b96f2907e7..36f850c0d4 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/OrderByCrossPartitionQueryPipelineStageTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/OrderByCrossPartitionQueryPipelineStageTests.cs @@ -238,7 +238,7 @@ FROM c [TestMethod] public async Task TestDrainFully_WithStateResume() { - int numItems = 1000; + int numItems = 100; IDocumentContainer documentContainer = await CreateDocumentContainerAsync(numItems); List documents = new List(); @@ -278,7 +278,7 @@ FROM c queryPage = tryGetQueryPage.Result; documents.AddRange(queryPage.Documents); queryState = queryPage.State; - } while (queryPage.Documents.Count == 0); + } while ((queryPage.Documents.Count == 0) && (queryState != null)); } while (queryState != null); Assert.AreEqual(numItems, documents.Count); From c63264d13c041f36c394040e2b9e7749f2bbf649 Mon Sep 17 00:00:00 2001 From: Brandon Chong Date: Fri, 7 Aug 2020 17:54:21 -0700 Subject: [PATCH 58/85] moved around some files --- .../OrderByContinuationToken.cs | 2 +- ...ggregateDocumentQueryExecutionComponent.cs | 2 +- ...DistinctDocumentQueryExecutionComponent.cs | 2 +- .../GroupByDocumentQueryExecutionComponent.cs | 1 + .../SkipDocumentQueryExecutionComponent.cs | 1 + .../TakeDocumentQueryExecutionComponent.cs | 1 + ...smosCrossPartitionQueryExecutionContext.cs | 682 ---------------- .../CosmosQueryExecutionContext.cs | 37 - .../ItemProducers/ItemProducer.cs | 439 ----------- .../ItemProducers/ItemProducerTree.cs | 742 ------------------ ...QueryExecutionContext.ContinuationToken.cs | 149 ---- ...sOrderByItemQueryExecutionContext.Drain.cs | 144 ---- ...OrderByItemQueryExecutionContext.Resume.cs | 69 -- .../CosmosOrderByItemQueryExecutionContext.cs | 71 -- .../OrderByItemProducerTreeComparer.cs | 142 ---- .../OrderBy/RangeFilterInitializationInfo.cs | 91 --- ...QueryExecutionContext.ContinuationToken.cs | 111 --- ...ParallelItemQueryExecutionContext.Drain.cs | 71 -- ...arallelItemQueryExecutionContext.Resume.cs | 180 ----- ...CosmosParallelItemQueryExecutionContext.cs | 58 -- ...inisticParallelItemProducerTreeComparer.cs | 62 -- ...inisticParallelItemProducerTreeComparer.cs | 65 -- .../Parallel/ParallelQueryConfig.cs | 77 -- .../QueryExecutionContextWithException.cs | 46 -- .../Aggregate/Aggregators/MinMaxAggregator.cs | 2 +- .../Pipeline/CatchAllQueryPipelineStage.cs | 1 - .../CosmosQueryExecutionContextFactory.cs | 51 +- .../ExecutionEnvironment.cs | 2 +- .../Query/Core/Pipeline/PipelineFactory.cs | 3 +- .../OrderBy/CosmosElementToQueryLiteral.cs | 2 +- .../Remote/OrderBy}/ItemComparer.cs | 4 +- ...OrderByCrossPartitionQueryPipelineStage.cs | 4 +- .../OrderByEnumeratorComparer.cs | 6 +- .../Remote}/OrderBy/OrderByItem.cs | 3 +- .../Remote/{ => OrderBy}/OrderByQueryPage.cs | 2 +- ...yQueryPartitionRangePageAsyncEnumerator.cs | 2 +- .../Remote}/OrderBy/OrderByQueryResult.cs | 3 +- .../Remote}/OrderBy/SortOrder.cs | 2 +- ...arallelCrossPartitionQueryPipelineStage.cs | 2 +- .../QueryPartitionRangePageAsyncEnumerator.cs | 2 +- .../Core/QueryClient/CosmosQueryClient.cs | 1 - .../Core/QueryClient/CosmosQueryContext.cs | 1 - .../src/Query/Core/QueryPlan/QueryInfo.cs | 2 +- .../src/RequestOptions/QueryRequestOptions.cs | 1 + .../Query/OrderByQueryTests.cs | 4 +- .../Query/QueryTestsBase.cs | 3 +- .../QueryTests.cs | 73 +- .../OrderByContinuationTokenTests.cs | 2 +- .../Query/Pipeline/FactoryTests.cs | 1 + ...ByCrossPartitionQueryPipelineStageTests.cs | 15 +- ...elCrossPartitionQueryPipelineStageTests.cs | 1 + .../QueryPartitionRangePageEnumeratorTests.cs | 1 + .../Query/QueryPlanBaselineTests.cs | 3 +- 53 files changed, 68 insertions(+), 3376 deletions(-) delete mode 100644 Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosCrossPartitionQueryExecutionContext.cs delete mode 100644 Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosQueryExecutionContext.cs delete mode 100644 Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/ItemProducers/ItemProducer.cs delete mode 100644 Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/ItemProducers/ItemProducerTree.cs delete mode 100644 Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/OrderBy/CosmosOrderByItemQueryExecutionContext.ContinuationToken.cs delete mode 100644 Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/OrderBy/CosmosOrderByItemQueryExecutionContext.Drain.cs delete mode 100644 Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/OrderBy/CosmosOrderByItemQueryExecutionContext.Resume.cs delete mode 100644 Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/OrderBy/CosmosOrderByItemQueryExecutionContext.cs delete mode 100644 Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/OrderBy/OrderByItemProducerTreeComparer.cs delete mode 100644 Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/OrderBy/RangeFilterInitializationInfo.cs delete mode 100644 Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/Parallel/CosmosParallelItemQueryExecutionContext.ContinuationToken.cs delete mode 100644 Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/Parallel/CosmosParallelItemQueryExecutionContext.Drain.cs delete mode 100644 Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/Parallel/CosmosParallelItemQueryExecutionContext.Resume.cs delete mode 100644 Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/Parallel/CosmosParallelItemQueryExecutionContext.cs delete mode 100644 Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/Parallel/DeterministicParallelItemProducerTreeComparer.cs delete mode 100644 Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/Parallel/NonDeterministicParallelItemProducerTreeComparer.cs delete mode 100644 Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/Parallel/ParallelQueryConfig.cs delete mode 100644 Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/QueryExecutionContextWithException.cs rename Microsoft.Azure.Cosmos/src/Query/Core/{ExecutionContext => Pipeline}/CosmosQueryExecutionContextFactory.cs (95%) rename Microsoft.Azure.Cosmos/src/Query/Core/{ExecutionContext => Pipeline}/ExecutionEnvironment.cs (90%) rename Microsoft.Azure.Cosmos/src/Query/Core/{ExecutionContext => Pipeline/Remote}/OrderBy/CosmosElementToQueryLiteral.cs (98%) rename Microsoft.Azure.Cosmos/src/Query/Core/{ExecutionContext/ItemProducers => Pipeline/Remote/OrderBy}/ItemComparer.cs (97%) rename Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/{ => OrderBy}/OrderByCrossPartitionQueryPipelineStage.cs (99%) rename Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/{ => OrderBy}/OrderByEnumeratorComparer.cs (96%) rename Microsoft.Azure.Cosmos/src/Query/Core/{ExecutionContext => Pipeline/Remote}/OrderBy/OrderByItem.cs (95%) rename Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/{ => OrderBy}/OrderByQueryPage.cs (92%) rename Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/{ => OrderBy}/OrderByQueryPartitionRangePageAsyncEnumerator.cs (97%) rename Microsoft.Azure.Cosmos/src/Query/Core/{ExecutionContext => Pipeline/Remote}/OrderBy/OrderByQueryResult.cs (98%) rename Microsoft.Azure.Cosmos/src/Query/Core/{ExecutionContext => Pipeline/Remote}/OrderBy/SortOrder.cs (79%) rename Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/{ => Parallel}/ParallelCrossPartitionQueryPipelineStage.cs (99%) rename Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/{ => Parallel}/QueryPartitionRangePageAsyncEnumerator.cs (96%) diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/OrderByContinuationToken.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/OrderByContinuationToken.cs index 74cb8fc3f6..c3482324be 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/OrderByContinuationToken.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/OrderByContinuationToken.cs @@ -9,8 +9,8 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ContinuationTokens using Microsoft.Azure.Cosmos.CosmosElements; using Microsoft.Azure.Cosmos.CosmosElements.Numbers; using Microsoft.Azure.Cosmos.Query.Core.Exceptions; - using Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.OrderBy; using Microsoft.Azure.Cosmos.Query.Core.Monads; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote.OrderBy; using Microsoft.Azure.Documents.Routing; using Newtonsoft.Json; diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/AggregateDocumentQueryExecutionComponent.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/AggregateDocumentQueryExecutionComponent.cs index 03fd609e8b..d7a88c5ea2 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/AggregateDocumentQueryExecutionComponent.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/AggregateDocumentQueryExecutionComponent.cs @@ -7,8 +7,8 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.Aggregate using System.Collections.Generic; using System.Threading.Tasks; using Microsoft.Azure.Cosmos.CosmosElements; - 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.Aggregate; using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Aggregate.Aggregators; diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctDocumentQueryExecutionComponent.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctDocumentQueryExecutionComponent.cs index 1e983601ff..79d3f45852 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctDocumentQueryExecutionComponent.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctDocumentQueryExecutionComponent.cs @@ -6,8 +6,8 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.Distinct using System; using System.Threading.Tasks; using Microsoft.Azure.Cosmos.CosmosElements; - 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.Distinct; internal abstract partial class DistinctDocumentQueryExecutionComponent : DocumentQueryExecutionComponentBase diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/GroupBy/GroupByDocumentQueryExecutionComponent.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/GroupBy/GroupByDocumentQueryExecutionComponent.cs index 7894065adc..1e65f8d4ae 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/GroupBy/GroupByDocumentQueryExecutionComponent.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/GroupBy/GroupByDocumentQueryExecutionComponent.cs @@ -12,6 +12,7 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.GroupBy using Microsoft.Azure.Cosmos.Query.Core.Exceptions; 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.Aggregate; using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Aggregate.Aggregators; using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Distinct; diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/SkipDocumentQueryExecutionComponent.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/SkipDocumentQueryExecutionComponent.cs index fc0539f3eb..5069cd529e 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/SkipDocumentQueryExecutionComponent.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/SkipDocumentQueryExecutionComponent.cs @@ -10,6 +10,7 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.SkipTake using Microsoft.Azure.Cosmos.Query.Core.ContinuationTokens; using Microsoft.Azure.Cosmos.Query.Core.ExecutionContext; using Microsoft.Azure.Cosmos.Query.Core.Monads; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline; internal abstract partial class SkipDocumentQueryExecutionComponent : DocumentQueryExecutionComponentBase { diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/TakeDocumentQueryExecutionComponent.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/TakeDocumentQueryExecutionComponent.cs index 07011fcb8d..3a56499c2d 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/TakeDocumentQueryExecutionComponent.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/TakeDocumentQueryExecutionComponent.cs @@ -10,6 +10,7 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.SkipTake using Microsoft.Azure.Cosmos.Query.Core.ContinuationTokens; using Microsoft.Azure.Cosmos.Query.Core.ExecutionContext; using Microsoft.Azure.Cosmos.Query.Core.Monads; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline; internal abstract partial class TakeDocumentQueryExecutionComponent : DocumentQueryExecutionComponentBase { diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosCrossPartitionQueryExecutionContext.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosCrossPartitionQueryExecutionContext.cs deleted file mode 100644 index 7e60f27d3e..0000000000 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosCrossPartitionQueryExecutionContext.cs +++ /dev/null @@ -1,682 +0,0 @@ -//------------------------------------------------------------ -// Copyright (c) Microsoft Corporation. All rights reserved. -//------------------------------------------------------------ -namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionContext -{ - using System; - using System.Collections.Concurrent; - using System.Collections.Generic; - using System.Linq; - using System.Threading; - using System.Threading.Tasks; - using Core.ExecutionComponent; - using Microsoft.Azure.Cosmos.CosmosElements; - using Microsoft.Azure.Cosmos.Diagnostics; - using Microsoft.Azure.Cosmos.Query.Core; - using Microsoft.Azure.Cosmos.Query.Core.Collections; - using Microsoft.Azure.Cosmos.Query.Core.ComparableTask; - using Microsoft.Azure.Cosmos.Query.Core.ContinuationTokens; - using Microsoft.Azure.Cosmos.Query.Core.Exceptions; - using Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.ItemProducers; - using Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.Parallel; - using Microsoft.Azure.Cosmos.Query.Core.Monads; - using Microsoft.Azure.Cosmos.Query.Core.QueryClient; - using Microsoft.Azure.Cosmos.Query.Core.QueryPlan; - using PartitionKeyRange = Documents.PartitionKeyRange; - using RequestChargeTracker = Documents.RequestChargeTracker; - using RMResources = Documents.RMResources; - - /// - /// This class is responsible for maintaining a forest of . - /// The trees in this forest are ordered using a priority queue and the nodes within the forest are internally ordered using a comparator. - /// The ordering is determine by the concrete derived class. - /// This class allows derived classes to iterate through the documents in the forest using Current and MoveNext semantics. - /// This class is also responsible for prefetching documents if necessary using whose ordering is also determined by the derived classes. - /// This class also aggregated all metrics from sending queries to individual partitions. - /// - internal abstract class CosmosCrossPartitionQueryExecutionContext : CosmosQueryExecutionComponent, IDocumentQueryExecutionComponent - { - /// - /// When a document producer tree successfully fetches a page we increase the page size by this factor so that any particular document producer will only ever make O(log(n)) roundtrips, while also only ever grabbing at most twice the number of documents needed. - /// - private const double DynamicPageSizeAdjustmentFactor = 1.6; - - /// - /// Request Charge Tracker used to atomically add request charges (doubles). - /// - protected readonly RequestChargeTracker requestChargeTracker; - - /// - /// Priority Queue of ItemProducerTrees that make a forest that can be iterated on. - /// - private readonly PriorityQueue itemProducerForest; - - /// - /// Function used to determine which document producer to fetch from first - /// - private readonly Func fetchPrioirtyFunction; - - /// - /// The task scheduler that kicks off all the prefetches behind the scenes. - /// - private readonly ComparableTaskScheduler comparableTaskScheduler; - - /// - /// The equality comparer used to determine whether a document producer needs it's continuation token to be part of the composite continuation token. - /// - private readonly IEqualityComparer equalityComparer; - /// - /// The actual max page size after all the optimizations have been made it in the create document query execution context layer. - /// - private readonly long actualMaxPageSize; - - /// - /// The actual max buffered item count after all the optimizations have been made it in the create document query execution context layer. - /// - private readonly long actualMaxBufferedItemCount; - - /// - /// Injections used to reproduce special failure cases. - /// Convert this to a mock in the future. - /// - private readonly TestInjections testSettings; - - private readonly CosmosQueryContext queryContext; - - private readonly bool returnResultsInDeterministicOrder; - - protected readonly CosmosQueryClient queryClient; - - /// - /// This stores the running query metrics. - /// When a feed response is returned he take a snapshot of this bag and store it in groupedQueryMetrics. - /// The bag is then emptied and available to store the query metric for future continuations. - /// - /// - /// Due to the nature of parallel queries and prefetches the query metrics you get for a single continuation does not always - /// map to how much work was done to get that continuation. - /// For example say for a simple cross partition query we return the first page of the results from the first partition, - /// but behind the scenes we prefetched from other partitions. - /// Another example is for an order by query we return one page of results but it only required us to use partial pages from each partition, - /// but we eventually used the whole page for the next continuation; which continuation reports the cost? - /// Basically the only thing we can ensure is if you drain a query fully you should get back the same query metrics by the end. - /// - private ConcurrentBag diagnosticsPages; - - /// - /// Total number of buffered items to determine if we can go for another prefetch while still honoring the MaxBufferedItemCount. - /// - private long totalBufferedItems; - - /// - /// The total response length. - /// - private long totalResponseLengthBytes; - - /// - /// Initializes a new instance of the CosmosCrossPartitionQueryExecutionContext class. - /// - /// Constructor parameters for the base class. - /// The max concurrency - /// The max buffered item count - /// Max item count - /// Comparer used to figure out that document producer tree to serve documents from next. - /// The priority function to determine which partition to fetch documents from next. - /// Used to determine whether we need to return the continuation token for a partition. - /// Whether or not to return results in deterministic order. - /// Test settings. - protected CosmosCrossPartitionQueryExecutionContext( - CosmosQueryContext queryContext, - int? maxConcurrency, - int? maxItemCount, - int? maxBufferedItemCount, - IComparer moveNextComparer, - Func fetchPrioirtyFunction, - IEqualityComparer equalityComparer, - bool returnResultsInDeterministicOrder, - TestInjections testSettings) - { - if (moveNextComparer == null) - { - throw new ArgumentNullException(nameof(moveNextComparer)); - } - - this.queryContext = queryContext ?? throw new ArgumentNullException(nameof(queryContext)); - this.queryClient = queryContext.QueryClient ?? throw new ArgumentNullException(nameof(queryContext.QueryClient)); - this.itemProducerForest = new PriorityQueue(moveNextComparer, isSynchronized: true); - this.fetchPrioirtyFunction = fetchPrioirtyFunction ?? throw new ArgumentNullException(nameof(fetchPrioirtyFunction)); - this.comparableTaskScheduler = new ComparableTaskScheduler(maxConcurrency.GetValueOrDefault(0)); - this.equalityComparer = equalityComparer ?? throw new ArgumentNullException(nameof(equalityComparer)); - this.testSettings = testSettings; - this.requestChargeTracker = new RequestChargeTracker(); - this.diagnosticsPages = new ConcurrentBag(); - this.actualMaxPageSize = maxItemCount.GetValueOrDefault(ParallelQueryConfig.GetConfig().ClientInternalMaxItemCount); - - if (this.actualMaxPageSize < 0) - { - throw new ArgumentOutOfRangeException("actualMaxPageSize should never be less than 0"); - } - - if (this.actualMaxPageSize > int.MaxValue) - { - throw new ArgumentOutOfRangeException("actualMaxPageSize should never be greater than int.MaxValue"); - } - - if (maxBufferedItemCount.HasValue) - { - this.actualMaxBufferedItemCount = maxBufferedItemCount.Value; - } - else - { - this.actualMaxBufferedItemCount = ParallelQueryConfig.GetConfig().DefaultMaximumBufferSize; - } - - if (this.actualMaxBufferedItemCount < 0) - { - throw new ArgumentOutOfRangeException("actualMaxBufferedItemCount should never be less than 0"); - } - - if (this.actualMaxBufferedItemCount > int.MaxValue) - { - throw new ArgumentOutOfRangeException("actualMaxBufferedItemCount should never be greater than int.MaxValue"); - } - - this.CanPrefetch = maxConcurrency.HasValue && maxConcurrency.Value != 0; - - this.returnResultsInDeterministicOrder = returnResultsInDeterministicOrder; - } - - /// - /// Gets a value indicating whether this context is done having documents drained. - /// - public override bool IsDone => !this.HasMoreResults; - - protected int ActualMaxBufferedItemCount => (int)this.actualMaxBufferedItemCount; - - protected int ActualMaxPageSize => (int)this.actualMaxPageSize; - - /// - /// Gets the continuation token for the context. - /// This method is overridden by the derived class, since they all have different continuation tokens. - /// - protected abstract string ContinuationToken - { - get; - } - - /// - /// Gets a value indicating whether we are allowed to prefetch. - /// - private bool CanPrefetch { get; } - - /// - /// Gets a value indicating whether the context still has more results. - /// - private bool HasMoreResults => (this.itemProducerForest.Count != 0) && this.CurrentItemProducerTree().HasMoreResults; - - /// - /// Gets the number of documents we can still buffer. - /// - private long FreeItemSpace => this.actualMaxBufferedItemCount - Interlocked.Read(ref this.totalBufferedItems); - - /// - /// After a split you need to maintain the continuation tokens for all the child document producers until a condition is met. - /// For example lets say that a document producer is at continuation X and it gets split, - /// then the children each get continuation X, but since you only drain from one of them at a time you are left with the first child having - /// continuation X + delta and the second child having continuation X (draw this out if you are following along). - /// At this point you have the answer the question: "Which continuation token do you return to the user?". - /// Let's say you return X, then when you come back to the first child you will be repeating work, thus returning some documents more than once. - /// Let's say you return X + delta, then you fine when you return to the first child, but when you get to the second child you don't have a continuation token - /// meaning that you will be repeating all the document for the second partition up until X and again you will be returning some documents more than once. - /// Thus you have to return the continuation token for both children. - /// Both this means you are returning more than 1 continuation token for the rest of the query. - /// Well a naive optimization is to flush the continuation for a child partition once you are done draining from it, which isn't bad for a parallel query, - /// but if you have an order by query you might not be done with a producer until the end of the query. - /// The next optimization for a parallel query is to flush the continuation token the moment you start reading from a child partition. - /// This works for a parallel query, but breaks for an order by query. - /// The final realization is that for an order by query you are only choosing between multiple child partitions when their is a tie, - /// so the key is that you can dump the continuation token the moment you come across a new order by item. - /// For order by queries that is determined by the order by field and for parallel queries that is the moment you come by a new rid (which is any document, since rids are unique within a partition). - /// So by passing an equality comparer to the document producers they can determine whether they are still "active". - /// - /// - /// Returns all document producers whose continuation token you have to return. - /// Only during a split will this list contain more than 1 item. - /// - public IEnumerable GetActiveItemProducers() - { - if (this.returnResultsInDeterministicOrder) - { - ItemProducerTree current = this.itemProducerForest.Peek().CurrentItemProducerTree; - if (current.HasMoreResults && !current.IsActive) - { - // If the current document producer tree has more results, but isn't active. - // then we still want to emit it, since it won't get picked up in the below for loop. - yield return current.Root; - } - - foreach (ItemProducerTree itemProducerTree in this.itemProducerForest) - { - foreach (ItemProducer itemProducer in itemProducerTree.GetActiveItemProducers()) - { - yield return itemProducer; - } - } - } - else - { - // Just return all item producers that have a continuation token - foreach (ItemProducerTree itemProducerTree in this.itemProducerForest) - { - foreach (ItemProducerTree leaf in itemProducerTree) - { - if (leaf.HasMoreResults) - { - yield return leaf.Root; - } - } - } - } - } - - /// - /// Gets the current document producer tree that should be drained from. - /// - /// The current document producer tree that should be drained from. - public ItemProducerTree CurrentItemProducerTree() - { - return this.itemProducerForest.Peek(); - } - - /// - /// Pushes a document producer back to the queue. - /// - public void PushCurrentItemProducerTree(ItemProducerTree itemProducerTree) - { - itemProducerTree.UpdatePriority(); - this.itemProducerForest.Enqueue(itemProducerTree); - } - - /// - /// Pops the current document producer tree that should be drained from. - /// - /// The current document producer tree that should be drained from. - public ItemProducerTree PopCurrentItemProducerTree() - { - return this.itemProducerForest.Dequeue(); - } - - /// - /// Disposes of the context and implements IDisposable. - /// - public override void Dispose() - { - this.comparableTaskScheduler.Dispose(); - } - - /// - /// Stops the execution context. - /// - public override void Stop() - { - this.comparableTaskScheduler.Stop(); - } - - protected async Task TryInitializeAsync( - string collectionRid, - int initialPageSize, - SqlQuerySpec querySpecForInit, - IReadOnlyDictionary targetRangeToContinuationMap, - bool deferFirstPage, - string filter, - Func> tryFilterAsync, - CancellationToken cancellationToken) - { - if (collectionRid == null) - { - throw new ArgumentNullException(nameof(collectionRid)); - } - - if (initialPageSize < 0) - { - throw new ArgumentOutOfRangeException(nameof(initialPageSize)); - } - - if (querySpecForInit == null) - { - throw new ArgumentNullException(nameof(querySpecForInit)); - } - - if (targetRangeToContinuationMap == null) - { - throw new ArgumentNullException(nameof(targetRangeToContinuationMap)); - } - - cancellationToken.ThrowIfCancellationRequested(); - - List itemProducerTrees = new List(); - foreach (KeyValuePair rangeAndContinuationToken in targetRangeToContinuationMap) - { - PartitionKeyRange partitionKeyRange = rangeAndContinuationToken.Key; - string continuationToken = rangeAndContinuationToken.Value; - ItemProducerTree itemProducerTree = new ItemProducerTree( - this.queryContext, - querySpecForInit, - partitionKeyRange, - this.OnItemProducerTreeCompleteFetching, - this.itemProducerForest.Comparer as IComparer, - this.equalityComparer, - this.testSettings, - deferFirstPage, - collectionRid, - initialPageSize, - continuationToken) - { - Filter = filter - }; - - // Prefetch if necessary, and populate consume queue. - if (this.CanPrefetch) - { - this.TryScheduleFetch(itemProducerTree); - } - - itemProducerTrees.Add(itemProducerTree); - } - - // Using loop fission so that we can load the document producers in parallel - foreach (ItemProducerTree itemProducerTree in itemProducerTrees) - { - if (!deferFirstPage) - { - while (true) - { - (bool movedToNextPage, QueryResponseCore? failureResponse) = await itemProducerTree.TryMoveNextPageAsync(cancellationToken); - - if (failureResponse.HasValue) - { - return TryCatch.FromException( - failureResponse.Value.CosmosException); - } - - if (!movedToNextPage) - { - break; - } - - if (itemProducerTree.IsAtBeginningOfPage) - { - break; - } - - if (itemProducerTree.TryMoveNextDocumentWithinPage()) - { - break; - } - } - } - - if (tryFilterAsync != null) - { - TryCatch tryFilter = await tryFilterAsync(itemProducerTree); - if (!tryFilter.Succeeded) - { - return tryFilter; - } - } - - this.itemProducerForest.Enqueue(itemProducerTree); - } - - return TryCatch.FromResult(); - } - - protected virtual long GetAndResetResponseLengthBytes() - { - return Interlocked.Exchange(ref this.totalResponseLengthBytes, 0); - } - - protected virtual long IncrementResponseLengthBytes(long incrementValue) - { - return Interlocked.Add(ref this.totalResponseLengthBytes, incrementValue); - } - - /// - /// Tries to schedule a fetch from the document producer tree. - /// - /// The document producer tree to schedule a fetch for. - /// Whether or not the fetch was successfully scheduled. - private bool TryScheduleFetch(ItemProducerTree itemProducerTree) - { - return this.comparableTaskScheduler.TryQueueTask( - new ItemProducerTreeComparableTask( - itemProducerTree, - this.fetchPrioirtyFunction), - default); - } - - /// - /// Function that is given to all the document producers to call on once they are done fetching. - /// This is so that the CosmosCrossPartitionQueryExecutionContext can aggregate metadata from them. - /// - /// The document producer that just finished fetching. - /// The number of items that the producer just fetched. - /// The amount of RUs that the producer just consumed. - /// The length of the response the producer just got back in bytes. - /// The cancellation token. - /// - /// This function is by nature a bit racy. - /// A query might be fully drained but a background task is still fetching documents so this will get called after the context is done. - /// - private void OnItemProducerTreeCompleteFetching( - ItemProducerTree producer, - int itemsBuffered, - double resourceUnitUsage, - long responseLengthBytes, - CancellationToken token) - { - // Update charge and states - this.requestChargeTracker.AddCharge(resourceUnitUsage); - Interlocked.Add(ref this.totalBufferedItems, itemsBuffered); - this.IncrementResponseLengthBytes(responseLengthBytes); - - // Adjust the producer page size so that we reach the optimal page size. - producer.PageSize = Math.Min((long)(producer.PageSize * DynamicPageSizeAdjustmentFactor), this.actualMaxPageSize); - - // Adjust Max Degree Of Parallelism if necessary - // (needs to wait for comparable task scheduler refactor). - - // Fetch again if necessary - if (producer.HasMoreBackendResults) - { - // 4mb is the max response size - long expectedResponseSize = Math.Min(producer.PageSize, 4 * 1024 * 1024); - if (this.CanPrefetch && this.FreeItemSpace > expectedResponseSize) - { - this.TryScheduleFetch(producer); - } - } - } - - public abstract CosmosElement GetCosmosElementContinuationToken(); - - /// - /// All CrossPartitionQueries need this information on top of the parameter for DocumentQueryExecutionContextBase. - /// I moved it out into it's own type, so that we don't have to keep passing around all the individual parameters in the factory pattern. - /// This also allows us to check the arguments once instead of in each of the constructors. - /// - public readonly struct CrossPartitionInitParams - { - /// - /// Initializes a new instance of the InitParams struct. - /// - /// The Sql query spec - /// The collection rid. - /// The partitioned query execution info. - /// The partition key ranges. - /// The initial page size. - /// The max concurrency - /// The max buffered item count - /// Max item count - /// Whether or not to return results in a deterministic order. - /// Test settings. - public CrossPartitionInitParams( - SqlQuerySpec sqlQuerySpec, - string collectionRid, - PartitionedQueryExecutionInfo partitionedQueryExecutionInfo, - List partitionKeyRanges, - int initialPageSize, - int? maxConcurrency, - int? maxItemCount, - int? maxBufferedItemCount, - bool returnResultsInDeterministicOrder, - TestInjections testSettings) - { - if (string.IsNullOrWhiteSpace(collectionRid)) - { - throw new ArgumentException($"{nameof(collectionRid)} can not be null, empty, or white space."); - } - - if (partitionKeyRanges == null) - { - throw new ArgumentNullException($"{nameof(partitionKeyRanges)} can not be null."); - } - - foreach (PartitionKeyRange partitionKeyRange in partitionKeyRanges) - { - if (partitionKeyRange == null) - { - throw new ArgumentNullException($"{nameof(partitionKeyRange)} can not be null."); - } - } - - if (initialPageSize <= 0) - { - throw new ArgumentOutOfRangeException($"{nameof(initialPageSize)} must be at least 1."); - } - - //// Request continuation is allowed to be null - this.SqlQuerySpec = sqlQuerySpec ?? throw new ArgumentNullException($"{nameof(sqlQuerySpec)} can not be null."); - this.CollectionRid = collectionRid; - this.PartitionedQueryExecutionInfo = partitionedQueryExecutionInfo ?? throw new ArgumentNullException($"{nameof(partitionedQueryExecutionInfo)} can not be null."); - this.PartitionKeyRanges = partitionKeyRanges; - this.InitialPageSize = initialPageSize; - this.MaxBufferedItemCount = maxBufferedItemCount; - this.MaxConcurrency = maxConcurrency; - this.MaxItemCount = maxItemCount; - this.ReturnResultsInDeterministicOrder = returnResultsInDeterministicOrder; - this.TestSettings = testSettings; - } - - /// - /// Get the sql query spec - /// - public SqlQuerySpec SqlQuerySpec { get; } - - /// - /// Gets the collection rid to drain documents from. - /// - public string CollectionRid { get; } - - /// - /// Gets the serialized version of the PipelinedDocumentQueryExecutionContext. - /// - public PartitionedQueryExecutionInfo PartitionedQueryExecutionInfo { get; } - - /// - /// Gets the partition key ranges to fan out to. - /// - public List PartitionKeyRanges { get; } - - /// - /// Gets the initial page size for each document producer. - /// - public int InitialPageSize { get; } - - /// - /// Gets the max concurrency - /// - public int? MaxConcurrency { get; } - - /// - /// Gets the max item count - /// - public int? MaxItemCount { get; } - - /// - /// Gets the max buffered item count - /// - public int? MaxBufferedItemCount { get; } - - public bool ReturnResultsInDeterministicOrder { get; } - - public TestInjections TestSettings { get; } - } - - #region ItemProducerTreeComparableTask - /// - /// Comparable task for the ComparableTaskScheduler. - /// This is specifically for tasks that fetch from partitions in a document producer tree. - /// - private sealed class ItemProducerTreeComparableTask : ComparableTask - { - /// - /// The producer to fetch from. - /// - private readonly ItemProducerTree producer; - - /// - /// Initializes a new instance of the ItemProducerTreeComparableTask class. - /// - /// The producer to fetch from. - /// The callback to determine the fetch priority of the document producer. - public ItemProducerTreeComparableTask( - ItemProducerTree producer, - Func taskPriorityFunction) - : base(taskPriorityFunction(producer)) - { - this.producer = producer; - } - - /// - /// Entry point for the function to start fetching. - /// - /// The cancellation token. - /// A task to await on. - public override Task StartAsync(CancellationToken token) - { - return this.producer.BufferMoreDocumentsAsync(token); - } - - /// - /// Determines whether this class is equal to another task. - /// - /// The other task - /// Whether this class is equal to another task. - public override bool Equals(IComparableTask other) - { - return this.Equals(other as ItemProducerTreeComparableTask); - } - - /// - /// Gets the hash code for this task. - /// - /// The hash code for this task. - public override int GetHashCode() - { - return this.producer.PartitionKeyRange.GetHashCode(); - } - - /// - /// Internal implementation of equality. - /// - /// The other comparable task to check for equality. - /// Whether or not the comparable tasks are equal. - private bool Equals(ItemProducerTreeComparableTask other) - { - return this.producer.PartitionKeyRange.Equals(other.producer.PartitionKeyRange); - } - } - #endregion - } -} diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosQueryExecutionContext.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosQueryExecutionContext.cs deleted file mode 100644 index 1cc7f29ab2..0000000000 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosQueryExecutionContext.cs +++ /dev/null @@ -1,37 +0,0 @@ -//------------------------------------------------------------ -// Copyright (c) Microsoft Corporation. All rights reserved. -//------------------------------------------------------------ -namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionContext -{ - using System; - using System.Threading; - using System.Threading.Tasks; - using Microsoft.Azure.Cosmos.CosmosElements; - using Microsoft.Azure.Cosmos.Json; - using Microsoft.Azure.Cosmos.Query.Core.QueryClient; - - /// - /// Interface for all document query execution contexts - /// - internal abstract class CosmosQueryExecutionContext : IDisposable - { - /// - /// Gets a value indicating whether or not the context is done serving documents. - /// - public abstract bool IsDone - { - get; - } - - public abstract void Dispose(); - - /// - /// Executes the context to feed the next page of results. - /// - /// The cancellation token. - /// A task to await on, which in return provides a DoucmentFeedResponse of documents. - public abstract Task ExecuteNextAsync(CancellationToken cancellationToken); - - public abstract CosmosElement GetCosmosElementContinuationToken(); - } -} diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/ItemProducers/ItemProducer.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/ItemProducers/ItemProducer.cs deleted file mode 100644 index 2ad20bebe2..0000000000 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/ItemProducers/ItemProducer.cs +++ /dev/null @@ -1,439 +0,0 @@ -//------------------------------------------------------------ -// Copyright (c) Microsoft Corporation. All rights reserved. -//------------------------------------------------------------ -namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.ItemProducers -{ - using System; - using System.Collections.Generic; - using System.Threading; - using System.Threading.Tasks; - using Microsoft.Azure.Cosmos; - using Microsoft.Azure.Cosmos.CosmosElements; - using Microsoft.Azure.Cosmos.Diagnostics; - using Microsoft.Azure.Cosmos.Query.Core; - using Microsoft.Azure.Cosmos.Query.Core.Collections; - using Microsoft.Azure.Cosmos.Query.Core.Metrics; - using Microsoft.Azure.Cosmos.Query.Core.Monads; - using Microsoft.Azure.Cosmos.Query.Core.Pipeline; - using Microsoft.Azure.Cosmos.Query.Core.QueryClient; - using Microsoft.Azure.Cosmos.Resource.CosmosExceptions; - using Microsoft.Azure.Documents; - using PartitionKeyRange = Documents.PartitionKeyRange; - using PartitionKeyRangeIdentity = Documents.PartitionKeyRangeIdentity; - - /// - /// The ItemProducer is the base unit of buffering and iterating through documents. - /// Note that a document producer will let you iterate through documents within the pages of a partition and maintain any state. - /// In pseudo code this works out to: - /// for page in partition: - /// for document in page: - /// yield document - /// update_state() - /// - internal sealed class ItemProducer - { - /// - /// The buffered pages that is thread safe, since the producer and consumer of the queue can be on different threads. - /// We buffer TryCatch of DoucmentFeedResponse of T, since we want to buffer exceptions, - /// so that the exception is thrown on the consumer thread (instead of the background producer thread), thus observing the exception. - /// - private readonly AsyncCollection bufferedPages; - - /// - /// The document producer can only be fetching one page at a time. - /// Since the fetch function can be called by the execution contexts or the scheduler, we use this semaphore to keep the fetch function thread safe. - /// - private readonly SemaphoreSlim fetchSemaphore; - - /// - /// Once a document producer tree finishes fetching document they should call on this function so that the higher level execution context can aggregate the number of documents fetched, the request charge, and the query metrics. - /// - private readonly ProduceAsyncCompleteDelegate produceAsyncCompleteCallback; - - /// - /// Equality comparer to determine if you have come across a distinct document according to the sort order. - /// - private readonly IEqualityComparer equalityComparer; - - private readonly CosmosQueryContext queryContext; - - private readonly SqlQuerySpec querySpecForInit; - - private readonly TestInjections testFlags; - - /// - /// Over the duration of the life time of a document producer the page size will change, since we have an adaptive page size. - /// - private readonly long pageSize; - - /// - /// The current page that is being enumerated. - /// - private IEnumerator CurrentPage; - - /// - /// The number of items left in the current page, which is used by parallel queries since they need to drain full pages. - /// - private int itemsLeftInCurrentPage; - - /// - /// The number of items currently buffered, which is used by the scheduler incase you want to implement give less full document producers a higher priority. - /// - private long bufferedItemCount; - - /// - /// Whether or not the document producer has started fetching. - /// - private bool hasStartedFetching; - - /// - /// Need this flag so that the document producer stops buffering more results after a fatal exception. - /// - private bool hitException; - - private bool enumeratorPrimed; - - /// - /// Initializes a new instance of the ItemProducer class. - /// - /// request context - /// query spec for initialization - /// The partition key range. - /// The callback to call once you are done fetching. - /// The comparer to use to determine whether the producer has seen a new document. - /// Flags used to help faciliate testing. - /// The initial page size. - /// The initial continuation token. - public ItemProducer( - CosmosQueryContext queryContext, - SqlQuerySpec querySpecForInit, - PartitionKeyRange partitionKeyRange, - ProduceAsyncCompleteDelegate produceAsyncCompleteCallback, - IEqualityComparer equalityComparer, - TestInjections testFlags, - long initialPageSize = 50, - string initialContinuationToken = null) - { - this.bufferedPages = new AsyncCollection(); - - // We use a binary semaphore to get the behavior of a mutex, - // since fetching documents from the backend using a continuation token is a critical section. - this.fetchSemaphore = new SemaphoreSlim(1, 1); - this.queryContext = queryContext; - this.querySpecForInit = querySpecForInit; - this.PartitionKeyRange = partitionKeyRange ?? throw new ArgumentNullException(nameof(partitionKeyRange)); - this.produceAsyncCompleteCallback = produceAsyncCompleteCallback ?? throw new ArgumentNullException(nameof(produceAsyncCompleteCallback)); - this.equalityComparer = equalityComparer ?? throw new ArgumentNullException(nameof(equalityComparer)); - this.pageSize = initialPageSize; - this.CurrentContinuationToken = initialContinuationToken; - this.BackendContinuationToken = initialContinuationToken; - this.PreviousContinuationToken = initialContinuationToken; - if (!string.IsNullOrEmpty(initialContinuationToken)) - { - this.hasStartedFetching = true; - this.IsActive = true; - } - - this.testFlags = testFlags; - - this.HasMoreResults = true; - } - - public delegate void ProduceAsyncCompleteDelegate( - int numberOfDocuments, - double requestCharge, - long responseLengthInBytes, - CancellationToken token); - - /// - /// Gets the for the partition that this document producer is fetching from. - /// - public PartitionKeyRange PartitionKeyRange - { - get; - } - - /// - /// Gets or sets the filter predicate for the document producer that is used by order by execution context. - /// - public string Filter { get; set; } - - /// - /// Gets the previous continuation token. - /// - public string PreviousContinuationToken { get; private set; } - - /// - /// The current continuation token that the user has read from the document producer tree. - /// This is used for determining whether there are more results. - /// - public string CurrentContinuationToken { get; private set; } - - /// - /// Gets the backend continuation token. - /// - public string BackendContinuationToken { get; private set; } - - /// - /// Gets a value indicating whether the continuation token for this producer needs to be given back as part of the composite continuation token. - /// - public bool IsActive { get; private set; } - - /// - /// Gets a value indicating whether this producer is at the beginning of the page. - /// - public bool IsAtBeginningOfPage { get; private set; } - - /// - /// Gets a value indicating whether this producer has more results. - /// - public bool HasMoreResults { get; private set; } - - /// - /// Gets a value indicating whether this producer has more backend results. - /// - public bool HasMoreBackendResults => this.hasStartedFetching == false - || (this.hasStartedFetching == true && !string.IsNullOrEmpty(this.BackendContinuationToken)); - - /// - /// Gets how many items are left in the current page. - /// - public int ItemsLeftInCurrentPage => this.itemsLeftInCurrentPage; - - /// - /// Gets how many documents are buffered in this producer. - /// - public int BufferedItemCount => (int)this.bufferedItemCount; - - /// - /// Gets or sets the page size of this producer. - /// - public long PageSize { get; set; } - - /// - /// Gets the activity for the last request made by this document producer. - /// - public Guid ActivityId { get; private set; } - - public CosmosQueryExecutionInfo CosmosQueryExecutionInfo { get; private set; } - - /// - /// Gets the current document in this producer. - /// - public CosmosElement Current => this.CurrentPage?.Current; - - /// - /// A static object representing that the move next operation succeeded, and was able to load the next page - /// - internal static readonly (bool successfullyMovedNext, QueryResponseCore? failureResponse) IsSuccessResponse = (true, null); - - /// - /// A static object representing that there is no more pages to load. - /// - internal static readonly (bool successfullyMovedNext, QueryResponseCore? failureResponse) IsDoneResponse = (false, null); - - /// - /// Buffers more documents if the producer is empty. - /// - /// The cancellation token. - /// A task to await on. - public async Task BufferMoreIfEmptyAsync(CancellationToken token) - { - token.ThrowIfCancellationRequested(); - - if (this.bufferedPages.Count == 0) - { - await this.BufferMoreDocumentsAsync(token); - } - } - - /// - /// Buffers more documents in the producer. - /// - /// The cancellation token. - /// A task to await on. - public async Task BufferMoreDocumentsAsync(CancellationToken token) - { - token.ThrowIfCancellationRequested(); - - try - { - await this.fetchSemaphore.WaitAsync(); - if (!this.HasMoreBackendResults || this.hitException) - { - // Just NOP - return; - } - - int pageSize = (int)Math.Min(this.pageSize, int.MaxValue); - - TryCatch tryGetQueryPage = await this.queryContext.ExecuteQueryAsync( - querySpecForInit: this.querySpecForInit, - continuationToken: this.BackendContinuationToken, - partitionKeyRange: new PartitionKeyRangeIdentity( - this.queryContext.ContainerResourceId, - this.PartitionKeyRange.Id), - isContinuationExpected: this.queryContext.IsContinuationExpected, - pageSize: pageSize, - cancellationToken: token); - - QueryResponseCore feedResponse; - if (tryGetQueryPage.Succeeded) - { - feedResponse = QueryResponseCore.CreateSuccess( - result: tryGetQueryPage.Result.Documents, - requestCharge: tryGetQueryPage.Result.RequestCharge, - activityId: tryGetQueryPage.Result.ActivityId, - responseLengthBytes: tryGetQueryPage.Result.ResponseLengthInBytes, - disallowContinuationTokenMessage: default, - continuationToken: tryGetQueryPage.Result.State == null ? null : ((CosmosString)tryGetQueryPage.Result.State.Value).Value, - cosmosQueryExecutionInfo: tryGetQueryPage.Result.CosmosQueryExecutionInfo); - } - else - { - feedResponse = QueryResponseCore.CreateFailure( - statusCode: ((CosmosException)tryGetQueryPage.Exception).StatusCode, - subStatusCodes: (SubStatusCodes)((CosmosException)tryGetQueryPage.Exception).SubStatusCode, - cosmosException: (CosmosException)tryGetQueryPage.Exception, - requestCharge: ((CosmosException)tryGetQueryPage.Exception).RequestCharge, - activityId: ((CosmosException)tryGetQueryPage.Exception).ActivityId); - } - - this.CosmosQueryExecutionInfo = feedResponse.CosmosQueryExecutionInfo; - - if ((this.testFlags != null) && this.testFlags.SimulateThrottles) - { - Random random = new Random(); - if (random.Next() % 2 == 0) - { - feedResponse = QueryResponseCore.CreateFailure( - statusCode: (System.Net.HttpStatusCode)429, - subStatusCodes: null, - cosmosException: CosmosExceptionFactory.CreateThrottledException("Request Rate Too Large"), - requestCharge: 0, - activityId: QueryResponseCore.EmptyGuidString); - } - } - - // Can not simulate an empty page on the first page, since we would return a null continuation token, which will end the query early. - if ((this.testFlags != null) && this.testFlags.SimulateEmptyPages && (this.BackendContinuationToken != null)) - { - Random random = new Random(); - if (random.Next() % 2 == 0) - { - feedResponse = QueryResponseCore.CreateSuccess( - result: new List(), - requestCharge: 0, - activityId: QueryResponseCore.EmptyGuidString, - responseLengthBytes: 0, - disallowContinuationTokenMessage: null, - continuationToken: this.BackendContinuationToken); - } - } - - this.hasStartedFetching = true; - this.ActivityId = Guid.Parse(feedResponse.ActivityId); - await this.bufferedPages.AddAsync(feedResponse); - if (!feedResponse.IsSuccess) - { - // set this flag so that people stop trying to buffer more on this producer. - this.hitException = true; - return; - } - - // The backend continuation token is used for the children on splits - // and should not be updated on exceptions - this.BackendContinuationToken = feedResponse.ContinuationToken; - - Interlocked.Add(ref this.bufferedItemCount, feedResponse.CosmosElements.Count); - - this.produceAsyncCompleteCallback( - feedResponse.CosmosElements.Count, - feedResponse.RequestCharge, - feedResponse.ResponseLengthBytes, - token); - - } - finally - { - this.fetchSemaphore.Release(); - } - } - - public void Shutdown() - { - this.HasMoreResults = false; - } - - public async Task<(bool movedToNextPage, QueryResponseCore? failureReponse)> TryMoveNextPageAsync(CancellationToken cancellationToken) - { - cancellationToken.ThrowIfCancellationRequested(); - if (this.itemsLeftInCurrentPage != 0) - { - throw new InvalidOperationException("Tried to move onto the next page before finishing the first page."); - } - - // We need to buffer pages if empty, since thats how we know there are no more pages left. - await this.BufferMoreIfEmptyAsync(cancellationToken); - if (this.bufferedPages.Count == 0) - { - this.HasMoreResults = false; - return ItemProducer.IsDoneResponse; - } - - // Pull a FeedResponse using TryCatch (we could have buffered an exception). - QueryResponseCore queryResponse = await this.bufferedPages.TakeAsync(cancellationToken); - if (!queryResponse.IsSuccess) - { - this.HasMoreResults = false; - return (false, queryResponse); - } - - // Update the state. - this.PreviousContinuationToken = this.CurrentContinuationToken; - this.CurrentContinuationToken = queryResponse.ContinuationToken; - this.CurrentPage = queryResponse.CosmosElements.GetEnumerator(); - this.itemsLeftInCurrentPage = queryResponse.CosmosElements.Count; - this.enumeratorPrimed = false; - this.IsAtBeginningOfPage = false; - - return ItemProducer.IsSuccessResponse; - } - - public bool TryMoveNextDocumentWithinPage() - { - if (this.CurrentPage == null) - { - return false; - } - - CosmosElement originalCurrent = this.Current; - bool movedNext = this.CurrentPage.MoveNext(); - - if (!movedNext || ((originalCurrent != null) && !this.equalityComparer.Equals(originalCurrent, this.Current))) - { - this.IsActive = false; - } - - if (!this.enumeratorPrimed) - { - this.IsAtBeginningOfPage = true; - this.enumeratorPrimed = true; - } - else - { - Interlocked.Decrement(ref this.bufferedItemCount); - Interlocked.Decrement(ref this.itemsLeftInCurrentPage); - - this.IsAtBeginningOfPage = false; - } - - if (!movedNext && (this.CurrentContinuationToken == null)) - { - this.HasMoreResults = false; - } - - return movedNext; - } - } -} \ No newline at end of file diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/ItemProducers/ItemProducerTree.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/ItemProducers/ItemProducerTree.cs deleted file mode 100644 index e70805add7..0000000000 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/ItemProducers/ItemProducerTree.cs +++ /dev/null @@ -1,742 +0,0 @@ -//------------------------------------------------------------ -// Copyright (c) Microsoft Corporation. All rights reserved. -//------------------------------------------------------------ -namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.ItemProducers -{ - using System; - using System.Collections; - using System.Collections.Generic; - using System.Globalization; - using System.Linq; - using System.Net; - using System.Threading; - using System.Threading.Tasks; - using Microsoft.Azure.Cosmos.CosmosElements; - using Microsoft.Azure.Cosmos.Diagnostics; - using Microsoft.Azure.Cosmos.Query.Core; - using Microsoft.Azure.Cosmos.Query.Core.Collections; - using Microsoft.Azure.Cosmos.Query.Core.QueryClient; - - /// - /// This class is responsible for fetching documents from a partition and all it's descendants, which is modeled as a tree of document producers. - /// The root node is responsible for buffering documents from the root partition and the children recursively buffer documents for their corresponding partitions. - /// The tree itself allows a user to iterate through it's documents using a comparator and Current / Move Next Async functions. - /// Note that if a user wants to determine the current document it will take the max of it's buffered documents and the recursive max of it's children. - /// Also note that if there are no buffered documents for any node in the recursive evaluation, then those nodes will go for a fetch. - /// Finally note that due to the tree structure of this class it is inherently split proof. - /// If any leaf node in the tree encounters a split exception it will spawn child document producer trees (any many as needed, so multiple splits is handled) and continue on as if the split never happened. - /// This code does not handle merges, but we will cross that bridge when we have to (I am currently thinking about a linked list where the nodes represent document producers and you can merge adjacent nodes). - /// As a implementation detail the documents are buffered and logically enumerated as a nested loop. The following is the pseudo code: - /// for partition in document_producer_tree: - /// for page in partition: - /// for document in page: - /// yield document. - /// And the way this is done is by buffering pages and updating the state of the ItemProducerTree whenever a user crosses a page boundary. - /// - internal sealed class ItemProducerTree : IEnumerable - { - /// - /// The child partitions of this node in the tree that are added after a split. - /// - private readonly PriorityQueue children; - - /// - /// Callback to create child document producer trees once a split happens. - /// - private readonly Func createItemProducerTreeCallback; - - /// - /// Whether or not to defer fetching the first page from all the partitions. - /// - private readonly bool deferFirstPage; - - /// - /// The collection rid to drain from. - /// - private readonly string collectionRid; - - /// - /// Semaphore to ensure mutual exclusion during fetching from a tree. - /// This is to ensure that there is no race conditions during splits. - /// - private readonly SemaphoreSlim executeWithSplitProofingSemaphore; - - private readonly CosmosQueryClient queryClient; - - /// - /// Initializes a new instance of the ItemProducerTree class. - /// - /// query context. - /// query spec init. - /// The partition key range. - /// Callback to invoke once a fetch finishes. - /// Comparer to determine, which tree to produce from. - /// Comparer to see if we need to return the continuation token for a partition. - /// Test flags. - /// Whether or not to defer fetching the first page. - /// The collection to drain from. - /// The initial page size. - /// The initial continuation token. - public ItemProducerTree( - CosmosQueryContext queryContext, - SqlQuerySpec querySpecForInit, - Documents.PartitionKeyRange partitionKeyRange, - ProduceAsyncCompleteDelegate produceAsyncCompleteCallback, - IComparer itemProducerTreeComparer, - IEqualityComparer equalityComparer, - TestInjections testSettings, - bool deferFirstPage, - string collectionRid, - long initialPageSize = 50, - string initialContinuationToken = null) - { - if (queryContext == null) - { - throw new ArgumentNullException($"{nameof(queryContext)}"); - } - - if (itemProducerTreeComparer == null) - { - throw new ArgumentNullException($"{nameof(itemProducerTreeComparer)}"); - } - - if (produceAsyncCompleteCallback == null) - { - throw new ArgumentNullException($"{nameof(produceAsyncCompleteCallback)}"); - } - - if (itemProducerTreeComparer == null) - { - throw new ArgumentNullException($"{nameof(itemProducerTreeComparer)}"); - } - - if (equalityComparer == null) - { - throw new ArgumentNullException($"{nameof(equalityComparer)}"); - } - - if (string.IsNullOrEmpty(collectionRid)) - { - throw new ArgumentException($"{nameof(collectionRid)} can not be null or empty."); - } - - this.Root = new ItemProducer( - queryContext, - querySpecForInit, - partitionKeyRange, - (itemsBuffered, resourceUnitUsage, requestLength, token) => produceAsyncCompleteCallback(this, itemsBuffered, resourceUnitUsage, requestLength, token), - equalityComparer, - testSettings, - initialPageSize, - initialContinuationToken); - - this.queryClient = queryContext.QueryClient; - this.children = new PriorityQueue(itemProducerTreeComparer, true); - this.deferFirstPage = deferFirstPage; - this.collectionRid = collectionRid; - this.createItemProducerTreeCallback = ItemProducerTree.CreateItemProducerTreeCallback( - queryContext, - querySpecForInit, - produceAsyncCompleteCallback, - itemProducerTreeComparer, - equalityComparer, - testSettings, - deferFirstPage, - collectionRid, - initialPageSize); - this.executeWithSplitProofingSemaphore = new SemaphoreSlim(1, 1); - } - - public delegate void ProduceAsyncCompleteDelegate( - ItemProducerTree itemProducerTree, - int numberOfDocuments, - double requestCharge, - long responseLengthInBytes, - CancellationToken token); - - /// - /// Gets the root document from the tree. - /// - public ItemProducer Root { get; } - - /// - /// Gets the partition key range from the current document producer tree. - /// - public Documents.PartitionKeyRange PartitionKeyRange - { - get - { - if (this.CurrentItemProducerTree == this) - { - return this.Root.PartitionKeyRange; - } - else - { - return this.CurrentItemProducerTree.PartitionKeyRange; - } - } - } - - /// - /// Gets or sets the filter for the current document producer tree. - /// - public string Filter - { - get - { - if (this.CurrentItemProducerTree == this) - { - return this.Root.Filter; - } - else - { - return this.CurrentItemProducerTree.Filter; - } - } - - set - { - if (this.CurrentItemProducerTree == this) - { - this.Root.Filter = value; - } - else - { - this.CurrentItemProducerTree.Filter = value; - } - } - } - - /// - /// Gets the current (highest priority) document producer tree from all subtrees. - /// - public ItemProducerTree CurrentItemProducerTree - { - get - { - if (this.HasSplit && !this.Root.HasMoreResults) - { - // If the partition has split and there are no more results in the parent buffer - // then just pull from the highest priority child (with recursive decent). - - return this.children.Peek().CurrentItemProducerTree; - } - else - { - // The partition has not split or there are still documents buffered, - // so keep trying to read from it. - return this; - } - } - } - - /// - /// Gets a value indicating whether the document producer tree is at the beginning of the page for the current document producer. - /// - public bool IsAtBeginningOfPage - { - get - { - if (this.CurrentItemProducerTree == this) - { - return this.Root.IsAtBeginningOfPage; - } - else - { - return this.CurrentItemProducerTree.IsAtBeginningOfPage; - } - } - } - - /// - /// Gets a value indicating whether the document producer tree has more results. - /// - public bool HasMoreResults => this.Root.HasMoreResults - || (this.HasSplit && this.children.Peek().HasMoreResults); - - /// - /// Gets a value indicating whether the document producer tree has more backend results. - /// - public bool HasMoreBackendResults => this.Root.HasMoreBackendResults - || (this.HasSplit && this.children.Peek().HasMoreBackendResults); - - /// - /// Gets whether there are items left in the current page of the document producer tree. - /// - public int ItemsLeftInCurrentPage - { - get - { - if (this.CurrentItemProducerTree == this) - { - return this.Root.ItemsLeftInCurrentPage; - } - else - { - return this.CurrentItemProducerTree.ItemsLeftInCurrentPage; - } - } - } - - /// - /// Gets the buffered item count in the current document producer tree. - /// - public int BufferedItemCount - { - get - { - if (this.CurrentItemProducerTree == this) - { - return this.Root.BufferedItemCount; - } - else - { - return this.CurrentItemProducerTree.BufferedItemCount; - } - } - } - - /// - /// Gets a value indicating whether the document producer tree is active. - /// - public bool IsActive => this.Root.IsActive || this.children.Any((child) => child.IsActive); - - /// - /// Gets or sets the page size for this document producer tree. - /// - public long PageSize - { - get - { - if (this.CurrentItemProducerTree == this) - { - return this.Root.PageSize; - } - else - { - return this.CurrentItemProducerTree.PageSize; - } - } - - set - { - if (this.CurrentItemProducerTree == this) - { - this.Root.PageSize = value; - } - else - { - this.CurrentItemProducerTree.PageSize = value; - } - } - } - - /// - /// Gets the activity id from the current document producer tree. - /// - public Guid ActivityId - { - get - { - if (this.CurrentItemProducerTree == this) - { - return this.Root.ActivityId; - } - else - { - return this.CurrentItemProducerTree.ActivityId; - } - } - } - - public CosmosQueryExecutionInfo CosmosQueryExecutionInfo - { - get - { - if (this.CurrentItemProducerTree == this) - { - return this.Root.CosmosQueryExecutionInfo; - } - else - { - return this.CurrentItemProducerTree.CosmosQueryExecutionInfo; - } - } - } - - /// - /// Gets the current item from the document producer tree. - /// - public CosmosElement Current - { - get - { - if (this.CurrentItemProducerTree == this) - { - return this.Root.Current; - } - else - { - return this.CurrentItemProducerTree.Current; - } - } - } - - /// - /// Gets a value indicating whether the document producer tree has split. - /// - private bool HasSplit => this.children.Count != 0; - - /// - /// Gets the document producers that need their continuation token return to the user. - /// - /// The document producers that need their continuation token return to the user. - public IEnumerable GetActiveItemProducers() - { - if (!this.HasSplit) - { - if (this.Root.IsActive) - { - yield return this.Root; - } - } - else - { - // A document producer is "active" if it resumed from a continuation token and has more buffered results. - if (this.Root.IsActive && this.Root.BufferedItemCount != 0) - { - // has split but need to check if parent is fully drained - yield return this.Root; - } - else - { - foreach (ItemProducerTree child in this.children) - { - foreach (ItemProducer activeItemProducer in child.GetActiveItemProducers()) - { - yield return activeItemProducer; - } - } - } - } - } - - /// - /// Gets the enumerator for all the leaf level document producers. - /// - /// The enumerator for all the leaf level document producers. - public IEnumerator GetEnumerator() - { - if (this.children.Count == 0) - { - yield return this; - } - - foreach (ItemProducerTree child in this.children) - { - foreach (ItemProducerTree itemProducer in child) - { - yield return itemProducer; - } - } - } - - /// - /// Gets the enumerator. - /// - /// The enumerator. - IEnumerator IEnumerable.GetEnumerator() - { - return this.GetEnumerator(); - } - - /// - /// Callback to create a child document producer tree based on the partition key range. - /// - /// request context - /// query spec for initialization - /// Callback to invoke once a fetch finishes. - /// Comparer to determine, which tree to produce from. - /// Comparer to see if we need to return the continuation token for a partition. - /// Test flags. - /// Whether or not to defer fetching the first page. - /// The collection to drain from. - /// The initial page size. - /// A function that given a partition key range and continuation token will create a document producer. - private static Func CreateItemProducerTreeCallback( - CosmosQueryContext queryContext, - SqlQuerySpec querySpecForInit, - ProduceAsyncCompleteDelegate produceAsyncCompleteCallback, - IComparer itemProducerTreeComparer, - IEqualityComparer equalityComparer, - TestInjections testFlags, - bool deferFirstPage, - string collectionRid, - long initialPageSize = 50) - { - return (partitionKeyRange, continuationToken) => - { - return new ItemProducerTree( - queryContext, - querySpecForInit, - partitionKeyRange, - produceAsyncCompleteCallback, - itemProducerTreeComparer, - equalityComparer, - testFlags, - deferFirstPage, - collectionRid, - initialPageSize, - continuationToken); - }; - } - - /// - /// Given a document client exception this function determines whether it was caused due to a split. - /// - /// The document client exception - /// Whether or not the exception was due to a split. - private static bool IsSplitException(QueryResponseCore ex) - { - return ex.StatusCode == HttpStatusCode.Gone && ex.SubStatusCode == Documents.SubStatusCodes.PartitionKeyRangeGone; - } - - public void UpdatePriority() - { - if (this.children.Count != 0) - { - this.children.Enqueue(this.children.Dequeue()); - } - } - - public async Task<(bool movedToNextPage, QueryResponseCore? failureResponse)> TryMoveNextPageAsync(CancellationToken cancellationToken) - { - cancellationToken.ThrowIfCancellationRequested(); - return await this.ExecuteWithSplitProofingAsync( - function: this.TryMoveNextPageImplementationAsync, - functionNeedsBeReexecuted: false, - cancellationToken: cancellationToken); - } - - public async Task<(bool movedToNextPage, QueryResponseCore? failureResponse)> TryMoveNextPageIfNotSplitAsync(CancellationToken cancellationToken) - { - cancellationToken.ThrowIfCancellationRequested(); - return await this.ExecuteWithSplitProofingAsync( - function: this.TryMoveNextPageIfNotSplitAsyncImplementationAsync, - functionNeedsBeReexecuted: false, - cancellationToken: cancellationToken); - } - - public async Task<(bool bufferedMoreDocuments, QueryResponseCore? failureResponse)> BufferMoreDocumentsAsync(CancellationToken cancellationToken) - { - cancellationToken.ThrowIfCancellationRequested(); - return await this.ExecuteWithSplitProofingAsync( - function: this.BufferMoreDocumentsImplementationAsync, - functionNeedsBeReexecuted: true, - cancellationToken: cancellationToken); - } - - public bool TryMoveNextDocumentWithinPage() - { - if (this.CurrentItemProducerTree == this) - { - return this.Root.TryMoveNextDocumentWithinPage(); - } - else - { - return this.CurrentItemProducerTree.TryMoveNextDocumentWithinPage(); - } - } - - private async Task<(bool succeeded, QueryResponseCore? failureResponse)> TryMoveNextPageImplementationAsync(CancellationToken cancellationToken) - { - cancellationToken.ThrowIfCancellationRequested(); - if (this.CurrentItemProducerTree == this) - { - return await this.Root.TryMoveNextPageAsync(cancellationToken); - } - else - { - return await this.CurrentItemProducerTree.TryMoveNextPageAsync(cancellationToken); - } - } - - private async Task<(bool successfullyMovedNext, QueryResponseCore? failureResponse)> TryMoveNextPageIfNotSplitAsyncImplementationAsync(CancellationToken cancellationToken) - { - cancellationToken.ThrowIfCancellationRequested(); - - if (this.HasSplit) - { - return ItemProducer.IsDoneResponse; - } - - return await this.TryMoveNextPageImplementationAsync(cancellationToken); - } - - private async Task<(bool successfullyMovedNext, QueryResponseCore? failureResponse)> BufferMoreDocumentsImplementationAsync(CancellationToken cancellationToken) - { - cancellationToken.ThrowIfCancellationRequested(); - - if (this.CurrentItemProducerTree == this) - { - if (!this.HasMoreBackendResults || this.HasSplit) - { - // Just no-op, since this method might be called by the scheduler, which doesn't know of the reconfiguration yet. - return ItemProducer.IsSuccessResponse; - } - - await this.Root.BufferMoreDocumentsAsync(cancellationToken); - } - else - { - await this.CurrentItemProducerTree.BufferMoreDocumentsAsync(cancellationToken); - } - - return ItemProducer.IsSuccessResponse; - } - - /// - /// This function will execute any function in a split proof manner. - /// What it does is it will try to execute the supplied function and catch any gone exceptions do to a split. - /// If a split happens when this function will - /// - /// The function to execute in a split proof manner. - /// If the function needs to be re-executed after split. - /// The cancellation token. - /// - /// - /// This function is thread safe meaning that if multiple functions want to execute in a split proof manner, - /// then they will need to go one after another. - /// This is required since you could have the follow scenario: - /// Time | CurrentItemProducer | Thread 1 | Thread2 - /// 0 | 0 | MoveNextAsync | BufferMore - /// 1 | 0 | Split | Split - /// - /// - /// Therefore thread 1 and thread 2 both think that document producer 0 got split and they both try to repair the execution context, - /// which is a race condition. - /// Note that this thread safety / serial behavior is only scoped to a single document producer tree - /// meaning this should not have a performance hit on the scheduler that is prefetching from other partitions. - /// - /// - /// The result of the function would have returned as if there were no splits. - private async Task<(bool successfullyMovedNext, QueryResponseCore? failureResponse)> ExecuteWithSplitProofingAsync( - Func> function, - bool functionNeedsBeReexecuted, - CancellationToken cancellationToken) - { - cancellationToken.ThrowIfCancellationRequested(); - - while (true) - { - try - { - await this.executeWithSplitProofingSemaphore.WaitAsync(); - (bool successfullyMovedNext, QueryResponseCore? failureResponse) response = await function(cancellationToken); - if (response.failureResponse == null || !ItemProducerTree.IsSplitException(response.failureResponse.Value)) - { - return response; - } - - // Split just happened - ItemProducerTree splitItemProducerTree = this.CurrentItemProducerTree; - - if (!functionNeedsBeReexecuted) - { - // If we run into a split for MoveNextAsync then that means we failed to buffer more meaning the document producer has no more results - splitItemProducerTree.Root.Shutdown(); - } - - // Repair the execution context: Get the replacement document producers and add them to the tree. - IReadOnlyList replacementRanges = await this.GetReplacementRangesAsync(splitItemProducerTree.PartitionKeyRange, this.collectionRid); - foreach (Documents.PartitionKeyRange replacementRange in replacementRanges) - { - ItemProducerTree replacementItemProducerTree = this.createItemProducerTreeCallback(replacementRange, splitItemProducerTree.Root.BackendContinuationToken); - - if (!this.deferFirstPage) - { - while (true) - { - (bool movedToNextPage, QueryResponseCore? failureResponse) = await replacementItemProducerTree.TryMoveNextPageAsync(cancellationToken); - if (!movedToNextPage) - { - if (failureResponse.HasValue) - { - return (movedToNextPage, failureResponse); - } - - break; - } - - // Order By requires at least one item in the tree so that the comparator will work. - if (replacementItemProducerTree.TryMoveNextDocumentWithinPage()) - { - break; - } - } - } - - replacementItemProducerTree.Filter = splitItemProducerTree.Root.Filter; - if (replacementItemProducerTree.HasMoreResults) - { - if (!splitItemProducerTree.children.TryAdd(replacementItemProducerTree)) - { - throw new InvalidOperationException("Unable to add child document producer tree"); - } - } - } - - if (!this.deferFirstPage) - { - (bool, QueryResponseCore?) fakeResponse; - if (splitItemProducerTree.children.Count == 0) - { - // None of the children had results - fakeResponse = ItemProducer.IsDoneResponse; - } - else - { - // We don't want to call move next async again, since we already did when creating the document producers - fakeResponse = ItemProducer.IsSuccessResponse; - } - - return fakeResponse; - } - } - finally - { - this.executeWithSplitProofingSemaphore.Release(); - } - } - } - - /// - /// Gets the replacement ranges for the target range that got split. - /// - /// The target range that got split. - /// The collection rid. - /// The replacement ranges for the target range that got split. - private async Task> GetReplacementRangesAsync(Documents.PartitionKeyRange targetRange, string collectionRid) - { - IReadOnlyList replacementRanges = await this.queryClient.TryGetOverlappingRangesAsync( - collectionRid, - targetRange.ToRange(), - true); - - string replaceMinInclusive = replacementRanges.First().MinInclusive; - string replaceMaxExclusive = replacementRanges.Last().MaxExclusive; - if (!replaceMinInclusive.Equals(targetRange.MinInclusive, StringComparison.Ordinal) || !replaceMaxExclusive.Equals(targetRange.MaxExclusive, StringComparison.Ordinal)) - { - throw new Documents.InternalServerErrorException(string.Format( - CultureInfo.InvariantCulture, - "Target range and Replacement range has mismatched min/max. Target range: [{0}, {1}). Replacement range: [{2}, {3}).", - targetRange.MinInclusive, - targetRange.MaxExclusive, - replaceMinInclusive, - replaceMaxExclusive)); - } - - return replacementRanges; - } - } -} \ No newline at end of file diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/OrderBy/CosmosOrderByItemQueryExecutionContext.ContinuationToken.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/OrderBy/CosmosOrderByItemQueryExecutionContext.ContinuationToken.cs deleted file mode 100644 index c0df108ecd..0000000000 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/OrderBy/CosmosOrderByItemQueryExecutionContext.ContinuationToken.cs +++ /dev/null @@ -1,149 +0,0 @@ -// ------------------------------------------------------------ -// Copyright (c) Microsoft Corporation. All rights reserved. -// ------------------------------------------------------------ - -namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.OrderBy -{ - using System; - using System.Collections.Generic; - using System.Linq; - using Microsoft.Azure.Cosmos.CosmosElements; - using Microsoft.Azure.Cosmos.Query.Core.ContinuationTokens; - using Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.ItemProducers; - - internal sealed partial class CosmosOrderByItemQueryExecutionContext - { - /// - /// Gets the continuation token for an order by query. - /// - protected override string ContinuationToken - { - // In general the continuation token for order by queries contains the following information: - // 1) What partition did we leave off on - // 2) What value did we leave off - // Along with the constraints that we get from how we drain the documents: - // Let mean that the last item we drained was item x from partition y. - // Then we know that for all partitions - // * < y that we have drained all items <= x - // * > y that we have drained all items < x - // * = y that we have drained all items <= x based on the backend continuation token for y - // With this information we have captured the progress for all partitions in a single continuation token. - get - { - IEnumerable activeItemProducers = this.GetActiveItemProducers(); - string continuationToken; - if (activeItemProducers.Any()) - { - IEnumerable orderByContinuationTokens = activeItemProducers.Select((itemProducer) => - { - OrderByQueryResult orderByQueryResult = new OrderByQueryResult(itemProducer.Current); - string filter = itemProducer.Filter; - OrderByContinuationToken orderByContinuationToken = new OrderByContinuationToken( - new CompositeContinuationToken - { - Token = itemProducer.PreviousContinuationToken, - Range = itemProducer.PartitionKeyRange.ToRange(), - }, - orderByQueryResult.OrderByItems, - orderByQueryResult.Rid, - this.ShouldIncrementSkipCount(itemProducer) ? this.skipCount + 1 : 0, - filter); - - return OrderByContinuationToken.ToCosmosElement(orderByContinuationToken); - }); - - continuationToken = CosmosArray.Create(orderByContinuationTokens).ToString(); - } - else - { - continuationToken = null; - } - - // Note we are no longer escaping non ascii continuation tokens. - // It is the callers job to encode a continuation token before adding it to a header in their service. - - return continuationToken; - } - } - - public override CosmosElement GetCosmosElementContinuationToken() - { - IEnumerable activeItemProducers = this.GetActiveItemProducers(); - if (!activeItemProducers.Any()) - { - return default; - } - - List orderByContinuationTokens = new List(); - foreach (ItemProducer activeItemProducer in activeItemProducers) - { - OrderByQueryResult orderByQueryResult = new OrderByQueryResult(activeItemProducer.Current); - OrderByContinuationToken orderByContinuationToken = new OrderByContinuationToken( - compositeContinuationToken: new CompositeContinuationToken() - { - Token = activeItemProducer.PreviousContinuationToken, - Range = new Documents.Routing.Range( - min: activeItemProducer.PartitionKeyRange.MinInclusive, - max: activeItemProducer.PartitionKeyRange.MaxExclusive, - isMinInclusive: true, - isMaxInclusive: false) - }, - orderByItems: orderByQueryResult.OrderByItems, - rid: orderByQueryResult.Rid, - skipCount: this.ShouldIncrementSkipCount(activeItemProducer) ? this.skipCount + 1 : 0, - filter: activeItemProducer.Filter); - - CosmosElement cosmosElementToken = OrderByContinuationToken.ToCosmosElement(orderByContinuationToken); - orderByContinuationTokens.Add(cosmosElementToken); - } - - return CosmosArray.Create(orderByContinuationTokens); - } - - /// - /// Equality comparer used to determine if a document producer needs it's continuation token returned. - /// Basically just says that the continuation token can be flushed once you stop seeing duplicates. - /// - private sealed class OrderByEqualityComparer : IEqualityComparer - { - /// - /// The order by comparer. - /// - private readonly OrderByItemProducerTreeComparer orderByConsumeComparer; - - /// - /// Initializes a new instance of the OrderByEqualityComparer class. - /// - /// The order by consume comparer. - public OrderByEqualityComparer(OrderByItemProducerTreeComparer orderByConsumeComparer) - { - this.orderByConsumeComparer = orderByConsumeComparer ?? throw new ArgumentNullException($"{nameof(orderByConsumeComparer)} can not be null."); - } - - /// - /// Gets whether two OrderByQueryResult instances are equal. - /// - /// The first. - /// The second. - /// Whether two OrderByQueryResult instances are equal. - public bool Equals(CosmosElement x, CosmosElement y) - { - OrderByQueryResult orderByQueryResultX = new OrderByQueryResult(x); - OrderByQueryResult orderByQueryResultY = new OrderByQueryResult(y); - return this.orderByConsumeComparer.CompareOrderByItems( - orderByQueryResultX.OrderByItems, - orderByQueryResultY.OrderByItems) == 0; - } - - /// - /// Gets the hash code for object. - /// - /// The object to hash. - /// The hash code for the OrderByQueryResult object. - public int GetHashCode(CosmosElement obj) - { - return 0; - } - } - } -} diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/OrderBy/CosmosOrderByItemQueryExecutionContext.Drain.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/OrderBy/CosmosOrderByItemQueryExecutionContext.Drain.cs deleted file mode 100644 index 0c3e3b1ff2..0000000000 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/OrderBy/CosmosOrderByItemQueryExecutionContext.Drain.cs +++ /dev/null @@ -1,144 +0,0 @@ -// ------------------------------------------------------------ -// Copyright (c) Microsoft Corporation. All rights reserved. -// ------------------------------------------------------------ - -namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.OrderBy -{ - using System; - using System.Collections.Generic; - using System.Threading; - using System.Threading.Tasks; - using Microsoft.Azure.Cosmos.CosmosElements; - using Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.ItemProducers; - using Microsoft.Azure.Cosmos.Query.Core.QueryClient; - - internal sealed partial class CosmosOrderByItemQueryExecutionContext - { - /// - /// Drains a page of documents from this context. - /// - /// The maximum number of elements. - /// The cancellation token. - /// A task that when awaited on return a page of documents. - public override async Task DrainAsync(int maxElements, CancellationToken cancellationToken) - { - cancellationToken.ThrowIfCancellationRequested(); - - //// In order to maintain the continuation token for the user we must drain with a few constraints - //// 1) We always drain from the partition, which has the highest priority item first - //// 2) If multiple partitions have the same priority item then we drain from the left most first - //// otherwise we would need to keep track of how many of each item we drained from each partition - //// (just like parallel queries). - //// Visually that look the following case where we have three partitions that are numbered and store letters. - //// For teaching purposes I have made each item a tuple of the following form: - //// - //// So that duplicates across partitions are distinct, but duplicates within partitions are indistinguishable. - //// |-------| |-------| |-------| - //// | | | | | | - //// | | | | | | - //// | | | | | | - //// | | | | | | - //// | | | | | | - //// | | | | | | - //// | | | | | | - //// |-------| |-------| |-------| - //// Now the correct drain order in this case is: - //// ,,,,,,,,,,, - //// ,,,,,,,,, - //// In more mathematical terms - //// 1) always comes before where x < z - //// 2) always come before where j < k - - List results = new List(); - while (results.Count < maxElements) - { - // Only drain from the highest priority document producer - // We need to pop and push back the document producer tree, since the priority changes according to the sort order. - ItemProducerTree currentItemProducerTree = this.PopCurrentItemProducerTree(); - try - { - if (!currentItemProducerTree.HasMoreResults) - { - // This means there are no more items to drain - break; - } - - OrderByQueryResult orderByQueryResult = new OrderByQueryResult(currentItemProducerTree.Current); - - // Only add the payload, since other stuff is garbage from the caller's perspective. - results.Add(orderByQueryResult.Payload); - - // If we are at the beginning of the page and seeing an rid from the previous page we should increment the skip count - // due to the fact that JOINs can make a document appear multiple times and across continuations, so we don't want to - // surface this more than needed. More information can be found in the continuation token docs. - if (this.ShouldIncrementSkipCount(currentItemProducerTree.CurrentItemProducerTree.Root)) - { - ++this.skipCount; - } - else - { - this.skipCount = 0; - } - - this.previousRid = orderByQueryResult.Rid; - this.previousOrderByItems = orderByQueryResult.OrderByItems; - - if (!currentItemProducerTree.TryMoveNextDocumentWithinPage()) - { - while (true) - { - (bool movedToNextPage, QueryResponseCore? failureResponse) = await currentItemProducerTree.TryMoveNextPageAsync(cancellationToken); - if (!movedToNextPage) - { - if (failureResponse.HasValue) - { - // TODO: We can buffer this failure so that the user can still get the pages we already got. - return failureResponse.Value; - } - - break; - } - - if (currentItemProducerTree.IsAtBeginningOfPage) - { - break; - } - - if (currentItemProducerTree.TryMoveNextDocumentWithinPage()) - { - break; - } - } - } - } - finally - { - this.PushCurrentItemProducerTree(currentItemProducerTree); - } - } - - return QueryResponseCore.CreateSuccess( - result: results, - requestCharge: this.requestChargeTracker.GetAndResetCharge(), - activityId: null, - responseLengthBytes: this.GetAndResetResponseLengthBytes(), - disallowContinuationTokenMessage: null, - continuationToken: this.ContinuationToken); - } - - /// - /// Gets whether or not we should increment the skip count based on the rid of the document. - /// - /// The current document producer. - /// Whether or not we should increment the skip count. - private bool ShouldIncrementSkipCount(ItemProducer currentItemProducer) - { - // If we are not at the beginning of the page and we saw the same rid again. - return !currentItemProducer.IsAtBeginningOfPage && - string.Equals( - this.previousRid, - new OrderByQueryResult(currentItemProducer.Current).Rid, - StringComparison.Ordinal); - } - } -} diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/OrderBy/CosmosOrderByItemQueryExecutionContext.Resume.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/OrderBy/CosmosOrderByItemQueryExecutionContext.Resume.cs deleted file mode 100644 index f099b268da..0000000000 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/OrderBy/CosmosOrderByItemQueryExecutionContext.Resume.cs +++ /dev/null @@ -1,69 +0,0 @@ -//------------------------------------------------------------ -// Copyright (c) Microsoft Corporation. All rights reserved. -//------------------------------------------------------------ -namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.OrderBy -{ - using System; - using System.Collections.Generic; - using System.Diagnostics; - using System.Linq; - using System.Text; - using System.Threading; - using System.Threading.Tasks; - using Microsoft.Azure.Cosmos.CosmosElements; - using Microsoft.Azure.Cosmos.Query.Core; - using Microsoft.Azure.Cosmos.Query.Core.ContinuationTokens; - using Microsoft.Azure.Cosmos.Query.Core.Exceptions; - using Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent; - using Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.ItemProducers; - using Microsoft.Azure.Cosmos.Query.Core.Monads; - using Microsoft.Azure.Cosmos.Query.Core.QueryClient; - using static Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote.OrderByCrossPartitionQueryPipelineStage; - using PartitionKeyRange = Documents.PartitionKeyRange; - using ResourceId = Documents.ResourceId; - - internal sealed partial class CosmosOrderByItemQueryExecutionContext : CosmosCrossPartitionQueryExecutionContext - { - public static async Task> MonadicCreateAsync( - CosmosQueryContext queryContext, - CosmosCrossPartitionQueryExecutionContext.CrossPartitionInitParams initParams, - CosmosElement requestContinuationToken, - CancellationToken cancellationToken) - { - Debug.Assert( - initParams.PartitionedQueryExecutionInfo.QueryInfo.HasOrderBy, - "OrderBy~Context must have order by query info."); - - if (queryContext == null) - { - throw new ArgumentNullException(nameof(queryContext)); - } - - cancellationToken.ThrowIfCancellationRequested(); - - OrderByItemProducerTreeComparer orderByItemProducerTreeComparer = new OrderByItemProducerTreeComparer(initParams.PartitionedQueryExecutionInfo.QueryInfo.OrderBy.ToArray()); - CosmosOrderByItemQueryExecutionContext context = new CosmosOrderByItemQueryExecutionContext( - initPararms: queryContext, - maxConcurrency: initParams.MaxConcurrency, - maxItemCount: initParams.MaxItemCount, - maxBufferedItemCount: initParams.MaxBufferedItemCount, - consumeComparer: orderByItemProducerTreeComparer, - testSettings: initParams.TestSettings); - - IReadOnlyList orderByExpressions = initParams.PartitionedQueryExecutionInfo.QueryInfo.OrderByExpressions; - IReadOnlyList sortOrders = initParams.PartitionedQueryExecutionInfo.QueryInfo.OrderBy; - if (orderByExpressions.Count != sortOrders.Count) - { - throw new ArgumentException("order by expressions count does not match sort order"); - } - - IReadOnlyList columns = orderByExpressions - .Zip(sortOrders, (expression, order) => new OrderByColumn(expression, order)) - .ToList(); - - await Task.Delay(0); - - return default; - } - } -} diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/OrderBy/CosmosOrderByItemQueryExecutionContext.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/OrderBy/CosmosOrderByItemQueryExecutionContext.cs deleted file mode 100644 index aa542a1e19..0000000000 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/OrderBy/CosmosOrderByItemQueryExecutionContext.cs +++ /dev/null @@ -1,71 +0,0 @@ -// ------------------------------------------------------------ -// Copyright (c) Microsoft Corporation. All rights reserved. -// ------------------------------------------------------------ - -namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.OrderBy -{ - using System; - using System.Collections.Generic; - using Microsoft.Azure.Cosmos.Query.Core; - using Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.ItemProducers; - using Microsoft.Azure.Cosmos.Query.Core.QueryClient; - - internal sealed partial class CosmosOrderByItemQueryExecutionContext - { - /// - /// Function to determine the priority of fetches. - /// Basically we are fetching from the partition with the least number of buffered documents first. - /// - private static readonly Func FetchPriorityFunction = itemProducerTree => itemProducerTree.BufferedItemCount; - - /// - /// Skip count used for JOIN queries. - /// You can read up more about this in the documentation for the continuation token. - /// - private int skipCount; - - /// - /// We need to keep track of the previousRid, since order by queries don't drain full pages. - /// - private string previousRid; - - private IReadOnlyList previousOrderByItems; - - /// - /// Initializes a new instance of the CosmosOrderByItemQueryExecutionContext class. - /// - /// The params used to construct the base class. - /// For cross partition order by queries a query like "SELECT c.id, c.field_0 ORDER BY r.field_7 gets rewritten as: - /// - /// This is needed because we need to add additional filters to the query when we resume from a continuation, - /// and it lets us easily parse out the _rid orderByItems, and payload without parsing the entire document (and having to remember the order by field). - /// The max concurrency - /// The max buffered item count - /// Max item count - /// Comparer used to internally compare documents from different sorted partitions. - /// Test settings. - private CosmosOrderByItemQueryExecutionContext( - CosmosQueryContext initPararms, - int? maxConcurrency, - int? maxItemCount, - int? maxBufferedItemCount, - OrderByItemProducerTreeComparer consumeComparer, - TestInjections testSettings) - : base( - queryContext: initPararms, - maxConcurrency: maxConcurrency, - maxItemCount: maxItemCount, - maxBufferedItemCount: maxBufferedItemCount, - moveNextComparer: consumeComparer, - fetchPrioirtyFunction: CosmosOrderByItemQueryExecutionContext.FetchPriorityFunction, - equalityComparer: new OrderByEqualityComparer(consumeComparer), - returnResultsInDeterministicOrder: true, - testSettings: testSettings) - { - } - } -} diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/OrderBy/OrderByItemProducerTreeComparer.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/OrderBy/OrderByItemProducerTreeComparer.cs deleted file mode 100644 index 0dc4977d8c..0000000000 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/OrderBy/OrderByItemProducerTreeComparer.cs +++ /dev/null @@ -1,142 +0,0 @@ -//------------------------------------------------------------ -// Copyright (c) Microsoft Corporation. All rights reserved. -//------------------------------------------------------------ -namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.OrderBy -{ - using System; - using System.Collections.Generic; - using System.Diagnostics; - using Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.ItemProducers; - - /// - /// For cross partition order by queries we serve documents from the partition - /// that has the next document in the sort order of the query. - /// If there is a tie, then we break the tie by picking the leftmost partition. - /// - internal sealed class OrderByItemProducerTreeComparer : IComparer - { - /// - /// The sort orders for the query (1 for each order by in the query). - /// Until composite indexing is released this will just be an array of length 1. - /// - private readonly IReadOnlyList sortOrders; - - /// - /// Initializes a new instance of the OrderByConsumeComparer class. - /// - /// The sort orders for the query. - public OrderByItemProducerTreeComparer(SortOrder[] sortOrders) - { - if (sortOrders == null) - { - throw new ArgumentNullException("Sort Orders array can not be null for an order by comparer."); - } - - if (sortOrders.Length == 0) - { - throw new ArgumentException("Sort Orders array can not be empty for an order by comparerer."); - } - - this.sortOrders = new List(sortOrders); - } - - /// - /// Compares two document producer trees and returns an integer with the relation of which has the document that comes first in the sort order. - /// - /// The first document producer tree. - /// The second document producer tree. - /// - /// Less than zero if the document in the first document producer comes first. - /// Zero if the documents are equivalent. - /// Greater than zero if the document in the second document producer comes first. - /// - public int Compare(ItemProducerTree producer1, ItemProducerTree producer2) - { - if (object.ReferenceEquals(producer1, producer2)) - { - return 0; - } - - if (producer1.HasMoreResults && !producer2.HasMoreResults) - { - return -1; - } - - if (!producer1.HasMoreResults && producer2.HasMoreResults) - { - return 1; - } - - if (!producer1.HasMoreResults && !producer2.HasMoreResults) - { - return string.CompareOrdinal(producer1.PartitionKeyRange.MinInclusive, producer2.PartitionKeyRange.MinInclusive); - } - - OrderByQueryResult result1 = new OrderByQueryResult(producer1.Current); - OrderByQueryResult result2 = new OrderByQueryResult(producer2.Current); - - // First compare the documents based on the sort order of the query. - int cmp = this.CompareOrderByItems(result1.OrderByItems, result2.OrderByItems); - if (cmp != 0) - { - // If there is no tie just return that. - return cmp; - } - - // If there is a tie, then break the tie by picking the one from the left most partition. - return string.CompareOrdinal(producer1.PartitionKeyRange.MinInclusive, producer2.PartitionKeyRange.MinInclusive); - - } - - /// - /// Takes the items relevant to the sort and return an integer defining the relationship. - /// - /// The items relevant to the sort from the first partition. - /// The items relevant to the sort from the second partition. - /// The sort relationship. - /// - /// Suppose the query was "SELECT * FROM c ORDER BY c.name asc, c.age desc", - /// then items1 could be ["Brandon", 22] and items2 could be ["Felix", 28] - /// Then we would first compare "Brandon" to "Felix" and say that "Brandon" comes first in an ascending lex order (we don't even have to look at age). - /// If items1 was ["Brandon", 22] and items2 was ["Brandon", 23] then we would say have to look at the age to break the tie and in this case 23 comes first in a descending order. - /// Some examples of composite order by: http://www.dofactory.com/sql/order-by - /// - public int CompareOrderByItems(IReadOnlyList items1, IReadOnlyList items2) - { - if (object.ReferenceEquals(items1, items2)) - { - return 0; - } - - Debug.Assert( - items1 != null && items2 != null, - "Order-by items must be present."); - - Debug.Assert( - items1.Count == items2.Count, - "OrderByResult instances should have the same number of order-by items."); - - Debug.Assert( - items1.Count > 0, - "OrderByResult instances should have at least 1 order-by item."); - - Debug.Assert( - this.sortOrders.Count == items1.Count, - "SortOrders must match size of order-by items."); - - for (int i = 0; i < this.sortOrders.Count; ++i) - { - int cmp = ItemComparer.Instance.Compare( - items1[i].Item, - items2[i].Item); - - if (cmp != 0) - { - return this.sortOrders[i] != SortOrder.Descending ? cmp : -cmp; - } - } - - return 0; - } - } -} \ No newline at end of file diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/OrderBy/RangeFilterInitializationInfo.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/OrderBy/RangeFilterInitializationInfo.cs deleted file mode 100644 index b68609f985..0000000000 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/OrderBy/RangeFilterInitializationInfo.cs +++ /dev/null @@ -1,91 +0,0 @@ -//------------------------------------------------------------ -// Copyright (c) Microsoft Corporation. All rights reserved. -//------------------------------------------------------------ -namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.OrderBy -{ - using System; - - /// - /// - /// InitializationInfo is a data structure to capture how the DocumentProducers are initialized - /// once we start a cross-partition OrderBy query execution from a continuation token. - /// - /// - /// Specifically, the data-structure captures the "filter" condition for all the partitions that - /// need to be visited as a part of the query. Please see the description of "filters" in the - /// OrderByContinuationToken class. - /// - /// - internal readonly struct RangeFilterInitializationInfo - { - /// - /// Initializes a new instance of the RangeFilterInitializationInfo struct. - /// - /// The filter to apply to the partitions. - /// The start index of the partitions. - /// The end index of the partitions. - public RangeFilterInitializationInfo( - string filter, - int startIndex, - int endIndex) - { - this.Filter = filter ?? throw new ArgumentNullException("filter"); - this.StartIndex = startIndex; - this.EndIndex = endIndex; - } - - /// - /// Gets the filter itself. - /// - /// - /// - /// For an order by query "select * from root order by root.key ASC", a filter string could be "root.key > 2", - /// considering "key" is an integer field. - /// - /// - /// The filter simply indicates that the order by query have already delivered all the root.key with "less than equal to 2". - /// - /// - public string Filter - { - get; - } - - /// - /// Gets the start index. - /// - /// Assuming that, at the beginning of the query execution, all the partitions, that needs to be visited, are ordered from 0 to n, - /// startIndex refers to the starting point of the contiguous block of partitions that requires the "filter" to be applied. - /// - /// - /// Typically, there would three such blocks - /// * one before the target range (Please study the OrderByContinuationToken class to understand what a target range is) - /// * the target range itself - /// * one after the target range. - /// Each of these block may have different filter conditions. - /// For example, if - /// (1) the target range has filter condition "root.key >= 2", then - /// (2) the preceding block will have condition "root.key > 2" and - /// (3) the succeeding block will have filter condition "root.key >= 2". - /// - /// - /// However, there could be more than one target ranges, in case of query execution across split, each leading to one more - /// blocks (typically containing one partition). - /// - /// - public int StartIndex - { - get; - } - - /// - /// Gets the end index. - /// Assuming that, at the beginning of the query execution, all the partitions, that needs to be visited, are ordered from 0 to n, - /// EndIndex refers to the end point of the contiguous block of partitions that requires the "filter" to be applied. - /// - public int EndIndex - { - get; - } - } -} diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/Parallel/CosmosParallelItemQueryExecutionContext.ContinuationToken.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/Parallel/CosmosParallelItemQueryExecutionContext.ContinuationToken.cs deleted file mode 100644 index cd34477f7f..0000000000 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/Parallel/CosmosParallelItemQueryExecutionContext.ContinuationToken.cs +++ /dev/null @@ -1,111 +0,0 @@ -//------------------------------------------------------------ -// Copyright (c) Microsoft Corporation. All rights reserved. -//------------------------------------------------------------ -namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.Parallel -{ - using System.Collections.Generic; - using System.Linq; - using Microsoft.Azure.Cosmos.CosmosElements; - using Microsoft.Azure.Cosmos.Query.Core.ContinuationTokens; - using Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.ItemProducers; - using Newtonsoft.Json; - - /// - /// CosmosParallelItemQueryExecutionContext is a concrete implementation for CrossPartitionQueryExecutionContext. - /// This class is responsible for draining cross partition queries that do not have order by conditions. - /// The way parallel queries work is that it drains from the left most partition first. - /// This class handles draining in the correct order and can also stop and resume the query - /// by generating a continuation token and resuming from said continuation token. - /// - internal sealed partial class CosmosParallelItemQueryExecutionContext : CosmosCrossPartitionQueryExecutionContext - { - /// - /// For parallel queries the continuation token semantically holds two pieces of information: - /// 1) What physical partition did the user read up to - /// 2) How far into said partition did they read up to - /// And since the client consumes queries strictly in a left to right order we can partition the documents: - /// 1) Documents left of the continuation token have been drained - /// 2) Documents to the right of the continuation token still need to be served. - /// This is useful since we can have a single continuation token for all partitions. - /// - protected override string ContinuationToken - { - get - { - IEnumerable activeItemProducers = this.GetActiveItemProducers(); - string continuationToken; - if (activeItemProducers.Any()) - { - IEnumerable compositeContinuationTokens = activeItemProducers.Select((documentProducer) => new CompositeContinuationToken - { - Token = documentProducer.CurrentContinuationToken, - Range = documentProducer.PartitionKeyRange.ToRange() - }); - continuationToken = JsonConvert.SerializeObject(compositeContinuationTokens, DefaultJsonSerializationSettings.Value); - } - else - { - continuationToken = null; - } - - return continuationToken; - } - } - - public override CosmosElement GetCosmosElementContinuationToken() - { - IEnumerable activeItemProducers = this.GetActiveItemProducers(); - if (!activeItemProducers.Any()) - { - return default; - } - - List compositeContinuationTokens = new List(); - foreach (ItemProducer activeItemProducer in activeItemProducers) - { - CompositeContinuationToken compositeToken = new CompositeContinuationToken() - { - Token = activeItemProducer.CurrentContinuationToken, - Range = new Documents.Routing.Range( - min: activeItemProducer.PartitionKeyRange.MinInclusive, - max: activeItemProducer.PartitionKeyRange.MaxExclusive, - isMinInclusive: true, - isMaxInclusive: false) - }; - - CosmosElement compositeContinuationToken = CompositeContinuationToken.ToCosmosElement(compositeToken); - compositeContinuationTokens.Add(compositeContinuationToken); - } - - return CosmosArray.Create(compositeContinuationTokens); - } - - /// - /// Comparer used to determine if we should return the continuation token to the user - /// - /// This basically just says that the two object are never equals, so that we don't return a continuation for a partition we have started draining. - private sealed class ParallelEqualityComparer : IEqualityComparer - { - /// - /// Returns whether two parallel query items are equal. - /// - /// The first item. - /// The second item. - /// Whether two parallel query items are equal. - public bool Equals(CosmosElement x, CosmosElement y) - { - return object.ReferenceEquals(x, y); - } - - /// - /// Gets the hash code of an object. - /// - /// The object to hash. - /// The hash code for the object. - public int GetHashCode(CosmosElement obj) - { - return obj == null ? 0 : obj.GetHashCode(); - } - } - } -} diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/Parallel/CosmosParallelItemQueryExecutionContext.Drain.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/Parallel/CosmosParallelItemQueryExecutionContext.Drain.cs deleted file mode 100644 index 63ea186e99..0000000000 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/Parallel/CosmosParallelItemQueryExecutionContext.Drain.cs +++ /dev/null @@ -1,71 +0,0 @@ -//------------------------------------------------------------ -// Copyright (c) Microsoft Corporation. All rights reserved. -//------------------------------------------------------------ -namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.Parallel -{ - using System; - using System.Collections.Generic; - using System.Threading; - using System.Threading.Tasks; - using Microsoft.Azure.Cosmos.CosmosElements; - using Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.ItemProducers; - using Microsoft.Azure.Cosmos.Query.Core.QueryClient; - - /// - /// CosmosParallelItemQueryExecutionContext is a concrete implementation for CrossPartitionQueryExecutionContext. - /// This class is responsible for draining cross partition queries that do not have order by conditions. - /// The way parallel queries work is that it drains from the left most partition first. - /// This class handles draining in the correct order and can also stop and resume the query - /// by generating a continuation token and resuming from said continuation token. - /// - internal sealed partial class CosmosParallelItemQueryExecutionContext : CosmosCrossPartitionQueryExecutionContext - { - public override async Task DrainAsync(int maxElements, CancellationToken cancellationToken) - { - cancellationToken.ThrowIfCancellationRequested(); - - // In order to maintain the continuation token for the user we must drain with a few constraints - // 1) We fully drain from the left most partition before moving on to the next partition - // 2) We drain only full pages from the document producer so we aren't left with a partial page - // otherwise we would need to add to the continuation token how many items to skip over on that page. - - // Only drain from the leftmost (current) document producer tree - ItemProducerTree currentItemProducerTree = this.PopCurrentItemProducerTree(); - List results = new List(); - try - { - (bool gotNextPage, QueryResponseCore? failureResponse) = await currentItemProducerTree.TryMoveNextPageAsync(cancellationToken); - if (failureResponse != null) - { - return failureResponse.Value; - } - - if (gotNextPage) - { - int itemsLeftInCurrentPage = currentItemProducerTree.ItemsLeftInCurrentPage; - - // Only drain full pages or less if this is a top query. - currentItemProducerTree.TryMoveNextDocumentWithinPage(); - int numberOfItemsToDrain = Math.Min(itemsLeftInCurrentPage, maxElements); - for (int i = 0; i < numberOfItemsToDrain; i++) - { - results.Add(currentItemProducerTree.Current); - currentItemProducerTree.TryMoveNextDocumentWithinPage(); - } - } - } - finally - { - this.PushCurrentItemProducerTree(currentItemProducerTree); - } - - return QueryResponseCore.CreateSuccess( - result: results, - requestCharge: this.requestChargeTracker.GetAndResetCharge(), - activityId: null, - responseLengthBytes: this.GetAndResetResponseLengthBytes(), - disallowContinuationTokenMessage: null, - continuationToken: this.ContinuationToken); - } - } -} diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/Parallel/CosmosParallelItemQueryExecutionContext.Resume.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/Parallel/CosmosParallelItemQueryExecutionContext.Resume.cs deleted file mode 100644 index 19670c41b7..0000000000 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/Parallel/CosmosParallelItemQueryExecutionContext.Resume.cs +++ /dev/null @@ -1,180 +0,0 @@ -// ------------------------------------------------------------ -// Copyright (c) Microsoft Corporation. All rights reserved. -// ------------------------------------------------------------ - -namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.Parallel -{ - using System; - using System.Collections.Generic; - using System.Linq; - using System.Threading; - using System.Threading.Tasks; - using Microsoft.Azure.Cosmos.CosmosElements; - using Microsoft.Azure.Cosmos.Query.Core.ContinuationTokens; - using Microsoft.Azure.Cosmos.Query.Core.Exceptions; - using Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent; - using Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.ItemProducers; - using Microsoft.Azure.Cosmos.Query.Core.Monads; - using Microsoft.Azure.Cosmos.Query.Core.QueryClient; - using static Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote.PartitionMapper; - using PartitionKeyRange = Documents.PartitionKeyRange; - - internal sealed partial class CosmosParallelItemQueryExecutionContext : CosmosCrossPartitionQueryExecutionContext - { - public static async Task> MonadicCreateAsync( - CosmosQueryContext queryContext, - CosmosCrossPartitionQueryExecutionContext.CrossPartitionInitParams initParams, - CosmosElement requestContinuationToken, - CancellationToken cancellationToken) - { - if (queryContext == null) - { - throw new ArgumentNullException(nameof(queryContext)); - } - - cancellationToken.ThrowIfCancellationRequested(); - - IComparer moveNextComparer; - if (initParams.ReturnResultsInDeterministicOrder) - { - moveNextComparer = DeterministicParallelItemProducerTreeComparer.Singleton; - } - else - { - moveNextComparer = NonDeterministicParallelItemProducerTreeComparer.Singleton; - } - - CosmosParallelItemQueryExecutionContext context = new CosmosParallelItemQueryExecutionContext( - queryContext: queryContext, - maxConcurrency: initParams.MaxConcurrency, - maxItemCount: initParams.MaxItemCount, - maxBufferedItemCount: initParams.MaxBufferedItemCount, - moveNextComparer: moveNextComparer, - returnResultsInDeterministicOrder: initParams.ReturnResultsInDeterministicOrder, - testSettings: initParams.TestSettings); - - return (await context.TryInitializeAsync( - sqlQuerySpec: initParams.SqlQuerySpec, - collectionRid: initParams.CollectionRid, - partitionKeyRanges: initParams.PartitionKeyRanges, - initialPageSize: initParams.InitialPageSize, - requestContinuation: requestContinuationToken, - cancellationToken: cancellationToken)).Try(x => x); - } - - /// - /// Initialize the execution context. - /// - /// SQL query spec. - /// The collection rid. - /// The partition key ranges to drain documents from. - /// The initial page size. - /// The continuation token to resume from. - /// The cancellation token. - /// A task to await on. - private async Task> TryInitializeAsync( - SqlQuerySpec sqlQuerySpec, - string collectionRid, - IReadOnlyList partitionKeyRanges, - int initialPageSize, - CosmosElement requestContinuation, - CancellationToken cancellationToken) - { - cancellationToken.ThrowIfCancellationRequested(); - - TryCatch> tryGetPartitionKeyRangeToCompositeContinuationToken = CosmosParallelItemQueryExecutionContext.TryGetPartitionKeyRangeToCompositeContinuationToken( - partitionKeyRanges, - requestContinuation); - if (!tryGetPartitionKeyRangeToCompositeContinuationToken.Succeeded) - { - return TryCatch.FromException(tryGetPartitionKeyRangeToCompositeContinuationToken.Exception); - } - - List> rangesToInitialize = new List>() - { - // Skip all the partitions left of the target range, since they have already been drained fully. - tryGetPartitionKeyRangeToCompositeContinuationToken.Result.TargetPartition, - tryGetPartitionKeyRangeToCompositeContinuationToken.Result.PartitionsRightOfTarget, - }; - - foreach (IReadOnlyDictionary rangeToCompositeToken in rangesToInitialize) - { - IReadOnlyDictionary partitionKeyRangeToContinuationToken = rangeToCompositeToken - .ToDictionary(kvp => kvp.Key, kvp => kvp.Value?.Token); - TryCatch tryInitialize = await base.TryInitializeAsync( - collectionRid, - initialPageSize, - sqlQuerySpec, - partitionKeyRangeToContinuationToken, - deferFirstPage: true, - filter: null, - tryFilterAsync: null, - cancellationToken); - if (!tryInitialize.Succeeded) - { - return TryCatch.FromException(tryInitialize.Exception); - } - } - - return TryCatch.FromResult(this); - } - - private static TryCatch> TryGetPartitionKeyRangeToCompositeContinuationToken( - IReadOnlyList partitionKeyRanges, - CosmosElement continuationToken) - { - if (continuationToken == null) - { - Dictionary dictionary = new Dictionary(); - foreach (PartitionKeyRange partitionKeyRange in partitionKeyRanges) - { - dictionary.Add(key: partitionKeyRange, value: null); - } - - return TryCatch>.FromResult( - new PartitionMapping( - partitionsLeftOfTarget: new Dictionary(), - targetPartition: dictionary, - partitionsRightOfTarget: new Dictionary())); - } - - TryCatch> tryParseCompositeContinuationTokens = TryParseCompositeContinuationList(continuationToken); - if (!tryParseCompositeContinuationTokens.Succeeded) - { - return TryCatch>.FromException(tryParseCompositeContinuationTokens.Exception); - } - - return default; - } - - private static TryCatch> TryParseCompositeContinuationList( - CosmosElement requestContinuationToken) - { - if (requestContinuationToken == null) - { - throw new ArgumentNullException(nameof(requestContinuationToken)); - } - - if (!(requestContinuationToken is CosmosArray compositeContinuationTokenListRaw)) - { - return TryCatch>.FromException( - new MalformedContinuationTokenException( - $"Invalid format for continuation token {requestContinuationToken} for {nameof(CosmosParallelItemQueryExecutionContext)}")); - } - - List compositeContinuationTokens = new List(); - foreach (CosmosElement compositeContinuationTokenRaw in compositeContinuationTokenListRaw) - { - TryCatch tryCreateCompositeContinuationToken = CompositeContinuationToken.TryCreateFromCosmosElement(compositeContinuationTokenRaw); - if (!tryCreateCompositeContinuationToken.Succeeded) - { - return TryCatch>.FromException(tryCreateCompositeContinuationToken.Exception); - } - - compositeContinuationTokens.Add(tryCreateCompositeContinuationToken.Result); - } - - return TryCatch>.FromResult(compositeContinuationTokens); - } - } -} diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/Parallel/CosmosParallelItemQueryExecutionContext.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/Parallel/CosmosParallelItemQueryExecutionContext.cs deleted file mode 100644 index bca7f6fe39..0000000000 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/Parallel/CosmosParallelItemQueryExecutionContext.cs +++ /dev/null @@ -1,58 +0,0 @@ -//------------------------------------------------------------ -// Copyright (c) Microsoft Corporation. All rights reserved. -//------------------------------------------------------------ -namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.Parallel -{ - using System; - using System.Collections.Generic; - using Microsoft.Azure.Cosmos.CosmosElements; - using Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.ItemProducers; - using Microsoft.Azure.Cosmos.Query.Core.QueryClient; - - internal sealed partial class CosmosParallelItemQueryExecutionContext : CosmosCrossPartitionQueryExecutionContext - { - /// - /// The comparer used to determine, which continuation tokens should be returned to the user. - /// - private static readonly IEqualityComparer EqualityComparer = new ParallelEqualityComparer(); - - /// - /// The function to determine which partition to fetch from first. - /// - private static readonly Func FetchPriorityFunction = documentProducerTree => int.Parse(documentProducerTree.PartitionKeyRange.Id); - - private readonly bool returnResultsInDeterministicOrder; - - /// - /// Initializes a new instance of the CosmosParallelItemQueryExecutionContext class. - /// - /// The parameters for constructing the base class. - /// The max concurrency - /// The max buffered item count - /// Max item count - /// The comparer to use for the priority queue. - /// Whether or not to return results in deterministic order. - /// Test settings. - private CosmosParallelItemQueryExecutionContext( - CosmosQueryContext queryContext, - int? maxConcurrency, - int? maxItemCount, - int? maxBufferedItemCount, - IComparer moveNextComparer, - bool returnResultsInDeterministicOrder, - TestInjections testSettings) - : base( - queryContext: queryContext, - maxConcurrency: maxConcurrency, - maxItemCount: maxItemCount, - maxBufferedItemCount: maxBufferedItemCount, - moveNextComparer: moveNextComparer, - fetchPrioirtyFunction: CosmosParallelItemQueryExecutionContext.FetchPriorityFunction, - equalityComparer: CosmosParallelItemQueryExecutionContext.EqualityComparer, - returnResultsInDeterministicOrder: returnResultsInDeterministicOrder, - testSettings: testSettings) - { - this.returnResultsInDeterministicOrder = returnResultsInDeterministicOrder; - } - } -} diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/Parallel/DeterministicParallelItemProducerTreeComparer.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/Parallel/DeterministicParallelItemProducerTreeComparer.cs deleted file mode 100644 index 232e4e8938..0000000000 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/Parallel/DeterministicParallelItemProducerTreeComparer.cs +++ /dev/null @@ -1,62 +0,0 @@ -//------------------------------------------------------------ -// Copyright (c) Microsoft Corporation. All rights reserved. -//------------------------------------------------------------ -namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.Parallel -{ - using System.Collections.Generic; - using Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.ItemProducers; - using PartitionKeyRange = Documents.PartitionKeyRange; - - /// - /// Implementation of that returns documents from the left partition first. - /// The documents within a partition are already sorted in _rid order. - /// - /// This comparer gaurentees that the query returns results in a deterministic order meaning that running the same query twice returns results in the same order. This also ensures that continuation token returned to the user is as small as possible, since we finish draining a partition before moving onto another partition. The only draw back is that all partitions can not be drained until all their partitions to the left are fully drained, which reduces the parallelism. - /// - internal sealed class DeterministicParallelItemProducerTreeComparer : IComparer - { - public static readonly DeterministicParallelItemProducerTreeComparer Singleton = new DeterministicParallelItemProducerTreeComparer(); - - private DeterministicParallelItemProducerTreeComparer() - { - // Singleton class, so leave the constructor private. - } - - /// - /// Compares two document producer trees in a parallel context and returns their comparison. - /// - /// The first document producer tree. - /// The second document producer tree. - /// - /// A negative number if the first comes before the second. - /// Zero if the two document producer trees are interchangeable. - /// A positive number if the second comes before the first. - /// - public int Compare( - ItemProducerTree documentProducerTree1, - ItemProducerTree documentProducerTree2) - { - if (object.ReferenceEquals(documentProducerTree1, documentProducerTree2)) - { - return 0; - } - - if (documentProducerTree1.HasMoreResults && !documentProducerTree2.HasMoreResults) - { - return -1; - } - - if (!documentProducerTree1.HasMoreResults && documentProducerTree2.HasMoreResults) - { - return 1; - } - - // Either both don't have results or both do. - PartitionKeyRange partitionKeyRange1 = documentProducerTree1.PartitionKeyRange; - PartitionKeyRange partitionKeyRange2 = documentProducerTree2.PartitionKeyRange; - return string.CompareOrdinal( - partitionKeyRange1.MinInclusive, - partitionKeyRange2.MinInclusive); - } - } -} \ No newline at end of file diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/Parallel/NonDeterministicParallelItemProducerTreeComparer.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/Parallel/NonDeterministicParallelItemProducerTreeComparer.cs deleted file mode 100644 index 632fa72324..0000000000 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/Parallel/NonDeterministicParallelItemProducerTreeComparer.cs +++ /dev/null @@ -1,65 +0,0 @@ -//------------------------------------------------------------ -// Copyright (c) Microsoft Corporation. All rights reserved. -//------------------------------------------------------------ -namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.Parallel -{ - using System; - using System.Collections.Generic; - using Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.ItemProducers; - using PartitionKeyRange = Documents.PartitionKeyRange; - - /// - /// Implementation of that returns documents from the partition that has the most documents buffered first. - /// - internal sealed class NonDeterministicParallelItemProducerTreeComparer : IComparer - { - public static readonly NonDeterministicParallelItemProducerTreeComparer Singleton = new NonDeterministicParallelItemProducerTreeComparer(); - - private NonDeterministicParallelItemProducerTreeComparer() - { - // Singleton class, so leave the constructor private. - } - - /// - /// Compares two document producer trees in a parallel context and returns their comparison. - /// - /// The first document producer tree. - /// The second document producer tree. - /// - /// A negative number if the first comes before the second. - /// Zero if the two document producer trees are interchangeable. - /// A positive number if the second comes before the first. - /// - public int Compare( - ItemProducerTree documentProducerTree1, - ItemProducerTree documentProducerTree2) - { - if (documentProducerTree1 == null) - { - throw new ArgumentNullException(nameof(documentProducerTree1)); - } - - if (documentProducerTree2 == null) - { - throw new ArgumentNullException(nameof(documentProducerTree2)); - } - - if (object.ReferenceEquals(documentProducerTree1, documentProducerTree2)) - { - return 0; - } - - if (documentProducerTree1.HasMoreResults && !documentProducerTree2.HasMoreResults) - { - return -1; - } - - if (!documentProducerTree1.HasMoreResults && documentProducerTree2.HasMoreResults) - { - return 1; - } - - return documentProducerTree2.BufferedItemCount.CompareTo(documentProducerTree1.BufferedItemCount); - } - } -} \ No newline at end of file diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/Parallel/ParallelQueryConfig.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/Parallel/ParallelQueryConfig.cs deleted file mode 100644 index ba717464a4..0000000000 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/Parallel/ParallelQueryConfig.cs +++ /dev/null @@ -1,77 +0,0 @@ -//------------------------------------------------------------ -// Copyright (c) Microsoft Corporation. All rights reserved. -//------------------------------------------------------------ -namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.Parallel -{ - /// - /// This class stores all the default values used during cross partition queries. - /// - internal sealed class ParallelQueryConfig - { - /// - /// If the client specifies a MaxItemCount as -1, we return documents in batches of 100 - /// - public readonly int ClientInternalPageSize; - - /// - /// Default maximum number of documents cached at the client side, if it is not specified in the feedOptions. - /// - public readonly long DefaultMaximumBufferSize; - - /// - /// We adaptively increase the number of threads as we see partitions are continuing to return results. 2 => we double the number of threads. - /// - public readonly int AutoModeTasksIncrementFactor; - - /// - /// This is the value we overwrite with in the above case. -1 => the server returns dynamic number of results. Overwriting -1 by very high number doesn't have any significant impact in performance. -1, tries to return maximum number of documents possible per roundtrip - /// - public readonly int ClientInternalMaxItemCount; - - /// - /// Making a backend call is equivalent to a Network Call. Here 1 indicates that if the client machine has 4 processor, we would allow at max (4*1) = 4 parallel calls to the backend. Of course, the number of parallel call won't exceed the number of partitions that needs to be visited for the query under consideration. - /// - public readonly int NumberOfNetworkCallsPerProcessor; - - /// - /// Singleton instance. - /// - private static readonly ParallelQueryConfig DefaultInstance = new ParallelQueryConfig( - clientInternalMaxItemCount: 100, - defaultMaximumBufferSize: 100, - clientInternalPageSize: 100, - autoModeTasksIncrementFactor: 2, - numberOfNetworkCallsPerProcessor: 1); - - /// - /// Initializes a new instance of the class. - /// - /// The client's internal max item count. - /// The default maximum buffer size. - /// The client's internal page size. - /// The increment factor for auto mode. - /// Number of network calls per processor. - private ParallelQueryConfig( - int clientInternalMaxItemCount, - int defaultMaximumBufferSize, - int clientInternalPageSize, - int autoModeTasksIncrementFactor, - int numberOfNetworkCallsPerProcessor) - { - this.ClientInternalMaxItemCount = clientInternalMaxItemCount; - this.DefaultMaximumBufferSize = defaultMaximumBufferSize; - this.ClientInternalPageSize = clientInternalPageSize; - this.AutoModeTasksIncrementFactor = autoModeTasksIncrementFactor; - this.NumberOfNetworkCallsPerProcessor = numberOfNetworkCallsPerProcessor; - } - - /// - /// Gets the configs for parallel queries. - /// - /// The configs for parallel queries. - public static ParallelQueryConfig GetConfig() - { - return ParallelQueryConfig.DefaultInstance; - } - } -} \ No newline at end of file diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/QueryExecutionContextWithException.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/QueryExecutionContextWithException.cs deleted file mode 100644 index 625acc6bbf..0000000000 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/QueryExecutionContextWithException.cs +++ /dev/null @@ -1,46 +0,0 @@ -// ------------------------------------------------------------ -// Copyright (c) Microsoft Corporation. All rights reserved. -// ------------------------------------------------------------ -namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionContext -{ - using System; - using System.Threading; - using System.Threading.Tasks; - using Microsoft.Azure.Cosmos.CosmosElements; - using Microsoft.Azure.Cosmos.Query.Core; - using Microsoft.Azure.Cosmos.Query.Core.QueryClient; - - internal sealed class QueryExecutionContextWithException : CosmosQueryExecutionContext - { - private readonly Exception exception; - private bool returnedErrorResponse; - - public QueryExecutionContextWithException(Exception exception) - { - if (exception == null) - { - throw new ArgumentNullException(nameof(exception)); - } - - this.exception = exception; - } - - public override bool IsDone => this.returnedErrorResponse; - - public override void Dispose() - { - } - - public override Task ExecuteNextAsync(CancellationToken cancellationToken) - { - QueryResponseCore queryResponse = QueryResponseFactory.CreateFromException(this.exception); - this.returnedErrorResponse = true; - return Task.FromResult(queryResponse); - } - - public override CosmosElement GetCosmosElementContinuationToken() - { - throw new NotImplementedException(); - } - } -} diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Aggregate/Aggregators/MinMaxAggregator.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Aggregate/Aggregators/MinMaxAggregator.cs index 25cce653dd..06e0c7e495 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Aggregate/Aggregators/MinMaxAggregator.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Aggregate/Aggregators/MinMaxAggregator.cs @@ -8,8 +8,8 @@ namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.Aggregate.Aggregators using System.Collections.Generic; using Microsoft.Azure.Cosmos.CosmosElements; using Microsoft.Azure.Cosmos.Query.Core.Exceptions; - using Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.ItemProducers; using Microsoft.Azure.Cosmos.Query.Core.Monads; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote.OrderBy; /// /// Concrete implementation of IAggregator that can take the global min/max from the local min/max of multiple partitions and continuations. diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CatchAllQueryPipelineStage.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CatchAllQueryPipelineStage.cs index 0a5af1767f..939ad9b708 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CatchAllQueryPipelineStage.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CatchAllQueryPipelineStage.cs @@ -8,7 +8,6 @@ namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline using System.Threading; using System.Threading.Tasks; using Microsoft.Azure.Cosmos.Query.Core.Monads; - using Microsoft.Azure.Cosmos.Query.Core.QueryClient; internal sealed class CatchAllQueryPipelineStage : QueryPipelineStageBase { diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosQueryExecutionContextFactory.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs similarity index 95% rename from Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosQueryExecutionContextFactory.cs rename to Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs index 6ff1f81bb5..1f9c3394bf 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosQueryExecutionContextFactory.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs @@ -19,6 +19,7 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionContext using Microsoft.Azure.Cosmos.Query.Core.Pipeline; using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Distinct; using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote.Parallel; using Microsoft.Azure.Cosmos.Query.Core.QueryClient; using Microsoft.Azure.Cosmos.Query.Core.QueryPlan; using Microsoft.Azure.Cosmos.SqlObjects; @@ -337,17 +338,17 @@ private static TryCatch TryCreatePassthroughQueryExecutionC List targetRanges, string collectionRid) { - CosmosCrossPartitionQueryExecutionContext.CrossPartitionInitParams initParams = new CosmosCrossPartitionQueryExecutionContext.CrossPartitionInitParams( - sqlQuerySpec: inputParameters.SqlQuerySpec, - collectionRid: collectionRid, - partitionedQueryExecutionInfo: partitionedQueryExecutionInfo, - partitionKeyRanges: targetRanges, - initialPageSize: inputParameters.MaxItemCount, - maxConcurrency: inputParameters.MaxConcurrency, - maxItemCount: inputParameters.MaxItemCount, - maxBufferedItemCount: inputParameters.MaxBufferedItemCount, - returnResultsInDeterministicOrder: inputParameters.ReturnResultsInDeterministicOrder, - testSettings: inputParameters.TestInjections); + //CosmosCrossPartitionQueryExecutionContext.CrossPartitionInitParams initParams = new CosmosCrossPartitionQueryExecutionContext.CrossPartitionInitParams( + // sqlQuerySpec: inputParameters.SqlQuerySpec, + // collectionRid: collectionRid, + // partitionedQueryExecutionInfo: partitionedQueryExecutionInfo, + // partitionKeyRanges: targetRanges, + // initialPageSize: inputParameters.MaxItemCount, + // maxConcurrency: inputParameters.MaxConcurrency, + // maxItemCount: inputParameters.MaxItemCount, + // maxBufferedItemCount: inputParameters.MaxBufferedItemCount, + // returnResultsInDeterministicOrder: inputParameters.ReturnResultsInDeterministicOrder, + // testSettings: inputParameters.TestInjections); // Modify query plan PartitionedQueryExecutionInfo passThroughQueryInfo = new PartitionedQueryExecutionInfo() @@ -367,26 +368,26 @@ private static TryCatch TryCreatePassthroughQueryExecutionC RewrittenQuery = null, Top = null, }, - QueryRanges = initParams.PartitionedQueryExecutionInfo.QueryRanges, + QueryRanges = partitionedQueryExecutionInfo.QueryRanges, }; - initParams = new CosmosCrossPartitionQueryExecutionContext.CrossPartitionInitParams( - sqlQuerySpec: initParams.SqlQuerySpec, - collectionRid: initParams.CollectionRid, - partitionedQueryExecutionInfo: passThroughQueryInfo, - partitionKeyRanges: initParams.PartitionKeyRanges, - initialPageSize: initParams.MaxItemCount.GetValueOrDefault(1000), - maxConcurrency: initParams.MaxConcurrency, - maxItemCount: initParams.MaxItemCount, - maxBufferedItemCount: initParams.MaxBufferedItemCount, - returnResultsInDeterministicOrder: initParams.ReturnResultsInDeterministicOrder, - testSettings: initParams.TestSettings); + //initParams = new CosmosCrossPartitionQueryExecutionContext.CrossPartitionInitParams( + // sqlQuerySpec: initParams.SqlQuerySpec, + // collectionRid: initParams.CollectionRid, + // partitionedQueryExecutionInfo: passThroughQueryInfo, + // partitionKeyRanges: initParams.PartitionKeyRanges, + // initialPageSize: initParams.MaxItemCount.GetValueOrDefault(1000), + // maxConcurrency: initParams.MaxConcurrency, + // maxItemCount: initParams.MaxItemCount, + // maxBufferedItemCount: initParams.MaxBufferedItemCount, + // returnResultsInDeterministicOrder: initParams.ReturnResultsInDeterministicOrder, + // testSettings: initParams.TestSettings); // Return a parallel context, since we still want to be able to handle splits and concurrency / buffering. return ParallelCrossPartitionQueryPipelineStage.MonadicCreate( documentContainer: documentContainer, sqlQuerySpec: inputParameters.SqlQuerySpec, - pageSize: initParams.InitialPageSize, + pageSize: inputParameters.MaxItemCount, continuationToken: inputParameters.InitialUserContinuationToken); } @@ -779,4 +780,4 @@ public override bool Visit(SqlParameterRefScalarExpression scalarExpression) } } } -} +} \ No newline at end of file diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/ExecutionEnvironment.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/ExecutionEnvironment.cs similarity index 90% rename from Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/ExecutionEnvironment.cs rename to Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/ExecutionEnvironment.cs index 46ae77d7bd..2daa905243 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/ExecutionEnvironment.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/ExecutionEnvironment.cs @@ -2,7 +2,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // ------------------------------------------------------------ -namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionContext +namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline { /// /// Environment the query is going to be executed on. diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/PipelineFactory.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/PipelineFactory.cs index d726656d82..0914388cb2 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/PipelineFactory.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/PipelineFactory.cs @@ -7,12 +7,11 @@ namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline using System; using Microsoft.Azure.Cosmos.CosmosElements; using Microsoft.Azure.Cosmos.Pagination; - using Microsoft.Azure.Cosmos.Query.Core.ExecutionContext; using Microsoft.Azure.Cosmos.Query.Core.Monads; using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Aggregate; using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Distinct; using Microsoft.Azure.Cosmos.Query.Core.Pipeline.GroupBy; - using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote.Parallel; using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Skip; using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Take; using Microsoft.Azure.Cosmos.Query.Core.QueryPlan; diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/OrderBy/CosmosElementToQueryLiteral.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderBy/CosmosElementToQueryLiteral.cs similarity index 98% rename from Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/OrderBy/CosmosElementToQueryLiteral.cs rename to Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderBy/CosmosElementToQueryLiteral.cs index 68eb125303..9af8bf31e4 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/OrderBy/CosmosElementToQueryLiteral.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderBy/CosmosElementToQueryLiteral.cs @@ -2,7 +2,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // ------------------------------------------------------------ -namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.OrderBy +namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote.OrderBy { using System; using System.Collections.Generic; diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/ItemProducers/ItemComparer.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderBy/ItemComparer.cs similarity index 97% rename from Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/ItemProducers/ItemComparer.cs rename to Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderBy/ItemComparer.cs index 3342a25e3d..64db04ed2c 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/ItemProducers/ItemComparer.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderBy/ItemComparer.cs @@ -1,13 +1,13 @@ //------------------------------------------------------------ // Copyright (c) Microsoft Corporation. All rights reserved. //------------------------------------------------------------ -namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.ItemProducers + +namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote.OrderBy { using System; using System.Collections.Generic; using Microsoft.Azure.Cosmos.CosmosElements; using Microsoft.Azure.Cosmos.Json; - using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Distinct; /// /// Utility class used to compare all items that we get back from a query. diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderByCrossPartitionQueryPipelineStage.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderBy/OrderByCrossPartitionQueryPipelineStage.cs similarity index 99% rename from Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderByCrossPartitionQueryPipelineStage.cs rename to Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderBy/OrderByCrossPartitionQueryPipelineStage.cs index 0c7b2ca31e..e4bc0b5b26 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderByCrossPartitionQueryPipelineStage.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderBy/OrderByCrossPartitionQueryPipelineStage.cs @@ -2,7 +2,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // ------------------------------------------------------------ -namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote +namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote.OrderBy { using System; using System.Collections.Generic; @@ -16,8 +16,6 @@ namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote using Microsoft.Azure.Cosmos.Query.Core.Collections; using Microsoft.Azure.Cosmos.Query.Core.ContinuationTokens; using Microsoft.Azure.Cosmos.Query.Core.Exceptions; - using Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.ItemProducers; - using Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.OrderBy; using Microsoft.Azure.Cosmos.Query.Core.Monads; using Microsoft.Azure.Documents; using static Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote.PartitionMapper; diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderByEnumeratorComparer.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderBy/OrderByEnumeratorComparer.cs similarity index 96% rename from Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderByEnumeratorComparer.cs rename to Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderBy/OrderByEnumeratorComparer.cs index 8e6125891b..c3d7ef81fd 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderByEnumeratorComparer.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderBy/OrderByEnumeratorComparer.cs @@ -2,15 +2,11 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // ------------------------------------------------------------ -namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote +namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote.OrderBy { using System; using System.Collections.Generic; using System.Diagnostics; - using System.Linq; - using System.Text; - using Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.ItemProducers; - using Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.OrderBy; /// /// For cross partition order by queries we serve documents from the partition diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/OrderBy/OrderByItem.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderBy/OrderByItem.cs similarity index 95% rename from Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/OrderBy/OrderByItem.cs rename to Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderBy/OrderByItem.cs index 4bab3fbffc..4b5c7bff1d 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/OrderBy/OrderByItem.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderBy/OrderByItem.cs @@ -1,7 +1,8 @@ //------------------------------------------------------------ // Copyright (c) Microsoft Corporation. All rights reserved. //------------------------------------------------------------ -namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.OrderBy + +namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote.OrderBy { using System; using Microsoft.Azure.Cosmos.CosmosElements; diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderByQueryPage.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderBy/OrderByQueryPage.cs similarity index 92% rename from Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderByQueryPage.cs rename to Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderBy/OrderByQueryPage.cs index 11b6218a87..78f19f67b6 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderByQueryPage.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderBy/OrderByQueryPage.cs @@ -2,7 +2,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // ------------------------------------------------------------ -namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote +namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote.OrderBy { using System; using System.Collections.Generic; diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderByQueryPartitionRangePageAsyncEnumerator.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderBy/OrderByQueryPartitionRangePageAsyncEnumerator.cs similarity index 97% rename from Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderByQueryPartitionRangePageAsyncEnumerator.cs rename to Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderBy/OrderByQueryPartitionRangePageAsyncEnumerator.cs index b646b97797..11ec9b899e 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderByQueryPartitionRangePageAsyncEnumerator.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderBy/OrderByQueryPartitionRangePageAsyncEnumerator.cs @@ -2,7 +2,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // ------------------------------------------------------------ -namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote +namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote.OrderBy { using System; using System.Threading; diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/OrderBy/OrderByQueryResult.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderBy/OrderByQueryResult.cs similarity index 98% rename from Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/OrderBy/OrderByQueryResult.cs rename to Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderBy/OrderByQueryResult.cs index 9eeff18e1f..1fcc4316d2 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/OrderBy/OrderByQueryResult.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderBy/OrderByQueryResult.cs @@ -1,7 +1,8 @@ //------------------------------------------------------------ // Copyright (c) Microsoft Corporation. All rights reserved. //------------------------------------------------------------ -namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.OrderBy + +namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote.OrderBy { using System; using System.Collections.Generic; diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/OrderBy/SortOrder.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderBy/SortOrder.cs similarity index 79% rename from Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/OrderBy/SortOrder.cs rename to Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderBy/SortOrder.cs index b9657245ca..df5b006325 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/OrderBy/SortOrder.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderBy/SortOrder.cs @@ -2,7 +2,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. //------------------------------------------------------------ -namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.OrderBy +namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote.OrderBy { internal enum SortOrder { diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/ParallelCrossPartitionQueryPipelineStage.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/Parallel/ParallelCrossPartitionQueryPipelineStage.cs similarity index 99% rename from Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/ParallelCrossPartitionQueryPipelineStage.cs rename to Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/Parallel/ParallelCrossPartitionQueryPipelineStage.cs index 04d498d827..f353b19364 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/ParallelCrossPartitionQueryPipelineStage.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/Parallel/ParallelCrossPartitionQueryPipelineStage.cs @@ -2,7 +2,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // ------------------------------------------------------------ -namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote +namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote.Parallel { using System; using System.Collections.Generic; diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/QueryPartitionRangePageAsyncEnumerator.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/Parallel/QueryPartitionRangePageAsyncEnumerator.cs similarity index 96% rename from Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/QueryPartitionRangePageAsyncEnumerator.cs rename to Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/Parallel/QueryPartitionRangePageAsyncEnumerator.cs index cb049dbce6..49983c0176 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/QueryPartitionRangePageAsyncEnumerator.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/Parallel/QueryPartitionRangePageAsyncEnumerator.cs @@ -2,7 +2,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // ------------------------------------------------------------ -namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote +namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote.Parallel { using System; using System.Threading; diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/QueryClient/CosmosQueryClient.cs b/Microsoft.Azure.Cosmos/src/Query/Core/QueryClient/CosmosQueryClient.cs index 66bb2639d9..bacb8d02f8 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/QueryClient/CosmosQueryClient.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/QueryClient/CosmosQueryClient.cs @@ -10,7 +10,6 @@ namespace Microsoft.Azure.Cosmos.Query.Core.QueryClient using System.Threading; using System.Threading.Tasks; using Microsoft.Azure.Cosmos.Diagnostics; - using Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.ItemProducers; using Microsoft.Azure.Cosmos.Query.Core.Monads; using Microsoft.Azure.Cosmos.Query.Core.Pipeline; using Microsoft.Azure.Cosmos.Query.Core.QueryPlan; diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/QueryClient/CosmosQueryContext.cs b/Microsoft.Azure.Cosmos/src/Query/Core/QueryClient/CosmosQueryContext.cs index aa8ddc76fe..1c6e1c8982 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/QueryClient/CosmosQueryContext.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/QueryClient/CosmosQueryContext.cs @@ -6,7 +6,6 @@ namespace Microsoft.Azure.Cosmos.Query.Core.QueryClient using System; using System.Threading; using System.Threading.Tasks; - using Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.ItemProducers; using Microsoft.Azure.Cosmos.Query.Core.Monads; using Microsoft.Azure.Cosmos.Query.Core.Pipeline; using Microsoft.Azure.Cosmos.Query.Core.QueryPlan; diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/QueryPlan/QueryInfo.cs b/Microsoft.Azure.Cosmos/src/Query/Core/QueryPlan/QueryInfo.cs index 9692b0f1d2..341fdc275b 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/QueryPlan/QueryInfo.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/QueryPlan/QueryInfo.cs @@ -6,9 +6,9 @@ namespace Microsoft.Azure.Cosmos.Query.Core.QueryPlan { using System.Collections.Generic; using System.Linq; - using Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.OrderBy; using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Aggregate; using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Distinct; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote.OrderBy; using Newtonsoft.Json; using Newtonsoft.Json.Converters; diff --git a/Microsoft.Azure.Cosmos/src/RequestOptions/QueryRequestOptions.cs b/Microsoft.Azure.Cosmos/src/RequestOptions/QueryRequestOptions.cs index 2bc739ff77..0ed1e25d27 100644 --- a/Microsoft.Azure.Cosmos/src/RequestOptions/QueryRequestOptions.cs +++ b/Microsoft.Azure.Cosmos/src/RequestOptions/QueryRequestOptions.cs @@ -10,6 +10,7 @@ namespace Microsoft.Azure.Cosmos using Microsoft.Azure.Cosmos.CosmosElements; using Microsoft.Azure.Cosmos.Query.Core; using Microsoft.Azure.Cosmos.Query.Core.ExecutionContext; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline; using Microsoft.Azure.Documents; /// diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/Query/OrderByQueryTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/Query/OrderByQueryTests.cs index 7c9780ddd1..e40e63e80e 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/Query/OrderByQueryTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/Query/OrderByQueryTests.cs @@ -7,20 +7,18 @@ using System.Globalization; using System.Linq; using System.Net; - using System.Runtime.ExceptionServices; using System.Text; using System.Threading.Tasks; using Microsoft.Azure.Cosmos.CosmosElements; using Microsoft.Azure.Cosmos.CosmosElements.Numbers; using Microsoft.Azure.Cosmos.Query.Core; - using Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.ItemProducers; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote.OrderBy; using Microsoft.Azure.Cosmos.Routing; using Microsoft.Azure.Cosmos.SDK.EmulatorTests.QueryOracle; using Microsoft.Azure.Documents; using Microsoft.Azure.Documents.Routing; using Microsoft.VisualStudio.TestTools.UnitTesting; using Newtonsoft.Json; - using Newtonsoft.Json.Linq; [TestClass] public sealed class OrderByQueryTests : QueryTestsBase diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/Query/QueryTestsBase.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/Query/QueryTestsBase.cs index aea6218df7..0f6a439ab1 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/Query/QueryTestsBase.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/Query/QueryTestsBase.cs @@ -15,6 +15,7 @@ namespace Microsoft.Azure.Cosmos.EmulatorTests.Query using System.Threading; using System.Threading.Tasks; using Microsoft.Azure.Cosmos.CosmosElements; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline; using Microsoft.Azure.Cosmos.Routing; using Microsoft.Azure.Cosmos.SDK.EmulatorTests; using Microsoft.Azure.Documents; @@ -584,7 +585,7 @@ internal static async Task> QueryWithCosmosElementContinuationTokenAsync do { QueryRequestOptions computeRequestOptions = queryRequestOptions.Clone(); - computeRequestOptions.ExecutionEnvironment = Cosmos.Query.Core.ExecutionContext.ExecutionEnvironment.Compute; + computeRequestOptions.ExecutionEnvironment = ExecutionEnvironment.Compute; computeRequestOptions.CosmosElementContinuationToken = continuationToken; using (FeedIteratorInternal itemQuery = (FeedIteratorInternal)container.GetItemQueryIterator( diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/QueryTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/QueryTests.cs index f69d9a0242..2027d609fd 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/QueryTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/QueryTests.cs @@ -19,7 +19,6 @@ namespace Microsoft.Azure.Cosmos.SDK.EmulatorTests using Microsoft.Azure.Cosmos.Linq; using Microsoft.Azure.Cosmos.Query; using Microsoft.Azure.Cosmos.Query.Core; - using Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.Parallel; using Microsoft.Azure.Cosmos.Query.Core.Metrics; using Microsoft.Azure.Cosmos.Routing; using Microsoft.Azure.Cosmos.Services.Management.Tests; @@ -1971,70 +1970,6 @@ private void TestForceQueryScanHeaders(Database database, bool partitionedCollec Assert.AreNotEqual(TimeSpan.Zero, queryMetrics.BackendMetrics.IndexLookupTime); } - [Ignore] // Need to use v3 pipeline - [TestMethod] - public void TestMaxDegreeOfParallelism() - { - List> inputOutputMaxDops = new List>() - { - new Tuple(null, 0), - new Tuple(-1, int.MaxValue), - new Tuple(-2, int.MaxValue), - new Tuple(0, 0), - new Tuple(1, 1), - new Tuple(int.MinValue, int.MaxValue), - new Tuple(int.MaxValue, int.MaxValue), - }; - - this.TestFeedOptionInput( - nameof(FeedOptions.MaxDegreeOfParallelism), - "MaxDegreeOfParallelism", - inputOutputMaxDops); - } - - [Ignore] // Need to use v3 pipeline - [TestMethod] - public void TestMaxBufferedItemCount() - { - List> inputOutputMaxBufferedItemCounts = new List>() - { - new Tuple(null, (int)ParallelQueryConfig.GetConfig().DefaultMaximumBufferSize), - new Tuple(-1, int.MaxValue), - new Tuple(-2, int.MaxValue), - new Tuple(0,(int)ParallelQueryConfig.GetConfig().DefaultMaximumBufferSize), - new Tuple(1, 1), - new Tuple(int.MinValue, int.MaxValue), - new Tuple(int.MaxValue, int.MaxValue), - }; - - this.TestFeedOptionInput( - nameof(FeedOptions.MaxBufferedItemCount), - "ActualMaxBufferedItemCount", - inputOutputMaxBufferedItemCounts); - } - - [Ignore] // Need to use v3 pipeline - [TestMethod] - public void TestMaxItemCount() - { - List> inputOutputMaxItemCounts = new List>() - { - new Tuple(null, ParallelQueryConfig.GetConfig().ClientInternalPageSize), - new Tuple(-1, int.MaxValue), - new Tuple(-2, int.MaxValue), - // 0 is not a valid MaxItemCount - // new Tuple(0,(int)ParallelQueryConfig.GetConfig().ClientInternalPageSize), - new Tuple(1, 1), - new Tuple(int.MinValue, int.MaxValue), - new Tuple(int.MaxValue, int.MaxValue), - }; - - this.TestFeedOptionInput( - nameof(FeedOptions.MaxItemCount), - "ActualMaxPageSize", - inputOutputMaxItemCounts); - } - private void TestFeedOptionInput( string feedOptionPropertyName, string componentPropertyName, @@ -2189,9 +2124,7 @@ private void TestContinuationLimitHeaders(Database database, bool partitionedCol } catch (AggregateException e) { - DocumentClientException exception = e.InnerException as DocumentClientException; - - if (exception == null) + if (!(e.InnerException is DocumentClientException exception)) { throw e; } @@ -2215,9 +2148,7 @@ private void TestContinuationLimitHeaders(Database database, bool partitionedCol } catch (AggregateException e) { - DocumentClientException exception = e.InnerException as DocumentClientException; - - if (exception == null) + if (!(e.InnerException is DocumentClientException exception)) { throw e; } diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/ContinuationTokens/OrderByContinuationTokenTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/ContinuationTokens/OrderByContinuationTokenTests.cs index 916e02ecd6..270548c081 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/ContinuationTokens/OrderByContinuationTokenTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/ContinuationTokens/OrderByContinuationTokenTests.cs @@ -10,8 +10,8 @@ namespace Microsoft.Azure.Cosmos.Query using Microsoft.Azure.Cosmos.CosmosElements; using Microsoft.Azure.Cosmos.CosmosElements.Numbers; using Microsoft.Azure.Cosmos.Query.Core.ContinuationTokens; - using Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.OrderBy; using Microsoft.Azure.Cosmos.Query.Core.Monads; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote.OrderBy; using Microsoft.Azure.Documents.Routing; using Newtonsoft.Json; using VisualStudio.TestTools.UnitTesting; 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 71af67e5c2..87adf78407 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 @@ -10,6 +10,7 @@ namespace Microsoft.Azure.Cosmos.Tests.Query.Pipeline using Microsoft.Azure.Cosmos.Query.Core.Monads; using Microsoft.Azure.Cosmos.Query.Core.Pipeline; using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote.Parallel; using Microsoft.Azure.Cosmos.Query.Core.QueryPlan; using Microsoft.VisualStudio.TestTools.UnitTesting; using Moq; diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/OrderByCrossPartitionQueryPipelineStageTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/OrderByCrossPartitionQueryPipelineStageTests.cs index 36f850c0d4..54a1d13237 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/OrderByCrossPartitionQueryPipelineStageTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/OrderByCrossPartitionQueryPipelineStageTests.cs @@ -12,10 +12,9 @@ namespace Microsoft.Azure.Cosmos.Tests.Query.Pipeline using Microsoft.Azure.Cosmos.Query.Core; using Microsoft.Azure.Cosmos.Query.Core.ContinuationTokens; using Microsoft.Azure.Cosmos.Query.Core.Exceptions; - using Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.OrderBy; using Microsoft.Azure.Cosmos.Query.Core.Monads; using Microsoft.Azure.Cosmos.Query.Core.Pipeline; - using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote.OrderBy; using Microsoft.Azure.Cosmos.Tests.Pagination; using Microsoft.Azure.Documents; using Microsoft.VisualStudio.TestTools.UnitTesting; @@ -35,7 +34,7 @@ public void MonadicCreate_NullContinuationToken() targetRanges: new List() { new PartitionKeyRange() }, orderByColumns: new List() { - new OrderByCrossPartitionQueryPipelineStage.OrderByColumn("_ts", Cosmos.Query.Core.ExecutionContext.OrderBy.SortOrder.Ascending) + new OrderByCrossPartitionQueryPipelineStage.OrderByColumn("_ts", SortOrder.Ascending) }, pageSize: 10, continuationToken: null); @@ -53,7 +52,7 @@ public void MonadicCreate_NonCosmosArrayContinuationToken() targetRanges: new List() { new PartitionKeyRange() }, orderByColumns: new List() { - new OrderByCrossPartitionQueryPipelineStage.OrderByColumn("_ts", Cosmos.Query.Core.ExecutionContext.OrderBy.SortOrder.Ascending) + new OrderByCrossPartitionQueryPipelineStage.OrderByColumn("_ts", SortOrder.Ascending) }, pageSize: 10, continuationToken: CosmosObject.Create(new Dictionary())); @@ -72,7 +71,7 @@ public void MonadicCreate_EmptyArrayContinuationToken() targetRanges: new List() { new PartitionKeyRange() }, orderByColumns: new List() { - new OrderByCrossPartitionQueryPipelineStage.OrderByColumn("_ts", Cosmos.Query.Core.ExecutionContext.OrderBy.SortOrder.Ascending) + new OrderByCrossPartitionQueryPipelineStage.OrderByColumn("_ts", SortOrder.Ascending) }, pageSize: 10, continuationToken: CosmosArray.Create(new List())); @@ -91,7 +90,7 @@ public void MonadicCreate_NonCompositeContinuationToken() targetRanges: new List() { new PartitionKeyRange() }, orderByColumns: new List() { - new OrderByCrossPartitionQueryPipelineStage.OrderByColumn("_ts", Cosmos.Query.Core.ExecutionContext.OrderBy.SortOrder.Ascending) + new OrderByCrossPartitionQueryPipelineStage.OrderByColumn("_ts", SortOrder.Ascending) }, pageSize: 10, continuationToken: CosmosArray.Create(new List() { CosmosString.Create("asdf") })); @@ -123,7 +122,7 @@ public void MonadicCreate_SingleOrderByContinuationToken() targetRanges: new List() { new PartitionKeyRange() { Id = "0", MinInclusive = "A", MaxExclusive = "B" } }, orderByColumns: new List() { - new OrderByCrossPartitionQueryPipelineStage.OrderByColumn("_ts", Cosmos.Query.Core.ExecutionContext.OrderBy.SortOrder.Ascending) + new OrderByCrossPartitionQueryPipelineStage.OrderByColumn("_ts", SortOrder.Ascending) }, pageSize: 10, continuationToken: CosmosString.Create( @@ -182,7 +181,7 @@ public void MonadicCreate_MultipleOrderByContinuationToken() }, orderByColumns: new List() { - new OrderByCrossPartitionQueryPipelineStage.OrderByColumn("_ts", Cosmos.Query.Core.ExecutionContext.OrderBy.SortOrder.Ascending) + new OrderByCrossPartitionQueryPipelineStage.OrderByColumn("_ts", SortOrder.Ascending) }, pageSize: 10, continuationToken: CosmosString.Create( 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 16bb1cab19..84dc537b3d 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 @@ -14,6 +14,7 @@ namespace Microsoft.Azure.Cosmos.Tests.Query.Pipeline using Microsoft.Azure.Cosmos.Query.Core.Monads; using Microsoft.Azure.Cosmos.Query.Core.Pipeline; using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote.Parallel; using Microsoft.Azure.Cosmos.Tests.Pagination; using Microsoft.Azure.Documents; using Microsoft.VisualStudio.TestTools.UnitTesting; 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 c99dcd5e0e..a86bc0d439 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 @@ -8,6 +8,7 @@ using Microsoft.Azure.Cosmos.Query.Core.Monads; using Microsoft.Azure.Cosmos.Query.Core.Pipeline; using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote.Parallel; using Microsoft.Azure.Cosmos.Tests.Pagination; using Microsoft.Azure.Documents; using Microsoft.VisualStudio.TestTools.UnitTesting; diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/QueryPlanBaselineTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/QueryPlanBaselineTests.cs index b1b6ee91ce..08d822cb5c 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/QueryPlanBaselineTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/QueryPlanBaselineTests.cs @@ -11,11 +11,10 @@ using Microsoft.VisualStudio.TestTools.UnitTesting; using Newtonsoft.Json; using Microsoft.Azure.Cosmos.Query.Core.QueryPlan; - using Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.OrderBy; - using Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.Aggregate; using Microsoft.Azure.Cosmos.Query.Core.Monads; using System.Linq; using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Aggregate; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote.OrderBy; /// /// Tests for . From a75df5ddc959ffda104f687c225dc911d3dfa1ae Mon Sep 17 00:00:00 2001 From: Brandon Chong Date: Fri, 7 Aug 2020 18:56:47 -0700 Subject: [PATCH 59/85] removed old execution component code --- .../src/CosmosElements/CosmosArray.cs | 1 - .../src/CosmosElements/CosmosObject.cs | 1 - ...ggregateDocumentQueryExecutionComponent.cs | 115 -------- .../CosmosQueryExecutionComponent.cs | 38 --- ...DistinctDocumentQueryExecutionComponent.cs | 37 --- .../DocumentQueryExecutionComponentBase.cs | 74 ----- .../GroupByDocumentQueryExecutionComponent.cs | 278 ------------------ .../IDocumentQueryExecutionComponent.cs | 38 --- .../SkipDocumentQueryExecutionComponent.cs | 39 --- .../TakeDocumentQueryExecutionComponent.cs | 39 --- .../CosmosQueryUnitTests.cs | 144 --------- 11 files changed, 804 deletions(-) delete mode 100644 Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/AggregateDocumentQueryExecutionComponent.cs delete mode 100644 Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/CosmosQueryExecutionComponent.cs delete mode 100644 Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctDocumentQueryExecutionComponent.cs delete mode 100644 Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/DocumentQueryExecutionComponentBase.cs delete mode 100644 Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/GroupBy/GroupByDocumentQueryExecutionComponent.cs delete mode 100644 Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/IDocumentQueryExecutionComponent.cs delete mode 100644 Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/SkipDocumentQueryExecutionComponent.cs delete mode 100644 Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/TakeDocumentQueryExecutionComponent.cs diff --git a/Microsoft.Azure.Cosmos/src/CosmosElements/CosmosArray.cs b/Microsoft.Azure.Cosmos/src/CosmosElements/CosmosArray.cs index 427da675aa..4f0dee8dfa 100644 --- a/Microsoft.Azure.Cosmos/src/CosmosElements/CosmosArray.cs +++ b/Microsoft.Azure.Cosmos/src/CosmosElements/CosmosArray.cs @@ -10,7 +10,6 @@ namespace Microsoft.Azure.Cosmos.CosmosElements using System.Collections.Generic; using System.Linq; using Microsoft.Azure.Cosmos.Json; - using Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.Distinct; using Microsoft.Azure.Cosmos.Query.Core.Monads; using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Distinct; diff --git a/Microsoft.Azure.Cosmos/src/CosmosElements/CosmosObject.cs b/Microsoft.Azure.Cosmos/src/CosmosElements/CosmosObject.cs index 9f4bcf2e48..367d6dc42b 100644 --- a/Microsoft.Azure.Cosmos/src/CosmosElements/CosmosObject.cs +++ b/Microsoft.Azure.Cosmos/src/CosmosElements/CosmosObject.cs @@ -10,7 +10,6 @@ namespace Microsoft.Azure.Cosmos.CosmosElements using System.Collections.Generic; using System.Linq; using Microsoft.Azure.Cosmos.Json; - using Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.Distinct; using Microsoft.Azure.Cosmos.Query.Core.Monads; using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Distinct; diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/AggregateDocumentQueryExecutionComponent.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/AggregateDocumentQueryExecutionComponent.cs deleted file mode 100644 index d7a88c5ea2..0000000000 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/AggregateDocumentQueryExecutionComponent.cs +++ /dev/null @@ -1,115 +0,0 @@ -//------------------------------------------------------------ -// Copyright (c) Microsoft Corporation. All rights reserved. -//------------------------------------------------------------ -namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.Aggregate -{ - using System; - using System.Collections.Generic; - using System.Threading.Tasks; - using Microsoft.Azure.Cosmos.CosmosElements; - using Microsoft.Azure.Cosmos.Query.Core.Monads; - using Microsoft.Azure.Cosmos.Query.Core.Pipeline; - using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Aggregate; - using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Aggregate.Aggregators; - - /// - /// Execution component that is able to aggregate local aggregates from multiple continuations and partitions. - /// At a high level aggregates queries only return a "partial" aggregate. - /// "partial" means that the result is only valid for that one continuation (and one partition). - /// For example suppose you have the query "SELECT COUNT(1) FROM c" and you have a single partition collection, - /// then you will get one count for each continuation of the query. - /// If you wanted the true result for this query, then you will have to take the sum of all continuations. - /// The reason why we have multiple continuations is because for a long running query we have to break up the results into multiple continuations. - /// Fortunately all the aggregates can be aggregated across continuations and partitions. - /// - internal abstract partial class AggregateDocumentQueryExecutionComponent : DocumentQueryExecutionComponentBase - { - /// - /// This class does most of the work, since a query like: - /// - /// SELECT VALUE AVG(c.age) - /// FROM c - /// - /// is really just an aggregation on a single grouping (the whole collection). - /// - private readonly SingleGroupAggregator singleGroupAggregator; - - /// - /// We need to keep track of whether the projection has the 'VALUE' keyword. - /// - private readonly bool isValueAggregateQuery; - - /// - /// Initializes a new instance of the AggregateDocumentQueryExecutionComponent class. - /// - /// The source component that will supply the local aggregates from multiple continuations and partitions. - /// The single group aggregator that we will feed results into. - /// Whether or not the query has the 'VALUE' keyword. - /// This constructor is private since there is some async initialization that needs to happen in CreateAsync(). - protected AggregateDocumentQueryExecutionComponent( - IDocumentQueryExecutionComponent source, - SingleGroupAggregator singleGroupAggregator, - bool isValueAggregateQuery) - : base(source) - { - this.singleGroupAggregator = singleGroupAggregator ?? throw new ArgumentNullException(nameof(singleGroupAggregator)); - this.isValueAggregateQuery = isValueAggregateQuery; - } - - public static Task> MonadicCreateAsync( - ExecutionEnvironment executionEnvironment, - IReadOnlyList aggregates, - IReadOnlyDictionary aliasToAggregateType, - IReadOnlyList orderedAliases, - bool hasSelectValue, - CosmosElement continuationToken, - Func>> monadicCreatePipelineStage) - { - return default; - } - - /// - /// Struct for getting the payload out of the rewritten projection. - /// - private readonly struct RewrittenAggregateProjections - { - public RewrittenAggregateProjections(bool isValueAggregateQuery, CosmosElement raw) - { - if (raw == null) - { - throw new ArgumentNullException(nameof(raw)); - } - - if (isValueAggregateQuery) - { - // SELECT VALUE [{"item": {"sum": SUM(c.blah), "count": COUNT(c.blah)}}] - if (!(raw is CosmosArray aggregates)) - { - throw new ArgumentException($"{nameof(RewrittenAggregateProjections)} was not an array for a value aggregate query. Type is: {raw.GetType()}"); - } - - this.Payload = aggregates[0]; - } - else - { - if (!(raw is CosmosObject cosmosObject)) - { - throw new ArgumentException($"{nameof(raw)} must not be an object."); - } - - if (!cosmosObject.TryGetValue("payload", out CosmosElement cosmosPayload)) - { - throw new InvalidOperationException($"Underlying object does not have an 'payload' field."); - } - - this.Payload = cosmosPayload ?? throw new ArgumentException($"{nameof(RewrittenAggregateProjections)} does not have a 'payload' property."); - } - } - - public CosmosElement Payload - { - get; - } - } - } -} \ No newline at end of file diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/CosmosQueryExecutionComponent.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/CosmosQueryExecutionComponent.cs deleted file mode 100644 index 7ba532d496..0000000000 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/CosmosQueryExecutionComponent.cs +++ /dev/null @@ -1,38 +0,0 @@ -//------------------------------------------------------------ -// Copyright (c) Microsoft Corporation. All rights reserved. -//------------------------------------------------------------ -namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent -{ - using System; - using System.Collections.Generic; - using System.Threading; - using System.Threading.Tasks; - using Microsoft.Azure.Cosmos; - using Microsoft.Azure.Cosmos.Query.Core.QueryClient; - - /// - /// Interface for all DocumentQueryExecutionComponents - /// - internal abstract class CosmosQueryExecutionComponent : IDisposable - { - /// - /// Gets a value indicating whether this component is done draining documents. - /// - public abstract bool IsDone { get; } - - /// - /// Drains documents from this component. - /// - /// The maximum number of documents to drain. - /// The cancellation token to cancel tasks. - /// A task that when awaited on returns a feed response. - public abstract Task DrainAsync(int maxElements, CancellationToken token); - - /// - /// Stops this document query execution component. - /// - public abstract void Stop(); - - public abstract void Dispose(); - } -} diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctDocumentQueryExecutionComponent.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctDocumentQueryExecutionComponent.cs deleted file mode 100644 index 79d3f45852..0000000000 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctDocumentQueryExecutionComponent.cs +++ /dev/null @@ -1,37 +0,0 @@ -//------------------------------------------------------------ -// Copyright (c) Microsoft Corporation. All rights reserved. -//------------------------------------------------------------ -namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.Distinct -{ - using System; - using System.Threading.Tasks; - using Microsoft.Azure.Cosmos.CosmosElements; - using Microsoft.Azure.Cosmos.Query.Core.Monads; - using Microsoft.Azure.Cosmos.Query.Core.Pipeline; - using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Distinct; - - internal abstract partial class DistinctDocumentQueryExecutionComponent : DocumentQueryExecutionComponentBase - { - /// - /// An DistinctMap that efficiently stores the documents that we have already seen. - /// - private readonly DistinctMap distinctMap; - - protected DistinctDocumentQueryExecutionComponent( - DistinctMap distinctMap, - IDocumentQueryExecutionComponent source) - : base(source) - { - this.distinctMap = distinctMap ?? throw new ArgumentNullException(nameof(distinctMap)); - } - - public static Task> MonadicCreateAsync( - ExecutionEnvironment executionEnvironment, - CosmosElement requestContinuation, - Func>> monadicCreatePipelineStage, - DistinctQueryType distinctQueryType) - { - return default; - } - } -} diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/DocumentQueryExecutionComponentBase.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/DocumentQueryExecutionComponentBase.cs deleted file mode 100644 index e3ab5efd67..0000000000 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/DocumentQueryExecutionComponentBase.cs +++ /dev/null @@ -1,74 +0,0 @@ -//------------------------------------------------------------ -// Copyright (c) Microsoft Corporation. All rights reserved. -//------------------------------------------------------------ -namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent -{ - using System; - using System.Threading; - using System.Threading.Tasks; - using Microsoft.Azure.Cosmos.CosmosElements; - using Microsoft.Azure.Cosmos.Query.Core.QueryClient; - - /// - /// Base class for all DocumentQueryExecutionComponents that implements and IDocumentQueryExecutionComponent - /// - internal abstract class DocumentQueryExecutionComponentBase : IDocumentQueryExecutionComponent - { - public static readonly string UseCosmosElementContinuationTokenInstead = $"Use Cosmos Element Continuation Token instead."; - - /// - /// Source DocumentQueryExecutionComponent that this component will drain from. - /// - protected readonly IDocumentQueryExecutionComponent Source; - - /// - /// Initializes a new instance of the DocumentQueryExecutionComponentBase class. - /// - /// The source to drain documents from. - protected DocumentQueryExecutionComponentBase(IDocumentQueryExecutionComponent source) - { - this.Source = source ?? throw new ArgumentNullException(nameof(source)); - } - - /// - /// Gets a value indicating whether or not this component is done draining documents. - /// - public virtual bool IsDone - { - get - { - return this.Source.IsDone; - } - } - - /// - /// Disposes this context. - /// - public virtual void Dispose() - { - this.Source.Dispose(); - } - - /// - /// Drains documents from this execution context. - /// - /// Upper bound for the number of documents you wish to receive. - /// The cancellation token to use. - /// A DoucmentFeedResponse of documents. - public virtual Task DrainAsync(int maxElements, CancellationToken token) - { - token.ThrowIfCancellationRequested(); - return this.Source.DrainAsync(maxElements, token); - } - - /// - /// Stops the execution component. - /// - public void Stop() - { - this.Source.Stop(); - } - - public abstract CosmosElement GetCosmosElementContinuationToken(); - } -} diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/GroupBy/GroupByDocumentQueryExecutionComponent.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/GroupBy/GroupByDocumentQueryExecutionComponent.cs deleted file mode 100644 index 1e65f8d4ae..0000000000 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/GroupBy/GroupByDocumentQueryExecutionComponent.cs +++ /dev/null @@ -1,278 +0,0 @@ -//------------------------------------------------------------ -// Copyright (c) Microsoft Corporation. All rights reserved. -//------------------------------------------------------------ -namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.GroupBy -{ - using System; - using System.Collections; - using System.Collections.Generic; - using System.Linq; - using System.Threading.Tasks; - using Microsoft.Azure.Cosmos.CosmosElements; - using Microsoft.Azure.Cosmos.Query.Core.Exceptions; - 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.Aggregate; - using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Aggregate.Aggregators; - using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Distinct; - - /// - /// Query execution component that groups groupings across continuations and pages. - /// The general idea is a query gets rewritten from this: - /// - /// SELECT c.team, c.name, COUNT(1) AS count - /// FROM c - /// GROUP BY c.team, c.name - /// - /// To this: - /// - /// SELECT - /// [{"item": c.team}, {"item": c.name}] AS groupByItems, - /// {"team": c.team, "name": c.name, "count": {"item": COUNT(1)}} AS payload - /// FROM c - /// GROUP BY c.team, c.name - /// - /// With the following dictionary: - /// - /// { - /// "team": null, - /// "name": null, - /// "count" COUNT - /// } - /// - /// So we know how to aggregate each column. - /// At the end the columns are stitched together to make the grouped document. - /// - internal abstract partial class GroupByDocumentQueryExecutionComponent : DocumentQueryExecutionComponentBase - { - private readonly GroupingTable groupingTable; - - protected GroupByDocumentQueryExecutionComponent( - IDocumentQueryExecutionComponent source, - GroupingTable groupingTable) - : base(source) - { - this.groupingTable = groupingTable ?? throw new ArgumentNullException(nameof(groupingTable)); - } - - public override bool IsDone => this.groupingTable.IsDone; - - public static Task> MonadicCreateAsync( - ExecutionEnvironment executionEnvironment, - CosmosElement continuationToken, - Func>> monadicCreatePipelineStage, - IReadOnlyDictionary groupByAliasToAggregateType, - IReadOnlyList orderedAliases, - bool hasSelectValue) - { - return default; - } - - protected void AggregateGroupings(IReadOnlyList cosmosElements) - { - foreach (CosmosElement result in cosmosElements) - { - // Aggregate the values for all groupings across all continuations. - RewrittenGroupByProjection groupByItem = new RewrittenGroupByProjection(result); - this.groupingTable.AddPayload(groupByItem); - } - } - - /// - /// When a group by query gets rewritten the projection looks like: - /// - /// SELECT - /// [{"item": c.age}, {"item": c.name}] AS groupByItems, - /// {"age": c.age, "name": c.name} AS payload - /// - /// This struct just lets us easily access the "groupByItems" and "payload" property. - /// - protected readonly struct RewrittenGroupByProjection - { - private const string GroupByItemsPropertyName = "groupByItems"; - private const string PayloadPropertyName = "payload"; - - private readonly CosmosObject cosmosObject; - - public RewrittenGroupByProjection(CosmosElement cosmosElement) - { - if (cosmosElement == null) - { - throw new ArgumentNullException(nameof(cosmosElement)); - } - - if (!(cosmosElement is CosmosObject cosmosObject)) - { - throw new ArgumentException($"{nameof(cosmosElement)} must not be an object."); - } - - this.cosmosObject = cosmosObject; - } - - public CosmosArray GroupByItems - { - get - { - if (!this.cosmosObject.TryGetValue(GroupByItemsPropertyName, out CosmosElement cosmosElement)) - { - throw new InvalidOperationException($"Underlying object does not have an 'groupByItems' field."); - } - - if (!(cosmosElement is CosmosArray cosmosArray)) - { - throw new ArgumentException($"{nameof(RewrittenGroupByProjection)}['groupByItems'] was not an array."); - } - - return cosmosArray; - } - } - - public bool TryGetPayload(out CosmosElement payload) => this.cosmosObject.TryGetValue(PayloadPropertyName, out payload); - } - - protected sealed class GroupingTable : IEnumerable> - { - private static readonly IReadOnlyList EmptyAggregateOperators = new AggregateOperator[] { }; - - private readonly Dictionary table; - private readonly IReadOnlyDictionary groupByAliasToAggregateType; - private readonly IReadOnlyList orderedAliases; - private readonly bool hasSelectValue; - - private GroupingTable( - IReadOnlyDictionary groupByAliasToAggregateType, - IReadOnlyList orderedAliases, - bool hasSelectValue) - { - this.groupByAliasToAggregateType = groupByAliasToAggregateType ?? throw new ArgumentNullException(nameof(groupByAliasToAggregateType)); - this.orderedAliases = orderedAliases; - this.hasSelectValue = hasSelectValue; - this.table = new Dictionary(); - } - - public int Count => this.table.Count; - - public bool IsDone { get; private set; } - - public void AddPayload(RewrittenGroupByProjection rewrittenGroupByProjection) - { - // For VALUE queries the payload will be undefined if the field was undefined. - if (rewrittenGroupByProjection.TryGetPayload(out CosmosElement payload)) - { - UInt128 groupByKeysHash = DistinctHash.GetHash(rewrittenGroupByProjection.GroupByItems); - - if (!this.table.TryGetValue(groupByKeysHash, out SingleGroupAggregator singleGroupAggregator)) - { - singleGroupAggregator = SingleGroupAggregator.TryCreate( - EmptyAggregateOperators, - this.groupByAliasToAggregateType, - this.orderedAliases, - this.hasSelectValue, - continuationToken: null).Result; - this.table[groupByKeysHash] = singleGroupAggregator; - } - - singleGroupAggregator.AddValues(payload); - } - } - - public IReadOnlyList Drain(int maxItemCount) - { - List keys = this.table.Keys.Take(maxItemCount).ToList(); - List singleGroupAggregators = new List(keys.Count); - foreach (UInt128 key in keys) - { - SingleGroupAggregator singleGroupAggregator = this.table[key]; - singleGroupAggregators.Add(singleGroupAggregator); - } - - foreach (UInt128 key in keys) - { - this.table.Remove(key); - } - - List results = new List(); - foreach (SingleGroupAggregator singleGroupAggregator in singleGroupAggregators) - { - results.Add(singleGroupAggregator.GetResult()); - } - - if (this.Count == 0) - { - this.IsDone = true; - } - - return results; - } - - public CosmosElement GetCosmosElementContinuationToken() - { - Dictionary dictionary = new Dictionary(); - foreach (KeyValuePair kvp in this.table) - { - dictionary.Add(kvp.Key.ToString(), kvp.Value.GetCosmosElementContinuationToken()); - } - - return CosmosObject.Create(dictionary); - } - - public IEnumerator> GetEnumerator => this.table.GetEnumerator(); - - public static TryCatch TryCreateFromContinuationToken( - IReadOnlyDictionary groupByAliasToAggregateType, - IReadOnlyList orderedAliases, - bool hasSelectValue, - CosmosElement continuationToken) - { - GroupingTable groupingTable = new GroupingTable( - groupByAliasToAggregateType, - orderedAliases, - hasSelectValue); - - if (continuationToken != null) - { - if (!(continuationToken is CosmosObject groupingTableContinuationToken)) - { - return TryCatch.FromException( - new MalformedContinuationTokenException($"Invalid GroupingTableContinuationToken")); - } - - foreach (KeyValuePair kvp in groupingTableContinuationToken) - { - string key = kvp.Key; - CosmosElement value = kvp.Value; - - if (!UInt128.TryParse(key, out UInt128 groupByKey)) - { - return TryCatch.FromException( - new MalformedContinuationTokenException($"Invalid GroupingTableContinuationToken")); - } - - TryCatch tryCreateSingleGroupAggregator = SingleGroupAggregator.TryCreate( - EmptyAggregateOperators, - groupByAliasToAggregateType, - orderedAliases, - hasSelectValue, - value); - - if (tryCreateSingleGroupAggregator.Succeeded) - { - groupingTable.table[groupByKey] = tryCreateSingleGroupAggregator.Result; - } - else - { - return TryCatch.FromException(tryCreateSingleGroupAggregator.Exception); - } - } - } - - return TryCatch.FromResult(groupingTable); - } - - IEnumerator> IEnumerable>.GetEnumerator() => this.table.GetEnumerator(); - - IEnumerator IEnumerable.GetEnumerator() => this.table.GetEnumerator(); - } - } -} \ No newline at end of file diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/IDocumentQueryExecutionComponent.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/IDocumentQueryExecutionComponent.cs deleted file mode 100644 index eb547dad68..0000000000 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/IDocumentQueryExecutionComponent.cs +++ /dev/null @@ -1,38 +0,0 @@ -//------------------------------------------------------------ -// Copyright (c) Microsoft Corporation. All rights reserved. -//------------------------------------------------------------ -namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent -{ - using System; - using System.Threading; - using System.Threading.Tasks; - using Microsoft.Azure.Cosmos.CosmosElements; - using Microsoft.Azure.Cosmos.Json; - using Microsoft.Azure.Cosmos.Query.Core.QueryClient; - - /// - /// Interface for all DocumentQueryExecutionComponents - /// - internal interface IDocumentQueryExecutionComponent : IDisposable - { - /// - /// Gets a value indicating whether this component is done draining documents. - /// - bool IsDone { get; } - - /// - /// Drains documents from this component. - /// - /// The maximum number of documents to drain. - /// The cancellation token to cancel tasks. - /// A task that when awaited on returns a feed response. - Task DrainAsync(int maxElements, CancellationToken token); - - /// - /// Stops this document query execution component. - /// - void Stop(); - - CosmosElement GetCosmosElementContinuationToken(); - } -} diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/SkipDocumentQueryExecutionComponent.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/SkipDocumentQueryExecutionComponent.cs deleted file mode 100644 index 5069cd529e..0000000000 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/SkipDocumentQueryExecutionComponent.cs +++ /dev/null @@ -1,39 +0,0 @@ -//------------------------------------------------------------ -// Copyright (c) Microsoft Corporation. All rights reserved. -//------------------------------------------------------------ - -namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.SkipTake -{ - using System; - using System.Threading.Tasks; - using Microsoft.Azure.Cosmos.CosmosElements; - using Microsoft.Azure.Cosmos.Query.Core.ContinuationTokens; - using Microsoft.Azure.Cosmos.Query.Core.ExecutionContext; - using Microsoft.Azure.Cosmos.Query.Core.Monads; - using Microsoft.Azure.Cosmos.Query.Core.Pipeline; - - internal abstract partial class SkipDocumentQueryExecutionComponent : DocumentQueryExecutionComponentBase - { - private int skipCount; - - protected SkipDocumentQueryExecutionComponent(IDocumentQueryExecutionComponent source, long skipCount) - : base(source) - { - if (skipCount > int.MaxValue) - { - throw new ArgumentOutOfRangeException(nameof(skipCount)); - } - - this.skipCount = (int)skipCount; - } - - public static Task> MonadicCreateAsync( - ExecutionEnvironment executionEnvironment, - int offsetCount, - CosmosElement continuationToken, - Func>> monadicCreatePipelineStage) - { - return default; - } - } -} \ No newline at end of file diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/TakeDocumentQueryExecutionComponent.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/TakeDocumentQueryExecutionComponent.cs deleted file mode 100644 index 3a56499c2d..0000000000 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/TakeDocumentQueryExecutionComponent.cs +++ /dev/null @@ -1,39 +0,0 @@ -//------------------------------------------------------------ -// Copyright (c) Microsoft Corporation. All rights reserved. -//------------------------------------------------------------ - -namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.SkipTake -{ - using System; - using System.Threading.Tasks; - using Microsoft.Azure.Cosmos.CosmosElements; - using Microsoft.Azure.Cosmos.Query.Core.ContinuationTokens; - using Microsoft.Azure.Cosmos.Query.Core.ExecutionContext; - using Microsoft.Azure.Cosmos.Query.Core.Monads; - using Microsoft.Azure.Cosmos.Query.Core.Pipeline; - - internal abstract partial class TakeDocumentQueryExecutionComponent : DocumentQueryExecutionComponentBase - { - private int skipCount; - - protected TakeDocumentQueryExecutionComponent(IDocumentQueryExecutionComponent source, long skipCount) - : base(source) - { - if (skipCount > int.MaxValue) - { - throw new ArgumentOutOfRangeException(nameof(skipCount)); - } - - this.skipCount = (int)skipCount; - } - - public static Task> MonadicCreateAsync( - ExecutionEnvironment executionEnvironment, - int offsetCount, - CosmosElement continuationToken, - Func>> monadicCreatePipelineStage) - { - return default; - } - } -} \ No newline at end of file diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/CosmosQueryUnitTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/CosmosQueryUnitTests.cs index 9a43d090dc..0901f234e5 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/CosmosQueryUnitTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/CosmosQueryUnitTests.cs @@ -8,24 +8,14 @@ namespace Microsoft.Azure.Cosmos.Tests using System.Collections.Generic; using System.IO; using System.Net; - using System.Net.Http; using System.Threading; using System.Threading.Tasks; using Microsoft.Azure.Cosmos.CosmosElements; - using Microsoft.Azure.Cosmos.Diagnostics; using Microsoft.Azure.Cosmos.Query; using Microsoft.Azure.Cosmos.Query.Core; - using Microsoft.Azure.Cosmos.Query.Core.ContinuationTokens; - using Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent; - using Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.Aggregate; - using Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.Distinct; - using Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.SkipTake; using Microsoft.Azure.Cosmos.Query.Core.ExecutionContext; - using Microsoft.Azure.Cosmos.Query.Core.Metrics; using Microsoft.Azure.Cosmos.Query.Core.Monads; using Microsoft.Azure.Cosmos.Query.Core.Pipeline; - using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Aggregate; - using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Distinct; using Microsoft.Azure.Cosmos.Query.Core.QueryClient; using Microsoft.Azure.Cosmos.Query.Core.QueryPlan; using Microsoft.Azure.Cosmos.Resource.CosmosExceptions; @@ -187,39 +177,6 @@ public void VerifyItemQueryResponseCosmosElements() } } - [TestMethod] - public async Task TestCosmosQueryExecutionComponentOnFailure() - { - (IList components, QueryResponseCore response) setupContext = await this.GetAllExecutionComponents(); - - foreach (DocumentQueryExecutionComponentBase component in setupContext.components) - { - QueryResponseCore response = await component.DrainAsync(1, default(CancellationToken)); - Assert.AreEqual(setupContext.response, response); - } - } - - [TestMethod] - public async Task TestCosmosQueryExecutionComponentCancellation() - { - (IList components, QueryResponseCore response) setupContext = await this.GetAllExecutionComponents(); - CancellationTokenSource cancellationTokenSource = new CancellationTokenSource(); - cancellationTokenSource.Cancel(); - - foreach (DocumentQueryExecutionComponentBase component in setupContext.components) - { - try - { - QueryResponseCore response = await component.DrainAsync(1, cancellationTokenSource.Token); - Assert.Fail("cancellation token should have thrown an exception"); - } - catch (OperationCanceledException e) - { - Assert.IsNotNull(e.Message); - } - } - } - [TestMethod] public async Task TestCosmosQueryPartitionKeyDefinition() { @@ -299,106 +256,5 @@ public async Task TestCosmosQueryPartitionKeyDefinition() (tryGetPage.Exception as CosmosException).ToString().Contains(exceptionMessage), "response error message did not contain the proper substring."); } - - private async Task<(IList components, QueryResponseCore response)> GetAllExecutionComponents() - { - (Func>> func, QueryResponseCore response) = this.SetupBaseContextToVerifyFailureScenario(); - - List components = new List(); - List operators = new List() - { - AggregateOperator.Average, - AggregateOperator.Count, - AggregateOperator.Max, - AggregateOperator.Min, - AggregateOperator.Sum - }; - - components.Add((await AggregateDocumentQueryExecutionComponent.MonadicCreateAsync( - ExecutionEnvironment.Client, - operators.ToArray(), - new Dictionary() - { - { "test", AggregateOperator.Count } - }, - new List() { "test" }, - false, - null, - func)).Result); - - components.Add((await DistinctDocumentQueryExecutionComponent.MonadicCreateAsync( - ExecutionEnvironment.Client, - null, - func, - DistinctQueryType.Ordered)).Result); - - components.Add((await SkipDocumentQueryExecutionComponent.MonadicCreateAsync( - ExecutionEnvironment.Client, - 5, - null, - func)).Result); - - components.Add((await TakeDocumentQueryExecutionComponent.MonadicCreateAsync( - ExecutionEnvironment.Client, - 5, - null, - func)).Result); - - components.Add((await TakeDocumentQueryExecutionComponent.MonadicCreateAsync( - ExecutionEnvironment.Client, - 5, - null, - func)).Result); - - return (components, response); - } - - private (Func>>, QueryResponseCore) SetupBaseContextToVerifyFailureScenario() - { - CosmosDiagnosticsContext diagnosticsContext = new CosmosDiagnosticsContextCore(); - diagnosticsContext.AddDiagnosticsInternal( new PointOperationStatistics( - Guid.NewGuid().ToString(), - System.Net.HttpStatusCode.Unauthorized, - subStatusCode: SubStatusCodes.PartitionKeyMismatch, - responseTimeUtc: DateTime.UtcNow, - requestCharge: 4, - errorMessage: null, - method: HttpMethod.Post, - requestUri: "http://localhost.com", - requestSessionToken: null, - responseSessionToken: null)); - IReadOnlyCollection diagnostics = new List() - { - new QueryPageDiagnostics( - Guid.NewGuid(), - "0", - "SomeQueryMetricText", - "SomeIndexUtilText", - diagnosticsContext) - }; - - QueryResponseCore failure = QueryResponseCore.CreateFailure( - System.Net.HttpStatusCode.Unauthorized, - SubStatusCodes.PartitionKeyMismatch, - new CosmosException( - statusCode: HttpStatusCode.Unauthorized, - message: "Random error message", - subStatusCode: default, - stackTrace: default, - activityId: "TestActivityId", - requestCharge: 42.89, - retryAfter: default, - headers: default, - diagnosticsContext: default, - error: default, - innerException: default), - 42.89, - "TestActivityId"); - - Mock baseContext = new Mock(); - baseContext.Setup(x => x.DrainAsync(It.IsAny(), It.IsAny())).Returns(Task.FromResult(failure)); - Task> callBack(CosmosElement x) => Task.FromResult>(TryCatch.FromResult(baseContext.Object)); - return (callBack, failure); - } } } \ No newline at end of file From c98095367c23eee76b5748f33dc869822271ca9d Mon Sep 17 00:00:00 2001 From: Brandon Chong Date: Fri, 7 Aug 2020 19:35:31 -0700 Subject: [PATCH 60/85] moved continuation tokens around --- .../FeedRangeCompositeContinuation.cs | 2 +- .../src/Handler/PartitionKeyRangeHandler.cs | 2 - .../CompositeContinuationToken.cs | 129 ------------------ .../Core/ContinuationTokens/RangeRefStruct.cs | 52 ------- .../CosmosQueryExecutionContextFactory.cs | 3 +- .../Core/Pipeline/Distinct/DistinctMap.cs | 30 +--- .../Remote}/IPartitionedToken.cs | 4 +- .../OrderBy}/OrderByContinuationToken.cs | 18 +-- ...OrderByCrossPartitionQueryPipelineStage.cs | 16 +-- .../Parallel/ParallelContinuationToken.cs | 110 +++++++++++++++ ...arallelCrossPartitionQueryPipelineStage.cs | 39 +++--- .../Core/Pipeline/Remote/PartitionMapper.cs | 11 +- .../Tokens}/PipelineContinuationToken.cs | 2 +- .../Tokens}/PipelineContinuationTokenV0.cs | 2 +- .../Tokens}/PipelineContinuationTokenV1.cs | 2 +- .../Tokens}/PipelineContinuationTokenV1_1.cs | 2 +- .../DefaultDocumentQueryExecutionContext.cs | 2 +- .../StandByFeedContinuationToken.cs | 2 +- .../QueryResponses/StandByFeedIteratorCore.cs | 3 - .../src/Routing/CompositeContinuationToken.cs | 47 +++++++ .../src/Routing/PartitionRoutingHelper.cs | 1 - ...FeedRangeCompositeContinuationConverter.cs | 2 +- .../CosmosItemChangeFeedTests.cs | 2 +- .../FeedToken/ChangeFeedIteratorCoreTests.cs | 2 +- .../ChangeFeedResultSetIteratorTests.cs | 2 - .../FeedRange/FeedRangeContinuationTests.cs | 2 +- .../PartitionKeyRangeHandlerTests.cs | 5 - .../Query/ContinuationResumeLogicTests.cs | 70 ++++------ .../CompositeContinuationTokenTests.cs | 48 ------- .../OrderByContinuationTokenTests.cs | 13 +- .../ParallelContinuationTokenTests.cs | 37 +++++ .../PipelineContinuationTokenTests.cs | 2 +- ...ByCrossPartitionQueryPipelineStageTests.cs | 42 +++--- ...elCrossPartitionQueryPipelineStageTests.cs | 30 ++-- .../Routing/PartitionRoutingHelperTest.cs | 5 - .../StandByFeedContinuationTokenTests.cs | 2 +- .../Utils/MockCosmosUtil.cs | 2 - 37 files changed, 318 insertions(+), 427 deletions(-) delete mode 100644 Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/CompositeContinuationToken.cs delete mode 100644 Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/RangeRefStruct.cs rename Microsoft.Azure.Cosmos/src/Query/Core/{ContinuationTokens => Pipeline/Remote}/IPartitionedToken.cs (66%) rename Microsoft.Azure.Cosmos/src/Query/Core/{ContinuationTokens => Pipeline/Remote/OrderBy}/OrderByContinuationToken.cs (94%) create mode 100644 Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/Parallel/ParallelContinuationToken.cs rename Microsoft.Azure.Cosmos/src/Query/Core/{ContinuationTokens => Pipeline/Tokens}/PipelineContinuationToken.cs (98%) rename Microsoft.Azure.Cosmos/src/Query/Core/{ContinuationTokens => Pipeline/Tokens}/PipelineContinuationTokenV0.cs (95%) rename Microsoft.Azure.Cosmos/src/Query/Core/{ContinuationTokens => Pipeline/Tokens}/PipelineContinuationTokenV1.cs (98%) rename Microsoft.Azure.Cosmos/src/Query/Core/{ContinuationTokens => Pipeline/Tokens}/PipelineContinuationTokenV1_1.cs (98%) create mode 100644 Microsoft.Azure.Cosmos/src/Routing/CompositeContinuationToken.cs delete mode 100644 Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/ContinuationTokens/CompositeContinuationTokenTests.cs create mode 100644 Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/ContinuationTokens/ParallelContinuationTokenTests.cs diff --git a/Microsoft.Azure.Cosmos/src/FeedRange/Continuations/FeedRangeCompositeContinuation.cs b/Microsoft.Azure.Cosmos/src/FeedRange/Continuations/FeedRangeCompositeContinuation.cs index 12130389f3..cacd4d856f 100644 --- a/Microsoft.Azure.Cosmos/src/FeedRange/Continuations/FeedRangeCompositeContinuation.cs +++ b/Microsoft.Azure.Cosmos/src/FeedRange/Continuations/FeedRangeCompositeContinuation.cs @@ -10,8 +10,8 @@ namespace Microsoft.Azure.Cosmos using System.Net; using System.Threading; using System.Threading.Tasks; - using Microsoft.Azure.Cosmos.Query.Core.ContinuationTokens; using Microsoft.Azure.Cosmos.Query.Core.Monads; + using Microsoft.Azure.Cosmos.Routing; using Newtonsoft.Json; /// diff --git a/Microsoft.Azure.Cosmos/src/Handler/PartitionKeyRangeHandler.cs b/Microsoft.Azure.Cosmos/src/Handler/PartitionKeyRangeHandler.cs index 4403c8c183..88797161e0 100644 --- a/Microsoft.Azure.Cosmos/src/Handler/PartitionKeyRangeHandler.cs +++ b/Microsoft.Azure.Cosmos/src/Handler/PartitionKeyRangeHandler.cs @@ -11,8 +11,6 @@ namespace Microsoft.Azure.Cosmos.Handlers using System.Threading; using System.Threading.Tasks; using Microsoft.Azure.Cosmos.Common; - using Microsoft.Azure.Cosmos.Query; - using Microsoft.Azure.Cosmos.Query.Core.ContinuationTokens; using Microsoft.Azure.Cosmos.Routing; using Microsoft.Azure.Documents; using Microsoft.Azure.Documents.Routing; diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/CompositeContinuationToken.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/CompositeContinuationToken.cs deleted file mode 100644 index c8b9474562..0000000000 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/CompositeContinuationToken.cs +++ /dev/null @@ -1,129 +0,0 @@ -//------------------------------------------------------------ -// Copyright (c) Microsoft Corporation. All rights reserved. -//------------------------------------------------------------ - -namespace Microsoft.Azure.Cosmos.Query.Core.ContinuationTokens -{ - using System.Collections.Generic; - using Microsoft.Azure.Cosmos.CosmosElements; - using Microsoft.Azure.Cosmos.Json; - using Microsoft.Azure.Cosmos.Query.Core.Exceptions; - using Microsoft.Azure.Cosmos.Query.Core.Monads; - using Microsoft.Azure.Cosmos.Routing; - using Microsoft.Azure.Documents.Routing; - using Newtonsoft.Json; - - /// - /// A composite continuation token that has both backend continuation token and partition range information. - /// - internal sealed class CompositeContinuationToken : IPartitionedToken - { - private static class PropertyNames - { - public const string Token = "token"; - public const string Range = "range"; - - public const string Min = "min"; - public const string Max = "max"; - } - - [JsonProperty(PropertyNames.Token)] - public string Token - { - get; - set; - } - - [JsonProperty(PropertyNames.Range)] - [JsonConverter(typeof(RangeJsonConverter))] - public Documents.Routing.Range Range - { - get; - set; - } - - [JsonIgnore] - public Range PartitionRange => this.Range; - - public object ShallowCopy() - { - return this.MemberwiseClone(); - } - - public static CosmosElement ToCosmosElement(CompositeContinuationToken compositeContinuationToken) - { - CosmosElement token = compositeContinuationToken.Token == null ? (CosmosElement)CosmosNull.Create() : (CosmosElement)CosmosString.Create(compositeContinuationToken.Token); - return CosmosObject.Create( - new Dictionary() - { - { CompositeContinuationToken.PropertyNames.Token, token }, - { - CompositeContinuationToken.PropertyNames.Range, - CosmosObject.Create( - new Dictionary() - { - { PropertyNames.Min, CosmosString.Create(compositeContinuationToken.Range.Min) }, - { PropertyNames.Max, CosmosString.Create(compositeContinuationToken.Range.Max) } - }) - }, - }); - } - - public static TryCatch TryCreateFromCosmosElement(CosmosElement cosmosElement) - { - if (!(cosmosElement is CosmosObject cosmosObject)) - { - return TryCatch.FromException( - new MalformedContinuationTokenException($"{nameof(CompositeContinuationToken)} is not an object: {cosmosElement}")); - } - - if (!cosmosObject.TryGetValue(PropertyNames.Token, out CosmosElement rawToken)) - { - return TryCatch.FromException( - new MalformedContinuationTokenException($"{nameof(CompositeContinuationToken)} is missing field: '{PropertyNames.Token}': {cosmosElement}")); - } - - string token; - if (rawToken is CosmosString rawTokenString) - { - token = rawTokenString.Value; - } - else - { - token = null; - } - - if (!cosmosObject.TryGetValue(PropertyNames.Range, out CosmosObject rawRange)) - { - return TryCatch.FromException( - new MalformedContinuationTokenException($"{nameof(CompositeContinuationToken)} is missing field: '{PropertyNames.Range}': {cosmosElement}")); - } - - if (!rawRange.TryGetValue(PropertyNames.Min, out CosmosString rawMin)) - { - return TryCatch.FromException( - new MalformedContinuationTokenException($"{nameof(CompositeContinuationToken)} is missing field: '{PropertyNames.Min}': {cosmosElement}")); - } - - string min = rawMin.Value; - - if (!rawRange.TryGetValue(PropertyNames.Max, out CosmosString rawMax)) - { - return TryCatch.FromException( - new MalformedContinuationTokenException($"{nameof(CompositeContinuationToken)} is missing field: '{PropertyNames.Max}': {cosmosElement}")); - } - - string max = rawMax.Value; - - Documents.Routing.Range range = new Documents.Routing.Range(min, max, true, false); - - CompositeContinuationToken compositeContinuationToken = new CompositeContinuationToken() - { - Token = token, - Range = range, - }; - - return TryCatch.FromResult(compositeContinuationToken); - } - } -} diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/RangeRefStruct.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/RangeRefStruct.cs deleted file mode 100644 index bd9856605a..0000000000 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/RangeRefStruct.cs +++ /dev/null @@ -1,52 +0,0 @@ -// ------------------------------------------------------------ -// Copyright (c) Microsoft Corporation. All rights reserved. -// ------------------------------------------------------------ - -namespace Microsoft.Azure.Cosmos.Query.Core.ContinuationTokens -{ - using System; - using Microsoft.Azure.Cosmos.Json; - - internal ref struct RangeRefStruct - { - private const string MinPropertyName = "min"; - private const string MaxPropertyName = "max"; - - public RangeRefStruct(string min, string max) - { - if (min == null) - { - throw new ArgumentNullException(nameof(min)); - } - - if (max == null) - { - throw new ArgumentNullException(nameof(max)); - } - - this.Min = min; - this.Max = max; - } - - public string Min { get; } - public string Max { get; } - - public void WriteTo(IJsonWriter jsonWriter) - { - if (jsonWriter == null) - { - throw new ArgumentNullException(nameof(jsonWriter)); - } - - jsonWriter.WriteObjectStart(); - - jsonWriter.WriteFieldName(RangeRefStruct.MinPropertyName); - jsonWriter.WriteStringValue(this.Min); - - jsonWriter.WriteFieldName(RangeRefStruct.MaxPropertyName); - jsonWriter.WriteStringValue(this.Max); - - jsonWriter.WriteObjectEnd(); - } - } -} diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs index 1f9c3394bf..26094a84c6 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs @@ -12,14 +12,13 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionContext using Microsoft.Azure.Cosmos.CosmosElements; using Microsoft.Azure.Cosmos.Pagination; using Microsoft.Azure.Cosmos.Query.Core; - using Microsoft.Azure.Cosmos.Query.Core.ContinuationTokens; using Microsoft.Azure.Cosmos.Query.Core.Exceptions; using Microsoft.Azure.Cosmos.Query.Core.Monads; using Microsoft.Azure.Cosmos.Query.Core.Parser; using Microsoft.Azure.Cosmos.Query.Core.Pipeline; using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Distinct; - using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote; using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote.Parallel; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Tokens; using Microsoft.Azure.Cosmos.Query.Core.QueryClient; using Microsoft.Azure.Cosmos.Query.Core.QueryPlan; using Microsoft.Azure.Cosmos.SqlObjects; diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Distinct/DistinctMap.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Distinct/DistinctMap.cs index 2621f5d393..821c5ff4bf 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Distinct/DistinctMap.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Distinct/DistinctMap.cs @@ -6,9 +6,6 @@ namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.Distinct { using System; using Microsoft.Azure.Cosmos.CosmosElements; - using Microsoft.Azure.Cosmos.Json; - using Microsoft.Azure.Cosmos.Query.Core; - using Microsoft.Azure.Cosmos.Query.Core.ContinuationTokens; using Microsoft.Azure.Cosmos.Query.Core.Monads; /// @@ -30,28 +27,13 @@ internal abstract partial class DistinctMap /// The appropriate IDistinctMap. public static TryCatch TryCreate( DistinctQueryType distinctQueryType, - CosmosElement distinctMapContinuationToken) - { - TryCatch tryCreateDistinctMap; - switch (distinctQueryType) + CosmosElement distinctMapContinuationToken) => distinctQueryType switch { - case DistinctQueryType.None: - throw new ArgumentException("distinctQueryType can not be None. This part of code is not supposed to be reachable. Please contact support to resolve this issue."); - - case DistinctQueryType.Unordered: - tryCreateDistinctMap = UnorderdDistinctMap.TryCreate(distinctMapContinuationToken); - break; - - case DistinctQueryType.Ordered: - tryCreateDistinctMap = OrderedDistinctMap.TryCreate(distinctMapContinuationToken); - break; - - default: - throw new ArgumentException($"Unrecognized DistinctQueryType: {distinctQueryType}."); - } - - return tryCreateDistinctMap; - } + DistinctQueryType.None => throw new ArgumentException("distinctQueryType can not be None. This part of code is not supposed to be reachable. Please contact support to resolve this issue."), + DistinctQueryType.Unordered => UnorderdDistinctMap.TryCreate(distinctMapContinuationToken), + DistinctQueryType.Ordered => OrderedDistinctMap.TryCreate(distinctMapContinuationToken), + _ => throw new ArgumentException($"Unrecognized DistinctQueryType: {distinctQueryType}."), + }; /// /// Adds a JToken to this DistinctMap. diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/IPartitionedToken.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/IPartitionedToken.cs similarity index 66% rename from Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/IPartitionedToken.cs rename to Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/IPartitionedToken.cs index 2e81a39fa1..5575bbee66 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/IPartitionedToken.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/IPartitionedToken.cs @@ -2,10 +2,10 @@ // Copyright (c) Microsoft Corporation. All rights reserved. //------------------------------------------------------------ -namespace Microsoft.Azure.Cosmos.Query.Core.ContinuationTokens +namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote { internal interface IPartitionedToken { - Documents.Routing.Range PartitionRange { get; } + Documents.Routing.Range Range { get; } } } diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/OrderByContinuationToken.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderBy/OrderByContinuationToken.cs similarity index 94% rename from Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/OrderByContinuationToken.cs rename to Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderBy/OrderByContinuationToken.cs index c3482324be..5bb2a20825 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/OrderByContinuationToken.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderBy/OrderByContinuationToken.cs @@ -1,7 +1,8 @@ //------------------------------------------------------------ // Copyright (c) Microsoft Corporation. All rights reserved. //------------------------------------------------------------ -namespace Microsoft.Azure.Cosmos.Query.Core.ContinuationTokens + +namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote.OrderBy { using System; using System.Collections.Generic; @@ -11,6 +12,7 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ContinuationTokens using Microsoft.Azure.Cosmos.Query.Core.Exceptions; using Microsoft.Azure.Cosmos.Query.Core.Monads; using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote.OrderBy; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote.Parallel; using Microsoft.Azure.Documents.Routing; using Newtonsoft.Json; @@ -74,7 +76,7 @@ private static class PropertyNames /// The skip count (refer to property documentation). /// The filter (refer to property documentation). public OrderByContinuationToken( - CompositeContinuationToken compositeContinuationToken, + ParallelContinuationToken compositeContinuationToken, IReadOnlyList orderByItems, string rid, int skipCount, @@ -96,7 +98,7 @@ public OrderByContinuationToken( } //// filter is allowed to be null. - this.CompositeContinuationToken = compositeContinuationToken ?? throw new ArgumentNullException(nameof(compositeContinuationToken)); + this.ParallelContinuationToken = compositeContinuationToken ?? throw new ArgumentNullException(nameof(compositeContinuationToken)); this.OrderByItems = orderByItems ?? throw new ArgumentNullException(nameof(orderByItems)); this.Rid = rid; this.SkipCount = skipCount; @@ -112,7 +114,7 @@ public OrderByContinuationToken( /// ]]> /// [JsonProperty(PropertyNames.CompositeToken)] - public CompositeContinuationToken CompositeContinuationToken + public ParallelContinuationToken ParallelContinuationToken { get; } @@ -208,11 +210,11 @@ public string Filter } [JsonIgnore] - public Range PartitionRange => this.CompositeContinuationToken.Range; + public Range Range => this.ParallelContinuationToken.Range; public static CosmosElement ToCosmosElement(OrderByContinuationToken orderByContinuationToken) { - CosmosElement compositeContinuationToken = CompositeContinuationToken.ToCosmosElement(orderByContinuationToken.CompositeContinuationToken); + CosmosElement compositeContinuationToken = ParallelContinuationToken.ToCosmosElement(orderByContinuationToken.ParallelContinuationToken); List orderByItemsRaw = new List(); foreach (OrderByItem orderByItem in orderByContinuationToken.OrderByItems) { @@ -250,13 +252,13 @@ public static TryCatch TryCreateFromCosmosElement(Cosm new MalformedContinuationTokenException($"{nameof(OrderByContinuationToken)} is missing field: '{PropertyNames.CompositeToken}': {cosmosElement}")); } - TryCatch tryCompositeContinuation = CompositeContinuationToken.TryCreateFromCosmosElement(compositeContinuationTokenElement); + TryCatch tryCompositeContinuation = ParallelContinuationToken.TryCreateFromCosmosElement(compositeContinuationTokenElement); if (!tryCompositeContinuation.Succeeded) { return TryCatch.FromException(tryCompositeContinuation.Exception); } - CompositeContinuationToken compositeContinuationToken = tryCompositeContinuation.Result; + ParallelContinuationToken compositeContinuationToken = tryCompositeContinuation.Result; if (!cosmosObject.TryGetValue(PropertyNames.OrderByItems, out CosmosArray orderByItemsRaw)) { diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderBy/OrderByCrossPartitionQueryPipelineStage.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderBy/OrderByCrossPartitionQueryPipelineStage.cs index e4bc0b5b26..4d86a30cca 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderBy/OrderByCrossPartitionQueryPipelineStage.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderBy/OrderByCrossPartitionQueryPipelineStage.cs @@ -14,9 +14,9 @@ namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote.OrderBy using Microsoft.Azure.Cosmos.CosmosElements; using Microsoft.Azure.Cosmos.Pagination; using Microsoft.Azure.Cosmos.Query.Core.Collections; - using Microsoft.Azure.Cosmos.Query.Core.ContinuationTokens; using Microsoft.Azure.Cosmos.Query.Core.Exceptions; using Microsoft.Azure.Cosmos.Query.Core.Monads; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote.Parallel; using Microsoft.Azure.Documents; using static Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote.PartitionMapper; @@ -259,11 +259,9 @@ public async ValueTask MoveNextAsync() else { OrderByContinuationToken orderByContinuationToken = new OrderByContinuationToken( - new CompositeContinuationToken() - { - Token = currentEnumerator.StartOfPageState != null ? ((CosmosString)currentEnumerator.StartOfPageState.Value).Value : null, - Range = currentEnumerator.Range.ToRange(), - }, + new ParallelContinuationToken( + token: currentEnumerator.StartOfPageState != null ? ((CosmosString)currentEnumerator.StartOfPageState.Value).Value : null, + range: currentEnumerator.Range.ToRange()), orderByQueryResult.OrderByItems, orderByQueryResult.Rid, skipCount: 0, @@ -413,7 +411,7 @@ public static TryCatch MonadicCreate( range, pageSize, filter, - state: token?.CompositeContinuationToken?.Token != null ? new QueryState(CosmosString.Create(token.CompositeContinuationToken.Token)) : null); + state: token?.ParallelContinuationToken?.Token != null ? new QueryState(CosmosString.Create(token.ParallelContinuationToken.Token)) : null); enumeratorsAndTokens.Add((remoteEnumerator, token)); } @@ -706,11 +704,9 @@ private static (string leftFilter, string targetFilter, string rightFilter) GetF { return TryCatch<(bool, TryCatch)>.FromException( new MalformedContinuationTokenException( - $"Invalid Rid in the continuation token {continuationToken.CompositeContinuationToken.Token} for OrderBy~Context.")); + $"Invalid Rid in the continuation token {continuationToken.ParallelContinuationToken.Token} for OrderBy~Context.")); } - int itemToSkip = continuationToken.SkipCount; - // Throw away documents until it matches the item from the continuation token. if (!await enumerator.MoveNextAsync()) { diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/Parallel/ParallelContinuationToken.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/Parallel/ParallelContinuationToken.cs new file mode 100644 index 0000000000..a8e20448d9 --- /dev/null +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/Parallel/ParallelContinuationToken.cs @@ -0,0 +1,110 @@ +//------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +//------------------------------------------------------------ + +namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote.Parallel +{ + using System; + using System.Collections.Generic; + using Microsoft.Azure.Cosmos.CosmosElements; + using Microsoft.Azure.Cosmos.Query.Core.Exceptions; + using Microsoft.Azure.Cosmos.Query.Core.Monads; + using Microsoft.Azure.Documents.Routing; + + /// + /// A composite continuation token that has both backend continuation token and partition range information. + /// + internal sealed class ParallelContinuationToken : IPartitionedToken + { + private static class PropertyNames + { + public const string Token = "token"; + public const string Range = "range"; + + public const string Min = "min"; + public const string Max = "max"; + } + + public ParallelContinuationToken(string token, Range range) + { + this.Token = token; + this.Range = range ?? throw new ArgumentNullException(nameof(range)); + } + + public string Token { get; } + + public Range Range { get; } + + public static CosmosElement ToCosmosElement(ParallelContinuationToken parallelContinuationToken) + { + CosmosElement token = parallelContinuationToken.Token == null ? (CosmosElement)CosmosNull.Create() : (CosmosElement)CosmosString.Create(parallelContinuationToken.Token); + return CosmosObject.Create( + new Dictionary() + { + { ParallelContinuationToken.PropertyNames.Token, token }, + { + ParallelContinuationToken.PropertyNames.Range, + CosmosObject.Create( + new Dictionary() + { + { PropertyNames.Min, CosmosString.Create(parallelContinuationToken.Range.Min) }, + { PropertyNames.Max, CosmosString.Create(parallelContinuationToken.Range.Max) } + }) + }, + }); + } + + public static TryCatch TryCreateFromCosmosElement(CosmosElement cosmosElement) + { + if (!(cosmosElement is CosmosObject cosmosObject)) + { + return TryCatch.FromException( + new MalformedContinuationTokenException($"{nameof(ParallelContinuationToken)} is not an object: {cosmosElement}")); + } + + if (!cosmosObject.TryGetValue(PropertyNames.Token, out CosmosElement rawToken)) + { + return TryCatch.FromException( + new MalformedContinuationTokenException($"{nameof(ParallelContinuationToken)} is missing field: '{PropertyNames.Token}': {cosmosElement}")); + } + + string token; + if (rawToken is CosmosString rawTokenString) + { + token = rawTokenString.Value; + } + else + { + token = null; + } + + if (!cosmosObject.TryGetValue(PropertyNames.Range, out CosmosObject rawRange)) + { + return TryCatch.FromException( + new MalformedContinuationTokenException($"{nameof(ParallelContinuationToken)} is missing field: '{PropertyNames.Range}': {cosmosElement}")); + } + + if (!rawRange.TryGetValue(PropertyNames.Min, out CosmosString rawMin)) + { + return TryCatch.FromException( + new MalformedContinuationTokenException($"{nameof(ParallelContinuationToken)} is missing field: '{PropertyNames.Min}': {cosmosElement}")); + } + + string min = rawMin.Value; + + if (!rawRange.TryGetValue(PropertyNames.Max, out CosmosString rawMax)) + { + return TryCatch.FromException( + new MalformedContinuationTokenException($"{nameof(ParallelContinuationToken)} is missing field: '{PropertyNames.Max}': {cosmosElement}")); + } + + string max = rawMax.Value; + + Range range = new Documents.Routing.Range(min, max, true, false); + + ParallelContinuationToken parallelContinuationToken = new ParallelContinuationToken(token, range); + + return TryCatch.FromResult(parallelContinuationToken); + } + } +} diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/Parallel/ParallelCrossPartitionQueryPipelineStage.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/Parallel/ParallelCrossPartitionQueryPipelineStage.cs index f353b19364..b95b076a9d 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/Parallel/ParallelCrossPartitionQueryPipelineStage.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/Parallel/ParallelCrossPartitionQueryPipelineStage.cs @@ -10,7 +10,6 @@ namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote.Parallel using System.Threading.Tasks; using Microsoft.Azure.Cosmos.CosmosElements; using Microsoft.Azure.Cosmos.Pagination; - using Microsoft.Azure.Cosmos.Query.Core.ContinuationTokens; using Microsoft.Azure.Cosmos.Query.Core.Exceptions; using Microsoft.Azure.Cosmos.Query.Core.Monads; using Microsoft.Azure.Documents; @@ -53,24 +52,22 @@ public TryCatch Current } else { - List compositeContinuationTokens = new List(crossPartitionState.Value.Count); + List parallelContinuationTokens = new List(crossPartitionState.Value.Count); foreach ((PartitionKeyRange range, QueryState state) in crossPartitionState.Value) { - CompositeContinuationToken compositeContinuationToken = new CompositeContinuationToken() - { - Range = range.ToRange(), - Token = state != null ? ((CosmosString)state.Value).Value : null, - }; + ParallelContinuationToken parallelContinuationToken = new ParallelContinuationToken( + token: state != null ? ((CosmosString)state.Value).Value : null, + range: range.ToRange()); - compositeContinuationTokens.Add(compositeContinuationToken); + parallelContinuationTokens.Add(parallelContinuationToken); } - List cosmosElementContinuationTokens = compositeContinuationTokens - .Select(token => CompositeContinuationToken.ToCosmosElement(token)) + List cosmosElementContinuationTokens = parallelContinuationTokens + .Select(token => ParallelContinuationToken.ToCosmosElement(token)) .ToList(); - CosmosArray cosmosElementCompositeContinuationTokens = CosmosArray.Create(cosmosElementContinuationTokens); + CosmosArray cosmosElementParallelContinuationTokens = CosmosArray.Create(cosmosElementContinuationTokens); - queryState = new QueryState(cosmosElementCompositeContinuationTokens); + queryState = new QueryState(cosmosElementParallelContinuationTokens); } QueryPage crossPartitionQueryPage = new QueryPage( @@ -127,34 +124,34 @@ private static TryCatch> MonadicExtractState( return TryCatch>.FromResult(default); } - if (!(continuationToken is CosmosArray compositeContinuationTokenListRaw)) + if (!(continuationToken is CosmosArray parallelContinuationTokenListRaw)) { return TryCatch>.FromException( new MalformedContinuationTokenException( $"Invalid format for continuation token {continuationToken} for {nameof(ParallelCrossPartitionQueryPipelineStage)}")); } - if (compositeContinuationTokenListRaw.Count == 0) + if (parallelContinuationTokenListRaw.Count == 0) { return TryCatch>.FromException( new MalformedContinuationTokenException( $"Invalid format for continuation token {continuationToken} for {nameof(ParallelCrossPartitionQueryPipelineStage)}")); } - List compositeContinuationTokens = new List(); - foreach (CosmosElement compositeContinuationTokenRaw in compositeContinuationTokenListRaw) + List parallelContinuationTokens = new List(); + foreach (CosmosElement parallelContinuationTokenRaw in parallelContinuationTokenListRaw) { - TryCatch tryCreateCompositeContinuationToken = CompositeContinuationToken.TryCreateFromCosmosElement(compositeContinuationTokenRaw); - if (tryCreateCompositeContinuationToken.Failed) + TryCatch tryCreateParallelContinuationToken = ParallelContinuationToken.TryCreateFromCosmosElement(parallelContinuationTokenRaw); + if (tryCreateParallelContinuationToken.Failed) { return TryCatch>.FromException( - tryCreateCompositeContinuationToken.Exception); + tryCreateParallelContinuationToken.Exception); } - compositeContinuationTokens.Add(tryCreateCompositeContinuationToken.Result); + parallelContinuationTokens.Add(tryCreateParallelContinuationToken.Result); } - List<(PartitionKeyRange, QueryState)> rangesAndStates = compositeContinuationTokens + List<(PartitionKeyRange, QueryState)> rangesAndStates = parallelContinuationTokens .Select(token => ( new PartitionKeyRange() { diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/PartitionMapper.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/PartitionMapper.cs index 5b5a9d8b2a..7523cc3209 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/PartitionMapper.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/PartitionMapper.cs @@ -7,7 +7,6 @@ namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote using System; using System.Collections.Generic; using System.Linq; - using Microsoft.Azure.Cosmos.Query.Core.ContinuationTokens; using Microsoft.Azure.Cosmos.Query.Core.Exceptions; using Microsoft.Azure.Cosmos.Query.Core.Monads; using Microsoft.Azure.Documents; @@ -46,7 +45,7 @@ public static TryCatch> MonadicGetPartitionMa // Find the continuation token for the partition we left off on: PartitionedToken firstContinuationToken = partitionedContinuationTokens - .OrderBy((partitionedToken) => partitionedToken.PartitionRange.Min) + .OrderBy((partitionedToken) => partitionedToken.Range.Min) .First(); // Segment the ranges based off that: @@ -56,8 +55,8 @@ public static TryCatch> MonadicGetPartitionMa PartitionKeyRange firstContinuationRange = new PartitionKeyRange { - MinInclusive = firstContinuationToken.PartitionRange.Min, - MaxExclusive = firstContinuationToken.PartitionRange.Max + MinInclusive = firstContinuationToken.Range.Min, + MaxExclusive = firstContinuationToken.Range.Max }; int matchedIndex = sortedRanges.Span.BinarySearch( @@ -119,8 +118,8 @@ public static IReadOnlyDictionary MatchRang foreach (PartitionedToken partitionedToken in partitionedContinuationTokens) { // See if continuation token includes the range - if ((partitionKeyRange.MinInclusive.CompareTo(partitionedToken.PartitionRange.Min) >= 0) - && (partitionKeyRange.MaxExclusive.CompareTo(partitionedToken.PartitionRange.Max) <= 0)) + if ((partitionKeyRange.MinInclusive.CompareTo(partitionedToken.Range.Min) >= 0) + && (partitionKeyRange.MaxExclusive.CompareTo(partitionedToken.Range.Max) <= 0)) { partitionKeyRangeToToken[partitionKeyRange] = partitionedToken; break; diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/PipelineContinuationToken.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Tokens/PipelineContinuationToken.cs similarity index 98% rename from Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/PipelineContinuationToken.cs rename to Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Tokens/PipelineContinuationToken.cs index b587288192..f0e9c1b2f5 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/PipelineContinuationToken.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Tokens/PipelineContinuationToken.cs @@ -2,7 +2,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // ------------------------------------------------------------ -namespace Microsoft.Azure.Cosmos.Query.Core.ContinuationTokens +namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.Tokens { using System; using Microsoft.Azure.Cosmos.CosmosElements; diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/PipelineContinuationTokenV0.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Tokens/PipelineContinuationTokenV0.cs similarity index 95% rename from Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/PipelineContinuationTokenV0.cs rename to Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Tokens/PipelineContinuationTokenV0.cs index 39c318e667..40589b8b00 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/PipelineContinuationTokenV0.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Tokens/PipelineContinuationTokenV0.cs @@ -2,7 +2,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // ------------------------------------------------------------ -namespace Microsoft.Azure.Cosmos.Query.Core.ContinuationTokens +namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.Tokens { using System; using Microsoft.Azure.Cosmos.CosmosElements; diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/PipelineContinuationTokenV1.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Tokens/PipelineContinuationTokenV1.cs similarity index 98% rename from Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/PipelineContinuationTokenV1.cs rename to Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Tokens/PipelineContinuationTokenV1.cs index 9594d6f5dc..4c5601906a 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/PipelineContinuationTokenV1.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Tokens/PipelineContinuationTokenV1.cs @@ -2,7 +2,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // ------------------------------------------------------------ -namespace Microsoft.Azure.Cosmos.Query.Core.ContinuationTokens +namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.Tokens { using System; using System.Collections.Generic; diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/PipelineContinuationTokenV1_1.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Tokens/PipelineContinuationTokenV1_1.cs similarity index 98% rename from Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/PipelineContinuationTokenV1_1.cs rename to Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Tokens/PipelineContinuationTokenV1_1.cs index c2ab949206..0ea72d22d9 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/PipelineContinuationTokenV1_1.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Tokens/PipelineContinuationTokenV1_1.cs @@ -2,7 +2,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // ------------------------------------------------------------ -namespace Microsoft.Azure.Cosmos.Query.Core.ContinuationTokens +namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.Tokens { using System; using System.Collections.Generic; diff --git a/Microsoft.Azure.Cosmos/src/Query/v2Query/DefaultDocumentQueryExecutionContext.cs b/Microsoft.Azure.Cosmos/src/Query/v2Query/DefaultDocumentQueryExecutionContext.cs index 76ba25382e..8fa1324109 100644 --- a/Microsoft.Azure.Cosmos/src/Query/v2Query/DefaultDocumentQueryExecutionContext.cs +++ b/Microsoft.Azure.Cosmos/src/Query/v2Query/DefaultDocumentQueryExecutionContext.cs @@ -11,8 +11,8 @@ namespace Microsoft.Azure.Cosmos.Query using Microsoft.Azure.Cosmos; using Microsoft.Azure.Cosmos.Common; using Microsoft.Azure.Cosmos.CosmosElements; - using Microsoft.Azure.Cosmos.Query.Core.ContinuationTokens; using Microsoft.Azure.Cosmos.Query.Core.Metrics; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote.Parallel; using Microsoft.Azure.Cosmos.Query.Core.QueryPlan; using Microsoft.Azure.Cosmos.Routing; using Microsoft.Azure.Documents; diff --git a/Microsoft.Azure.Cosmos/src/Resource/QueryResponses/StandByFeedContinuationToken.cs b/Microsoft.Azure.Cosmos/src/Resource/QueryResponses/StandByFeedContinuationToken.cs index afec119dac..0c3cdd98da 100644 --- a/Microsoft.Azure.Cosmos/src/Resource/QueryResponses/StandByFeedContinuationToken.cs +++ b/Microsoft.Azure.Cosmos/src/Resource/QueryResponses/StandByFeedContinuationToken.cs @@ -9,7 +9,7 @@ namespace Microsoft.Azure.Cosmos.Query using System.Diagnostics; using System.Linq; using System.Threading.Tasks; - using Microsoft.Azure.Cosmos.Query.Core.ContinuationTokens; + using Microsoft.Azure.Cosmos.Routing; using Newtonsoft.Json; /// diff --git a/Microsoft.Azure.Cosmos/src/Resource/QueryResponses/StandByFeedIteratorCore.cs b/Microsoft.Azure.Cosmos/src/Resource/QueryResponses/StandByFeedIteratorCore.cs index 278c63bfeb..507b1700a1 100644 --- a/Microsoft.Azure.Cosmos/src/Resource/QueryResponses/StandByFeedIteratorCore.cs +++ b/Microsoft.Azure.Cosmos/src/Resource/QueryResponses/StandByFeedIteratorCore.cs @@ -6,14 +6,11 @@ namespace Microsoft.Azure.Cosmos { using System; using System.Net; - using System.Text; using System.Threading; using System.Threading.Tasks; using Microsoft.Azure.Cosmos.CosmosElements; using Microsoft.Azure.Cosmos.Query; - using Microsoft.Azure.Cosmos.Query.Core.ContinuationTokens; using Microsoft.Azure.Cosmos.Routing; - using Microsoft.Azure.Documents.Routing; /// /// Cosmos Stand-By Feed iterator implementing Composite Continuation Token diff --git a/Microsoft.Azure.Cosmos/src/Routing/CompositeContinuationToken.cs b/Microsoft.Azure.Cosmos/src/Routing/CompositeContinuationToken.cs new file mode 100644 index 0000000000..6ee0ac4525 --- /dev/null +++ b/Microsoft.Azure.Cosmos/src/Routing/CompositeContinuationToken.cs @@ -0,0 +1,47 @@ +//------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +//------------------------------------------------------------ + +namespace Microsoft.Azure.Cosmos.Routing +{ + using Microsoft.Azure.Documents.Routing; + using Newtonsoft.Json; + + /// + /// A composite continuation token that has both backend continuation token and partition range information. + /// + internal sealed class CompositeContinuationToken + { + private static class PropertyNames + { + public const string Token = "token"; + public const string Range = "range"; + + public const string Min = "min"; + public const string Max = "max"; + } + + [JsonProperty(PropertyNames.Token)] + public string Token + { + get; + set; + } + + [JsonProperty(PropertyNames.Range)] + [JsonConverter(typeof(RangeJsonConverter))] + public Documents.Routing.Range Range + { + get; + set; + } + + [JsonIgnore] + public Range PartitionRange => this.Range; + + public object ShallowCopy() + { + return this.MemberwiseClone(); + } + } +} diff --git a/Microsoft.Azure.Cosmos/src/Routing/PartitionRoutingHelper.cs b/Microsoft.Azure.Cosmos/src/Routing/PartitionRoutingHelper.cs index 9147aa979f..a8cb7906c1 100644 --- a/Microsoft.Azure.Cosmos/src/Routing/PartitionRoutingHelper.cs +++ b/Microsoft.Azure.Cosmos/src/Routing/PartitionRoutingHelper.cs @@ -14,7 +14,6 @@ namespace Microsoft.Azure.Cosmos.Routing using System.Threading.Tasks; using Microsoft.Azure.Cosmos.Core.Trace; using Microsoft.Azure.Cosmos.Query.Core; - using Microsoft.Azure.Cosmos.Query.Core.ContinuationTokens; using Microsoft.Azure.Cosmos.Query.Core.Monads; using Microsoft.Azure.Cosmos.Query.Core.QueryPlan; using Microsoft.Azure.Documents; diff --git a/Microsoft.Azure.Cosmos/src/Serializer/FeedRange/FeedRangeCompositeContinuationConverter.cs b/Microsoft.Azure.Cosmos/src/Serializer/FeedRange/FeedRangeCompositeContinuationConverter.cs index f764b6aa61..a9bc342244 100644 --- a/Microsoft.Azure.Cosmos/src/Serializer/FeedRange/FeedRangeCompositeContinuationConverter.cs +++ b/Microsoft.Azure.Cosmos/src/Serializer/FeedRange/FeedRangeCompositeContinuationConverter.cs @@ -6,7 +6,7 @@ namespace Microsoft.Azure.Cosmos { using System; using System.Collections.Generic; - using Microsoft.Azure.Cosmos.Query.Core.ContinuationTokens; + using Microsoft.Azure.Cosmos.Routing; using Newtonsoft.Json; using Newtonsoft.Json.Linq; diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CosmosItemChangeFeedTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CosmosItemChangeFeedTests.cs index 9d2d637788..27c871757c 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CosmosItemChangeFeedTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CosmosItemChangeFeedTests.cs @@ -12,7 +12,7 @@ namespace Microsoft.Azure.Cosmos.SDK.EmulatorTests using System.Threading; using System.Threading.Tasks; using Microsoft.Azure.Cosmos.Query; - using Microsoft.Azure.Cosmos.Query.Core.ContinuationTokens; + using Microsoft.Azure.Cosmos.Routing; using Microsoft.VisualStudio.TestTools.UnitTesting; using Newtonsoft.Json; diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/FeedToken/ChangeFeedIteratorCoreTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/FeedToken/ChangeFeedIteratorCoreTests.cs index 3ea3f8ace6..6c002a4e52 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/FeedToken/ChangeFeedIteratorCoreTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/FeedToken/ChangeFeedIteratorCoreTests.cs @@ -11,7 +11,7 @@ namespace Microsoft.Azure.Cosmos.EmulatorTests.FeedRanges using System.Net; using System.Threading; using System.Threading.Tasks; - using Microsoft.Azure.Cosmos.Query.Core.ContinuationTokens; + using Microsoft.Azure.Cosmos.Routing; using Microsoft.Azure.Cosmos.SDK.EmulatorTests; using Microsoft.VisualStudio.TestTools.UnitTesting; diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/ChangeFeedResultSetIteratorTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/ChangeFeedResultSetIteratorTests.cs index 825f449408..b36d331834 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/ChangeFeedResultSetIteratorTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/ChangeFeedResultSetIteratorTests.cs @@ -11,8 +11,6 @@ namespace Microsoft.Azure.Cosmos.Tests using System.Net; using System.Threading; using System.Threading.Tasks; - using Microsoft.Azure.Cosmos.Query; - using Microsoft.Azure.Cosmos.Query.Core.ContinuationTokens; using Microsoft.Azure.Cosmos.Routing; using Microsoft.Azure.Cosmos.Resource.CosmosExceptions; using Microsoft.VisualStudio.TestTools.UnitTesting; diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/FeedRange/FeedRangeContinuationTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/FeedRange/FeedRangeContinuationTests.cs index 26ea121f75..b5978d76a7 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/FeedRange/FeedRangeContinuationTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/FeedRange/FeedRangeContinuationTests.cs @@ -9,7 +9,7 @@ namespace Microsoft.Azure.Cosmos.Tests.FeedRange using System.Net; using System.Threading; using System.Threading.Tasks; - using Microsoft.Azure.Cosmos.Query.Core.ContinuationTokens; + using Microsoft.Azure.Cosmos.Routing; using Microsoft.VisualStudio.TestTools.UnitTesting; using Moq; using Newtonsoft.Json; diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/PartitionKeyRangeHandlerTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/PartitionKeyRangeHandlerTests.cs index 87d54e6991..d0396f2c14 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/PartitionKeyRangeHandlerTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/PartitionKeyRangeHandlerTests.cs @@ -11,12 +11,7 @@ namespace Microsoft.Azure.Cosmos.Tests using System.Net.Http; using System.Threading; using System.Threading.Tasks; - using Microsoft.Azure.Cosmos.Collections; - using Microsoft.Azure.Cosmos.Common; using Microsoft.Azure.Cosmos.Handlers; - using Microsoft.Azure.Cosmos.Internal; - using Microsoft.Azure.Cosmos.Query; - using Microsoft.Azure.Cosmos.Query.Core.ContinuationTokens; using Microsoft.Azure.Cosmos.Routing; using Microsoft.Azure.Documents; using Microsoft.Azure.Documents.Collections; diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/ContinuationResumeLogicTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/ContinuationResumeLogicTests.cs index 5a589533e0..d7b5004885 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/ContinuationResumeLogicTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/ContinuationResumeLogicTests.cs @@ -3,10 +3,10 @@ using System; using System.Collections.Generic; using System.Linq; - using Microsoft.Azure.Cosmos.Query.Core.ContinuationTokens; - using Microsoft.Azure.Cosmos.Query.Core.ExecutionContext; using Microsoft.Azure.Cosmos.Query.Core.Monads; using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote.Parallel; + using Microsoft.Azure.Cosmos.Routing; using Microsoft.Azure.Documents; using Microsoft.VisualStudio.TestTools.UnitTesting; using Newtonsoft.Json; @@ -25,15 +25,13 @@ public void TestMatchRangesTocontinuationTokens_OneToOne() Id = "0" }; - CompositeContinuationToken token = new CompositeContinuationToken() - { - Range = new Documents.Routing.Range( + ParallelContinuationToken token = new ParallelContinuationToken( + token: "asdf", + range: new Documents.Routing.Range( min: string.Empty, max: "FF", isMinInclusive: true, - isMaxInclusive: false), - Token = "asdf" - }; + isMaxInclusive: false)); IReadOnlyDictionary expectedMapping = new Dictionary() { @@ -43,7 +41,7 @@ public void TestMatchRangesTocontinuationTokens_OneToOne() ContinuationResumeLogicTests.RunMatchRangesToContinuationTokens( expectedMapping, new PartitionKeyRange[] { partitionKeyRange }, - new CompositeContinuationToken[] { token }); + new ParallelContinuationToken[] { token }); } [TestMethod] @@ -63,15 +61,13 @@ public void TestMatchRangesTocontinuationTokens_OneToMany() Id = "1" }; - CompositeContinuationToken token = new CompositeContinuationToken() - { - Range = new Documents.Routing.Range( + ParallelContinuationToken token = new ParallelContinuationToken( + token: "asdf", + range: new Documents.Routing.Range( min: string.Empty, max: "B", isMinInclusive: true, - isMaxInclusive: false), - Token = "asdf" - }; + isMaxInclusive: false)); IReadOnlyDictionary expectedMapping = new Dictionary() { @@ -82,7 +78,7 @@ public void TestMatchRangesTocontinuationTokens_OneToMany() ContinuationResumeLogicTests.RunMatchRangesToContinuationTokens( expectedMapping, new PartitionKeyRange[] { partitionKeyRange1, partitionKeyRange2 }, - new CompositeContinuationToken[] { token }); + new ParallelContinuationToken[] { token }); } [TestMethod] @@ -95,15 +91,13 @@ public void TestMatchRangesTocontinuationTokens_OneToNone() Id = "1" }; - CompositeContinuationToken token = new CompositeContinuationToken() - { - Range = new Documents.Routing.Range( + ParallelContinuationToken token = new ParallelContinuationToken( + token: "asdf", + range: new Documents.Routing.Range( min: "B", max: "C", isMinInclusive: true, - isMaxInclusive: false), - Token = "asdf" - }; + isMaxInclusive: false)); IReadOnlyDictionary expectedMapping = new Dictionary() { @@ -113,7 +107,7 @@ public void TestMatchRangesTocontinuationTokens_OneToNone() ContinuationResumeLogicTests.RunMatchRangesToContinuationTokens( expectedMapping, new PartitionKeyRange[] { partitionKeyRange }, - new CompositeContinuationToken[] { token }); + new ParallelContinuationToken[] { token }); } [TestMethod] @@ -150,15 +144,13 @@ public void TestTryGetInitializationInfo_ResumeLeftMostPartition() Id = "3" }; - CompositeContinuationToken token = new CompositeContinuationToken() - { - Range = new Documents.Routing.Range( + ParallelContinuationToken token = new ParallelContinuationToken( + token: "asdf", + range: new Documents.Routing.Range( min: string.Empty, max: "A", isMinInclusive: true, - isMaxInclusive: false), - Token = "asdf" - }; + isMaxInclusive: false)); IReadOnlyDictionary expectedMappingLeftPartitions = new Dictionary() { @@ -207,15 +199,13 @@ public void TestTryGetInitializationInfo_ResumeMiddlePartition() Id = "3" }; - CompositeContinuationToken token = new CompositeContinuationToken() - { - Range = new Documents.Routing.Range( + ParallelContinuationToken token = new ParallelContinuationToken( + token: "asdf", + range: new Documents.Routing.Range( min: "A", max: "B", isMinInclusive: true, - isMaxInclusive: false), - Token = "asdf" - }; + isMaxInclusive: false)); IReadOnlyDictionary expectedMappingLeftPartitions = new Dictionary() { @@ -264,15 +254,13 @@ public void TestTryGetInitializationInfo_ResumeRightPartition() Id = "3" }; - CompositeContinuationToken token = new CompositeContinuationToken() - { - Range = new Documents.Routing.Range( + ParallelContinuationToken token = new ParallelContinuationToken( + token: "asdf", + range: new Documents.Routing.Range( min: "B", max: "C", isMinInclusive: true, - isMaxInclusive: false), - Token = "asdf" - }; + isMaxInclusive: false)); IReadOnlyDictionary expectedMappingLeftPartitions = new Dictionary() { diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/ContinuationTokens/CompositeContinuationTokenTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/ContinuationTokens/CompositeContinuationTokenTests.cs deleted file mode 100644 index bb25d3975f..0000000000 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/ContinuationTokens/CompositeContinuationTokenTests.cs +++ /dev/null @@ -1,48 +0,0 @@ -//------------------------------------------------------------ -// Copyright (c) Microsoft Corporation. All rights reserved. -//------------------------------------------------------------ - -namespace Microsoft.Azure.Cosmos.Query -{ - using Microsoft.Azure.Cosmos.CosmosElements; - using Microsoft.Azure.Cosmos.Query.Core.ContinuationTokens; - using Microsoft.Azure.Cosmos.Query.Core.Monads; - using Newtonsoft.Json; - using VisualStudio.TestTools.UnitTesting; - - [TestClass] - public class CompositeContinuationTokenTests - { - [TestMethod] - public void TestNewtonsoftConvertsions() - { - string serializedContinuationToken = "[{\"token\":null,\"range\":{\"min\":\"05C1C9CD673398\",\"max\":\"05C1D9CD673398\"}}]"; - CompositeContinuationToken[] deserializedTokens = JsonConvert.DeserializeObject(serializedContinuationToken); - Assert.IsNotNull(deserializedTokens); - Assert.IsTrue(deserializedTokens.Length == 1); - CompositeContinuationToken deserializedToken = deserializedTokens[0]; - Assert.IsNull(deserializedToken.Token); - Assert.AreEqual("05C1C9CD673398", deserializedToken.Range.Min); - Assert.AreEqual("05C1D9CD673398", deserializedToken.Range.Max); - Assert.AreEqual(serializedContinuationToken, JsonConvert.SerializeObject(deserializedTokens, Formatting.None)); - } - - [TestMethod] - [DataRow(null, DisplayName = "null token")] - [DataRow("some token", DisplayName = "some token")] - public void TestRoundTripAsCosmosElement(string token) - { - CompositeContinuationToken compositeContinuationToken = new CompositeContinuationToken() - { - Token = token, - Range = new Documents.Routing.Range("asdf", "asdf", false, false), - }; - - CosmosElement cosmosElementToken = CompositeContinuationToken.ToCosmosElement(compositeContinuationToken); - TryCatch tryCompositeContinuationTokenFromCosmosElement = CompositeContinuationToken.TryCreateFromCosmosElement(cosmosElementToken); - Assert.IsTrue(tryCompositeContinuationTokenFromCosmosElement.Succeeded); - CompositeContinuationToken compositeContinuationTokenFromCosmosElement = tryCompositeContinuationTokenFromCosmosElement.Result; - Assert.AreEqual(JsonConvert.SerializeObject(compositeContinuationToken), JsonConvert.SerializeObject(compositeContinuationTokenFromCosmosElement)); - } - } -} diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/ContinuationTokens/OrderByContinuationTokenTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/ContinuationTokens/OrderByContinuationTokenTests.cs index 270548c081..bb8584da2c 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/ContinuationTokens/OrderByContinuationTokenTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/ContinuationTokens/OrderByContinuationTokenTests.cs @@ -9,11 +9,10 @@ namespace Microsoft.Azure.Cosmos.Query using System.Text; using Microsoft.Azure.Cosmos.CosmosElements; using Microsoft.Azure.Cosmos.CosmosElements.Numbers; - using Microsoft.Azure.Cosmos.Query.Core.ContinuationTokens; using Microsoft.Azure.Cosmos.Query.Core.Monads; using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote.OrderBy; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote.Parallel; using Microsoft.Azure.Documents.Routing; - using Newtonsoft.Json; using VisualStudio.TestTools.UnitTesting; [TestClass] @@ -22,11 +21,9 @@ public class OrderByContinuationTokenTests [TestMethod] public void TestRoundTripAsCosmosElement() { - CompositeContinuationToken compositeContinuationToken = new CompositeContinuationToken() - { - Token = "someToken", - Range = new Documents.Routing.Range("asdf", "asdf", false, false), - }; + ParallelContinuationToken parallelContinuationToken = new ParallelContinuationToken( + token: "someToken", + range: new Range("asdf", "asdf", false, false)); List orderByItems = new List() { @@ -39,7 +36,7 @@ public void TestRoundTripAsCosmosElement() int skipCount = 42; string filter = "someFilter"; OrderByContinuationToken orderByContinuationToken = new OrderByContinuationToken( - compositeContinuationToken, + parallelContinuationToken, orderByItems, rid, skipCount, diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/ContinuationTokens/ParallelContinuationTokenTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/ContinuationTokens/ParallelContinuationTokenTests.cs new file mode 100644 index 0000000000..7b39ca66ce --- /dev/null +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/ContinuationTokens/ParallelContinuationTokenTests.cs @@ -0,0 +1,37 @@ +//------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +//------------------------------------------------------------ + +namespace Microsoft.Azure.Cosmos.Query +{ + using Microsoft.Azure.Cosmos.CosmosElements; + using Microsoft.Azure.Cosmos.Query.Core.Monads; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote.Parallel; + using Newtonsoft.Json; + using VisualStudio.TestTools.UnitTesting; + + [TestClass] + public class ParallelContinuationTokenTests + { + [TestMethod] + [DataRow(null, DisplayName = "null token")] + [DataRow("some token", DisplayName = "some token")] + public void TestRoundTripAsCosmosElement(string token) + { + ParallelContinuationToken compositeContinuationToken = new ParallelContinuationToken( + token, + new Documents.Routing.Range("asdf", "asdf", false, false)); + + CosmosElement cosmosElementToken = ParallelContinuationToken.ToCosmosElement(compositeContinuationToken); + TryCatch tryCompositeContinuationTokenFromCosmosElement = ParallelContinuationToken.TryCreateFromCosmosElement(cosmosElementToken); + Assert.IsTrue(tryCompositeContinuationTokenFromCosmosElement.Succeeded); + ParallelContinuationToken compositeContinuationTokenFromCosmosElement = tryCompositeContinuationTokenFromCosmosElement.Result; + + Assert.AreEqual(compositeContinuationToken.Token, compositeContinuationTokenFromCosmosElement.Token); + Assert.AreEqual(compositeContinuationToken.Range.Min, compositeContinuationTokenFromCosmosElement.Range.Min); + Assert.AreEqual(compositeContinuationToken.Range.Max, compositeContinuationTokenFromCosmosElement.Range.Max); + Assert.AreEqual(compositeContinuationToken.Range.IsMinInclusive, compositeContinuationTokenFromCosmosElement.Range.IsMinInclusive); + Assert.AreEqual(compositeContinuationToken.Range.IsMaxInclusive, compositeContinuationTokenFromCosmosElement.Range.IsMaxInclusive); + } + } +} diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/ContinuationTokens/PipelineContinuationTokenTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/ContinuationTokens/PipelineContinuationTokenTests.cs index d7a7e500cd..08cae02df8 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/ContinuationTokens/PipelineContinuationTokenTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/ContinuationTokens/PipelineContinuationTokenTests.cs @@ -7,8 +7,8 @@ namespace Microsoft.Azure.Cosmos.Query using System.Collections.Generic; using System.Xml; using Microsoft.Azure.Cosmos.CosmosElements; - using Microsoft.Azure.Cosmos.Query.Core.ContinuationTokens; using Microsoft.Azure.Cosmos.Query.Core.Monads; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Tokens; using Microsoft.Azure.Cosmos.Test.BaselineTest; using VisualStudio.TestTools.UnitTesting; diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/OrderByCrossPartitionQueryPipelineStageTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/OrderByCrossPartitionQueryPipelineStageTests.cs index 54a1d13237..f311851844 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/OrderByCrossPartitionQueryPipelineStageTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/OrderByCrossPartitionQueryPipelineStageTests.cs @@ -10,11 +10,11 @@ namespace Microsoft.Azure.Cosmos.Tests.Query.Pipeline using Microsoft.Azure.Cosmos.CosmosElements; using Microsoft.Azure.Cosmos.Pagination; using Microsoft.Azure.Cosmos.Query.Core; - using Microsoft.Azure.Cosmos.Query.Core.ContinuationTokens; using Microsoft.Azure.Cosmos.Query.Core.Exceptions; using Microsoft.Azure.Cosmos.Query.Core.Monads; using Microsoft.Azure.Cosmos.Query.Core.Pipeline; using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote.OrderBy; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote.Parallel; using Microsoft.Azure.Cosmos.Tests.Pagination; using Microsoft.Azure.Documents; using Microsoft.VisualStudio.TestTools.UnitTesting; @@ -80,7 +80,7 @@ public void MonadicCreate_EmptyArrayContinuationToken() } [TestMethod] - public void MonadicCreate_NonCompositeContinuationToken() + public void MonadicCreate_NonParallelContinuationToken() { Mock mockDocumentContainer = new Mock(); @@ -103,14 +103,12 @@ public void MonadicCreate_SingleOrderByContinuationToken() { Mock mockDocumentContainer = new Mock(); - CompositeContinuationToken compositeContinuationToken = new CompositeContinuationToken() - { - Range = new Documents.Routing.Range("A", "B", true, false), - Token = "asdf", - }; + ParallelContinuationToken parallelContinuationToken = new ParallelContinuationToken( + token: "asdf", + range: new Documents.Routing.Range("A", "B", true, false)); OrderByContinuationToken orderByContinuationToken = new OrderByContinuationToken( - compositeContinuationToken, + parallelContinuationToken, new List() { new OrderByItem(CosmosObject.Create(new Dictionary() { { "item", CosmosString.Create("asdf") } })) }, rid: "rid", skipCount: 42, @@ -139,33 +137,27 @@ public void MonadicCreate_MultipleOrderByContinuationToken() { Mock mockDocumentContainer = new Mock(); - CompositeContinuationToken token = new CompositeContinuationToken() - { - Range = new Documents.Routing.Range("A", "B", true, false), - Token = "asdf", - }; + ParallelContinuationToken token = new ParallelContinuationToken( + token: "asdf", + range: new Documents.Routing.Range("A", "B", true, false)); - CompositeContinuationToken compositeContinuationToken1 = new CompositeContinuationToken() - { - Range = new Documents.Routing.Range("A", "B", true, false), - Token = "asdf", - }; + ParallelContinuationToken parallelContinuationToken1 = new ParallelContinuationToken( + token: "asdf", + range: new Documents.Routing.Range("A", "B", true, false)); OrderByContinuationToken orderByContinuationToken1 = new OrderByContinuationToken( - compositeContinuationToken1, + parallelContinuationToken1, new List() { new OrderByItem(CosmosObject.Create(new Dictionary() { { "item", CosmosString.Create("asdf") } })) }, rid: "rid", skipCount: 42, filter: "filter"); - CompositeContinuationToken compositeContinuationToken2 = new CompositeContinuationToken() - { - Range = new Documents.Routing.Range("B", "C", true, false), - Token = "asdf", - }; + ParallelContinuationToken parallelContinuationToken2 = new ParallelContinuationToken( + token: "asdf", + range: new Documents.Routing.Range("B", "C", true, false)); OrderByContinuationToken orderByContinuationToken2 = new OrderByContinuationToken( - compositeContinuationToken2, + parallelContinuationToken2, new List() { new OrderByItem(CosmosObject.Create(new Dictionary() { { "item", CosmosString.Create("asdf") } })) }, rid: "rid", skipCount: 42, 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 84dc537b3d..a9f61b0574 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 @@ -9,11 +9,9 @@ namespace Microsoft.Azure.Cosmos.Tests.Query.Pipeline using Microsoft.Azure.Cosmos.CosmosElements; using Microsoft.Azure.Cosmos.Pagination; using Microsoft.Azure.Cosmos.Query.Core; - using Microsoft.Azure.Cosmos.Query.Core.ContinuationTokens; using Microsoft.Azure.Cosmos.Query.Core.Exceptions; using Microsoft.Azure.Cosmos.Query.Core.Monads; using Microsoft.Azure.Cosmos.Query.Core.Pipeline; - using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote; using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote.Parallel; using Microsoft.Azure.Cosmos.Tests.Pagination; using Microsoft.Azure.Documents; @@ -65,7 +63,7 @@ public void MonadicCreate_EmptyArrayContinuationToken() } [TestMethod] - public void MonadicCreate_NonCompositeContinuationToken() + public void MonadicCreate_NonParallelContinuationToken() { Mock mockDocumentContainer = new Mock(); @@ -79,34 +77,30 @@ public void MonadicCreate_NonCompositeContinuationToken() } [TestMethod] - public void MonadicCreate_SingleCompositeContinuationToken() + public void MonadicCreate_SingleParallelContinuationToken() { Mock mockDocumentContainer = new Mock(); - CompositeContinuationToken token = new CompositeContinuationToken() - { - Range = new Documents.Routing.Range("A", "B", true, false), - Token = "asdf", - }; + ParallelContinuationToken token = new ParallelContinuationToken( + token: "asdf", + range: new Documents.Routing.Range("A", "B", true, false)); TryCatch monadicCreate = ParallelCrossPartitionQueryPipelineStage.MonadicCreate( documentContainer: mockDocumentContainer.Object, sqlQuerySpec: new SqlQuerySpec("SELECT * FROM c"), pageSize: 10, - continuationToken: CosmosArray.Create(new List() { CompositeContinuationToken.ToCosmosElement(token) })); + continuationToken: CosmosArray.Create(new List() { ParallelContinuationToken.ToCosmosElement(token) })); Assert.IsTrue(monadicCreate.Succeeded); } [TestMethod] - public void MonadicCreate_MultipleCompositeContinuationToken() + public void MonadicCreate_MultipleParallelContinuationToken() { Mock mockDocumentContainer = new Mock(); - CompositeContinuationToken token = new CompositeContinuationToken() - { - Range = new Documents.Routing.Range("A", "B", true, false), - Token = "asdf", - }; + ParallelContinuationToken token = new ParallelContinuationToken( + token: "asdf", + range: new Documents.Routing.Range("A", "B", true, false)); TryCatch monadicCreate = ParallelCrossPartitionQueryPipelineStage.MonadicCreate( documentContainer: mockDocumentContainer.Object, @@ -115,8 +109,8 @@ public void MonadicCreate_MultipleCompositeContinuationToken() continuationToken: CosmosArray.Create( new List() { - CompositeContinuationToken.ToCosmosElement(token), - CompositeContinuationToken.ToCosmosElement(token) + ParallelContinuationToken.ToCosmosElement(token), + ParallelContinuationToken.ToCosmosElement(token) })); Assert.IsTrue(monadicCreate.Succeeded); } diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Routing/PartitionRoutingHelperTest.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Routing/PartitionRoutingHelperTest.cs index f4d24c2699..a192358e3b 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Routing/PartitionRoutingHelperTest.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Routing/PartitionRoutingHelperTest.cs @@ -6,19 +6,14 @@ namespace Microsoft.Azure.Cosmos.Tests.Routing { using System; using System.Collections.Generic; - using System.Collections.Specialized; using System.IO; using System.Linq; using System.Threading.Tasks; using Microsoft.VisualStudio.TestTools.UnitTesting; - using Microsoft.Azure.Cosmos.Collections; - using Microsoft.Azure.Cosmos.Query; using Newtonsoft.Json; - using Microsoft.Azure.Cosmos.Internal; using Microsoft.Azure.Documents.Routing; using Microsoft.Azure.Documents; using Microsoft.Azure.Documents.Collections; - using Microsoft.Azure.Cosmos.Query.Core.ContinuationTokens; using Microsoft.Azure.Cosmos.Routing; /// diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/StandByFeedContinuationTokenTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/StandByFeedContinuationTokenTests.cs index 66564e856c..9225b8b870 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/StandByFeedContinuationTokenTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/StandByFeedContinuationTokenTests.cs @@ -9,7 +9,7 @@ namespace Microsoft.Azure.Cosmos using System.Linq; using System.Threading.Tasks; using Microsoft.Azure.Cosmos.Query; - using Microsoft.Azure.Cosmos.Query.Core.ContinuationTokens; + using Microsoft.Azure.Cosmos.Routing; using Microsoft.VisualStudio.TestTools.UnitTesting; using Newtonsoft.Json; diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Utils/MockCosmosUtil.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Utils/MockCosmosUtil.cs index da0098e7ac..5c7c6cbcf5 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Utils/MockCosmosUtil.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Utils/MockCosmosUtil.cs @@ -10,8 +10,6 @@ namespace Microsoft.Azure.Cosmos.Tests using System.Threading.Tasks; using Microsoft.Azure.Cosmos; using Microsoft.Azure.Cosmos.Fluent; - using Microsoft.Azure.Cosmos.Query; - using Microsoft.Azure.Cosmos.Query.Core.ContinuationTokens; using Microsoft.Azure.Cosmos.Routing; using Microsoft.Azure.Documents; using Microsoft.Azure.Documents.Collections; From c4953c603521ed012f94f912b396d17ae303d858 Mon Sep 17 00:00:00 2001 From: Brandon Chong Date: Sat, 8 Aug 2020 15:26:11 -0700 Subject: [PATCH 61/85] got basic order by working end to end --- .../src/Handler/RequestInvokerHandler.cs | 3 - .../CosmosQueryExecutionContextFactory.cs | 1 + .../Query/Core/Pipeline/PipelineFactory.cs | 18 ++++- .../Pipeline/Remote/OrderBy/OrderByColumn.cs | 20 ++++++ ...OrderByCrossPartitionQueryPipelineStage.cs | 12 ---- .../SkipEmptyPageQueryPipelineStage.cs | 67 +++++++++++++++++++ .../Query/Pipeline/FactoryTests.cs | 5 +- .../Query/Pipeline/FullPipelineTests.cs | 2 + ...ByCrossPartitionQueryPipelineStageTests.cs | 32 ++++----- 9 files changed, 125 insertions(+), 35 deletions(-) create mode 100644 Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderBy/OrderByColumn.cs create mode 100644 Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/SkipEmptyPageQueryPipelineStage.cs diff --git a/Microsoft.Azure.Cosmos/src/Handler/RequestInvokerHandler.cs b/Microsoft.Azure.Cosmos/src/Handler/RequestInvokerHandler.cs index 9fcf5cdf7a..decdefde67 100644 --- a/Microsoft.Azure.Cosmos/src/Handler/RequestInvokerHandler.cs +++ b/Microsoft.Azure.Cosmos/src/Handler/RequestInvokerHandler.cs @@ -125,9 +125,6 @@ public virtual async Task SendAsync( // user request might span multiple backend requests. // This will still have a single request id for retry scenarios ActivityScope activityScope = ActivityScope.CreateIfDefaultActivityId(); - Debug.Assert(activityScope != null && - (operationType != OperationType.SqlQuery || operationType != OperationType.Query || operationType != OperationType.QueryPlan), - "There should be an activity id already set"); try { diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs index 26094a84c6..0b88c638a3 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs @@ -430,6 +430,7 @@ private static TryCatch TryCreateSpecializedDocumentQueryEx executionEnvironment: inputParameters.ExecutionEnvironment, documentContainer: documentContainer, sqlQuerySpec: inputParameters.SqlQuerySpec, + targetRanges: targetRanges, queryInfo: partitionedQueryExecutionInfo.QueryInfo, pageSize: (int)optimalPageSize, requestContinuationToken: inputParameters.InitialUserContinuationToken); diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/PipelineFactory.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/PipelineFactory.cs index 0914388cb2..3cfadb6b40 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/PipelineFactory.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/PipelineFactory.cs @@ -5,16 +5,20 @@ namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline { using System; + using System.Collections.Generic; + using System.Linq; using Microsoft.Azure.Cosmos.CosmosElements; using Microsoft.Azure.Cosmos.Pagination; using Microsoft.Azure.Cosmos.Query.Core.Monads; using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Aggregate; using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Distinct; using Microsoft.Azure.Cosmos.Query.Core.Pipeline.GroupBy; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote.OrderBy; using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote.Parallel; using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Skip; using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Take; using Microsoft.Azure.Cosmos.Query.Core.QueryPlan; + using Microsoft.Azure.Documents; internal static class PipelineFactory { @@ -22,6 +26,7 @@ public static TryCatch MonadicCreate( ExecutionEnvironment executionEnvironment, IDocumentContainer documentContainer, SqlQuerySpec sqlQuerySpec, + IReadOnlyList targetRanges, QueryInfo queryInfo, int pageSize, CosmosElement requestContinuationToken) @@ -29,7 +34,15 @@ public static TryCatch MonadicCreate( MonadicCreatePipelineStage monadicCreatePipelineStage; if (queryInfo.HasOrderBy) { - throw new NotImplementedException(); + monadicCreatePipelineStage = (continuationToken) => OrderByCrossPartitionQueryPipelineStage.MonadicCreate( + documentContainer: documentContainer, + sqlQuerySpec: sqlQuerySpec, + targetRanges: targetRanges, + orderByColumns: queryInfo + .OrderByExpressions + .Zip(queryInfo.OrderBy, (expression, sortOrder) => new OrderByColumn(expression, sortOrder)).ToList(), + pageSize: pageSize, + continuationToken: requestContinuationToken); } else { @@ -105,7 +118,8 @@ public static TryCatch MonadicCreate( monadicCreateSourceStage); } - return monadicCreatePipelineStage(requestContinuationToken); + return monadicCreatePipelineStage(requestContinuationToken) + .Try(onSuccess: (stage) => new SkipEmptyPageQueryPipelineStage(stage)); } } } diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderBy/OrderByColumn.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderBy/OrderByColumn.cs new file mode 100644 index 0000000000..46253c76dc --- /dev/null +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderBy/OrderByColumn.cs @@ -0,0 +1,20 @@ +// ------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +// ------------------------------------------------------------ + +namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote.OrderBy +{ + using System; + + internal readonly struct OrderByColumn + { + public OrderByColumn(string expression, SortOrder sortOrder) + { + this.Expression = expression ?? throw new ArgumentNullException(nameof(expression)); + this.SortOrder = sortOrder; + } + + public string Expression { get; } + public SortOrder SortOrder { get; } + } +} diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderBy/OrderByCrossPartitionQueryPipelineStage.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderBy/OrderByCrossPartitionQueryPipelineStage.cs index 4d86a30cca..3c8c997d9e 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderBy/OrderByCrossPartitionQueryPipelineStage.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderBy/OrderByCrossPartitionQueryPipelineStage.cs @@ -783,17 +783,5 @@ private static bool IsSplitException(Exception exeception) && (cosmosException.StatusCode == HttpStatusCode.Gone) && (cosmosException.SubStatusCode == (int)Documents.SubStatusCodes.PartitionKeyRangeGone); } - - public readonly struct OrderByColumn - { - public OrderByColumn(string expression, SortOrder sortOrder) - { - this.Expression = expression ?? throw new ArgumentNullException(nameof(expression)); - this.SortOrder = sortOrder; - } - - public string Expression { get; } - public SortOrder SortOrder { get; } - } } } diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/SkipEmptyPageQueryPipelineStage.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/SkipEmptyPageQueryPipelineStage.cs new file mode 100644 index 0000000000..2295fe01bd --- /dev/null +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/SkipEmptyPageQueryPipelineStage.cs @@ -0,0 +1,67 @@ +// ------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +// ------------------------------------------------------------ + +namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline +{ + using System; + using System.Threading.Tasks; + using Microsoft.Azure.Cosmos.Query.Core.Monads; + + internal sealed class SkipEmptyPageQueryPipelineStage : IQueryPipelineStage + { + private readonly IQueryPipelineStage inputStage; + private double cumulativeRequestCharge; + private long cumulativeResponseLengthInBytes; + + public SkipEmptyPageQueryPipelineStage(IQueryPipelineStage inputStage) + { + this.inputStage = inputStage ?? throw new ArgumentNullException(nameof(inputStage)); + } + + public TryCatch Current { get; private set; } + + public ValueTask DisposeAsync() => this.inputStage.DisposeAsync(); + + public async ValueTask MoveNextAsync() + { + await this.inputStage.MoveNextAsync(); + TryCatch tryGetSourcePage = this.inputStage.Current; + if (tryGetSourcePage.Failed) + { + this.Current = tryGetSourcePage; + return true; + } + + QueryPage sourcePage = tryGetSourcePage.Result; + if (sourcePage.Documents.Count == 0) + { + this.cumulativeRequestCharge += sourcePage.RequestCharge; + this.cumulativeResponseLengthInBytes += sourcePage.ResponseLengthInBytes; + return await this.MoveNextAsync(); + } + + QueryPage cumulativeQueryPage; + if (this.cumulativeRequestCharge != 0) + { + cumulativeQueryPage = new QueryPage( + documents: sourcePage.Documents, + requestCharge: sourcePage.RequestCharge + this.cumulativeRequestCharge, + activityId: sourcePage.ActivityId, + responseLengthInBytes: sourcePage.ResponseLengthInBytes + this.cumulativeResponseLengthInBytes, + cosmosQueryExecutionInfo: sourcePage.CosmosQueryExecutionInfo, + disallowContinuationTokenMessage: sourcePage.DisallowContinuationTokenMessage, + state: sourcePage.State); + this.cumulativeRequestCharge = 0; + this.cumulativeResponseLengthInBytes = 0; + } + else + { + cumulativeQueryPage = sourcePage; + } + + this.Current = TryCatch.FromResult(cumulativeQueryPage); + return true; + } + } +} 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 87adf78407..68e831beaa 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 @@ -4,14 +4,14 @@ namespace Microsoft.Azure.Cosmos.Tests.Query.Pipeline { + using System.Collections.Generic; using Microsoft.Azure.Cosmos.Pagination; 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.Remote; using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote.Parallel; using Microsoft.Azure.Cosmos.Query.Core.QueryPlan; + using Microsoft.Azure.Documents; using Microsoft.VisualStudio.TestTools.UnitTesting; using Moq; @@ -27,6 +27,7 @@ public void TestCreate() ExecutionEnvironment.Compute, documentContainer: mockDocumentContainer.Object, sqlQuerySpec: new SqlQuerySpec("SELECT * FROM c"), + targetRanges: new List(), queryInfo: new QueryInfo() { }, pageSize: 10, requestContinuationToken: default); 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 f83308f425..a1f2e6315d 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 @@ -7,6 +7,7 @@ namespace Microsoft.Azure.Cosmos.Tests.Query.Pipeline using System.Collections.Generic; using System.Collections.ObjectModel; using System.Linq; + using System.Threading; using System.Threading.Tasks; using Microsoft.Azure.Cosmos.CosmosElements; using Microsoft.Azure.Cosmos.Pagination; @@ -170,6 +171,7 @@ private static IQueryPipelineStage CreatePipeline(IDocumentContainer documentCon ExecutionEnvironment.Compute, documentContainer, new SqlQuerySpec(query), + documentContainer.GetFeedRangesAsync(default(CancellationToken)).Result, GetQueryPlan(query), pageSize: 10, requestContinuationToken: state); diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/OrderByCrossPartitionQueryPipelineStageTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/OrderByCrossPartitionQueryPipelineStageTests.cs index f311851844..cce832a4ac 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/OrderByCrossPartitionQueryPipelineStageTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/OrderByCrossPartitionQueryPipelineStageTests.cs @@ -32,9 +32,9 @@ public void MonadicCreate_NullContinuationToken() documentContainer: mockDocumentContainer.Object, sqlQuerySpec: new SqlQuerySpec("SELECT * FROM c ORDER BY c._ts"), targetRanges: new List() { new PartitionKeyRange() }, - orderByColumns: new List() + orderByColumns: new List() { - new OrderByCrossPartitionQueryPipelineStage.OrderByColumn("_ts", SortOrder.Ascending) + new OrderByColumn("_ts", SortOrder.Ascending) }, pageSize: 10, continuationToken: null); @@ -50,9 +50,9 @@ public void MonadicCreate_NonCosmosArrayContinuationToken() documentContainer: mockDocumentContainer.Object, sqlQuerySpec: new SqlQuerySpec("SELECT * FROM c ORDER BY c._ts"), targetRanges: new List() { new PartitionKeyRange() }, - orderByColumns: new List() + orderByColumns: new List() { - new OrderByCrossPartitionQueryPipelineStage.OrderByColumn("_ts", SortOrder.Ascending) + new OrderByColumn("_ts", SortOrder.Ascending) }, pageSize: 10, continuationToken: CosmosObject.Create(new Dictionary())); @@ -69,9 +69,9 @@ public void MonadicCreate_EmptyArrayContinuationToken() documentContainer: mockDocumentContainer.Object, sqlQuerySpec: new SqlQuerySpec("SELECT * FROM c ORDER BY c._ts"), targetRanges: new List() { new PartitionKeyRange() }, - orderByColumns: new List() + orderByColumns: new List() { - new OrderByCrossPartitionQueryPipelineStage.OrderByColumn("_ts", SortOrder.Ascending) + new OrderByColumn("_ts", SortOrder.Ascending) }, pageSize: 10, continuationToken: CosmosArray.Create(new List())); @@ -88,9 +88,9 @@ public void MonadicCreate_NonParallelContinuationToken() documentContainer: mockDocumentContainer.Object, sqlQuerySpec: new SqlQuerySpec("SELECT * FROM c ORDER BY c._ts"), targetRanges: new List() { new PartitionKeyRange() }, - orderByColumns: new List() + orderByColumns: new List() { - new OrderByCrossPartitionQueryPipelineStage.OrderByColumn("_ts", SortOrder.Ascending) + new OrderByColumn("_ts", SortOrder.Ascending) }, pageSize: 10, continuationToken: CosmosArray.Create(new List() { CosmosString.Create("asdf") })); @@ -118,9 +118,9 @@ public void MonadicCreate_SingleOrderByContinuationToken() documentContainer: mockDocumentContainer.Object, sqlQuerySpec: new SqlQuerySpec("SELECT * FROM c ORDER BY c._ts"), targetRanges: new List() { new PartitionKeyRange() { Id = "0", MinInclusive = "A", MaxExclusive = "B" } }, - orderByColumns: new List() + orderByColumns: new List() { - new OrderByCrossPartitionQueryPipelineStage.OrderByColumn("_ts", SortOrder.Ascending) + new OrderByColumn("_ts", SortOrder.Ascending) }, pageSize: 10, continuationToken: CosmosString.Create( @@ -171,9 +171,9 @@ public void MonadicCreate_MultipleOrderByContinuationToken() new PartitionKeyRange() { Id = "0", MinInclusive = "A", MaxExclusive = "B" }, new PartitionKeyRange() { Id = "1", MinInclusive = "B", MaxExclusive = "C" } }, - orderByColumns: new List() + orderByColumns: new List() { - new OrderByCrossPartitionQueryPipelineStage.OrderByColumn("_ts", SortOrder.Ascending) + new OrderByColumn("_ts", SortOrder.Ascending) }, pageSize: 10, continuationToken: CosmosString.Create( @@ -200,9 +200,9 @@ FROM c WHERE {documentdb-formattableorderbyquery-filter} ORDER BY c._ts"), targetRanges: await documentContainer.GetFeedRangesAsync(cancellationToken: default), - orderByColumns: new List() + orderByColumns: new List() { - new OrderByCrossPartitionQueryPipelineStage.OrderByColumn("c._ts", SortOrder.Ascending) + new OrderByColumn("c._ts", SortOrder.Ascending) }, pageSize: 10, continuationToken: null); @@ -245,9 +245,9 @@ FROM c WHERE {documentdb-formattableorderbyquery-filter} ORDER BY c._ts"), targetRanges: await documentContainer.GetFeedRangesAsync(cancellationToken: default), - orderByColumns: new List() + orderByColumns: new List() { - new OrderByCrossPartitionQueryPipelineStage.OrderByColumn("c._ts", SortOrder.Ascending) + new OrderByColumn("c._ts", SortOrder.Ascending) }, pageSize: 10, continuationToken: queryState?.Value); From 243128cc58a7b4ab4aefe97af6b597d479447ccf Mon Sep 17 00:00:00 2001 From: Brandon Chong Date: Thu, 27 Aug 2020 11:41:15 -0700 Subject: [PATCH 62/85] need to handle splits --- ...OrderByCrossPartitionQueryPipelineStage.cs | 19 +++++++++++++++++-- 1 file changed, 17 insertions(+), 2 deletions(-) diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderBy/OrderByCrossPartitionQueryPipelineStage.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderBy/OrderByCrossPartitionQueryPipelineStage.cs index 3c8c997d9e..411b8f3735 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderBy/OrderByCrossPartitionQueryPipelineStage.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderBy/OrderByCrossPartitionQueryPipelineStage.cs @@ -109,7 +109,12 @@ public async ValueTask MoveNextAsync() if (uninitializedEnumerator.Current.Failed) { - //TODO HANDLE SPLIT + if (IsSplitException(uninitializedEnumerator.Current.Exception)) + { + // Handle split + throw new NotImplementedException(); + } + this.uninitializedEnumeratorsAndTokens.Enqueue((uninitializedEnumerator, token)); this.Current = TryCatch.FromException(uninitializedEnumerator.Current.Exception); } @@ -169,7 +174,12 @@ public async ValueTask MoveNextAsync() if (filterMonad.Failed) { - //TODO HANDLE SPLIT + if (IsSplitException(filterMonad.Exception)) + { + // Handle split + throw new NotImplementedException(); + } + this.Current = TryCatch.FromException(filterMonad.Exception); return true; } @@ -199,6 +209,11 @@ public async ValueTask MoveNextAsync() } else { + if (monadicQueryByPage.Failed) + { + // Handle split + throw new NotImplementedException(); + } this.uninitializedEnumeratorsAndTokens.Enqueue((uninitializedEnumerator, token)); } From f93cdf070b3e97d21c7ab1a00d28e680c1fd3907 Mon Sep 17 00:00:00 2001 From: Brandon Chong Date: Fri, 28 Aug 2020 12:24:47 -0700 Subject: [PATCH 63/85] got split working for order by --- ...OrderByCrossPartitionQueryPipelineStage.cs | 51 ++++++-- ...yQueryPartitionRangePageAsyncEnumerator.cs | 16 +-- .../Pagination/DocumentContainerTests.cs | 1 - .../Pagination/InMemoryContainer.cs | 40 ++++--- ...ByCrossPartitionQueryPipelineStageTests.cs | 112 +++++++++++++++++- .../Spatial/CommonSerializationTest.cs | 14 +-- 6 files changed, 195 insertions(+), 39 deletions(-) diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderBy/OrderByCrossPartitionQueryPipelineStage.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderBy/OrderByCrossPartitionQueryPipelineStage.cs index 411b8f3735..3da9cc2029 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderBy/OrderByCrossPartitionQueryPipelineStage.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderBy/OrderByCrossPartitionQueryPipelineStage.cs @@ -111,8 +111,23 @@ public async ValueTask MoveNextAsync() { if (IsSplitException(uninitializedEnumerator.Current.Exception)) { - // Handle split - throw new NotImplementedException(); + IEnumerable childRanges = await this.documentContainer.GetChildRangeAsync( + uninitializedEnumerator.Range, + cancellationToken: default); + foreach (PartitionKeyRange childRange in childRanges) + { + OrderByQueryPartitionRangePageAsyncEnumerator childPaginator = new OrderByQueryPartitionRangePageAsyncEnumerator( + this.documentContainer, + uninitializedEnumerator.SqlQuerySpec, + childRange, + uninitializedEnumerator.PageSize, + uninitializedEnumerator.Filter, + state: uninitializedEnumerator.StartOfPageState); + this.uninitializedEnumeratorsAndTokens.Enqueue((childPaginator, token)); + } + + // Recursively retry + return await this.MoveNextAsync(); } this.uninitializedEnumeratorsAndTokens.Enqueue((uninitializedEnumerator, token)); @@ -176,8 +191,23 @@ public async ValueTask MoveNextAsync() { if (IsSplitException(filterMonad.Exception)) { - // Handle split - throw new NotImplementedException(); + IEnumerable childRanges = await this.documentContainer.GetChildRangeAsync( + uninitializedEnumerator.Range, + cancellationToken: default); + foreach (PartitionKeyRange childRange in childRanges) + { + OrderByQueryPartitionRangePageAsyncEnumerator childPaginator = new OrderByQueryPartitionRangePageAsyncEnumerator( + this.documentContainer, + uninitializedEnumerator.SqlQuerySpec, + childRange, + uninitializedEnumerator.PageSize, + uninitializedEnumerator.Filter, + state: uninitializedEnumerator.StartOfPageState); + this.uninitializedEnumeratorsAndTokens.Enqueue((childPaginator, token)); + } + + // Recursively retry + return await this.MoveNextAsync(); } this.Current = TryCatch.FromException(filterMonad.Exception); @@ -288,6 +318,8 @@ public async ValueTask MoveNextAsync() continuationTokenString = continuationTokenList.ToString(); } + this.state = continuationTokenString != null ? new QueryState(CosmosString.Create(continuationTokenString)) : null; + // Return a page of results // No stats to report, since we already reported it when we moved to this page. this.Current = TryCatch.FromResult( @@ -298,7 +330,7 @@ public async ValueTask MoveNextAsync() responseLengthInBytes: 0, cosmosQueryExecutionInfo: default, disallowContinuationTokenMessage: default, - state: continuationTokenString != null ? new QueryState(CosmosString.Create(continuationTokenString)) : null)); + state: this.state)); return true; } @@ -792,9 +824,14 @@ private static (string leftFilter, string targetFilter, string rightFilter) GetF return TryCatch<(bool, TryCatch)>.FromResult((false, enumerator.Current)); } - private static bool IsSplitException(Exception exeception) + private static bool IsSplitException(Exception exception) { - return exeception is CosmosException cosmosException + while (exception.InnerException != null) + { + exception = exception.InnerException; + } + + return exception is CosmosException cosmosException && (cosmosException.StatusCode == HttpStatusCode.Gone) && (cosmosException.SubStatusCode == (int)Documents.SubStatusCodes.PartitionKeyRangeGone); } diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderBy/OrderByQueryPartitionRangePageAsyncEnumerator.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderBy/OrderByQueryPartitionRangePageAsyncEnumerator.cs index 11ec9b899e..a5746feee2 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderBy/OrderByQueryPartitionRangePageAsyncEnumerator.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderBy/OrderByQueryPartitionRangePageAsyncEnumerator.cs @@ -15,8 +15,6 @@ namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote.OrderBy internal sealed class OrderByQueryPartitionRangePageAsyncEnumerator : PartitionRangePageAsyncEnumerator { private readonly IQueryDataSource queryDataSource; - private readonly SqlQuerySpec sqlQuerySpec; - private readonly int pageSize; public OrderByQueryPartitionRangePageAsyncEnumerator( IQueryDataSource queryDataSource, @@ -28,12 +26,16 @@ public OrderByQueryPartitionRangePageAsyncEnumerator( : base(feedRange, state) { this.queryDataSource = queryDataSource ?? throw new ArgumentNullException(nameof(queryDataSource)); - this.sqlQuerySpec = sqlQuerySpec ?? throw new ArgumentNullException(nameof(sqlQuerySpec)); - this.pageSize = pageSize; + this.SqlQuerySpec = sqlQuerySpec ?? throw new ArgumentNullException(nameof(sqlQuerySpec)); + this.PageSize = pageSize; this.Filter = filter; this.StartOfPageState = state; } + public SqlQuerySpec SqlQuerySpec { get; } + + public int PageSize { get; } + public string Filter { get; } public QueryState StartOfPageState { get; private set; } @@ -45,10 +47,10 @@ protected override Task> GetNextPageAsync(Cancellatio this.StartOfPageState = this.State; return this.queryDataSource .MonadicQueryAsync( - sqlQuerySpec: this.sqlQuerySpec, + sqlQuerySpec: this.SqlQuerySpec, continuationToken: this.State == null ? null : ((CosmosString)this.State.Value).Value, - feedRange: new FeedRangeEpk(this.Range.ToRange()), - pageSize: this.pageSize, + feedRange: new FeedRangePartitionKeyRange(this.Range.Id), + pageSize: this.PageSize, cancellationToken) .ContinueWith>(antecedent => { diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/DocumentContainerTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/DocumentContainerTests.cs index a2b0d23994..9ff25b0024 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/DocumentContainerTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/DocumentContainerTests.cs @@ -124,7 +124,6 @@ public async Task TestCrudAsync() Record record = await documentContainer.CreateItemAsync(item, cancellationToken: default); Assert.IsNotNull(record); Assert.AreNotEqual(Guid.Empty, record.Identifier); - Assert.AreEqual(1, record.ResourceIdentifier); // Try to read it back Record readRecord = await documentContainer.ReadItemAsync( 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 9d38bc4663..57340d9ffa 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 @@ -220,7 +220,20 @@ static Task> CreateNotFoundException(CosmosElement partitionKey CosmosElement candidatePartitionKey = GetPartitionKeyFromPayload( candidate.Payload, this.partitionKeyDefinition); - bool partitionKeyMatches = candidatePartitionKey.Equals(partitionKey); + + bool partitionKeyMatches; + if (candidatePartitionKey is null && partitionKey is null) + { + partitionKeyMatches = true; + } + else if ((candidatePartitionKey != null) && (partitionKey != null)) + { + partitionKeyMatches = candidatePartitionKey.Equals(partitionKey); + } + else + { + partitionKeyMatches = false; + } if (identifierMatches && partitionKeyMatches) { @@ -347,37 +360,34 @@ public Task> MonadicQueryAsync( IEnumerable queryResults = SqlInterpreter.ExecuteQuery(documents, sqlQuery); - int index; + IEnumerable queryPageResults = queryResults; if (continuationToken != null) { - if (!int.TryParse(continuationToken, out index)) + queryPageResults = queryPageResults.Where(c => { - return Task.FromResult( - TryCatch.FromException( - new Exception("FAILED TO PARSE CONTINUATION TOKEN"))); - } - } - else - { - index = 0; + ResourceId continuationResourceId = ResourceId.Parse(continuationToken); + ResourceId documentResourceId = ResourceId.Parse(((CosmosString)((CosmosObject)c)["_rid"]).Value); + return documentResourceId.Document > continuationResourceId.Document; + }); } - List queryPageResults = queryResults.Skip(index).Take(pageSize).ToList(); + queryPageResults = queryPageResults.Take(pageSize); + List queryPageResultList = queryPageResults.ToList(); QueryState queryState; - if (queryPageResults.Count == 0) + if (queryPageResultList.Count == 0) { queryState = null; } else { - queryState = new QueryState(CosmosString.Create((index + pageSize).ToString())); + queryState = new QueryState(((CosmosObject)queryPageResultList.Last())["_rid"]); } return Task.FromResult( TryCatch.FromResult( new QueryPage( - queryPageResults, + queryPageResultList, requestCharge: 42, activityId: Guid.NewGuid().ToString(), responseLengthInBytes: 1337, diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/OrderByCrossPartitionQueryPipelineStageTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/OrderByCrossPartitionQueryPipelineStageTests.cs index cce832a4ac..72c14ddaf7 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/OrderByCrossPartitionQueryPipelineStageTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/OrderByCrossPartitionQueryPipelineStageTests.cs @@ -4,6 +4,7 @@ namespace Microsoft.Azure.Cosmos.Tests.Query.Pipeline { + using System; using System.Collections.Generic; using System.Linq; using System.Threading.Tasks; @@ -17,6 +18,7 @@ namespace Microsoft.Azure.Cosmos.Tests.Query.Pipeline using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote.Parallel; using Microsoft.Azure.Cosmos.Tests.Pagination; using Microsoft.Azure.Documents; + using Microsoft.Azure.Documents.Routing; using Microsoft.VisualStudio.TestTools.UnitTesting; using Moq; @@ -189,7 +191,7 @@ public void MonadicCreate_MultipleOrderByContinuationToken() [TestMethod] public async Task TestDrainFully_StartFromBeginingAsync() { - int numItems = 100; + int numItems = 1000; IDocumentContainer documentContainer = await CreateDocumentContainerAsync(numItems); TryCatch monadicCreate = OrderByCrossPartitionQueryPipelineStage.MonadicCreate( @@ -229,7 +231,7 @@ FROM c [TestMethod] public async Task TestDrainFully_WithStateResume() { - int numItems = 100; + int numItems = 1000; IDocumentContainer documentContainer = await CreateDocumentContainerAsync(numItems); List documents = new List(); @@ -273,6 +275,112 @@ FROM c } while (queryState != null); Assert.AreEqual(numItems, documents.Count); + Assert.IsTrue(documents.OrderBy(document => ((CosmosObject)document)["_ts"]).ToList().SequenceEqual(documents)); + } + + [TestMethod] + public async Task TestDrainFully_WithSplits() + { + int numItems = 1000; + IDocumentContainer documentContainer = await CreateDocumentContainerAsync(numItems); + + TryCatch monadicCreate = OrderByCrossPartitionQueryPipelineStage.MonadicCreate( + documentContainer: documentContainer, + sqlQuerySpec: new SqlQuerySpec(@" + SELECT c._rid AS _rid, [{""item"": c._ts}] AS orderByItems, c AS payload + FROM c + WHERE {documentdb-formattableorderbyquery-filter} + ORDER BY c._ts"), + targetRanges: await documentContainer.GetFeedRangesAsync(cancellationToken: default), + orderByColumns: new List() + { + new OrderByColumn("c._ts", SortOrder.Ascending) + }, + pageSize: 10, + continuationToken: null); + Assert.IsTrue(monadicCreate.Succeeded); + IQueryPipelineStage queryPipelineStage = monadicCreate.Result; + + Random random = new Random(); + List documents = new List(); + while (await queryPipelineStage.MoveNextAsync()) + { + TryCatch tryGetQueryPage = queryPipelineStage.Current; + if (tryGetQueryPage.Failed) + { + Assert.Fail(tryGetQueryPage.Exception.ToString()); + } + + QueryPage queryPage = tryGetQueryPage.Result; + documents.AddRange(queryPage.Documents); + + if (random.Next() % 4 == 0) + { + // Can not always split otherwise the split handling code will livelock trying to split proof every partition in a cycle. + List ranges = documentContainer.GetFeedRangesAsync(cancellationToken: default).Result; + PartitionKeyRange randomRange = ranges[random.Next(ranges.Count)]; + documentContainer.SplitAsync(int.Parse(randomRange.Id), cancellationToken: default).Wait(); + } + } + + Assert.AreEqual(numItems, documents.Count); + Assert.IsTrue(documents.OrderBy(document => ((CosmosObject)document)["_ts"]).ToList().SequenceEqual(documents)); + } + + [TestMethod] + public async Task TestDrainFully_WithSplit_WithStateResume() + { + int numItems = 1000; + IDocumentContainer documentContainer = await CreateDocumentContainerAsync(numItems); + + Random random = new Random(); + List documents = new List(); + QueryState queryState = null; + + do + { + TryCatch monadicCreate = OrderByCrossPartitionQueryPipelineStage.MonadicCreate( + documentContainer: documentContainer, + sqlQuerySpec: new SqlQuerySpec(@" + SELECT c._rid AS _rid, [{""item"": c._ts}] AS orderByItems, c AS payload + FROM c + WHERE {documentdb-formattableorderbyquery-filter} + ORDER BY c._ts"), + targetRanges: await documentContainer.GetFeedRangesAsync(cancellationToken: default), + orderByColumns: new List() + { + new OrderByColumn("c._ts", SortOrder.Ascending) + }, + pageSize: 10, + continuationToken: queryState?.Value); + Assert.IsTrue(monadicCreate.Succeeded); + IQueryPipelineStage queryPipelineStage = monadicCreate.Result; + + QueryPage queryPage; + do + { + // We need to drain out all the initial empty pages, + // since they are non resumable state. + Assert.IsTrue(await queryPipelineStage.MoveNextAsync()); + TryCatch tryGetQueryPage = queryPipelineStage.Current; + if (tryGetQueryPage.Failed) + { + Assert.Fail(tryGetQueryPage.Exception.ToString()); + } + + queryPage = tryGetQueryPage.Result; + documents.AddRange(queryPage.Documents); + queryState = queryPage.State; + } while ((queryPage.Documents.Count == 0) && (queryState != null)); + + // Split + List ranges = documentContainer.GetFeedRangesAsync(cancellationToken: default).Result; + PartitionKeyRange randomRange = ranges[random.Next(ranges.Count)]; + documentContainer.SplitAsync(int.Parse(randomRange.Id), cancellationToken: default).Wait(); + } while (queryState != null); + + Assert.AreEqual(numItems, documents.Count); + Assert.IsTrue(documents.OrderBy(document => ((CosmosObject)document)["_ts"]).ToList().SequenceEqual(documents)); } private static async Task CreateDocumentContainerAsync( diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Spatial/CommonSerializationTest.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Spatial/CommonSerializationTest.cs index 473022ea34..b71a0a07fe 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Spatial/CommonSerializationTest.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Spatial/CommonSerializationTest.cs @@ -26,7 +26,7 @@ public class CommonSerializationTest public void TestInvalidJson() { string json = @"{""type"":""Poi}"; - var point = JsonConvert.DeserializeObject(json); + Point point = JsonConvert.DeserializeObject(json); } /// @@ -37,7 +37,7 @@ public void TestInvalidJson() public void TestNullCoordinates() { string json = @"{""type"":""Point"",""coordinates"":null}"; - var point = JsonConvert.DeserializeObject(json); + Point point = JsonConvert.DeserializeObject(json); } /// @@ -48,7 +48,7 @@ public void TestNullCoordinates() public void TestNullType() { string json = @"{""type"":null, ""coordinates"":[20, 30]}"; - var point = JsonConvert.DeserializeObject(json); + Point point = JsonConvert.DeserializeObject(json); } /// @@ -59,7 +59,7 @@ public void TestNullType() public void TestNullTypeGeometry() { string json = @"{""type"":null, ""coordinates"":[20, 30]}"; - var point = JsonConvert.DeserializeObject(json); + Geometry point = JsonConvert.DeserializeObject(json); } /// @@ -70,7 +70,7 @@ public void TestNullTypeGeometry() public void TestBoundingBoxWithNonEvenNumberOfCoordinates() { string json = @"{""type"":""Point"", ""coordinates"":[20, 30], ""bbox"":[0, 0, 0, 5, 5]}"; - var point = JsonConvert.DeserializeObject(json); + Point point = JsonConvert.DeserializeObject(json); } /// @@ -81,7 +81,7 @@ public void TestBoundingBoxWithNonEvenNumberOfCoordinates() public void TestBoundingBoxWithNotEnoughCoordinates() { string json = @"{""type"":""Point"", ""coordinates"":[20, 30], ""bbox"":[0, 0]}"; - var point = JsonConvert.DeserializeObject(json); + Point point = JsonConvert.DeserializeObject(json); } /// @@ -91,7 +91,7 @@ public void TestBoundingBoxWithNotEnoughCoordinates() public void TestNullBoundingBox() { string json = @"{""type"":""Point"", ""coordinates"":[20, 30], ""bbox"":null}"; - var point = JsonConvert.DeserializeObject(json); + Point point = JsonConvert.DeserializeObject(json); Assert.IsNull(point.BoundingBox); } From 3097cce57217f981d961ab58916c0bf6fdb3d97e Mon Sep 17 00:00:00 2001 From: Brandon Chong Date: Fri, 28 Aug 2020 13:19:38 -0700 Subject: [PATCH 64/85] refactored to be a state machine --- ...OrderByCrossPartitionQueryPipelineStage.cs | 339 ++++++++++-------- 1 file changed, 185 insertions(+), 154 deletions(-) diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderBy/OrderByCrossPartitionQueryPipelineStage.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderBy/OrderByCrossPartitionQueryPipelineStage.cs index 3da9cc2029..dbafe64f74 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderBy/OrderByCrossPartitionQueryPipelineStage.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderBy/OrderByCrossPartitionQueryPipelineStage.cs @@ -84,191 +84,205 @@ private OrderByCrossPartitionQueryPipelineStage( public ValueTask DisposeAsync() => default; - public async ValueTask MoveNextAsync() + private async ValueTask MoveNextAsync_Initialize_FromBeginningAsync( + OrderByQueryPartitionRangePageAsyncEnumerator uninitializedEnumerator, + OrderByContinuationToken token) { - if (this.uninitializedEnumeratorsAndTokens.Count != 0) + if (token != null) + { + throw new ArgumentException(nameof(token)); + } + + // We need to prime the page + if (!await uninitializedEnumerator.MoveNextAsync()) + { + // No more documents, so just return an empty page + this.Current = TryCatch.FromResult( + new QueryPage( + documents: EmptyPage, + requestCharge: 0, + activityId: string.Empty, + responseLengthInBytes: 0, + cosmosQueryExecutionInfo: default, + disallowContinuationTokenMessage: default, + state: this.state)); + return true; + } + + if (uninitializedEnumerator.Current.Failed) + { + if (IsSplitException(uninitializedEnumerator.Current.Exception)) + { + return await this.MoveNextAsync_InitializeAsync_HandleSplitAsync(uninitializedEnumerator, token); + } + + this.uninitializedEnumeratorsAndTokens.Enqueue((uninitializedEnumerator, token)); + this.Current = TryCatch.FromException(uninitializedEnumerator.Current.Exception); + } + else { - (OrderByQueryPartitionRangePageAsyncEnumerator uninitializedEnumerator, OrderByContinuationToken token) = this.uninitializedEnumeratorsAndTokens.Dequeue(); - if (token is null) + if (!uninitializedEnumerator.Current.Result.Enumerator.MoveNext()) { - // We need to prime the page - if (!await uninitializedEnumerator.MoveNextAsync()) + // Page was empty + if (uninitializedEnumerator.State != null) + { + this.uninitializedEnumeratorsAndTokens.Enqueue((uninitializedEnumerator, token)); + } + + if ((this.uninitializedEnumeratorsAndTokens.Count == 0) && (this.enumerators.Count == 0)) { - // No more documents, so just return an empty page + // Query did not match any results. We need to emit a fake empty page with null continuation this.Current = TryCatch.FromResult( new QueryPage( documents: EmptyPage, requestCharge: 0, - activityId: string.Empty, + activityId: Guid.NewGuid().ToString(), responseLengthInBytes: 0, cosmosQueryExecutionInfo: default, disallowContinuationTokenMessage: default, - state: this.state)); + state: null)); return true; } + } + else + { + this.enumerators.Enqueue(uninitializedEnumerator); + } - if (uninitializedEnumerator.Current.Failed) - { - if (IsSplitException(uninitializedEnumerator.Current.Exception)) - { - IEnumerable childRanges = await this.documentContainer.GetChildRangeAsync( - uninitializedEnumerator.Range, - cancellationToken: default); - foreach (PartitionKeyRange childRange in childRanges) - { - OrderByQueryPartitionRangePageAsyncEnumerator childPaginator = new OrderByQueryPartitionRangePageAsyncEnumerator( - this.documentContainer, - uninitializedEnumerator.SqlQuerySpec, - childRange, - uninitializedEnumerator.PageSize, - uninitializedEnumerator.Filter, - state: uninitializedEnumerator.StartOfPageState); - this.uninitializedEnumeratorsAndTokens.Enqueue((childPaginator, token)); - } + QueryPage page = uninitializedEnumerator.Current.Result.Page; + // Just return an empty page with the stats + this.Current = TryCatch.FromResult( + new QueryPage( + documents: EmptyPage, + requestCharge: page.RequestCharge, + activityId: page.ActivityId, + responseLengthInBytes: page.ResponseLengthInBytes, + cosmosQueryExecutionInfo: page.CosmosQueryExecutionInfo, + disallowContinuationTokenMessage: page.DisallowContinuationTokenMessage, + state: this.state)); + } - // Recursively retry - return await this.MoveNextAsync(); - } + return true; + } - this.uninitializedEnumeratorsAndTokens.Enqueue((uninitializedEnumerator, token)); - this.Current = TryCatch.FromException(uninitializedEnumerator.Current.Exception); - } - else - { - if (!uninitializedEnumerator.Current.Result.Enumerator.MoveNext()) - { - // Page was empty - if (uninitializedEnumerator.State != null) - { - this.uninitializedEnumeratorsAndTokens.Enqueue((uninitializedEnumerator, token)); - } + private async ValueTask MoveNextAsync_Iniialize_FilterAsync( + OrderByQueryPartitionRangePageAsyncEnumerator uninitializedEnumerator, + OrderByContinuationToken token) + { + if (token == null) + { + throw new ArgumentNullException(nameof(token)); + } - if ((this.uninitializedEnumeratorsAndTokens.Count == 0) && (this.enumerators.Count == 0)) - { - // Query did not match any results. We need to emit a fake empty page with null continuation - this.Current = TryCatch.FromResult( - new QueryPage( - documents: EmptyPage, - requestCharge: 0, - activityId: Guid.NewGuid().ToString(), - responseLengthInBytes: 0, - cosmosQueryExecutionInfo: default, - disallowContinuationTokenMessage: default, - state: null)); - return true; - } - } - else - { - this.enumerators.Enqueue(uninitializedEnumerator); - } + TryCatch<(bool, TryCatch)> filterMonad = await FilterNextAsync( + uninitializedEnumerator, + this.sortOrders, + token, + cancellationToken: default); - QueryPage page = uninitializedEnumerator.Current.Result.Page; - // Just return an empty page with the stats - this.Current = TryCatch.FromResult( - new QueryPage( - documents: EmptyPage, - requestCharge: page.RequestCharge, - activityId: page.ActivityId, - responseLengthInBytes: page.ResponseLengthInBytes, - cosmosQueryExecutionInfo: page.CosmosQueryExecutionInfo, - disallowContinuationTokenMessage: page.DisallowContinuationTokenMessage, - state: this.state)); - } + if (filterMonad.Failed) + { + if (IsSplitException(filterMonad.Exception)) + { + return await this.MoveNextAsync_InitializeAsync_HandleSplitAsync(uninitializedEnumerator, token); + } + this.Current = TryCatch.FromException(filterMonad.Exception); + return true; + } + + (bool doneFiltering, TryCatch monadicQueryByPage) = filterMonad.Result; + if (doneFiltering) + { + if (uninitializedEnumerator.Current.Result.Enumerator.Current != null) + { + this.enumerators.Enqueue(uninitializedEnumerator); + } + else if ((this.uninitializedEnumeratorsAndTokens.Count == 0) && (this.enumerators.Count == 0)) + { + // Query did not match any results. + // We need to emit a fake empty page with null continuation + this.Current = TryCatch.FromResult( + new QueryPage( + documents: EmptyPage, + requestCharge: 0, + activityId: Guid.NewGuid().ToString(), + responseLengthInBytes: 0, + cosmosQueryExecutionInfo: default, + disallowContinuationTokenMessage: default, + state: null)); return true; } - else + } + else + { + if (monadicQueryByPage.Failed) { - // We need to actually filter the enumerator - TryCatch<(bool, TryCatch)> filterMonad = await FilterNextAsync( - uninitializedEnumerator, - this.sortOrders, - token, - cancellationToken: default); - - if (filterMonad.Failed) + if (IsSplitException(filterMonad.Exception)) { - if (IsSplitException(filterMonad.Exception)) - { - IEnumerable childRanges = await this.documentContainer.GetChildRangeAsync( - uninitializedEnumerator.Range, - cancellationToken: default); - foreach (PartitionKeyRange childRange in childRanges) - { - OrderByQueryPartitionRangePageAsyncEnumerator childPaginator = new OrderByQueryPartitionRangePageAsyncEnumerator( - this.documentContainer, - uninitializedEnumerator.SqlQuerySpec, - childRange, - uninitializedEnumerator.PageSize, - uninitializedEnumerator.Filter, - state: uninitializedEnumerator.StartOfPageState); - this.uninitializedEnumeratorsAndTokens.Enqueue((childPaginator, token)); - } - - // Recursively retry - return await this.MoveNextAsync(); - } - - this.Current = TryCatch.FromException(filterMonad.Exception); - return true; + return await this.MoveNextAsync_InitializeAsync_HandleSplitAsync(uninitializedEnumerator, token); } + } - (bool doneFiltering, TryCatch monadicQueryByPage) = filterMonad.Result; - if (doneFiltering) - { - if (uninitializedEnumerator.Current.Result.Enumerator.Current != null) - { - this.enumerators.Enqueue(uninitializedEnumerator); - } - else if ((this.uninitializedEnumeratorsAndTokens.Count == 0) && (this.enumerators.Count == 0)) - { - // Query did not match any results. - // We need to emit a fake empty page with null continuation - this.Current = TryCatch.FromResult( - new QueryPage( - documents: EmptyPage, - requestCharge: 0, - activityId: Guid.NewGuid().ToString(), - responseLengthInBytes: 0, - cosmosQueryExecutionInfo: default, - disallowContinuationTokenMessage: default, - state: null)); - return true; - } - } - else - { - if (monadicQueryByPage.Failed) - { - // Handle split - throw new NotImplementedException(); - } - this.uninitializedEnumeratorsAndTokens.Enqueue((uninitializedEnumerator, token)); - } + this.uninitializedEnumeratorsAndTokens.Enqueue((uninitializedEnumerator, token)); + } - QueryPage page = uninitializedEnumerator.Current.Result.Page; - // Just return an empty page with the stats - this.Current = TryCatch.FromResult( - new QueryPage( - documents: EmptyPage, - requestCharge: page.RequestCharge, - activityId: page.ActivityId, - responseLengthInBytes: page.ResponseLengthInBytes, - cosmosQueryExecutionInfo: page.CosmosQueryExecutionInfo, - disallowContinuationTokenMessage: page.DisallowContinuationTokenMessage, - state: this.state)); + QueryPage page = uninitializedEnumerator.Current.Result.Page; + // Just return an empty page with the stats + this.Current = TryCatch.FromResult( + new QueryPage( + documents: EmptyPage, + requestCharge: page.RequestCharge, + activityId: page.ActivityId, + responseLengthInBytes: page.ResponseLengthInBytes, + cosmosQueryExecutionInfo: page.CosmosQueryExecutionInfo, + disallowContinuationTokenMessage: page.DisallowContinuationTokenMessage, + state: this.state)); - return true; - } + return true; + } + + private async ValueTask MoveNextAsync_InitializeAsync_HandleSplitAsync( + OrderByQueryPartitionRangePageAsyncEnumerator uninitializedEnumerator, + OrderByContinuationToken token) + { + IEnumerable childRanges = await this.documentContainer.GetChildRangeAsync( + uninitializedEnumerator.Range, + cancellationToken: default); + foreach (PartitionKeyRange childRange in childRanges) + { + OrderByQueryPartitionRangePageAsyncEnumerator childPaginator = new OrderByQueryPartitionRangePageAsyncEnumerator( + this.documentContainer, + uninitializedEnumerator.SqlQuerySpec, + childRange, + uninitializedEnumerator.PageSize, + uninitializedEnumerator.Filter, + state: uninitializedEnumerator.StartOfPageState); + this.uninitializedEnumeratorsAndTokens.Enqueue((childPaginator, token)); } - if (this.enumerators.Count == 0) + // Recursively retry + return await this.MoveNextAsync(); + } + + private ValueTask MoveNextAsync_InitializeAsync() + { + (OrderByQueryPartitionRangePageAsyncEnumerator uninitializedEnumerator, OrderByContinuationToken token) = this.uninitializedEnumeratorsAndTokens.Dequeue(); + ValueTask task; + if (token is null) { - // Finished draining. - return false; + task = this.MoveNextAsync_Initialize_FromBeginningAsync(uninitializedEnumerator, token); + } + else + { + task = this.MoveNextAsync_Iniialize_FilterAsync(uninitializedEnumerator, token); } + return task; + } + + private ValueTask MoveNextAsync_DrainPageAsync() + { OrderByQueryPartitionRangePageAsyncEnumerator currentEnumerator = default; OrderByQueryResult orderByQueryResult = default; @@ -296,6 +310,7 @@ public async ValueTask MoveNextAsync() this.enumerators.Enqueue(currentEnumerator); } + // Create the continuation token. string continuationTokenString; if ((this.enumerators.Count == 0) && (this.uninitializedEnumeratorsAndTokens.Count == 0)) { @@ -331,7 +346,23 @@ public async ValueTask MoveNextAsync() cosmosQueryExecutionInfo: default, disallowContinuationTokenMessage: default, state: this.state)); - return true; + return new ValueTask(true); + } + + public ValueTask MoveNextAsync() + { + if (this.uninitializedEnumeratorsAndTokens.Count != 0) + { + return this.MoveNextAsync_InitializeAsync(); + } + + if (this.enumerators.Count == 0) + { + // Finished draining. + return new ValueTask(false); + } + + return this.MoveNextAsync_DrainPageAsync(); } public static TryCatch MonadicCreate( From ee74df3db0ef395181fb880a9a2958ba81f10b8b Mon Sep 17 00:00:00 2001 From: Brandon Chong Date: Fri, 28 Aug 2020 17:12:29 -0700 Subject: [PATCH 65/85] fixed bug where partiton key ranges where matched to a continuation due to string.Empty start and ends --- ...OrderByCrossPartitionQueryPipelineStage.cs | 2 +- .../Core/Pipeline/Remote/PartitionMapper.cs | 8 +- .../PartitionKeyHashRangeSplitterAndMerger.cs | 19 +-- .../Query/ContinuationResumeLogicTests.cs | 111 ++++++++++++++++++ 4 files changed, 124 insertions(+), 16 deletions(-) diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderBy/OrderByCrossPartitionQueryPipelineStage.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderBy/OrderByCrossPartitionQueryPipelineStage.cs index dbafe64f74..342177a22f 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderBy/OrderByCrossPartitionQueryPipelineStage.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderBy/OrderByCrossPartitionQueryPipelineStage.cs @@ -442,7 +442,6 @@ public static TryCatch MonadicCreate( } PartitionMapping partitionMapping = monadicGetOrderByContinuationTokenMapping.Result; - IReadOnlyList orderByItems = partitionMapping .TargetPartition .Values @@ -450,6 +449,7 @@ public static TryCatch MonadicCreate( .OrderByItems .Select(x => x.Item) .ToList(); + if (orderByItems.Count != orderByColumns.Count) { return TryCatch.FromException( diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/PartitionMapper.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/PartitionMapper.cs index 7523cc3209..fd2aff8608 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/PartitionMapper.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/PartitionMapper.cs @@ -9,6 +9,7 @@ namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote using System.Linq; using Microsoft.Azure.Cosmos.Query.Core.Exceptions; using Microsoft.Azure.Cosmos.Query.Core.Monads; + using Microsoft.Azure.Cosmos.Routing; using Microsoft.Azure.Documents; internal static class PartitionMapper @@ -117,9 +118,12 @@ public static IReadOnlyDictionary MatchRang PartitionKeyRange partitionKeyRange = partitionKeyRangeSpan[i]; foreach (PartitionedToken partitionedToken in partitionedContinuationTokens) { + bool rightOfStart = (partitionedToken.Range.Min == string.Empty) + || ((partitionKeyRange.MinInclusive != string.Empty) && (partitionKeyRange.MinInclusive.CompareTo(partitionedToken.Range.Min) >= 0)); + bool leftOfEnd = (partitionedToken.Range.Max == string.Empty) + || ((partitionKeyRange.MaxExclusive != string.Empty) && (partitionKeyRange.MaxExclusive.CompareTo(partitionedToken.Range.Max) <= 0)); // See if continuation token includes the range - if ((partitionKeyRange.MinInclusive.CompareTo(partitionedToken.Range.Min) >= 0) - && (partitionKeyRange.MaxExclusive.CompareTo(partitionedToken.Range.Max) <= 0)) + if (rightOfStart && leftOfEnd) { partitionKeyRangeToToken[partitionKeyRange] = partitionedToken; break; diff --git a/Microsoft.Azure.Cosmos/src/Routing/PartitionKeyHashRangeSplitterAndMerger.cs b/Microsoft.Azure.Cosmos/src/Routing/PartitionKeyHashRangeSplitterAndMerger.cs index cccc9363b6..54b72ae3a2 100644 --- a/Microsoft.Azure.Cosmos/src/Routing/PartitionKeyHashRangeSplitterAndMerger.cs +++ b/Microsoft.Azure.Cosmos/src/Routing/PartitionKeyHashRangeSplitterAndMerger.cs @@ -22,20 +22,13 @@ public static PartitionKeyHashRanges SplitRange(PartitionKeyHashRange partitionK rangeCount, out PartitionKeyHashRanges splitRanges); - switch (splitOutcome) + return splitOutcome switch { - case SplitOutcome.Success: - return splitRanges; - - case SplitOutcome.NumRangesNeedsToBeGreaterThanZero: - throw new ArgumentOutOfRangeException($"{nameof(rangeCount)} must be a positive integer"); - - case SplitOutcome.RangeNotWideEnough: - throw new ArgumentOutOfRangeException($"{nameof(partitionKeyHashRange)} is not wide enough to split into {rangeCount} ranges."); - - default: - throw new ArgumentOutOfRangeException($"Unknown {nameof(SplitOutcome)}: {splitOutcome}."); - } + SplitOutcome.Success => splitRanges, + SplitOutcome.NumRangesNeedsToBeGreaterThanZero => throw new ArgumentOutOfRangeException($"{nameof(rangeCount)} must be a positive integer"), + SplitOutcome.RangeNotWideEnough => throw new ArgumentOutOfRangeException($"{nameof(partitionKeyHashRange)} is not wide enough to split into {rangeCount} ranges."), + _ => throw new ArgumentOutOfRangeException($"Unknown {nameof(SplitOutcome)}: {splitOutcome}."), + }; } public static SplitOutcome TrySplitRange(PartitionKeyHashRange partitionKeyHashRange, int rangeCount, out PartitionKeyHashRanges splitRanges) diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/ContinuationResumeLogicTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/ContinuationResumeLogicTests.cs index 3102a09490..cfcdec5b0c 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/ContinuationResumeLogicTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/ContinuationResumeLogicTests.cs @@ -121,6 +121,117 @@ public void TestMatchRangesTocontinuationTokens_ArgumentNullException() partitionedTokens: null); } + [TestMethod] + public void TestTryGetInitializationInfo_ResumeEmptyStart() + { + PartitionKeyRange pkRange1 = new PartitionKeyRange() + { + MinInclusive = string.Empty, + MaxExclusive = "A", + Id = "1" + }; + + PartitionKeyRange pkRange2 = new PartitionKeyRange() + { + MinInclusive = "A", + MaxExclusive = "B", + Id = "2" + }; + + PartitionKeyRange pkRange3 = new PartitionKeyRange() + { + MinInclusive = "B", + MaxExclusive = string.Empty, + Id = "3" + }; + + ParallelContinuationToken token = new ParallelContinuationToken( + token: "asdf", + range: new Documents.Routing.Range( + min: string.Empty, + max: "B", + isMinInclusive: true, + isMaxInclusive: false)); + + IReadOnlyDictionary expectedMappingLeftPartitions = new Dictionary() + { + }; + + IReadOnlyDictionary expectedMappingTargetPartition = new Dictionary() + { + { pkRange1, token }, + }; + + IReadOnlyDictionary expectedMappingRightPartitions = new Dictionary() + { + { pkRange2, token }, + { pkRange3, null}, + }; + + RunTryGetInitializationInfo( + expectedMappingLeftPartitions, + expectedMappingTargetPartition, + expectedMappingRightPartitions, + new PartitionKeyRange[] { pkRange1, pkRange2, pkRange3 }, + new IPartitionedToken[] { token }); + } + + [TestMethod] + public void TestTryGetInitializationInfo_ResumeEmptyEnd() + { + PartitionKeyRange pkRange1 = new PartitionKeyRange() + { + MinInclusive = string.Empty, + MaxExclusive = "A", + Id = "1" + }; + + PartitionKeyRange pkRange2 = new PartitionKeyRange() + { + MinInclusive = "A", + MaxExclusive = "B", + Id = "2" + }; + + PartitionKeyRange pkRange3 = new PartitionKeyRange() + { + MinInclusive = "B", + MaxExclusive = string.Empty, + Id = "3" + }; + + ParallelContinuationToken token = new ParallelContinuationToken( + token: "asdf", + range: new Documents.Routing.Range( + min: "A", + max: string.Empty, + isMinInclusive: true, + isMaxInclusive: false)); + + IReadOnlyDictionary expectedMappingLeftPartitions = new Dictionary() + { + { pkRange1, null }, + }; + + IReadOnlyDictionary expectedMappingTargetPartition = new Dictionary() + { + { pkRange2, token }, + }; + + IReadOnlyDictionary expectedMappingRightPartitions = new Dictionary() + { + + { pkRange3, token }, + }; + + RunTryGetInitializationInfo( + expectedMappingLeftPartitions, + expectedMappingTargetPartition, + expectedMappingRightPartitions, + new PartitionKeyRange[] { pkRange1, pkRange2, pkRange3 }, + new IPartitionedToken[] { token }); + } + [TestMethod] public void TestTryGetInitializationInfo_ResumeLeftMostPartition() { From ba264d0125df56c3a8786624c4d1a589585eeb51 Mon Sep 17 00:00:00 2001 From: Brandon Chong Date: Tue, 1 Sep 2020 11:44:20 -0700 Subject: [PATCH 66/85] added parallel buffering --- ...fferedPartitionRangePageAsyncEnumerator.cs | 70 ++++++ .../CrossPartitionRangePageAsyncEnumerable.cs | 4 + .../CrossPartitionRangePageAsyncEnumerator.cs | 36 ++- .../src/Pagination/IBufferable.cs | 14 ++ .../src/Pagination/ParallelBuffering.cs | 56 +++++ .../PartitionRangePageAsyncEnumerator.cs | 3 +- .../src/Pagination/StartedState.cs | 15 ++ .../src/Pagination/State.cs | 4 - .../CosmosQueryExecutionContextFactory.cs | 2 + .../Query/Core/Pipeline/PipelineFactory.cs | 3 + ...OrderByCrossPartitionQueryPipelineStage.cs | 30 ++- ...yQueryPartitionRangePageAsyncEnumerator.cs | 96 +++++--- ...arallelCrossPartitionQueryPipelineStage.cs | 2 + .../BufferedPartitionRangeEnumeratorTests.cs | 222 ++++++++++++++++++ ...sPartitionPartitionRangeEnumeratorTests.cs | 2 + .../Query/Pipeline/FactoryTests.cs | 1 + .../Query/Pipeline/FullPipelineTests.cs | 1 + ...ByCrossPartitionQueryPipelineStageTests.cs | 10 + ...elCrossPartitionQueryPipelineStageTests.cs | 8 + 19 files changed, 521 insertions(+), 58 deletions(-) create mode 100644 Microsoft.Azure.Cosmos/src/Pagination/BufferedPartitionRangePageAsyncEnumerator.cs create mode 100644 Microsoft.Azure.Cosmos/src/Pagination/IBufferable.cs create mode 100644 Microsoft.Azure.Cosmos/src/Pagination/ParallelBuffering.cs create mode 100644 Microsoft.Azure.Cosmos/src/Pagination/StartedState.cs create mode 100644 Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/BufferedPartitionRangeEnumeratorTests.cs diff --git a/Microsoft.Azure.Cosmos/src/Pagination/BufferedPartitionRangePageAsyncEnumerator.cs b/Microsoft.Azure.Cosmos/src/Pagination/BufferedPartitionRangePageAsyncEnumerator.cs new file mode 100644 index 0000000000..dc11a14d58 --- /dev/null +++ b/Microsoft.Azure.Cosmos/src/Pagination/BufferedPartitionRangePageAsyncEnumerator.cs @@ -0,0 +1,70 @@ +// ------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +// ------------------------------------------------------------ + +namespace Microsoft.Azure.Cosmos.Pagination +{ + using System; + using System.Threading; + using System.Threading.Tasks; + using Microsoft.Azure.Cosmos.Query.Core.Monads; + + internal sealed class BufferedPartitionRangePageAsyncEnumerator : PartitionRangePageAsyncEnumerator, IBufferable + where TPage : Page + where TState : State + { + private readonly PartitionRangePageAsyncEnumerator enumerator; + private readonly SemaphoreSlim mutex; + private TryCatch? bufferedPage; + + public BufferedPartitionRangePageAsyncEnumerator(PartitionRangePageAsyncEnumerator enumerator) + : base(enumerator.Range, enumerator.State) + { + this.enumerator = enumerator ?? throw new ArgumentNullException(nameof(enumerator)); + this.mutex = new SemaphoreSlim(initialCount: 1, maxCount: 1); + } + + public override ValueTask DisposeAsync() => this.enumerator.DisposeAsync(); + + protected override async Task> GetNextPageAsync(CancellationToken cancellationToken) + { + try + { + await this.mutex.WaitAsync(cancellationToken); + if (this.bufferedPage.HasValue) + { + // Serve from the buffered page first. + TryCatch returnValue = this.bufferedPage.Value; + this.bufferedPage = null; + return returnValue; + } + + await this.enumerator.MoveNextAsync(); + return this.enumerator.Current; + } + finally + { + this.mutex.Release(); + } + } + + public async ValueTask BufferAsync() + { + try + { + await this.mutex.WaitAsync(); + if (this.bufferedPage.HasValue) + { + return; + } + + await this.enumerator.MoveNextAsync(); + this.bufferedPage = this.enumerator.Current; + } + finally + { + this.mutex.Release(); + } + } + } +} diff --git a/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionRangePageAsyncEnumerable.cs b/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionRangePageAsyncEnumerable.cs index c4cb0cb669..b5d0894512 100644 --- a/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionRangePageAsyncEnumerable.cs +++ b/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionRangePageAsyncEnumerable.cs @@ -17,17 +17,20 @@ internal sealed class CrossPartitionRangePageAsyncEnumerable : IA private readonly CreatePartitionRangePageAsyncEnumerator createPartitionRangeEnumerator; private readonly IComparer> comparer; private readonly IFeedRangeProvider feedRangeProvider; + private readonly int maxConcurrency; public CrossPartitionRangePageAsyncEnumerable( IFeedRangeProvider feedRangeProvider, CreatePartitionRangePageAsyncEnumerator createPartitionRangeEnumerator, IComparer> comparer, + int maxConcurrency, CrossPartitionState state = default) { this.feedRangeProvider = feedRangeProvider ?? throw new ArgumentNullException(nameof(comparer)); this.createPartitionRangeEnumerator = createPartitionRangeEnumerator ?? throw new ArgumentNullException(nameof(createPartitionRangeEnumerator)); this.comparer = comparer ?? throw new ArgumentNullException(nameof(comparer)); this.state = state; + this.maxConcurrency = maxConcurrency < 0 ? throw new ArgumentOutOfRangeException(nameof(maxConcurrency)) : maxConcurrency; } public IAsyncEnumerator>> GetAsyncEnumerator(CancellationToken cancellationToken = default) @@ -38,6 +41,7 @@ public IAsyncEnumerator>> GetAsyncEnu this.feedRangeProvider, this.createPartitionRangeEnumerator, this.comparer, + this.maxConcurrency, this.state); } } diff --git a/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionRangePageAsyncEnumerator.cs b/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionRangePageAsyncEnumerator.cs index 459e6a1f14..a07c570e71 100644 --- a/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionRangePageAsyncEnumerator.cs +++ b/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionRangePageAsyncEnumerator.cs @@ -6,6 +6,7 @@ namespace Microsoft.Azure.Cosmos.Pagination { using System; using System.Collections.Generic; + using System.Linq; using System.Net; using System.Threading; using System.Threading.Tasks; @@ -29,6 +30,7 @@ public CrossPartitionRangePageAsyncEnumerator( IFeedRangeProvider feedRangeProvider, CreatePartitionRangePageAsyncEnumerator createPartitionRangeEnumerator, IComparer> comparer, + int? maxConcurrency, CrossPartitionState state = default) { this.feedRangeProvider = feedRangeProvider ?? throw new ArgumentNullException(nameof(feedRangeProvider)); @@ -60,13 +62,23 @@ public CrossPartitionRangePageAsyncEnumerator( rangeAndStates = rangesAndStatesBuilder; } - PriorityQueue> enumerators = new PriorityQueue>(comparer); - foreach ((PartitionKeyRange range, TState rangeState) in rangeAndStates) + List> bufferedEnumerators = rangeAndStates + .Select(rangeAndState => + { + PartitionRangePageAsyncEnumerator enumerator = createPartitionRangeEnumerator(rangeAndState.Item1, rangeAndState.Item2); + BufferedPartitionRangePageAsyncEnumerator bufferedEnumerator = new BufferedPartitionRangePageAsyncEnumerator(enumerator); + return bufferedEnumerator; + }) + .ToList(); + + if (maxConcurrency.HasValue) { - PartitionRangePageAsyncEnumerator enumerator = createPartitionRangeEnumerator(range, rangeState); - enumerators.Enqueue(enumerator); + await ParallelBuffering.BufferInParallelAsync(bufferedEnumerators, maxConcurrency.Value); } + PriorityQueue> enumerators = new PriorityQueue>( + bufferedEnumerators, + comparer); return enumerators; }); } @@ -120,18 +132,18 @@ public async ValueTask MoveNextAsync() { throw new NotImplementedException(); } - } - if (currentPaginator.State != null) - { + // Just enqueue the paginator and the user can decide if they want to retry. enumerators.Enqueue(currentPaginator); + + this.Current = TryCatch>.FromException(currentPaginator.Current.Exception); + return true; } - TryCatch backendPage = currentPaginator.Current; - if (backendPage.Failed) + if (currentPaginator.State != default) { - this.Current = TryCatch>.FromException(backendPage.Exception); - return true; + // Don't enqueue the paginator otherwise it's an infinite loop. + enumerators.Enqueue(currentPaginator); } CrossPartitionState crossPartitionState; @@ -151,7 +163,7 @@ public async ValueTask MoveNextAsync() } this.Current = TryCatch>.FromResult( - new CrossPartitionPage(backendPage.Result, crossPartitionState)); + new CrossPartitionPage(currentPaginator.Current.Result, crossPartitionState)); return true; } diff --git a/Microsoft.Azure.Cosmos/src/Pagination/IBufferable.cs b/Microsoft.Azure.Cosmos/src/Pagination/IBufferable.cs new file mode 100644 index 0000000000..23625a8fc9 --- /dev/null +++ b/Microsoft.Azure.Cosmos/src/Pagination/IBufferable.cs @@ -0,0 +1,14 @@ +// ------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +// ------------------------------------------------------------ + +namespace Microsoft.Azure.Cosmos.Pagination +{ + using System.Collections.Generic; + using System.Threading.Tasks; + + internal interface IBufferable + { + ValueTask BufferAsync(); + } +} diff --git a/Microsoft.Azure.Cosmos/src/Pagination/ParallelBuffering.cs b/Microsoft.Azure.Cosmos/src/Pagination/ParallelBuffering.cs new file mode 100644 index 0000000000..0b8e7988ce --- /dev/null +++ b/Microsoft.Azure.Cosmos/src/Pagination/ParallelBuffering.cs @@ -0,0 +1,56 @@ +// ------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +// ------------------------------------------------------------ + +namespace Microsoft.Azure.Cosmos.Pagination +{ + using System; + using System.Collections.Generic; + using System.Threading.Tasks; + + internal static class ParallelBuffering + { + public static async Task BufferInParallelAsync(IEnumerable bufferables, int maxConcurrency) + { + HashSet tasks = new HashSet(); + IEnumerator bufferablesEnumerator = bufferables.GetEnumerator(); + for (int i = 0; i < maxConcurrency; i++) + { + if (bufferablesEnumerator.MoveNext()) + { + IBufferable bufferable = bufferablesEnumerator.Current; + tasks.Add(Task.Run(async () => await bufferable.BufferAsync())); + } + } + + while (tasks.Count != 0) + { + Task completedTask = await Task.WhenAny(tasks); + tasks.Remove(completedTask); + try + { + await completedTask; + } + catch + { + // Observe the remaining tasks + try + { + await Task.WhenAll(tasks); + } + catch + { + } + + throw; + } + + if (bufferablesEnumerator.MoveNext()) + { + IBufferable bufferable = bufferablesEnumerator.Current; + tasks.Add(Task.Run(async () => await bufferable.BufferAsync())); + } + } + } + } +} diff --git a/Microsoft.Azure.Cosmos/src/Pagination/PartitionRangePageAsyncEnumerator.cs b/Microsoft.Azure.Cosmos/src/Pagination/PartitionRangePageAsyncEnumerator.cs index 1be4829fb7..b46436c737 100644 --- a/Microsoft.Azure.Cosmos/src/Pagination/PartitionRangePageAsyncEnumerator.cs +++ b/Microsoft.Azure.Cosmos/src/Pagination/PartitionRangePageAsyncEnumerator.cs @@ -44,10 +44,9 @@ public async ValueTask MoveNextAsync() if (this.Current.Succeeded) { this.State = this.Current.Result.State; + this.HasStarted = true; } - this.HasStarted = true; - return true; } diff --git a/Microsoft.Azure.Cosmos/src/Pagination/StartedState.cs b/Microsoft.Azure.Cosmos/src/Pagination/StartedState.cs new file mode 100644 index 0000000000..640cef70d2 --- /dev/null +++ b/Microsoft.Azure.Cosmos/src/Pagination/StartedState.cs @@ -0,0 +1,15 @@ +// ------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +// ------------------------------------------------------------ + +namespace Microsoft.Azure.Cosmos.Pagination +{ + internal sealed class StartedState : State + { + public static readonly StartedState Singleton = new StartedState(); + + private StartedState() + { + } + } +} diff --git a/Microsoft.Azure.Cosmos/src/Pagination/State.cs b/Microsoft.Azure.Cosmos/src/Pagination/State.cs index f91cf1278c..ac0a800cb9 100644 --- a/Microsoft.Azure.Cosmos/src/Pagination/State.cs +++ b/Microsoft.Azure.Cosmos/src/Pagination/State.cs @@ -4,10 +4,6 @@ namespace Microsoft.Azure.Cosmos.Pagination { - using System; - using System.Collections.Generic; - using System.Text; - internal abstract class State { } diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs index 7ecaeaa07a..feb2261077 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs @@ -387,6 +387,7 @@ private static TryCatch TryCreatePassthroughQueryExecutionC documentContainer: documentContainer, sqlQuerySpec: inputParameters.SqlQuerySpec, pageSize: inputParameters.MaxItemCount, + maxConcurrency: inputParameters.MaxConcurrency, continuationToken: inputParameters.InitialUserContinuationToken); } @@ -433,6 +434,7 @@ private static TryCatch TryCreateSpecializedDocumentQueryEx targetRanges: targetRanges, queryInfo: partitionedQueryExecutionInfo.QueryInfo, pageSize: (int)optimalPageSize, + maxConcurrency: inputParameters.MaxConcurrency, requestContinuationToken: inputParameters.InitialUserContinuationToken); } diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/PipelineFactory.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/PipelineFactory.cs index 3cfadb6b40..c8196a2529 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/PipelineFactory.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/PipelineFactory.cs @@ -29,6 +29,7 @@ public static TryCatch MonadicCreate( IReadOnlyList targetRanges, QueryInfo queryInfo, int pageSize, + int maxConcurrency, CosmosElement requestContinuationToken) { MonadicCreatePipelineStage monadicCreatePipelineStage; @@ -42,6 +43,7 @@ public static TryCatch MonadicCreate( .OrderByExpressions .Zip(queryInfo.OrderBy, (expression, sortOrder) => new OrderByColumn(expression, sortOrder)).ToList(), pageSize: pageSize, + maxConcurrency: maxConcurrency, continuationToken: requestContinuationToken); } else @@ -50,6 +52,7 @@ public static TryCatch MonadicCreate( documentContainer: documentContainer, sqlQuerySpec: sqlQuerySpec, pageSize: pageSize, + maxConcurrency: maxConcurrency, continuationToken: continuationToken); } diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderBy/OrderByCrossPartitionQueryPipelineStage.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderBy/OrderByCrossPartitionQueryPipelineStage.cs index 342177a22f..e08744fb47 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderBy/OrderByCrossPartitionQueryPipelineStage.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderBy/OrderByCrossPartitionQueryPipelineStage.cs @@ -46,11 +46,11 @@ internal sealed class OrderByCrossPartitionQueryPipelineStage : IQueryPipelineSt private static readonly IReadOnlyList EmptyPage = new List(); private readonly IDocumentContainer documentContainer; - private readonly SqlQuerySpec sqlQuerySpec; private readonly IReadOnlyList sortOrders; private readonly PriorityQueue enumerators; private readonly Queue<(OrderByQueryPartitionRangePageAsyncEnumerator, OrderByContinuationToken)> uninitializedEnumeratorsAndTokens; private readonly int pageSize; + private readonly int maxConcurrency; private QueryState state; @@ -65,17 +65,17 @@ private static class Expressions private OrderByCrossPartitionQueryPipelineStage( IDocumentContainer documentContainer, - SqlQuerySpec sqlQuerySpec, IReadOnlyList sortOrders, int pageSize, + int maxConcurrency, IEnumerable<(OrderByQueryPartitionRangePageAsyncEnumerator, OrderByContinuationToken)> uninitializedEnumeratorsAndTokens, QueryState state) { this.documentContainer = documentContainer ?? throw new ArgumentNullException(nameof(documentContainer)); - this.sqlQuerySpec = sqlQuerySpec ?? throw new ArgumentNullException(nameof(sqlQuerySpec)); this.sortOrders = sortOrders ?? throw new ArgumentNullException(nameof(sortOrders)); this.enumerators = new PriorityQueue(new OrderByEnumeratorComparer(this.sortOrders)); this.pageSize = pageSize < 0 ? throw new ArgumentOutOfRangeException($"{nameof(pageSize)} must be a non negative number.") : pageSize; + this.maxConcurrency = maxConcurrency < 0 ? throw new ArgumentOutOfRangeException($"{nameof(maxConcurrency)} must be a non negative number.") : maxConcurrency; this.uninitializedEnumeratorsAndTokens = new Queue<(OrderByQueryPartitionRangePageAsyncEnumerator, OrderByContinuationToken)>(uninitializedEnumeratorsAndTokens ?? throw new ArgumentNullException(nameof(uninitializedEnumeratorsAndTokens))); this.state = state ?? InitializingQueryState; } @@ -175,10 +175,10 @@ private async ValueTask MoveNextAsync_Iniialize_FilterAsync( } TryCatch<(bool, TryCatch)> filterMonad = await FilterNextAsync( - uninitializedEnumerator, - this.sortOrders, - token, - cancellationToken: default); + uninitializedEnumerator, + this.sortOrders, + token, + cancellationToken: default); if (filterMonad.Failed) { @@ -265,20 +265,23 @@ private async ValueTask MoveNextAsync_InitializeAsync_HandleSplitAsync( return await this.MoveNextAsync(); } - private ValueTask MoveNextAsync_InitializeAsync() + private async ValueTask MoveNextAsync_InitializeAsync() { + await ParallelBuffering.BufferInParallelAsync( + this.uninitializedEnumeratorsAndTokens.Select(value => value.Item1), + this.maxConcurrency); (OrderByQueryPartitionRangePageAsyncEnumerator uninitializedEnumerator, OrderByContinuationToken token) = this.uninitializedEnumeratorsAndTokens.Dequeue(); - ValueTask task; + bool movedNext; if (token is null) { - task = this.MoveNextAsync_Initialize_FromBeginningAsync(uninitializedEnumerator, token); + movedNext = await this.MoveNextAsync_Initialize_FromBeginningAsync(uninitializedEnumerator, token); } else { - task = this.MoveNextAsync_Iniialize_FilterAsync(uninitializedEnumerator, token); + movedNext = await this.MoveNextAsync_Iniialize_FilterAsync(uninitializedEnumerator, token); } - return task; + return movedNext; } private ValueTask MoveNextAsync_DrainPageAsync() @@ -371,6 +374,7 @@ public static TryCatch MonadicCreate( IReadOnlyList targetRanges, IReadOnlyList orderByColumns, int pageSize, + int maxConcurrency, CosmosElement continuationToken) { // TODO (brchon): For now we are not honoring non deterministic ORDER BY queries, since there is a bug in the continuation logic. @@ -498,9 +502,9 @@ public static TryCatch MonadicCreate( OrderByCrossPartitionQueryPipelineStage stage = new OrderByCrossPartitionQueryPipelineStage( documentContainer, - sqlQuerySpec, orderByColumns.Select(column => column.SortOrder).ToList(), pageSize, + maxConcurrency, enumeratorsAndTokens, continuationToken == null ? null : new QueryState(continuationToken)); return TryCatch.FromResult(stage); diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderBy/OrderByQueryPartitionRangePageAsyncEnumerator.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderBy/OrderByQueryPartitionRangePageAsyncEnumerator.cs index a5746feee2..f1032a5bef 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderBy/OrderByQueryPartitionRangePageAsyncEnumerator.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderBy/OrderByQueryPartitionRangePageAsyncEnumerator.cs @@ -12,9 +12,10 @@ namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote.OrderBy using Microsoft.Azure.Cosmos.Query.Core.Monads; using Microsoft.Azure.Documents; - internal sealed class OrderByQueryPartitionRangePageAsyncEnumerator : PartitionRangePageAsyncEnumerator + internal sealed class OrderByQueryPartitionRangePageAsyncEnumerator : PartitionRangePageAsyncEnumerator, IBufferable { - private readonly IQueryDataSource queryDataSource; + private readonly InnerEnumerator innerEnumerator; + private readonly BufferedPartitionRangePageAsyncEnumerator bufferedEnumerator; public OrderByQueryPartitionRangePageAsyncEnumerator( IQueryDataSource queryDataSource, @@ -25,44 +26,85 @@ public OrderByQueryPartitionRangePageAsyncEnumerator( QueryState state = default) : base(feedRange, state) { - this.queryDataSource = queryDataSource ?? throw new ArgumentNullException(nameof(queryDataSource)); - this.SqlQuerySpec = sqlQuerySpec ?? throw new ArgumentNullException(nameof(sqlQuerySpec)); - this.PageSize = pageSize; - this.Filter = filter; this.StartOfPageState = state; + this.innerEnumerator = new InnerEnumerator( + queryDataSource, + sqlQuerySpec, + feedRange, + pageSize, + filter, + state); + this.bufferedEnumerator = new BufferedPartitionRangePageAsyncEnumerator( + this.innerEnumerator); } - public SqlQuerySpec SqlQuerySpec { get; } + public SqlQuerySpec SqlQuerySpec => this.innerEnumerator.SqlQuerySpec; - public int PageSize { get; } + public int PageSize => this.innerEnumerator.PageSize; - public string Filter { get; } + public string Filter => this.innerEnumerator.Filter; - public QueryState StartOfPageState { get; private set; } + public QueryState StartOfPageState { get; set; } public override ValueTask DisposeAsync() => default; - protected override Task> GetNextPageAsync(CancellationToken cancellationToken) + protected override async Task> GetNextPageAsync(CancellationToken cancellationToken) { this.StartOfPageState = this.State; - return this.queryDataSource - .MonadicQueryAsync( - sqlQuerySpec: this.SqlQuerySpec, - continuationToken: this.State == null ? null : ((CosmosString)this.State.Value).Value, - feedRange: new FeedRangePartitionKeyRange(this.Range.Id), - pageSize: this.PageSize, - cancellationToken) - .ContinueWith>(antecedent => - { - TryCatch monadicQueryPage = antecedent.Result; - if (monadicQueryPage.Failed) + await this.bufferedEnumerator.MoveNextAsync(); + return this.bufferedEnumerator.Current; + } + + public ValueTask BufferAsync() => this.bufferedEnumerator.BufferAsync(); + + private sealed class InnerEnumerator : PartitionRangePageAsyncEnumerator + { + private readonly IQueryDataSource queryDataSource; + + public InnerEnumerator( + IQueryDataSource queryDataSource, + SqlQuerySpec sqlQuerySpec, + PartitionKeyRange feedRange, + int pageSize, + string filter, + QueryState state = default) + : base(feedRange, state) + { + this.queryDataSource = queryDataSource ?? throw new ArgumentNullException(nameof(queryDataSource)); + this.SqlQuerySpec = sqlQuerySpec ?? throw new ArgumentNullException(nameof(sqlQuerySpec)); + this.PageSize = pageSize; + this.Filter = filter; + } + + public SqlQuerySpec SqlQuerySpec { get; } + + public int PageSize { get; } + + public string Filter { get; } + + public override ValueTask DisposeAsync() => default; + + protected override Task> GetNextPageAsync(CancellationToken cancellationToken) + { + return this.queryDataSource + .MonadicQueryAsync( + sqlQuerySpec: this.SqlQuerySpec, + continuationToken: this.State == null ? null : ((CosmosString)this.State.Value).Value, + feedRange: new FeedRangePartitionKeyRange(this.Range.Id), + pageSize: this.PageSize, + cancellationToken) + .ContinueWith>(antecedent => { - return TryCatch.FromException(monadicQueryPage.Exception); - } + TryCatch monadicQueryPage = antecedent.Result; + if (monadicQueryPage.Failed) + { + return TryCatch.FromException(monadicQueryPage.Exception); + } - QueryPage queryPage = monadicQueryPage.Result; - return TryCatch.FromResult(new OrderByQueryPage(queryPage)); - }); + QueryPage queryPage = monadicQueryPage.Result; + return TryCatch.FromResult(new OrderByQueryPage(queryPage)); + }); + } } } } diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/Parallel/ParallelCrossPartitionQueryPipelineStage.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/Parallel/ParallelCrossPartitionQueryPipelineStage.cs index b95b076a9d..ee542b077c 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/Parallel/ParallelCrossPartitionQueryPipelineStage.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/Parallel/ParallelCrossPartitionQueryPipelineStage.cs @@ -91,6 +91,7 @@ public static TryCatch MonadicCreate( IDocumentContainer documentContainer, SqlQuerySpec sqlQuerySpec, int pageSize, + int maxConcurrency, CosmosElement continuationToken) { if (pageSize <= 0) @@ -110,6 +111,7 @@ public static TryCatch MonadicCreate( documentContainer, ParallelCrossPartitionQueryPipelineStage.MakeCreateFunction(documentContainer, sqlQuerySpec, pageSize), Comparer.Singleton, + maxConcurrency, state: state); ParallelCrossPartitionQueryPipelineStage stage = new ParallelCrossPartitionQueryPipelineStage(crossPartitionPageEnumerator); diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/BufferedPartitionRangeEnumeratorTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/BufferedPartitionRangeEnumeratorTests.cs new file mode 100644 index 0000000000..b9fac57090 --- /dev/null +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/BufferedPartitionRangeEnumeratorTests.cs @@ -0,0 +1,222 @@ +namespace Microsoft.Azure.Cosmos.Tests.Pagination +{ + using System.Collections.Generic; + using System.Threading.Tasks; + using Microsoft.Azure.Cosmos.Pagination; + using Microsoft.Azure.Cosmos.Query.Core.Monads; + using Microsoft.VisualStudio.TestTools.UnitTesting; + using Microsoft.Azure.Documents; + using System.Runtime.CompilerServices; + using System; + + [TestClass] + public sealed class BufferedPartitionPartitionRangeEnumeratorTests + { + [TestMethod] + public async Task Test429sAsync() + { + Implementation implementation = new Implementation(); + await implementation.Test429sAsync(); + } + + [TestMethod] + public async Task Test429sWithContinuationsAsync() + { + Implementation implementation = new Implementation(); + await implementation.Test429sWithContinuationsAsync(); + } + + [TestMethod] + public async Task TestDrainFullyAsync() + { + Implementation implementation = new Implementation(); + await implementation.TestDrainFullyAsync(); + } + + [TestMethod] + public async Task TestEmptyPages() + { + Implementation implementation = new Implementation(); + await implementation.TestEmptyPages(); + } + + [TestMethod] + public async Task TestResumingFromStateAsync() + { + Implementation implementation = new Implementation(); + await implementation.TestResumingFromStateAsync(); + } + + [TestMethod] + public async Task TestSplitAsync() + { + Implementation implementation = new Implementation(); + await implementation.TestSplitAsync(); + } + + [TestMethod] + public async Task TestBufferPageAsync() + { + Implementation implementation = new Implementation(); + await implementation.TestBufferPageAsync(); + } + + [TestMethod] + public async Task TestConcurrentMoveNextAndBufferPageAsync() + { + Implementation implementation = new Implementation(); + await implementation.TestConcurrentMoveNextAndBufferPageAsync(); + } + + [TestClass] + private sealed class Implementation : PartitionRangeEnumeratorTests + { + private static readonly int iterations = 1; + + public Implementation() + : base(singlePartition: true) + { + } + + [TestMethod] + public async Task TestSplitAsync() + { + int numItems = 100; + IDocumentContainer inMemoryCollection = await this.CreateDocumentContainerAsync(numItems); + IAsyncEnumerator> enumerator = this.CreateEnumerator(inMemoryCollection); + + (HashSet parentIdentifiers, DocumentContainerState state) = await this.PartialDrainAsync(enumerator, numIterations: 3); + + // Split the partition + await inMemoryCollection.SplitAsync(partitionKeyRangeId: 0, cancellationToken: default); + + // Try To read from the partition that is gone. + await enumerator.MoveNextAsync(); + Assert.IsTrue(enumerator.Current.Failed); + + // Resume on the children using the parent continuaiton token + HashSet childIdentifiers = new HashSet(); + foreach (int partitionKeyRangeId in new int[] { 1, 2 }) + { + PartitionRangePageAsyncEnumerable enumerable = new PartitionRangePageAsyncEnumerable( + range: new PartitionKeyRange() { Id = partitionKeyRangeId.ToString() }, + state: state, + (range, state) => new DocumentContainerPartitionRangeEnumerator( + inMemoryCollection, + partitionKeyRangeId: int.Parse(range.Id), + pageSize: 10, + state: state)); + HashSet resourceIdentifiers = await this.DrainFullyAsync(enumerable); + + childIdentifiers.UnionWith(resourceIdentifiers); + } + + Assert.AreEqual(numItems, parentIdentifiers.Count + childIdentifiers.Count); + } + + [TestMethod] + public async Task TestBufferPageAsync() + { + int numItems = 100; + IDocumentContainer inMemoryCollection = await this.CreateDocumentContainerAsync(numItems); + BufferedPartitionRangePageAsyncEnumerator enumerator = new BufferedPartitionRangePageAsyncEnumerator( + new DocumentContainerPartitionRangeEnumerator( + inMemoryCollection, + partitionKeyRangeId: 0, + pageSize: 10)); + + int count = 0; + + for (int i = 0; i < 10; i++) + { + // This call is idempotent; + await enumerator.BufferAsync(); + } + + Random random = new Random(); + while (await enumerator.MoveNextAsync()) + { + count += enumerator.Current.Result.Records.Count; + if (random.Next() % 2 == 0) + { + for (int i = 0; i < 10; i++) + { + // This call is idempotent; + await enumerator.BufferAsync(); + } + } + } + + Assert.AreEqual(numItems, count); + } + + [TestMethod] + public async Task TestConcurrentMoveNextAndBufferPageAsync() + { + int numItems = 100; + IDocumentContainer inMemoryCollection = await this.CreateDocumentContainerAsync(numItems); + + Random random = new Random(); + for (int iteration = 0; iteration < iterations; iteration++) + { + BufferedPartitionRangePageAsyncEnumerator enumerator = new BufferedPartitionRangePageAsyncEnumerator( + new DocumentContainerPartitionRangeEnumerator( + inMemoryCollection, + partitionKeyRangeId: 0, + pageSize: 10)); + + // BufferMore + // Fire and Forget this task. +#pragma warning disable 4014 + Task.Run(() => this.BufferMoreInBackground(enumerator)); +#pragma warning restore 4014 + + // MoveNextAsync + int count = 0; + while (await enumerator.MoveNextAsync()) + { + count += enumerator.Current.Result.Records.Count; + await Task.Delay(random.Next(0, 10)); + } + + Assert.AreEqual(numItems, count); + } + } + + public override IReadOnlyList GetRecordsFromPage(DocumentContainerPage page) + { + return page.Records; + } + + public override IAsyncEnumerable> CreateEnumerable( + IDocumentContainer documentContainer, + DocumentContainerState state = null) => new PartitionRangePageAsyncEnumerable( + range: new PartitionKeyRange() { Id = "0" }, + state: state, + (range, state) => new BufferedPartitionRangePageAsyncEnumerator( + new DocumentContainerPartitionRangeEnumerator( + documentContainer, + partitionKeyRangeId: int.Parse(range.Id), + pageSize: 10, + state: state))); + + public override IAsyncEnumerator> CreateEnumerator( + IDocumentContainer inMemoryCollection, + DocumentContainerState state = null) => new BufferedPartitionRangePageAsyncEnumerator( + new DocumentContainerPartitionRangeEnumerator( + inMemoryCollection, + partitionKeyRangeId: 0, + pageSize: 10, + state: state)); + + private async Task BufferMoreInBackground(BufferedPartitionRangePageAsyncEnumerator enumerator) + { + while (true) + { + await enumerator.BufferAsync(); + await Task.Delay(10); + } + } + } + } +} diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/CrossPartitionPartitionRangeEnumeratorTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/CrossPartitionPartitionRangeEnumeratorTests.cs index d41e228915..cb9f642786 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/CrossPartitionPartitionRangeEnumeratorTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/CrossPartitionPartitionRangeEnumeratorTests.cs @@ -147,6 +147,7 @@ PartitionRangePageAsyncEnumerator feedRangeProvider: inMemoryCollection, createPartitionRangeEnumerator: createEnumerator, comparer: PartitionRangePageAsyncEnumeratorComparer.Singleton, + maxConcurrency: 10, state: state); } @@ -166,6 +167,7 @@ PartitionRangePageAsyncEnumerator feedRangeProvider: inMemoryCollection, createPartitionRangeEnumerator: createEnumerator, comparer: PartitionRangePageAsyncEnumeratorComparer.Singleton, + maxConcurrency: 10, state: state); return enumerator; 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 68e831beaa..66bc7c6bad 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 @@ -30,6 +30,7 @@ public void TestCreate() targetRanges: new List(), queryInfo: new QueryInfo() { }, pageSize: 10, + maxConcurrency: 10, requestContinuationToken: default); Assert.IsTrue(monadicCreatePipeline.Succeeded); 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 a1f2e6315d..51c8ce2e81 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 @@ -174,6 +174,7 @@ private static IQueryPipelineStage CreatePipeline(IDocumentContainer documentCon documentContainer.GetFeedRangesAsync(default(CancellationToken)).Result, GetQueryPlan(query), pageSize: 10, + maxConcurrency: 10, requestContinuationToken: state); tryCreatePipeline.ThrowIfFailed(); diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/OrderByCrossPartitionQueryPipelineStageTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/OrderByCrossPartitionQueryPipelineStageTests.cs index 72c14ddaf7..0394906523 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/OrderByCrossPartitionQueryPipelineStageTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/OrderByCrossPartitionQueryPipelineStageTests.cs @@ -39,6 +39,7 @@ public void MonadicCreate_NullContinuationToken() new OrderByColumn("_ts", SortOrder.Ascending) }, pageSize: 10, + maxConcurrency: 10, continuationToken: null); Assert.IsTrue(monadicCreate.Succeeded); } @@ -57,6 +58,7 @@ public void MonadicCreate_NonCosmosArrayContinuationToken() new OrderByColumn("_ts", SortOrder.Ascending) }, pageSize: 10, + maxConcurrency: 10, continuationToken: CosmosObject.Create(new Dictionary())); Assert.IsTrue(monadicCreate.Failed); Assert.IsTrue(monadicCreate.InnerMostException is MalformedContinuationTokenException); @@ -76,6 +78,7 @@ public void MonadicCreate_EmptyArrayContinuationToken() new OrderByColumn("_ts", SortOrder.Ascending) }, pageSize: 10, + maxConcurrency: 10, continuationToken: CosmosArray.Create(new List())); Assert.IsTrue(monadicCreate.Failed); Assert.IsTrue(monadicCreate.InnerMostException is MalformedContinuationTokenException); @@ -95,6 +98,7 @@ public void MonadicCreate_NonParallelContinuationToken() new OrderByColumn("_ts", SortOrder.Ascending) }, pageSize: 10, + maxConcurrency: 10, continuationToken: CosmosArray.Create(new List() { CosmosString.Create("asdf") })); Assert.IsTrue(monadicCreate.Failed); Assert.IsTrue(monadicCreate.InnerMostException is MalformedContinuationTokenException); @@ -125,6 +129,7 @@ public void MonadicCreate_SingleOrderByContinuationToken() new OrderByColumn("_ts", SortOrder.Ascending) }, pageSize: 10, + maxConcurrency: 10, continuationToken: CosmosString.Create( CosmosArray.Create( new List() @@ -178,6 +183,7 @@ public void MonadicCreate_MultipleOrderByContinuationToken() new OrderByColumn("_ts", SortOrder.Ascending) }, pageSize: 10, + maxConcurrency: 10, continuationToken: CosmosString.Create( CosmosArray.Create( new List() @@ -207,6 +213,7 @@ FROM c new OrderByColumn("c._ts", SortOrder.Ascending) }, pageSize: 10, + maxConcurrency: 10, continuationToken: null); Assert.IsTrue(monadicCreate.Succeeded); IQueryPipelineStage queryPipelineStage = monadicCreate.Result; @@ -252,6 +259,7 @@ FROM c new OrderByColumn("c._ts", SortOrder.Ascending) }, pageSize: 10, + maxConcurrency: 10, continuationToken: queryState?.Value); Assert.IsTrue(monadicCreate.Succeeded); IQueryPipelineStage queryPipelineStage = monadicCreate.Result; @@ -297,6 +305,7 @@ FROM c new OrderByColumn("c._ts", SortOrder.Ascending) }, pageSize: 10, + maxConcurrency: 10, continuationToken: null); Assert.IsTrue(monadicCreate.Succeeded); IQueryPipelineStage queryPipelineStage = monadicCreate.Result; @@ -352,6 +361,7 @@ FROM c new OrderByColumn("c._ts", SortOrder.Ascending) }, pageSize: 10, + maxConcurrency: 10, continuationToken: queryState?.Value); Assert.IsTrue(monadicCreate.Succeeded); IQueryPipelineStage queryPipelineStage = monadicCreate.Result; 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 a9f61b0574..b758fd3b44 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 @@ -30,6 +30,7 @@ public void MonadicCreate_NullContinuationToken() documentContainer: mockDocumentContainer.Object, sqlQuerySpec: new SqlQuerySpec("SELECT * FROM c"), pageSize: 10, + maxConcurrency: 10, continuationToken: null); Assert.IsTrue(monadicCreate.Succeeded); } @@ -43,6 +44,7 @@ public void MonadicCreate_NonCosmosArrayContinuationToken() documentContainer: mockDocumentContainer.Object, sqlQuerySpec: new SqlQuerySpec("SELECT * FROM c"), pageSize: 10, + maxConcurrency: 10, continuationToken: CosmosObject.Create(new Dictionary())); Assert.IsTrue(monadicCreate.Failed); Assert.IsTrue(monadicCreate.InnerMostException is MalformedContinuationTokenException); @@ -57,6 +59,7 @@ public void MonadicCreate_EmptyArrayContinuationToken() documentContainer: mockDocumentContainer.Object, sqlQuerySpec: new SqlQuerySpec("SELECT * FROM c"), pageSize: 10, + maxConcurrency: 10, continuationToken: CosmosArray.Create(new List())); Assert.IsTrue(monadicCreate.Failed); Assert.IsTrue(monadicCreate.InnerMostException is MalformedContinuationTokenException); @@ -71,6 +74,7 @@ public void MonadicCreate_NonParallelContinuationToken() documentContainer: mockDocumentContainer.Object, sqlQuerySpec: new SqlQuerySpec("SELECT * FROM c"), pageSize: 10, + maxConcurrency: 10, continuationToken: CosmosArray.Create(new List() { CosmosString.Create("asdf") })); Assert.IsTrue(monadicCreate.Failed); Assert.IsTrue(monadicCreate.InnerMostException is MalformedContinuationTokenException); @@ -89,6 +93,7 @@ public void MonadicCreate_SingleParallelContinuationToken() documentContainer: mockDocumentContainer.Object, sqlQuerySpec: new SqlQuerySpec("SELECT * FROM c"), pageSize: 10, + maxConcurrency: 10, continuationToken: CosmosArray.Create(new List() { ParallelContinuationToken.ToCosmosElement(token) })); Assert.IsTrue(monadicCreate.Succeeded); } @@ -106,6 +111,7 @@ public void MonadicCreate_MultipleParallelContinuationToken() documentContainer: mockDocumentContainer.Object, sqlQuerySpec: new SqlQuerySpec("SELECT * FROM c"), pageSize: 10, + maxConcurrency: 10, continuationToken: CosmosArray.Create( new List() { @@ -125,6 +131,7 @@ public async Task TestDrainFully_StartFromBeginingAsync() documentContainer: documentContainer, sqlQuerySpec: new SqlQuerySpec("SELECT * FROM c"), pageSize: 10, + maxConcurrency: 10, continuationToken: default); Assert.IsTrue(monadicCreate.Succeeded); IQueryPipelineStage queryPipelineStage = monadicCreate.Result; @@ -157,6 +164,7 @@ public async Task TestDrainFully_WithStateResume() documentContainer: documentContainer, sqlQuerySpec: new SqlQuerySpec("SELECT * FROM c"), pageSize: 10, + maxConcurrency: 10, continuationToken: queryState?.Value); Assert.IsTrue(monadicCreate.Succeeded); IQueryPipelineStage queryPipelineStage = monadicCreate.Result; From 7bec2f3b2d0bfc1c1f6a6f3a0b03a556b58b1463 Mon Sep 17 00:00:00 2001 From: Brandon Chong Date: Fri, 25 Sep 2020 23:47:59 -0700 Subject: [PATCH 67/85] resolved iteration comments --- .../src/{Authorization => }/MurmurHash3.cs | 0 ...fferedPartitionRangePageAsyncEnumerator.cs | 46 ++----- .../CrossPartitionRangePageAsyncEnumerator.cs | 2 +- .../src/Pagination/IDocumentContainer.cs | 2 +- .../Pagination/IMonadicDocumentContainer.cs | 2 +- .../{IBufferable.cs => IPrefetcher.cs} | 5 +- .../NetworkAttachedDocumentContainer.cs | 2 +- ...rallelBuffering.cs => ParallelPrefetch.cs} | 23 ++-- .../src/Query/Core/Parser/QueryParser.cs | 116 ------------------ .../AggregateQueryPipelineStage.Compute.cs | 2 +- .../Aggregate/Aggregators/MinMaxAggregator.cs | 2 +- .../CosmosQueryExecutionContextFactory.cs | 2 +- .../ExecuteQueryBasedOnFeedRangeVisitor.cs | 2 +- .../IMonadicQueryDataSource.cs | 2 +- .../IPartitionedToken.cs | 2 +- .../IQueryDataSource.cs | 2 +- .../OrderBy/CosmosElementToQueryLiteral.cs | 2 +- .../OrderBy/ItemComparer.cs | 2 +- .../OrderBy/OrderByColumn.cs | 2 +- .../OrderBy/OrderByContinuationToken.cs | 6 +- ...OrderByCrossPartitionQueryPipelineStage.cs | 30 +++-- .../OrderBy/OrderByEnumeratorComparer.cs | 2 +- .../OrderBy/OrderByItem.cs | 2 +- .../OrderBy/OrderByQueryPage.cs | 2 +- ...yQueryPartitionRangePageAsyncEnumerator.cs | 6 +- .../OrderBy/OrderByQueryResult.cs | 2 +- .../OrderBy/SortOrder.cs | 2 +- .../Parallel/ParallelContinuationToken.cs | 2 +- ...arallelCrossPartitionQueryPipelineStage.cs | 95 +++++++------- .../QueryPartitionRangePageAsyncEnumerator.cs | 2 +- .../PartitionMapper.cs | 2 +- .../Core/Pipeline/EmptyQueryPipelineStage.cs | 2 +- .../GroupByQueryPipelineStage.Compute.cs | 2 +- .../Query/Core/Pipeline/PipelineFactory.cs | 4 +- .../src/Query/Core/QueryPlan/QueryInfo.cs | 2 +- .../DefaultDocumentQueryExecutionContext.cs | 2 +- .../Query/OrderByQueryTests.cs | 2 +- .../BufferedPartitionRangeEnumeratorTests.cs | 6 +- .../Pagination/InMemoryContainer.cs | 26 +--- .../Query/ContinuationResumeLogicTests.cs | 6 +- .../OrderByContinuationTokenTests.cs | 4 +- .../ParallelContinuationTokenTests.cs | 2 +- .../Query/Pipeline/FactoryTests.cs | 5 +- .../Query/Pipeline/MockQueryPipelineStage.cs | 2 +- ...ByCrossPartitionQueryPipelineStageTests.cs | 4 +- ...elCrossPartitionQueryPipelineStageTests.cs | 2 +- .../QueryPartitionRangePageEnumeratorTests.cs | 4 +- .../Query/QueryPlanBaselineTests.cs | 2 +- 48 files changed, 151 insertions(+), 297 deletions(-) rename Microsoft.Azure.Cosmos/src/{Authorization => }/MurmurHash3.cs (100%) rename Microsoft.Azure.Cosmos/src/Pagination/{IBufferable.cs => IPrefetcher.cs} (73%) rename Microsoft.Azure.Cosmos/src/Pagination/{ParallelBuffering.cs => ParallelPrefetch.cs} (63%) delete mode 100644 Microsoft.Azure.Cosmos/src/Query/Core/Parser/QueryParser.cs rename Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/{Remote => CrossPartition}/ExecuteQueryBasedOnFeedRangeVisitor.cs (98%) rename Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/{Remote => CrossPartition}/IMonadicQueryDataSource.cs (94%) rename Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/{Remote => CrossPartition}/IPartitionedToken.cs (81%) rename Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/{Remote => CrossPartition}/IQueryDataSource.cs (94%) rename Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/{Remote => CrossPartition}/OrderBy/CosmosElementToQueryLiteral.cs (98%) rename Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/{Remote => CrossPartition}/OrderBy/ItemComparer.cs (98%) rename Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/{Remote => CrossPartition}/OrderBy/OrderByColumn.cs (88%) rename Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/{Remote => CrossPartition}/OrderBy/OrderByContinuationToken.cs (98%) rename Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/{Remote => CrossPartition}/OrderBy/OrderByCrossPartitionQueryPipelineStage.cs (98%) rename Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/{Remote => CrossPartition}/OrderBy/OrderByEnumeratorComparer.cs (98%) rename Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/{Remote => CrossPartition}/OrderBy/OrderByItem.cs (95%) rename Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/{Remote => CrossPartition}/OrderBy/OrderByQueryPage.cs (91%) rename Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/{Remote => CrossPartition}/OrderBy/OrderByQueryPartitionRangePageAsyncEnumerator.cs (95%) rename Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/{Remote => CrossPartition}/OrderBy/OrderByQueryResult.cs (97%) rename Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/{Remote => CrossPartition}/OrderBy/SortOrder.cs (78%) rename Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/{Remote => CrossPartition}/Parallel/ParallelContinuationToken.cs (98%) rename Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/{Remote => CrossPartition}/Parallel/ParallelCrossPartitionQueryPipelineStage.cs (72%) rename Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/{Remote => CrossPartition}/Parallel/QueryPartitionRangePageAsyncEnumerator.cs (95%) rename Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/{Remote => CrossPartition}/PartitionMapper.cs (99%) diff --git a/Microsoft.Azure.Cosmos/src/Authorization/MurmurHash3.cs b/Microsoft.Azure.Cosmos/src/MurmurHash3.cs similarity index 100% rename from Microsoft.Azure.Cosmos/src/Authorization/MurmurHash3.cs rename to Microsoft.Azure.Cosmos/src/MurmurHash3.cs diff --git a/Microsoft.Azure.Cosmos/src/Pagination/BufferedPartitionRangePageAsyncEnumerator.cs b/Microsoft.Azure.Cosmos/src/Pagination/BufferedPartitionRangePageAsyncEnumerator.cs index dc11a14d58..f73f21ee42 100644 --- a/Microsoft.Azure.Cosmos/src/Pagination/BufferedPartitionRangePageAsyncEnumerator.cs +++ b/Microsoft.Azure.Cosmos/src/Pagination/BufferedPartitionRangePageAsyncEnumerator.cs @@ -9,62 +9,40 @@ namespace Microsoft.Azure.Cosmos.Pagination using System.Threading.Tasks; using Microsoft.Azure.Cosmos.Query.Core.Monads; - internal sealed class BufferedPartitionRangePageAsyncEnumerator : PartitionRangePageAsyncEnumerator, IBufferable + internal sealed class BufferedPartitionRangePageAsyncEnumerator : PartitionRangePageAsyncEnumerator, IPrefetcher where TPage : Page where TState : State { private readonly PartitionRangePageAsyncEnumerator enumerator; - private readonly SemaphoreSlim mutex; private TryCatch? bufferedPage; public BufferedPartitionRangePageAsyncEnumerator(PartitionRangePageAsyncEnumerator enumerator) : base(enumerator.Range, enumerator.State) { this.enumerator = enumerator ?? throw new ArgumentNullException(nameof(enumerator)); - this.mutex = new SemaphoreSlim(initialCount: 1, maxCount: 1); } public override ValueTask DisposeAsync() => this.enumerator.DisposeAsync(); protected override async Task> GetNextPageAsync(CancellationToken cancellationToken) { - try - { - await this.mutex.WaitAsync(cancellationToken); - if (this.bufferedPage.HasValue) - { - // Serve from the buffered page first. - TryCatch returnValue = this.bufferedPage.Value; - this.bufferedPage = null; - return returnValue; - } + await this.PrefetchAsync(); - await this.enumerator.MoveNextAsync(); - return this.enumerator.Current; - } - finally - { - this.mutex.Release(); - } + // Serve from the buffered page first. + TryCatch returnValue = this.bufferedPage.Value; + this.bufferedPage = null; + return returnValue; } - public async ValueTask BufferAsync() + public async ValueTask PrefetchAsync() { - try + if (this.bufferedPage.HasValue) { - await this.mutex.WaitAsync(); - if (this.bufferedPage.HasValue) - { - return; - } - - await this.enumerator.MoveNextAsync(); - this.bufferedPage = this.enumerator.Current; - } - finally - { - this.mutex.Release(); + return; } + + await this.enumerator.MoveNextAsync(); + this.bufferedPage = this.enumerator.Current; } } } diff --git a/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionRangePageAsyncEnumerator.cs b/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionRangePageAsyncEnumerator.cs index a07c570e71..251f59b9c6 100644 --- a/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionRangePageAsyncEnumerator.cs +++ b/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionRangePageAsyncEnumerator.cs @@ -73,7 +73,7 @@ public CrossPartitionRangePageAsyncEnumerator( if (maxConcurrency.HasValue) { - await ParallelBuffering.BufferInParallelAsync(bufferedEnumerators, maxConcurrency.Value); + await ParallelPrefetch.PrefetchInParallelAsync(bufferedEnumerators, maxConcurrency.Value); } PriorityQueue> enumerators = new PriorityQueue>( diff --git a/Microsoft.Azure.Cosmos/src/Pagination/IDocumentContainer.cs b/Microsoft.Azure.Cosmos/src/Pagination/IDocumentContainer.cs index 42d223beb3..2e63b9f30d 100644 --- a/Microsoft.Azure.Cosmos/src/Pagination/IDocumentContainer.cs +++ b/Microsoft.Azure.Cosmos/src/Pagination/IDocumentContainer.cs @@ -7,7 +7,7 @@ namespace Microsoft.Azure.Cosmos.Pagination using System.Threading; using System.Threading.Tasks; using Microsoft.Azure.Cosmos.CosmosElements; - using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline.CrossPartition; using Microsoft.Azure.Documents; internal interface IDocumentContainer : IMonadicDocumentContainer, IFeedRangeProvider, IQueryDataSource diff --git a/Microsoft.Azure.Cosmos/src/Pagination/IMonadicDocumentContainer.cs b/Microsoft.Azure.Cosmos/src/Pagination/IMonadicDocumentContainer.cs index fc63f3760f..41f8885386 100644 --- a/Microsoft.Azure.Cosmos/src/Pagination/IMonadicDocumentContainer.cs +++ b/Microsoft.Azure.Cosmos/src/Pagination/IMonadicDocumentContainer.cs @@ -12,7 +12,7 @@ namespace Microsoft.Azure.Cosmos.Pagination using Microsoft.Azure.Cosmos.Query.Core; using Microsoft.Azure.Cosmos.Query.Core.Monads; using Microsoft.Azure.Cosmos.Query.Core.Pipeline; - using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline.CrossPartition; using Microsoft.Azure.Documents; internal interface IMonadicDocumentContainer : IMonadicFeedRangeProvider, IMonadicQueryDataSource diff --git a/Microsoft.Azure.Cosmos/src/Pagination/IBufferable.cs b/Microsoft.Azure.Cosmos/src/Pagination/IPrefetcher.cs similarity index 73% rename from Microsoft.Azure.Cosmos/src/Pagination/IBufferable.cs rename to Microsoft.Azure.Cosmos/src/Pagination/IPrefetcher.cs index 23625a8fc9..8bb032089f 100644 --- a/Microsoft.Azure.Cosmos/src/Pagination/IBufferable.cs +++ b/Microsoft.Azure.Cosmos/src/Pagination/IPrefetcher.cs @@ -4,11 +4,10 @@ namespace Microsoft.Azure.Cosmos.Pagination { - using System.Collections.Generic; using System.Threading.Tasks; - internal interface IBufferable + internal interface IPrefetcher { - ValueTask BufferAsync(); + ValueTask PrefetchAsync(); } } diff --git a/Microsoft.Azure.Cosmos/src/Pagination/NetworkAttachedDocumentContainer.cs b/Microsoft.Azure.Cosmos/src/Pagination/NetworkAttachedDocumentContainer.cs index 1428941d60..9b8b6019a4 100644 --- a/Microsoft.Azure.Cosmos/src/Pagination/NetworkAttachedDocumentContainer.cs +++ b/Microsoft.Azure.Cosmos/src/Pagination/NetworkAttachedDocumentContainer.cs @@ -14,7 +14,7 @@ namespace Microsoft.Azure.Cosmos.Pagination using Microsoft.Azure.Cosmos.Query.Core; using Microsoft.Azure.Cosmos.Query.Core.Monads; using Microsoft.Azure.Cosmos.Query.Core.Pipeline; - using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline.CrossPartition; using Microsoft.Azure.Cosmos.Query.Core.QueryClient; using Microsoft.Azure.Documents; diff --git a/Microsoft.Azure.Cosmos/src/Pagination/ParallelBuffering.cs b/Microsoft.Azure.Cosmos/src/Pagination/ParallelPrefetch.cs similarity index 63% rename from Microsoft.Azure.Cosmos/src/Pagination/ParallelBuffering.cs rename to Microsoft.Azure.Cosmos/src/Pagination/ParallelPrefetch.cs index 0b8e7988ce..2aae921978 100644 --- a/Microsoft.Azure.Cosmos/src/Pagination/ParallelBuffering.cs +++ b/Microsoft.Azure.Cosmos/src/Pagination/ParallelPrefetch.cs @@ -8,18 +8,23 @@ namespace Microsoft.Azure.Cosmos.Pagination using System.Collections.Generic; using System.Threading.Tasks; - internal static class ParallelBuffering + internal static class ParallelPrefetch { - public static async Task BufferInParallelAsync(IEnumerable bufferables, int maxConcurrency) + public static async Task PrefetchInParallelAsync(IEnumerable prefetchers, int maxConcurrency) { + if (prefetchers == null) + { + throw new ArgumentNullException(nameof(prefetchers)); + } + HashSet tasks = new HashSet(); - IEnumerator bufferablesEnumerator = bufferables.GetEnumerator(); + IEnumerator prefetchersEnumerator = prefetchers.GetEnumerator(); for (int i = 0; i < maxConcurrency; i++) { - if (bufferablesEnumerator.MoveNext()) + if (prefetchersEnumerator.MoveNext()) { - IBufferable bufferable = bufferablesEnumerator.Current; - tasks.Add(Task.Run(async () => await bufferable.BufferAsync())); + IPrefetcher prefetcher = prefetchersEnumerator.Current; + tasks.Add(Task.Run(async () => await prefetcher.PrefetchAsync())); } } @@ -45,10 +50,10 @@ public static async Task BufferInParallelAsync(IEnumerable bufferab throw; } - if (bufferablesEnumerator.MoveNext()) + if (prefetchersEnumerator.MoveNext()) { - IBufferable bufferable = bufferablesEnumerator.Current; - tasks.Add(Task.Run(async () => await bufferable.BufferAsync())); + IPrefetcher bufferable = prefetchersEnumerator.Current; + tasks.Add(Task.Run(async () => await bufferable.PrefetchAsync())); } } } diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Parser/QueryParser.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Parser/QueryParser.cs deleted file mode 100644 index 05c0ee0ee4..0000000000 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Parser/QueryParser.cs +++ /dev/null @@ -1,116 +0,0 @@ -// ------------------------------------------------------------ -// Copyright (c) Microsoft Corporation. All rights reserved. -// ------------------------------------------------------------ - -namespace Microsoft.Azure.Cosmos.Query.Core.Parser -{ - using System; - using System.Runtime.ExceptionServices; - using Antlr4.Runtime; - using Antlr4.Runtime.Misc; - using Microsoft.Azure.Cosmos.Query.Core.Monads; - using Microsoft.Azure.Cosmos.SqlObjects; - - internal static class QueryParser - { - public static class Monadic - { - public static TryCatch Parse(string text) - { - if (text == null) - { - throw new ArgumentNullException(nameof(text)); - } - - AntlrInputStream str = new AntlrInputStream(text); - sqlLexer lexer = new sqlLexer(str); - CommonTokenStream tokens = new CommonTokenStream(lexer); - sqlParser parser = new sqlParser(tokens) - { - ErrorHandler = ThrowExceptionOnErrors.Singleton, - }; - ErrorListener listener = new ErrorListener(parser, lexer, tokens); - parser.AddErrorListener(listener); - - sqlParser.ProgramContext programContext; - try - { - programContext = parser.program(); - } - catch (Exception ex) - { - return TryCatch.FromException(ex); - } - - if (listener.parseException != null) - { - return TryCatch.FromException(listener.parseException); - } - - SqlQuery sqlQuery = (SqlQuery)CstToAstVisitor.Singleton.Visit(programContext); - return TryCatch.FromResult(sqlQuery); - } - - private sealed class ThrowExceptionOnErrors : IAntlrErrorStrategy - { - public static readonly ThrowExceptionOnErrors Singleton = new ThrowExceptionOnErrors(); - - public bool InErrorRecoveryMode(Parser recognizer) - { - return false; - } - - public void Recover(Parser recognizer, RecognitionException e) - { - ExceptionDispatchInfo.Capture(e).Throw(); - } - - [return: NotNull] - public IToken RecoverInline(Parser recognizer) - { - throw new NotSupportedException("can not recover."); - } - - public void ReportError(Parser recognizer, RecognitionException e) - { - ExceptionDispatchInfo.Capture(e).Throw(); - } - - public void ReportMatch(Parser recognizer) - { - // Do nothing - } - - public void Reset(Parser recognizer) - { - // Do nothing - } - - public void Sync(Parser recognizer) - { - // Do nothing - } - } - } - - public static bool TryParse(string text, out SqlQuery sqlQuery) - { - TryCatch monadicParse = Monadic.Parse(text); - if (monadicParse.Failed) - { - sqlQuery = default; - return false; - } - - sqlQuery = monadicParse.Result; - return true; - } - - public static SqlQuery Parse(string text) - { - TryCatch monadicParse = Monadic.Parse(text); - monadicParse.ThrowIfFailed(); - return monadicParse.Result; - } - } -} diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Aggregate/AggregateQueryPipelineStage.Compute.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Aggregate/AggregateQueryPipelineStage.Compute.cs index 66712e79c1..e4c78d3f1e 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Aggregate/AggregateQueryPipelineStage.Compute.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Aggregate/AggregateQueryPipelineStage.Compute.cs @@ -69,7 +69,7 @@ public static TryCatch MonadicCreate( TryCatch tryCreateSource; if (aggregateContinuationToken.SourceContinuationToken is CosmosString stringToken && (stringToken.Value == DoneSourceToken.Value)) { - tryCreateSource = TryCatch.FromResult(EmptyQueryPipelineStage.Value); + tryCreateSource = TryCatch.FromResult(EmptyQueryPipelineStage.Singleton); } else { diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Aggregate/Aggregators/MinMaxAggregator.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Aggregate/Aggregators/MinMaxAggregator.cs index 7d76bc73bd..9e76fa4197 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Aggregate/Aggregators/MinMaxAggregator.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Aggregate/Aggregators/MinMaxAggregator.cs @@ -9,7 +9,7 @@ namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.Aggregate.Aggregators using Microsoft.Azure.Cosmos.CosmosElements; using Microsoft.Azure.Cosmos.Query.Core.Exceptions; using Microsoft.Azure.Cosmos.Query.Core.Monads; - using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote.OrderBy; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline.CrossPartition.OrderBy; /// /// Concrete implementation of IAggregator that can take the global min/max from the local min/max of multiple partitions and continuations. diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs index 27c6a8fca7..ac30d81f8c 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs @@ -16,8 +16,8 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionContext using Microsoft.Azure.Cosmos.Query.Core.Monads; using Microsoft.Azure.Cosmos.Query.Core.Parser; using Microsoft.Azure.Cosmos.Query.Core.Pipeline; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline.CrossPartition.Parallel; using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Distinct; - using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote.Parallel; using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Tokens; using Microsoft.Azure.Cosmos.Query.Core.QueryClient; using Microsoft.Azure.Cosmos.Query.Core.QueryPlan; diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/ExecuteQueryBasedOnFeedRangeVisitor.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/ExecuteQueryBasedOnFeedRangeVisitor.cs similarity index 98% rename from Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/ExecuteQueryBasedOnFeedRangeVisitor.cs rename to Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/ExecuteQueryBasedOnFeedRangeVisitor.cs index 668a7af6f5..eeea5402d6 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/ExecuteQueryBasedOnFeedRangeVisitor.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/ExecuteQueryBasedOnFeedRangeVisitor.cs @@ -2,7 +2,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // ------------------------------------------------------------ -namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote +namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.CrossPartition { using System; using System.Collections.Generic; diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/IMonadicQueryDataSource.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/IMonadicQueryDataSource.cs similarity index 94% rename from Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/IMonadicQueryDataSource.cs rename to Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/IMonadicQueryDataSource.cs index 8e547394a3..28e7b456cc 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/IMonadicQueryDataSource.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/IMonadicQueryDataSource.cs @@ -2,7 +2,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // ------------------------------------------------------------ -namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote +namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.CrossPartition { using System; using System.Threading; diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/IPartitionedToken.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/IPartitionedToken.cs similarity index 81% rename from Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/IPartitionedToken.cs rename to Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/IPartitionedToken.cs index 5575bbee66..bc719a7fbf 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/IPartitionedToken.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/IPartitionedToken.cs @@ -2,7 +2,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. //------------------------------------------------------------ -namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote +namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.CrossPartition { internal interface IPartitionedToken { diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/IQueryDataSource.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/IQueryDataSource.cs similarity index 94% rename from Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/IQueryDataSource.cs rename to Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/IQueryDataSource.cs index d813e4ad22..6b35d22821 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/IQueryDataSource.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/IQueryDataSource.cs @@ -2,7 +2,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // ------------------------------------------------------------ -namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote +namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.CrossPartition { using System.Threading; using System.Threading.Tasks; diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderBy/CosmosElementToQueryLiteral.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/OrderBy/CosmosElementToQueryLiteral.cs similarity index 98% rename from Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderBy/CosmosElementToQueryLiteral.cs rename to Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/OrderBy/CosmosElementToQueryLiteral.cs index 75790b6194..eeb15f591a 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderBy/CosmosElementToQueryLiteral.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/OrderBy/CosmosElementToQueryLiteral.cs @@ -2,7 +2,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // ------------------------------------------------------------ -namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote.OrderBy +namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.CrossPartition.OrderBy { using System; using System.Collections.Generic; diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderBy/ItemComparer.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/OrderBy/ItemComparer.cs similarity index 98% rename from Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderBy/ItemComparer.cs rename to Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/OrderBy/ItemComparer.cs index 64db04ed2c..60e66f10ff 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderBy/ItemComparer.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/OrderBy/ItemComparer.cs @@ -2,7 +2,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. //------------------------------------------------------------ -namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote.OrderBy +namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.CrossPartition.OrderBy { using System; using System.Collections.Generic; diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderBy/OrderByColumn.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/OrderBy/OrderByColumn.cs similarity index 88% rename from Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderBy/OrderByColumn.cs rename to Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/OrderBy/OrderByColumn.cs index 46253c76dc..fc09fe9109 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderBy/OrderByColumn.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/OrderBy/OrderByColumn.cs @@ -2,7 +2,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // ------------------------------------------------------------ -namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote.OrderBy +namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.CrossPartition.OrderBy { using System; diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderBy/OrderByContinuationToken.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/OrderBy/OrderByContinuationToken.cs similarity index 98% rename from Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderBy/OrderByContinuationToken.cs rename to Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/OrderBy/OrderByContinuationToken.cs index 1ff54544ff..3c638f27ac 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderBy/OrderByContinuationToken.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/OrderBy/OrderByContinuationToken.cs @@ -2,7 +2,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. //------------------------------------------------------------ -namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote.OrderBy +namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.CrossPartition.OrderBy { using System; using System.Collections.Generic; @@ -11,8 +11,8 @@ namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote.OrderBy using Microsoft.Azure.Cosmos.CosmosElements.Numbers; using Microsoft.Azure.Cosmos.Query.Core.Exceptions; using Microsoft.Azure.Cosmos.Query.Core.Monads; - using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote.OrderBy; - using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote.Parallel; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline.CrossPartition.OrderBy; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline.CrossPartition.Parallel; using Microsoft.Azure.Documents.Routing; using Newtonsoft.Json; diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderBy/OrderByCrossPartitionQueryPipelineStage.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/OrderBy/OrderByCrossPartitionQueryPipelineStage.cs similarity index 98% rename from Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderBy/OrderByCrossPartitionQueryPipelineStage.cs rename to Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/OrderBy/OrderByCrossPartitionQueryPipelineStage.cs index e08744fb47..fb425cef7d 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderBy/OrderByCrossPartitionQueryPipelineStage.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/OrderBy/OrderByCrossPartitionQueryPipelineStage.cs @@ -2,7 +2,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // ------------------------------------------------------------ -namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote.OrderBy +namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.CrossPartition.OrderBy { using System; using System.Collections.Generic; @@ -16,9 +16,9 @@ namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote.OrderBy using Microsoft.Azure.Cosmos.Query.Core.Collections; using Microsoft.Azure.Cosmos.Query.Core.Exceptions; using Microsoft.Azure.Cosmos.Query.Core.Monads; - using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote.Parallel; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline.CrossPartition.Parallel; using Microsoft.Azure.Documents; - using static Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote.PartitionMapper; + using static Microsoft.Azure.Cosmos.Query.Core.Pipeline.CrossPartition.PartitionMapper; /// /// CosmosOrderByItemQueryExecutionContext is a concrete implementation for CrossPartitionQueryExecutionContext. @@ -48,7 +48,7 @@ internal sealed class OrderByCrossPartitionQueryPipelineStage : IQueryPipelineSt private readonly IDocumentContainer documentContainer; private readonly IReadOnlyList sortOrders; private readonly PriorityQueue enumerators; - private readonly Queue<(OrderByQueryPartitionRangePageAsyncEnumerator, OrderByContinuationToken)> uninitializedEnumeratorsAndTokens; + private readonly Queue<(OrderByQueryPartitionRangePageAsyncEnumerator enumerator, OrderByContinuationToken token)> uninitializedEnumeratorsAndTokens; private readonly int pageSize; private readonly int maxConcurrency; @@ -85,12 +85,11 @@ private OrderByCrossPartitionQueryPipelineStage( public ValueTask DisposeAsync() => default; private async ValueTask MoveNextAsync_Initialize_FromBeginningAsync( - OrderByQueryPartitionRangePageAsyncEnumerator uninitializedEnumerator, - OrderByContinuationToken token) + OrderByQueryPartitionRangePageAsyncEnumerator uninitializedEnumerator) { - if (token != null) + if (uninitializedEnumerator == null) { - throw new ArgumentException(nameof(token)); + throw new ArgumentNullException(nameof(uninitializedEnumerator)); } // We need to prime the page @@ -113,10 +112,10 @@ private async ValueTask MoveNextAsync_Initialize_FromBeginningAsync( { if (IsSplitException(uninitializedEnumerator.Current.Exception)) { - return await this.MoveNextAsync_InitializeAsync_HandleSplitAsync(uninitializedEnumerator, token); + return await this.MoveNextAsync_InitializeAsync_HandleSplitAsync(uninitializedEnumerator, token: null); } - this.uninitializedEnumeratorsAndTokens.Enqueue((uninitializedEnumerator, token)); + this.uninitializedEnumeratorsAndTokens.Enqueue((uninitializedEnumerator, token: null)); this.Current = TryCatch.FromException(uninitializedEnumerator.Current.Exception); } else @@ -126,7 +125,7 @@ private async ValueTask MoveNextAsync_Initialize_FromBeginningAsync( // Page was empty if (uninitializedEnumerator.State != null) { - this.uninitializedEnumeratorsAndTokens.Enqueue((uninitializedEnumerator, token)); + this.uninitializedEnumeratorsAndTokens.Enqueue((uninitializedEnumerator, token: null)); } if ((this.uninitializedEnumeratorsAndTokens.Count == 0) && (this.enumerators.Count == 0)) @@ -169,6 +168,11 @@ private async ValueTask MoveNextAsync_Iniialize_FilterAsync( OrderByQueryPartitionRangePageAsyncEnumerator uninitializedEnumerator, OrderByContinuationToken token) { + if (uninitializedEnumerator == null) + { + throw new ArgumentNullException(nameof(uninitializedEnumerator)); + } + if (token == null) { throw new ArgumentNullException(nameof(token)); @@ -267,14 +271,14 @@ private async ValueTask MoveNextAsync_InitializeAsync_HandleSplitAsync( private async ValueTask MoveNextAsync_InitializeAsync() { - await ParallelBuffering.BufferInParallelAsync( + await ParallelPrefetch.PrefetchInParallelAsync( this.uninitializedEnumeratorsAndTokens.Select(value => value.Item1), this.maxConcurrency); (OrderByQueryPartitionRangePageAsyncEnumerator uninitializedEnumerator, OrderByContinuationToken token) = this.uninitializedEnumeratorsAndTokens.Dequeue(); bool movedNext; if (token is null) { - movedNext = await this.MoveNextAsync_Initialize_FromBeginningAsync(uninitializedEnumerator, token); + movedNext = await this.MoveNextAsync_Initialize_FromBeginningAsync(uninitializedEnumerator); } else { diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderBy/OrderByEnumeratorComparer.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/OrderBy/OrderByEnumeratorComparer.cs similarity index 98% rename from Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderBy/OrderByEnumeratorComparer.cs rename to Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/OrderBy/OrderByEnumeratorComparer.cs index c3d7ef81fd..56f14abc3a 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderBy/OrderByEnumeratorComparer.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/OrderBy/OrderByEnumeratorComparer.cs @@ -2,7 +2,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // ------------------------------------------------------------ -namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote.OrderBy +namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.CrossPartition.OrderBy { using System; using System.Collections.Generic; diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderBy/OrderByItem.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/OrderBy/OrderByItem.cs similarity index 95% rename from Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderBy/OrderByItem.cs rename to Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/OrderBy/OrderByItem.cs index 4b5c7bff1d..42b90785b4 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderBy/OrderByItem.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/OrderBy/OrderByItem.cs @@ -2,7 +2,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. //------------------------------------------------------------ -namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote.OrderBy +namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.CrossPartition.OrderBy { using System; using Microsoft.Azure.Cosmos.CosmosElements; diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderBy/OrderByQueryPage.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/OrderBy/OrderByQueryPage.cs similarity index 91% rename from Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderBy/OrderByQueryPage.cs rename to Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/OrderBy/OrderByQueryPage.cs index 78f19f67b6..ca8706019c 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderBy/OrderByQueryPage.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/OrderBy/OrderByQueryPage.cs @@ -2,7 +2,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // ------------------------------------------------------------ -namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote.OrderBy +namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.CrossPartition.OrderBy { using System; using System.Collections.Generic; diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderBy/OrderByQueryPartitionRangePageAsyncEnumerator.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/OrderBy/OrderByQueryPartitionRangePageAsyncEnumerator.cs similarity index 95% rename from Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderBy/OrderByQueryPartitionRangePageAsyncEnumerator.cs rename to Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/OrderBy/OrderByQueryPartitionRangePageAsyncEnumerator.cs index f1032a5bef..be824625b0 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderBy/OrderByQueryPartitionRangePageAsyncEnumerator.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/OrderBy/OrderByQueryPartitionRangePageAsyncEnumerator.cs @@ -2,7 +2,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // ------------------------------------------------------------ -namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote.OrderBy +namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.CrossPartition.OrderBy { using System; using System.Threading; @@ -12,7 +12,7 @@ namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote.OrderBy using Microsoft.Azure.Cosmos.Query.Core.Monads; using Microsoft.Azure.Documents; - internal sealed class OrderByQueryPartitionRangePageAsyncEnumerator : PartitionRangePageAsyncEnumerator, IBufferable + internal sealed class OrderByQueryPartitionRangePageAsyncEnumerator : PartitionRangePageAsyncEnumerator, IPrefetcher { private readonly InnerEnumerator innerEnumerator; private readonly BufferedPartitionRangePageAsyncEnumerator bufferedEnumerator; @@ -55,7 +55,7 @@ protected override async Task> GetNextPageAsync(Cance return this.bufferedEnumerator.Current; } - public ValueTask BufferAsync() => this.bufferedEnumerator.BufferAsync(); + public ValueTask PrefetchAsync() => this.bufferedEnumerator.PrefetchAsync(); private sealed class InnerEnumerator : PartitionRangePageAsyncEnumerator { diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderBy/OrderByQueryResult.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/OrderBy/OrderByQueryResult.cs similarity index 97% rename from Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderBy/OrderByQueryResult.cs rename to Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/OrderBy/OrderByQueryResult.cs index 1fcc4316d2..398a06478a 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderBy/OrderByQueryResult.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/OrderBy/OrderByQueryResult.cs @@ -2,7 +2,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. //------------------------------------------------------------ -namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote.OrderBy +namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.CrossPartition.OrderBy { using System; using System.Collections.Generic; diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderBy/SortOrder.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/OrderBy/SortOrder.cs similarity index 78% rename from Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderBy/SortOrder.cs rename to Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/OrderBy/SortOrder.cs index df5b006325..84b27440e0 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/OrderBy/SortOrder.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/OrderBy/SortOrder.cs @@ -2,7 +2,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. //------------------------------------------------------------ -namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote.OrderBy +namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.CrossPartition.OrderBy { internal enum SortOrder { diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/Parallel/ParallelContinuationToken.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/Parallel/ParallelContinuationToken.cs similarity index 98% rename from Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/Parallel/ParallelContinuationToken.cs rename to Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/Parallel/ParallelContinuationToken.cs index a8e20448d9..7808612600 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/Parallel/ParallelContinuationToken.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/Parallel/ParallelContinuationToken.cs @@ -2,7 +2,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. //------------------------------------------------------------ -namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote.Parallel +namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.CrossPartition.Parallel { using System; using System.Collections.Generic; diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/Parallel/ParallelCrossPartitionQueryPipelineStage.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/Parallel/ParallelCrossPartitionQueryPipelineStage.cs similarity index 72% rename from Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/Parallel/ParallelCrossPartitionQueryPipelineStage.cs rename to Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/Parallel/ParallelCrossPartitionQueryPipelineStage.cs index ee542b077c..28b2b5894a 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/Parallel/ParallelCrossPartitionQueryPipelineStage.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/Parallel/ParallelCrossPartitionQueryPipelineStage.cs @@ -2,7 +2,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // ------------------------------------------------------------ -namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote.Parallel +namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.CrossPartition.Parallel { using System; using System.Collections.Generic; @@ -31,61 +31,66 @@ private ParallelCrossPartitionQueryPipelineStage( this.crossPartitionRangePageAsyncEnumerator = crossPartitionRangePageAsyncEnumerator ?? throw new ArgumentNullException(nameof(crossPartitionRangePageAsyncEnumerator)); } - public TryCatch Current + public TryCatch Current { get; private set; } + + public ValueTask DisposeAsync() => this.crossPartitionRangePageAsyncEnumerator.DisposeAsync(); + + public async ValueTask MoveNextAsync() { - get + if (!await this.crossPartitionRangePageAsyncEnumerator.MoveNextAsync()) { - TryCatch> currentCrossPartitionPage = this.crossPartitionRangePageAsyncEnumerator.Current; - if (currentCrossPartitionPage.Failed) - { - return TryCatch.FromException(currentCrossPartitionPage.Exception); - } - - CrossPartitionPage crossPartitionPageResult = currentCrossPartitionPage.Result; - QueryPage backendQueryPage = crossPartitionPageResult.Page; - CrossPartitionState crossPartitionState = crossPartitionPageResult.State; + this.Current = default; + return false; + } - QueryState queryState; - if (crossPartitionState == null) - { - queryState = null; - } - else - { - List parallelContinuationTokens = new List(crossPartitionState.Value.Count); - foreach ((PartitionKeyRange range, QueryState state) in crossPartitionState.Value) - { - ParallelContinuationToken parallelContinuationToken = new ParallelContinuationToken( - token: state != null ? ((CosmosString)state.Value).Value : null, - range: range.ToRange()); + TryCatch> currentCrossPartitionPage = this.crossPartitionRangePageAsyncEnumerator.Current; + if (currentCrossPartitionPage.Failed) + { + this.Current = TryCatch.FromException(currentCrossPartitionPage.Exception); + return true; + } - parallelContinuationTokens.Add(parallelContinuationToken); - } + CrossPartitionPage crossPartitionPageResult = currentCrossPartitionPage.Result; + QueryPage backendQueryPage = crossPartitionPageResult.Page; + CrossPartitionState crossPartitionState = crossPartitionPageResult.State; - List cosmosElementContinuationTokens = parallelContinuationTokens - .Select(token => ParallelContinuationToken.ToCosmosElement(token)) - .ToList(); - CosmosArray cosmosElementParallelContinuationTokens = CosmosArray.Create(cosmosElementContinuationTokens); + QueryState queryState; + if (crossPartitionState == null) + { + queryState = null; + } + else + { + List parallelContinuationTokens = new List(crossPartitionState.Value.Count); + foreach ((PartitionKeyRange range, QueryState state) in crossPartitionState.Value) + { + ParallelContinuationToken parallelContinuationToken = new ParallelContinuationToken( + token: state != null ? ((CosmosString)state.Value).Value : null, + range: range.ToRange()); - queryState = new QueryState(cosmosElementParallelContinuationTokens); + parallelContinuationTokens.Add(parallelContinuationToken); } - QueryPage crossPartitionQueryPage = new QueryPage( - backendQueryPage.Documents, - backendQueryPage.RequestCharge, - backendQueryPage.ActivityId, - backendQueryPage.ResponseLengthInBytes, - backendQueryPage.CosmosQueryExecutionInfo, - backendQueryPage.DisallowContinuationTokenMessage, - queryState); + List cosmosElementContinuationTokens = parallelContinuationTokens + .Select(token => ParallelContinuationToken.ToCosmosElement(token)) + .ToList(); + CosmosArray cosmosElementParallelContinuationTokens = CosmosArray.Create(cosmosElementContinuationTokens); - return TryCatch.FromResult(crossPartitionQueryPage); + queryState = new QueryState(cosmosElementParallelContinuationTokens); } - } - public ValueTask DisposeAsync() => this.crossPartitionRangePageAsyncEnumerator.DisposeAsync(); - - public ValueTask MoveNextAsync() => this.crossPartitionRangePageAsyncEnumerator.MoveNextAsync(); + QueryPage crossPartitionQueryPage = new QueryPage( + backendQueryPage.Documents, + backendQueryPage.RequestCharge, + backendQueryPage.ActivityId, + backendQueryPage.ResponseLengthInBytes, + backendQueryPage.CosmosQueryExecutionInfo, + backendQueryPage.DisallowContinuationTokenMessage, + queryState); + + this.Current = TryCatch.FromResult(crossPartitionQueryPage); + return true; + } public static TryCatch MonadicCreate( IDocumentContainer documentContainer, diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/Parallel/QueryPartitionRangePageAsyncEnumerator.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/Parallel/QueryPartitionRangePageAsyncEnumerator.cs similarity index 95% rename from Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/Parallel/QueryPartitionRangePageAsyncEnumerator.cs rename to Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/Parallel/QueryPartitionRangePageAsyncEnumerator.cs index 49983c0176..d25770adab 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/Parallel/QueryPartitionRangePageAsyncEnumerator.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/Parallel/QueryPartitionRangePageAsyncEnumerator.cs @@ -2,7 +2,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // ------------------------------------------------------------ -namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote.Parallel +namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.CrossPartition.Parallel { using System; using System.Threading; diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/PartitionMapper.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/PartitionMapper.cs similarity index 99% rename from Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/PartitionMapper.cs rename to Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/PartitionMapper.cs index fd2aff8608..93cb8efad6 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Remote/PartitionMapper.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/PartitionMapper.cs @@ -2,7 +2,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // ------------------------------------------------------------ -namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote +namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.CrossPartition { using System; using System.Collections.Generic; diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/EmptyQueryPipelineStage.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/EmptyQueryPipelineStage.cs index 265eeb3440..036e91e011 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/EmptyQueryPipelineStage.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/EmptyQueryPipelineStage.cs @@ -10,7 +10,7 @@ namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline internal sealed class EmptyQueryPipelineStage : IQueryPipelineStage { - public static readonly EmptyQueryPipelineStage Value = new EmptyQueryPipelineStage(); + public static readonly EmptyQueryPipelineStage Singleton = new EmptyQueryPipelineStage(); private readonly EmptyAsyncEnumerator> emptyAsyncEnumerator; diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/GroupBy/GroupByQueryPipelineStage.Compute.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/GroupBy/GroupByQueryPipelineStage.Compute.cs index 55be638e72..81e63b88d7 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/GroupBy/GroupByQueryPipelineStage.Compute.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/GroupBy/GroupByQueryPipelineStage.Compute.cs @@ -61,7 +61,7 @@ public static TryCatch MonadicCreate( if ((groupByContinuationToken.SourceContinuationToken is CosmosString sourceContinuationToken) && (sourceContinuationToken.Value == ComputeGroupByQueryPipelineStage.DoneReadingGroupingsContinuationToken)) { - tryCreateSource = TryCatch.FromResult(EmptyQueryPipelineStage.Value); + tryCreateSource = TryCatch.FromResult(EmptyQueryPipelineStage.Singleton); } else { diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/PipelineFactory.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/PipelineFactory.cs index c8196a2529..cef0436e55 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/PipelineFactory.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/PipelineFactory.cs @@ -11,10 +11,10 @@ namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline using Microsoft.Azure.Cosmos.Pagination; using Microsoft.Azure.Cosmos.Query.Core.Monads; using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Aggregate; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline.CrossPartition.OrderBy; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline.CrossPartition.Parallel; using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Distinct; using Microsoft.Azure.Cosmos.Query.Core.Pipeline.GroupBy; - using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote.OrderBy; - using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote.Parallel; using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Skip; using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Take; using Microsoft.Azure.Cosmos.Query.Core.QueryPlan; diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/QueryPlan/QueryInfo.cs b/Microsoft.Azure.Cosmos/src/Query/Core/QueryPlan/QueryInfo.cs index 341fdc275b..aac9da655e 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/QueryPlan/QueryInfo.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/QueryPlan/QueryInfo.cs @@ -7,8 +7,8 @@ namespace Microsoft.Azure.Cosmos.Query.Core.QueryPlan using System.Collections.Generic; using System.Linq; using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Aggregate; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline.CrossPartition.OrderBy; using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Distinct; - using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote.OrderBy; using Newtonsoft.Json; using Newtonsoft.Json.Converters; diff --git a/Microsoft.Azure.Cosmos/src/Query/v2Query/DefaultDocumentQueryExecutionContext.cs b/Microsoft.Azure.Cosmos/src/Query/v2Query/DefaultDocumentQueryExecutionContext.cs index 3c2435b069..66f9a1edb2 100644 --- a/Microsoft.Azure.Cosmos/src/Query/v2Query/DefaultDocumentQueryExecutionContext.cs +++ b/Microsoft.Azure.Cosmos/src/Query/v2Query/DefaultDocumentQueryExecutionContext.cs @@ -12,7 +12,7 @@ namespace Microsoft.Azure.Cosmos.Query using Microsoft.Azure.Cosmos.Common; using Microsoft.Azure.Cosmos.CosmosElements; using Microsoft.Azure.Cosmos.Query.Core.Metrics; - using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote.Parallel; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline.CrossPartition.Parallel; using Microsoft.Azure.Cosmos.Query.Core.QueryPlan; using Microsoft.Azure.Cosmos.Routing; using Microsoft.Azure.Documents; diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/Query/OrderByQueryTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/Query/OrderByQueryTests.cs index e40e63e80e..71fcabf791 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/Query/OrderByQueryTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/Query/OrderByQueryTests.cs @@ -12,7 +12,7 @@ using Microsoft.Azure.Cosmos.CosmosElements; using Microsoft.Azure.Cosmos.CosmosElements.Numbers; using Microsoft.Azure.Cosmos.Query.Core; - using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote.OrderBy; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline.CrossPartition.OrderBy; using Microsoft.Azure.Cosmos.Routing; using Microsoft.Azure.Cosmos.SDK.EmulatorTests.QueryOracle; using Microsoft.Azure.Documents; diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/BufferedPartitionRangeEnumeratorTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/BufferedPartitionRangeEnumeratorTests.cs index b9fac57090..cd0734263e 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/BufferedPartitionRangeEnumeratorTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/BufferedPartitionRangeEnumeratorTests.cs @@ -130,7 +130,7 @@ public async Task TestBufferPageAsync() for (int i = 0; i < 10; i++) { // This call is idempotent; - await enumerator.BufferAsync(); + await enumerator.PrefetchAsync(); } Random random = new Random(); @@ -142,7 +142,7 @@ public async Task TestBufferPageAsync() for (int i = 0; i < 10; i++) { // This call is idempotent; - await enumerator.BufferAsync(); + await enumerator.PrefetchAsync(); } } } @@ -213,7 +213,7 @@ private async Task BufferMoreInBackground(BufferedPartitionRangePageAsyncEnumera { while (true) { - await enumerator.BufferAsync(); + await enumerator.PrefetchAsync(); await Task.Delay(10); } } 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 57340d9ffa..4be4a82e20 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 @@ -18,7 +18,7 @@ namespace Microsoft.Azure.Cosmos.Tests.Pagination using Microsoft.Azure.Cosmos.Query.Core.Monads; using Microsoft.Azure.Cosmos.Query.Core.Parser; using Microsoft.Azure.Cosmos.Query.Core.Pipeline; - using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline.CrossPartition; using Microsoft.Azure.Cosmos.Routing; using Microsoft.Azure.Cosmos.SqlObjects; using Microsoft.Azure.Cosmos.Tests.Query.OfflineEngine; @@ -314,7 +314,7 @@ public Task> MonadicQueryAsync( throw new ArgumentNullException(nameof(sqlQuerySpec)); } - TryCatch monadicParse = QueryParser.Monadic.Parse(sqlQuerySpec.QueryText); + TryCatch monadicParse = SqlQueryParser.Monadic.Parse(sqlQuerySpec.QueryText); if (monadicParse.Failed) { return Task.FromResult(TryCatch.FromException(monadicParse.Exception)); @@ -374,16 +374,7 @@ public Task> MonadicQueryAsync( queryPageResults = queryPageResults.Take(pageSize); List queryPageResultList = queryPageResults.ToList(); - QueryState queryState; - if (queryPageResultList.Count == 0) - { - queryState = null; - } - else - { - queryState = new QueryState(((CosmosObject)queryPageResultList.Last())["_rid"]); - } - + QueryState queryState = queryPageResultList.Count == 0 ? null : new QueryState(((CosmosObject)queryPageResultList.Last())["_rid"]); return Task.FromResult( TryCatch.FromResult( new QueryPage( @@ -588,16 +579,7 @@ public Records() public Record Add(CosmosObject payload) { - ResourceId previousResourceId; - if (this.Count == 0) - { - previousResourceId = ResourceId.Empty; - } - else - { - previousResourceId = this.storage[this.storage.Count - 1].ResourceIdentifier; - } - + ResourceId previousResourceId = this.Count == 0 ? ResourceId.Empty : this.storage[this.storage.Count - 1].ResourceIdentifier; Record record = Record.Create(previousResourceId, payload); this.storage.Add(record); return record; diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/ContinuationResumeLogicTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/ContinuationResumeLogicTests.cs index cfcdec5b0c..0a6ed26b75 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/ContinuationResumeLogicTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/ContinuationResumeLogicTests.cs @@ -4,13 +4,13 @@ using System.Collections.Generic; using System.Linq; using Microsoft.Azure.Cosmos.Query.Core.Monads; - using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote; - using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote.Parallel; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline.CrossPartition; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline.CrossPartition.Parallel; using Microsoft.Azure.Cosmos.Routing; using Microsoft.Azure.Documents; using Microsoft.VisualStudio.TestTools.UnitTesting; using Newtonsoft.Json; - using static Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote.PartitionMapper; + using static Microsoft.Azure.Cosmos.Query.Core.Pipeline.CrossPartition.PartitionMapper; using Microsoft.Azure.Cosmos.Query.Core.ExecutionContext; [TestClass] diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/ContinuationTokens/OrderByContinuationTokenTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/ContinuationTokens/OrderByContinuationTokenTests.cs index bb8584da2c..969d370565 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/ContinuationTokens/OrderByContinuationTokenTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/ContinuationTokens/OrderByContinuationTokenTests.cs @@ -10,8 +10,8 @@ namespace Microsoft.Azure.Cosmos.Query using Microsoft.Azure.Cosmos.CosmosElements; using Microsoft.Azure.Cosmos.CosmosElements.Numbers; using Microsoft.Azure.Cosmos.Query.Core.Monads; - using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote.OrderBy; - using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote.Parallel; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline.CrossPartition.OrderBy; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline.CrossPartition.Parallel; using Microsoft.Azure.Documents.Routing; using VisualStudio.TestTools.UnitTesting; diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/ContinuationTokens/ParallelContinuationTokenTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/ContinuationTokens/ParallelContinuationTokenTests.cs index 7b39ca66ce..ffe376d710 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/ContinuationTokens/ParallelContinuationTokenTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/ContinuationTokens/ParallelContinuationTokenTests.cs @@ -6,7 +6,7 @@ namespace Microsoft.Azure.Cosmos.Query { using Microsoft.Azure.Cosmos.CosmosElements; using Microsoft.Azure.Cosmos.Query.Core.Monads; - using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote.Parallel; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline.CrossPartition.Parallel; using Newtonsoft.Json; using VisualStudio.TestTools.UnitTesting; 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 66bc7c6bad..4b4c8ce26a 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 @@ -9,7 +9,7 @@ namespace Microsoft.Azure.Cosmos.Tests.Query.Pipeline using Microsoft.Azure.Cosmos.Query.Core; using Microsoft.Azure.Cosmos.Query.Core.Monads; using Microsoft.Azure.Cosmos.Query.Core.Pipeline; - using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote.Parallel; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline.CrossPartition.Parallel; using Microsoft.Azure.Cosmos.Query.Core.QueryPlan; using Microsoft.Azure.Documents; using Microsoft.VisualStudio.TestTools.UnitTesting; @@ -33,9 +33,6 @@ public void TestCreate() maxConcurrency: 10, requestContinuationToken: default); Assert.IsTrue(monadicCreatePipeline.Succeeded); - - IQueryPipelineStage pipelineStage = monadicCreatePipeline.Result; - Assert.IsTrue(pipelineStage is ParallelCrossPartitionQueryPipelineStage); } } } diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/MockQueryPipelineStage.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/MockQueryPipelineStage.cs index 67943111c1..faef51ef64 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/MockQueryPipelineStage.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/MockQueryPipelineStage.cs @@ -18,7 +18,7 @@ internal sealed class MockQueryPipelineStage : QueryPipelineStageBase private int pageIndex; public MockQueryPipelineStage(IReadOnlyList> pages) - : base(EmptyQueryPipelineStage.Value) + : base(EmptyQueryPipelineStage.Singleton) { this.pages = pages ?? throw new ArgumentNullException(nameof(pages)); } diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/OrderByCrossPartitionQueryPipelineStageTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/OrderByCrossPartitionQueryPipelineStageTests.cs index 0394906523..f8db0e1f73 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/OrderByCrossPartitionQueryPipelineStageTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/OrderByCrossPartitionQueryPipelineStageTests.cs @@ -14,8 +14,8 @@ namespace Microsoft.Azure.Cosmos.Tests.Query.Pipeline using Microsoft.Azure.Cosmos.Query.Core.Exceptions; using Microsoft.Azure.Cosmos.Query.Core.Monads; using Microsoft.Azure.Cosmos.Query.Core.Pipeline; - using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote.OrderBy; - using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote.Parallel; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline.CrossPartition.OrderBy; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline.CrossPartition.Parallel; using Microsoft.Azure.Cosmos.Tests.Pagination; using Microsoft.Azure.Documents; using Microsoft.Azure.Documents.Routing; 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 b758fd3b44..2a771127ed 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 @@ -12,7 +12,7 @@ namespace Microsoft.Azure.Cosmos.Tests.Query.Pipeline using Microsoft.Azure.Cosmos.Query.Core.Exceptions; using Microsoft.Azure.Cosmos.Query.Core.Monads; using Microsoft.Azure.Cosmos.Query.Core.Pipeline; - using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote.Parallel; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline.CrossPartition.Parallel; using Microsoft.Azure.Cosmos.Tests.Pagination; using Microsoft.Azure.Documents; using Microsoft.VisualStudio.TestTools.UnitTesting; 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 a86bc0d439..8411d7434b 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 @@ -7,8 +7,8 @@ using Microsoft.Azure.Cosmos.Pagination; using Microsoft.Azure.Cosmos.Query.Core.Monads; using Microsoft.Azure.Cosmos.Query.Core.Pipeline; - using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote; - using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote.Parallel; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline.CrossPartition; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline.CrossPartition.Parallel; using Microsoft.Azure.Cosmos.Tests.Pagination; using Microsoft.Azure.Documents; using Microsoft.VisualStudio.TestTools.UnitTesting; diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/QueryPlanBaselineTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/QueryPlanBaselineTests.cs index 90edb29bd5..3304fbcd63 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/QueryPlanBaselineTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/QueryPlanBaselineTests.cs @@ -14,7 +14,7 @@ using Microsoft.Azure.Cosmos.Query.Core.Monads; using System.Linq; using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Aggregate; - using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Remote.OrderBy; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline.CrossPartition.OrderBy; /// /// Tests for . From e905d16484ff553ea51d7b469a64ef9c6d0e4b16 Mon Sep 17 00:00:00 2001 From: Brandon Chong Date: Sat, 26 Sep 2020 13:19:03 -0700 Subject: [PATCH 68/85] fixed some non baseline tests --- .../DistinctQueryPipelineStage.Client.cs | 2 +- .../DistinctQueryPipelineStage.Compute.cs | 2 +- .../SkipEmptyPageQueryPipelineStage.cs | 7 +++++- .../CosmosQueryUnitTests.cs | 4 ++-- .../BufferedPartitionRangeEnumeratorTests.cs | 22 ++++++++++--------- .../Query/Pipeline/FullPipelineTests.cs | 6 ++++- 6 files changed, 27 insertions(+), 16 deletions(-) diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Distinct/DistinctQueryPipelineStage.Client.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Distinct/DistinctQueryPipelineStage.Client.cs index e2179679e5..a35cb1b9d2 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Distinct/DistinctQueryPipelineStage.Client.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Distinct/DistinctQueryPipelineStage.Client.cs @@ -133,7 +133,7 @@ protected override async Task> GetNextPageAsync(Cancellation List distinctResults = new List(); foreach (CosmosElement document in sourcePage.Documents) { - if (this.distinctMap.Add(document, out UInt128 hash)) + if (this.distinctMap.Add(document, out UInt128 _)) { distinctResults.Add(document); } diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Distinct/DistinctQueryPipelineStage.Compute.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Distinct/DistinctQueryPipelineStage.Compute.cs index 06992788f6..1da9ccc576 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Distinct/DistinctQueryPipelineStage.Compute.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Distinct/DistinctQueryPipelineStage.Compute.cs @@ -92,7 +92,7 @@ protected override async Task> GetNextPageAsync(Cancellation List distinctResults = new List(); foreach (CosmosElement document in sourcePage.Documents) { - if (this.distinctMap.Add(document, out UInt128 hash)) + if (this.distinctMap.Add(document, out UInt128 _)) { distinctResults.Add(document); } diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/SkipEmptyPageQueryPipelineStage.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/SkipEmptyPageQueryPipelineStage.cs index 2295fe01bd..7714e9a061 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/SkipEmptyPageQueryPipelineStage.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/SkipEmptyPageQueryPipelineStage.cs @@ -25,7 +25,12 @@ public SkipEmptyPageQueryPipelineStage(IQueryPipelineStage inputStage) public async ValueTask MoveNextAsync() { - await this.inputStage.MoveNextAsync(); + if (!await this.inputStage.MoveNextAsync()) + { + this.Current = default; + return false; + } + TryCatch tryGetSourcePage = this.inputStage.Current; if (tryGetSourcePage.Failed) { diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/CosmosQueryUnitTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/CosmosQueryUnitTests.cs index 0901f234e5..6564f50211 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/CosmosQueryUnitTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/CosmosQueryUnitTests.cs @@ -251,9 +251,9 @@ public async Task TestCosmosQueryPartitionKeyDefinition() Assert.IsTrue(await pipelineStage.MoveNextAsync()); TryCatch tryGetPage = pipelineStage.Current; Assert.IsTrue(tryGetPage.Failed); - Assert.AreEqual(HttpStatusCode.BadRequest, (tryGetPage.Exception as CosmosException).StatusCode); + Assert.AreEqual(HttpStatusCode.BadRequest, (tryGetPage.InnerMostException as CosmosException).StatusCode); Assert.IsTrue( - (tryGetPage.Exception as CosmosException).ToString().Contains(exceptionMessage), + (tryGetPage.InnerMostException as CosmosException).ToString().Contains(exceptionMessage), "response error message did not contain the proper substring."); } } diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/BufferedPartitionRangeEnumeratorTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/BufferedPartitionRangeEnumeratorTests.cs index cd0734263e..cd648d6ddc 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/BufferedPartitionRangeEnumeratorTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/BufferedPartitionRangeEnumeratorTests.cs @@ -62,10 +62,10 @@ public async Task TestBufferPageAsync() } [TestMethod] - public async Task TestConcurrentMoveNextAndBufferPageAsync() + public async Task TestMoveNextAndBufferPageAsync() { Implementation implementation = new Implementation(); - await implementation.TestConcurrentMoveNextAndBufferPageAsync(); + await implementation.TestMoveNextAndBufferPageAsync(); } [TestClass] @@ -151,7 +151,7 @@ public async Task TestBufferPageAsync() } [TestMethod] - public async Task TestConcurrentMoveNextAndBufferPageAsync() + public async Task TestMoveNextAndBufferPageAsync() { int numItems = 100; IDocumentContainer inMemoryCollection = await this.CreateDocumentContainerAsync(numItems); @@ -165,18 +165,20 @@ public async Task TestConcurrentMoveNextAndBufferPageAsync() partitionKeyRangeId: 0, pageSize: 10)); - // BufferMore - // Fire and Forget this task. -#pragma warning disable 4014 - Task.Run(() => this.BufferMoreInBackground(enumerator)); -#pragma warning restore 4014 + if ((random.Next() % 2) == 0) + { + await enumerator.PrefetchAsync(); + } - // MoveNextAsync int count = 0; while (await enumerator.MoveNextAsync()) { count += enumerator.Current.Result.Records.Count; - await Task.Delay(random.Next(0, 10)); + + if ((random.Next() % 2) == 0) + { + await enumerator.PrefetchAsync(); + } } Assert.AreEqual(numItems, count); 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 51c8ce2e81..839978e754 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 @@ -113,7 +113,11 @@ private static async Task> DrainWithStateAsync(string query, { pipelineStage = CreatePipeline(documentContainer, query, state); - await pipelineStage.MoveNextAsync(); + if (!await pipelineStage.MoveNextAsync()) + { + break; + } + TryCatch tryGetQueryPage = pipelineStage.Current; tryGetQueryPage.ThrowIfFailed(); From c0fec4ed2ddfb5b8047c823702a5467134cf0ae4 Mon Sep 17 00:00:00 2001 From: Brandon Chong Date: Sat, 26 Sep 2020 13:27:07 -0700 Subject: [PATCH 69/85] updated baselines --- .../QueryPlanBaselineTests.Aggregates.xml | 29 +- .../QueryPlanBaselineTests.GroupBy.xml | 317 +++++-- ...ryPlanBaselineTests.NonValueAggregates.xml | 796 +++++++++++++----- .../QueryPlanBaselineTests.OffsetLimit.xml | 69 +- .../QueryPlanBaselineTests.OrderBy.xml | 13 +- .../QueryPlanBaselineTests.Top.xml | 6 +- 6 files changed, 917 insertions(+), 313 deletions(-) diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/BaselineTest/TestBaseline/QueryPlanBaselineTests.Aggregates.xml b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/BaselineTest/TestBaseline/QueryPlanBaselineTests.Aggregates.xml index 18f5bcfd10..7648693ffa 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/BaselineTest/TestBaseline/QueryPlanBaselineTests.Aggregates.xml +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/BaselineTest/TestBaseline/QueryPlanBaselineTests.Aggregates.xml @@ -16,7 +16,9 @@ - + + Average + True @@ -26,7 +28,8 @@ [[],"Infinity") - + @@ -48,8 +51,15 @@ - - + + + $1 + Average + + + + $1 + False @@ -57,7 +67,8 @@ [[],"Infinity") - + @@ -368,7 +379,9 @@ FROM c]]> - + + Min + True @@ -378,7 +391,9 @@ FROM c]]> [[1.0],[1.0]] - + diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/BaselineTest/TestBaseline/QueryPlanBaselineTests.GroupBy.xml b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/BaselineTest/TestBaseline/QueryPlanBaselineTests.GroupBy.xml index 675ab24ef0..0810cfd30e 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/BaselineTest/TestBaseline/QueryPlanBaselineTests.GroupBy.xml +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/BaselineTest/TestBaseline/QueryPlanBaselineTests.GroupBy.xml @@ -10,27 +10,7 @@ Hash - - - None - - - - - - - - - - False - - - - [[],"Infinity") - - - - + TryCatch resulted in an exception. @@ -44,27 +24,7 @@ Hash - - - None - - - - - - - - - - False - - - - [[],"Infinity") - - - - + TryCatch resulted in an exception. @@ -98,7 +58,9 @@ - + + c.age + @@ -111,7 +73,9 @@ [[],"Infinity") - + @@ -132,7 +96,9 @@ - + + {"age": c.age} + @@ -145,7 +111,9 @@ [[],"Infinity") - + @@ -166,12 +134,27 @@ - + + c.age + c.name + - - + + + age + null + + + name + null + + + + age + name + False @@ -179,7 +162,9 @@ [[],"Infinity") - + @@ -200,12 +185,47 @@ - + + c.team + c.name + - - + + + min_age + Min + + + max_age + Max + + + count + Count + + + name + null + + + team + null + + + avg_age + Average + + + + team + name + count + avg_age + min_age + max_age + False @@ -213,7 +233,9 @@ [[],"Infinity") - + @@ -234,12 +256,32 @@ - + + UPPER(c.name) + SUBSTRING(c.address.city, 0, 3) + - - + + + city + null + + + name + null + + + income + Average + + + + name + city + income + False @@ -247,7 +289,9 @@ [[],"Infinity") - + @@ -334,12 +378,32 @@ - + + info + g.name + - - + + + group_name + null + + + count + Count + + + info + null + + + + info + count + group_name + False @@ -347,7 +411,26 @@ [[],"Infinity") - + @@ -368,12 +451,26 @@ - + + c.name + - - + + + name + null + + + $1 + null + + + + name + $1 + False @@ -381,7 +478,9 @@ [[],"Infinity") - + @@ -402,12 +501,26 @@ - + + c.name + - - + + + name + null + + + $1 + Average + + + + name + $1 + False @@ -415,7 +528,9 @@ [[],"Infinity") - + @@ -436,10 +551,14 @@ - + + c.name + - + + Average + True @@ -449,7 +568,9 @@ [[],"Infinity") - + @@ -470,12 +591,52 @@ - + + c.team + c.name + - - + + + max_age + Max + + + name + null + + + count + Count + + + $1 + Min + + + team + null + + + $2 + Max + + + avg_age + Average + + + + team + count + name + avg_age + $1 + max_age + $2 + False @@ -483,7 +644,9 @@ [[],"Infinity") - + diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/BaselineTest/TestBaseline/QueryPlanBaselineTests.NonValueAggregates.xml b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/BaselineTest/TestBaseline/QueryPlanBaselineTests.NonValueAggregates.xml index fb845f9ae0..df0a03b653 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/BaselineTest/TestBaseline/QueryPlanBaselineTests.NonValueAggregates.xml +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/BaselineTest/TestBaseline/QueryPlanBaselineTests.NonValueAggregates.xml @@ -17,8 +17,15 @@ - - + + + $1 + Sum + + + + $1 + False @@ -26,7 +33,8 @@ [[],"Infinity") - + @@ -48,8 +56,15 @@ - - + + + $1 + Count + + + + $1 + False @@ -57,7 +72,8 @@ [[],"Infinity") - + @@ -79,8 +95,15 @@ - - + + + $1 + Min + + + + $1 + False @@ -88,7 +111,8 @@ [[],"Infinity") - + @@ -110,8 +134,15 @@ - - + + + $1 + Max + + + + $1 + False @@ -119,7 +150,8 @@ [[],"Infinity") - + @@ -141,8 +173,15 @@ - - + + + $1 + Average + + + + $1 + False @@ -150,7 +189,8 @@ [[],"Infinity") - + @@ -172,8 +212,15 @@ - - + + + sum_blah + Sum + + + + sum_blah + False @@ -181,7 +228,8 @@ [[],"Infinity") - + @@ -203,8 +251,15 @@ - - + + + count_blah + Count + + + + count_blah + False @@ -212,7 +267,8 @@ [[],"Infinity") - + @@ -234,8 +290,15 @@ - - + + + min_blah + Min + + + + min_blah + False @@ -243,7 +306,8 @@ [[],"Infinity") - + @@ -265,8 +329,15 @@ - - + + + max_blah + Max + + + + max_blah + False @@ -274,7 +345,8 @@ [[],"Infinity") - + @@ -296,8 +368,15 @@ - - + + + avg_blah + Average + + + + avg_blah + False @@ -305,7 +384,8 @@ [[],"Infinity") - + @@ -331,8 +411,20 @@ - - + + + sum_blah + Sum + + + sum_blah2 + Sum + + + + sum_blah + sum_blah2 + False @@ -340,7 +432,8 @@ [[],"Infinity") - + @@ -366,8 +459,20 @@ - - + + + count_blah + Count + + + count_blah2 + Count + + + + count_blah + count_blah2 + False @@ -375,7 +480,8 @@ [[],"Infinity") - + @@ -401,8 +507,20 @@ - - + + + min_blah2 + Min + + + min_blah + Min + + + + min_blah + min_blah2 + False @@ -410,7 +528,8 @@ [[],"Infinity") - + @@ -436,8 +555,20 @@ - - + + + max_blah + Max + + + max_blah2 + Max + + + + max_blah + max_blah2 + False @@ -445,7 +576,8 @@ [[],"Infinity") - + @@ -471,8 +603,20 @@ - - + + + avg_blah2 + Average + + + avg_blah + Average + + + + avg_blah + avg_blah2 + False @@ -480,7 +624,8 @@ [[],"Infinity") - + @@ -506,8 +651,20 @@ - - + + + $1 + Sum + + + $2 + Sum + + + + $1 + $2 + False @@ -515,7 +672,8 @@ [[],"Infinity") - + @@ -541,8 +699,20 @@ - - + + + $1 + Count + + + $2 + Count + + + + $1 + $2 + False @@ -550,7 +720,8 @@ [[],"Infinity") - + @@ -576,8 +747,20 @@ - - + + + $1 + Min + + + $2 + Min + + + + $1 + $2 + False @@ -585,7 +768,8 @@ [[],"Infinity") - + @@ -611,8 +795,20 @@ - - + + + $1 + Max + + + $2 + Max + + + + $1 + $2 + False @@ -620,7 +816,8 @@ [[],"Infinity") - + @@ -646,8 +843,20 @@ - - + + + $1 + Average + + + $2 + Average + + + + $1 + $2 + False @@ -655,7 +864,8 @@ [[],"Infinity") - + @@ -681,8 +891,20 @@ - - + + + $1 + Sum + + + sum_blah + Sum + + + + sum_blah + $1 + False @@ -690,7 +912,8 @@ [[],"Infinity") - + @@ -716,8 +939,20 @@ - - + + + $1 + Sum + + + sum_blah + Sum + + + + $1 + sum_blah + False @@ -725,7 +960,8 @@ [[],"Infinity") - + @@ -751,8 +987,20 @@ - - + + + $1 + Count + + + count_blah + Count + + + + count_blah + $1 + False @@ -760,7 +1008,8 @@ [[],"Infinity") - + @@ -786,8 +1035,20 @@ - - + + + count_blah + Count + + + $1 + Count + + + + $1 + count_blah + False @@ -795,7 +1056,8 @@ [[],"Infinity") - + @@ -821,8 +1083,20 @@ - - + + + $1 + Min + + + min_blah + Min + + + + min_blah + $1 + False @@ -830,7 +1104,8 @@ [[],"Infinity") - + @@ -856,8 +1131,20 @@ - - + + + $1 + Min + + + min_blah + Min + + + + $1 + min_blah + False @@ -865,7 +1152,8 @@ [[],"Infinity") - + @@ -891,8 +1179,20 @@ - - + + + max_blah + Max + + + $1 + Max + + + + max_blah + $1 + False @@ -900,7 +1200,8 @@ [[],"Infinity") - + @@ -926,8 +1227,20 @@ - - + + + max_blah + Max + + + $1 + Max + + + + $1 + max_blah + False @@ -935,7 +1248,8 @@ [[],"Infinity") - + @@ -961,8 +1275,20 @@ - - + + + $1 + Average + + + avg_blah + Average + + + + avg_blah + $1 + False @@ -970,7 +1296,8 @@ [[],"Infinity") - + @@ -996,8 +1323,20 @@ - - + + + $1 + Average + + + avg_blah + Average + + + + $1 + avg_blah + False @@ -1005,7 +1344,8 @@ [[],"Infinity") - + @@ -1032,8 +1372,25 @@ - - + + + $1 + Sum + + + sum_blah + Sum + + + sum_blah2 + Sum + + + + sum_blah + $1 + sum_blah2 + False @@ -1041,7 +1398,8 @@ [[],"Infinity") - + @@ -1068,8 +1426,25 @@ - - + + + $1 + Count + + + count_blah + Count + + + count_blah2 + Count + + + + count_blah + $1 + count_blah2 + False @@ -1077,7 +1452,8 @@ [[],"Infinity") - + @@ -1104,8 +1480,25 @@ - - + + + $1 + Min + + + min_blah2 + Min + + + min_blah + Min + + + + min_blah + $1 + min_blah2 + False @@ -1113,7 +1506,8 @@ [[],"Infinity") - + @@ -1140,8 +1534,25 @@ - - + + + max_blah + Max + + + max_blah2 + Max + + + $1 + Max + + + + max_blah + $1 + max_blah2 + False @@ -1149,7 +1560,8 @@ [[],"Infinity") - + @@ -1176,8 +1588,25 @@ - - + + + $1 + Average + + + avg_blah2 + Average + + + avg_blah + Average + + + + avg_blah + $1 + avg_blah2 + False @@ -1185,7 +1614,8 @@ [[],"Infinity") - + @@ -1209,8 +1639,15 @@ - - + + + $1 + Min + + + + $1 + False @@ -1218,7 +1655,9 @@ [[1.0],[1.0]] - + @@ -1242,8 +1681,20 @@ - - + + + $1 + Min + + + $2 + Max + + + + $1 + $2 + False @@ -1251,7 +1702,9 @@ [[1.0],[1.0]] - + @@ -1263,27 +1716,7 @@ Hash - - - None - - - - - - - - - - False - - - - [[],"Infinity") - - - - + TryCatch resulted in an exception. @@ -1304,8 +1737,15 @@ - - + + + $1 + Min + + + + $1 + False @@ -1313,7 +1753,8 @@ [[],"Infinity") - + @@ -1325,27 +1766,7 @@ Hash - - - None - - - - - - - - - - False - - - - [[],"Infinity") - - - - + TryCatch resulted in an exception. @@ -1356,27 +1777,7 @@ Hash - - - None - - - - - - - - - - False - - - - [[],"Infinity") - - - - + TryCatch resulted in an exception. @@ -1397,8 +1798,20 @@ - - + + + $1 + Min + + + $2 + Max + + + + $1 + $2 + False @@ -1406,7 +1819,8 @@ [[],"Infinity") - + @@ -1418,27 +1832,7 @@ Hash - - - None - - - - - - - - - - False - - - - [[],"Infinity") - - - - + TryCatch resulted in an exception. \ No newline at end of file diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/BaselineTest/TestBaseline/QueryPlanBaselineTests.OffsetLimit.xml b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/BaselineTest/TestBaseline/QueryPlanBaselineTests.OffsetLimit.xml index 0e70715efe..84df99f0a5 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/BaselineTest/TestBaseline/QueryPlanBaselineTests.OffsetLimit.xml +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/BaselineTest/TestBaseline/QueryPlanBaselineTests.OffsetLimit.xml @@ -13,8 +13,8 @@ None - - + 1 + 2 @@ -28,7 +28,9 @@ [[],"Infinity") - + @@ -56,8 +58,8 @@ None - - + 42 + 1337 @@ -71,7 +73,9 @@ [[],"Infinity") - + @@ -89,8 +93,8 @@ None - - + 1 + 2 @@ -104,7 +108,10 @@ [[],"Infinity") - + @@ -155,8 +162,8 @@ None - - + 0 + 0 @@ -170,7 +177,9 @@ [[],"Infinity") - + @@ -221,12 +230,14 @@ None - - + 1 + 2 - + + Count + True @@ -236,7 +247,9 @@ [[5.0],[5.0]] - + @@ -289,14 +302,23 @@ WHERE (c.key = 5)]]> None - - - + 1 + 2 + + c.name + - - + + + name + null + + + + name + False @@ -304,7 +326,10 @@ WHERE (c.key = 5)]]> [[5.0],[5.0]] - + diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/BaselineTest/TestBaseline/QueryPlanBaselineTests.OrderBy.xml b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/BaselineTest/TestBaseline/QueryPlanBaselineTests.OrderBy.xml index 6364a1f34f..1766f3c729 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/BaselineTest/TestBaseline/QueryPlanBaselineTests.OrderBy.xml +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/BaselineTest/TestBaseline/QueryPlanBaselineTests.OrderBy.xml @@ -176,8 +176,12 @@ ORDER BY r.key DESC]]> - - + + Ascending + + + r.key + @@ -188,7 +192,10 @@ ORDER BY r.key DESC]]> [["key"],["key"]] - + diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/BaselineTest/TestBaseline/QueryPlanBaselineTests.Top.xml b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/BaselineTest/TestBaseline/QueryPlanBaselineTests.Top.xml index 7c25a797c4..101e60d336 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/BaselineTest/TestBaseline/QueryPlanBaselineTests.Top.xml +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/BaselineTest/TestBaseline/QueryPlanBaselineTests.Top.xml @@ -117,7 +117,7 @@ None - + 5 @@ -150,7 +150,7 @@ None - + 5 @@ -216,7 +216,7 @@ None - + 0 From 7746b898be9157f1c3d92d0a0bd0812fe3436703 Mon Sep 17 00:00:00 2001 From: Brandon Chong Date: Sat, 26 Sep 2020 16:43:59 -0700 Subject: [PATCH 70/85] honoring ranges and compressed continuation tokens for parallel --- .../CosmosQueryExecutionContextFactory.cs | 75 ++-------------- ...OrderByCrossPartitionQueryPipelineStage.cs | 2 +- .../Parallel/ParallelContinuationToken.cs | 2 +- ...arallelCrossPartitionQueryPipelineStage.cs | 85 ++++++++++++++----- .../Parallel/PartitionKeyRangeComparer.cs | 44 ++++++++++ .../CrossPartition/PartitionMapper.cs | 1 - .../Query/Core/Pipeline/PipelineFactory.cs | 1 + .../ParallelContinuationTokenTests.cs | 2 +- ...elCrossPartitionQueryPipelineStageTests.cs | 28 +++++- 9 files changed, 143 insertions(+), 97 deletions(-) create mode 100644 Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/Parallel/PartitionKeyRangeComparer.cs diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs index ac30d81f8c..14c205d373 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs @@ -185,11 +185,8 @@ private static async Task> TryCreateCoreContextAsy return CosmosQueryExecutionContextFactory.TryCreatePassthroughQueryExecutionContext( documentContainer, - cosmosQueryContext, inputParameters, - partitionedQueryExecutionInfo: new PartitionedQueryExecutionInfo(), - targetRanges, - containerQueryProperties.ResourceId); + targetRanges); } } } @@ -278,11 +275,8 @@ public static async Task> TryCreateFromPartitioned tryCreatePipelineStage = CosmosQueryExecutionContextFactory.TryCreatePassthroughQueryExecutionContext( documentContainer, - cosmosQueryContext, inputParameters, - partitionedQueryExecutionInfo, - targetRanges, - containerQueryProperties.ResourceId); + targetRanges); } else { @@ -322,8 +316,7 @@ public static async Task> TryCreateFromPartitioned cosmosQueryContext, inputParameters, partitionedQueryExecutionInfo, - targetRanges, - containerQueryProperties.ResourceId); + targetRanges); } return tryCreatePipelineStage; @@ -331,61 +324,14 @@ public static async Task> TryCreateFromPartitioned private static TryCatch TryCreatePassthroughQueryExecutionContext( DocumentContainer documentContainer, - CosmosQueryContext cosmosQueryContext, InputParameters inputParameters, - PartitionedQueryExecutionInfo partitionedQueryExecutionInfo, - List targetRanges, - string collectionRid) + List targetRanges) { - //CosmosCrossPartitionQueryExecutionContext.CrossPartitionInitParams initParams = new CosmosCrossPartitionQueryExecutionContext.CrossPartitionInitParams( - // sqlQuerySpec: inputParameters.SqlQuerySpec, - // collectionRid: collectionRid, - // partitionedQueryExecutionInfo: partitionedQueryExecutionInfo, - // partitionKeyRanges: targetRanges, - // initialPageSize: inputParameters.MaxItemCount, - // maxConcurrency: inputParameters.MaxConcurrency, - // maxItemCount: inputParameters.MaxItemCount, - // maxBufferedItemCount: inputParameters.MaxBufferedItemCount, - // returnResultsInDeterministicOrder: inputParameters.ReturnResultsInDeterministicOrder, - // testSettings: inputParameters.TestInjections); - - // Modify query plan - PartitionedQueryExecutionInfo passThroughQueryInfo = new PartitionedQueryExecutionInfo() - { - QueryInfo = new QueryInfo() - { - Aggregates = null, - DistinctType = DistinctQueryType.None, - GroupByAliases = null, - GroupByAliasToAggregateType = null, - GroupByExpressions = null, - HasSelectValue = false, - Limit = null, - Offset = null, - OrderBy = null, - OrderByExpressions = null, - RewrittenQuery = null, - Top = null, - }, - QueryRanges = partitionedQueryExecutionInfo.QueryRanges, - }; - - //initParams = new CosmosCrossPartitionQueryExecutionContext.CrossPartitionInitParams( - // sqlQuerySpec: initParams.SqlQuerySpec, - // collectionRid: initParams.CollectionRid, - // partitionedQueryExecutionInfo: passThroughQueryInfo, - // partitionKeyRanges: initParams.PartitionKeyRanges, - // initialPageSize: initParams.MaxItemCount.GetValueOrDefault(1000), - // maxConcurrency: initParams.MaxConcurrency, - // maxItemCount: initParams.MaxItemCount, - // maxBufferedItemCount: initParams.MaxBufferedItemCount, - // returnResultsInDeterministicOrder: initParams.ReturnResultsInDeterministicOrder, - // testSettings: initParams.TestSettings); - // Return a parallel context, since we still want to be able to handle splits and concurrency / buffering. return ParallelCrossPartitionQueryPipelineStage.MonadicCreate( documentContainer: documentContainer, sqlQuerySpec: inputParameters.SqlQuerySpec, + targetRanges: targetRanges, pageSize: inputParameters.MaxItemCount, maxConcurrency: inputParameters.MaxConcurrency, continuationToken: inputParameters.InitialUserContinuationToken); @@ -396,8 +342,7 @@ private static TryCatch TryCreateSpecializedDocumentQueryEx CosmosQueryContext cosmosQueryContext, InputParameters inputParameters, PartitionedQueryExecutionInfo partitionedQueryExecutionInfo, - List targetRanges, - string collectionRid) + List targetRanges) { QueryInfo queryInfo = partitionedQueryExecutionInfo.QueryInfo; @@ -520,16 +465,14 @@ private static Documents.PartitionKeyDefinition GetPartitionKeyDefinition(InputP if ((inputParameters.Properties != null) && inputParameters.Properties.TryGetValue(InternalPartitionKeyDefinitionProperty, out object partitionKeyDefinitionObject)) { - if (partitionKeyDefinitionObject is Documents.PartitionKeyDefinition definition) - { - partitionKeyDefinition = definition; - } - else + if (!(partitionKeyDefinitionObject is Documents.PartitionKeyDefinition definition)) { throw new ArgumentException( "partitionkeydefinition has invalid type", nameof(partitionKeyDefinitionObject)); } + + partitionKeyDefinition = definition; } else { diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/OrderBy/OrderByCrossPartitionQueryPipelineStage.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/OrderBy/OrderByCrossPartitionQueryPipelineStage.cs index fb425cef7d..2b984d52f2 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/OrderBy/OrderByCrossPartitionQueryPipelineStage.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/OrderBy/OrderByCrossPartitionQueryPipelineStage.cs @@ -540,7 +540,7 @@ private static TryCatch> MonadicGetOr return TryCatch>.FromException(monadicExtractContinuationTokens.Exception); } - return MonadicGetPartitionMapping( + return PartitionMapper.MonadicGetPartitionMapping( partitionKeyRanges, monadicExtractContinuationTokens.Result); } diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/Parallel/ParallelContinuationToken.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/Parallel/ParallelContinuationToken.cs index 7808612600..6303d3352f 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/Parallel/ParallelContinuationToken.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/Parallel/ParallelContinuationToken.cs @@ -100,7 +100,7 @@ public static TryCatch TryCreateFromCosmosElement(Cos string max = rawMax.Value; - Range range = new Documents.Routing.Range(min, max, true, false); + Range range = new Documents.Routing.Range(min, max, isMinInclusive: true, isMaxInclusive: false); ParallelContinuationToken parallelContinuationToken = new ParallelContinuationToken(token, range); 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 28b2b5894a..7ff518f320 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 @@ -6,6 +6,7 @@ namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.CrossPartition.Parallel { using System; using System.Collections.Generic; + using System.Collections.Immutable; using System.Linq; using System.Threading.Tasks; using Microsoft.Azure.Cosmos.CosmosElements; @@ -13,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.Documents; + using static Microsoft.Azure.Cosmos.Query.Core.Pipeline.CrossPartition.PartitionMapper; /// /// is an implementation of that drain results from multiple remote nodes. @@ -61,19 +63,24 @@ public async ValueTask MoveNextAsync() } else { - List parallelContinuationTokens = new List(crossPartitionState.Value.Count); - foreach ((PartitionKeyRange range, QueryState state) in crossPartitionState.Value) + // left most and any non null continuations + List<(PartitionKeyRange, QueryState)> rangesAndStates = crossPartitionState.Value.OrderBy(tuple => tuple.Item1, PartitionKeyRangeComparer.Singleton).ToList(); + List activeParallelContinuationTokens = new List(); + for (int i = 0; i < rangesAndStates.Count; i++) { - ParallelContinuationToken parallelContinuationToken = new ParallelContinuationToken( - token: state != null ? ((CosmosString)state.Value).Value : null, - range: range.ToRange()); + (PartitionKeyRange range, QueryState state) = rangesAndStates[i]; + if ((i == 0) || (state != null)) + { + ParallelContinuationToken parallelContinuationToken = new ParallelContinuationToken( + token: state != null ? ((CosmosString)state.Value).Value : null, + range: range.ToRange()); - parallelContinuationTokens.Add(parallelContinuationToken); + activeParallelContinuationTokens.Add(parallelContinuationToken); + } } - List cosmosElementContinuationTokens = parallelContinuationTokens - .Select(token => ParallelContinuationToken.ToCosmosElement(token)) - .ToList(); + IEnumerable cosmosElementContinuationTokens = activeParallelContinuationTokens + .Select(token => ParallelContinuationToken.ToCosmosElement(token)); CosmosArray cosmosElementParallelContinuationTokens = CosmosArray.Create(cosmosElementContinuationTokens); queryState = new QueryState(cosmosElementParallelContinuationTokens); @@ -95,16 +102,27 @@ public async ValueTask MoveNextAsync() public static TryCatch MonadicCreate( IDocumentContainer documentContainer, SqlQuerySpec sqlQuerySpec, + IReadOnlyList targetRanges, int pageSize, int maxConcurrency, CosmosElement continuationToken) { + if (targetRanges == null) + { + throw new ArgumentNullException(nameof(targetRanges)); + } + + if (targetRanges.Count == 0) + { + throw new ArgumentException($"{nameof(targetRanges)} must have some elements"); + } + if (pageSize <= 0) { throw new ArgumentOutOfRangeException(nameof(pageSize)); } - TryCatch> monadicExtractState = MonadicExtractState(continuationToken); + TryCatch> monadicExtractState = MonadicExtractState(continuationToken, targetRanges); if (monadicExtractState.Failed) { return TryCatch.FromException(monadicExtractState.Exception); @@ -124,11 +142,14 @@ public static TryCatch MonadicCreate( } private static TryCatch> MonadicExtractState( - CosmosElement continuationToken) + CosmosElement continuationToken, + IReadOnlyList ranges) { if (continuationToken == null) { - return TryCatch>.FromResult(default); + // Full fan out to the ranges with null continuations + CrossPartitionState fullFanOutState = new CrossPartitionState(ranges.Select(range => (range, (QueryState)null)).ToArray()); + return TryCatch>.FromResult(fullFanOutState); } if (!(continuationToken is CosmosArray parallelContinuationTokenListRaw)) @@ -158,19 +179,37 @@ private static TryCatch> MonadicExtractState( parallelContinuationTokens.Add(tryCreateParallelContinuationToken.Result); } - List<(PartitionKeyRange, QueryState)> rangesAndStates = parallelContinuationTokens - .Select(token => ( - new PartitionKeyRange() - { - MinInclusive = token.Range.Min, - MaxExclusive = token.Range.Max, - }, - token.Token != null ? new QueryState(CosmosString.Create(token.Token)) : null)) - .ToList(); + TryCatch> partitionMappingMonad = PartitionMapper.MonadicGetPartitionMapping( + ranges, + parallelContinuationTokens); + if (partitionMappingMonad.Failed) + { + return TryCatch>.FromException( + partitionMappingMonad.Exception); + } + + PartitionMapping partitionMapping = partitionMappingMonad.Result; + List<(PartitionKeyRange, QueryState)> rangesAndStates = new List<(PartitionKeyRange, QueryState)>(); + + List> rangesToInitialize = new List>() + { + // Skip all the partitions left of the target range, since they have already been drained fully. + partitionMapping.TargetPartition, + partitionMapping.PartitionsRightOfTarget, + }; + + foreach (IReadOnlyDictionary rangeToInitalize in rangesToInitialize) + { + foreach (KeyValuePair kvp in rangeToInitalize) + { + (PartitionKeyRange, QueryState) rangeAndState = (kvp.Key, kvp.Value?.Token != null ? new QueryState(CosmosString.Create(kvp.Value.Token)) : null); + rangesAndStates.Add(rangeAndState); + } + } - CrossPartitionState state = new CrossPartitionState(rangesAndStates); + CrossPartitionState crossPartitionState = new CrossPartitionState(rangesAndStates); - return TryCatch>.FromResult(state); + return TryCatch>.FromResult(crossPartitionState); } private static CreatePartitionRangePageAsyncEnumerator MakeCreateFunction( diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/Parallel/PartitionKeyRangeComparer.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/Parallel/PartitionKeyRangeComparer.cs new file mode 100644 index 0000000000..1376b75207 --- /dev/null +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/Parallel/PartitionKeyRangeComparer.cs @@ -0,0 +1,44 @@ +// ------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +// ------------------------------------------------------------ + +namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.CrossPartition.Parallel +{ + using System; + using System.Collections.Generic; + using Microsoft.Azure.Documents; + + internal sealed class PartitionKeyRangeComparer : IComparer + { + public static readonly PartitionKeyRangeComparer Singleton = new PartitionKeyRangeComparer(); + + private PartitionKeyRangeComparer() + { + } + + public int Compare(PartitionKeyRange x, PartitionKeyRange y) + { + if (x == null) + { + throw new ArgumentNullException(nameof(x)); + } + + if (y == null) + { + throw new ArgumentNullException(nameof(y)); + } + + if (x.MinInclusive.Length == 0) + { + return -1; + } + + if (y.MinInclusive.Length == 0) + { + return 1; + } + + return x.MinInclusive.CompareTo(y.MinInclusive); + } + } +} diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/PartitionMapper.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/PartitionMapper.cs index 93cb8efad6..dc9550feb8 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/PartitionMapper.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/PartitionMapper.cs @@ -9,7 +9,6 @@ namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.CrossPartition using System.Linq; using Microsoft.Azure.Cosmos.Query.Core.Exceptions; using Microsoft.Azure.Cosmos.Query.Core.Monads; - using Microsoft.Azure.Cosmos.Routing; using Microsoft.Azure.Documents; internal static class PartitionMapper diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/PipelineFactory.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/PipelineFactory.cs index cef0436e55..11537c1a5c 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/PipelineFactory.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/PipelineFactory.cs @@ -51,6 +51,7 @@ public static TryCatch MonadicCreate( monadicCreatePipelineStage = (continuationToken) => ParallelCrossPartitionQueryPipelineStage.MonadicCreate( documentContainer: documentContainer, sqlQuerySpec: sqlQuerySpec, + targetRanges: targetRanges, pageSize: pageSize, maxConcurrency: maxConcurrency, continuationToken: continuationToken); diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/ContinuationTokens/ParallelContinuationTokenTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/ContinuationTokens/ParallelContinuationTokenTests.cs index ffe376d710..2fe9ab15f3 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/ContinuationTokens/ParallelContinuationTokenTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/ContinuationTokens/ParallelContinuationTokenTests.cs @@ -20,7 +20,7 @@ public void TestRoundTripAsCosmosElement(string token) { ParallelContinuationToken compositeContinuationToken = new ParallelContinuationToken( token, - new Documents.Routing.Range("asdf", "asdf", false, false)); + new Documents.Routing.Range("asdf", "asdf", true, false)); CosmosElement cosmosElementToken = ParallelContinuationToken.ToCosmosElement(compositeContinuationToken); TryCatch tryCompositeContinuationTokenFromCosmosElement = ParallelContinuationToken.TryCreateFromCosmosElement(cosmosElementToken); 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 2a771127ed..ac02d0379c 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 @@ -29,6 +29,7 @@ public void MonadicCreate_NullContinuationToken() TryCatch monadicCreate = ParallelCrossPartitionQueryPipelineStage.MonadicCreate( documentContainer: mockDocumentContainer.Object, sqlQuerySpec: new SqlQuerySpec("SELECT * FROM c"), + targetRanges: new List() { new PartitionKeyRange() }, pageSize: 10, maxConcurrency: 10, continuationToken: null); @@ -43,6 +44,7 @@ public void MonadicCreate_NonCosmosArrayContinuationToken() TryCatch monadicCreate = ParallelCrossPartitionQueryPipelineStage.MonadicCreate( documentContainer: mockDocumentContainer.Object, sqlQuerySpec: new SqlQuerySpec("SELECT * FROM c"), + targetRanges: new List() { new PartitionKeyRange() }, pageSize: 10, maxConcurrency: 10, continuationToken: CosmosObject.Create(new Dictionary())); @@ -58,6 +60,7 @@ public void MonadicCreate_EmptyArrayContinuationToken() TryCatch monadicCreate = ParallelCrossPartitionQueryPipelineStage.MonadicCreate( documentContainer: mockDocumentContainer.Object, sqlQuerySpec: new SqlQuerySpec("SELECT * FROM c"), + targetRanges: new List() { new PartitionKeyRange() }, pageSize: 10, maxConcurrency: 10, continuationToken: CosmosArray.Create(new List())); @@ -73,6 +76,7 @@ public void MonadicCreate_NonParallelContinuationToken() TryCatch monadicCreate = ParallelCrossPartitionQueryPipelineStage.MonadicCreate( documentContainer: mockDocumentContainer.Object, sqlQuerySpec: new SqlQuerySpec("SELECT * FROM c"), + targetRanges: new List() { new PartitionKeyRange() }, pageSize: 10, maxConcurrency: 10, continuationToken: CosmosArray.Create(new List() { CosmosString.Create("asdf") })); @@ -92,6 +96,7 @@ public void MonadicCreate_SingleParallelContinuationToken() TryCatch monadicCreate = ParallelCrossPartitionQueryPipelineStage.MonadicCreate( documentContainer: mockDocumentContainer.Object, sqlQuerySpec: new SqlQuerySpec("SELECT * FROM c"), + targetRanges: new List() { new PartitionKeyRange() { Id = "0", MinInclusive = "A", MaxExclusive = "B" } }, pageSize: 10, maxConcurrency: 10, continuationToken: CosmosArray.Create(new List() { ParallelContinuationToken.ToCosmosElement(token) })); @@ -103,20 +108,29 @@ public void MonadicCreate_MultipleParallelContinuationToken() { Mock mockDocumentContainer = new Mock(); - ParallelContinuationToken token = new ParallelContinuationToken( + ParallelContinuationToken token1 = new ParallelContinuationToken( token: "asdf", - range: new Documents.Routing.Range("A", "B", true, false)); + range: new Documents.Routing.Range("A", "B", true, false)); + + ParallelContinuationToken token2 = new ParallelContinuationToken( + token: "asdf", + range: new Documents.Routing.Range("B", "C", true, false)); TryCatch monadicCreate = ParallelCrossPartitionQueryPipelineStage.MonadicCreate( documentContainer: mockDocumentContainer.Object, sqlQuerySpec: new SqlQuerySpec("SELECT * FROM c"), + targetRanges: new List() + { + new PartitionKeyRange() { Id = "0", MinInclusive = "A", MaxExclusive = "B" }, + new PartitionKeyRange() { Id = "0", MinInclusive = "B", MaxExclusive = "C" }, + }, pageSize: 10, maxConcurrency: 10, continuationToken: CosmosArray.Create( new List() { - ParallelContinuationToken.ToCosmosElement(token), - ParallelContinuationToken.ToCosmosElement(token) + ParallelContinuationToken.ToCosmosElement(token1), + ParallelContinuationToken.ToCosmosElement(token2) })); Assert.IsTrue(monadicCreate.Succeeded); } @@ -130,6 +144,7 @@ public async Task TestDrainFully_StartFromBeginingAsync() TryCatch monadicCreate = ParallelCrossPartitionQueryPipelineStage.MonadicCreate( documentContainer: documentContainer, sqlQuerySpec: new SqlQuerySpec("SELECT * FROM c"), + targetRanges: await documentContainer.GetFeedRangesAsync(cancellationToken: default), pageSize: 10, maxConcurrency: 10, continuationToken: default); @@ -163,9 +178,14 @@ public async Task TestDrainFully_WithStateResume() TryCatch monadicCreate = ParallelCrossPartitionQueryPipelineStage.MonadicCreate( documentContainer: documentContainer, sqlQuerySpec: new SqlQuerySpec("SELECT * FROM c"), + targetRanges: await documentContainer.GetFeedRangesAsync(cancellationToken: default), pageSize: 10, maxConcurrency: 10, continuationToken: queryState?.Value); + if (monadicCreate.Failed) + { + Assert.Fail(); + } Assert.IsTrue(monadicCreate.Succeeded); IQueryPipelineStage queryPipelineStage = monadicCreate.Result; From fa4ef5acbf988d11a78f7db535a500358bf641bd Mon Sep 17 00:00:00 2001 From: Brandon Chong Date: Sat, 26 Sep 2020 23:33:42 -0700 Subject: [PATCH 71/85] wired through cancellation token --- ...fferedPartitionRangePageAsyncEnumerator.cs | 10 +-- .../CrossPartitionRangePageAsyncEnumerable.cs | 1 + .../CrossPartitionRangePageAsyncEnumerator.cs | 24 +++++-- .../src/Pagination/IPrefetcher.cs | 3 +- .../src/Pagination/ParallelPrefetch.cs | 7 +- .../PartitionRangePageAsyncEnumerator.cs | 12 +++- .../AggregateQueryPipelineStage.Client.cs | 13 ++-- .../AggregateQueryPipelineStage.Compute.cs | 17 +++-- .../Aggregate/AggregateQueryPipelineStage.cs | 12 ++-- .../Pipeline/CatchAllQueryPipelineStage.cs | 4 +- .../CosmosQueryExecutionContextFactory.cs | 8 ++- ...OrderByCrossPartitionQueryPipelineStage.cs | 72 +++++++++++++------ ...yQueryPartitionRangePageAsyncEnumerator.cs | 12 ++-- ...arallelCrossPartitionQueryPipelineStage.cs | 25 +++++-- .../QueryPartitionRangePageAsyncEnumerator.cs | 3 +- .../DistinctQueryPipelineStage.Client.cs | 26 ++++--- .../DistinctQueryPipelineStage.Compute.cs | 13 ++-- .../Distinct/DistinctQueryPipelineStage.cs | 8 ++- .../Core/Pipeline/EmptyQueryPipelineStage.cs | 6 ++ .../Pipeline/FaultedQueryPipelineStage.cs | 6 ++ .../GroupByQueryPipelineStage.Client.cs | 11 +-- .../GroupByQueryPipelineStage.Compute.cs | 9 +-- .../GroupBy/GroupByQueryPipelineStage.cs | 6 +- .../Core/Pipeline/IQueryPipelineStage.cs | 2 + .../Core/Pipeline/LazyQueryPipelineStage.cs | 14 +++- .../Pipeline/MonadicCreatePipelineStage.cs | 4 +- .../NameCacheStaleRetryQueryPipelineStage.cs | 5 ++ .../Query/Core/Pipeline/PipelineFactory.cs | 37 ++++++---- .../Core/Pipeline/QueryPipelineStageBase.cs | 14 +++- .../Skip/SkipQueryPipelineStage.Client.cs | 11 ++- .../Skip/SkipQueryPipelineStage.Compute.cs | 10 +-- .../Pipeline/Skip/SkipQueryPipelineStage.cs | 6 +- .../SkipEmptyPageQueryPipelineStage.cs | 10 ++- .../Take/TakeQueryPipelineStage.Client.cs | 11 ++- .../Take/TakeQueryPipelineStage.Compute.cs | 12 +++- .../Pipeline/Take/TakeQueryPipelineStage.cs | 11 ++- .../src/Query/v3Query/QueryIterator.cs | 1 + .../BufferedPartitionRangeEnumeratorTests.cs | 27 ++++--- ...sPartitionPartitionRangeEnumeratorTests.cs | 3 + ...cumentContainerPartitionRangeEnumerator.cs | 5 +- ...ePartitionPartitionRangeEnumeratorTests.cs | 12 ++-- .../AggregateQueryPipelineStageTests.cs | 3 +- .../DistinctQueryPipelineStageTests.cs | 4 +- .../Query/Pipeline/FactoryTests.cs | 5 +- .../Query/Pipeline/FullPipelineTests.cs | 1 + .../GroupByQueryPipelineStageTests.cs | 4 +- .../Query/Pipeline/MockQueryPipelineStage.cs | 2 +- ...ByCrossPartitionQueryPipelineStageTests.cs | 14 ++-- ...elCrossPartitionQueryPipelineStageTests.cs | 8 +++ .../Pipeline/SkipQueryPipelineStageTests.cs | 4 +- .../Pipeline/TakeQueryPipelineStageTests.cs | 5 +- .../QueryPartitionRangePageEnumeratorTests.cs | 9 ++- 52 files changed, 397 insertions(+), 165 deletions(-) diff --git a/Microsoft.Azure.Cosmos/src/Pagination/BufferedPartitionRangePageAsyncEnumerator.cs b/Microsoft.Azure.Cosmos/src/Pagination/BufferedPartitionRangePageAsyncEnumerator.cs index f73f21ee42..600ff8c0f7 100644 --- a/Microsoft.Azure.Cosmos/src/Pagination/BufferedPartitionRangePageAsyncEnumerator.cs +++ b/Microsoft.Azure.Cosmos/src/Pagination/BufferedPartitionRangePageAsyncEnumerator.cs @@ -16,8 +16,8 @@ internal sealed class BufferedPartitionRangePageAsyncEnumerator : private readonly PartitionRangePageAsyncEnumerator enumerator; private TryCatch? bufferedPage; - public BufferedPartitionRangePageAsyncEnumerator(PartitionRangePageAsyncEnumerator enumerator) - : base(enumerator.Range, enumerator.State) + public BufferedPartitionRangePageAsyncEnumerator(PartitionRangePageAsyncEnumerator enumerator, CancellationToken cancellationToken) + : base(enumerator.Range, cancellationToken, enumerator.State) { this.enumerator = enumerator ?? throw new ArgumentNullException(nameof(enumerator)); } @@ -26,7 +26,7 @@ public BufferedPartitionRangePageAsyncEnumerator(PartitionRangePageAsyncEnumerat protected override async Task> GetNextPageAsync(CancellationToken cancellationToken) { - await this.PrefetchAsync(); + await this.PrefetchAsync(cancellationToken); // Serve from the buffered page first. TryCatch returnValue = this.bufferedPage.Value; @@ -34,8 +34,10 @@ protected override async Task> GetNextPageAsync(CancellationToke return returnValue; } - public async ValueTask PrefetchAsync() + public async ValueTask PrefetchAsync(CancellationToken cancellationToken) { + cancellationToken.ThrowIfCancellationRequested(); + if (this.bufferedPage.HasValue) { return; diff --git a/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionRangePageAsyncEnumerable.cs b/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionRangePageAsyncEnumerable.cs index b5d0894512..6f8dc4277b 100644 --- a/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionRangePageAsyncEnumerable.cs +++ b/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionRangePageAsyncEnumerable.cs @@ -42,6 +42,7 @@ public IAsyncEnumerator>> GetAsyncEnu this.createPartitionRangeEnumerator, this.comparer, this.maxConcurrency, + cancellationToken, this.state); } } diff --git a/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionRangePageAsyncEnumerator.cs b/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionRangePageAsyncEnumerator.cs index 251f59b9c6..569e696ce8 100644 --- a/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionRangePageAsyncEnumerator.cs +++ b/Microsoft.Azure.Cosmos/src/Pagination/CrossPartitionRangePageAsyncEnumerator.cs @@ -25,22 +25,25 @@ internal sealed class CrossPartitionRangePageAsyncEnumerator : IA private readonly IFeedRangeProvider feedRangeProvider; private readonly CreatePartitionRangePageAsyncEnumerator createPartitionRangeEnumerator; private readonly AsyncLazy>> lazyEnumerators; + private CancellationToken cancellationToken; public CrossPartitionRangePageAsyncEnumerator( IFeedRangeProvider feedRangeProvider, CreatePartitionRangePageAsyncEnumerator createPartitionRangeEnumerator, IComparer> comparer, int? maxConcurrency, + CancellationToken cancellationToken, CrossPartitionState state = default) { - this.feedRangeProvider = feedRangeProvider ?? throw new ArgumentNullException(nameof(feedRangeProvider)); - this.createPartitionRangeEnumerator = createPartitionRangeEnumerator ?? throw new ArgumentNullException(nameof(createPartitionRangeEnumerator)); - if (comparer == null) { throw new ArgumentNullException(nameof(comparer)); } + this.feedRangeProvider = feedRangeProvider ?? throw new ArgumentNullException(nameof(feedRangeProvider)); + this.createPartitionRangeEnumerator = createPartitionRangeEnumerator ?? throw new ArgumentNullException(nameof(createPartitionRangeEnumerator)); + this.cancellationToken = cancellationToken; + this.lazyEnumerators = new AsyncLazy>>(async (CancellationToken token) => { IReadOnlyList<(PartitionKeyRange, TState)> rangeAndStates; @@ -66,14 +69,14 @@ public CrossPartitionRangePageAsyncEnumerator( .Select(rangeAndState => { PartitionRangePageAsyncEnumerator enumerator = createPartitionRangeEnumerator(rangeAndState.Item1, rangeAndState.Item2); - BufferedPartitionRangePageAsyncEnumerator bufferedEnumerator = new BufferedPartitionRangePageAsyncEnumerator(enumerator); + BufferedPartitionRangePageAsyncEnumerator bufferedEnumerator = new BufferedPartitionRangePageAsyncEnumerator(enumerator, cancellationToken); return bufferedEnumerator; }) .ToList(); if (maxConcurrency.HasValue) { - await ParallelPrefetch.PrefetchInParallelAsync(bufferedEnumerators, maxConcurrency.Value); + await ParallelPrefetch.PrefetchInParallelAsync(bufferedEnumerators, maxConcurrency.Value, token); } PriorityQueue> enumerators = new PriorityQueue>( @@ -87,7 +90,9 @@ public CrossPartitionRangePageAsyncEnumerator( public async ValueTask MoveNextAsync() { - PriorityQueue> enumerators = await this.lazyEnumerators.GetValueAsync(cancellationToken: default); + this.cancellationToken.ThrowIfCancellationRequested(); + + PriorityQueue> enumerators = await this.lazyEnumerators.GetValueAsync(cancellationToken: this.cancellationToken); if (enumerators.Count == 0) { return false; @@ -115,7 +120,7 @@ public async ValueTask MoveNextAsync() // Handle split IEnumerable childRanges = await this.feedRangeProvider.GetChildRangeAsync( currentPaginator.Range, - cancellationToken: default); + cancellationToken: this.cancellationToken); foreach (PartitionKeyRange childRange in childRanges) { PartitionRangePageAsyncEnumerator childPaginator = this.createPartitionRangeEnumerator( @@ -173,6 +178,11 @@ public ValueTask DisposeAsync() return default; } + public void SetCancellationToken(CancellationToken cancellationToken) + { + this.cancellationToken = cancellationToken; + } + private static bool IsSplitException(Exception exeception) { return exeception is CosmosException cosmosException diff --git a/Microsoft.Azure.Cosmos/src/Pagination/IPrefetcher.cs b/Microsoft.Azure.Cosmos/src/Pagination/IPrefetcher.cs index 8bb032089f..02d524dc79 100644 --- a/Microsoft.Azure.Cosmos/src/Pagination/IPrefetcher.cs +++ b/Microsoft.Azure.Cosmos/src/Pagination/IPrefetcher.cs @@ -4,10 +4,11 @@ namespace Microsoft.Azure.Cosmos.Pagination { + using System.Threading; using System.Threading.Tasks; internal interface IPrefetcher { - ValueTask PrefetchAsync(); + ValueTask PrefetchAsync(CancellationToken cancellationToken); } } diff --git a/Microsoft.Azure.Cosmos/src/Pagination/ParallelPrefetch.cs b/Microsoft.Azure.Cosmos/src/Pagination/ParallelPrefetch.cs index 2aae921978..e1138d9432 100644 --- a/Microsoft.Azure.Cosmos/src/Pagination/ParallelPrefetch.cs +++ b/Microsoft.Azure.Cosmos/src/Pagination/ParallelPrefetch.cs @@ -6,11 +6,12 @@ namespace Microsoft.Azure.Cosmos.Pagination { using System; using System.Collections.Generic; + using System.Threading; using System.Threading.Tasks; internal static class ParallelPrefetch { - public static async Task PrefetchInParallelAsync(IEnumerable prefetchers, int maxConcurrency) + public static async Task PrefetchInParallelAsync(IEnumerable prefetchers, int maxConcurrency, CancellationToken cancellationToken) { if (prefetchers == null) { @@ -24,7 +25,7 @@ public static async Task PrefetchInParallelAsync(IEnumerable prefet if (prefetchersEnumerator.MoveNext()) { IPrefetcher prefetcher = prefetchersEnumerator.Current; - tasks.Add(Task.Run(async () => await prefetcher.PrefetchAsync())); + tasks.Add(Task.Run(async () => await prefetcher.PrefetchAsync(cancellationToken))); } } @@ -53,7 +54,7 @@ public static async Task PrefetchInParallelAsync(IEnumerable prefet if (prefetchersEnumerator.MoveNext()) { IPrefetcher bufferable = prefetchersEnumerator.Current; - tasks.Add(Task.Run(async () => await bufferable.PrefetchAsync())); + tasks.Add(Task.Run(async () => await bufferable.PrefetchAsync(cancellationToken))); } } } diff --git a/Microsoft.Azure.Cosmos/src/Pagination/PartitionRangePageAsyncEnumerator.cs b/Microsoft.Azure.Cosmos/src/Pagination/PartitionRangePageAsyncEnumerator.cs index b46436c737..d0ba5c6786 100644 --- a/Microsoft.Azure.Cosmos/src/Pagination/PartitionRangePageAsyncEnumerator.cs +++ b/Microsoft.Azure.Cosmos/src/Pagination/PartitionRangePageAsyncEnumerator.cs @@ -17,10 +17,13 @@ internal abstract class PartitionRangePageAsyncEnumerator : IAsyn where TPage : Page where TState : State { - protected PartitionRangePageAsyncEnumerator(PartitionKeyRange range, TState state = default) + private CancellationToken cancellationToken; + + protected PartitionRangePageAsyncEnumerator(PartitionKeyRange range, CancellationToken cancellationToken, TState state = default) { this.Range = range; this.State = state; + this.cancellationToken = cancellationToken; } public PartitionKeyRange Range { get; } @@ -40,7 +43,7 @@ public async ValueTask MoveNextAsync() return false; } - this.Current = await this.GetNextPageAsync(cancellationToken: default); + this.Current = await this.GetNextPageAsync(cancellationToken: this.cancellationToken); if (this.Current.Succeeded) { this.State = this.Current.Result.State; @@ -53,5 +56,10 @@ public async ValueTask MoveNextAsync() protected abstract Task> GetNextPageAsync(CancellationToken cancellationToken); public abstract ValueTask DisposeAsync(); + + public void SetCancellationToken(CancellationToken cancellationToken) + { + this.cancellationToken = cancellationToken; + } } } diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Aggregate/AggregateQueryPipelineStage.Client.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Aggregate/AggregateQueryPipelineStage.Client.cs index e27fe469dc..7f918206f8 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Aggregate/AggregateQueryPipelineStage.Client.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Aggregate/AggregateQueryPipelineStage.Client.cs @@ -20,8 +20,9 @@ private sealed class ClientAggregateQueryPipelineStage : AggregateQueryPipelineS private ClientAggregateQueryPipelineStage( IQueryPipelineStage source, SingleGroupAggregator singleGroupAggregator, - bool isValueAggregateQuery) - : base(source, singleGroupAggregator, isValueAggregateQuery) + bool isValueAggregateQuery, + CancellationToken cancellationToken) + : base(source, singleGroupAggregator, isValueAggregateQuery, cancellationToken) { // all the work is done in the base constructor. } @@ -32,6 +33,7 @@ public static TryCatch MonadicCreate( IReadOnlyList orderedAliases, bool hasSelectValue, CosmosElement continuationToken, + CancellationToken cancellationToken, MonadicCreatePipelineStage monadicCreatePipelineStage) { if (monadicCreatePipelineStage == null) @@ -50,7 +52,7 @@ public static TryCatch MonadicCreate( return TryCatch.FromException(tryCreateSingleGroupAggregator.Exception); } - TryCatch tryCreateSource = monadicCreatePipelineStage(continuationToken); + TryCatch tryCreateSource = monadicCreatePipelineStage(continuationToken, cancellationToken); if (tryCreateSource.Failed) { return tryCreateSource; @@ -59,7 +61,8 @@ public static TryCatch MonadicCreate( ClientAggregateQueryPipelineStage stage = new ClientAggregateQueryPipelineStage( tryCreateSource.Result, tryCreateSingleGroupAggregator.Result, - hasSelectValue); + hasSelectValue, + cancellationToken); return TryCatch.FromResult(stage); } @@ -90,6 +93,8 @@ protected override async Task> GetNextPageAsync(Cancellation foreach (CosmosElement element in sourcePage.Documents) { + cancellationToken.ThrowIfCancellationRequested(); + RewrittenAggregateProjections rewrittenAggregateProjections = new RewrittenAggregateProjections( this.isValueQuery, element); diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Aggregate/AggregateQueryPipelineStage.Compute.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Aggregate/AggregateQueryPipelineStage.Compute.cs index e4c78d3f1e..8e9a5fed29 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Aggregate/AggregateQueryPipelineStage.Compute.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Aggregate/AggregateQueryPipelineStage.Compute.cs @@ -24,8 +24,9 @@ private sealed class ComputeAggregateQueryPipelineStage : AggregateQueryPipeline private ComputeAggregateQueryPipelineStage( IQueryPipelineStage source, SingleGroupAggregator singleGroupAggregator, - bool isValueAggregateQuery) - : base(source, singleGroupAggregator, isValueAggregateQuery) + bool isValueAggregateQuery, + CancellationToken cancellationToken) + : base(source, singleGroupAggregator, isValueAggregateQuery, cancellationToken) { // all the work is done in the base constructor. } @@ -36,8 +37,11 @@ public static TryCatch MonadicCreate( IReadOnlyList orderedAliases, bool hasSelectValue, CosmosElement continuationToken, + CancellationToken cancellationToken, MonadicCreatePipelineStage monadicCreatePipelineStage) { + cancellationToken.ThrowIfCancellationRequested(); + AggregateContinuationToken aggregateContinuationToken; if (continuationToken != null) { @@ -52,7 +56,7 @@ public static TryCatch MonadicCreate( } else { - aggregateContinuationToken = new AggregateContinuationToken(null, null); + aggregateContinuationToken = new AggregateContinuationToken(singleGroupAggregatorContinuationToken: null, sourceContinuationToken: null); } TryCatch tryCreateSingleGroupAggregator = SingleGroupAggregator.TryCreate( @@ -73,7 +77,7 @@ public static TryCatch MonadicCreate( } else { - tryCreateSource = monadicCreatePipelineStage(aggregateContinuationToken.SourceContinuationToken); + tryCreateSource = monadicCreatePipelineStage(aggregateContinuationToken.SourceContinuationToken, cancellationToken); } if (tryCreateSource.Failed) @@ -84,7 +88,8 @@ public static TryCatch MonadicCreate( ComputeAggregateQueryPipelineStage stage = new ComputeAggregateQueryPipelineStage( tryCreateSource.Result, tryCreateSingleGroupAggregator.Result, - hasSelectValue); + hasSelectValue, + cancellationToken); return TryCatch.FromResult(stage); } @@ -109,6 +114,8 @@ protected override async Task> GetNextPageAsync(Cancellation QueryPage sourcePage = tryGetSourcePage.Result; foreach (CosmosElement element in sourcePage.Documents) { + cancellationToken.ThrowIfCancellationRequested(); + RewrittenAggregateProjections rewrittenAggregateProjections = new RewrittenAggregateProjections( this.isValueQuery, element); diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Aggregate/AggregateQueryPipelineStage.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Aggregate/AggregateQueryPipelineStage.cs index 03ac6e4aa4..0530ca5c21 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Aggregate/AggregateQueryPipelineStage.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Aggregate/AggregateQueryPipelineStage.cs @@ -7,12 +7,9 @@ namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.Aggregate using System; using System.Collections.Generic; using System.Threading; - using System.Threading.Tasks; using Microsoft.Azure.Cosmos.CosmosElements; - using Microsoft.Azure.Cosmos.Query.Core.ExecutionContext; using Microsoft.Azure.Cosmos.Query.Core.Monads; using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Aggregate.Aggregators; - using static Microsoft.Azure.Cosmos.Query.Core.Pipeline.PipelineFactory; /// /// Stage that is able to aggregate local aggregates from multiple continuations and partitions. @@ -47,12 +44,14 @@ internal abstract partial class AggregateQueryPipelineStage : QueryPipelineStage /// The source component that will supply the local aggregates from multiple continuations and partitions. /// The single group aggregator that we will feed results into. /// Whether or not the query has the 'VALUE' keyword. + /// The cancellation token for cooperative yeilding. /// This constructor is private since there is some async initialization that needs to happen in CreateAsync(). public AggregateQueryPipelineStage( IQueryPipelineStage source, SingleGroupAggregator singleGroupAggregator, - bool isValueQuery) - : base(source) + bool isValueQuery, + CancellationToken cancellationToken) + : base(source, cancellationToken) { this.singleGroupAggregator = singleGroupAggregator ?? throw new ArgumentNullException(nameof(singleGroupAggregator)); this.isValueQuery = isValueQuery; @@ -65,6 +64,7 @@ public static TryCatch MonadicCreate( IReadOnlyList orderedAliases, bool hasSelectValue, CosmosElement continuationToken, + CancellationToken cancellationToken, MonadicCreatePipelineStage monadicCreatePipelineStage) => executionEnvironment switch { ExecutionEnvironment.Client => ClientAggregateQueryPipelineStage.MonadicCreate( @@ -73,6 +73,7 @@ public static TryCatch MonadicCreate( orderedAliases, hasSelectValue, continuationToken, + cancellationToken, monadicCreatePipelineStage), ExecutionEnvironment.Compute => ComputeAggregateQueryPipelineStage.MonadicCreate( aggregates, @@ -80,6 +81,7 @@ public static TryCatch MonadicCreate( orderedAliases, hasSelectValue, continuationToken, + cancellationToken, monadicCreatePipelineStage), _ => throw new ArgumentException($"Unknown {nameof(ExecutionEnvironment)}: {executionEnvironment}."), }; diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CatchAllQueryPipelineStage.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CatchAllQueryPipelineStage.cs index 939ad9b708..5c0d11aae5 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CatchAllQueryPipelineStage.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CatchAllQueryPipelineStage.cs @@ -11,8 +11,8 @@ namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline internal sealed class CatchAllQueryPipelineStage : QueryPipelineStageBase { - public CatchAllQueryPipelineStage(IQueryPipelineStage inputStage) - : base(inputStage) + public CatchAllQueryPipelineStage(IQueryPipelineStage inputStage, CancellationToken cancellationToken) + : base(inputStage, cancellationToken) { } diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs index 14c205d373..0ffd01df35 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs @@ -56,11 +56,11 @@ public static IQueryPipelineStage Create( inputParameters, innerCancellationToken)); - LazyQueryPipelineStage lazyQueryPipelineStage = new LazyQueryPipelineStage(lazyTryCreateStage: lazyTryCreateStage); + LazyQueryPipelineStage lazyQueryPipelineStage = new LazyQueryPipelineStage(lazyTryCreateStage: lazyTryCreateStage, cancellationToken: default); return lazyQueryPipelineStage; }); - CatchAllQueryPipelineStage catchAllQueryPipelineStage = new CatchAllQueryPipelineStage(nameCacheStaleRetryQueryPipelineStage); + CatchAllQueryPipelineStage catchAllQueryPipelineStage = new CatchAllQueryPipelineStage(nameCacheStaleRetryQueryPipelineStage, cancellationToken: default); return catchAllQueryPipelineStage; } @@ -334,6 +334,7 @@ private static TryCatch TryCreatePassthroughQueryExecutionC targetRanges: targetRanges, pageSize: inputParameters.MaxItemCount, maxConcurrency: inputParameters.MaxConcurrency, + cancellationToken: default, continuationToken: inputParameters.InitialUserContinuationToken); } @@ -380,7 +381,8 @@ private static TryCatch TryCreateSpecializedDocumentQueryEx queryInfo: partitionedQueryExecutionInfo.QueryInfo, pageSize: (int)optimalPageSize, maxConcurrency: inputParameters.MaxConcurrency, - requestContinuationToken: inputParameters.InitialUserContinuationToken); + requestContinuationToken: inputParameters.InitialUserContinuationToken, + requestCancellationToken: default); } /// diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/OrderBy/OrderByCrossPartitionQueryPipelineStage.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/OrderBy/OrderByCrossPartitionQueryPipelineStage.cs index 2b984d52f2..4cf20480c8 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/OrderBy/OrderByCrossPartitionQueryPipelineStage.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/OrderBy/OrderByCrossPartitionQueryPipelineStage.cs @@ -8,6 +8,7 @@ namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.CrossPartition.OrderBy using System.Collections.Generic; using System.Linq; using System.Net; + using System.Runtime.CompilerServices; using System.Text; using System.Threading; using System.Threading.Tasks; @@ -18,7 +19,6 @@ namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.CrossPartition.OrderBy using Microsoft.Azure.Cosmos.Query.Core.Monads; using Microsoft.Azure.Cosmos.Query.Core.Pipeline.CrossPartition.Parallel; using Microsoft.Azure.Documents; - using static Microsoft.Azure.Cosmos.Query.Core.Pipeline.CrossPartition.PartitionMapper; /// /// CosmosOrderByItemQueryExecutionContext is a concrete implementation for CrossPartitionQueryExecutionContext. @@ -51,6 +51,7 @@ internal sealed class OrderByCrossPartitionQueryPipelineStage : IQueryPipelineSt private readonly Queue<(OrderByQueryPartitionRangePageAsyncEnumerator enumerator, OrderByContinuationToken token)> uninitializedEnumeratorsAndTokens; private readonly int pageSize; private readonly int maxConcurrency; + private CancellationToken cancellationToken; private QueryState state; @@ -69,7 +70,8 @@ private OrderByCrossPartitionQueryPipelineStage( int pageSize, int maxConcurrency, IEnumerable<(OrderByQueryPartitionRangePageAsyncEnumerator, OrderByContinuationToken)> uninitializedEnumeratorsAndTokens, - QueryState state) + QueryState state, + CancellationToken cancellationToken) { this.documentContainer = documentContainer ?? throw new ArgumentNullException(nameof(documentContainer)); this.sortOrders = sortOrders ?? throw new ArgumentNullException(nameof(sortOrders)); @@ -78,6 +80,7 @@ private OrderByCrossPartitionQueryPipelineStage( this.maxConcurrency = maxConcurrency < 0 ? throw new ArgumentOutOfRangeException($"{nameof(maxConcurrency)} must be a non negative number.") : maxConcurrency; this.uninitializedEnumeratorsAndTokens = new Queue<(OrderByQueryPartitionRangePageAsyncEnumerator, OrderByContinuationToken)>(uninitializedEnumeratorsAndTokens ?? throw new ArgumentNullException(nameof(uninitializedEnumeratorsAndTokens))); this.state = state ?? InitializingQueryState; + this.cancellationToken = cancellationToken; } public TryCatch Current { get; private set; } @@ -87,6 +90,8 @@ private OrderByCrossPartitionQueryPipelineStage( private async ValueTask MoveNextAsync_Initialize_FromBeginningAsync( OrderByQueryPartitionRangePageAsyncEnumerator uninitializedEnumerator) { + this.cancellationToken.ThrowIfCancellationRequested(); + if (uninitializedEnumerator == null) { throw new ArgumentNullException(nameof(uninitializedEnumerator)); @@ -168,6 +173,8 @@ private async ValueTask MoveNextAsync_Iniialize_FilterAsync( OrderByQueryPartitionRangePageAsyncEnumerator uninitializedEnumerator, OrderByContinuationToken token) { + this.cancellationToken.ThrowIfCancellationRequested(); + if (uninitializedEnumerator == null) { throw new ArgumentNullException(nameof(uninitializedEnumerator)); @@ -250,17 +257,22 @@ private async ValueTask MoveNextAsync_InitializeAsync_HandleSplitAsync( OrderByQueryPartitionRangePageAsyncEnumerator uninitializedEnumerator, OrderByContinuationToken token) { + this.cancellationToken.ThrowIfCancellationRequested(); + IEnumerable childRanges = await this.documentContainer.GetChildRangeAsync( uninitializedEnumerator.Range, - cancellationToken: default); + cancellationToken: this.cancellationToken); foreach (PartitionKeyRange childRange in childRanges) { + this.cancellationToken.ThrowIfCancellationRequested(); + OrderByQueryPartitionRangePageAsyncEnumerator childPaginator = new OrderByQueryPartitionRangePageAsyncEnumerator( this.documentContainer, uninitializedEnumerator.SqlQuerySpec, childRange, uninitializedEnumerator.PageSize, uninitializedEnumerator.Filter, + this.cancellationToken, state: uninitializedEnumerator.StartOfPageState); this.uninitializedEnumeratorsAndTokens.Enqueue((childPaginator, token)); } @@ -271,25 +283,23 @@ private async ValueTask MoveNextAsync_InitializeAsync_HandleSplitAsync( private async ValueTask MoveNextAsync_InitializeAsync() { + this.cancellationToken.ThrowIfCancellationRequested(); + await ParallelPrefetch.PrefetchInParallelAsync( - this.uninitializedEnumeratorsAndTokens.Select(value => value.Item1), - this.maxConcurrency); + this.uninitializedEnumeratorsAndTokens.Select(value => value.enumerator), + this.maxConcurrency, + this.cancellationToken); (OrderByQueryPartitionRangePageAsyncEnumerator uninitializedEnumerator, OrderByContinuationToken token) = this.uninitializedEnumeratorsAndTokens.Dequeue(); - bool movedNext; - if (token is null) - { - movedNext = await this.MoveNextAsync_Initialize_FromBeginningAsync(uninitializedEnumerator); - } - else - { - movedNext = await this.MoveNextAsync_Iniialize_FilterAsync(uninitializedEnumerator, token); - } - + bool movedNext = token is null + ? await this.MoveNextAsync_Initialize_FromBeginningAsync(uninitializedEnumerator) + : await this.MoveNextAsync_Iniialize_FilterAsync(uninitializedEnumerator, token); return movedNext; } private ValueTask MoveNextAsync_DrainPageAsync() { + this.cancellationToken.ThrowIfCancellationRequested(); + OrderByQueryPartitionRangePageAsyncEnumerator currentEnumerator = default; OrderByQueryResult orderByQueryResult = default; @@ -358,6 +368,8 @@ private ValueTask MoveNextAsync_DrainPageAsync() public ValueTask MoveNextAsync() { + this.cancellationToken.ThrowIfCancellationRequested(); + if (this.uninitializedEnumeratorsAndTokens.Count != 0) { return this.MoveNextAsync_InitializeAsync(); @@ -379,7 +391,8 @@ public static TryCatch MonadicCreate( IReadOnlyList orderByColumns, int pageSize, int maxConcurrency, - CosmosElement continuationToken) + CosmosElement continuationToken, + CancellationToken cancellationToken) { // TODO (brchon): For now we are not honoring non deterministic ORDER BY queries, since there is a bug in the continuation logic. // We can turn it back on once the bug is fixed. @@ -435,12 +448,13 @@ public static TryCatch MonadicCreate( range, pageSize, TrueFilter, + cancellationToken, state: default), (OrderByContinuationToken)null)) .ToList(); } else { - TryCatch> monadicGetOrderByContinuationTokenMapping = MonadicGetOrderByContinuationTokenMapping( + TryCatch> monadicGetOrderByContinuationTokenMapping = MonadicGetOrderByContinuationTokenMapping( targetRanges, continuationToken, orderByColumns.Count); @@ -449,7 +463,7 @@ public static TryCatch MonadicCreate( return TryCatch.FromException(monadicGetOrderByContinuationTokenMapping.Exception); } - PartitionMapping partitionMapping = monadicGetOrderByContinuationTokenMapping.Result; + PartitionMapper.PartitionMapping partitionMapping = monadicGetOrderByContinuationTokenMapping.Result; IReadOnlyList orderByItems = partitionMapping .TargetPartition .Values @@ -497,6 +511,7 @@ public static TryCatch MonadicCreate( range, pageSize, filter, + cancellationToken, state: token?.ParallelContinuationToken?.Token != null ? new QueryState(CosmosString.Create(token.ParallelContinuationToken.Token)) : null); enumeratorsAndTokens.Add((remoteEnumerator, token)); @@ -510,11 +525,12 @@ public static TryCatch MonadicCreate( pageSize, maxConcurrency, enumeratorsAndTokens, - continuationToken == null ? null : new QueryState(continuationToken)); + continuationToken == null ? null : new QueryState(continuationToken), + cancellationToken); return TryCatch.FromResult(stage); } - private static TryCatch> MonadicGetOrderByContinuationTokenMapping( + private static TryCatch> MonadicGetOrderByContinuationTokenMapping( IReadOnlyList partitionKeyRanges, CosmosElement continuationToken, int numOrderByItems) @@ -537,7 +553,7 @@ private static TryCatch> MonadicGetOr TryCatch> monadicExtractContinuationTokens = MonadicExtractOrderByTokens(continuationToken, numOrderByItems); if (monadicExtractContinuationTokens.Failed) { - return TryCatch>.FromException(monadicExtractContinuationTokens.Exception); + return TryCatch>.FromException(monadicExtractContinuationTokens.Exception); } return PartitionMapper.MonadicGetPartitionMapping( @@ -874,5 +890,19 @@ private static bool IsSplitException(Exception exception) && (cosmosException.StatusCode == HttpStatusCode.Gone) && (cosmosException.SubStatusCode == (int)Documents.SubStatusCodes.PartitionKeyRangeGone); } + + public void SetCancellationToken(CancellationToken cancellationToken) + { + this.cancellationToken = cancellationToken; + foreach (OrderByQueryPartitionRangePageAsyncEnumerator enumerator in this.enumerators) + { + enumerator.SetCancellationToken(cancellationToken); + } + + foreach ((OrderByQueryPartitionRangePageAsyncEnumerator, OrderByContinuationToken) enumeratorAndToken in this.uninitializedEnumeratorsAndTokens) + { + enumeratorAndToken.Item1.SetCancellationToken(cancellationToken); + } + } } } diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/OrderBy/OrderByQueryPartitionRangePageAsyncEnumerator.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/OrderBy/OrderByQueryPartitionRangePageAsyncEnumerator.cs index be824625b0..49c02e1858 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/OrderBy/OrderByQueryPartitionRangePageAsyncEnumerator.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/OrderBy/OrderByQueryPartitionRangePageAsyncEnumerator.cs @@ -23,8 +23,9 @@ public OrderByQueryPartitionRangePageAsyncEnumerator( PartitionKeyRange feedRange, int pageSize, string filter, + CancellationToken cancellationToken, QueryState state = default) - : base(feedRange, state) + : base(feedRange, cancellationToken, state) { this.StartOfPageState = state; this.innerEnumerator = new InnerEnumerator( @@ -33,9 +34,11 @@ public OrderByQueryPartitionRangePageAsyncEnumerator( feedRange, pageSize, filter, + cancellationToken, state); this.bufferedEnumerator = new BufferedPartitionRangePageAsyncEnumerator( - this.innerEnumerator); + this.innerEnumerator, + cancellationToken); } public SqlQuerySpec SqlQuerySpec => this.innerEnumerator.SqlQuerySpec; @@ -55,7 +58,7 @@ protected override async Task> GetNextPageAsync(Cance return this.bufferedEnumerator.Current; } - public ValueTask PrefetchAsync() => this.bufferedEnumerator.PrefetchAsync(); + public ValueTask PrefetchAsync(CancellationToken cancellationToken) => this.bufferedEnumerator.PrefetchAsync(cancellationToken); private sealed class InnerEnumerator : PartitionRangePageAsyncEnumerator { @@ -67,8 +70,9 @@ public InnerEnumerator( PartitionKeyRange feedRange, int pageSize, string filter, + CancellationToken cancellationToken, QueryState state = default) - : base(feedRange, state) + : base(feedRange, cancellationToken, state) { this.queryDataSource = queryDataSource ?? throw new ArgumentNullException(nameof(queryDataSource)); this.SqlQuerySpec = sqlQuerySpec ?? throw new ArgumentNullException(nameof(sqlQuerySpec)); 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 7ff518f320..ee81e0c35c 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 @@ -8,6 +8,7 @@ namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.CrossPartition.Parallel using System.Collections.Generic; using System.Collections.Immutable; using System.Linq; + using System.Threading; using System.Threading.Tasks; using Microsoft.Azure.Cosmos.CosmosElements; using Microsoft.Azure.Cosmos.Pagination; @@ -26,11 +27,14 @@ namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.CrossPartition.Parallel internal sealed class ParallelCrossPartitionQueryPipelineStage : IQueryPipelineStage { private readonly CrossPartitionRangePageAsyncEnumerator crossPartitionRangePageAsyncEnumerator; + private CancellationToken cancellationToken; private ParallelCrossPartitionQueryPipelineStage( - CrossPartitionRangePageAsyncEnumerator crossPartitionRangePageAsyncEnumerator) + CrossPartitionRangePageAsyncEnumerator crossPartitionRangePageAsyncEnumerator, + CancellationToken cancellationToken) { this.crossPartitionRangePageAsyncEnumerator = crossPartitionRangePageAsyncEnumerator ?? throw new ArgumentNullException(nameof(crossPartitionRangePageAsyncEnumerator)); + this.cancellationToken = cancellationToken; } public TryCatch Current { get; private set; } @@ -39,6 +43,8 @@ private ParallelCrossPartitionQueryPipelineStage( public async ValueTask MoveNextAsync() { + this.cancellationToken.ThrowIfCancellationRequested(); + if (!await this.crossPartitionRangePageAsyncEnumerator.MoveNextAsync()) { this.Current = default; @@ -68,6 +74,8 @@ public async ValueTask MoveNextAsync() List activeParallelContinuationTokens = new List(); for (int i = 0; i < rangesAndStates.Count; i++) { + this.cancellationToken.ThrowIfCancellationRequested(); + (PartitionKeyRange range, QueryState state) = rangesAndStates[i]; if ((i == 0) || (state != null)) { @@ -105,7 +113,8 @@ public static TryCatch MonadicCreate( IReadOnlyList targetRanges, int pageSize, int maxConcurrency, - CosmosElement continuationToken) + CosmosElement continuationToken, + CancellationToken cancellationToken) { if (targetRanges == null) { @@ -135,9 +144,10 @@ public static TryCatch MonadicCreate( ParallelCrossPartitionQueryPipelineStage.MakeCreateFunction(documentContainer, sqlQuerySpec, pageSize), Comparer.Singleton, maxConcurrency, - state: state); + state: state, + cancellationToken: cancellationToken); - ParallelCrossPartitionQueryPipelineStage stage = new ParallelCrossPartitionQueryPipelineStage(crossPartitionPageEnumerator); + ParallelCrossPartitionQueryPipelineStage stage = new ParallelCrossPartitionQueryPipelineStage(crossPartitionPageEnumerator, cancellationToken); return TryCatch.FromResult(stage); } @@ -220,8 +230,15 @@ private static CreatePartitionRangePageAsyncEnumerator Ma sqlQuerySpec, range, pageSize, + cancellationToken: default, state); + public void SetCancellationToken(CancellationToken cancellationToken) + { + this.cancellationToken = cancellationToken; + + } + private sealed class Comparer : IComparer> { public static readonly Comparer Singleton = new Comparer(); 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 d25770adab..bffa599fc5 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 @@ -23,8 +23,9 @@ public QueryPartitionRangePageAsyncEnumerator( SqlQuerySpec sqlQuerySpec, PartitionKeyRange feedRange, int pageSize, + CancellationToken cancellationToken, QueryState state = default) - : base(feedRange, state) + : base(feedRange, cancellationToken, state) { this.queryDataSource = queryDataSource ?? throw new ArgumentNullException(nameof(queryDataSource)); this.sqlQuerySpec = sqlQuerySpec ?? throw new ArgumentNullException(nameof(sqlQuerySpec)); diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Distinct/DistinctQueryPipelineStage.Client.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Distinct/DistinctQueryPipelineStage.Client.cs index a35cb1b9d2..39ba2b9584 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Distinct/DistinctQueryPipelineStage.Client.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Distinct/DistinctQueryPipelineStage.Client.cs @@ -29,8 +29,9 @@ private sealed class ClientDistinctQueryPipelineStage : DistinctQueryPipelineSta private ClientDistinctQueryPipelineStage( DistinctQueryType distinctQueryType, DistinctMap distinctMap, - IQueryPipelineStage source) - : base(distinctMap, source) + IQueryPipelineStage source, + CancellationToken cancellationToken) + : base(distinctMap, source, cancellationToken) { if ((distinctQueryType != DistinctQueryType.Unordered) && (distinctQueryType != DistinctQueryType.Ordered)) { @@ -42,6 +43,7 @@ private ClientDistinctQueryPipelineStage( public static TryCatch MonadicCreate( CosmosElement requestContinuation, + CancellationToken cancellationToken, MonadicCreatePipelineStage monadicCreatePipelineStage, DistinctQueryType distinctQueryType) { @@ -67,16 +69,9 @@ public static TryCatch MonadicCreate( distinctMapToken: null); } - CosmosElement distinctMapToken; - if (distinctContinuationToken.DistinctMapToken != null) - { - distinctMapToken = CosmosString.Create(distinctContinuationToken.DistinctMapToken); - } - else - { - distinctMapToken = null; - } - + CosmosElement distinctMapToken = distinctContinuationToken.DistinctMapToken != null + ? CosmosString.Create(distinctContinuationToken.DistinctMapToken) + : null; TryCatch tryCreateDistinctMap = DistinctMap.TryCreate( distinctQueryType, distinctMapToken); @@ -104,7 +99,7 @@ public static TryCatch MonadicCreate( sourceToken = null; } - TryCatch tryCreateSource = monadicCreatePipelineStage(sourceToken); + TryCatch tryCreateSource = monadicCreatePipelineStage(sourceToken, cancellationToken); if (!tryCreateSource.Succeeded) { return TryCatch.FromException(tryCreateSource.Exception); @@ -114,7 +109,8 @@ public static TryCatch MonadicCreate( new ClientDistinctQueryPipelineStage( distinctQueryType, tryCreateDistinctMap.Result, - tryCreateSource.Result)); + tryCreateSource.Result, + cancellationToken)); } protected override async Task> GetNextPageAsync(CancellationToken cancellationToken) @@ -133,6 +129,8 @@ protected override async Task> GetNextPageAsync(Cancellation List distinctResults = new List(); foreach (CosmosElement document in sourcePage.Documents) { + cancellationToken.ThrowIfCancellationRequested(); + if (this.distinctMap.Add(document, out UInt128 _)) { distinctResults.Add(document); diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Distinct/DistinctQueryPipelineStage.Compute.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Distinct/DistinctQueryPipelineStage.Compute.cs index 1da9ccc576..dc43cdbbb5 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Distinct/DistinctQueryPipelineStage.Compute.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Distinct/DistinctQueryPipelineStage.Compute.cs @@ -25,15 +25,16 @@ private sealed class ComputeDistinctQueryPipelineStage : DistinctQueryPipelineSt private static readonly string UseTryGetContinuationTokenMessage = $"Use TryGetContinuationToken"; private ComputeDistinctQueryPipelineStage( - DistinctQueryType distinctQueryType, DistinctMap distinctMap, - IQueryPipelineStage source) - : base(distinctMap, source) + IQueryPipelineStage source, + CancellationToken cancellationToken) + : base(distinctMap, source, cancellationToken) { } public static TryCatch MonadicCreate( CosmosElement requestContinuation, + CancellationToken cancellationToken, MonadicCreatePipelineStage monadicCreatePipelineStage, DistinctQueryType distinctQueryType) { @@ -65,7 +66,7 @@ public static TryCatch MonadicCreate( return TryCatch.FromException(tryCreateDistinctMap.Exception); } - TryCatch tryCreateSource = monadicCreatePipelineStage(distinctContinuationToken.SourceToken); + TryCatch tryCreateSource = monadicCreatePipelineStage(distinctContinuationToken.SourceToken, cancellationToken); if (!tryCreateSource.Succeeded) { return TryCatch.FromException(tryCreateSource.Exception); @@ -73,9 +74,9 @@ public static TryCatch MonadicCreate( return TryCatch.FromResult( new ComputeDistinctQueryPipelineStage( - distinctQueryType, tryCreateDistinctMap.Result, - tryCreateSource.Result)); + tryCreateSource.Result, + cancellationToken)); } protected override async Task> GetNextPageAsync(CancellationToken cancellationToken) diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Distinct/DistinctQueryPipelineStage.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Distinct/DistinctQueryPipelineStage.cs index 7f1b348324..088c64cc23 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Distinct/DistinctQueryPipelineStage.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Distinct/DistinctQueryPipelineStage.cs @@ -30,8 +30,9 @@ internal abstract partial class DistinctQueryPipelineStage : QueryPipelineStageB protected DistinctQueryPipelineStage( DistinctMap distinctMap, - IQueryPipelineStage source) - : base(source) + IQueryPipelineStage source, + CancellationToken cancellationToken) + : base(source, cancellationToken) { this.distinctMap = distinctMap ?? throw new ArgumentNullException(nameof(distinctMap)); } @@ -39,15 +40,18 @@ protected DistinctQueryPipelineStage( public static TryCatch MonadicCreate( ExecutionEnvironment executionEnvironment, CosmosElement requestContinuation, + CancellationToken cancellationToken, MonadicCreatePipelineStage monadicCreatePipelineStage, DistinctQueryType distinctQueryType) => executionEnvironment switch { ExecutionEnvironment.Client => ClientDistinctQueryPipelineStage.MonadicCreate( requestContinuation, + cancellationToken, monadicCreatePipelineStage, distinctQueryType), ExecutionEnvironment.Compute => ComputeDistinctQueryPipelineStage.MonadicCreate( requestContinuation, + cancellationToken, monadicCreatePipelineStage, distinctQueryType), _ => throw new ArgumentException($"Unknown {nameof(ExecutionEnvironment)}: {executionEnvironment}."), diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/EmptyQueryPipelineStage.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/EmptyQueryPipelineStage.cs index 036e91e011..e3a7aa79aa 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/EmptyQueryPipelineStage.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/EmptyQueryPipelineStage.cs @@ -4,6 +4,7 @@ namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline { + using System.Threading; using System.Threading.Tasks; using Microsoft.Azure.Cosmos.Query.Core.Monads; using Microsoft.Azure.Cosmos.Reactive; @@ -24,5 +25,10 @@ public EmptyQueryPipelineStage() public ValueTask DisposeAsync() => this.emptyAsyncEnumerator.DisposeAsync(); public ValueTask MoveNextAsync() => this.emptyAsyncEnumerator.MoveNextAsync(); + + public void SetCancellationToken(CancellationToken cancellationToken) + { + // No work to do since this enumerator is fully sync. + } } } diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/FaultedQueryPipelineStage.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/FaultedQueryPipelineStage.cs index 8761aea123..7dbc79e329 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/FaultedQueryPipelineStage.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/FaultedQueryPipelineStage.cs @@ -5,6 +5,7 @@ namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline { using System; + using System.Threading; using System.Threading.Tasks; using Microsoft.Azure.Cosmos.Query.Core.Monads; using Microsoft.Azure.Cosmos.Reactive; @@ -28,5 +29,10 @@ public FaultedQueryPipelineStage(Exception exception) public ValueTask DisposeAsync() => this.justAsyncEnumerator.DisposeAsync(); public ValueTask MoveNextAsync() => this.justAsyncEnumerator.MoveNextAsync(); + + public void SetCancellationToken(CancellationToken cancellationToken) + { + // No work to do with since this enumerator is fully sync. + } } } diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/GroupBy/GroupByQueryPipelineStage.Client.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/GroupBy/GroupByQueryPipelineStage.Client.cs index 37079609bd..b9a7416f23 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/GroupBy/GroupByQueryPipelineStage.Client.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/GroupBy/GroupByQueryPipelineStage.Client.cs @@ -20,15 +20,15 @@ private sealed class ClientGroupByQueryPipelineStage : GroupByQueryPipelineStage private ClientGroupByQueryPipelineStage( IQueryPipelineStage source, + CancellationToken cancellationToken, GroupingTable groupingTable) - : base( - source, - groupingTable) + : base(source, cancellationToken, groupingTable) { } public static TryCatch MonadicCreate( CosmosElement requestContinuation, + CancellationToken cancellationToken, MonadicCreatePipelineStage monadicCreatePipelineStage, IReadOnlyDictionary groupByAliasToAggregateType, IReadOnlyList orderedAliases, @@ -45,7 +45,7 @@ public static TryCatch MonadicCreate( return TryCatch.FromException(tryCreateGroupingTable.Exception); } - TryCatch tryCreateSource = monadicCreatePipelineStage(requestContinuation); + TryCatch tryCreateSource = monadicCreatePipelineStage(requestContinuation, cancellationToken); if (tryCreateSource.Failed) { return tryCreateSource; @@ -53,6 +53,7 @@ public static TryCatch MonadicCreate( IQueryPipelineStage stage = new ClientGroupByQueryPipelineStage( tryCreateSource.Result, + cancellationToken, tryCreateGroupingTable.Result); return TryCatch.FromResult(stage); @@ -71,6 +72,8 @@ protected override async Task> GetNextPageAsync(Cancellation while (await this.inputStage.MoveNextAsync()) { + cancellationToken.ThrowIfCancellationRequested(); + // Stage 1: // Drain the groupings fully from all continuation and all partitions TryCatch tryGetSourcePage = this.inputStage.Current; diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/GroupBy/GroupByQueryPipelineStage.Compute.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/GroupBy/GroupByQueryPipelineStage.Compute.cs index 81e63b88d7..bc277b0c53 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/GroupBy/GroupByQueryPipelineStage.Compute.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/GroupBy/GroupByQueryPipelineStage.Compute.cs @@ -26,15 +26,15 @@ private sealed class ComputeGroupByQueryPipelineStage : GroupByQueryPipelineStag private ComputeGroupByQueryPipelineStage( IQueryPipelineStage source, + CancellationToken cancellationToken, GroupingTable groupingTable) - : base( - source, - groupingTable) + : base(source, cancellationToken, groupingTable) { } public static TryCatch MonadicCreate( CosmosElement requestContinuation, + CancellationToken cancellationToken, MonadicCreatePipelineStage monadicCreatePipelineStage, IReadOnlyDictionary groupByAliasToAggregateType, IReadOnlyList orderedAliases, @@ -65,7 +65,7 @@ public static TryCatch MonadicCreate( } else { - tryCreateSource = monadicCreatePipelineStage(groupByContinuationToken.SourceContinuationToken); + tryCreateSource = monadicCreatePipelineStage(groupByContinuationToken.SourceContinuationToken, cancellationToken); } if (!tryCreateSource.Succeeded) @@ -87,6 +87,7 @@ public static TryCatch MonadicCreate( return TryCatch.FromResult( new ComputeGroupByQueryPipelineStage( tryCreateSource.Result, + cancellationToken, tryCreateGroupingTable.Result)); } diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/GroupBy/GroupByQueryPipelineStage.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/GroupBy/GroupByQueryPipelineStage.cs index 1fd23a2192..f50c1abbc6 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/GroupBy/GroupByQueryPipelineStage.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/GroupBy/GroupByQueryPipelineStage.cs @@ -51,8 +51,9 @@ internal abstract partial class GroupByQueryPipelineStage : QueryPipelineStageBa protected GroupByQueryPipelineStage( IQueryPipelineStage source, + CancellationToken cancellationToken, GroupingTable groupingTable) - : base(source) + : base(source, cancellationToken) { this.groupingTable = groupingTable ?? throw new ArgumentNullException(nameof(groupingTable)); } @@ -60,6 +61,7 @@ protected GroupByQueryPipelineStage( public static TryCatch MonadicCreate( ExecutionEnvironment executionEnvironment, CosmosElement continuationToken, + CancellationToken cancellationToken, MonadicCreatePipelineStage monadicCreatePipelineStage, IReadOnlyDictionary groupByAliasToAggregateType, IReadOnlyList orderedAliases, @@ -67,12 +69,14 @@ public static TryCatch MonadicCreate( { ExecutionEnvironment.Client => ClientGroupByQueryPipelineStage.MonadicCreate( continuationToken, + cancellationToken, monadicCreatePipelineStage, groupByAliasToAggregateType, orderedAliases, hasSelectValue), ExecutionEnvironment.Compute => ComputeGroupByQueryPipelineStage.MonadicCreate( continuationToken, + cancellationToken, monadicCreatePipelineStage, groupByAliasToAggregateType, orderedAliases, diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/IQueryPipelineStage.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/IQueryPipelineStage.cs index 49c5832f5c..fba3faec12 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/IQueryPipelineStage.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/IQueryPipelineStage.cs @@ -5,9 +5,11 @@ namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline { using System.Collections.Generic; + using System.Threading; using Microsoft.Azure.Cosmos.Query.Core.Monads; internal interface IQueryPipelineStage : IAsyncEnumerator> { + void SetCancellationToken(CancellationToken cancellationToken); } } diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/LazyQueryPipelineStage.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/LazyQueryPipelineStage.cs index 3757c21c6c..746d36c277 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/LazyQueryPipelineStage.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/LazyQueryPipelineStage.cs @@ -5,16 +5,19 @@ namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline { using System; + using System.Threading; using System.Threading.Tasks; using Microsoft.Azure.Cosmos.Query.Core.Monads; internal sealed class LazyQueryPipelineStage : IQueryPipelineStage { private readonly AsyncLazy> lazyTryCreateStage; + private CancellationToken cancellationToken; - public LazyQueryPipelineStage(AsyncLazy> lazyTryCreateStage) + public LazyQueryPipelineStage(AsyncLazy> lazyTryCreateStage, CancellationToken cancellationToken) { this.lazyTryCreateStage = lazyTryCreateStage ?? throw new ArgumentNullException(nameof(lazyTryCreateStage)); + this.cancellationToken = cancellationToken; } public TryCatch Current { get; private set; } @@ -35,7 +38,9 @@ public ValueTask DisposeAsync() public async ValueTask MoveNextAsync() { - TryCatch tryCreateStage = await this.lazyTryCreateStage.GetValueAsync(default); + this.cancellationToken.ThrowIfCancellationRequested(); + + TryCatch tryCreateStage = await this.lazyTryCreateStage.GetValueAsync(this.cancellationToken); if (tryCreateStage.Failed) { this.Current = TryCatch.FromException(tryCreateStage.Exception); @@ -51,5 +56,10 @@ public async ValueTask MoveNextAsync() this.Current = stage.Current; return true; } + + public void SetCancellationToken(CancellationToken cancellationToken) + { + this.cancellationToken = cancellationToken; + } } } diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/MonadicCreatePipelineStage.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/MonadicCreatePipelineStage.cs index 3f26975d14..9067eb471b 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/MonadicCreatePipelineStage.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/MonadicCreatePipelineStage.cs @@ -4,9 +4,9 @@ namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline { - using System.Threading.Tasks; + using System.Threading; using Microsoft.Azure.Cosmos.CosmosElements; using Microsoft.Azure.Cosmos.Query.Core.Monads; - internal delegate TryCatch MonadicCreatePipelineStage(CosmosElement continuationToken); + internal delegate TryCatch MonadicCreatePipelineStage(CosmosElement continuationToken, CancellationToken cancellationToken); } diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/NameCacheStaleRetryQueryPipelineStage.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/NameCacheStaleRetryQueryPipelineStage.cs index 20b30e7fc7..67eeb76a9f 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/NameCacheStaleRetryQueryPipelineStage.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/NameCacheStaleRetryQueryPipelineStage.cs @@ -62,5 +62,10 @@ await this.cosmosQueryContext.QueryClient.ForceRefreshCollectionCacheAsync( return true; } + + public void SetCancellationToken(CancellationToken cancellationToken) + { + // No work to do since this enumerator is fully sync. + } } } diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/PipelineFactory.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/PipelineFactory.cs index 11537c1a5c..741bda4680 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/PipelineFactory.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/PipelineFactory.cs @@ -7,6 +7,8 @@ namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline using System; using System.Collections.Generic; using System.Linq; + using System.Net.Cache; + using System.Threading; using Microsoft.Azure.Cosmos.CosmosElements; using Microsoft.Azure.Cosmos.Pagination; using Microsoft.Azure.Cosmos.Query.Core.Monads; @@ -30,12 +32,13 @@ public static TryCatch MonadicCreate( QueryInfo queryInfo, int pageSize, int maxConcurrency, - CosmosElement requestContinuationToken) + CosmosElement requestContinuationToken, + CancellationToken requestCancellationToken) { MonadicCreatePipelineStage monadicCreatePipelineStage; if (queryInfo.HasOrderBy) { - monadicCreatePipelineStage = (continuationToken) => OrderByCrossPartitionQueryPipelineStage.MonadicCreate( + monadicCreatePipelineStage = (continuationToken, cancellationToken) => OrderByCrossPartitionQueryPipelineStage.MonadicCreate( documentContainer: documentContainer, sqlQuerySpec: sqlQuerySpec, targetRanges: targetRanges, @@ -44,38 +47,42 @@ public static TryCatch MonadicCreate( .Zip(queryInfo.OrderBy, (expression, sortOrder) => new OrderByColumn(expression, sortOrder)).ToList(), pageSize: pageSize, maxConcurrency: maxConcurrency, - continuationToken: requestContinuationToken); + continuationToken: requestContinuationToken, + cancellationToken: cancellationToken); } else { - monadicCreatePipelineStage = (continuationToken) => ParallelCrossPartitionQueryPipelineStage.MonadicCreate( + monadicCreatePipelineStage = (continuationToken, cancellationToken) => ParallelCrossPartitionQueryPipelineStage.MonadicCreate( documentContainer: documentContainer, sqlQuerySpec: sqlQuerySpec, targetRanges: targetRanges, pageSize: pageSize, maxConcurrency: maxConcurrency, - continuationToken: continuationToken); + continuationToken: continuationToken, + cancellationToken: cancellationToken); } if (queryInfo.HasAggregates && !queryInfo.HasGroupBy) { MonadicCreatePipelineStage monadicCreateSourceStage = monadicCreatePipelineStage; - monadicCreatePipelineStage = (continuationToken) => AggregateQueryPipelineStage.MonadicCreate( + monadicCreatePipelineStage = (continuationToken, cancellationToken) => AggregateQueryPipelineStage.MonadicCreate( executionEnvironment, queryInfo.Aggregates, queryInfo.GroupByAliasToAggregateType, queryInfo.GroupByAliases, queryInfo.HasSelectValue, continuationToken, + cancellationToken, monadicCreateSourceStage); } if (queryInfo.HasDistinct) { MonadicCreatePipelineStage monadicCreateSourceStage = monadicCreatePipelineStage; - monadicCreatePipelineStage = (continuationToken) => DistinctQueryPipelineStage.MonadicCreate( + monadicCreatePipelineStage = (continuationToken, cancellationToken) => DistinctQueryPipelineStage.MonadicCreate( executionEnvironment, continuationToken, + cancellationToken, monadicCreateSourceStage, queryInfo.DistinctType); } @@ -83,9 +90,10 @@ public static TryCatch MonadicCreate( if (queryInfo.HasGroupBy) { MonadicCreatePipelineStage monadicCreateSourceStage = monadicCreatePipelineStage; - monadicCreatePipelineStage = (continuationToken) => GroupByQueryPipelineStage.MonadicCreate( + monadicCreatePipelineStage = (continuationToken, cancellationToken) => GroupByQueryPipelineStage.MonadicCreate( executionEnvironment, continuationToken, + cancellationToken, monadicCreateSourceStage, queryInfo.GroupByAliasToAggregateType, queryInfo.GroupByAliases, @@ -95,35 +103,38 @@ public static TryCatch MonadicCreate( if (queryInfo.HasOffset) { MonadicCreatePipelineStage monadicCreateSourceStage = monadicCreatePipelineStage; - monadicCreatePipelineStage = (continuationToken) => SkipQueryPipelineStage.MonadicCreate( + monadicCreatePipelineStage = (continuationToken, cancellationToken) => SkipQueryPipelineStage.MonadicCreate( executionEnvironment, queryInfo.Offset.Value, continuationToken, + cancellationToken, monadicCreateSourceStage); } if (queryInfo.HasLimit) { MonadicCreatePipelineStage monadicCreateSourceStage = monadicCreatePipelineStage; - monadicCreatePipelineStage = (continuationToken) => TakeQueryPipelineStage.MonadicCreateLimitStage( + monadicCreatePipelineStage = (continuationToken, cancellationToken) => TakeQueryPipelineStage.MonadicCreateLimitStage( executionEnvironment, queryInfo.Limit.Value, continuationToken, + cancellationToken, monadicCreateSourceStage); } if (queryInfo.HasTop) { MonadicCreatePipelineStage monadicCreateSourceStage = monadicCreatePipelineStage; - monadicCreatePipelineStage = (continuationToken) => TakeQueryPipelineStage.MonadicCreateTopStage( + monadicCreatePipelineStage = (continuationToken, cancellationToken) => TakeQueryPipelineStage.MonadicCreateTopStage( executionEnvironment, queryInfo.Top.Value, continuationToken, + cancellationToken, monadicCreateSourceStage); } - return monadicCreatePipelineStage(requestContinuationToken) - .Try(onSuccess: (stage) => new SkipEmptyPageQueryPipelineStage(stage)); + return monadicCreatePipelineStage(requestContinuationToken, requestCancellationToken) + .Try(onSuccess: (stage) => new SkipEmptyPageQueryPipelineStage(stage, requestCancellationToken)); } } } diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/QueryPipelineStageBase.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/QueryPipelineStageBase.cs index 1864ee7e99..78206754df 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/QueryPipelineStageBase.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/QueryPipelineStageBase.cs @@ -12,11 +12,13 @@ namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline internal abstract class QueryPipelineStageBase : IQueryPipelineStage { protected readonly IQueryPipelineStage inputStage; + protected CancellationToken cancellationToken; private bool hasStarted; - protected QueryPipelineStageBase(IQueryPipelineStage inputStage) + protected QueryPipelineStageBase(IQueryPipelineStage inputStage, CancellationToken cancellationToken) { this.inputStage = inputStage ?? throw new ArgumentNullException(nameof(inputStage)); + this.cancellationToken = cancellationToken; } public TryCatch Current { get; private set; } @@ -38,7 +40,7 @@ public async ValueTask MoveNextAsync() this.hasStarted = true; - this.Current = await this.GetNextPageAsync(default); + this.Current = await this.GetNextPageAsync(this.cancellationToken); if (this.Current.Succeeded) { this.State = this.Current.Result.State; @@ -46,5 +48,13 @@ public async ValueTask MoveNextAsync() return true; } + + public void SetCancellationToken(CancellationToken cancellationToken) + { + // Only here to support legacy query iterator and ExecuteNextAsync + // can be removed only we only expose IAsyncEnumerable in v4 sdk. + this.cancellationToken = cancellationToken; + this.inputStage.SetCancellationToken(cancellationToken); + } } } diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Skip/SkipQueryPipelineStage.Client.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Skip/SkipQueryPipelineStage.Client.cs index a8bf2b3676..5447d0df48 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Skip/SkipQueryPipelineStage.Client.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Skip/SkipQueryPipelineStage.Client.cs @@ -19,8 +19,11 @@ internal abstract partial class SkipQueryPipelineStage : QueryPipelineStageBase { private sealed class ClientSkipQueryPipelineStage : SkipQueryPipelineStage { - private ClientSkipQueryPipelineStage(IQueryPipelineStage source, long skipCount) - : base(source, skipCount) + private ClientSkipQueryPipelineStage( + IQueryPipelineStage source, + CancellationToken cancellationToken, + long skipCount) + : base(source, cancellationToken, skipCount) { // Work is done in base constructor. } @@ -28,6 +31,7 @@ private ClientSkipQueryPipelineStage(IQueryPipelineStage source, long skipCount) public static TryCatch MonadicCreate( int offsetCount, CosmosElement continuationToken, + CancellationToken cancellationToken, MonadicCreatePipelineStage monadicCreatePipelineStage) { if (monadicCreatePipelineStage == null) @@ -76,7 +80,7 @@ public static TryCatch MonadicCreate( sourceToken = null; } - TryCatch tryCreateSource = monadicCreatePipelineStage(sourceToken); + TryCatch tryCreateSource = monadicCreatePipelineStage(sourceToken, cancellationToken); if (tryCreateSource.Failed) { return tryCreateSource; @@ -84,6 +88,7 @@ public static TryCatch MonadicCreate( IQueryPipelineStage stage = new ClientSkipQueryPipelineStage( tryCreateSource.Result, + cancellationToken, offsetContinuationToken.Offset); return TryCatch.FromResult(stage); diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Skip/SkipQueryPipelineStage.Compute.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Skip/SkipQueryPipelineStage.Compute.cs index fbfa845766..0b54c0dab8 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Skip/SkipQueryPipelineStage.Compute.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Skip/SkipQueryPipelineStage.Compute.cs @@ -18,15 +18,16 @@ internal abstract partial class SkipQueryPipelineStage : QueryPipelineStageBase { private sealed class ComputeSkipQueryPipelineStage : SkipQueryPipelineStage { - private ComputeSkipQueryPipelineStage(IQueryPipelineStage source, long skipCount) - : base(source, skipCount) + private ComputeSkipQueryPipelineStage(IQueryPipelineStage source, CancellationToken cancellationToken, long skipCount) + : base(source, cancellationToken, skipCount) { // Work is done in base constructor. } public static TryCatch MonadicCreate( int offsetCount, - CosmosElement continuationToken, + CosmosElement continuationToken, + CancellationToken cancellationToken, MonadicCreatePipelineStage monadicCreatePipelineStage) { if (monadicCreatePipelineStage == null) @@ -58,7 +59,7 @@ public static TryCatch MonadicCreate( "offset count in continuation token can not be greater than the offsetcount in the query.")); } - TryCatch tryCreateSource = monadicCreatePipelineStage(offsetContinuationToken.SourceToken); + TryCatch tryCreateSource = monadicCreatePipelineStage(offsetContinuationToken.SourceToken, cancellationToken); if (tryCreateSource.Failed) { return tryCreateSource; @@ -66,6 +67,7 @@ public static TryCatch MonadicCreate( IQueryPipelineStage stage = new ComputeSkipQueryPipelineStage( tryCreateSource.Result, + cancellationToken, offsetContinuationToken.Offset); return TryCatch.FromResult(stage); diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Skip/SkipQueryPipelineStage.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Skip/SkipQueryPipelineStage.cs index adf84a9fb6..ef4d279ecf 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Skip/SkipQueryPipelineStage.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Skip/SkipQueryPipelineStage.cs @@ -19,8 +19,9 @@ internal abstract partial class SkipQueryPipelineStage : QueryPipelineStageBase protected SkipQueryPipelineStage( IQueryPipelineStage source, + CancellationToken cancellationToken, long skipCount) - : base(source) + : base(source, cancellationToken) { if (skipCount > int.MaxValue) { @@ -34,6 +35,7 @@ public static TryCatch MonadicCreate( ExecutionEnvironment executionEnvironment, int offsetCount, CosmosElement continuationToken, + CancellationToken cancellationToken, MonadicCreatePipelineStage monadicCreatePipelineStage) { TryCatch tryCreate = executionEnvironment switch @@ -41,10 +43,12 @@ public static TryCatch MonadicCreate( ExecutionEnvironment.Client => ClientSkipQueryPipelineStage.MonadicCreate( offsetCount, continuationToken, + cancellationToken, monadicCreatePipelineStage), ExecutionEnvironment.Compute => ComputeSkipQueryPipelineStage.MonadicCreate( offsetCount, continuationToken, + cancellationToken, monadicCreatePipelineStage), _ => throw new ArgumentException($"Unknown {nameof(ExecutionEnvironment)}: {executionEnvironment}"), }; diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/SkipEmptyPageQueryPipelineStage.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/SkipEmptyPageQueryPipelineStage.cs index 7714e9a061..a97aa2a82b 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/SkipEmptyPageQueryPipelineStage.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/SkipEmptyPageQueryPipelineStage.cs @@ -5,6 +5,7 @@ namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline { using System; + using System.Threading; using System.Threading.Tasks; using Microsoft.Azure.Cosmos.Query.Core.Monads; @@ -13,10 +14,12 @@ internal sealed class SkipEmptyPageQueryPipelineStage : IQueryPipelineStage private readonly IQueryPipelineStage inputStage; private double cumulativeRequestCharge; private long cumulativeResponseLengthInBytes; + private CancellationToken cancellationToken; - public SkipEmptyPageQueryPipelineStage(IQueryPipelineStage inputStage) + public SkipEmptyPageQueryPipelineStage(IQueryPipelineStage inputStage, CancellationToken cancellationToken) { this.inputStage = inputStage ?? throw new ArgumentNullException(nameof(inputStage)); + this.cancellationToken = cancellationToken; } public TryCatch Current { get; private set; } @@ -68,5 +71,10 @@ public async ValueTask MoveNextAsync() this.Current = TryCatch.FromResult(cumulativeQueryPage); return true; } + + public void SetCancellationToken(CancellationToken cancellationToken) + { + this.cancellationToken = cancellationToken; + } } } diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Take/TakeQueryPipelineStage.Client.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Take/TakeQueryPipelineStage.Client.cs index b760d14895..ea5c57c9f2 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Take/TakeQueryPipelineStage.Client.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Take/TakeQueryPipelineStage.Client.cs @@ -22,9 +22,10 @@ private sealed class ClientTakeQueryPipelineStage : TakeQueryPipelineStage private ClientTakeQueryPipelineStage( IQueryPipelineStage source, + CancellationToken cancellationToken, int takeCount, TakeEnum takeEnum) - : base(source, takeCount) + : base(source, cancellationToken, takeCount) { this.takeEnum = takeEnum; } @@ -32,6 +33,7 @@ private ClientTakeQueryPipelineStage( public static TryCatch MonadicCreateLimitStage( int limitCount, CosmosElement requestContinuationToken, + CancellationToken cancellationToken, MonadicCreatePipelineStage monadicCreatePipelineStage) { if (limitCount < 0) @@ -85,7 +87,7 @@ public static TryCatch MonadicCreateLimitStage( sourceToken = null; } - TryCatch tryCreateSource = monadicCreatePipelineStage(sourceToken); + TryCatch tryCreateSource = monadicCreatePipelineStage(sourceToken, cancellationToken); if (tryCreateSource.Failed) { return tryCreateSource; @@ -93,6 +95,7 @@ public static TryCatch MonadicCreateLimitStage( IQueryPipelineStage stage = new ClientTakeQueryPipelineStage( tryCreateSource.Result, + cancellationToken, limitContinuationToken.Limit, TakeEnum.Limit); @@ -102,6 +105,7 @@ public static TryCatch MonadicCreateLimitStage( public static TryCatch MonadicCreateTopStage( int topCount, CosmosElement requestContinuationToken, + CancellationToken cancellationToken, MonadicCreatePipelineStage monadicCreatePipelineStage) { if (topCount < 0) @@ -155,7 +159,7 @@ public static TryCatch MonadicCreateTopStage( sourceToken = null; } - TryCatch tryCreateSource = monadicCreatePipelineStage(sourceToken); + TryCatch tryCreateSource = monadicCreatePipelineStage(sourceToken, cancellationToken); if (tryCreateSource.Failed) { return tryCreateSource; @@ -163,6 +167,7 @@ public static TryCatch MonadicCreateTopStage( IQueryPipelineStage stage = new ClientTakeQueryPipelineStage( tryCreateSource.Result, + cancellationToken, topContinuationToken.Top, TakeEnum.Top); diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Take/TakeQueryPipelineStage.Compute.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Take/TakeQueryPipelineStage.Compute.cs index ca1057a815..d22e8e7a32 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Take/TakeQueryPipelineStage.Compute.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Take/TakeQueryPipelineStage.Compute.cs @@ -13,7 +13,6 @@ namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.Take using Microsoft.Azure.Cosmos.CosmosElements.Numbers; using Microsoft.Azure.Cosmos.Query.Core.Exceptions; using Microsoft.Azure.Cosmos.Query.Core.Monads; - using Newtonsoft.Json; internal abstract partial class TakeQueryPipelineStage : QueryPipelineStageBase { @@ -21,8 +20,9 @@ private sealed class ComputeTakeQueryPipelineStage : TakeQueryPipelineStage { private ComputeTakeQueryPipelineStage( IQueryPipelineStage source, + CancellationToken cancellationToken, int takeCount) - : base(source, takeCount) + : base(source, cancellationToken, takeCount) { // Work is done in the base class. } @@ -30,22 +30,27 @@ private ComputeTakeQueryPipelineStage( public static TryCatch MonadicCreateLimitStage( int takeCount, CosmosElement requestContinuationToken, + CancellationToken cancellationToken, MonadicCreatePipelineStage monadicCreatePipelineStage) => ComputeTakeQueryPipelineStage.MonadicCreate( takeCount, requestContinuationToken, + cancellationToken, monadicCreatePipelineStage); public static TryCatch MonadicCreateTopStage( int takeCount, CosmosElement requestContinuationToken, + CancellationToken cancellationToken, MonadicCreatePipelineStage monadicCreatePipelineStage) => ComputeTakeQueryPipelineStage.MonadicCreate( takeCount, requestContinuationToken, + cancellationToken, monadicCreatePipelineStage); private static TryCatch MonadicCreate( int takeCount, CosmosElement requestContinuationToken, + CancellationToken cancellationToken, MonadicCreatePipelineStage monadicCreatePipelineStage) { if (takeCount < 0) @@ -80,7 +85,7 @@ private static TryCatch MonadicCreate( $"{nameof(TakeContinuationToken.TakeCount)} in {nameof(TakeContinuationToken)}: {requestContinuationToken}: {takeContinuationToken.TakeCount} can not be greater than the limit count in the query: {takeCount}.")); } - TryCatch tryCreateSource = monadicCreatePipelineStage(takeContinuationToken.SourceToken); + TryCatch tryCreateSource = monadicCreatePipelineStage(takeContinuationToken.SourceToken, cancellationToken); if (tryCreateSource.Failed) { return tryCreateSource; @@ -88,6 +93,7 @@ private static TryCatch MonadicCreate( IQueryPipelineStage stage = new ComputeTakeQueryPipelineStage( tryCreateSource.Result, + cancellationToken, takeContinuationToken.TakeCount); return TryCatch.FromResult(stage); diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Take/TakeQueryPipelineStage.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Take/TakeQueryPipelineStage.cs index 2eea20bee8..107d510747 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Take/TakeQueryPipelineStage.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Take/TakeQueryPipelineStage.cs @@ -6,9 +6,7 @@ namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.Take { using System; using System.Threading; - using System.Threading.Tasks; using Microsoft.Azure.Cosmos.CosmosElements; - using Microsoft.Azure.Cosmos.Query.Core.ExecutionContext; using Microsoft.Azure.Cosmos.Query.Core.Monads; internal abstract partial class TakeQueryPipelineStage : QueryPipelineStageBase @@ -17,8 +15,9 @@ internal abstract partial class TakeQueryPipelineStage : QueryPipelineStageBase protected TakeQueryPipelineStage( IQueryPipelineStage source, + CancellationToken cancellationToken, int takeCount) - : base(source) + : base(source, cancellationToken) { this.takeCount = takeCount; } @@ -27,15 +26,18 @@ public static TryCatch MonadicCreateLimitStage( ExecutionEnvironment executionEnvironment, int limitCount, CosmosElement requestContinuationToken, + CancellationToken cancellationToken, MonadicCreatePipelineStage monadicCreatePipelineStage) => executionEnvironment switch { ExecutionEnvironment.Client => ClientTakeQueryPipelineStage.MonadicCreateLimitStage( limitCount, requestContinuationToken, + cancellationToken, monadicCreatePipelineStage), ExecutionEnvironment.Compute => ComputeTakeQueryPipelineStage.MonadicCreateLimitStage( limitCount, requestContinuationToken, + cancellationToken, monadicCreatePipelineStage), _ => throw new ArgumentOutOfRangeException($"Unknown {nameof(ExecutionEnvironment)}: {executionEnvironment}."), }; @@ -44,15 +46,18 @@ public static TryCatch MonadicCreateTopStage( ExecutionEnvironment executionEnvironment, int limitCount, CosmosElement requestContinuationToken, + CancellationToken cancellationToken, MonadicCreatePipelineStage monadicCreatePipelineStage) => executionEnvironment switch { ExecutionEnvironment.Client => ClientTakeQueryPipelineStage.MonadicCreateTopStage( limitCount, requestContinuationToken, + cancellationToken, monadicCreatePipelineStage), ExecutionEnvironment.Compute => ComputeTakeQueryPipelineStage.MonadicCreateTopStage( limitCount, requestContinuationToken, + cancellationToken, monadicCreatePipelineStage), _ => throw new ArgumentOutOfRangeException($"Unknown {nameof(ExecutionEnvironment)}: {executionEnvironment}."), }; diff --git a/Microsoft.Azure.Cosmos/src/Query/v3Query/QueryIterator.cs b/Microsoft.Azure.Cosmos/src/Query/v3Query/QueryIterator.cs index b97f02e718..9732b50a3a 100644 --- a/Microsoft.Azure.Cosmos/src/Query/v3Query/QueryIterator.cs +++ b/Microsoft.Azure.Cosmos/src/Query/v3Query/QueryIterator.cs @@ -151,6 +151,7 @@ public override async Task ReadNextAsync(CancellationToken canc try { // This catches exception thrown by the pipeline and converts it to QueryResponse + this.queryPipelineStage.SetCancellationToken(cancellationToken); await this.queryPipelineStage.MoveNextAsync(); tryGetQueryPage = this.queryPipelineStage.Current; } diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/BufferedPartitionRangeEnumeratorTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/BufferedPartitionRangeEnumeratorTests.cs index cd648d6ddc..c4aecfab55 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/BufferedPartitionRangeEnumeratorTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/BufferedPartitionRangeEnumeratorTests.cs @@ -105,6 +105,7 @@ public async Task TestSplitAsync() inMemoryCollection, partitionKeyRangeId: int.Parse(range.Id), pageSize: 10, + cancellationToken: default, state: state)); HashSet resourceIdentifiers = await this.DrainFullyAsync(enumerable); @@ -123,14 +124,16 @@ public async Task TestBufferPageAsync() new DocumentContainerPartitionRangeEnumerator( inMemoryCollection, partitionKeyRangeId: 0, - pageSize: 10)); + pageSize: 10, + cancellationToken: default), + cancellationToken: default); int count = 0; for (int i = 0; i < 10; i++) { // This call is idempotent; - await enumerator.PrefetchAsync(); + await enumerator.PrefetchAsync(default); } Random random = new Random(); @@ -142,7 +145,7 @@ public async Task TestBufferPageAsync() for (int i = 0; i < 10; i++) { // This call is idempotent; - await enumerator.PrefetchAsync(); + await enumerator.PrefetchAsync(default); } } } @@ -163,11 +166,13 @@ public async Task TestMoveNextAndBufferPageAsync() new DocumentContainerPartitionRangeEnumerator( inMemoryCollection, partitionKeyRangeId: 0, - pageSize: 10)); + pageSize: 10, + cancellationToken: default), + cancellationToken: default); if ((random.Next() % 2) == 0) { - await enumerator.PrefetchAsync(); + await enumerator.PrefetchAsync(default); } int count = 0; @@ -177,7 +182,7 @@ public async Task TestMoveNextAndBufferPageAsync() if ((random.Next() % 2) == 0) { - await enumerator.PrefetchAsync(); + await enumerator.PrefetchAsync(default); } } @@ -200,7 +205,9 @@ public override IAsyncEnumerable> CreateEnumerab documentContainer, partitionKeyRangeId: int.Parse(range.Id), pageSize: 10, - state: state))); + cancellationToken: default, + state: state), + cancellationToken: default)); public override IAsyncEnumerator> CreateEnumerator( IDocumentContainer inMemoryCollection, @@ -209,13 +216,15 @@ public override IAsyncEnumerator> CreateEnumerat inMemoryCollection, partitionKeyRangeId: 0, pageSize: 10, - state: state)); + cancellationToken: default, + state: state), + cancellationToken: default); private async Task BufferMoreInBackground(BufferedPartitionRangePageAsyncEnumerator enumerator) { while (true) { - await enumerator.PrefetchAsync(); + await enumerator.PrefetchAsync(default); await Task.Delay(10); } } diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/CrossPartitionPartitionRangeEnumeratorTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/CrossPartitionPartitionRangeEnumeratorTests.cs index cb9f642786..ec17d8968e 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/CrossPartitionPartitionRangeEnumeratorTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/CrossPartitionPartitionRangeEnumeratorTests.cs @@ -141,6 +141,7 @@ PartitionRangePageAsyncEnumerator inMemoryCollection, partitionKeyRangeId: int.Parse(range.Id), pageSize: 10, + cancellationToken: default, state: state); return new CrossPartitionRangePageAsyncEnumerable( @@ -161,6 +162,7 @@ PartitionRangePageAsyncEnumerator inMemoryCollection, partitionKeyRangeId: int.Parse(range.Id), pageSize: 10, + cancellationToken: default, state: state); CrossPartitionRangePageAsyncEnumerator enumerator = new CrossPartitionRangePageAsyncEnumerator( @@ -168,6 +170,7 @@ PartitionRangePageAsyncEnumerator createPartitionRangeEnumerator: createEnumerator, comparer: PartitionRangePageAsyncEnumeratorComparer.Singleton, maxConcurrency: 10, + cancellationToken: default, state: state); return enumerator; diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/DocumentContainerPartitionRangeEnumerator.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/DocumentContainerPartitionRangeEnumerator.cs index 41910827e0..706d62f991 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/DocumentContainerPartitionRangeEnumerator.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/DocumentContainerPartitionRangeEnumerator.cs @@ -5,7 +5,6 @@ namespace Microsoft.Azure.Cosmos.Tests.Pagination { using System; - using System.Collections.Generic; using System.Threading; using System.Threading.Tasks; using Microsoft.Azure.Cosmos.Pagination; @@ -22,6 +21,7 @@ public DocumentContainerPartitionRangeEnumerator( IDocumentContainer documentContainer, int partitionKeyRangeId, int pageSize, + CancellationToken cancellationToken, DocumentContainerState state = null) : base( new PartitionKeyRange() @@ -30,6 +30,7 @@ public DocumentContainerPartitionRangeEnumerator( MinInclusive = partitionKeyRangeId.ToString(), MaxExclusive = partitionKeyRangeId.ToString() }, + cancellationToken, state ?? new DocumentContainerState(resourceIdentifier: ResourceId.Empty)) { this.documentContainer = documentContainer ?? throw new ArgumentNullException(nameof(documentContainer)); @@ -43,6 +44,6 @@ protected override Task> GetNextPageAsync(Cancel partitionKeyRangeId: this.partitionKeyRangeId, resourceIdentifer: this.State.ResourceIdentifer, pageSize: this.pageSize, - cancellationToken: default); + cancellationToken: cancellationToken); } } diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/SinglePartitionPartitionRangeEnumeratorTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/SinglePartitionPartitionRangeEnumeratorTests.cs index 3a63276275..597ee9c75c 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/SinglePartitionPartitionRangeEnumeratorTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/SinglePartitionPartitionRangeEnumeratorTests.cs @@ -74,7 +74,8 @@ public async Task TestSplitAsync() DocumentContainerPartitionRangeEnumerator enumerator = new DocumentContainerPartitionRangeEnumerator( inMemoryCollection, partitionKeyRangeId: 0, - pageSize: 10); + pageSize: 10, + cancellationToken: default); (HashSet parentIdentifiers, DocumentContainerState state) = await this.PartialDrainAsync(enumerator, numIterations: 3); @@ -96,7 +97,8 @@ public async Task TestSplitAsync() inMemoryCollection, partitionKeyRangeId: int.Parse(range.Id), pageSize: 10, - state: state)); + state: state, + cancellationToken: default)); HashSet resourceIdentifiers = await this.DrainFullyAsync(enumerable); childIdentifiers.UnionWith(resourceIdentifiers); @@ -119,7 +121,8 @@ public override IAsyncEnumerable> CreateEnumerab documentContainer, partitionKeyRangeId: int.Parse(range.Id), pageSize: 10, - state: state)); + state: state, + cancellationToken: default)); public override IAsyncEnumerator> CreateEnumerator( IDocumentContainer inMemoryCollection, @@ -127,7 +130,8 @@ public override IAsyncEnumerator> CreateEnumerat inMemoryCollection, partitionKeyRangeId: 0, pageSize: 10, - state: state); + state: state, + cancellationToken: default); } } } diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/AggregateQueryPipelineStageTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/AggregateQueryPipelineStageTests.cs index fbd9d6283a..093f56669a 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/AggregateQueryPipelineStageTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/AggregateQueryPipelineStageTests.cs @@ -109,7 +109,8 @@ private static async Task> CreateAndDrain( orderedAliases: orderedAliases, hasSelectValue: hasSelectValue, continuationToken: continuationToken, - monadicCreatePipelineStage: (CosmosElement continuationToken) => TryCatch.FromResult(source)); + cancellationToken: default, + monadicCreatePipelineStage: (CosmosElement continuationToken, CancellationToken cancellationToken) => TryCatch.FromResult(source)); Assert.IsTrue(tryCreateAggregateQueryPipelineStage.Succeeded); IQueryPipelineStage aggregateQueryPipelineStage = tryCreateAggregateQueryPipelineStage.Result; diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/DistinctQueryPipelineStageTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/DistinctQueryPipelineStageTests.cs index 974669e1c0..530b91552d 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/DistinctQueryPipelineStageTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/DistinctQueryPipelineStageTests.cs @@ -3,6 +3,7 @@ using System; using System.Collections.Generic; using System.Linq; + using System.Threading; using System.Threading.Tasks; using Microsoft.Azure.Cosmos.CosmosElements; using Microsoft.Azure.Cosmos.Query.Core.ExecutionContext; @@ -46,7 +47,8 @@ private static async Task> CreateAndDrainAsync( executionEnvironment: executionEnvironment, requestContinuation: continuationToken, distinctQueryType: distinctQueryType, - monadicCreatePipelineStage: (CosmosElement continuationToken) => TryCatch.FromResult(source)); + cancellationToken: default, + monadicCreatePipelineStage: (CosmosElement continuationToken, CancellationToken cancellationToken) => TryCatch.FromResult(source)); Assert.IsTrue(tryCreateDistinctQueryPipelineStage.Succeeded); IQueryPipelineStage distinctQueryPipelineStage = tryCreateDistinctQueryPipelineStage.Result; 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 4b4c8ce26a..14672628cd 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 @@ -27,11 +27,12 @@ public void TestCreate() ExecutionEnvironment.Compute, documentContainer: mockDocumentContainer.Object, sqlQuerySpec: new SqlQuerySpec("SELECT * FROM c"), - targetRanges: new List(), + targetRanges: new List() { new PartitionKeyRange() }, queryInfo: new QueryInfo() { }, pageSize: 10, maxConcurrency: 10, - requestContinuationToken: default); + requestCancellationToken: default, + requestContinuationToken: default); ; Assert.IsTrue(monadicCreatePipeline.Succeeded); } } 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 839978e754..9f328767c9 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 @@ -179,6 +179,7 @@ private static IQueryPipelineStage CreatePipeline(IDocumentContainer documentCon GetQueryPlan(query), pageSize: 10, maxConcurrency: 10, + requestCancellationToken: default, requestContinuationToken: state); tryCreatePipeline.ThrowIfFailed(); diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/GroupByQueryPipelineStageTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/GroupByQueryPipelineStageTests.cs index 63535acea0..2b5ab5902c 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/GroupByQueryPipelineStageTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/GroupByQueryPipelineStageTests.cs @@ -6,6 +6,7 @@ namespace Microsoft.Azure.Cosmos.Tests.Query.Pipeline { using System; using System.Collections.Generic; + using System.Threading; using System.Threading.Tasks; using Microsoft.Azure.Cosmos.CosmosElements; using Microsoft.Azure.Cosmos.Query.Core.ExecutionContext; @@ -55,7 +56,8 @@ private static async Task> CreateAndDrainAsync( TryCatch tryCreateGroupByStage = GroupByQueryPipelineStage.MonadicCreate( executionEnvironment: executionEnvironment, continuationToken: continuationToken, - monadicCreatePipelineStage: (CosmosElement continuationToken) => TryCatch.FromResult(source), + cancellationToken: default, + monadicCreatePipelineStage: (CosmosElement continuationToken, CancellationToken cancellationToken) => TryCatch.FromResult(source), groupByAliasToAggregateType: groupByAliasToAggregateType, orderedAliases: orderedAliases, hasSelectValue: hasSelectValue); diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/MockQueryPipelineStage.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/MockQueryPipelineStage.cs index faef51ef64..2b1b8cc87b 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/MockQueryPipelineStage.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/MockQueryPipelineStage.cs @@ -18,7 +18,7 @@ internal sealed class MockQueryPipelineStage : QueryPipelineStageBase private int pageIndex; public MockQueryPipelineStage(IReadOnlyList> pages) - : base(EmptyQueryPipelineStage.Singleton) + : base(EmptyQueryPipelineStage.Singleton, cancellationToken: default) { this.pages = pages ?? throw new ArgumentNullException(nameof(pages)); } diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/OrderByCrossPartitionQueryPipelineStageTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/OrderByCrossPartitionQueryPipelineStageTests.cs index f8db0e1f73..25e71a68c5 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/OrderByCrossPartitionQueryPipelineStageTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/OrderByCrossPartitionQueryPipelineStageTests.cs @@ -40,6 +40,7 @@ public void MonadicCreate_NullContinuationToken() }, pageSize: 10, maxConcurrency: 10, + cancellationToken: default, continuationToken: null); Assert.IsTrue(monadicCreate.Succeeded); } @@ -59,6 +60,7 @@ public void MonadicCreate_NonCosmosArrayContinuationToken() }, pageSize: 10, maxConcurrency: 10, + cancellationToken: default, continuationToken: CosmosObject.Create(new Dictionary())); Assert.IsTrue(monadicCreate.Failed); Assert.IsTrue(monadicCreate.InnerMostException is MalformedContinuationTokenException); @@ -79,6 +81,7 @@ public void MonadicCreate_EmptyArrayContinuationToken() }, pageSize: 10, maxConcurrency: 10, + cancellationToken: default, continuationToken: CosmosArray.Create(new List())); Assert.IsTrue(monadicCreate.Failed); Assert.IsTrue(monadicCreate.InnerMostException is MalformedContinuationTokenException); @@ -99,6 +102,7 @@ public void MonadicCreate_NonParallelContinuationToken() }, pageSize: 10, maxConcurrency: 10, + cancellationToken: default, continuationToken: CosmosArray.Create(new List() { CosmosString.Create("asdf") })); Assert.IsTrue(monadicCreate.Failed); Assert.IsTrue(monadicCreate.InnerMostException is MalformedContinuationTokenException); @@ -130,6 +134,7 @@ public void MonadicCreate_SingleOrderByContinuationToken() }, pageSize: 10, maxConcurrency: 10, + cancellationToken: default, continuationToken: CosmosString.Create( CosmosArray.Create( new List() @@ -144,10 +149,6 @@ public void MonadicCreate_MultipleOrderByContinuationToken() { Mock mockDocumentContainer = new Mock(); - ParallelContinuationToken token = new ParallelContinuationToken( - token: "asdf", - range: new Documents.Routing.Range("A", "B", true, false)); - ParallelContinuationToken parallelContinuationToken1 = new ParallelContinuationToken( token: "asdf", range: new Documents.Routing.Range("A", "B", true, false)); @@ -184,6 +185,7 @@ public void MonadicCreate_MultipleOrderByContinuationToken() }, pageSize: 10, maxConcurrency: 10, + cancellationToken: default, continuationToken: CosmosString.Create( CosmosArray.Create( new List() @@ -214,6 +216,7 @@ FROM c }, pageSize: 10, maxConcurrency: 10, + cancellationToken: default, continuationToken: null); Assert.IsTrue(monadicCreate.Succeeded); IQueryPipelineStage queryPipelineStage = monadicCreate.Result; @@ -260,6 +263,7 @@ FROM c }, pageSize: 10, maxConcurrency: 10, + cancellationToken: default, continuationToken: queryState?.Value); Assert.IsTrue(monadicCreate.Succeeded); IQueryPipelineStage queryPipelineStage = monadicCreate.Result; @@ -306,6 +310,7 @@ FROM c }, pageSize: 10, maxConcurrency: 10, + cancellationToken: default, continuationToken: null); Assert.IsTrue(monadicCreate.Succeeded); IQueryPipelineStage queryPipelineStage = monadicCreate.Result; @@ -362,6 +367,7 @@ FROM c }, pageSize: 10, maxConcurrency: 10, + cancellationToken: default, continuationToken: queryState?.Value); Assert.IsTrue(monadicCreate.Succeeded); IQueryPipelineStage queryPipelineStage = monadicCreate.Result; 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 ac02d0379c..fee189643a 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 @@ -32,6 +32,7 @@ public void MonadicCreate_NullContinuationToken() targetRanges: new List() { new PartitionKeyRange() }, pageSize: 10, maxConcurrency: 10, + cancellationToken: default, continuationToken: null); Assert.IsTrue(monadicCreate.Succeeded); } @@ -47,6 +48,7 @@ public void MonadicCreate_NonCosmosArrayContinuationToken() targetRanges: new List() { new PartitionKeyRange() }, pageSize: 10, maxConcurrency: 10, + cancellationToken: default, continuationToken: CosmosObject.Create(new Dictionary())); Assert.IsTrue(monadicCreate.Failed); Assert.IsTrue(monadicCreate.InnerMostException is MalformedContinuationTokenException); @@ -63,6 +65,7 @@ public void MonadicCreate_EmptyArrayContinuationToken() targetRanges: new List() { new PartitionKeyRange() }, pageSize: 10, maxConcurrency: 10, + cancellationToken: default, continuationToken: CosmosArray.Create(new List())); Assert.IsTrue(monadicCreate.Failed); Assert.IsTrue(monadicCreate.InnerMostException is MalformedContinuationTokenException); @@ -79,6 +82,7 @@ public void MonadicCreate_NonParallelContinuationToken() targetRanges: new List() { new PartitionKeyRange() }, pageSize: 10, maxConcurrency: 10, + cancellationToken: default, continuationToken: CosmosArray.Create(new List() { CosmosString.Create("asdf") })); Assert.IsTrue(monadicCreate.Failed); Assert.IsTrue(monadicCreate.InnerMostException is MalformedContinuationTokenException); @@ -99,6 +103,7 @@ public void MonadicCreate_SingleParallelContinuationToken() targetRanges: new List() { new PartitionKeyRange() { Id = "0", MinInclusive = "A", MaxExclusive = "B" } }, pageSize: 10, maxConcurrency: 10, + cancellationToken: default, continuationToken: CosmosArray.Create(new List() { ParallelContinuationToken.ToCosmosElement(token) })); Assert.IsTrue(monadicCreate.Succeeded); } @@ -126,6 +131,7 @@ public void MonadicCreate_MultipleParallelContinuationToken() }, pageSize: 10, maxConcurrency: 10, + cancellationToken: default, continuationToken: CosmosArray.Create( new List() { @@ -147,6 +153,7 @@ public async Task TestDrainFully_StartFromBeginingAsync() targetRanges: await documentContainer.GetFeedRangesAsync(cancellationToken: default), pageSize: 10, maxConcurrency: 10, + cancellationToken: default, continuationToken: default); Assert.IsTrue(monadicCreate.Succeeded); IQueryPipelineStage queryPipelineStage = monadicCreate.Result; @@ -181,6 +188,7 @@ public async Task TestDrainFully_WithStateResume() targetRanges: await documentContainer.GetFeedRangesAsync(cancellationToken: default), pageSize: 10, maxConcurrency: 10, + cancellationToken: default, continuationToken: queryState?.Value); if (monadicCreate.Failed) { diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/SkipQueryPipelineStageTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/SkipQueryPipelineStageTests.cs index 80820a5de1..2a341690da 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/SkipQueryPipelineStageTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/SkipQueryPipelineStageTests.cs @@ -3,6 +3,7 @@ using System; using System.Collections.Generic; using System.Linq; + using System.Threading; using System.Threading.Tasks; using Microsoft.Azure.Cosmos.CosmosElements; using Microsoft.Azure.Cosmos.Query.Core.ExecutionContext; @@ -49,7 +50,8 @@ private static async Task> CreateAndDrainAsync( executionEnvironment: executionEnvironment, offsetCount: offsetCount, continuationToken: continuationToken, - monadicCreatePipelineStage: (CosmosElement continuationToken) => TryCatch.FromResult(source)); + cancellationToken: default, + monadicCreatePipelineStage: (CosmosElement continuationToken, CancellationToken token) => TryCatch.FromResult(source)); Assert.IsTrue(tryCreateSkipQueryPipelineStage.Succeeded); IQueryPipelineStage aggregateQueryPipelineStage = tryCreateSkipQueryPipelineStage.Result; diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/TakeQueryPipelineStageTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/TakeQueryPipelineStageTests.cs index 9fa35f087c..41adc07489 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/TakeQueryPipelineStageTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/TakeQueryPipelineStageTests.cs @@ -3,9 +3,9 @@ using System; using System.Collections.Generic; using System.Linq; + using System.Threading; using System.Threading.Tasks; using Microsoft.Azure.Cosmos.CosmosElements; - 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.Take; @@ -49,7 +49,8 @@ private static async Task> CreateAndDrainAsync( executionEnvironment: executionEnvironment, limitCount: takeCount, requestContinuationToken: continuationToken, - monadicCreatePipelineStage: (CosmosElement continuationToken) => TryCatch.FromResult(source)); + cancellationToken: default, + monadicCreatePipelineStage: (CosmosElement continuationToken, CancellationToken cancellationToken) => TryCatch.FromResult(source)); Assert.IsTrue(tryCreateSkipQueryPipelineStage.Succeeded); IQueryPipelineStage takeQueryPipelineStage = tryCreateSkipQueryPipelineStage.Result; 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 8411d7434b..c4f754a8ee 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 @@ -96,7 +96,8 @@ public async Task TestSplitAsync() sqlQuerySpec: new Cosmos.Query.Core.SqlQuerySpec("SELECT * FROM c"), feedRange: range, pageSize: 10, - state: state)); + state: state, + cancellationToken: default)); HashSet resourceIdentifiers = await this.DrainFullyAsync(enumerable); childIdentifiers.UnionWith(resourceIdentifiers); @@ -135,7 +136,8 @@ public override IAsyncEnumerable> CreateEnumerable( sqlQuerySpec: new Cosmos.Query.Core.SqlQuerySpec("SELECT * FROM c"), feedRange: range, pageSize: 10, - state: state)); + state: state, + cancellationToken: default)); } public override IAsyncEnumerator> CreateEnumerator( @@ -149,7 +151,8 @@ public override IAsyncEnumerator> CreateEnumerator( sqlQuerySpec: new Cosmos.Query.Core.SqlQuerySpec("SELECT * FROM c"), feedRange: ranges[0], pageSize: 10, - state: state); + state: state, + cancellationToken: default); } } } From 14b2f011cd1c21e6eb2f5c063f15d6e1ba2822dc Mon Sep 17 00:00:00 2001 From: Brandon Chong Date: Mon, 28 Sep 2020 19:08:36 -0700 Subject: [PATCH 72/85] fixed random integration bugs --- .../AggregateQueryPipelineStage.Client.cs | 20 ++++-- .../AggregateQueryPipelineStage.Compute.cs | 19 ++++-- .../Aggregate/AggregateQueryPipelineStage.cs | 2 + .../Pipeline/CatchAllQueryPipelineStage.cs | 19 ++++-- ...OrderByCrossPartitionQueryPipelineStage.cs | 24 ++----- .../DistinctQueryPipelineStage.Client.cs | 27 +++++--- .../DistinctQueryPipelineStage.Compute.cs | 15 +++-- .../GroupByQueryPipelineStage.Client.cs | 22 +++++-- .../GroupByQueryPipelineStage.Compute.cs | 17 +++-- .../GroupBy/GroupByQueryPipelineStage.cs | 1 + .../Core/Pipeline/LazyQueryPipelineStage.cs | 1 + .../Query/Core/Pipeline/PipelineFactory.cs | 2 +- .../Core/Pipeline/QueryPipelineStageBase.cs | 27 +------- .../Skip/SkipQueryPipelineStage.Client.cs | 22 ++++--- .../Skip/SkipQueryPipelineStage.Compute.cs | 19 ++++-- .../SkipEmptyPageQueryPipelineStage.cs | 22 +++++++ .../Take/TakeQueryPipelineStage.Client.cs | 17 +++-- .../Take/TakeQueryPipelineStage.Compute.cs | 17 +++-- .../Query/Pipeline/MockQueryPipelineStage.cs | 11 +++- ...ByCrossPartitionQueryPipelineStageTests.cs | 66 +++++++++++++++---- 20 files changed, 247 insertions(+), 123 deletions(-) diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Aggregate/AggregateQueryPipelineStage.Client.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Aggregate/AggregateQueryPipelineStage.Client.cs index 7f918206f8..51d66e45eb 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Aggregate/AggregateQueryPipelineStage.Client.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Aggregate/AggregateQueryPipelineStage.Client.cs @@ -11,7 +11,6 @@ namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.Aggregate using Microsoft.Azure.Cosmos.CosmosElements; using Microsoft.Azure.Cosmos.Query.Core.Monads; using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Aggregate.Aggregators; - using static Microsoft.Azure.Cosmos.Query.Core.Pipeline.PipelineFactory; internal abstract partial class AggregateQueryPipelineStage : QueryPipelineStageBase { @@ -67,9 +66,14 @@ public static TryCatch MonadicCreate( return TryCatch.FromResult(stage); } - protected override async Task> GetNextPageAsync(CancellationToken cancellationToken) + public override async ValueTask MoveNextAsync() { - cancellationToken.ThrowIfCancellationRequested(); + this.cancellationToken.ThrowIfCancellationRequested(); + + if (this.returnedFinalPage) + { + return false; + } // Note-2016-10-25-felixfan: Given what we support now, we should expect to return only 1 document. // Note-2019-07-11-brchon: We can return empty pages until all the documents are drained, @@ -80,10 +84,10 @@ protected override async Task> GetNextPageAsync(Cancellation while (await this.inputStage.MoveNextAsync()) { TryCatch tryGetPageFromSource = this.inputStage.Current; - if (tryGetPageFromSource.Failed) { - return tryGetPageFromSource; + this.Current = tryGetPageFromSource; + return true; } QueryPage sourcePage = tryGetPageFromSource.Result; @@ -93,7 +97,7 @@ protected override async Task> GetNextPageAsync(Cancellation foreach (CosmosElement element in sourcePage.Documents) { - cancellationToken.ThrowIfCancellationRequested(); + this.cancellationToken.ThrowIfCancellationRequested(); RewrittenAggregateProjections rewrittenAggregateProjections = new RewrittenAggregateProjections( this.isValueQuery, @@ -118,7 +122,9 @@ protected override async Task> GetNextPageAsync(Cancellation disallowContinuationTokenMessage: default, state: default); - return TryCatch.FromResult(queryPage); + this.Current = TryCatch.FromResult(queryPage); + this.returnedFinalPage = true; + return true; } } } diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Aggregate/AggregateQueryPipelineStage.Compute.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Aggregate/AggregateQueryPipelineStage.Compute.cs index 8e9a5fed29..7e1348b1fb 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Aggregate/AggregateQueryPipelineStage.Compute.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Aggregate/AggregateQueryPipelineStage.Compute.cs @@ -94,9 +94,15 @@ public static TryCatch MonadicCreate( return TryCatch.FromResult(stage); } - protected override async Task> GetNextPageAsync(CancellationToken cancellationToken) + public override async ValueTask MoveNextAsync() { - cancellationToken.ThrowIfCancellationRequested(); + this.cancellationToken.ThrowIfCancellationRequested(); + + if (this.returnedFinalPage) + { + this.Current = default; + return false; + } // Draining aggregates is broken down into two stages QueryPage queryPage; @@ -108,13 +114,14 @@ protected override async Task> GetNextPageAsync(Cancellation TryCatch tryGetSourcePage = this.inputStage.Current; if (tryGetSourcePage.Failed) { - return tryGetSourcePage; + this.Current = tryGetSourcePage; + return true; } QueryPage sourcePage = tryGetSourcePage.Result; foreach (CosmosElement element in sourcePage.Documents) { - cancellationToken.ThrowIfCancellationRequested(); + this.cancellationToken.ThrowIfCancellationRequested(); RewrittenAggregateProjections rewrittenAggregateProjections = new RewrittenAggregateProjections( this.isValueQuery, @@ -158,9 +165,11 @@ protected override async Task> GetNextPageAsync(Cancellation state: default); queryPage = finalPage; + this.returnedFinalPage = true; } - return TryCatch.FromResult(queryPage); + this.Current = TryCatch.FromResult(queryPage); + return true; } private sealed class AggregateContinuationToken diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Aggregate/AggregateQueryPipelineStage.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Aggregate/AggregateQueryPipelineStage.cs index 0530ca5c21..0146c3b15e 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Aggregate/AggregateQueryPipelineStage.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Aggregate/AggregateQueryPipelineStage.cs @@ -38,6 +38,8 @@ internal abstract partial class AggregateQueryPipelineStage : QueryPipelineStage /// private readonly bool isValueQuery; + protected bool returnedFinalPage; + /// /// Initializes a new instance of the AggregateDocumentQueryExecutionComponent class. /// diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CatchAllQueryPipelineStage.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CatchAllQueryPipelineStage.cs index 5c0d11aae5..b134b1a480 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CatchAllQueryPipelineStage.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CatchAllQueryPipelineStage.cs @@ -16,14 +16,22 @@ public CatchAllQueryPipelineStage(IQueryPipelineStage inputStage, CancellationTo { } - protected override async Task> GetNextPageAsync(CancellationToken cancellationToken) + public override async ValueTask MoveNextAsync() { + this.cancellationToken.ThrowIfCancellationRequested(); + try { - await this.inputStage.MoveNextAsync(); - return this.inputStage.Current; + if (!await this.inputStage.MoveNextAsync()) + { + this.Current = default; + return false; + } + + this.Current = this.inputStage.Current; + return true; } - catch (OperationCanceledException) when (cancellationToken.IsCancellationRequested) + catch (OperationCanceledException) when (this.cancellationToken.IsCancellationRequested) { // Per cancellationToken.ThrowIfCancellationRequested(); line above, this function should still throw OperationCanceledException. throw; @@ -31,7 +39,8 @@ protected override async Task> GetNextPageAsync(Cancellation catch (Exception ex) { CosmosException cosmosException = ExceptionToCosmosException.CreateFromException(ex); - return TryCatch.FromException(cosmosException); + this.Current = TryCatch.FromException(cosmosException); + return true; } } } diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/OrderBy/OrderByCrossPartitionQueryPipelineStage.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/OrderBy/OrderByCrossPartitionQueryPipelineStage.cs index 4cf20480c8..cdb82d66a5 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/OrderBy/OrderByCrossPartitionQueryPipelineStage.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/OrderBy/OrderByCrossPartitionQueryPipelineStage.cs @@ -328,10 +328,10 @@ private ValueTask MoveNextAsync_DrainPageAsync() } // Create the continuation token. - string continuationTokenString; + CosmosElement state; if ((this.enumerators.Count == 0) && (this.uninitializedEnumeratorsAndTokens.Count == 0)) { - continuationTokenString = null; + state = null; } else { @@ -347,10 +347,10 @@ private ValueTask MoveNextAsync_DrainPageAsync() CosmosElement cosmosElementOrderByContinuationToken = OrderByContinuationToken.ToCosmosElement(orderByContinuationToken); CosmosArray continuationTokenList = CosmosArray.Create(new List() { cosmosElementOrderByContinuationToken }); - continuationTokenString = continuationTokenList.ToString(); + state = continuationTokenList; } - this.state = continuationTokenString != null ? new QueryState(CosmosString.Create(continuationTokenString)) : null; + this.state = state != null ? new QueryState(state) : null; // Return a page of results // No stats to report, since we already reported it when we moved to this page. @@ -570,25 +570,13 @@ private static TryCatch> MonadicExtractOrderByTok return TryCatch>.FromResult(default); } - if (!(continuationToken is CosmosString continuationTokenString)) + if (!(continuationToken is CosmosArray cosmosArray)) { return TryCatch>.FromException( new MalformedContinuationTokenException( - $"Order by continuation token must be a string: {continuationToken}.")); + $"Order by continuation token must be an array: {continuationToken}.")); } - string rawJson = continuationTokenString.Value; - - TryCatch monadicCosmosArray = CosmosArray.Monadic.Parse(rawJson); - if (monadicCosmosArray.Failed) - { - return TryCatch>.FromException( - new MalformedContinuationTokenException( - $"Order by continuation token must be an array: {continuationToken}.", - monadicCosmosArray.Exception)); - } - - CosmosArray cosmosArray = monadicCosmosArray.Result; if (cosmosArray.Count == 0) { return TryCatch>.FromException( diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Distinct/DistinctQueryPipelineStage.Client.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Distinct/DistinctQueryPipelineStage.Client.cs index 39ba2b9584..b9d14ae3a3 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Distinct/DistinctQueryPipelineStage.Client.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Distinct/DistinctQueryPipelineStage.Client.cs @@ -69,8 +69,8 @@ public static TryCatch MonadicCreate( distinctMapToken: null); } - CosmosElement distinctMapToken = distinctContinuationToken.DistinctMapToken != null - ? CosmosString.Create(distinctContinuationToken.DistinctMapToken) + CosmosElement distinctMapToken = distinctContinuationToken.DistinctMapToken != null + ? CosmosString.Create(distinctContinuationToken.DistinctMapToken) : null; TryCatch tryCreateDistinctMap = DistinctMap.TryCreate( distinctQueryType, @@ -113,15 +113,21 @@ public static TryCatch MonadicCreate( cancellationToken)); } - protected override async Task> GetNextPageAsync(CancellationToken cancellationToken) + public override async ValueTask MoveNextAsync() { - cancellationToken.ThrowIfCancellationRequested(); + this.cancellationToken.ThrowIfCancellationRequested(); + + if (!await this.inputStage.MoveNextAsync()) + { + this.Current = default; + return false; + } - await this.inputStage.MoveNextAsync(); TryCatch tryGetSourcePage = this.inputStage.Current; if (tryGetSourcePage.Failed) { - return tryGetSourcePage; + this.Current = tryGetSourcePage; + return true; } QueryPage sourcePage = tryGetSourcePage.Result; @@ -129,7 +135,7 @@ protected override async Task> GetNextPageAsync(Cancellation List distinctResults = new List(); foreach (CosmosElement document in sourcePage.Documents) { - cancellationToken.ThrowIfCancellationRequested(); + this.cancellationToken.ThrowIfCancellationRequested(); if (this.distinctMap.Add(document, out UInt128 _)) { @@ -145,9 +151,9 @@ protected override async Task> GetNextPageAsync(Cancellation if (sourcePage.State != null) { string updatedContinuationToken = new DistinctContinuationToken( - sourceToken: ((CosmosString)sourcePage.State.Value).Value, + sourceToken: sourcePage.State.Value.ToString(), distinctMapToken: this.distinctMap.GetContinuationToken()).ToString(); - state = new QueryState(CosmosString.Create(updatedContinuationToken)); + state = new QueryState(CosmosElement.Parse(updatedContinuationToken)); } else { @@ -175,7 +181,8 @@ protected override async Task> GetNextPageAsync(Cancellation state: null); } - return TryCatch.FromResult(queryPage); + this.Current = TryCatch.FromResult(queryPage); + return true; } /// diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Distinct/DistinctQueryPipelineStage.Compute.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Distinct/DistinctQueryPipelineStage.Compute.cs index dc43cdbbb5..4529228b6e 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Distinct/DistinctQueryPipelineStage.Compute.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Distinct/DistinctQueryPipelineStage.Compute.cs @@ -79,13 +79,19 @@ public static TryCatch MonadicCreate( cancellationToken)); } - protected override async Task> GetNextPageAsync(CancellationToken cancellationToken) + public override async ValueTask MoveNextAsync() { - await this.inputStage.MoveNextAsync(); + if (!await this.inputStage.MoveNextAsync()) + { + this.Current = default; + return false; + } + TryCatch tryGetSourcePage = this.inputStage.Current; if (tryGetSourcePage.Failed) { - return tryGetSourcePage; + this.Current = tryGetSourcePage; + return true; } QueryPage sourcePage = tryGetSourcePage.Result; @@ -121,7 +127,8 @@ protected override async Task> GetNextPageAsync(Cancellation disallowContinuationTokenMessage: ComputeDistinctQueryPipelineStage.UseTryGetContinuationTokenMessage, state: queryState); - return TryCatch.FromResult(queryPage); + this.Current = TryCatch.FromResult(queryPage); + return true; } private readonly struct DistinctContinuationToken diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/GroupBy/GroupByQueryPipelineStage.Client.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/GroupBy/GroupByQueryPipelineStage.Client.cs index b9a7416f23..abe295e471 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/GroupBy/GroupByQueryPipelineStage.Client.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/GroupBy/GroupByQueryPipelineStage.Client.cs @@ -59,9 +59,15 @@ public static TryCatch MonadicCreate( return TryCatch.FromResult(stage); } - protected override async Task> GetNextPageAsync(CancellationToken cancellationToken) + public override async ValueTask MoveNextAsync() { - cancellationToken.ThrowIfCancellationRequested(); + this.cancellationToken.ThrowIfCancellationRequested(); + + if (this.returnedLastPage) + { + this.Current = default; + return false; + } // Draining GROUP BY is broken down into two stages: @@ -72,14 +78,15 @@ protected override async Task> GetNextPageAsync(Cancellation while (await this.inputStage.MoveNextAsync()) { - cancellationToken.ThrowIfCancellationRequested(); + this.cancellationToken.ThrowIfCancellationRequested(); // Stage 1: // Drain the groupings fully from all continuation and all partitions TryCatch tryGetSourcePage = this.inputStage.Current; if (tryGetSourcePage.Failed) { - return tryGetSourcePage; + this.Current = tryGetSourcePage; + return true; } QueryPage sourcePage = tryGetSourcePage.Result; @@ -94,6 +101,10 @@ protected override async Task> GetNextPageAsync(Cancellation // Stage 2: // Emit the results from the grouping table page by page IReadOnlyList results = this.groupingTable.Drain(maxPageSize); + if (this.groupingTable.Count == 0) + { + this.returnedLastPage = true; + } QueryPage queryPage = new QueryPage( documents: results, @@ -104,7 +115,8 @@ protected override async Task> GetNextPageAsync(Cancellation disallowContinuationTokenMessage: ClientGroupByQueryPipelineStage.ContinuationTokenNotSupportedWithGroupBy, state: null); - return TryCatch.FromResult(queryPage); + this.Current = TryCatch.FromResult(queryPage); + return true; } } } diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/GroupBy/GroupByQueryPipelineStage.Compute.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/GroupBy/GroupByQueryPipelineStage.Compute.cs index bc277b0c53..2aab36e64f 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/GroupBy/GroupByQueryPipelineStage.Compute.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/GroupBy/GroupByQueryPipelineStage.Compute.cs @@ -91,9 +91,15 @@ public static TryCatch MonadicCreate( tryCreateGroupingTable.Result)); } - protected override async Task> GetNextPageAsync(CancellationToken cancellationToken) + public override async ValueTask MoveNextAsync() { - cancellationToken.ThrowIfCancellationRequested(); + this.cancellationToken.ThrowIfCancellationRequested(); + + if (this.returnedLastPage) + { + this.Current = default; + return false; + } // Draining GROUP BY is broken down into two stages: QueryPage queryPage; @@ -104,7 +110,8 @@ protected override async Task> GetNextPageAsync(Cancellation TryCatch tryGetSourcePage = this.inputStage.Current; if (tryGetSourcePage.Failed) { - return tryGetSourcePage; + this.Current = tryGetSourcePage; + return true; } QueryPage sourcePage = tryGetSourcePage.Result; @@ -137,6 +144,7 @@ protected override async Task> GetNextPageAsync(Cancellation if (this.groupingTable.IsDone) { state = default; + this.returnedLastPage = true; } else { @@ -156,7 +164,8 @@ protected override async Task> GetNextPageAsync(Cancellation state: state); } - return TryCatch.FromResult(queryPage); + this.Current = TryCatch.FromResult(queryPage); + return true; } private readonly struct GroupByContinuationToken diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/GroupBy/GroupByQueryPipelineStage.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/GroupBy/GroupByQueryPipelineStage.cs index f50c1abbc6..77f2a98dea 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/GroupBy/GroupByQueryPipelineStage.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/GroupBy/GroupByQueryPipelineStage.cs @@ -48,6 +48,7 @@ namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.GroupBy internal abstract partial class GroupByQueryPipelineStage : QueryPipelineStageBase { private readonly GroupingTable groupingTable; + protected bool returnedLastPage; protected GroupByQueryPipelineStage( IQueryPipelineStage source, diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/LazyQueryPipelineStage.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/LazyQueryPipelineStage.cs index 746d36c277..93d235a440 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/LazyQueryPipelineStage.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/LazyQueryPipelineStage.cs @@ -50,6 +50,7 @@ public async ValueTask MoveNextAsync() IQueryPipelineStage stage = tryCreateStage.Result; if (!await stage.MoveNextAsync()) { + this.Current = default; return false; } diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/PipelineFactory.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/PipelineFactory.cs index 741bda4680..a5b93f07cd 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/PipelineFactory.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/PipelineFactory.cs @@ -47,7 +47,7 @@ public static TryCatch MonadicCreate( .Zip(queryInfo.OrderBy, (expression, sortOrder) => new OrderByColumn(expression, sortOrder)).ToList(), pageSize: pageSize, maxConcurrency: maxConcurrency, - continuationToken: requestContinuationToken, + continuationToken: continuationToken, cancellationToken: cancellationToken); } else diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/QueryPipelineStageBase.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/QueryPipelineStageBase.cs index 78206754df..552321630d 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/QueryPipelineStageBase.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/QueryPipelineStageBase.cs @@ -13,7 +13,6 @@ internal abstract class QueryPipelineStageBase : IQueryPipelineStage { protected readonly IQueryPipelineStage inputStage; protected CancellationToken cancellationToken; - private bool hasStarted; protected QueryPipelineStageBase(IQueryPipelineStage inputStage, CancellationToken cancellationToken) { @@ -21,33 +20,11 @@ protected QueryPipelineStageBase(IQueryPipelineStage inputStage, CancellationTok this.cancellationToken = cancellationToken; } - public TryCatch Current { get; private set; } - - public QueryState State { get; protected set; } - - public bool HasMoreResults => !this.hasStarted || (this.State != default); + public TryCatch Current { get; protected set; } public ValueTask DisposeAsync() => this.inputStage.DisposeAsync(); - protected abstract Task> GetNextPageAsync(CancellationToken cancellationToken); - - public async ValueTask MoveNextAsync() - { - if (!this.HasMoreResults) - { - return false; - } - - this.hasStarted = true; - - this.Current = await this.GetNextPageAsync(this.cancellationToken); - if (this.Current.Succeeded) - { - this.State = this.Current.Result.State; - } - - return true; - } + public abstract ValueTask MoveNextAsync(); public void SetCancellationToken(CancellationToken cancellationToken) { diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Skip/SkipQueryPipelineStage.Client.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Skip/SkipQueryPipelineStage.Client.cs index 5447d0df48..0fe7cdf8c9 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Skip/SkipQueryPipelineStage.Client.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Skip/SkipQueryPipelineStage.Client.cs @@ -7,7 +7,6 @@ namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.Skip using System; using System.Collections.Generic; using System.Linq; - using System.Text; using System.Threading; using System.Threading.Tasks; using Microsoft.Azure.Cosmos.CosmosElements; @@ -20,8 +19,8 @@ internal abstract partial class SkipQueryPipelineStage : QueryPipelineStageBase private sealed class ClientSkipQueryPipelineStage : SkipQueryPipelineStage { private ClientSkipQueryPipelineStage( - IQueryPipelineStage source, - CancellationToken cancellationToken, + IQueryPipelineStage source, + CancellationToken cancellationToken, long skipCount) : base(source, cancellationToken, skipCount) { @@ -94,15 +93,21 @@ public static TryCatch MonadicCreate( return TryCatch.FromResult(stage); } - protected override async Task> GetNextPageAsync(CancellationToken cancellationToken) + public override async ValueTask MoveNextAsync() { - cancellationToken.ThrowIfCancellationRequested(); + this.cancellationToken.ThrowIfCancellationRequested(); + + if (!await this.inputStage.MoveNextAsync()) + { + this.Current = default; + return false; + } - await this.inputStage.MoveNextAsync(); TryCatch tryGetSourcePage = this.inputStage.Current; if (tryGetSourcePage.Failed) { - return tryGetSourcePage; + this.Current = tryGetSourcePage; + return true; } QueryPage sourcePage = tryGetSourcePage.Result; @@ -135,7 +140,8 @@ protected override async Task> GetNextPageAsync(Cancellation disallowContinuationTokenMessage: sourcePage.DisallowContinuationTokenMessage, state: state); - return TryCatch.FromResult(queryPage); + this.Current = TryCatch.FromResult(queryPage); + return true; } /// diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Skip/SkipQueryPipelineStage.Compute.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Skip/SkipQueryPipelineStage.Compute.cs index 0b54c0dab8..8b233ab6ce 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Skip/SkipQueryPipelineStage.Compute.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Skip/SkipQueryPipelineStage.Compute.cs @@ -26,7 +26,7 @@ private ComputeSkipQueryPipelineStage(IQueryPipelineStage source, CancellationTo public static TryCatch MonadicCreate( int offsetCount, - CosmosElement continuationToken, + CosmosElement continuationToken, CancellationToken cancellationToken, MonadicCreatePipelineStage monadicCreatePipelineStage) { @@ -73,15 +73,21 @@ public static TryCatch MonadicCreate( return TryCatch.FromResult(stage); } - protected override async Task> GetNextPageAsync(CancellationToken cancellationToken) + public override async ValueTask MoveNextAsync() { - cancellationToken.ThrowIfCancellationRequested(); + this.cancellationToken.ThrowIfCancellationRequested(); + + if (!await this.inputStage.MoveNextAsync()) + { + this.Current = default; + return false; + } - await this.inputStage.MoveNextAsync(); TryCatch tryGetSourcePage = this.inputStage.Current; if (tryGetSourcePage.Failed) { - return tryGetSourcePage; + this.Current = tryGetSourcePage; + return true; } QueryPage sourcePage = tryGetSourcePage.Result; @@ -115,7 +121,8 @@ protected override async Task> GetNextPageAsync(Cancellation disallowContinuationTokenMessage: sourcePage.DisallowContinuationTokenMessage, state: state); - return TryCatch.FromResult(queryPage); + this.Current = TryCatch.FromResult(queryPage); + return true; } /// diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/SkipEmptyPageQueryPipelineStage.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/SkipEmptyPageQueryPipelineStage.cs index a97aa2a82b..dd7bfbfda5 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/SkipEmptyPageQueryPipelineStage.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/SkipEmptyPageQueryPipelineStage.cs @@ -5,12 +5,16 @@ namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline { using System; + using System.Collections.Generic; using System.Threading; using System.Threading.Tasks; + using Microsoft.Azure.Cosmos.CosmosElements; using Microsoft.Azure.Cosmos.Query.Core.Monads; internal sealed class SkipEmptyPageQueryPipelineStage : IQueryPipelineStage { + private static readonly IReadOnlyList EmptyPage = new List(); + private readonly IQueryPipelineStage inputStage; private double cumulativeRequestCharge; private long cumulativeResponseLengthInBytes; @@ -28,6 +32,8 @@ public SkipEmptyPageQueryPipelineStage(IQueryPipelineStage inputStage, Cancellat public async ValueTask MoveNextAsync() { + this.cancellationToken.ThrowIfCancellationRequested(); + if (!await this.inputStage.MoveNextAsync()) { this.Current = default; @@ -46,6 +52,22 @@ public async ValueTask MoveNextAsync() { this.cumulativeRequestCharge += sourcePage.RequestCharge; this.cumulativeResponseLengthInBytes += sourcePage.ResponseLengthInBytes; + if (sourcePage.State == null) + { + QueryPage queryPage = new QueryPage( + documents: EmptyPage, + requestCharge: sourcePage.RequestCharge + this.cumulativeRequestCharge, + activityId: sourcePage.ActivityId, + responseLengthInBytes: sourcePage.ResponseLengthInBytes + this.cumulativeResponseLengthInBytes, + cosmosQueryExecutionInfo: sourcePage.CosmosQueryExecutionInfo, + disallowContinuationTokenMessage: sourcePage.DisallowContinuationTokenMessage, + state: default); + this.cumulativeRequestCharge = 0; + this.cumulativeResponseLengthInBytes = 0; + this.Current = TryCatch.FromResult(queryPage); + return true; + } + return await this.MoveNextAsync(); } diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Take/TakeQueryPipelineStage.Client.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Take/TakeQueryPipelineStage.Client.cs index ea5c57c9f2..ffc2bc50a7 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Take/TakeQueryPipelineStage.Client.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Take/TakeQueryPipelineStage.Client.cs @@ -174,15 +174,21 @@ public static TryCatch MonadicCreateTopStage( return TryCatch.FromResult(stage); } - protected override async Task> GetNextPageAsync(CancellationToken cancellationToken) + public override async ValueTask MoveNextAsync() { - cancellationToken.ThrowIfCancellationRequested(); + this.cancellationToken.ThrowIfCancellationRequested(); + + if (!await this.inputStage.MoveNextAsync()) + { + this.Current = default; + return false; + } - await this.inputStage.MoveNextAsync(); TryCatch tryGetSourcePage = this.inputStage.Current; if (tryGetSourcePage.Failed) { - return tryGetSourcePage; + this.Current = tryGetSourcePage; + return true; } QueryPage sourcePage = tryGetSourcePage.Result; @@ -220,7 +226,8 @@ protected override async Task> GetNextPageAsync(Cancellation disallowContinuationTokenMessage: sourcePage.DisallowContinuationTokenMessage, state: state); - return TryCatch.FromResult(queryPage); + this.Current = TryCatch.FromResult(queryPage); + return true; } private enum TakeEnum diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Take/TakeQueryPipelineStage.Compute.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Take/TakeQueryPipelineStage.Compute.cs index d22e8e7a32..37dfff4cd3 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Take/TakeQueryPipelineStage.Compute.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Take/TakeQueryPipelineStage.Compute.cs @@ -99,15 +99,21 @@ private static TryCatch MonadicCreate( return TryCatch.FromResult(stage); } - protected override async Task> GetNextPageAsync(CancellationToken cancellationToken) + public override async ValueTask MoveNextAsync() { - cancellationToken.ThrowIfCancellationRequested(); + this.cancellationToken.ThrowIfCancellationRequested(); + + if (!await this.inputStage.MoveNextAsync()) + { + this.Current = default; + return false; + } - await this.inputStage.MoveNextAsync(); TryCatch tryGetSourcePage = this.inputStage.Current; if (tryGetSourcePage.Failed) { - return tryGetSourcePage; + this.Current = tryGetSourcePage; + return true; } QueryPage sourcePage = tryGetSourcePage.Result; @@ -137,7 +143,8 @@ protected override async Task> GetNextPageAsync(Cancellation disallowContinuationTokenMessage: sourcePage.DisallowContinuationTokenMessage, state: queryState); - return TryCatch.FromResult(queryPage); + this.Current = TryCatch.FromResult(queryPage); + return true; } private readonly struct TakeContinuationToken diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/MockQueryPipelineStage.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/MockQueryPipelineStage.cs index 2b1b8cc87b..5881c16a61 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/MockQueryPipelineStage.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/MockQueryPipelineStage.cs @@ -23,8 +23,14 @@ public MockQueryPipelineStage(IReadOnlyList> pages) this.pages = pages ?? throw new ArgumentNullException(nameof(pages)); } - protected override Task> GetNextPageAsync(CancellationToken cancellationToken) + public override ValueTask MoveNextAsync() { + if (this.pageIndex == this.pages.Count) + { + this.Current = default; + return new ValueTask(false); + } + IReadOnlyList documents = this.pages[this.pageIndex++]; QueryState state = (this.pageIndex == this.pages.Count) ? null : new QueryState(CosmosString.Create(this.pageIndex.ToString())); QueryPage page = new QueryPage( @@ -35,7 +41,8 @@ protected override Task> GetNextPageAsync(CancellationToken cosmosQueryExecutionInfo: default, disallowContinuationTokenMessage: default, state: state); - return Task.FromResult(TryCatch.FromResult(page)); + this.Current = TryCatch.FromResult(page); + return new ValueTask(true); } } } diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/OrderByCrossPartitionQueryPipelineStageTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/OrderByCrossPartitionQueryPipelineStageTests.cs index 25e71a68c5..8007d51b91 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/OrderByCrossPartitionQueryPipelineStageTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/OrderByCrossPartitionQueryPipelineStageTests.cs @@ -135,12 +135,11 @@ public void MonadicCreate_SingleOrderByContinuationToken() pageSize: 10, maxConcurrency: 10, cancellationToken: default, - continuationToken: CosmosString.Create( - CosmosArray.Create( - new List() - { - OrderByContinuationToken.ToCosmosElement(orderByContinuationToken) - }).ToString())); + continuationToken: CosmosArray.Create( + new List() + { + OrderByContinuationToken.ToCosmosElement(orderByContinuationToken) + })); Assert.IsTrue(monadicCreate.Succeeded); } @@ -186,13 +185,12 @@ public void MonadicCreate_MultipleOrderByContinuationToken() pageSize: 10, maxConcurrency: 10, cancellationToken: default, - continuationToken: CosmosString.Create( - CosmosArray.Create( - new List() - { - OrderByContinuationToken.ToCosmosElement(orderByContinuationToken1), - OrderByContinuationToken.ToCosmosElement(orderByContinuationToken2) - }).ToString())); + continuationToken: CosmosArray.Create( + new List() + { + OrderByContinuationToken.ToCosmosElement(orderByContinuationToken1), + OrderByContinuationToken.ToCosmosElement(orderByContinuationToken2) + })); Assert.IsTrue(monadicCreate.Succeeded); } @@ -238,6 +236,48 @@ FROM c Assert.IsTrue(documents.OrderBy(document => ((CosmosObject)document)["_ts"]).ToList().SequenceEqual(documents)); } + [TestMethod] + public async Task TestDrainFully_StartFromBeginingAsync_NoDocuments() + { + int numItems = 0; + IDocumentContainer documentContainer = await CreateDocumentContainerAsync(numItems); + + TryCatch monadicCreate = OrderByCrossPartitionQueryPipelineStage.MonadicCreate( + documentContainer: documentContainer, + sqlQuerySpec: new SqlQuerySpec(@" + SELECT c._rid AS _rid, [{""item"": c._ts}] AS orderByItems, c AS payload + FROM c + WHERE {documentdb-formattableorderbyquery-filter} + ORDER BY c._ts"), + targetRanges: await documentContainer.GetFeedRangesAsync(cancellationToken: default), + orderByColumns: new List() + { + new OrderByColumn("c._ts", SortOrder.Ascending) + }, + pageSize: 10, + maxConcurrency: 10, + cancellationToken: default, + continuationToken: null); + Assert.IsTrue(monadicCreate.Succeeded); + IQueryPipelineStage queryPipelineStage = monadicCreate.Result; + + List documents = new List(); + while (await queryPipelineStage.MoveNextAsync()) + { + TryCatch tryGetQueryPage = queryPipelineStage.Current; + if (tryGetQueryPage.Failed) + { + Assert.Fail(tryGetQueryPage.Exception.ToString()); + } + + QueryPage queryPage = tryGetQueryPage.Result; + documents.AddRange(queryPage.Documents); + } + + Assert.AreEqual(numItems, documents.Count); + Assert.IsTrue(documents.OrderBy(document => ((CosmosObject)document)["_ts"]).ToList().SequenceEqual(documents)); + } + [TestMethod] public async Task TestDrainFully_WithStateResume() { From 2ba7a5acb57b110526d09872442ffeebcdf18fbb Mon Sep 17 00:00:00 2001 From: Brandon Chong Date: Tue, 29 Sep 2020 23:12:55 -0700 Subject: [PATCH 73/85] fixed some test cases --- .../Query/Core/Pipeline/PipelineFactory.cs | 27 +++++ .../Skip/SkipQueryPipelineStage.Client.cs | 6 +- .../SkipEmptyPageQueryPipelineStage.cs | 18 ++++ .../Take/TakeQueryPipelineStage.Client.cs | 6 +- .../src/Query/v3Query/QueryIterator.cs | 27 ++++- .../src/Reactive/JustAsyncEnumerator.cs | 3 +- .../Pagination/InMemoryContainer.cs | 14 ++- .../Query/OfflineEngine/Projector.cs | 35 +++--- .../Query/Pipeline/FullPipelineTests.cs | 100 +++++++++++++++++- ...ByCrossPartitionQueryPipelineStageTests.cs | 7 +- 10 files changed, 208 insertions(+), 35 deletions(-) diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/PipelineFactory.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/PipelineFactory.cs index a5b93f07cd..ae17717c20 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/PipelineFactory.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/PipelineFactory.cs @@ -35,6 +35,33 @@ public static TryCatch MonadicCreate( CosmosElement requestContinuationToken, CancellationToken requestCancellationToken) { + if (documentContainer == null) + { + throw new ArgumentNullException(nameof(documentContainer)); + } + + if (sqlQuerySpec == null) + { + throw new ArgumentNullException(nameof(sqlQuerySpec)); + } + + if (targetRanges == null) + { + throw new ArgumentNullException(nameof(targetRanges)); + } + + if (targetRanges.Count == 0) + { + throw new ArgumentException($"{nameof(targetRanges)} must not be empty."); + } + + if (queryInfo == null) + { + throw new ArgumentNullException(nameof(queryInfo)); + } + + sqlQuerySpec = !string.IsNullOrEmpty(queryInfo.RewrittenQuery) ? new SqlQuerySpec(queryInfo.RewrittenQuery) : sqlQuerySpec; + MonadicCreatePipelineStage monadicCreatePipelineStage; if (queryInfo.HasOrderBy) { diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Skip/SkipQueryPipelineStage.Client.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Skip/SkipQueryPipelineStage.Client.cs index 0fe7cdf8c9..4d991c2eb2 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Skip/SkipQueryPipelineStage.Client.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Skip/SkipQueryPipelineStage.Client.cs @@ -119,12 +119,12 @@ public override async ValueTask MoveNextAsync() this.skipCount -= numberOfDocumentsSkipped; QueryState state; - if (sourcePage.DisallowContinuationTokenMessage == null) + if ((sourcePage.State != null) && (sourcePage.DisallowContinuationTokenMessage == null)) { string token = new OffsetContinuationToken( offset: this.skipCount, - sourceToken: sourcePage.State != null ? ((CosmosString)sourcePage.State.Value).Value : null).ToString(); - state = new QueryState(CosmosString.Create(token)); + sourceToken: sourcePage.State?.Value.ToString()).ToString(); + state = new QueryState(CosmosElement.Parse(token)); } else { diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/SkipEmptyPageQueryPipelineStage.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/SkipEmptyPageQueryPipelineStage.cs index dd7bfbfda5..b330bfe3ab 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/SkipEmptyPageQueryPipelineStage.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/SkipEmptyPageQueryPipelineStage.cs @@ -19,6 +19,7 @@ internal sealed class SkipEmptyPageQueryPipelineStage : IQueryPipelineStage private double cumulativeRequestCharge; private long cumulativeResponseLengthInBytes; private CancellationToken cancellationToken; + private bool returnedFinalStats; public SkipEmptyPageQueryPipelineStage(IQueryPipelineStage inputStage, CancellationToken cancellationToken) { @@ -36,6 +37,23 @@ public async ValueTask MoveNextAsync() if (!await this.inputStage.MoveNextAsync()) { + if (!this.returnedFinalStats) + { + QueryPage queryPage = new QueryPage( + documents: EmptyPage, + requestCharge: this.cumulativeRequestCharge, + activityId: Guid.Empty.ToString(), + responseLengthInBytes: this.cumulativeResponseLengthInBytes, + cosmosQueryExecutionInfo: default, + disallowContinuationTokenMessage: default, + state: default); + this.cumulativeRequestCharge = 0; + this.cumulativeResponseLengthInBytes = 0; + this.returnedFinalStats = true; + this.Current = TryCatch.FromResult(queryPage); + return true; + } + this.Current = default; return false; } diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Take/TakeQueryPipelineStage.Client.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Take/TakeQueryPipelineStage.Client.cs index ffc2bc50a7..77b96af3d5 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Take/TakeQueryPipelineStage.Client.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Take/TakeQueryPipelineStage.Client.cs @@ -203,14 +203,14 @@ public override async ValueTask MoveNextAsync() { TakeEnum.Limit => new LimitContinuationToken( limit: this.takeCount, - sourceToken: ((CosmosString)sourcePage.State.Value).Value).ToString(), + sourceToken: sourcePage.State?.Value.ToString()).ToString(), TakeEnum.Top => new TopContinuationToken( top: this.takeCount, - sourceToken: ((CosmosString)sourcePage.State.Value).Value).ToString(), + sourceToken: sourcePage.State?.Value.ToString()).ToString(), _ => throw new ArgumentOutOfRangeException($"Unknown {nameof(TakeEnum)}: {this.takeEnum}."), }; - state = new QueryState(CosmosString.Create(updatedContinuationToken)); + state = new QueryState(CosmosElement.Parse(updatedContinuationToken)); } else { diff --git a/Microsoft.Azure.Cosmos/src/Query/v3Query/QueryIterator.cs b/Microsoft.Azure.Cosmos/src/Query/v3Query/QueryIterator.cs index 9732b50a3a..1694ca6de0 100644 --- a/Microsoft.Azure.Cosmos/src/Query/v3Query/QueryIterator.cs +++ b/Microsoft.Azure.Cosmos/src/Query/v3Query/QueryIterator.cs @@ -5,6 +5,7 @@ namespace Microsoft.Azure.Cosmos.Query { using System; + using System.Collections.Generic; using System.Threading; using System.Threading.Tasks; using Microsoft.Azure.Cosmos.CosmosElements; @@ -20,6 +21,8 @@ namespace Microsoft.Azure.Cosmos.Query internal sealed class QueryIterator : FeedIteratorInternal { + private static readonly IReadOnlyList EmptyPage = new List(); + private readonly CosmosQueryContextCore cosmosQueryContext; private readonly IQueryPipelineStage queryPipelineStage; private readonly CosmosSerializationFormatOptions cosmosSerializationFormatOptions; @@ -152,7 +155,27 @@ public override async Task ReadNextAsync(CancellationToken canc { // This catches exception thrown by the pipeline and converts it to QueryResponse this.queryPipelineStage.SetCancellationToken(cancellationToken); - await this.queryPipelineStage.MoveNextAsync(); + if (!await this.queryPipelineStage.MoveNextAsync()) + { + this.hasMoreResults = false; + return QueryResponse.CreateSuccess( + result: EmptyPage, + count: EmptyPage.Count, + responseLengthBytes: default, + diagnostics: default, + serializationOptions: this.cosmosSerializationFormatOptions, + responseHeaders: new CosmosQueryResponseMessageHeaders( + continauationToken: default, + disallowContinuationTokenMessage: default, + this.cosmosQueryContext.ResourceTypeEnum, + this.cosmosQueryContext.ContainerResourceId) + { + RequestCharge = default, + ActivityId = Guid.Empty.ToString(), + SubStatusCode = Documents.SubStatusCodes.Unknown + }); + } + tryGetQueryPage = this.queryPipelineStage.Current; } catch (OperationCanceledException ex) when (!(ex is CosmosOperationCanceledException)) @@ -168,7 +191,7 @@ public override async Task ReadNextAsync(CancellationToken canc if (tryGetQueryPage.Succeeded) { - if (tryGetQueryPage.Result.State == null) + if ((tryGetQueryPage.Result.State == null) && (tryGetQueryPage.Result.DisallowContinuationTokenMessage == null)) { this.hasMoreResults = false; } diff --git a/Microsoft.Azure.Cosmos/src/Reactive/JustAsyncEnumerator.cs b/Microsoft.Azure.Cosmos/src/Reactive/JustAsyncEnumerator.cs index cd24a96e73..ea36359759 100644 --- a/Microsoft.Azure.Cosmos/src/Reactive/JustAsyncEnumerator.cs +++ b/Microsoft.Azure.Cosmos/src/Reactive/JustAsyncEnumerator.cs @@ -5,6 +5,7 @@ namespace Microsoft.Azure.Cosmos.Reactive { using System.Collections.Generic; + using System.Linq; using System.Threading.Tasks; /// @@ -18,7 +19,7 @@ internal sealed class JustAsyncEnumerator : IAsyncEnumerator public JustAsyncEnumerator(params T[] items) { - this.enumerator = (IEnumerator)items.GetEnumerator(); + this.enumerator = items.ToList().GetEnumerator(); } public T Current => this.enumerator.Current; 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 4be4a82e20..2c358e13a6 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 @@ -373,8 +373,20 @@ public Task> MonadicQueryAsync( queryPageResults = queryPageResults.Take(pageSize); List queryPageResultList = queryPageResults.ToList(); + QueryState queryState; + if (queryPageResultList.Count == 0) + { + queryState = default; + } + else if (queryPageResultList.Last() is CosmosObject lastDocument) + { + queryState = new QueryState(lastDocument["_rid"]); + } + else + { + queryState = default; + } - QueryState queryState = queryPageResultList.Count == 0 ? null : new QueryState(((CosmosObject)queryPageResultList.Last())["_rid"]); return Task.FromResult( TryCatch.FromResult( new QueryPage( diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OfflineEngine/Projector.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OfflineEngine/Projector.cs index c1751e56bf..bf2a182b46 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OfflineEngine/Projector.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OfflineEngine/Projector.cs @@ -40,23 +40,23 @@ public override CosmosElement Visit(SqlSelectListSpec selectSpec, CosmosElement { key = sqlSelectItem.Alias.Value; } - else + else if (sqlSelectItem.Expression is SqlMemberIndexerScalarExpression memberIndexer) { - if (sqlSelectItem.Expression is SqlMemberIndexerScalarExpression memberIndexer) - { - SqlScalarExpression lastToken = this.GetLastMemberIndexerToken(memberIndexer); + SqlScalarExpression lastToken = Projector.GetLastMemberIndexerToken(memberIndexer); - if (lastToken is SqlLiteralScalarExpression literalScalarExpression) + if (lastToken is SqlLiteralScalarExpression literalScalarExpression) + { + if (literalScalarExpression.Literal is SqlStringLiteral stringLiteral) { - if (literalScalarExpression.Literal is SqlStringLiteral stringLiteral) - { - key = stringLiteral.Value; - } + key = stringLiteral.Value; } } } - - if (key == default) + else if (sqlSelectItem.Expression is SqlPropertyRefScalarExpression propertyRef) + { + key = propertyRef.Identifier.Value; + } + else { key = $"${aliasCounter++}"; } @@ -107,19 +107,14 @@ public override CosmosElement Visit(SqlSelectValueSpec selectSpec, CosmosElement return selectSpec.Expression.Accept(ScalarExpressionEvaluator.Singleton, document); } - private SqlScalarExpression GetLastMemberIndexerToken(SqlMemberIndexerScalarExpression memberIndexer) + private static SqlScalarExpression GetLastMemberIndexerToken(SqlMemberIndexerScalarExpression memberIndexer) { - SqlScalarExpression last; - if (!(memberIndexer.Indexer is SqlMemberIndexerScalarExpression memberIndexerExpression)) - { - last = memberIndexer.Indexer; - } - else + if (memberIndexer.Indexer is SqlMemberIndexerScalarExpression memberIndexerExpression) { - last = this.GetLastMemberIndexerToken(memberIndexerExpression); + return Projector.GetLastMemberIndexerToken(memberIndexerExpression); } - return last; + return memberIndexer.Indexer; } } } 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 9f328767c9..8e72f95eb1 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 @@ -12,7 +12,6 @@ namespace Microsoft.Azure.Cosmos.Tests.Query.Pipeline using Microsoft.Azure.Cosmos.CosmosElements; using Microsoft.Azure.Cosmos.Pagination; 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.QueryPlan; @@ -73,6 +72,88 @@ public async Task SelectStar() Assert.AreEqual(expected: documents.Count, actual: documentsQueried.Count); } + [TestMethod] + public async Task OrderBy() + { + List documents = new List(); + for (int i = 0; i < 250; i++) + { + documents.Add(CosmosObject.Parse($"{{\"pk\" : {i} }}")); + } + + List documentsQueried = await ExecuteQueryAsync( + query: "SELECT * FROM c ORDER BY c._ts", + documents: documents); + + Assert.AreEqual(expected: documents.Count, actual: documentsQueried.Count); + } + + [TestMethod] + public async Task Top() + { + List documents = new List(); + for (int i = 0; i < 250; i++) + { + documents.Add(CosmosObject.Parse($"{{\"pk\" : {i} }}")); + } + + List documentsQueried = await ExecuteQueryAsync( + query: "SELECT TOP 10 * FROM c", + documents: documents); + + Assert.AreEqual(expected: 10, actual: documentsQueried.Count); + } + + [TestMethod] + public async Task OffsetLimit() + { + List documents = new List(); + for (int i = 0; i < 250; i++) + { + documents.Add(CosmosObject.Parse($"{{\"pk\" : {i} }}")); + } + + List documentsQueried = await ExecuteQueryAsync( + query: "SELECT * FROM c OFFSET 10 LIMIT 103", + documents: documents); + + Assert.AreEqual(expected: 103, actual: documentsQueried.Count); + } + + [TestMethod] + public async Task Aggregates() + { + List documents = new List(); + for (int i = 0; i < 250; i++) + { + documents.Add(CosmosObject.Parse($"{{\"pk\" : {i} }}")); + } + + List documentsQueried = await ExecuteQueryAsync( + query: "SELECT VALUE COUNT(1) FROM c", + documents: documents); + + Assert.AreEqual(expected: 1, actual: documentsQueried.Count); + } + + [TestMethod] + [Ignore] + // Need to implement group by continuation token on the in memory collection. + public async Task GroupBy() + { + List documents = new List(); + for (int i = 0; i < 250; i++) + { + documents.Add(CosmosObject.Parse($"{{\"pk\" : {i} }}")); + } + + List documentsQueried = await ExecuteQueryAsync( + query: "SELECT VALUE COUNT(1) FROM c GROUP BY c.pk", + documents: documents); + + Assert.AreEqual(expected: documents.Count, actual: documentsQueried.Count); + } + private static async Task> ExecuteQueryAsync( string query, IReadOnlyList documents) @@ -87,6 +168,22 @@ private static async Task> ExecuteQueryAsync( return resultsFromDrainWithoutState; } + [TestMethod] + public async Task Fuzz() + { + List documents = new List(); + for (int i = 0; i < 250; i++) + { + documents.Add(CosmosObject.Parse($"{{\"pk\" : {i} }}")); + } + + List documentsQueried = await ExecuteQueryAsync( + query: "SELECT * FROM c ORDER BY c._ts OFFSET 1 LIMIT 500", + documents: documents); + + Assert.AreEqual(expected: 249, actual: documentsQueried.Count); + } + private static async Task> DrainWithoutStateAsync(string query, IDocumentContainer documentContainer) { IQueryPipelineStage pipelineStage = CreatePipeline(documentContainer, query); @@ -119,7 +216,6 @@ private static async Task> DrainWithStateAsync(string query, } TryCatch tryGetQueryPage = pipelineStage.Current; - tryGetQueryPage.ThrowIfFailed(); elements.AddRange(tryGetQueryPage.Result.Documents); diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/OrderByCrossPartitionQueryPipelineStageTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/OrderByCrossPartitionQueryPipelineStageTests.cs index 8007d51b91..e00cb7e866 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/OrderByCrossPartitionQueryPipelineStageTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/OrderByCrossPartitionQueryPipelineStageTests.cs @@ -387,7 +387,8 @@ public async Task TestDrainFully_WithSplit_WithStateResume() int numItems = 1000; IDocumentContainer documentContainer = await CreateDocumentContainerAsync(numItems); - Random random = new Random(); + int seed = new Random().Next(); + Random random = new Random(seed); List documents = new List(); QueryState queryState = null; @@ -403,7 +404,7 @@ FROM c targetRanges: await documentContainer.GetFeedRangesAsync(cancellationToken: default), orderByColumns: new List() { - new OrderByColumn("c._ts", SortOrder.Ascending) + new OrderByColumn("c._ts", SortOrder.Ascending) }, pageSize: 10, maxConcurrency: 10, @@ -436,7 +437,7 @@ FROM c } while (queryState != null); Assert.AreEqual(numItems, documents.Count); - Assert.IsTrue(documents.OrderBy(document => ((CosmosObject)document)["_ts"]).ToList().SequenceEqual(documents)); + Assert.IsTrue(documents.OrderBy(document => ((CosmosObject)document)["_ts"]).ToList().SequenceEqual(documents), $"Failed with seed: {seed}"); } private static async Task CreateDocumentContainerAsync( From 78c2f98be017705c6a65c38c5feb4dd6989d622d Mon Sep 17 00:00:00 2001 From: Brandon Chong Date: Wed, 30 Sep 2020 13:39:35 -0700 Subject: [PATCH 74/85] fixed test --- .../Query/OfflineEngine/Projector.cs | 17 ++++++----------- .../Query/OfflineEngineTests/ProjectorTests.cs | 2 -- 2 files changed, 6 insertions(+), 13 deletions(-) diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OfflineEngine/Projector.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OfflineEngine/Projector.cs index bf2a182b46..368d88f6e7 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OfflineEngine/Projector.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OfflineEngine/Projector.cs @@ -35,22 +35,17 @@ public override CosmosElement Visit(SqlSelectListSpec selectSpec, CosmosElement CosmosElement value = sqlSelectItem.Expression.Accept(ScalarExpressionEvaluator.Singleton, document); if (value != null) { - string key = default; + string key; if (sqlSelectItem.Alias != null) { key = sqlSelectItem.Alias.Value; } - else if (sqlSelectItem.Expression is SqlMemberIndexerScalarExpression memberIndexer) + else if ( + sqlSelectItem.Expression is SqlMemberIndexerScalarExpression memberIndexer + && Projector.GetLastMemberIndexerToken(memberIndexer) is SqlLiteralScalarExpression literalScalarExpression + && literalScalarExpression.Literal is SqlStringLiteral stringLiteral) { - SqlScalarExpression lastToken = Projector.GetLastMemberIndexerToken(memberIndexer); - - if (lastToken is SqlLiteralScalarExpression literalScalarExpression) - { - if (literalScalarExpression.Literal is SqlStringLiteral stringLiteral) - { - key = stringLiteral.Value; - } - } + key = stringLiteral.Value; } else if (sqlSelectItem.Expression is SqlPropertyRefScalarExpression propertyRef) { diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OfflineEngineTests/ProjectorTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OfflineEngineTests/ProjectorTests.cs index cca380f2bb..77e1128c95 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OfflineEngineTests/ProjectorTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OfflineEngineTests/ProjectorTests.cs @@ -6,8 +6,6 @@ namespace Microsoft.Azure.Cosmos.Tests.Query.OfflineEngineTests { using System.Collections.Generic; - using Newtonsoft.Json; - using Newtonsoft.Json.Linq; using SqlObjects; using OfflineEngine; using VisualStudio.TestTools.UnitTesting; From 0ebb26c1b8e926c7c990aa66aff7c8ccd64e07fe Mon Sep 17 00:00:00 2001 From: Brandon Chong Date: Thu, 1 Oct 2020 17:59:42 -0700 Subject: [PATCH 75/85] fixed page size bugs with group by --- .../GroupByQueryPipelineStage.Client.cs | 18 ++- .../GroupByQueryPipelineStage.Compute.cs | 14 +- .../GroupBy/GroupByQueryPipelineStage.cs | 16 ++- .../Query/Core/Pipeline/PipelineFactory.cs | 3 +- .../Query/QueryTestsBase.cs | 131 ++++++------------ .../Query/SanityQueryTests.cs | 25 ++-- 6 files changed, 82 insertions(+), 125 deletions(-) diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/GroupBy/GroupByQueryPipelineStage.Client.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/GroupBy/GroupByQueryPipelineStage.Client.cs index abe295e471..59f3165100 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/GroupBy/GroupByQueryPipelineStage.Client.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/GroupBy/GroupByQueryPipelineStage.Client.cs @@ -17,12 +17,12 @@ internal abstract partial class GroupByQueryPipelineStage : QueryPipelineStageBa private sealed class ClientGroupByQueryPipelineStage : GroupByQueryPipelineStage { public const string ContinuationTokenNotSupportedWithGroupBy = "Continuation token is not supported for queries with GROUP BY. Do not use FeedResponse.ResponseContinuation or remove the GROUP BY from the query."; - private ClientGroupByQueryPipelineStage( IQueryPipelineStage source, CancellationToken cancellationToken, - GroupingTable groupingTable) - : base(source, cancellationToken, groupingTable) + GroupingTable groupingTable, + int pageSize) + : base(source, cancellationToken, groupingTable, pageSize) { } @@ -32,7 +32,8 @@ public static TryCatch MonadicCreate( MonadicCreatePipelineStage monadicCreatePipelineStage, IReadOnlyDictionary groupByAliasToAggregateType, IReadOnlyList orderedAliases, - bool hasSelectValue) + bool hasSelectValue, + int pageSize) { TryCatch tryCreateGroupingTable = GroupingTable.TryCreateFromContinuationToken( groupByAliasToAggregateType, @@ -54,7 +55,8 @@ public static TryCatch MonadicCreate( IQueryPipelineStage stage = new ClientGroupByQueryPipelineStage( tryCreateSource.Result, cancellationToken, - tryCreateGroupingTable.Result); + tryCreateGroupingTable.Result, + pageSize); return TryCatch.FromResult(stage); } @@ -74,8 +76,6 @@ public override async ValueTask MoveNextAsync() double requestCharge = 0.0; long responseLengthInBytes = 0; - int maxPageSize = 0; - while (await this.inputStage.MoveNextAsync()) { this.cancellationToken.ThrowIfCancellationRequested(); @@ -93,14 +93,12 @@ public override async ValueTask MoveNextAsync() requestCharge += sourcePage.RequestCharge; responseLengthInBytes += sourcePage.ResponseLengthInBytes; - maxPageSize = Math.Max(sourcePage.Documents.Count, maxPageSize); - this.AggregateGroupings(sourcePage.Documents); } // Stage 2: // Emit the results from the grouping table page by page - IReadOnlyList results = this.groupingTable.Drain(maxPageSize); + IReadOnlyList results = this.groupingTable.Drain(this.pageSize); if (this.groupingTable.Count == 0) { this.returnedLastPage = true; diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/GroupBy/GroupByQueryPipelineStage.Compute.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/GroupBy/GroupByQueryPipelineStage.Compute.cs index 2aab36e64f..165e841787 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/GroupBy/GroupByQueryPipelineStage.Compute.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/GroupBy/GroupByQueryPipelineStage.Compute.cs @@ -9,6 +9,7 @@ namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.GroupBy using System.Threading; using System.Threading.Tasks; using Microsoft.Azure.Cosmos.CosmosElements; + using Microsoft.Azure.Cosmos.CosmosElements.Numbers; using Microsoft.Azure.Cosmos.Query.Core.Exceptions; using Microsoft.Azure.Cosmos.Query.Core.Metrics; using Microsoft.Azure.Cosmos.Query.Core.Monads; @@ -27,8 +28,9 @@ private sealed class ComputeGroupByQueryPipelineStage : GroupByQueryPipelineStag private ComputeGroupByQueryPipelineStage( IQueryPipelineStage source, CancellationToken cancellationToken, - GroupingTable groupingTable) - : base(source, cancellationToken, groupingTable) + GroupingTable groupingTable, + int pageSize) + : base(source, cancellationToken, groupingTable, pageSize) { } @@ -38,7 +40,8 @@ public static TryCatch MonadicCreate( MonadicCreatePipelineStage monadicCreatePipelineStage, IReadOnlyDictionary groupByAliasToAggregateType, IReadOnlyList orderedAliases, - bool hasSelectValue) + bool hasSelectValue, + int pageSize) { GroupByContinuationToken groupByContinuationToken; if (requestContinuation != null) @@ -88,7 +91,8 @@ public static TryCatch MonadicCreate( new ComputeGroupByQueryPipelineStage( tryCreateSource.Result, cancellationToken, - tryCreateGroupingTable.Result)); + tryCreateGroupingTable.Result, + pageSize)); } public override async ValueTask MoveNextAsync() @@ -138,7 +142,7 @@ public override async ValueTask MoveNextAsync() { // Stage 2: // Emit the results from the grouping table page by page - IReadOnlyList results = this.groupingTable.Drain(10/*FIX THIS*/); + IReadOnlyList results = this.groupingTable.Drain(this.pageSize); QueryState state; if (this.groupingTable.IsDone) diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/GroupBy/GroupByQueryPipelineStage.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/GroupBy/GroupByQueryPipelineStage.cs index 77f2a98dea..71a9ddb33f 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/GroupBy/GroupByQueryPipelineStage.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/GroupBy/GroupByQueryPipelineStage.cs @@ -48,15 +48,18 @@ namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.GroupBy internal abstract partial class GroupByQueryPipelineStage : QueryPipelineStageBase { private readonly GroupingTable groupingTable; - protected bool returnedLastPage; + protected readonly int pageSize; + protected bool returnedLastPage; protected GroupByQueryPipelineStage( IQueryPipelineStage source, CancellationToken cancellationToken, - GroupingTable groupingTable) + GroupingTable groupingTable, + int pageSize) : base(source, cancellationToken) { this.groupingTable = groupingTable ?? throw new ArgumentNullException(nameof(groupingTable)); + this.pageSize = pageSize; } public static TryCatch MonadicCreate( @@ -66,7 +69,8 @@ public static TryCatch MonadicCreate( MonadicCreatePipelineStage monadicCreatePipelineStage, IReadOnlyDictionary groupByAliasToAggregateType, IReadOnlyList orderedAliases, - bool hasSelectValue) => executionEnvironment switch + bool hasSelectValue, + int pageSize) => executionEnvironment switch { ExecutionEnvironment.Client => ClientGroupByQueryPipelineStage.MonadicCreate( continuationToken, @@ -74,14 +78,16 @@ public static TryCatch MonadicCreate( monadicCreatePipelineStage, groupByAliasToAggregateType, orderedAliases, - hasSelectValue), + hasSelectValue, + pageSize), ExecutionEnvironment.Compute => ComputeGroupByQueryPipelineStage.MonadicCreate( continuationToken, cancellationToken, monadicCreatePipelineStage, groupByAliasToAggregateType, orderedAliases, - hasSelectValue), + hasSelectValue, + pageSize), _ => throw new ArgumentException($"Unknown {nameof(ExecutionEnvironment)}: {executionEnvironment}"), }; diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/PipelineFactory.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/PipelineFactory.cs index ae17717c20..bd82a12d0c 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/PipelineFactory.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/PipelineFactory.cs @@ -124,7 +124,8 @@ public static TryCatch MonadicCreate( monadicCreateSourceStage, queryInfo.GroupByAliasToAggregateType, queryInfo.GroupByAliases, - queryInfo.HasSelectValue); + queryInfo.HasSelectValue, + pageSize); } if (queryInfo.HasOffset) diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/Query/QueryTestsBase.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/Query/QueryTestsBase.cs index 0f6a439ab1..118864e459 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/Query/QueryTestsBase.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/Query/QueryTestsBase.cs @@ -234,25 +234,13 @@ private async Task CreatePartitionedContainer( string partitionKey = "/id", Cosmos.IndexingPolicy indexingPolicy = null) { - Container container; - switch (collectionType) + Container container = collectionType switch { - case CollectionTypes.NonPartitioned: - container = await this.CreateNonPartitionedContainerAsync(indexingPolicy); - break; - - case CollectionTypes.SinglePartition: - container = await this.CreateSinglePartitionContainer(partitionKey, indexingPolicy); - break; - - case CollectionTypes.MultiPartition: - container = await this.CreateMultiPartitionContainer(partitionKey, indexingPolicy); - break; - - default: - throw new ArgumentException($"Unknown {nameof(CollectionTypes)} : {collectionType}"); - } - + CollectionTypes.NonPartitioned => await this.CreateNonPartitionedContainerAsync(indexingPolicy), + CollectionTypes.SinglePartition => await this.CreateSinglePartitionContainer(partitionKey, indexingPolicy), + CollectionTypes.MultiPartition => await this.CreateMultiPartitionContainer(partitionKey, indexingPolicy), + _ => throw new ArgumentException($"Unknown {nameof(CollectionTypes)} : {collectionType}"), + }; List insertedDocuments = new List(); foreach (string document in documents) { @@ -319,7 +307,7 @@ private static async Task CleanUp(CosmosClient client) await client.GetDatabase(db.Id).DeleteAsync(); } } - } + } } internal async Task RunWithApiVersion(string apiVersion, Func function) @@ -445,33 +433,21 @@ internal async Task CreateIngestQueryDeleteAsync( continue; } - Task<(Container, IReadOnlyList)> createContainerTask; - switch (collectionType) + Task<(Container, IReadOnlyList)> createContainerTask = collectionType switch { - case CollectionTypes.NonPartitioned: - createContainerTask = this.CreateNonPartitionedContainerAndIngestDocumentsAsync( - documents, - indexingPolicy); - break; - - case CollectionTypes.SinglePartition: - createContainerTask = this.CreateSinglePartitionContainerAndIngestDocumentsAsync( - documents, - partitionKey, - indexingPolicy); - break; - - case CollectionTypes.MultiPartition: - createContainerTask = this.CreateMultiPartitionContainerAndIngestDocumentsAsync( - documents, - partitionKey, - indexingPolicy); - break; - - default: - throw new ArgumentException($"Unknown {nameof(CollectionTypes)} : {collectionType}"); - } - + CollectionTypes.NonPartitioned => this.CreateNonPartitionedContainerAndIngestDocumentsAsync( + documents, + indexingPolicy), + CollectionTypes.SinglePartition => this.CreateSinglePartitionContainerAndIngestDocumentsAsync( + documents, + partitionKey, + indexingPolicy), + CollectionTypes.MultiPartition => this.CreateMultiPartitionContainerAndIngestDocumentsAsync( + documents, + partitionKey, + indexingPolicy), + _ => throw new ArgumentException($"Unknown {nameof(CollectionTypes)} : {collectionType}"), + }; collectionsAndDocuments.Add(await createContainerTask); } @@ -526,48 +502,32 @@ internal async Task CreateIngestQueryDeleteAsync( private static ConnectionMode GetTargetConnectionMode(ConnectionModes connectionMode) { - ConnectionMode targetConnectionMode; - switch (connectionMode) + return connectionMode switch { - case ConnectionModes.Gateway: - targetConnectionMode = ConnectionMode.Gateway; - break; - - case ConnectionModes.Direct: - targetConnectionMode = ConnectionMode.Direct; - break; - - default: - throw new ArgumentException($"Unexpected connection mode: {connectionMode}"); - } - - return targetConnectionMode; + ConnectionModes.Gateway => ConnectionMode.Gateway, + ConnectionModes.Direct => ConnectionMode.Direct, + _ => throw new ArgumentException($"Unexpected connection mode: {connectionMode}"), + }; } internal CosmosClient CreateDefaultCosmosClient(ConnectionMode connectionMode) { - switch (connectionMode) + return connectionMode switch { - case ConnectionMode.Gateway: - return this.GatewayClient; - case ConnectionMode.Direct: - return this.Client; - default: - throw new ArgumentException($"Unexpected connection mode: {connectionMode}"); - } + ConnectionMode.Gateway => this.GatewayClient, + ConnectionMode.Direct => this.Client, + _ => throw new ArgumentException($"Unexpected connection mode: {connectionMode}"), + }; } internal CosmosClient CreateNewCosmosClient(ConnectionMode connectionMode) { - switch (connectionMode) + return connectionMode switch { - case ConnectionMode.Gateway: - return TestCommon.CreateCosmosClient(true); - case ConnectionMode.Direct: - return TestCommon.CreateCosmosClient(false); - default: - throw new ArgumentException($"Unexpected connection mode: {connectionMode}"); - } + ConnectionMode.Gateway => TestCommon.CreateCosmosClient(true), + ConnectionMode.Direct => TestCommon.CreateCosmosClient(false), + _ => throw new ArgumentException($"Unexpected connection mode: {connectionMode}"), + }; } internal static async Task> QueryWithCosmosElementContinuationTokenAsync( @@ -599,7 +559,7 @@ internal static async Task> QueryWithCosmosElementContinuationTokenAsync { Assert.IsTrue( cosmosQueryResponse.Count <= queryRequestOptions.MaxItemCount.Value, - "Max Item Count is not being honored"); + $"Max Item Count is not being honored. Got {cosmosQueryResponse.Count} documents when {queryRequestOptions.MaxItemCount.Value} is the max."); } resultsFromCosmosElementContinuationToken.AddRange(cosmosQueryResponse); @@ -608,14 +568,7 @@ internal static async Task> QueryWithCosmosElementContinuationTokenAsync // There was a bug where resuming from double.NaN lead to an exception, // since we parsed the type assuming it was always a double and not a string. CosmosElement originalContinuationToken = itemQuery.GetCosmosElementContinuationToken(); - if (originalContinuationToken != null) - { - continuationToken = CosmosElement.Parse(originalContinuationToken.ToString()); - } - else - { - continuationToken = null; - } + continuationToken = originalContinuationToken != null ? CosmosElement.Parse(originalContinuationToken.ToString()) : null; } catch (CosmosException cosmosException) when (cosmosException.StatusCode == (HttpStatusCode)429) { @@ -656,7 +609,7 @@ internal static async Task> QueryWithContinuationTokensAsync( { Assert.IsTrue( cosmosQueryResponse.Count <= queryRequestOptions.MaxItemCount.Value, - "Max Item Count is not being honored"); + $"Max Item Count is not being honored. Got {cosmosQueryResponse.Count} when {queryRequestOptions.MaxItemCount.Value} is the max."); } resultsFromContinuationToken.AddRange(cosmosQueryResponse); @@ -708,9 +661,13 @@ internal static async Task> QueryWithoutContinuationTokensAsync( if (queryRequestOptions.MaxItemCount.HasValue) { + if (page.Count > queryRequestOptions.MaxItemCount.Value) + { + Console.WriteLine(); + } Assert.IsTrue( page.Count <= queryRequestOptions.MaxItemCount.Value, - "Max Item Count is not being honored"); + $"Max Item Count is not being honored. Got {page.Count} documents when the max is {queryRequestOptions.MaxItemCount.Value}."); } try diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/Query/SanityQueryTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/Query/SanityQueryTests.cs index bedb3d7277..a7d8f620c3 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/Query/SanityQueryTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/Query/SanityQueryTests.cs @@ -34,7 +34,7 @@ await this.CreateIngestQueryDeleteAsync( inputDocuments, ImplementationAsync); - async Task ImplementationAsync(Container container, IReadOnlyList documents) + static async Task ImplementationAsync(Container container, IReadOnlyList documents) { List queryResults = await QueryTestsBase.RunQueryAsync( container, @@ -60,7 +60,7 @@ await this.CreateIngestQueryDeleteAsync( inputDocuments, ImplementationAsync); - async Task ImplementationAsync(Container container, IReadOnlyList documents) + static async Task ImplementationAsync(Container container, IReadOnlyList documents) { foreach (int maxDegreeOfParallelism in new int[] { 1, 100 }) { @@ -105,7 +105,7 @@ await this.CreateIngestQueryDeleteAsync( inputDocuments, ImplementationAsync); - async Task ImplementationAsync(Container container, IReadOnlyList documents) + static async Task ImplementationAsync(Container container, IReadOnlyList documents) { List weakReferences = await CreateWeakReferenceToFeedIterator(container); GC.Collect(); @@ -222,16 +222,7 @@ async Task ImplementationAsync(Container container, IReadOnlyList { foreach (bool useOrderBy in new bool[] { false, true }) { - string query; - if (useOrderBy) - { - query = "SELECT c._ts, c.id FROM c ORDER BY c._ts"; - } - else - { - query = "SELECT c.id FROM c"; - } - + string query = useOrderBy ? "SELECT c._ts, c.id FROM c ORDER BY c._ts" : "SELECT c.id FROM c"; QueryRequestOptions queryRequestOptions = new QueryRequestOptions { MaxBufferedItemCount = 7000, @@ -284,7 +275,7 @@ await this.CreateIngestQueryDeleteAsync( inputDocuments, ImplementationAsync); - async Task ImplementationAsync(Container container, IReadOnlyList documents) + static async Task ImplementationAsync(Container container, IReadOnlyList documents) { foreach (int maxItemCount in new int[] { 10, 100 }) { @@ -326,7 +317,7 @@ await this.CreateIngestQueryDeleteAsync( inputDocuments, ImplementationAsync); - async Task ImplementationAsync(Container container, IReadOnlyList documents) + static async Task ImplementationAsync(Container container, IReadOnlyList documents) { foreach (int maxItemCount in new int[] { 10, 100 }) { @@ -368,7 +359,7 @@ await this.CreateIngestQueryDeleteAsync( inputDocuments, ImplementationAsync); - async Task ImplementationAsync(Container container, IReadOnlyList documents) + static async Task ImplementationAsync(Container container, IReadOnlyList documents) { ContainerInternal containerCore = (ContainerInlineCore)container; @@ -428,7 +419,7 @@ await this.CreateIngestQueryDeleteAsync( NoDocuments, ImplementationAsync); - async Task ImplementationAsync(Container container, IReadOnlyList documents) + static async Task ImplementationAsync(Container container, IReadOnlyList documents) { QueryRequestOptions feedOptions = new QueryRequestOptions { From c12323ef948e65ba5244e4bc73ca02fc6a572d8d Mon Sep 17 00:00:00 2001 From: Brandon Chong Date: Thu, 1 Oct 2020 18:37:30 -0700 Subject: [PATCH 76/85] removed unused collections --- .../Query/Core/Collections/AsyncCollection.cs | 199 ----------------- .../Core/Collections/PartialReadOnlyList.cs | 76 ------- .../Core/ComparableTask/ComparableTask.cs | 57 ----- .../ComparableTask/ComparableTaskScheduler.cs | 170 --------------- .../Core/ComparableTask/CompletedTask.cs | 34 --- .../Core/ComparableTask/IComparableTask.cs | 15 -- .../src/Query/v2Query/DocumentQueryClient.cs | 3 +- .../Concurrent/AsyncCollectionTest.cs | 200 ------------------ .../Generic/PartialReadOnlyListTest.cs | 106 ---------- .../ComparableTaskSchedulerTests.cs | 169 --------------- .../GroupByQueryPipelineStageTests.cs | 3 +- 11 files changed, 3 insertions(+), 1029 deletions(-) delete mode 100644 Microsoft.Azure.Cosmos/src/Query/Core/Collections/AsyncCollection.cs delete mode 100644 Microsoft.Azure.Cosmos/src/Query/Core/Collections/PartialReadOnlyList.cs delete mode 100644 Microsoft.Azure.Cosmos/src/Query/Core/ComparableTask/ComparableTask.cs delete mode 100644 Microsoft.Azure.Cosmos/src/Query/Core/ComparableTask/ComparableTaskScheduler.cs delete mode 100644 Microsoft.Azure.Cosmos/src/Query/Core/ComparableTask/CompletedTask.cs delete mode 100644 Microsoft.Azure.Cosmos/src/Query/Core/ComparableTask/IComparableTask.cs delete mode 100644 Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Collections/Concurrent/AsyncCollectionTest.cs delete mode 100644 Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Collections/Generic/PartialReadOnlyListTest.cs delete mode 100644 Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/ComparableTaskSchedulerTests.cs diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Collections/AsyncCollection.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Collections/AsyncCollection.cs deleted file mode 100644 index bdb056354d..0000000000 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Collections/AsyncCollection.cs +++ /dev/null @@ -1,199 +0,0 @@ -//------------------------------------------------------------ -// Copyright (c) Microsoft Corporation. All rights reserved. -//------------------------------------------------------------ - -namespace Microsoft.Azure.Cosmos.Query.Core.Collections -{ - using System; - using System.Collections.Concurrent; - using System.Collections.Generic; - using System.Diagnostics; - using System.Threading; - using System.Threading.Tasks; - - /// - /// Provides awaitable and bounding capabilities for thread-safe collections that implement IProducerConsumerCollection<T>. - /// - internal sealed class AsyncCollection - { - private delegate bool TryPeekDelegate(out T item); - private readonly IProducerConsumerCollection collection; - private readonly int boundingCapacity; - private readonly SemaphoreSlim notFull; - private readonly SemaphoreSlim notEmpty; - private readonly TryPeekDelegate tryPeekDelegate; - - public AsyncCollection() - : this(new ConcurrentQueue(), int.MaxValue) - { - } - - public AsyncCollection(int boundingCapacity) - : this(new ConcurrentQueue(), boundingCapacity) - { - } - - public AsyncCollection(IProducerConsumerCollection collection) - : this(collection, int.MaxValue) - { - } - - public AsyncCollection(IProducerConsumerCollection collection, int boundingCapacity) - { - if (collection == null) - { - throw new ArgumentNullException("collection"); - } - - if (boundingCapacity < 1) - { - throw new ArgumentOutOfRangeException("boundedCapacity is not a positive value."); - } - - int count = collection.Count; - - if (boundingCapacity < count) - { - throw new ArgumentOutOfRangeException("boundedCapacity is less than the size of collection."); - } - - this.collection = collection; - this.boundingCapacity = boundingCapacity; - this.notFull = this.IsUnbounded ? null : new SemaphoreSlim(boundingCapacity - count, boundingCapacity); - this.notEmpty = new SemaphoreSlim(count); - if (collection is ConcurrentQueue concurrentQueue) - { - this.tryPeekDelegate = concurrentQueue.TryPeek; - return; - } - - if (collection is PriorityQueue priorityQueue) - { - this.tryPeekDelegate = priorityQueue.TryPeek; - return; - } - - throw new NotSupportedException($"The IProducerConsumerCollection type of {typeof(T)} is not supported."); - } - - public int Count => this.collection.Count; - - public bool IsUnbounded => this.boundingCapacity >= int.MaxValue; - - public async Task AddAsync(T item, CancellationToken token = default) - { - if (!this.IsUnbounded) - { - await this.notFull.WaitAsync(token); - } - - if (this.collection.TryAdd(item)) - { - this.notEmpty.Release(); - } - } - - public async Task AddRangeAsync(IEnumerable items, CancellationToken token = default) - { - if (!this.IsUnbounded) - { - foreach (T item in items) - { - await this.AddAsync(item); - } - } - else - { - int count = 0; - foreach (T item in items) - { - if (this.collection.TryAdd(item)) - { - ++count; - } - } - - if (count > 0) - { - this.notEmpty.Release(count); - } - } - } - - public async Task TakeAsync(CancellationToken token = default) - { - await this.notEmpty.WaitAsync(token); - if (this.collection.TryTake(out T item)) - { - if (!this.IsUnbounded) - { - this.notFull.Release(); - } - } - - return item; - } - - public async Task PeekAsync(CancellationToken token = default) - { - if (this.tryPeekDelegate == null) - { - throw new NotImplementedException(); - } - - await this.notEmpty.WaitAsync(token); - // Do nothing if tryPeekFunc returns false - this.tryPeekDelegate(out T item); - this.notEmpty.Release(); - - return item; - } - - public bool TryPeek(out T item) - { - if (this.tryPeekDelegate == null) - { - throw new NotImplementedException(); - } - - return this.tryPeekDelegate(out item); - } - - public async Task> DrainAsync( - int maxElements = int.MaxValue, - TimeSpan timeout = default, - Func callback = null, - CancellationToken token = default) - { - if (maxElements < 1) - { - throw new ArgumentOutOfRangeException("maxElements is not a positive value."); - } - - List elements = new List(); - - Stopwatch stopWatch = Stopwatch.StartNew(); - while (elements.Count < maxElements && await this.notEmpty.WaitAsync(timeout, token)) - { - if (this.collection.TryTake(out T item) && (callback == null || callback(item))) - { - elements.Add(item); - } - else - { - break; - } - - timeout.Subtract(TimeSpan.FromTicks(Math.Min(stopWatch.ElapsedTicks, timeout.Ticks))); - stopWatch.Restart(); - } - - if (!this.IsUnbounded && elements.Count > 0) - { - this.notFull.Release(elements.Count); - } - - return elements; - } - } -} \ No newline at end of file diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Collections/PartialReadOnlyList.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Collections/PartialReadOnlyList.cs deleted file mode 100644 index fc27c236f8..0000000000 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Collections/PartialReadOnlyList.cs +++ /dev/null @@ -1,76 +0,0 @@ -//------------------------------------------------------------ -// Copyright (c) Microsoft Corporation. All rights reserved. -//------------------------------------------------------------ - -namespace Microsoft.Azure.Cosmos.Query.Core.Collections -{ - using System; - using System.Collections; - using System.Collections.Generic; - - internal sealed class PartialReadOnlyList : IReadOnlyList - { - private readonly IReadOnlyList list; - private readonly int startIndex; - - public PartialReadOnlyList( - IReadOnlyList list, - int count) - : this(list, 0, count) - { - } - - public PartialReadOnlyList( - IReadOnlyList list, - int startIndex, - int count) - { - if (list == null) - { - throw new ArgumentNullException("list"); - } - - if (count <= 0 || count > list.Count) - { - throw new ArgumentOutOfRangeException("count"); - } - - if (startIndex < 0 || (startIndex + count) > list.Count) - { - throw new ArgumentOutOfRangeException("startIndex"); - } - - this.list = list; - this.startIndex = startIndex; - this.Count = count; - } - - public T this[int index] - { - get - { - if (index < 0 || index >= this.Count) - { - throw new ArgumentOutOfRangeException("index"); - } - - return this.list[checked(this.startIndex + index)]; - } - } - - public int Count { get; } - - public IEnumerator GetEnumerator() - { - for (int i = 0; i < this.Count; ++i) - { - yield return this.list[i + this.startIndex]; - } - } - - IEnumerator IEnumerable.GetEnumerator() - { - return this.GetEnumerator(); - } - } -} diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ComparableTask/ComparableTask.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ComparableTask/ComparableTask.cs deleted file mode 100644 index 9c3ed13f9c..0000000000 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ComparableTask/ComparableTask.cs +++ /dev/null @@ -1,57 +0,0 @@ -//------------------------------------------------------------ -// Copyright (c) Microsoft Corporation. All rights reserved. -//------------------------------------------------------------ - -namespace Microsoft.Azure.Cosmos.Query.Core.ComparableTask -{ - using System; - using System.Threading; - using System.Threading.Tasks; - - internal abstract class ComparableTask : IComparableTask - { - protected readonly int schedulePriority; - - protected ComparableTask(int schedulePriority) - { - this.schedulePriority = schedulePriority; - } - - public abstract Task StartAsync(CancellationToken token); - - public virtual int CompareTo(IComparableTask other) - { - if (other == null) - { - throw new ArgumentNullException(nameof(other)); - } - - return this.CompareToByPriority(other as ComparableTask); - } - - public override bool Equals(object obj) - { - return this.Equals(obj as IComparableTask); - } - - public abstract bool Equals(IComparableTask other); - - public abstract override int GetHashCode(); - - protected int CompareToByPriority(ComparableTask other) - { - if (other == null) - { - // This task is not ComparableTask, assume it has a higher priority. - return 1; - } - - if (object.ReferenceEquals(this, other)) - { - return 0; - } - - return this.schedulePriority.CompareTo(other.schedulePriority); - } - } -} \ No newline at end of file diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ComparableTask/ComparableTaskScheduler.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ComparableTask/ComparableTaskScheduler.cs deleted file mode 100644 index e2b2098139..0000000000 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ComparableTask/ComparableTaskScheduler.cs +++ /dev/null @@ -1,170 +0,0 @@ -//------------------------------------------------------------ -// Copyright (c) Microsoft Corporation. All rights reserved. -//------------------------------------------------------------ - -namespace Microsoft.Azure.Cosmos.Query.Core.ComparableTask -{ - using System; - using System.Collections.Concurrent; - using System.Collections.Generic; - using System.Linq; - using System.Threading; - using System.Threading.Tasks; - using Microsoft.Azure.Cosmos.Query.Core.Collections; - using Microsoft.Azure.Documents; - - internal sealed class ComparableTaskScheduler : IDisposable - { - private const int MinimumBatchSize = 1; - private readonly AsyncCollection taskQueue; - private readonly ConcurrentDictionary delayedTasks; - private readonly CancellationTokenSource tokenSource; - private readonly SemaphoreSlim canRunTaskSemaphoreSlim; - private readonly Task schedulerTask; - private volatile bool isStopped; - - public ComparableTaskScheduler() - : this(Environment.ProcessorCount) - { - } - - public ComparableTaskScheduler(int maximumConcurrencyLevel) - : this(Enumerable.Empty(), maximumConcurrencyLevel) - { - } - - public ComparableTaskScheduler(IEnumerable tasks, int maximumConcurrencyLevel) - { - this.taskQueue = new AsyncCollection(new PriorityQueue(tasks, true)); - this.delayedTasks = new ConcurrentDictionary(); - this.MaximumConcurrencyLevel = maximumConcurrencyLevel; - this.tokenSource = new CancellationTokenSource(); - this.canRunTaskSemaphoreSlim = new SemaphoreSlim(maximumConcurrencyLevel); - this.schedulerTask = this.ScheduleAsync(); - } - - public int MaximumConcurrencyLevel { get; private set; } - - public int CurrentRunningTaskCount => this.MaximumConcurrencyLevel - Math.Max(0, this.canRunTaskSemaphoreSlim.CurrentCount); - - public bool IsStopped => this.isStopped; - - private CancellationToken CancellationToken => this.tokenSource.Token; - - public void IncreaseMaximumConcurrencyLevel(int delta) - { - if (delta <= 0) - { - throw new ArgumentOutOfRangeException("delta must be a positive number."); - } - - this.canRunTaskSemaphoreSlim.Release(delta); - this.MaximumConcurrencyLevel += delta; - } - - public void Dispose() - { - this.Stop(); - - this.canRunTaskSemaphoreSlim.Dispose(); - this.tokenSource.Dispose(); - } - - public void Stop() - { - this.isStopped = true; - this.tokenSource.Cancel(); - this.delayedTasks.Clear(); - } - - public bool TryQueueTask(IComparableTask comparableTask, TimeSpan delay = default) - { - if (comparableTask == null) - { - throw new ArgumentNullException("task"); - } - - if (this.isStopped) - { - return false; - } - - Task newTask = new Task(() => this.QueueDelayedTaskAsync(comparableTask, delay), this.CancellationToken); - - if (this.delayedTasks.TryAdd(comparableTask, newTask)) - { - newTask.Start(); - return true; - } - - return false; - } - - private async Task QueueDelayedTaskAsync(IComparableTask comparableTask, TimeSpan delay) - { - if (this.delayedTasks.TryRemove(comparableTask, out Task task) && !task.IsCanceled) - { - if (delay > default(TimeSpan)) - { - await Task.Delay(delay, this.CancellationToken); - } - - if (this.taskQueue.TryPeek(out IComparableTask firstComparableTask) && (comparableTask.CompareTo(firstComparableTask) <= 0)) - { - await this.ExecuteComparableTaskAsync(comparableTask); - } - else - { - await this.taskQueue.AddAsync(comparableTask, this.CancellationToken); - } - } - } - - private async Task ScheduleAsync() - { - while (!this.isStopped) - { - await this.ExecuteComparableTaskAsync(await this.taskQueue.TakeAsync(this.CancellationToken)); - } - } - - private async Task ExecuteComparableTaskAsync(IComparableTask comparableTask) - { - await this.canRunTaskSemaphoreSlim.WaitAsync(this.CancellationToken); - -#pragma warning disable 4014 - // Schedule execution on current .NET task scheduler. - // Compute gateway uses custom task scheduler to track tenant resource utilization. - // Task.Run() switches to default task scheduler for entire sub-tree of tasks making compute gateway incapable of tracking resource usage accurately. - // Task.Factory.StartNew() allows specifying task scheduler to use. - Task.Factory - .StartNewOnCurrentTaskSchedulerAsync( - function: () => comparableTask - .StartAsync(this.CancellationToken) - .ContinueWith((antecendent) => - { - // Observing the exception. - Exception exception = antecendent.Exception; - Extensions.TraceException(exception); - - // Semaphore.Release can also throw an exception. - try - { - this.canRunTaskSemaphoreSlim.Release(); - } - catch (Exception releaseException) - { - Extensions.TraceException(releaseException); - } - }, TaskScheduler.Current), - cancellationToken: this.CancellationToken) - .ContinueWith((antecendent) => - { - // StartNew can have a task cancelled exception - Exception exception = antecendent.Exception; - Extensions.TraceException(exception); - }); -#pragma warning restore 4014 - } - } -} \ No newline at end of file diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ComparableTask/CompletedTask.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ComparableTask/CompletedTask.cs deleted file mode 100644 index 3381899215..0000000000 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ComparableTask/CompletedTask.cs +++ /dev/null @@ -1,34 +0,0 @@ -//------------------------------------------------------------ -// Copyright (c) Microsoft Corporation. All rights reserved. -//------------------------------------------------------------ -namespace Microsoft.Azure.Cosmos.Query.Core.ComparableTask -{ - using System; - using System.Threading.Tasks; - - internal static class CompletedTask - { - private static Task instance; - - public static Task Instance - { - get - { - if (CompletedTask.instance == null) - { - TaskCompletionSource completionSource = new TaskCompletionSource(); - completionSource.SetResult(true); - CompletedTask.instance = completionSource.Task; - } - return CompletedTask.instance; - } - } - - public static Task SetExceptionAsync(Exception exception) - { - TaskCompletionSource completionSource = new TaskCompletionSource(); - completionSource.SetException(exception); - return completionSource.Task; - } - } -} diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ComparableTask/IComparableTask.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ComparableTask/IComparableTask.cs deleted file mode 100644 index f49fcd93c1..0000000000 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ComparableTask/IComparableTask.cs +++ /dev/null @@ -1,15 +0,0 @@ -//------------------------------------------------------------ -// Copyright (c) Microsoft Corporation. All rights reserved. -//------------------------------------------------------------ - -namespace Microsoft.Azure.Cosmos.Query.Core.ComparableTask -{ - using System; - using System.Threading; - using System.Threading.Tasks; - - internal interface IComparableTask : IComparable, IEquatable - { - Task StartAsync(CancellationToken token); - } -} \ No newline at end of file diff --git a/Microsoft.Azure.Cosmos/src/Query/v2Query/DocumentQueryClient.cs b/Microsoft.Azure.Cosmos/src/Query/v2Query/DocumentQueryClient.cs index 22d7af54fc..778f322adc 100644 --- a/Microsoft.Azure.Cosmos/src/Query/v2Query/DocumentQueryClient.cs +++ b/Microsoft.Azure.Cosmos/src/Query/v2Query/DocumentQueryClient.cs @@ -9,7 +9,6 @@ namespace Microsoft.Azure.Cosmos.Query using System.Threading; using System.Threading.Tasks; using Microsoft.Azure.Cosmos.Common; - using Microsoft.Azure.Cosmos.Query.Core.ComparableTask; using Microsoft.Azure.Cosmos.Query.Core.QueryPlan; using Microsoft.Azure.Cosmos.Routing; using Microsoft.Azure.Documents; @@ -113,7 +112,7 @@ public async Task GetDefaultConsistencyLevelAsync() public Task EnsureValidOverwriteAsync(ConsistencyLevel requestedConsistencyLevel) { this.innerClient.EnsureValidOverwrite(requestedConsistencyLevel); - return CompletedTask.Instance; + return Task.CompletedTask; } public Task GetPartitionKeyRangeCacheAsync() diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Collections/Concurrent/AsyncCollectionTest.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Collections/Concurrent/AsyncCollectionTest.cs deleted file mode 100644 index 4ed6138e1b..0000000000 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Collections/Concurrent/AsyncCollectionTest.cs +++ /dev/null @@ -1,200 +0,0 @@ -//------------------------------------------------------------ -// Copyright (c) Microsoft Corporation. All rights reserved. -//------------------------------------------------------------ - -namespace Microsoft.Azure.Cosmos.Collections.Generic -{ - using System; - using System.Collections.Concurrent; - using System.Collections.Generic; - using System.Threading.Tasks; - using Microsoft.Azure.Cosmos.Query.Core.Collections; - using Microsoft.VisualStudio.TestTools.UnitTesting; - - /// - /// Tests for "PriorityQueue" class. - /// - [TestClass] - public class AsyncCollectionTest - { - private const int DelayInMilliSeconds = 100; - /// - /// Test for "AddAsync". - /// - [TestMethod] - public async Task TestAddAsync() - { - ConcurrentQueue queue = new ConcurrentQueue(); - queue.Enqueue(0); - AsyncCollection asyncCollection = new AsyncCollection(queue, 1); - int item = 1; - Task task = asyncCollection.AddAsync(item); - await Task.Delay(DelayInMilliSeconds); - Assert.AreEqual(false, task.IsCompleted); - Assert.AreEqual(0, await asyncCollection.TakeAsync()); - await task; - Assert.AreEqual(item, await asyncCollection.TakeAsync()); - } - - /// - /// Test for "AddRangeAsync". - /// - [TestMethod] - public async Task TestAddRangeAsync() - { - AsyncCollection asyncCollection = new AsyncCollection(2); - Task task = asyncCollection.AddRangeAsync(new[] { 0, 1, 2 }); - await Task.Delay(DelayInMilliSeconds); - Assert.AreEqual(false, task.IsCompleted); - Assert.AreEqual(0, await asyncCollection.TakeAsync()); - await task; - Assert.AreEqual(1, await asyncCollection.TakeAsync()); - } - - /// - /// Test for "TakeAsync". - /// - [TestMethod] - public async Task TestTakeAsync() - { - AsyncCollection asyncCollection = new AsyncCollection(); - Task task = asyncCollection.TakeAsync(); - await Task.Delay(DelayInMilliSeconds); - Assert.AreEqual(false, task.IsCompleted); - int item = 1; - await asyncCollection.AddAsync(item); - Assert.AreEqual(item, await task); - } - - /// - /// Test for "PeekAsync". - /// - [TestMethod] - public async Task TestPeekAsync() - { - AsyncCollection asyncCollection = new AsyncCollection(); - Task task = asyncCollection.PeekAsync(); - await Task.Delay(DelayInMilliSeconds); - Assert.AreEqual(false, task.IsCompleted); - int item = 1; - await asyncCollection.AddAsync(item); - Assert.AreEqual(item, await task); - } - - /// - /// Simple Test for "AsyncCollection". - /// - [TestMethod] - public async Task SimpleTest() - { - int size = 100; - AsyncCollection asyncCollection = new AsyncCollection(); - Assert.AreEqual(0, asyncCollection.Count); - List list = new List(); - for (int i = 0; i < size; ++i) - { - list.Add(i); - } - - foreach (bool addOneByOne in new[] { true, false }) - { - foreach (bool takeOneByOne in new[] { true, false }) - { - if (addOneByOne) - { - foreach (int i in list) - { - await asyncCollection.AddAsync(i); - Assert.AreEqual(i + 1, asyncCollection.Count); - } - } - else - { - await asyncCollection.AddRangeAsync(list); - Assert.AreEqual(size, asyncCollection.Count); - } - - if (takeOneByOne) - { - foreach (int i in list) - { - Assert.AreEqual(i, await asyncCollection.PeekAsync()); - Assert.AreEqual(size - i, asyncCollection.Count); - Assert.AreEqual(i, await asyncCollection.TakeAsync()); - Assert.AreEqual(size - i - 1, asyncCollection.Count); - } - } - else - { - Assert.AreEqual( - string.Join(",", list), - string.Join(",", await asyncCollection.DrainAsync())); - } - } - } - } - - /// - /// Test for AsyncCollection.PeekAsync NotImplementedException - /// - [TestMethod] - [ExpectedException(typeof(NotImplementedException))] - public async Task TestNotImplmenetedPeekAsync() - { - await new AsyncCollection(new TestProducerConsumerCollection()).PeekAsync(); - } - - private class TestProducerConsumerCollection : IProducerConsumerCollection - { - public void CopyTo(T[] array, int index) - { - throw new NotImplementedException(); - } - - public T[] ToArray() - { - throw new NotImplementedException(); - } - - public bool TryAdd(T item) - { - throw new NotImplementedException(); - } - - public bool TryTake(out T item) - { - throw new NotImplementedException(); - } - - public IEnumerator GetEnumerator() - { - throw new NotImplementedException(); - } - - System.Collections.IEnumerator System.Collections.IEnumerable.GetEnumerator() - { - throw new NotImplementedException(); - } - - public void CopyTo(Array array, int index) - { - throw new NotImplementedException(); - } - - public int Count - { - get { throw new NotImplementedException(); } - } - - public bool IsSynchronized - { - get { throw new NotImplementedException(); } - } - - public object SyncRoot - { - get { throw new NotImplementedException(); } - } - } - } -} diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Collections/Generic/PartialReadOnlyListTest.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Collections/Generic/PartialReadOnlyListTest.cs deleted file mode 100644 index 35ed7ecdfb..0000000000 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Collections/Generic/PartialReadOnlyListTest.cs +++ /dev/null @@ -1,106 +0,0 @@ -//------------------------------------------------------------ -// Copyright (c) Microsoft Corporation. All rights reserved. -//------------------------------------------------------------ - -namespace Microsoft.Azure.Cosmos.Collections.Generic -{ - using System; - using System.Linq; - using System.Collections.Generic; - using Microsoft.VisualStudio.TestTools.UnitTesting; - using Microsoft.Azure.Cosmos.Query.Core.Collections; - - /// - /// Tests for "PartialReadOnlyList" class. - /// - [TestClass] - public class PartialReadOnlyListTest - { - /// - /// Tests for "PartialReadOnlyList" constructors. - /// - [TestMethod] - public void TestConstructors() - { - // Constructor Test 1 - Assert.AreEqual(1, new PartialReadOnlyList(new[] { 1 }.ToList(), 1).Count); - - // Constructor Test 2 - Assert.AreEqual(1, new PartialReadOnlyList(new[] { 1, 2 }.ToList(), 0, 1).Count); - - // Constructor Test 3 - try - { - new PartialReadOnlyList(null, 1); - Assert.Fail("ArgumentNullException should have been thrown."); - } - catch (ArgumentNullException) - { } - - // Constructor Test 4 - try - { - new PartialReadOnlyList(new[] { 1 }.ToList(), 2); - Assert.Fail("ArgumentOutOfRangeException should have been thrown."); - } - catch (ArgumentOutOfRangeException) - { } - - // Constructor Test 5 - try - { - new PartialReadOnlyList(new[] { 1 }.ToList(), -1, 1); - Assert.Fail("ArgumentOutOfRangeException should have been thrown."); - } - catch (ArgumentOutOfRangeException) - { } - - // Constructor Test 6 - try - { - new PartialReadOnlyList(new[] { 1 }.ToList(), -1, 1); - Assert.Fail("ArgumentOutOfRangeException should have been thrown."); - } - catch (ArgumentOutOfRangeException) - { } - - // Constructor Test 7 - try - { - new PartialReadOnlyList(new[] { 1 }.ToList(), 1, 1); - Assert.Fail("ArgumentOutOfRangeException should have been thrown."); - } - catch (ArgumentOutOfRangeException) - { } - } - - - /// - /// Simple test for "PartialReadOnlyList". - /// - [TestMethod] - public void SimpleTest() - { - PartialReadOnlyList list = new PartialReadOnlyList(new[] { 1, 2, 3, 4, 5 }.ToList(), 3, 2); - Assert.AreEqual(4, list[0]); - Assert.AreEqual(5, list[1]); - Assert.AreEqual(string.Join(", ", 4, 5), string.Join(", ", list)); - - try - { - int i = list[-1]; - Assert.Fail("ArgumentOutOfRangeException should have been thrown."); - } - catch (ArgumentOutOfRangeException) - { } - - try - { - int i = list[2]; - Assert.Fail("ArgumentOutOfRangeException should have been thrown."); - } - catch (ArgumentOutOfRangeException) - { } - } - } -} \ No newline at end of file diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/ComparableTaskSchedulerTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/ComparableTaskSchedulerTests.cs deleted file mode 100644 index 14ea4da804..0000000000 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/ComparableTaskSchedulerTests.cs +++ /dev/null @@ -1,169 +0,0 @@ -//------------------------------------------------------------ -// Copyright (c) Microsoft Corporation. All rights reserved. -//------------------------------------------------------------ - -namespace Microsoft.Azure.Cosmos.Test -{ - using System; - using System.Collections.Generic; - using System.Linq; - using System.Threading; - using System.Threading.Tasks; - using Microsoft.Azure.Cosmos.Query.Core.ComparableTask; - using Microsoft.VisualStudio.TestTools.UnitTesting; - - - [TestClass] - public class ComparableTaskSchedulerTests - { - [TestMethod] - public async Task SimpleTestAsync() - { - foreach (bool useConstructorToAddTasks in new[] { true, false }) - { - List tasks = new List(); - int maximumConcurrencyLevel = 10; - - for (int i = 0; i < maximumConcurrencyLevel; ++i) - { - tasks.Add(new Task(() => { })); - } - - await Task.Delay(1); - - foreach (Task task in tasks) - { - Assert.AreEqual(false, task.IsCompleted); - } - - ComparableTaskScheduler scheduler; - if (useConstructorToAddTasks) - { - scheduler = new ComparableTaskScheduler( - tasks.Select(task => new TestComparableTask(tasks.IndexOf(task), task)), - maximumConcurrencyLevel); - } - else - { - scheduler = new ComparableTaskScheduler(maximumConcurrencyLevel); - for (int i = 0; i < maximumConcurrencyLevel; ++i) - { - Assert.AreEqual(true, scheduler.TryQueueTask(new TestComparableTask(i, tasks[i]))); - } - } - - bool completionStatus = Task.WaitAll(tasks.ToArray(), TimeSpan.FromSeconds(10)); - Assert.IsTrue(completionStatus); - - foreach (Task task in tasks) - { - Assert.AreEqual(true, task.IsCompleted, $"Is overloaded constructor {useConstructorToAddTasks} and status {task.Status.ToString()}"); - } - } - } - - [TestMethod] - public void TestMaximumConcurrencyLevel() - { - ComparableTaskScheduler scheduler = new ComparableTaskScheduler(10); - Assert.AreEqual(10, scheduler.MaximumConcurrencyLevel); - - scheduler = new ComparableTaskScheduler(); - Assert.AreEqual(Environment.ProcessorCount, scheduler.MaximumConcurrencyLevel); - - scheduler.IncreaseMaximumConcurrencyLevel(1); - Assert.AreEqual(Environment.ProcessorCount + 1, scheduler.MaximumConcurrencyLevel); - - try - { - scheduler.IncreaseMaximumConcurrencyLevel(-1); - Assert.Fail("Expect ArgumentOutOfRangeException"); - } - catch (ArgumentOutOfRangeException) - { - } - } - - [TestMethod] - public void TestStop() - { - ComparableTaskScheduler scheduler = new ComparableTaskScheduler(); - Assert.AreEqual(true, scheduler.TryQueueTask(new TestComparableTask(0, Task.FromResult(false)))); - scheduler.Stop(); - Assert.AreEqual(false, scheduler.TryQueueTask(new TestComparableTask(0, Task.FromResult(false)))); - } - - - [TestMethod] - public async Task TestDelayedQueueTaskAsync() - { - ComparableTaskScheduler scheduler = new ComparableTaskScheduler(); - - Task task = new Task(() => - { - Assert.AreEqual(1, scheduler.CurrentRunningTaskCount); - }); - - Task delayedTask = new Task(() => - { - Assert.AreEqual(1, scheduler.CurrentRunningTaskCount); - }); - - Assert.AreEqual( - true, - scheduler.TryQueueTask(new TestComparableTask(schedulePriority: 0, delayedTask), TimeSpan.FromMilliseconds(200))); - Assert.AreEqual( - false, - scheduler.TryQueueTask(new TestComparableTask(schedulePriority: 0, delayedTask), TimeSpan.FromMilliseconds(200))); - Assert.AreEqual( - false, - scheduler.TryQueueTask(new TestComparableTask(schedulePriority: 0, task))); - Assert.AreEqual( - true, - scheduler.TryQueueTask(new TestComparableTask(schedulePriority: 1, task))); - - await Task.Delay(150); - - Assert.AreEqual(true, task.IsCompleted); - Assert.AreEqual(false, delayedTask.IsCompleted); - Assert.AreEqual(0, scheduler.CurrentRunningTaskCount); - - await Task.Delay(400); - - Assert.AreEqual(true, delayedTask.IsCompleted); - } - - private sealed class TestComparableTask : ComparableTask - { - private readonly Task task; - public TestComparableTask(int schedulePriority, Task task) : - base(schedulePriority) - { - this.task = task; - } - - public override Task StartAsync(CancellationToken token) - { - try - { - this.task.Start(); - } - catch (InvalidOperationException) - { - } - - return this.task; - } - - public override int GetHashCode() - { - return this.schedulePriority; - } - - public override bool Equals(IComparableTask other) - { - return this.CompareTo(other) == 0; - } - } - } -} \ No newline at end of file diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/GroupByQueryPipelineStageTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/GroupByQueryPipelineStageTests.cs index 2b5ab5902c..10d3d200da 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/GroupByQueryPipelineStageTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/GroupByQueryPipelineStageTests.cs @@ -60,7 +60,8 @@ private static async Task> CreateAndDrainAsync( monadicCreatePipelineStage: (CosmosElement continuationToken, CancellationToken cancellationToken) => TryCatch.FromResult(source), groupByAliasToAggregateType: groupByAliasToAggregateType, orderedAliases: orderedAliases, - hasSelectValue: hasSelectValue); + hasSelectValue: hasSelectValue, + pageSize: int.MaxValue); Assert.IsTrue(tryCreateGroupByStage.Succeeded); IQueryPipelineStage groupByQueryPipelineStage = tryCreateGroupByStage.Result; From 645f3e40aa4906be2c04d47253796a06c6e8be68 Mon Sep 17 00:00:00 2001 From: Brandon Chong Date: Sat, 3 Oct 2020 00:32:25 -0700 Subject: [PATCH 77/85] opting for simple resource id: --- .../src/Pagination/DocumentContainer.cs | 4 +- .../src/Pagination/DocumentContainerState.cs | 4 +- .../src/Pagination/IDocumentContainer.cs | 3 +- .../Pagination/IMonadicDocumentContainer.cs | 7 +- .../NetworkAttachedDocumentContainer.cs | 4 +- .../src/Pagination/Record.cs | 24 +- .../src/Pagination/ResourceIdentifier.cs | 560 ++++++++++++++++++ ...OrderByCrossPartitionQueryPipelineStage.cs | 157 +++-- ...yQueryPartitionRangePageAsyncEnumerator.cs | 2 +- ...arallelCrossPartitionQueryPipelineStage.cs | 4 + ...cumentContainerPartitionRangeEnumerator.cs | 2 +- .../Pagination/DocumentContainerTests.cs | 4 +- .../Pagination/FlakyDocumentContainer.cs | 3 +- .../Pagination/InMemoryContainer.cs | 95 ++- .../Query/Pipeline/FullPipelineTests.cs | 44 +- .../QueryPartitionRangePageEnumeratorTests.cs | 2 +- 16 files changed, 816 insertions(+), 103 deletions(-) create mode 100644 Microsoft.Azure.Cosmos/src/Pagination/ResourceIdentifier.cs diff --git a/Microsoft.Azure.Cosmos/src/Pagination/DocumentContainer.cs b/Microsoft.Azure.Cosmos/src/Pagination/DocumentContainer.cs index fc3ee75fe2..65805d33fd 100644 --- a/Microsoft.Azure.Cosmos/src/Pagination/DocumentContainer.cs +++ b/Microsoft.Azure.Cosmos/src/Pagination/DocumentContainer.cs @@ -84,7 +84,7 @@ public Task ReadItemAsync( public Task> MonadicReadFeedAsync( int partitionKeyRangeId, - ResourceId resourceIdentifer, + ResourceIdentifier resourceIdentifer, int pageSize, CancellationToken cancellationToken) => this.monadicDocumentContainer.MonadicReadFeedAsync( partitionKeyRangeId, @@ -94,7 +94,7 @@ public Task> MonadicReadFeedAsync( public Task ReadFeedAsync( int partitionKeyRangeId, - ResourceId resourceIdentifier, + ResourceIdentifier resourceIdentifier, int pageSize, CancellationToken cancellationToken) => TryCatch.UnsafeGetResultAsync( this.MonadicReadFeedAsync( diff --git a/Microsoft.Azure.Cosmos/src/Pagination/DocumentContainerState.cs b/Microsoft.Azure.Cosmos/src/Pagination/DocumentContainerState.cs index dd187d638d..5c97e4b62b 100644 --- a/Microsoft.Azure.Cosmos/src/Pagination/DocumentContainerState.cs +++ b/Microsoft.Azure.Cosmos/src/Pagination/DocumentContainerState.cs @@ -8,11 +8,11 @@ namespace Microsoft.Azure.Cosmos.Pagination internal sealed class DocumentContainerState : State { - public DocumentContainerState(ResourceId resourceIdentifier) + public DocumentContainerState(ResourceIdentifier resourceIdentifier) { this.ResourceIdentifer = resourceIdentifier; } - public ResourceId ResourceIdentifer { get; } + public ResourceIdentifier ResourceIdentifer { get; } } } diff --git a/Microsoft.Azure.Cosmos/src/Pagination/IDocumentContainer.cs b/Microsoft.Azure.Cosmos/src/Pagination/IDocumentContainer.cs index 2e63b9f30d..ee4fd24657 100644 --- a/Microsoft.Azure.Cosmos/src/Pagination/IDocumentContainer.cs +++ b/Microsoft.Azure.Cosmos/src/Pagination/IDocumentContainer.cs @@ -8,7 +8,6 @@ namespace Microsoft.Azure.Cosmos.Pagination using System.Threading.Tasks; using Microsoft.Azure.Cosmos.CosmosElements; using Microsoft.Azure.Cosmos.Query.Core.Pipeline.CrossPartition; - using Microsoft.Azure.Documents; internal interface IDocumentContainer : IMonadicDocumentContainer, IFeedRangeProvider, IQueryDataSource { @@ -23,7 +22,7 @@ Task ReadItemAsync( Task ReadFeedAsync( int partitionKeyRangeId, - ResourceId resourceIdentifier, + ResourceIdentifier resourceIdentifier, int pageSize, CancellationToken cancellationToken); diff --git a/Microsoft.Azure.Cosmos/src/Pagination/IMonadicDocumentContainer.cs b/Microsoft.Azure.Cosmos/src/Pagination/IMonadicDocumentContainer.cs index 41f8885386..ce559b8eb1 100644 --- a/Microsoft.Azure.Cosmos/src/Pagination/IMonadicDocumentContainer.cs +++ b/Microsoft.Azure.Cosmos/src/Pagination/IMonadicDocumentContainer.cs @@ -4,16 +4,11 @@ namespace Microsoft.Azure.Cosmos.Pagination { - using System; - using System.Collections.Generic; using System.Threading; using System.Threading.Tasks; using Microsoft.Azure.Cosmos.CosmosElements; - using Microsoft.Azure.Cosmos.Query.Core; using Microsoft.Azure.Cosmos.Query.Core.Monads; - using Microsoft.Azure.Cosmos.Query.Core.Pipeline; using Microsoft.Azure.Cosmos.Query.Core.Pipeline.CrossPartition; - using Microsoft.Azure.Documents; internal interface IMonadicDocumentContainer : IMonadicFeedRangeProvider, IMonadicQueryDataSource { @@ -28,7 +23,7 @@ Task> MonadicReadItemAsync( Task> MonadicReadFeedAsync( int partitionKeyRangeId, - ResourceId resourceIdentifer, + ResourceIdentifier resourceIdentifer, int pageSize, CancellationToken cancellationToken); diff --git a/Microsoft.Azure.Cosmos/src/Pagination/NetworkAttachedDocumentContainer.cs b/Microsoft.Azure.Cosmos/src/Pagination/NetworkAttachedDocumentContainer.cs index 9b8b6019a4..07ce6a4886 100644 --- a/Microsoft.Azure.Cosmos/src/Pagination/NetworkAttachedDocumentContainer.cs +++ b/Microsoft.Azure.Cosmos/src/Pagination/NetworkAttachedDocumentContainer.cs @@ -63,7 +63,7 @@ public async Task> MonadicCreateItemAsync( CosmosObject insertedDocument = tryInsertDocument.Resource; string identifier = ((CosmosString)insertedDocument["id"]).Value; - ResourceId resourceIdentifier = ResourceId.Parse(((CosmosString)insertedDocument["_rid"]).Value); + ResourceIdentifier resourceIdentifier = ResourceIdentifier.Parse(((CosmosString)insertedDocument["_rid"]).Value); long timestamp = Number64.ToLong(((CosmosNumber)insertedDocument["_ts"]).Value); Record record = new Record(resourceIdentifier, timestamp, identifier, insertedDocument); @@ -92,7 +92,7 @@ public Task>> MonadicGetChildRangeAsync( public Task> MonadicReadFeedAsync( int partitionKeyRangeId, - ResourceId resourceIdentifer, + ResourceIdentifier resourceIdentifer, int pageSize, CancellationToken cancellationToken) { diff --git a/Microsoft.Azure.Cosmos/src/Pagination/Record.cs b/Microsoft.Azure.Cosmos/src/Pagination/Record.cs index 5496256e3b..63812c69a6 100644 --- a/Microsoft.Azure.Cosmos/src/Pagination/Record.cs +++ b/Microsoft.Azure.Cosmos/src/Pagination/Record.cs @@ -5,14 +5,12 @@ namespace Microsoft.Azure.Cosmos.Pagination { using System; - using System.Reflection; using Microsoft.Azure.Cosmos.CosmosElements; - using Microsoft.Azure.Documents; internal sealed class Record { public Record( - ResourceId resourceIdentifier, + ResourceIdentifier resourceIdentifier, long timestamp, string identifier, CosmosObject payload) @@ -23,7 +21,7 @@ public Record( this.Payload = payload ?? throw new ArgumentNullException(nameof(payload)); } - public ResourceId ResourceIdentifier { get; } + public ResourceIdentifier ResourceIdentifier { get; } public long Timestamp { get; } @@ -31,14 +29,18 @@ public Record( public CosmosObject Payload { get; } - public static Record Create(ResourceId previousResourceIdentifier, CosmosObject payload) + public static Record Create(ResourceIdentifier previousResourceIdentifier, CosmosObject payload) { - const string dummyRidString = "AYIMAMmFOw8YAAAAAAAAAA=="; - ResourceId resourceId = ResourceId.Parse(dummyRidString); - PropertyInfo prop = resourceId - .GetType() - .GetProperty("Document", BindingFlags.NonPublic | BindingFlags.Public | BindingFlags.Instance); - prop.SetValue(resourceId, previousResourceIdentifier.Document + 1); + ResourceIdentifier resourceId = new ResourceIdentifier( + previousResourceIdentifier.Offer, previousResourceIdentifier.Database, + previousResourceIdentifier.DocumentCollection, previousResourceIdentifier.StoredProcedure, + previousResourceIdentifier.Trigger, previousResourceIdentifier.UserDefinedFunction, + previousResourceIdentifier.Conflict, previousResourceIdentifier.Document + 1, + previousResourceIdentifier.PartitionKeyRange, previousResourceIdentifier.User, + previousResourceIdentifier.ClientEncryptionKey, previousResourceIdentifier.UserDefinedType, + previousResourceIdentifier.Permission, previousResourceIdentifier.Attachment, + previousResourceIdentifier.Schema, previousResourceIdentifier.Snapshot, + previousResourceIdentifier.RoleAssignment, previousResourceIdentifier.RoleDefinition); return new Record( resourceId, diff --git a/Microsoft.Azure.Cosmos/src/Pagination/ResourceIdentifier.cs b/Microsoft.Azure.Cosmos/src/Pagination/ResourceIdentifier.cs new file mode 100644 index 0000000000..356fd3d9d6 --- /dev/null +++ b/Microsoft.Azure.Cosmos/src/Pagination/ResourceIdentifier.cs @@ -0,0 +1,560 @@ +// ------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +// ------------------------------------------------------------ + +namespace Microsoft.Azure.Cosmos.Pagination +{ + using System; + using System.Linq; + + //Resource ID is 20 byte number. It is not a guid. + //First 4 bytes for DB id -> 2^32 DBs per tenant --> uint + // Next 4 bytes for coll id OR user id -> 2^31 coll/users per DB. the first bit indicates if this is a collection or a user --> uint + // Next 8 bytes for doc id OR permission or sproc/trigger/function or conflicts -> 2^64 permission per user. Note that the disambiguation between collection + // based resources (document/sproc/trigger/function/conflict) and user based resource i.e. permission is based on ownerid. --> ulong + // Disambiguation between document/sproc/conflict is based upon first 4 bits of the highest byte. 0x00 is document, 0x08 is sproc, 0x07 is + // trigger, 0x06 is function and 0x04 is conflict. 2^60 document/sproc/trigger/function/conflict per collection. + // Last 4 bytes for attachment id -> 2^32 attachments per "Document". These bits are used only for document's children. Permission/Sproc/ + // Conflict RID hierarchy is only 16 bytes. + + internal sealed class ResourceIdentifier : IEquatable + { + private const int OfferIdLength = 3; + private const int RbacResourceIdLength = 6; + private const int SnapshotIdLength = 7; + public static readonly ushort Length = 20; + public static readonly ushort MaxPathFragment = 8; // for public resource + public static readonly ResourceIdentifier Empty = new ResourceIdentifier(); + + public ResourceIdentifier( + uint offer = 0, + uint database = 0, + uint documentCollection = 0, + ulong storedProcedure = 0, + ulong trigger = 0, + ulong userDefinedFunction = 0, + ulong conflict = 0, + ulong document = 0, + ulong partitionKeyRange = 0, + uint user = 0, + uint clientEncryptionKey = 0, + uint userDefinedType = 0, + ulong permission = 0, + uint attachment = 0, + ulong schema = 0, + ulong snapshot = 0, + ulong roleAssignment = 0, + ulong roleDefinition = 0) + { + this.Offer = offer; + this.Database = database; + this.DocumentCollection = documentCollection; + this.StoredProcedure = storedProcedure; + this.Trigger = trigger; + this.UserDefinedFunction = userDefinedFunction; + this.Conflict = conflict; + this.Document = document; + this.PartitionKeyRange = partitionKeyRange; + this.User = user; + this.ClientEncryptionKey = clientEncryptionKey; + this.Permission = permission; + this.Attachment = attachment; + this.Schema = schema; + this.UserDefinedType = userDefinedType; + this.Snapshot = snapshot; + this.RoleAssignment = roleAssignment; + this.RoleDefinition = roleDefinition; + } + + public uint Offer { get; } + public uint Database { get; } + public uint DocumentCollection { get; } + public ulong StoredProcedure { get; } + public ulong Trigger { get; } + public ulong UserDefinedFunction { get; } + public ulong Conflict { get; } + public ulong Document { get; } + public ulong PartitionKeyRange { get; } + public uint User { get; } + public uint ClientEncryptionKey { get; } + public uint UserDefinedType { get; } + public ulong Permission { get; } + public uint Attachment { get; } + public ulong Schema { get; } + public ulong Snapshot { get; } + public ulong RoleAssignment { get; } + public ulong RoleDefinition { get; } + + public byte[] ToByteArray() + { + int len = 0; + if (this.Offer > 0) + { + len += ResourceIdentifier.OfferIdLength; + } + else if (this.Snapshot > 0) + { + len += ResourceIdentifier.SnapshotIdLength; + } + else if (this.RoleAssignment > 0) + { + len += ResourceIdentifier.RbacResourceIdLength; + } + else if (this.RoleDefinition > 0) + { + len += ResourceIdentifier.RbacResourceIdLength; + } + else if (this.Database > 0) + { + len += 4; + } + + if (this.DocumentCollection > 0 || this.User > 0 || this.UserDefinedType > 0 || this.ClientEncryptionKey > 0) + { + len += 4; + } + + if (this.Document > 0 || this.Permission > 0 || this.StoredProcedure > 0 || this.Trigger > 0 + || this.UserDefinedFunction > 0 || this.Conflict > 0 || this.PartitionKeyRange > 0 || this.Schema > 0 + || this.UserDefinedType > 0 || this.ClientEncryptionKey > 0) + { + len += 8; + } + + if (this.Attachment > 0) + { + len += 4; + } + + byte[] val = new byte[len]; + + if (this.Offer > 0) + { + ResourceIdentifier.BlockCopy(BitConverter.GetBytes(this.Offer), 0, val, 0, ResourceIdentifier.OfferIdLength); + } + else if (this.Database > 0) + { + ResourceIdentifier.BlockCopy(BitConverter.GetBytes(this.Database), 0, val, 0, 4); + } + else if (this.Snapshot > 0) + { + ResourceIdentifier.BlockCopy(BitConverter.GetBytes(this.Snapshot), 0, val, 0, ResourceIdentifier.SnapshotIdLength); + } + else if (this.RoleAssignment > 0) + { + ResourceIdentifier.BlockCopy(BitConverter.GetBytes(this.RoleAssignment), 0, val, 0, 4); + ResourceIdentifier.BlockCopy(BitConverter.GetBytes(0x1000), 0, val, 4, 2); + } + else if (this.RoleDefinition > 0) + { + ResourceIdentifier.BlockCopy(BitConverter.GetBytes(this.RoleDefinition), 0, val, 0, ResourceIdentifier.RbacResourceIdLength); + } + + if (this.DocumentCollection > 0) + { + ResourceIdentifier.BlockCopy(BitConverter.GetBytes(this.DocumentCollection), 0, val, 4, 4); + } + else if (this.User > 0) + { + ResourceIdentifier.BlockCopy(BitConverter.GetBytes(this.User), 0, val, 4, 4); + } + + if (this.StoredProcedure > 0) + { + ResourceIdentifier.BlockCopy(BitConverter.GetBytes(this.StoredProcedure), 0, val, 8, 8); + } + else if (this.Trigger > 0) + { + ResourceIdentifier.BlockCopy(BitConverter.GetBytes(this.Trigger), 0, val, 8, 8); + } + else if (this.UserDefinedFunction > 0) + { + ResourceIdentifier.BlockCopy(BitConverter.GetBytes(this.UserDefinedFunction), 0, val, 8, 8); + } + else if (this.Conflict > 0) + { + ResourceIdentifier.BlockCopy(BitConverter.GetBytes(this.Conflict), 0, val, 8, 8); + } + else if (this.Document > 0) + { + ResourceIdentifier.BlockCopy(BitConverter.GetBytes(this.Document), 0, val, 8, 8); + } + else if (this.PartitionKeyRange > 0) + { + ResourceIdentifier.BlockCopy(BitConverter.GetBytes(this.PartitionKeyRange), 0, val, 8, 8); + } + else if (this.Permission > 0) + { + ResourceIdentifier.BlockCopy(BitConverter.GetBytes(this.Permission), 0, val, 8, 8); + } + else if (this.Schema > 0) + { + ResourceIdentifier.BlockCopy(BitConverter.GetBytes(this.Schema), 0, val, 8, 8); + } + else if (this.UserDefinedType > 0) + { + ResourceIdentifier.BlockCopy(BitConverter.GetBytes(this.UserDefinedType), 0, val, 8, 4); + ResourceIdentifier.BlockCopy(BitConverter.GetBytes((uint)ExtendedDatabaseChildResourceType.UserDefinedType), 0, val, 12, 4); + } + else if (this.ClientEncryptionKey > 0) + { + ResourceIdentifier.BlockCopy(BitConverter.GetBytes(this.ClientEncryptionKey), 0, val, 8, 4); + ResourceIdentifier.BlockCopy(BitConverter.GetBytes((uint)ExtendedDatabaseChildResourceType.ClientEncryptionKey), 0, val, 12, 4); + } + + if (this.Attachment > 0) + { + ResourceIdentifier.BlockCopy(BitConverter.GetBytes(this.Attachment), 0, val, 16, 4); + } + + return val; + } + + public static ResourceIdentifier Parse(string id) + { + if (!ResourceIdentifier.TryParse(id, out ResourceIdentifier rid)) + { + throw new FormatException("Failed to parse ResourceId"); + } + + return rid; + } + + public static bool TryParse(string id, out ResourceIdentifier rid) + { + try + { + uint offer = 0; + uint database = 0; + uint documentCollection = 0; + ulong storedProcedure = 0; + ulong trigger = 0; + ulong userDefinedFunction = 0; + ulong conflict = 0; + ulong document = 0; + ulong partitionKeyRange = 0; + uint user = 0; + uint clientEncryptionKey = 0; + ulong permission = 0; + uint attachment = 0; + ulong schema = 0; + uint userDefinedType = 0; + ulong snapshot = 0; + ulong roleAssignment = 0; + ulong roleDefinition = 0; + + if (string.IsNullOrEmpty(id)) + { + rid = default; + return false; + } + + if (id.Length % 4 != 0) + { + // our resourceId string is always padded + rid = default; + return false; + } + + if (ResourceIdentifier.Verify(id, out byte[] buffer) == false) + { + rid = default; + return false; + } + + if (buffer.Length % 4 != 0 && + buffer.Length != ResourceIdentifier.OfferIdLength && + buffer.Length != ResourceIdentifier.SnapshotIdLength && + buffer.Length != ResourceIdentifier.RbacResourceIdLength) + { + rid = default; + return false; + } + + if (buffer.Length == ResourceIdentifier.OfferIdLength) + { + offer = (uint)ResourceIdentifier.ToUnsignedLong(buffer); + rid = new ResourceIdentifier(offer: offer); + return true; + } + + if (buffer.Length == ResourceIdentifier.SnapshotIdLength) + { + snapshot = ResourceIdentifier.ToUnsignedLong(buffer); + rid = new ResourceIdentifier(snapshot: snapshot); + return true; + } + + if (buffer.Length == ResourceIdentifier.RbacResourceIdLength) + { + byte rbacResourceType = buffer[ResourceIdentifier.RbacResourceIdLength - 1]; + ulong rbacResourceId = ResourceIdentifier.ToUnsignedLong(buffer, 4); + + switch ((RbacResourceType)rbacResourceType) + { + case RbacResourceType.RbacResourceType_RoleDefinition: + roleDefinition = rbacResourceId; + rid = new ResourceIdentifier( + offer, database, documentCollection, storedProcedure, + trigger, userDefinedFunction, conflict, document, + partitionKeyRange, user, clientEncryptionKey, + userDefinedType, permission, attachment, schema, + snapshot, roleAssignment, roleDefinition); + return true; + + case RbacResourceType.RbacResourceType_RoleAssignment: + roleAssignment = rbacResourceId; + rid = new ResourceIdentifier( + offer, database, documentCollection, storedProcedure, + trigger, userDefinedFunction, conflict, document, + partitionKeyRange, user, clientEncryptionKey, + userDefinedType, permission, attachment, schema, + snapshot, roleAssignment, roleDefinition); + return true; + + default: + rid = default; + return false; + } + } + + if (buffer.Length >= 4) + database = BitConverter.ToUInt32(buffer, 0); + + if (buffer.Length >= 8) + { + byte[] temp = new byte[4]; + ResourceIdentifier.BlockCopy(buffer, 4, temp, 0, 4); + + bool isCollection = (temp[0] & (128)) > 0; + + if (isCollection) + { + documentCollection = BitConverter.ToUInt32(temp, 0); + + if (buffer.Length >= 16) + { + byte[] subCollRes = new byte[8]; + ResourceIdentifier.BlockCopy(buffer, 8, subCollRes, 0, 8); + + UInt64 subCollectionResource = BitConverter.ToUInt64(buffer, 8); + if ((subCollRes[7] >> 4) == (byte)CollectionChildResourceType.Document) + { + document = subCollectionResource; + + if (buffer.Length == 20) + { + attachment = BitConverter.ToUInt32(buffer, 16); + } + } + else if ((subCollRes[7] >> 4) == (byte)CollectionChildResourceType.StoredProcedure) + { + storedProcedure = subCollectionResource; + } + else if ((subCollRes[7] >> 4) == (byte)CollectionChildResourceType.Trigger) + { + trigger = subCollectionResource; + } + else if ((subCollRes[7] >> 4) == (byte)CollectionChildResourceType.UserDefinedFunction) + { + userDefinedFunction = subCollectionResource; + } + else if ((subCollRes[7] >> 4) == (byte)CollectionChildResourceType.Conflict) + { + conflict = subCollectionResource; + } + else if ((subCollRes[7] >> 4) == (byte)CollectionChildResourceType.PartitionKeyRange) + { + partitionKeyRange = subCollectionResource; + } + else if ((subCollRes[7] >> 4) == (byte)CollectionChildResourceType.Schema) + { + schema = subCollectionResource; + } + else + { + rid = default; + return false; + } + } + else if (buffer.Length != 8) + { + rid = default; + return false; + } + } + else + { + user = BitConverter.ToUInt32(temp, 0); + + if (buffer.Length == 16) + { + if (user > 0) + { + permission = BitConverter.ToUInt64(buffer, 8); + } + else + { + uint exDatabaseChildResourceId = BitConverter.ToUInt32(buffer, 8); + ExtendedDatabaseChildResourceType exDatabaseChildResType = (ExtendedDatabaseChildResourceType)BitConverter.ToUInt32(buffer, 12); + + if (exDatabaseChildResType == ExtendedDatabaseChildResourceType.UserDefinedType) + { + userDefinedType = exDatabaseChildResourceId; + } + else if (exDatabaseChildResType == ExtendedDatabaseChildResourceType.ClientEncryptionKey) + { + clientEncryptionKey = exDatabaseChildResourceId; + } + else + { + rid = default; + return false; + } + } + } + else if (buffer.Length != 8) + { + rid = default; + return false; + } + } + } + + rid = new ResourceIdentifier( + offer, database, documentCollection, storedProcedure, + trigger, userDefinedFunction, conflict, document, + partitionKeyRange, user, clientEncryptionKey, + userDefinedType, permission, attachment, schema, + snapshot, roleAssignment, roleDefinition); + return true; + } + catch (Exception) + { + rid = default; + return false; + } + } + + public static bool Verify(string id, out byte[] buffer) + { + if (string.IsNullOrEmpty(id)) + throw new ArgumentNullException("id"); + + buffer = null; + + try + { + buffer = ResourceIdentifier.FromBase64String(id); + } + catch (FormatException) + { + } + + if (buffer == null || buffer.Length > ResourceIdentifier.Length) + { + buffer = null; + return false; + } + + return true; + } + + public override string ToString() + { + return ResourceIdentifier.ToBase64String(this.ToByteArray()); + } + + public bool Equals(ResourceIdentifier other) + { + if (other == null) + { + return false; + } + + return Enumerable.SequenceEqual(this.ToByteArray(), other.ToByteArray()); + } + + public override bool Equals(object obj) + { + if (obj is null) + { + return false; + } + if (ReferenceEquals(this, obj)) + { + return true; + } + return obj is ResourceIdentifier id && this.Equals(id); + } + + public override int GetHashCode() + { + throw new NotImplementedException(); + } + + public static byte[] FromBase64String(string s) + { + return Convert.FromBase64String(s.Replace('-', '/')); + } + + public static ulong ToUnsignedLong(byte[] buffer) + { + return ResourceIdentifier.ToUnsignedLong(buffer, buffer.Length); + } + + public static ulong ToUnsignedLong(byte[] buffer, int length) + { + ulong value = 0; + + for (int index = 0; index < length; index++) + { + value |= (uint)(buffer[index] << (index * 8)); + } + + return value; + } + + public static string ToBase64String(byte[] buffer) + { + return ResourceIdentifier.ToBase64String(buffer, 0, buffer.Length); + } + + public static string ToBase64String(byte[] buffer, int offset, int length) + { + return Convert.ToBase64String(buffer, offset, length).Replace('/', '-'); + } + + // Copy the bytes provided with a for loop, faster when there are only a few bytes to copy + public static void BlockCopy(byte[] src, int srcOffset, byte[] dst, int dstOffset, int count) + { + int stop = srcOffset + count; + for (int i = srcOffset; i < stop; i++) + dst[dstOffset++] = src[i]; + } + + // Using a byte however, we only need nibble here. + private enum CollectionChildResourceType : byte + { + Document = 0x00, + StoredProcedure = 0x08, + Trigger = 0x07, + UserDefinedFunction = 0x06, + Conflict = 0x04, + PartitionKeyRange = 0x05, + Schema = 0x09, + } + + private enum ExtendedDatabaseChildResourceType + { + UserDefinedType = 0x01, + ClientEncryptionKey = 0x02 + } + + internal enum RbacResourceType : byte + { + RbacResourceType_RoleDefinition = 0x00, + RbacResourceType_RoleAssignment = 0x10, + } + } +} diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/OrderBy/OrderByCrossPartitionQueryPipelineStage.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/OrderBy/OrderByCrossPartitionQueryPipelineStage.cs index cdb82d66a5..d7e9ff14c1 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/OrderBy/OrderByCrossPartitionQueryPipelineStage.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/OrderBy/OrderByCrossPartitionQueryPipelineStage.cs @@ -19,6 +19,7 @@ namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.CrossPartition.OrderBy using Microsoft.Azure.Cosmos.Query.Core.Monads; using Microsoft.Azure.Cosmos.Query.Core.Pipeline.CrossPartition.Parallel; using Microsoft.Azure.Documents; + using ResourceId = Documents.ResourceId; /// /// CosmosOrderByItemQueryExecutionContext is a concrete implementation for CrossPartitionQueryExecutionContext. @@ -51,8 +52,8 @@ internal sealed class OrderByCrossPartitionQueryPipelineStage : IQueryPipelineSt private readonly Queue<(OrderByQueryPartitionRangePageAsyncEnumerator enumerator, OrderByContinuationToken token)> uninitializedEnumeratorsAndTokens; private readonly int pageSize; private readonly int maxConcurrency; - private CancellationToken cancellationToken; + private CancellationToken cancellationToken; private QueryState state; private static class Expressions @@ -185,7 +186,7 @@ private async ValueTask MoveNextAsync_Iniialize_FilterAsync( throw new ArgumentNullException(nameof(token)); } - TryCatch<(bool, TryCatch)> filterMonad = await FilterNextAsync( + TryCatch<(bool, int, TryCatch)> filterMonad = await FilterNextAsync( uninitializedEnumerator, this.sortOrders, token, @@ -202,7 +203,7 @@ private async ValueTask MoveNextAsync_Iniialize_FilterAsync( return true; } - (bool doneFiltering, TryCatch monadicQueryByPage) = filterMonad.Result; + (bool doneFiltering, int itemsLeftToSkip, TryCatch monadicQueryByPage) = filterMonad.Result; if (doneFiltering) { if (uninitializedEnumerator.Current.Result.Enumerator.Current != null) @@ -235,7 +236,22 @@ private async ValueTask MoveNextAsync_Iniialize_FilterAsync( } } - this.uninitializedEnumeratorsAndTokens.Enqueue((uninitializedEnumerator, token)); + if (uninitializedEnumerator.State != default) + { + // We need to update the token + OrderByContinuationToken modifiedToken = new OrderByContinuationToken( + new ParallelContinuationToken( + ((CosmosString)uninitializedEnumerator.State.Value).Value, + uninitializedEnumerator.Range.ToRange()), + token.OrderByItems, + token.Rid, + itemsLeftToSkip, + token.Filter); + this.uninitializedEnumeratorsAndTokens.Enqueue((uninitializedEnumerator, modifiedToken)); + CosmosElement cosmosElementOrderByContinuationToken = OrderByContinuationToken.ToCosmosElement(modifiedToken); + CosmosArray continuationTokenList = CosmosArray.Create(new List() { cosmosElementOrderByContinuationToken }); + this.state = new QueryState(continuationTokenList); + } } QueryPage page = uninitializedEnumerator.Current.Result.Page; @@ -248,7 +264,7 @@ private async ValueTask MoveNextAsync_Iniialize_FilterAsync( responseLengthInBytes: page.ResponseLengthInBytes, cosmosQueryExecutionInfo: page.CosmosQueryExecutionInfo, disallowContinuationTokenMessage: page.DisallowContinuationTokenMessage, - state: this.state)); + state: InitializingQueryState)); return true; } @@ -304,12 +320,12 @@ private ValueTask MoveNextAsync_DrainPageAsync() OrderByQueryResult orderByQueryResult = default; // Try to form a page with as many items in the sorted order without having to do async work. - List documents = new List(); - while (documents.Count < this.pageSize) + List results = new List(); + while (results.Count < this.pageSize) { currentEnumerator = this.enumerators.Dequeue(); orderByQueryResult = new OrderByQueryResult(currentEnumerator.Current.Result.Enumerator.Current); - documents.Add(orderByQueryResult.Payload); + results.Add(orderByQueryResult); if (!currentEnumerator.Current.Result.Enumerator.MoveNext()) { @@ -327,6 +343,11 @@ private ValueTask MoveNextAsync_DrainPageAsync() this.enumerators.Enqueue(currentEnumerator); } + // It is possible that we emit multiple documents with the same rid due to JOIN queries. + // This means it is not enough to serialize the rid that we left on to resume the query. + // We need to also serialize the number of documents with that rid, so we can skip it when resuming + int skipCount = results.Where(result => string.Equals(result.Rid, orderByQueryResult.Rid)).Count(); + // Create the continuation token. CosmosElement state; if ((this.enumerators.Count == 0) && (this.uninitializedEnumeratorsAndTokens.Count == 0)) @@ -341,7 +362,7 @@ private ValueTask MoveNextAsync_DrainPageAsync() range: currentEnumerator.Range.ToRange()), orderByQueryResult.OrderByItems, orderByQueryResult.Rid, - skipCount: 0, + skipCount: skipCount, filter: currentEnumerator.Filter); CosmosElement cosmosElementOrderByContinuationToken = OrderByContinuationToken.ToCosmosElement(orderByContinuationToken); @@ -356,7 +377,7 @@ private ValueTask MoveNextAsync_DrainPageAsync() // No stats to report, since we already reported it when we moved to this page. this.Current = TryCatch.FromResult( new QueryPage( - documents: documents, + documents: results.Select(result => result.Payload).ToList(), requestCharge: 0, activityId: default, responseLengthInBytes: 0, @@ -366,6 +387,30 @@ private ValueTask MoveNextAsync_DrainPageAsync() return new ValueTask(true); } + //// In order to maintain the continuation token for the user we must drain with a few constraints + //// 1) We always drain from the partition, which has the highest priority item first + //// 2) If multiple partitions have the same priority item then we drain from the left most first + //// otherwise we would need to keep track of how many of each item we drained from each partition + //// (just like parallel queries). + //// Visually that look the following case where we have three partitions that are numbered and store letters. + //// For teaching purposes I have made each item a tuple of the following form: + //// + //// So that duplicates across partitions are distinct, but duplicates within partitions are indistinguishable. + //// |-------| |-------| |-------| + //// | | | | | | + //// | | | | | | + //// | | | | | | + //// | | | | | | + //// | | | | | | + //// | | | | | | + //// | | | | | | + //// |-------| |-------| |-------| + //// Now the correct drain order in this case is: + //// ,,,,,,,,,,, + //// ,,,,,,,,, + //// In more mathematical terms + //// 1) always comes before where x < z + //// 2) always come before where j < k public ValueTask MoveNextAsync() { this.cancellationToken.ThrowIfCancellationRequested(); @@ -776,7 +821,7 @@ private static (string leftFilter, string targetFilter, string rightFilter) GetF /// The continuation token. /// The cancellation token. /// A task to await on. - private static async Task)>> FilterNextAsync( + private static async Task monadicQueryByPage)>> FilterNextAsync( OrderByQueryPartitionRangePageAsyncEnumerator enumerator, IReadOnlyList sortOrders, OrderByContinuationToken continuationToken, @@ -789,82 +834,104 @@ private static (string leftFilter, string targetFilter, string rightFilter) GetF // The key is to seek until we get an order by value that matches the order by value we left off on. // Once we do that we need to seek to the correct _rid within the term, // since there might be many documents with the same order by value we left off on. + // Finally we need to skip some duplicate _rids, since JOINs emit multiples documents with the same rid and we read a partial page. + // You can also think about this as a seek on a composite index where the columns are [sort_order, rid, skip_count] + int itemsToSkip = continuationToken.SkipCount; if (!ResourceId.TryParse(continuationToken.Rid, out ResourceId continuationRid)) { - return TryCatch<(bool, TryCatch)>.FromException( + return TryCatch<(bool, int, TryCatch)>.FromException( new MalformedContinuationTokenException( $"Invalid Rid in the continuation token {continuationToken.ParallelContinuationToken.Token} for OrderBy~Context.")); } - // Throw away documents until it matches the item from the continuation token. if (!await enumerator.MoveNextAsync()) { - return TryCatch<(bool, TryCatch)>.FromResult((true, enumerator.Current)); + return TryCatch<(bool, int, TryCatch)>.FromResult((true, 0, enumerator.Current)); } TryCatch monadicOrderByQueryPage = enumerator.Current; if (monadicOrderByQueryPage.Failed) { - return TryCatch<(bool, TryCatch)>.FromException(monadicOrderByQueryPage.Exception); + return TryCatch<(bool, int, TryCatch)>.FromException(monadicOrderByQueryPage.Exception); } OrderByQueryPage orderByQueryPage = monadicOrderByQueryPage.Result; IEnumerator documents = orderByQueryPage.Enumerator; + while (documents.MoveNext()) { + int sortOrderCompare = 0; + // Filter out documents until we find something that matches the sort order. OrderByQueryResult orderByResult = new OrderByQueryResult(documents.Current); - ResourceId rid = ResourceId.Parse(orderByResult.Rid); - int cmp = 0; - for (int i = 0; (i < sortOrders.Count) && (cmp == 0); ++i) + for (int i = 0; (i < sortOrders.Count) && (sortOrderCompare == 0); ++i) { - cmp = ItemComparer.Instance.Compare( + sortOrderCompare = ItemComparer.Instance.Compare( continuationToken.OrderByItems[i].Item, orderByResult.OrderByItems[i].Item); - if (cmp != 0) + if (sortOrderCompare != 0) { - cmp = sortOrders[i] == SortOrder.Ascending ? cmp : -cmp; + sortOrderCompare = sortOrders[i] == SortOrder.Ascending ? sortOrderCompare : -sortOrderCompare; } } - if (cmp < 0) + if (sortOrderCompare < 0) { // We might have passed the item due to deletions and filters. - return TryCatch<(bool, TryCatch)>.FromResult((true, enumerator.Current)); + return TryCatch<(bool, int, TryCatch)>.FromResult((true, 0, enumerator.Current)); } - if (cmp == 0) + if (sortOrderCompare > 0) { - // Once the item matches the order by items from the continuation tokens - // We still need to remove all the documents that have a lower or same rid in the rid sort order. - // If there is a tie in the sort order the documents should be in _rid order in the same direction as the index (given by the backend) - cmp = continuationRid.Document.CompareTo(rid.Document); - if ((orderByQueryPage.Page.CosmosQueryExecutionInfo == null) || orderByQueryPage.Page.CosmosQueryExecutionInfo.ReverseRidEnabled) + // This document does not match the sort order, so skip it. + continue; + } + + // Once the item matches the order by items from the continuation tokens + // We still need to remove all the documents that have a lower or same rid in the rid sort order. + // If there is a tie in the sort order the documents should be in _rid order in the same direction as the index (given by the backend) + ResourceId rid = ResourceId.Parse(orderByResult.Rid); + int ridOrderCompare = continuationRid.Document.CompareTo(rid.Document); + if ((orderByQueryPage.Page.CosmosQueryExecutionInfo == null) || orderByQueryPage.Page.CosmosQueryExecutionInfo.ReverseRidEnabled) + { + // If reverse rid is enabled on the backend then fallback to the old way of doing it. + if (sortOrders[0] == SortOrder.Descending) { - // If reverse rid is enabled on the backend then fallback to the old way of doing it. - if (sortOrders[0] == SortOrder.Descending) - { - cmp = -cmp; - } + ridOrderCompare = -ridOrderCompare; } - else + } + else + { + // Go by the whatever order the index wants + if (orderByQueryPage.Page.CosmosQueryExecutionInfo.ReverseIndexScan) { - // Go by the whatever order the index wants - if (orderByQueryPage.Page.CosmosQueryExecutionInfo.ReverseIndexScan) - { - cmp = -cmp; - } + ridOrderCompare = -ridOrderCompare; } + } - if (cmp < 0) - { - return TryCatch<(bool, TryCatch)>.FromResult((true, enumerator.Current)); - } + if (ridOrderCompare < 0) + { + // We might have passed the rid due to deletions and filters. + return TryCatch<(bool, int, TryCatch)>.FromResult((true, 0, enumerator.Current)); + } + + if (ridOrderCompare > 0) + { + // This document does not match the rid order, so skip it. + continue; + } + + // At this point we need to skip due to joins + if (--itemsToSkip < 0) + { + return TryCatch<(bool, int, TryCatch)>.FromResult((true, 0, enumerator.Current)); } } - return TryCatch<(bool, TryCatch)>.FromResult((false, enumerator.Current)); + // If we made it here it means we failed to find the resume order by item which is possible + // if the user added documents inbetween continuations, so we need to yield and filter the next page of results also. + return TryCatch<(bool, int, TryCatch)>.FromResult((false, itemsToSkip, enumerator.Current)); } private static bool IsSplitException(Exception exception) diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/OrderBy/OrderByQueryPartitionRangePageAsyncEnumerator.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/OrderBy/OrderByQueryPartitionRangePageAsyncEnumerator.cs index 49c02e1858..d247a97e42 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/OrderBy/OrderByQueryPartitionRangePageAsyncEnumerator.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/OrderBy/OrderByQueryPartitionRangePageAsyncEnumerator.cs @@ -47,7 +47,7 @@ public OrderByQueryPartitionRangePageAsyncEnumerator( public string Filter => this.innerEnumerator.Filter; - public QueryState StartOfPageState { get; set; } + public QueryState StartOfPageState { get; private set; } public override ValueTask DisposeAsync() => default; 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 ee81e0c35c..e3b507bcd7 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 @@ -41,6 +41,10 @@ private ParallelCrossPartitionQueryPipelineStage( public ValueTask DisposeAsync() => this.crossPartitionRangePageAsyncEnumerator.DisposeAsync(); + // In order to maintain the continuation token for the user we must drain with a few constraints + // 1) We fully drain from the left most partition before moving on to the next partition + // 2) We drain only full pages from the document producer so we aren't left with a partial page + // otherwise we would need to add to the continuation token how many items to skip over on that page. public async ValueTask MoveNextAsync() { this.cancellationToken.ThrowIfCancellationRequested(); diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/DocumentContainerPartitionRangeEnumerator.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/DocumentContainerPartitionRangeEnumerator.cs index 706d62f991..0a02a55a46 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/DocumentContainerPartitionRangeEnumerator.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/DocumentContainerPartitionRangeEnumerator.cs @@ -31,7 +31,7 @@ public DocumentContainerPartitionRangeEnumerator( MaxExclusive = partitionKeyRangeId.ToString() }, cancellationToken, - state ?? new DocumentContainerState(resourceIdentifier: ResourceId.Empty)) + state ?? new DocumentContainerState(resourceIdentifier: ResourceIdentifier.Empty)) { this.documentContainer = documentContainer ?? throw new ArgumentNullException(nameof(documentContainer)); this.partitionKeyRangeId = partitionKeyRangeId; diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/DocumentContainerTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/DocumentContainerTests.cs index 9ff25b0024..54b6a7cf02 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/DocumentContainerTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/DocumentContainerTests.cs @@ -204,7 +204,7 @@ async Task AssertChildPartitionAsync(int partitionKeyRangeId) { DocumentContainerPage readFeedPage = await documentContainer.ReadFeedAsync( partitionKeyRangeId: partitionKeyRangeId, - resourceIdentifier: ResourceId.Empty, + resourceIdentifier: ResourceIdentifier.Empty, pageSize: 100, cancellationToken: default); @@ -296,7 +296,7 @@ async Task AssertChildPartitionAsync(int partitionKeyRangeId) { DocumentContainerPage page = await documentContainer.ReadFeedAsync( partitionKeyRangeId: partitionKeyRangeId, - resourceIdentifier: ResourceId.Empty, + resourceIdentifier: ResourceIdentifier.Empty, pageSize: 100, cancellationToken: default); diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/FlakyDocumentContainer.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/FlakyDocumentContainer.cs index c1d31583dd..bb786dde8b 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/FlakyDocumentContainer.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/FlakyDocumentContainer.cs @@ -15,6 +15,7 @@ namespace Microsoft.Azure.Cosmos.Tests.Pagination using Microsoft.Azure.Cosmos.Query.Core.Monads; using Microsoft.Azure.Cosmos.Query.Core.Pipeline; using Microsoft.Azure.Documents; + using ResourceIdentifier = Cosmos.Pagination.ResourceIdentifier; /// /// Implementation of that composes another and randomly adds in exceptions. @@ -89,7 +90,7 @@ public Task> MonadicReadItemAsync( public Task> MonadicReadFeedAsync( int partitionKeyRangeId, - ResourceId resourceIdentifer, + ResourceIdentifier resourceIdentifer, int pageSize, CancellationToken cancellationToken) { 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 2c358e13a6..fd1d7ff68f 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 @@ -23,6 +23,7 @@ namespace Microsoft.Azure.Cosmos.Tests.Pagination using Microsoft.Azure.Cosmos.SqlObjects; using Microsoft.Azure.Cosmos.Tests.Query.OfflineEngine; using Microsoft.Azure.Documents; + using ResourceIdentifier = Cosmos.Pagination.ResourceIdentifier; // Collection useful for mocking requests and repartitioning (splits / merge). internal sealed class InMemoryContainer : IMonadicDocumentContainer @@ -180,7 +181,21 @@ public Task> MonadicCreateItemAsync( this.partitionedRecords[partitionKeyHash] = records; } - Record recordAdded = records.Add(payload); + int? pkrangeid = null; + foreach (KeyValuePair kvp in this.partitionKeyRangeIdToHashRange) + { + if (kvp.Value.Contains(partitionKeyHash)) + { + pkrangeid = kvp.Key; + } + } + + if (!pkrangeid.HasValue) + { + throw new InvalidOperationException(); + } + + Record recordAdded = records.Add(pkrangeid.Value, payload); return Task.FromResult(TryCatch.FromResult(recordAdded)); } @@ -246,7 +261,7 @@ static Task> CreateNotFoundException(CosmosElement partitionKey public Task> MonadicReadFeedAsync( int partitionKeyRangeId, - ResourceId resourceIdentifer, + ResourceIdentifier resourceIdentifer, int pageSize, CancellationToken cancellationToken) { @@ -361,26 +376,59 @@ public Task> MonadicQueryAsync( IEnumerable queryResults = SqlInterpreter.ExecuteQuery(documents, sqlQuery); IEnumerable queryPageResults = queryResults; + + string continuationResourceId; + int continuationSkipCount; + if (continuationToken != null) { - queryPageResults = queryPageResults.Where(c => + CosmosObject parsedContinuationToken = CosmosObject.Parse(continuationToken); + continuationResourceId = ((CosmosString)parsedContinuationToken["resourceId"]).Value; + continuationSkipCount = (int)Number64.ToLong(((CosmosNumber64)parsedContinuationToken["skipCount"]).Value); + + queryPageResults = queryPageResults.Where((Func)(c => { - ResourceId continuationResourceId = ResourceId.Parse(continuationToken); - ResourceId documentResourceId = ResourceId.Parse(((CosmosString)((CosmosObject)c)["_rid"]).Value); - return documentResourceId.Document > continuationResourceId.Document; - }); + ResourceIdentifier continuationParsedResourceId = ResourceIdentifier.Parse(continuationResourceId); + ResourceIdentifier documentResourceId = ResourceIdentifier.Parse(((CosmosString)((CosmosObject)c)["_rid"]).Value); + return documentResourceId.Document >= continuationParsedResourceId.Document; + })); + + if (queryPageResults.FirstOrDefault() is CosmosObject firstDocument) + { + string currentResourceId = ((CosmosString)firstDocument["_rid"]).Value; + if (currentResourceId == continuationResourceId) + { + queryPageResults = queryPageResults.Skip(continuationSkipCount); + } + } + } + else + { + continuationResourceId = null; + continuationSkipCount = 0; } queryPageResults = queryPageResults.Take(pageSize); List queryPageResultList = queryPageResults.ToList(); QueryState queryState; - if (queryPageResultList.Count == 0) - { - queryState = default; - } - else if (queryPageResultList.Last() is CosmosObject lastDocument) + if (queryPageResultList.LastOrDefault() is CosmosObject lastDocument) { - queryState = new QueryState(lastDocument["_rid"]); + string currentResourceId = ((CosmosString)lastDocument["_rid"]).Value; + int currentSkipCount = queryPageResultList + .Where(document => ((CosmosString)((CosmosObject)document)["_rid"]).Value == currentResourceId) + .Count(); + if (currentResourceId == continuationResourceId) + { + currentSkipCount += continuationSkipCount; + } + + CosmosObject queryStateValue = CosmosObject.Create(new Dictionary() + { + { "resourceId", CosmosString.Create(currentResourceId) }, + { "skipCount", CosmosNumber64.Create(currentSkipCount) }, + }); + + queryState = new QueryState(CosmosString.Create(queryStateValue.ToString())); } else { @@ -589,10 +637,25 @@ public Records() IEnumerator IEnumerable.GetEnumerator() => this.storage.GetEnumerator(); - public Record Add(CosmosObject payload) + public Record Add(int pkrangeid, CosmosObject payload) { - ResourceId previousResourceId = this.Count == 0 ? ResourceId.Empty : this.storage[this.storage.Count - 1].ResourceIdentifier; - Record record = Record.Create(previousResourceId, payload); + ResourceIdentifier currentResourceId; + if (this.Count == 0) + { + currentResourceId = new ResourceIdentifier(database: 1, documentCollection: 1, partitionKeyRange: (ulong)pkrangeid, document: 1); + } + else + { + currentResourceId = this.storage[this.storage.Count - 1].ResourceIdentifier; + } + + ResourceIdentifier nextResourceId = new ResourceIdentifier( + database: currentResourceId.Database, + documentCollection: currentResourceId.DocumentCollection, + partitionKeyRange: (ulong)pkrangeid, + document: currentResourceId.Document); + + Record record = Record.Create(nextResourceId, payload); this.storage.Add(record); return record; } 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 8e72f95eb1..5752d9603a 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 @@ -4,6 +4,7 @@ namespace Microsoft.Azure.Cosmos.Tests.Query.Pipeline { + using System; using System.Collections.Generic; using System.Collections.ObjectModel; using System.Linq; @@ -88,6 +89,26 @@ public async Task OrderBy() Assert.AreEqual(expected: documents.Count, actual: documentsQueried.Count); } + [TestMethod] + public async Task OrderByWithJoins() + { + Random random = new Random(42); + List documents = new List() + { + CosmosObject.Parse($"{{\"pk\" : {random.Next()}, \"children\" : [\"Alice\", \"Bob\", \"Charlie\"]}}"), + CosmosObject.Parse($"{{\"pk\" : {random.Next()}, \"children\" : [\"Dave\", \"Eve\", \"Fancy\"]}}"), + CosmosObject.Parse($"{{\"pk\" : {random.Next()}, \"children\" : [\"George\", \"Henry\", \"Igor\"]}}"), + CosmosObject.Parse($"{{\"pk\" : {random.Next()}, \"children\" : [\"Jack\", \"Kim\", \"Levin\"]}}"), + }; + + List documentsQueried = await ExecuteQueryAsync( + query: "SELECT d FROM c JOIN d in c.children ORDER BY c._ts", + documents: documents, + pageSize: 2); + + Assert.AreEqual(expected: documents.Count * 3, actual: documentsQueried.Count); + } + [TestMethod] public async Task Top() { @@ -156,16 +177,17 @@ public async Task GroupBy() private static async Task> ExecuteQueryAsync( string query, - IReadOnlyList documents) + IReadOnlyList documents, + int pageSize = 10) { IDocumentContainer documentContainer = await CreateDocumentContainerAsync(documents); - List resultsFromDrainWithoutState = await DrainWithoutStateAsync(query, documentContainer); - List resultsFromDrainWithState = await DrainWithStateAsync(query, documentContainer); + //List resultsFromDrainWithoutState = await DrainWithoutStateAsync(query, documentContainer, pageSize); + List resultsFromDrainWithState = await DrainWithStateAsync(query, documentContainer, pageSize); - Assert.IsTrue(resultsFromDrainWithoutState.SequenceEqual(resultsFromDrainWithState)); + //Assert.IsTrue(resultsFromDrainWithoutState.SequenceEqual(resultsFromDrainWithState)); - return resultsFromDrainWithoutState; + return resultsFromDrainWithState; } [TestMethod] @@ -184,9 +206,9 @@ public async Task Fuzz() Assert.AreEqual(expected: 249, actual: documentsQueried.Count); } - private static async Task> DrainWithoutStateAsync(string query, IDocumentContainer documentContainer) + private static async Task> DrainWithoutStateAsync(string query, IDocumentContainer documentContainer, int pageSize = 10) { - IQueryPipelineStage pipelineStage = CreatePipeline(documentContainer, query); + IQueryPipelineStage pipelineStage = CreatePipeline(documentContainer, query, pageSize); List elements = new List(); while (await pipelineStage.MoveNextAsync()) @@ -200,7 +222,7 @@ private static async Task> DrainWithoutStateAsync(string que return elements; } - private static async Task> DrainWithStateAsync(string query, IDocumentContainer documentContainer) + private static async Task> DrainWithStateAsync(string query, IDocumentContainer documentContainer, int pageSize = 10) { IQueryPipelineStage pipelineStage; CosmosElement state = null; @@ -208,7 +230,7 @@ private static async Task> DrainWithStateAsync(string query, List elements = new List(); do { - pipelineStage = CreatePipeline(documentContainer, query, state); + pipelineStage = CreatePipeline(documentContainer, query, pageSize, state); if (!await pipelineStage.MoveNextAsync()) { @@ -265,7 +287,7 @@ private static async Task CreateDocumentContainerAsync( return documentContainer; } - private static IQueryPipelineStage CreatePipeline(IDocumentContainer documentContainer, string query, CosmosElement state = null) + private static IQueryPipelineStage CreatePipeline(IDocumentContainer documentContainer, string query, int pageSize = 10, CosmosElement state = null) { TryCatch tryCreatePipeline = PipelineFactory.MonadicCreate( ExecutionEnvironment.Compute, @@ -273,7 +295,7 @@ private static IQueryPipelineStage CreatePipeline(IDocumentContainer documentCon new SqlQuerySpec(query), documentContainer.GetFeedRangesAsync(default(CancellationToken)).Result, GetQueryPlan(query), - pageSize: 10, + pageSize: pageSize, maxConcurrency: 10, requestCancellationToken: default, requestContinuationToken: state); 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 c4f754a8ee..b2e68ee272 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 @@ -112,7 +112,7 @@ public override IReadOnlyList GetRecordsFromPage(QueryPage page) foreach (CosmosElement element in page.Documents) { CosmosObject document = (CosmosObject)element; - ResourceId resourceIdentifier = ResourceId.Parse(((CosmosString)document["_rid"]).Value); + ResourceIdentifier resourceIdentifier = ResourceIdentifier.Parse(((CosmosString)document["_rid"]).Value); long timestamp = Number64.ToLong(((CosmosNumber)document["_ts"]).Value); string identifer = ((CosmosString)document["id"]).Value; From 6f9c2af325f8dee07b120caa09e720c4201e71a1 Mon Sep 17 00:00:00 2001 From: Brandon Chong Date: Sat, 3 Oct 2020 02:19:55 -0700 Subject: [PATCH 78/85] fixed bug with JOINs + ORDER BY + Continuation tokens --- .../src/Pagination/DocumentContainer.cs | 4 +- .../src/Pagination/DocumentContainerState.cs | 4 +- .../src/Pagination/IDocumentContainer.cs | 3 +- .../Pagination/IMonadicDocumentContainer.cs | 3 +- .../NetworkAttachedDocumentContainer.cs | 4 +- .../src/Pagination/Record.cs | 25 ++------- .../src/Pagination/ResourceIdentifier.cs | 30 +---------- ...OrderByCrossPartitionQueryPipelineStage.cs | 32 +++++++++++- ...cumentContainerPartitionRangeEnumerator.cs | 2 +- .../Pagination/DocumentContainerTests.cs | 4 +- .../Pagination/FlakyDocumentContainer.cs | 2 +- .../Pagination/InMemoryContainer.cs | 52 +++++++++++++++---- .../Query/Pipeline/FullPipelineTests.cs | 17 +++--- .../QueryPartitionRangePageEnumeratorTests.cs | 2 +- 14 files changed, 98 insertions(+), 86 deletions(-) diff --git a/Microsoft.Azure.Cosmos/src/Pagination/DocumentContainer.cs b/Microsoft.Azure.Cosmos/src/Pagination/DocumentContainer.cs index 65805d33fd..fc3ee75fe2 100644 --- a/Microsoft.Azure.Cosmos/src/Pagination/DocumentContainer.cs +++ b/Microsoft.Azure.Cosmos/src/Pagination/DocumentContainer.cs @@ -84,7 +84,7 @@ public Task ReadItemAsync( public Task> MonadicReadFeedAsync( int partitionKeyRangeId, - ResourceIdentifier resourceIdentifer, + ResourceId resourceIdentifer, int pageSize, CancellationToken cancellationToken) => this.monadicDocumentContainer.MonadicReadFeedAsync( partitionKeyRangeId, @@ -94,7 +94,7 @@ public Task> MonadicReadFeedAsync( public Task ReadFeedAsync( int partitionKeyRangeId, - ResourceIdentifier resourceIdentifier, + ResourceId resourceIdentifier, int pageSize, CancellationToken cancellationToken) => TryCatch.UnsafeGetResultAsync( this.MonadicReadFeedAsync( diff --git a/Microsoft.Azure.Cosmos/src/Pagination/DocumentContainerState.cs b/Microsoft.Azure.Cosmos/src/Pagination/DocumentContainerState.cs index 5c97e4b62b..dd187d638d 100644 --- a/Microsoft.Azure.Cosmos/src/Pagination/DocumentContainerState.cs +++ b/Microsoft.Azure.Cosmos/src/Pagination/DocumentContainerState.cs @@ -8,11 +8,11 @@ namespace Microsoft.Azure.Cosmos.Pagination internal sealed class DocumentContainerState : State { - public DocumentContainerState(ResourceIdentifier resourceIdentifier) + public DocumentContainerState(ResourceId resourceIdentifier) { this.ResourceIdentifer = resourceIdentifier; } - public ResourceIdentifier ResourceIdentifer { get; } + public ResourceId ResourceIdentifer { get; } } } diff --git a/Microsoft.Azure.Cosmos/src/Pagination/IDocumentContainer.cs b/Microsoft.Azure.Cosmos/src/Pagination/IDocumentContainer.cs index ee4fd24657..2e63b9f30d 100644 --- a/Microsoft.Azure.Cosmos/src/Pagination/IDocumentContainer.cs +++ b/Microsoft.Azure.Cosmos/src/Pagination/IDocumentContainer.cs @@ -8,6 +8,7 @@ namespace Microsoft.Azure.Cosmos.Pagination using System.Threading.Tasks; using Microsoft.Azure.Cosmos.CosmosElements; using Microsoft.Azure.Cosmos.Query.Core.Pipeline.CrossPartition; + using Microsoft.Azure.Documents; internal interface IDocumentContainer : IMonadicDocumentContainer, IFeedRangeProvider, IQueryDataSource { @@ -22,7 +23,7 @@ Task ReadItemAsync( Task ReadFeedAsync( int partitionKeyRangeId, - ResourceIdentifier resourceIdentifier, + ResourceId resourceIdentifier, int pageSize, CancellationToken cancellationToken); diff --git a/Microsoft.Azure.Cosmos/src/Pagination/IMonadicDocumentContainer.cs b/Microsoft.Azure.Cosmos/src/Pagination/IMonadicDocumentContainer.cs index ce559b8eb1..6886765eb4 100644 --- a/Microsoft.Azure.Cosmos/src/Pagination/IMonadicDocumentContainer.cs +++ b/Microsoft.Azure.Cosmos/src/Pagination/IMonadicDocumentContainer.cs @@ -9,6 +9,7 @@ namespace Microsoft.Azure.Cosmos.Pagination using Microsoft.Azure.Cosmos.CosmosElements; using Microsoft.Azure.Cosmos.Query.Core.Monads; using Microsoft.Azure.Cosmos.Query.Core.Pipeline.CrossPartition; + using Microsoft.Azure.Documents; internal interface IMonadicDocumentContainer : IMonadicFeedRangeProvider, IMonadicQueryDataSource { @@ -23,7 +24,7 @@ Task> MonadicReadItemAsync( Task> MonadicReadFeedAsync( int partitionKeyRangeId, - ResourceIdentifier resourceIdentifer, + ResourceId resourceIdentifer, int pageSize, CancellationToken cancellationToken); diff --git a/Microsoft.Azure.Cosmos/src/Pagination/NetworkAttachedDocumentContainer.cs b/Microsoft.Azure.Cosmos/src/Pagination/NetworkAttachedDocumentContainer.cs index 07ce6a4886..9b8b6019a4 100644 --- a/Microsoft.Azure.Cosmos/src/Pagination/NetworkAttachedDocumentContainer.cs +++ b/Microsoft.Azure.Cosmos/src/Pagination/NetworkAttachedDocumentContainer.cs @@ -63,7 +63,7 @@ public async Task> MonadicCreateItemAsync( CosmosObject insertedDocument = tryInsertDocument.Resource; string identifier = ((CosmosString)insertedDocument["id"]).Value; - ResourceIdentifier resourceIdentifier = ResourceIdentifier.Parse(((CosmosString)insertedDocument["_rid"]).Value); + ResourceId resourceIdentifier = ResourceId.Parse(((CosmosString)insertedDocument["_rid"]).Value); long timestamp = Number64.ToLong(((CosmosNumber)insertedDocument["_ts"]).Value); Record record = new Record(resourceIdentifier, timestamp, identifier, insertedDocument); @@ -92,7 +92,7 @@ public Task>> MonadicGetChildRangeAsync( public Task> MonadicReadFeedAsync( int partitionKeyRangeId, - ResourceIdentifier resourceIdentifer, + ResourceId resourceIdentifer, int pageSize, CancellationToken cancellationToken) { diff --git a/Microsoft.Azure.Cosmos/src/Pagination/Record.cs b/Microsoft.Azure.Cosmos/src/Pagination/Record.cs index 63812c69a6..8563ee1ceb 100644 --- a/Microsoft.Azure.Cosmos/src/Pagination/Record.cs +++ b/Microsoft.Azure.Cosmos/src/Pagination/Record.cs @@ -6,11 +6,12 @@ namespace Microsoft.Azure.Cosmos.Pagination { using System; using Microsoft.Azure.Cosmos.CosmosElements; + using Microsoft.Azure.Documents; internal sealed class Record { public Record( - ResourceIdentifier resourceIdentifier, + ResourceId resourceIdentifier, long timestamp, string identifier, CosmosObject payload) @@ -21,32 +22,12 @@ public Record( this.Payload = payload ?? throw new ArgumentNullException(nameof(payload)); } - public ResourceIdentifier ResourceIdentifier { get; } + public ResourceId ResourceIdentifier { get; } public long Timestamp { get; } public string Identifier { get; } public CosmosObject Payload { get; } - - public static Record Create(ResourceIdentifier previousResourceIdentifier, CosmosObject payload) - { - ResourceIdentifier resourceId = new ResourceIdentifier( - previousResourceIdentifier.Offer, previousResourceIdentifier.Database, - previousResourceIdentifier.DocumentCollection, previousResourceIdentifier.StoredProcedure, - previousResourceIdentifier.Trigger, previousResourceIdentifier.UserDefinedFunction, - previousResourceIdentifier.Conflict, previousResourceIdentifier.Document + 1, - previousResourceIdentifier.PartitionKeyRange, previousResourceIdentifier.User, - previousResourceIdentifier.ClientEncryptionKey, previousResourceIdentifier.UserDefinedType, - previousResourceIdentifier.Permission, previousResourceIdentifier.Attachment, - previousResourceIdentifier.Schema, previousResourceIdentifier.Snapshot, - previousResourceIdentifier.RoleAssignment, previousResourceIdentifier.RoleDefinition); - - return new Record( - resourceId, - DateTime.UtcNow.Ticks, - Guid.NewGuid().ToString(), - payload); - } } } diff --git a/Microsoft.Azure.Cosmos/src/Pagination/ResourceIdentifier.cs b/Microsoft.Azure.Cosmos/src/Pagination/ResourceIdentifier.cs index 356fd3d9d6..9acb9ee72b 100644 --- a/Microsoft.Azure.Cosmos/src/Pagination/ResourceIdentifier.cs +++ b/Microsoft.Azure.Cosmos/src/Pagination/ResourceIdentifier.cs @@ -17,7 +17,7 @@ namespace Microsoft.Azure.Cosmos.Pagination // Last 4 bytes for attachment id -> 2^32 attachments per "Document". These bits are used only for document's children. Permission/Sproc/ // Conflict RID hierarchy is only 16 bytes. - internal sealed class ResourceIdentifier : IEquatable + internal sealed class ResourceIdentifier { private const int OfferIdLength = 3; private const int RbacResourceIdLength = 6; @@ -465,34 +465,6 @@ public override string ToString() return ResourceIdentifier.ToBase64String(this.ToByteArray()); } - public bool Equals(ResourceIdentifier other) - { - if (other == null) - { - return false; - } - - return Enumerable.SequenceEqual(this.ToByteArray(), other.ToByteArray()); - } - - public override bool Equals(object obj) - { - if (obj is null) - { - return false; - } - if (ReferenceEquals(this, obj)) - { - return true; - } - return obj is ResourceIdentifier id && this.Equals(id); - } - - public override int GetHashCode() - { - throw new NotImplementedException(); - } - public static byte[] FromBase64String(string s) { return Convert.FromBase64String(s.Replace('-', '/')); diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/OrderBy/OrderByCrossPartitionQueryPipelineStage.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/OrderBy/OrderByCrossPartitionQueryPipelineStage.cs index d7e9ff14c1..a755e1e5c9 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/OrderBy/OrderByCrossPartitionQueryPipelineStage.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/OrderBy/OrderByCrossPartitionQueryPipelineStage.cs @@ -170,7 +170,7 @@ private async ValueTask MoveNextAsync_Initialize_FromBeginningAsync( return true; } - private async ValueTask MoveNextAsync_Iniialize_FilterAsync( + private async ValueTask MoveNextAsync_Initialize_FilterAsync( OrderByQueryPartitionRangePageAsyncEnumerator uninitializedEnumerator, OrderByContinuationToken token) { @@ -308,7 +308,7 @@ await ParallelPrefetch.PrefetchInParallelAsync( (OrderByQueryPartitionRangePageAsyncEnumerator uninitializedEnumerator, OrderByContinuationToken token) = this.uninitializedEnumeratorsAndTokens.Dequeue(); bool movedNext = token is null ? await this.MoveNextAsync_Initialize_FromBeginningAsync(uninitializedEnumerator) - : await this.MoveNextAsync_Iniialize_FilterAsync(uninitializedEnumerator, token); + : await this.MoveNextAsync_Initialize_FilterAsync(uninitializedEnumerator, token); return movedNext; } @@ -335,6 +335,34 @@ private ValueTask MoveNextAsync_DrainPageAsync() // If the continuation isn't null // then mark the enumerator as unitialized and it will get requeueed on the next iteration with a fresh page. this.uninitializedEnumeratorsAndTokens.Enqueue((currentEnumerator, (OrderByContinuationToken)null)); + + // Use the token for the next page, since we fully drained the enumerator. + OrderByContinuationToken orderByContinuationToken = new OrderByContinuationToken( + new ParallelContinuationToken( + token: ((CosmosString)currentEnumerator.State.Value).Value, + range: currentEnumerator.Range.ToRange()), + orderByQueryResult.OrderByItems, + orderByQueryResult.Rid, + skipCount: 0, + filter: currentEnumerator.Filter); + + CosmosElement cosmosElementOrderByContinuationToken = OrderByContinuationToken.ToCosmosElement(orderByContinuationToken); + CosmosArray continuationTokenList = CosmosArray.Create(new List() { cosmosElementOrderByContinuationToken }); + + this.state = new QueryState(continuationTokenList); + + // Return a page of results + // No stats to report, since we already reported it when we moved to this page. + this.Current = TryCatch.FromResult( + new QueryPage( + documents: results.Select(result => result.Payload).ToList(), + requestCharge: 0, + activityId: default, + responseLengthInBytes: 0, + cosmosQueryExecutionInfo: default, + disallowContinuationTokenMessage: default, + state: this.state)); + return new ValueTask(true); } break; diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/DocumentContainerPartitionRangeEnumerator.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/DocumentContainerPartitionRangeEnumerator.cs index 0a02a55a46..706d62f991 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/DocumentContainerPartitionRangeEnumerator.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/DocumentContainerPartitionRangeEnumerator.cs @@ -31,7 +31,7 @@ public DocumentContainerPartitionRangeEnumerator( MaxExclusive = partitionKeyRangeId.ToString() }, cancellationToken, - state ?? new DocumentContainerState(resourceIdentifier: ResourceIdentifier.Empty)) + state ?? new DocumentContainerState(resourceIdentifier: ResourceId.Empty)) { this.documentContainer = documentContainer ?? throw new ArgumentNullException(nameof(documentContainer)); this.partitionKeyRangeId = partitionKeyRangeId; diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/DocumentContainerTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/DocumentContainerTests.cs index 54b6a7cf02..9ff25b0024 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/DocumentContainerTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/DocumentContainerTests.cs @@ -204,7 +204,7 @@ async Task AssertChildPartitionAsync(int partitionKeyRangeId) { DocumentContainerPage readFeedPage = await documentContainer.ReadFeedAsync( partitionKeyRangeId: partitionKeyRangeId, - resourceIdentifier: ResourceIdentifier.Empty, + resourceIdentifier: ResourceId.Empty, pageSize: 100, cancellationToken: default); @@ -296,7 +296,7 @@ async Task AssertChildPartitionAsync(int partitionKeyRangeId) { DocumentContainerPage page = await documentContainer.ReadFeedAsync( partitionKeyRangeId: partitionKeyRangeId, - resourceIdentifier: ResourceIdentifier.Empty, + resourceIdentifier: ResourceId.Empty, pageSize: 100, cancellationToken: default); diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/FlakyDocumentContainer.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/FlakyDocumentContainer.cs index bb786dde8b..315e83a66e 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/FlakyDocumentContainer.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/FlakyDocumentContainer.cs @@ -90,7 +90,7 @@ public Task> MonadicReadItemAsync( public Task> MonadicReadFeedAsync( int partitionKeyRangeId, - ResourceIdentifier resourceIdentifer, + ResourceId resourceIdentifer, int pageSize, CancellationToken cancellationToken) { 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 fd1d7ff68f..ee406e70aa 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,6 +8,7 @@ namespace Microsoft.Azure.Cosmos.Tests.Pagination using System.Collections; using System.Collections.Generic; using System.Linq; + using System.Reflection; using System.Threading; using System.Threading.Tasks; using Microsoft.Azure.Cosmos; @@ -28,6 +29,7 @@ namespace Microsoft.Azure.Cosmos.Tests.Pagination // Collection useful for mocking requests and repartitioning (splits / merge). internal sealed class InMemoryContainer : IMonadicDocumentContainer { + private static readonly ResourceIdentifier SeedIdentifier = ResourceIdentifier.Parse("AYIMAMmFOw8YAAAAAAAAAA=="); private static readonly PartitionKeyRange FullRange = new PartitionKeyRange() { MinInclusive = Documents.Routing.PartitionKeyInternal.MinimumInclusiveEffectivePartitionKey, @@ -261,7 +263,7 @@ static Task> CreateNotFoundException(CosmosElement partitionKey public Task> MonadicReadFeedAsync( int partitionKeyRangeId, - ResourceIdentifier resourceIdentifer, + ResourceId resourceIdentifer, int pageSize, CancellationToken cancellationToken) { @@ -386,10 +388,10 @@ public Task> MonadicQueryAsync( continuationResourceId = ((CosmosString)parsedContinuationToken["resourceId"]).Value; continuationSkipCount = (int)Number64.ToLong(((CosmosNumber64)parsedContinuationToken["skipCount"]).Value); + ResourceIdentifier continuationParsedResourceId = ResourceIdentifier.Parse(continuationResourceId); queryPageResults = queryPageResults.Where((Func)(c => { - ResourceIdentifier continuationParsedResourceId = ResourceIdentifier.Parse(continuationResourceId); - ResourceIdentifier documentResourceId = ResourceIdentifier.Parse(((CosmosString)((CosmosObject)c)["_rid"]).Value); + ResourceId documentResourceId = ResourceId.Parse(((CosmosString)((CosmosObject)c)["_rid"]).Value); return documentResourceId.Document >= continuationParsedResourceId.Document; })); @@ -639,23 +641,51 @@ public Records() public Record Add(int pkrangeid, CosmosObject payload) { - ResourceIdentifier currentResourceId; + // using pkrangeid for document collection since resource id doesnt serialize both document and pkrangeid. + ResourceId currentResourceId; if (this.Count == 0) { - currentResourceId = new ResourceIdentifier(database: 1, documentCollection: 1, partitionKeyRange: (ulong)pkrangeid, document: 1); + currentResourceId = ResourceId.Parse("AYIMAMmFOw8YAAAAAAAAAA=="); + + PropertyInfo documentProp = currentResourceId + .GetType() + .GetProperty("Document", BindingFlags.NonPublic | BindingFlags.Public | BindingFlags.Instance); + documentProp.SetValue(currentResourceId, (ulong)1); + + //PropertyInfo documentCollectionProp = currentResourceId + // .GetType() + // .GetProperty("DocumentCollection", BindingFlags.NonPublic | BindingFlags.Public | BindingFlags.Instance); + //documentCollectionProp.SetValue(currentResourceId, (uint)pkrangeid); + + PropertyInfo databaseProp = currentResourceId + .GetType() + .GetProperty("Database", BindingFlags.NonPublic | BindingFlags.Public | BindingFlags.Instance); + databaseProp.SetValue(currentResourceId, (uint)pkrangeid); } else { currentResourceId = this.storage[this.storage.Count - 1].ResourceIdentifier; } - ResourceIdentifier nextResourceId = new ResourceIdentifier( - database: currentResourceId.Database, - documentCollection: currentResourceId.DocumentCollection, - partitionKeyRange: (ulong)pkrangeid, - document: currentResourceId.Document); + ResourceId nextResourceId = ResourceId.Parse("AYIMAMmFOw8YAAAAAAAAAA=="); + { + PropertyInfo documentProp = nextResourceId + .GetType() + .GetProperty("Document", BindingFlags.NonPublic | BindingFlags.Public | BindingFlags.Instance); + documentProp.SetValue(nextResourceId, (ulong)(currentResourceId.Document + 1)); + + //PropertyInfo documentCollectionProp = nextResourceId + // .GetType() + // .GetProperty("DocumentCollection", BindingFlags.NonPublic | BindingFlags.Public | BindingFlags.Instance); + //documentCollectionProp.SetValue(nextResourceId, (uint)pkrangeid); + + PropertyInfo databaseProp = nextResourceId + .GetType() + .GetProperty("Database", BindingFlags.NonPublic | BindingFlags.Public | BindingFlags.Instance); + databaseProp.SetValue(nextResourceId, (uint)pkrangeid); + } - Record record = Record.Create(nextResourceId, payload); + Record record = new Record(nextResourceId, DateTime.UtcNow.Ticks, Guid.NewGuid().ToString(), payload); this.storage.Add(record); return record; } 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 5752d9603a..bbc8cfcdc5 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 @@ -92,17 +92,16 @@ public async Task OrderBy() [TestMethod] public async Task OrderByWithJoins() { - Random random = new Random(42); List documents = new List() { - CosmosObject.Parse($"{{\"pk\" : {random.Next()}, \"children\" : [\"Alice\", \"Bob\", \"Charlie\"]}}"), - CosmosObject.Parse($"{{\"pk\" : {random.Next()}, \"children\" : [\"Dave\", \"Eve\", \"Fancy\"]}}"), - CosmosObject.Parse($"{{\"pk\" : {random.Next()}, \"children\" : [\"George\", \"Henry\", \"Igor\"]}}"), - CosmosObject.Parse($"{{\"pk\" : {random.Next()}, \"children\" : [\"Jack\", \"Kim\", \"Levin\"]}}"), + CosmosObject.Parse($"{{\"pk\" : {1}, \"children\" : [\"Alice\", \"Bob\", \"Charlie\"]}}"), + CosmosObject.Parse($"{{\"pk\" : {2}, \"children\" : [\"Dave\", \"Eve\", \"Fancy\"]}}"), + CosmosObject.Parse($"{{\"pk\" : {3}, \"children\" : [\"George\", \"Henry\", \"Igor\"]}}"), + CosmosObject.Parse($"{{\"pk\" : {4}, \"children\" : [\"Jack\", \"Kim\", \"Levin\"]}}"), }; List documentsQueried = await ExecuteQueryAsync( - query: "SELECT d FROM c JOIN d in c.children ORDER BY c._ts", + query: "SELECT d FROM c JOIN d in c.children ORDER BY c.pk", documents: documents, pageSize: 2); @@ -182,12 +181,12 @@ private static async Task> ExecuteQueryAsync( { IDocumentContainer documentContainer = await CreateDocumentContainerAsync(documents); - //List resultsFromDrainWithoutState = await DrainWithoutStateAsync(query, documentContainer, pageSize); + List resultsFromDrainWithoutState = await DrainWithoutStateAsync(query, documentContainer, pageSize); List resultsFromDrainWithState = await DrainWithStateAsync(query, documentContainer, pageSize); - //Assert.IsTrue(resultsFromDrainWithoutState.SequenceEqual(resultsFromDrainWithState)); + Assert.IsTrue(resultsFromDrainWithoutState.SequenceEqual(resultsFromDrainWithState)); - return resultsFromDrainWithState; + return resultsFromDrainWithoutState; } [TestMethod] 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 b2e68ee272..c4f754a8ee 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 @@ -112,7 +112,7 @@ public override IReadOnlyList GetRecordsFromPage(QueryPage page) foreach (CosmosElement element in page.Documents) { CosmosObject document = (CosmosObject)element; - ResourceIdentifier resourceIdentifier = ResourceIdentifier.Parse(((CosmosString)document["_rid"]).Value); + ResourceId resourceIdentifier = ResourceId.Parse(((CosmosString)document["_rid"]).Value); long timestamp = Number64.ToLong(((CosmosNumber)document["_ts"]).Value); string identifer = ((CosmosString)document["id"]).Value; From eba6d9189a5bf563f4f3332ccf38ca4a2ed4ed4d Mon Sep 17 00:00:00 2001 From: Brandon Chong Date: Sat, 3 Oct 2020 02:22:47 -0700 Subject: [PATCH 79/85] left some comments --- .../OrderByCrossPartitionQueryPipelineStage.cs | 4 ++++ .../Pagination/InMemoryContainer.cs | 12 +----------- 2 files changed, 5 insertions(+), 11 deletions(-) diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/OrderBy/OrderByCrossPartitionQueryPipelineStage.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/OrderBy/OrderByCrossPartitionQueryPipelineStage.cs index a755e1e5c9..dd91f6f090 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/OrderBy/OrderByCrossPartitionQueryPipelineStage.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/OrderBy/OrderByCrossPartitionQueryPipelineStage.cs @@ -365,6 +365,10 @@ private ValueTask MoveNextAsync_DrainPageAsync() return new ValueTask(true); } + // Todo: we can optimize this by having a special "Done" continuation token + // so we don't grab a full page and filter it through + // but this would break older clients, so wait for a compute only fork. + break; } 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 ee406e70aa..2ece7482c1 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 @@ -641,7 +641,7 @@ public Records() public Record Add(int pkrangeid, CosmosObject payload) { - // using pkrangeid for document collection since resource id doesnt serialize both document and pkrangeid. + // using pkrangeid for database since resource id doesnt serialize both document and pkrangeid. ResourceId currentResourceId; if (this.Count == 0) { @@ -652,11 +652,6 @@ public Record Add(int pkrangeid, CosmosObject payload) .GetProperty("Document", BindingFlags.NonPublic | BindingFlags.Public | BindingFlags.Instance); documentProp.SetValue(currentResourceId, (ulong)1); - //PropertyInfo documentCollectionProp = currentResourceId - // .GetType() - // .GetProperty("DocumentCollection", BindingFlags.NonPublic | BindingFlags.Public | BindingFlags.Instance); - //documentCollectionProp.SetValue(currentResourceId, (uint)pkrangeid); - PropertyInfo databaseProp = currentResourceId .GetType() .GetProperty("Database", BindingFlags.NonPublic | BindingFlags.Public | BindingFlags.Instance); @@ -674,11 +669,6 @@ public Record Add(int pkrangeid, CosmosObject payload) .GetProperty("Document", BindingFlags.NonPublic | BindingFlags.Public | BindingFlags.Instance); documentProp.SetValue(nextResourceId, (ulong)(currentResourceId.Document + 1)); - //PropertyInfo documentCollectionProp = nextResourceId - // .GetType() - // .GetProperty("DocumentCollection", BindingFlags.NonPublic | BindingFlags.Public | BindingFlags.Instance); - //documentCollectionProp.SetValue(nextResourceId, (uint)pkrangeid); - PropertyInfo databaseProp = nextResourceId .GetType() .GetProperty("Database", BindingFlags.NonPublic | BindingFlags.Public | BindingFlags.Instance); From 181e0af95d88763328b0245fcaf26961779a837a Mon Sep 17 00:00:00 2001 From: Brandon Chong Date: Sat, 3 Oct 2020 22:00:34 -0700 Subject: [PATCH 80/85] fixed bugs --- .../src/Pagination/ParallelPrefetch.cs | 8 +++--- ...OrderByCrossPartitionQueryPipelineStage.cs | 25 +++++++++++++++++++ .../Query/Core/Pipeline/PipelineFactory.cs | 2 +- .../src/Query/v3Query/QueryIterator.cs | 21 +++------------- Microsoft.Azure.Cosmos/src/UInt128.cs | 2 ++ .../CosmosItemLinqTests.cs | 7 ++++-- .../Query/OrderByQueryTests.cs | 8 ++++-- .../Pagination/InMemoryContainer.cs | 16 ++++++------ ...ByCrossPartitionQueryPipelineStageTests.cs | 8 ++++-- 9 files changed, 62 insertions(+), 35 deletions(-) diff --git a/Microsoft.Azure.Cosmos/src/Pagination/ParallelPrefetch.cs b/Microsoft.Azure.Cosmos/src/Pagination/ParallelPrefetch.cs index e1138d9432..f2edaa0d64 100644 --- a/Microsoft.Azure.Cosmos/src/Pagination/ParallelPrefetch.cs +++ b/Microsoft.Azure.Cosmos/src/Pagination/ParallelPrefetch.cs @@ -22,11 +22,13 @@ public static async Task PrefetchInParallelAsync(IEnumerable prefet IEnumerator prefetchersEnumerator = prefetchers.GetEnumerator(); for (int i = 0; i < maxConcurrency; i++) { - if (prefetchersEnumerator.MoveNext()) + if (!prefetchersEnumerator.MoveNext()) { - IPrefetcher prefetcher = prefetchersEnumerator.Current; - tasks.Add(Task.Run(async () => await prefetcher.PrefetchAsync(cancellationToken))); + break; } + + IPrefetcher prefetcher = prefetchersEnumerator.Current; + tasks.Add(Task.Run(async () => await prefetcher.PrefetchAsync(cancellationToken))); } while (tasks.Count != 0) diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/OrderBy/OrderByCrossPartitionQueryPipelineStage.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/OrderBy/OrderByCrossPartitionQueryPipelineStage.cs index dd91f6f090..3512ec6191 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/OrderBy/OrderByCrossPartitionQueryPipelineStage.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/OrderBy/OrderByCrossPartitionQueryPipelineStage.cs @@ -55,6 +55,7 @@ internal sealed class OrderByCrossPartitionQueryPipelineStage : IQueryPipelineSt private CancellationToken cancellationToken; private QueryState state; + private bool returnedFinalPage; private static class Expressions { @@ -146,6 +147,7 @@ private async ValueTask MoveNextAsync_Initialize_FromBeginningAsync( cosmosQueryExecutionInfo: default, disallowContinuationTokenMessage: default, state: null)); + this.returnedFinalPage = true; return true; } } @@ -223,6 +225,7 @@ private async ValueTask MoveNextAsync_Initialize_FilterAsync( cosmosQueryExecutionInfo: default, disallowContinuationTokenMessage: default, state: null)); + this.returnedFinalPage = true; return true; } } @@ -416,6 +419,12 @@ private ValueTask MoveNextAsync_DrainPageAsync() cosmosQueryExecutionInfo: default, disallowContinuationTokenMessage: default, state: this.state)); + + if (state == null) + { + this.returnedFinalPage = true; + } + return new ValueTask(true); } @@ -454,6 +463,22 @@ public ValueTask MoveNextAsync() if (this.enumerators.Count == 0) { + if (!this.returnedFinalPage) + { + // return a empty page with null continuation token + this.Current = TryCatch.FromResult( + new QueryPage( + documents: EmptyPage, + requestCharge: 0, + activityId: Guid.NewGuid().ToString(), + responseLengthInBytes: 0, + cosmosQueryExecutionInfo: default, + disallowContinuationTokenMessage: default, + state: null)); + this.returnedFinalPage = true; + return new ValueTask(true); + } + // Finished draining. return new ValueTask(false); } diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/PipelineFactory.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/PipelineFactory.cs index bd82a12d0c..a9fb9b78fb 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/PipelineFactory.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/PipelineFactory.cs @@ -60,7 +60,7 @@ public static TryCatch MonadicCreate( throw new ArgumentNullException(nameof(queryInfo)); } - sqlQuerySpec = !string.IsNullOrEmpty(queryInfo.RewrittenQuery) ? new SqlQuerySpec(queryInfo.RewrittenQuery) : sqlQuerySpec; + sqlQuerySpec = !string.IsNullOrEmpty(queryInfo.RewrittenQuery) ? new SqlQuerySpec(queryInfo.RewrittenQuery, sqlQuerySpec.Parameters) : sqlQuerySpec; MonadicCreatePipelineStage monadicCreatePipelineStage; if (queryInfo.HasOrderBy) diff --git a/Microsoft.Azure.Cosmos/src/Query/v3Query/QueryIterator.cs b/Microsoft.Azure.Cosmos/src/Query/v3Query/QueryIterator.cs index 1694ca6de0..6a6cb6390a 100644 --- a/Microsoft.Azure.Cosmos/src/Query/v3Query/QueryIterator.cs +++ b/Microsoft.Azure.Cosmos/src/Query/v3Query/QueryIterator.cs @@ -215,31 +215,16 @@ public override async Task ReadNextAsync(CancellationToken canc } CosmosException cosmosException = ExceptionToCosmosException.CreateFromException(tryGetQueryPage.Exception); - SubStatusCodes subStatusCode; - if (Enum.IsDefined(typeof(SubStatusCodes), cosmosException.SubStatusCode)) - { - subStatusCode = (SubStatusCodes)cosmosException.SubStatusCode; - } - else - { - subStatusCode = Documents.SubStatusCodes.Unknown; - } return QueryResponse.CreateFailure( statusCode: cosmosException.StatusCode, cosmosException: cosmosException, requestMessage: null, diagnostics: diagnostics, - responseHeaders: new CosmosQueryResponseMessageHeaders( - continauationToken: default, - disallowContinuationTokenMessage: default, + responseHeaders: CosmosQueryResponseMessageHeaders.ConvertToQueryHeaders( + cosmosException.Headers, this.cosmosQueryContext.ResourceTypeEnum, - this.cosmosQueryContext.ContainerResourceId) - { - RequestCharge = cosmosException.RequestCharge, - ActivityId = cosmosException.ActivityId, - SubStatusCode = subStatusCode, - }); + this.cosmosQueryContext.ContainerResourceId)); } } diff --git a/Microsoft.Azure.Cosmos/src/UInt128.cs b/Microsoft.Azure.Cosmos/src/UInt128.cs index d11b5c9025..fb76ee004b 100644 --- a/Microsoft.Azure.Cosmos/src/UInt128.cs +++ b/Microsoft.Azure.Cosmos/src/UInt128.cs @@ -499,6 +499,8 @@ public override int GetHashCode() public override string ToString() { byte[] bytes = UInt128.ToByteArray(this); + // Reverse the bytes and make it big endian so that the lex sort is equivalent to the numeric sort. + Array.Reverse(bytes, 0, bytes.Length); return BitConverter.ToString(bytes); } diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CosmosItemLinqTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CosmosItemLinqTests.cs index 4879691737..1b10ee0bad 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CosmosItemLinqTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CosmosItemLinqTests.cs @@ -577,8 +577,11 @@ public async Task LinqParameterisedTest2() Assert.AreEqual(0, (await this.FetchResults(queryDefinition)).Count); //Test orderby, skip, take, distinct, these will not get parameterized. - queryText = "SELECT VALUE root FROM root WHERE (root[\"CamelCase\"] = @param1) ORDER BY" + - " root[\"taskNum\"] ASC OFFSET @param2 LIMIT @param3"; + queryText = "SELECT VALUE root " + + "FROM root " + + "WHERE (root[\"CamelCase\"] = @param1) " + + "ORDER BY root[\"taskNum\"] ASC " + + "OFFSET @param2 LIMIT @param3"; queriable = linqQueryable .Where(item => item.CamelCase == camelCase) .OrderBy(item => item.taskNum) diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/Query/OrderByQueryTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/Query/OrderByQueryTests.cs index 71fcabf791..56ab5ff3ec 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/Query/OrderByQueryTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/Query/OrderByQueryTests.cs @@ -427,8 +427,12 @@ await container.GetItemQueryIterator( #endregion FeedResponse responseWithEmptyContinuationExpected = await container.GetItemQueryIterator( - string.Format(CultureInfo.InvariantCulture, "SELECT TOP 1 * FROM r ORDER BY r.{0}", partitionKey), - requestOptions: new QueryRequestOptions() { MaxConcurrency = 10, MaxItemCount = -1 }).ReadNextAsync(); + $"SELECT TOP 0 * FROM r", + requestOptions: new QueryRequestOptions() + { + MaxConcurrency = 10, + MaxItemCount = -1 + }).ReadNextAsync(); Assert.AreEqual(null, responseWithEmptyContinuationExpected.ContinuationToken); 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 2ece7482c1..8dca3c8683 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 @@ -29,7 +29,6 @@ namespace Microsoft.Azure.Cosmos.Tests.Pagination // Collection useful for mocking requests and repartitioning (splits / merge). internal sealed class InMemoryContainer : IMonadicDocumentContainer { - private static readonly ResourceIdentifier SeedIdentifier = ResourceIdentifier.Parse("AYIMAMmFOw8YAAAAAAAAAA=="); private static readonly PartitionKeyRange FullRange = new PartitionKeyRange() { MinInclusive = Documents.Routing.PartitionKeyInternal.MinimumInclusiveEffectivePartitionKey, @@ -395,12 +394,15 @@ public Task> MonadicQueryAsync( return documentResourceId.Document >= continuationParsedResourceId.Document; })); - if (queryPageResults.FirstOrDefault() is CosmosObject firstDocument) + for (int i = 0; i < continuationSkipCount; i++) { - string currentResourceId = ((CosmosString)firstDocument["_rid"]).Value; - if (currentResourceId == continuationResourceId) + if (queryPageResults.FirstOrDefault() is CosmosObject firstDocument) { - queryPageResults = queryPageResults.Skip(continuationSkipCount); + string currentResourceId = ((CosmosString)firstDocument["_rid"]).Value; + if (currentResourceId == continuationResourceId) + { + queryPageResults = queryPageResults.Skip(1); + } } } } @@ -655,7 +657,7 @@ public Record Add(int pkrangeid, CosmosObject payload) PropertyInfo databaseProp = currentResourceId .GetType() .GetProperty("Database", BindingFlags.NonPublic | BindingFlags.Public | BindingFlags.Instance); - databaseProp.SetValue(currentResourceId, (uint)pkrangeid); + databaseProp.SetValue(currentResourceId, (uint)pkrangeid + 1); } else { @@ -672,7 +674,7 @@ public Record Add(int pkrangeid, CosmosObject payload) PropertyInfo databaseProp = nextResourceId .GetType() .GetProperty("Database", BindingFlags.NonPublic | BindingFlags.Public | BindingFlags.Instance); - databaseProp.SetValue(nextResourceId, (uint)pkrangeid); + databaseProp.SetValue(nextResourceId, (uint)pkrangeid + 1); } Record record = new Record(nextResourceId, DateTime.UtcNow.Ticks, Guid.NewGuid().ToString(), payload); diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/OrderByCrossPartitionQueryPipelineStageTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/OrderByCrossPartitionQueryPipelineStageTests.cs index e00cb7e866..acf94817ad 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/OrderByCrossPartitionQueryPipelineStageTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/OrderByCrossPartitionQueryPipelineStageTests.cs @@ -410,7 +410,11 @@ FROM c maxConcurrency: 10, cancellationToken: default, continuationToken: queryState?.Value); - Assert.IsTrue(monadicCreate.Succeeded); + if (monadicCreate.Failed) + { + Assert.Fail(monadicCreate.Exception.ToString()); + } + IQueryPipelineStage queryPipelineStage = monadicCreate.Result; QueryPage queryPage; @@ -436,7 +440,7 @@ FROM c documentContainer.SplitAsync(int.Parse(randomRange.Id), cancellationToken: default).Wait(); } while (queryState != null); - Assert.AreEqual(numItems, documents.Count); + Assert.AreEqual(numItems, documents.Count, $"Failed with seed: {seed}. got {documents.Count} documents when {numItems} was expected."); Assert.IsTrue(documents.OrderBy(document => ((CosmosObject)document)["_ts"]).ToList().SequenceEqual(documents), $"Failed with seed: {seed}"); } From 0cdb13600a7b7f523d871d2753d3fa12544d952d Mon Sep 17 00:00:00 2001 From: Brandon Chong Date: Sun, 4 Oct 2020 00:26:58 -0700 Subject: [PATCH 81/85] added split tests for parallel and patched bugs in the inmemorycontainer --- .../src/Routing/PartitionKeyHashRange.cs | 7 ++ Microsoft.Azure.Cosmos/src/UInt128.cs | 4 +- .../Pagination/InMemoryContainer.cs | 24 ++--- ...elCrossPartitionQueryPipelineStageTests.cs | 91 +++++++++++++++++++ 4 files changed, 114 insertions(+), 12 deletions(-) diff --git a/Microsoft.Azure.Cosmos/src/Routing/PartitionKeyHashRange.cs b/Microsoft.Azure.Cosmos/src/Routing/PartitionKeyHashRange.cs index 34d16bee88..8327102cbd 100644 --- a/Microsoft.Azure.Cosmos/src/Routing/PartitionKeyHashRange.cs +++ b/Microsoft.Azure.Cosmos/src/Routing/PartitionKeyHashRange.cs @@ -37,6 +37,13 @@ public bool Contains(PartitionKeyHash partitionKeyHash) return rangeStartsBefore && rangeEndsAfter; } + public bool Contains(PartitionKeyHashRange partitionKeyHashRange) + { + bool rangeStartsBefore = !this.StartInclusive.HasValue || (partitionKeyHashRange.StartInclusive.HasValue && (this.StartInclusive.Value <= partitionKeyHashRange.StartInclusive.Value)); + bool rangeEndsAfter = !this.EndExclusive.HasValue || (partitionKeyHashRange.EndExclusive.HasValue && (partitionKeyHashRange.EndExclusive.Value <= this.EndExclusive.Value)); + return rangeStartsBefore && rangeEndsAfter; + } + public int CompareTo(PartitionKeyHashRange other) { // Provide a total sort order by first comparing on the start and then going to the end. diff --git a/Microsoft.Azure.Cosmos/src/UInt128.cs b/Microsoft.Azure.Cosmos/src/UInt128.cs index fb76ee004b..5eb43009f8 100644 --- a/Microsoft.Azure.Cosmos/src/UInt128.cs +++ b/Microsoft.Azure.Cosmos/src/UInt128.cs @@ -536,7 +536,7 @@ public static bool TryParse(string value, out UInt128 uInt128) return false; } - byte[] bytes = new byte[UInt128.Length]; + Span bytes = stackalloc byte[UInt128.Length]; for (int index = 0; index < UInt128.Length; index++) { if (!byte.TryParse(hexPairs[index], System.Globalization.NumberStyles.HexNumber, null, out byte parsedBytes)) @@ -548,6 +548,8 @@ public static bool TryParse(string value, out UInt128 uInt128) bytes[index] = parsedBytes; } + bytes.Reverse(); + uInt128 = UInt128.FromByteArray(bytes); return true; } 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 8dca3c8683..25a1ecb4af 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 @@ -93,30 +93,32 @@ PartitionKeyRange CreateRangeFromId(int id) return TryCatch>.FromResult(ranges); } - int partitionKeyRangeId; + if (partitionKeyRange.Id == null) { + // look for overlapping epk ranges. PartitionKeyHash? start = partitionKeyRange.MinInclusive == string.Empty ? (PartitionKeyHash?)null : PartitionKeyHash.Parse(partitionKeyRange.MinInclusive); PartitionKeyHash? end = partitionKeyRange.MaxExclusive == string.Empty ? (PartitionKeyHash?)null : PartitionKeyHash.Parse(partitionKeyRange.MaxExclusive); PartitionKeyHashRange hashRange = new PartitionKeyHashRange(start, end); - IEnumerable> kvps = this.partitionKeyRangeIdToHashRange.Where(kvp => kvp.Value.Equals(hashRange)); - if (!kvps.Any()) + List overlappedIds = this.partitionKeyRangeIdToHashRange + .Where(kvp => hashRange.Contains(kvp.Value)) + .Select(kvp => CreateRangeFromId(kvp.Key)) + .ToList(); + if (overlappedIds.Count == 0) { return TryCatch>.FromException( new KeyNotFoundException( $"PartitionKeyRangeId: {hashRange} does not exist.")); } - partitionKeyRangeId = kvps.First().Key; + return TryCatch>.FromResult(overlappedIds); } - else + + if (!int.TryParse(partitionKeyRange.Id, out int partitionKeyRangeId)) { - if (!int.TryParse(partitionKeyRange.Id, out partitionKeyRangeId)) - { - return TryCatch>.FromException( - new FormatException( - $"PartitionKeyRangeId: {partitionKeyRange.Id} is not an integer.")); - } + return TryCatch>.FromException( + new FormatException( + $"PartitionKeyRangeId: {partitionKeyRange.Id} is not an integer.")); } if (!this.parentToChildMapping.TryGetValue(partitionKeyRangeId, out (int left, int right) children)) 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 fee189643a..3f40dc87d1 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 @@ -4,6 +4,7 @@ namespace Microsoft.Azure.Cosmos.Tests.Query.Pipeline { + using System; using System.Collections.Generic; using System.Threading.Tasks; using Microsoft.Azure.Cosmos.CosmosElements; @@ -210,6 +211,96 @@ public async Task TestDrainFully_WithStateResume() Assert.AreEqual(numItems, documents.Count); } + [TestMethod] + public async Task TestDrainFully_WithStateResume_WithSplitAsync() + { + int numItems = 1000; + IDocumentContainer documentContainer = await CreateDocumentContainerAsync(numItems); + + Random random = new Random(); + List documents = new List(); + + QueryState queryState = null; + do + { + TryCatch monadicCreate = ParallelCrossPartitionQueryPipelineStage.MonadicCreate( + documentContainer: documentContainer, + sqlQuerySpec: new SqlQuerySpec("SELECT * FROM c"), + targetRanges: await documentContainer.GetFeedRangesAsync(cancellationToken: default), + pageSize: 10, + maxConcurrency: 10, + cancellationToken: default, + continuationToken: queryState?.Value); + if (monadicCreate.Failed) + { + Assert.Fail(); + } + Assert.IsTrue(monadicCreate.Succeeded); + IQueryPipelineStage queryPipelineStage = monadicCreate.Result; + + Assert.IsTrue(await queryPipelineStage.MoveNextAsync()); + TryCatch tryGetQueryPage = queryPipelineStage.Current; + Assert.IsTrue(tryGetQueryPage.Succeeded); + + QueryPage queryPage = tryGetQueryPage.Result; + documents.AddRange(queryPage.Documents); + + queryState = queryPage.State; + + if (random.Next() % 4 == 0) + { + // Can not always split otherwise the split handling code will livelock trying to split proof every partition in a cycle. + List ranges = documentContainer.GetFeedRangesAsync(cancellationToken: default).Result; + PartitionKeyRange randomRange = ranges[random.Next(ranges.Count)]; + documentContainer.SplitAsync(int.Parse(randomRange.Id), cancellationToken: default).Wait(); + } + } while (queryState != null); + + Assert.AreEqual(numItems, documents.Count); + } + + [TestMethod] + public async Task TestDrainFully_StartFromBegining_WithSplits_Async() + { + int numItems = 1000; + IDocumentContainer documentContainer = await CreateDocumentContainerAsync(numItems); + + TryCatch monadicCreate = ParallelCrossPartitionQueryPipelineStage.MonadicCreate( + documentContainer: documentContainer, + sqlQuerySpec: new SqlQuerySpec("SELECT * FROM c"), + targetRanges: await documentContainer.GetFeedRangesAsync(cancellationToken: default), + pageSize: 10, + maxConcurrency: 10, + cancellationToken: default, + continuationToken: default); + Assert.IsTrue(monadicCreate.Succeeded); + IQueryPipelineStage queryPipelineStage = monadicCreate.Result; + + Random random = new Random(); + List documents = new List(); + while (await queryPipelineStage.MoveNextAsync()) + { + TryCatch tryGetQueryPage = queryPipelineStage.Current; + if(tryGetQueryPage.Failed) + { + Assert.Fail(tryGetQueryPage.Exception.ToString()); + } + + QueryPage queryPage = tryGetQueryPage.Result; + documents.AddRange(queryPage.Documents); + + if (random.Next() % 4 == 0) + { + // Can not always split otherwise the split handling code will livelock trying to split proof every partition in a cycle. + List ranges = documentContainer.GetFeedRangesAsync(cancellationToken: default).Result; + PartitionKeyRange randomRange = ranges[random.Next(ranges.Count)]; + documentContainer.SplitAsync(int.Parse(randomRange.Id), cancellationToken: default).Wait(); + } + } + + Assert.AreEqual(numItems, documents.Count); + } + private static async Task CreateDocumentContainerAsync( int numItems, FlakyDocumentContainer.FailureConfigs failureConfigs = null) From 88f7ca2da303b9320a18bf36702cecdcce5def3b Mon Sep 17 00:00:00 2001 From: Brandon Chong Date: Sun, 4 Oct 2020 03:42:04 -0700 Subject: [PATCH 82/85] fixed more tests --- .../CosmosQueryResponseMessageHeaders.cs | 8 +- .../NetworkAttachedDocumentContainer.cs | 20 ++- .../src/Query/v3Query/QueryIterator.cs | 4 +- .../CosmosNotFoundTests.cs | 1 - ...DistinctHashBaselineTests.ElementsHash.xml | 18 +-- ...nctHashBaselineTests.ExtendedTypesHash.xml | 6 +- .../DistinctHashBaselineTests.NumbersHash.xml | 16 +-- ...tHashBaselineTests.WrappedElementsHash.xml | 120 +++++++++--------- .../PartitionKeyHashBaselineTest.Numbers.xml | 64 +++++----- ...artitionKeyHashBaselineTest.Singletons.xml | 16 +-- .../PartitionKeyHashBaselineTest.Strings.xml | 24 ++-- 11 files changed, 156 insertions(+), 141 deletions(-) diff --git a/Microsoft.Azure.Cosmos/src/Headers/CosmosQueryResponseMessageHeaders.cs b/Microsoft.Azure.Cosmos/src/Headers/CosmosQueryResponseMessageHeaders.cs index 0fca37356b..9324474f53 100644 --- a/Microsoft.Azure.Cosmos/src/Headers/CosmosQueryResponseMessageHeaders.cs +++ b/Microsoft.Azure.Cosmos/src/Headers/CosmosQueryResponseMessageHeaders.cs @@ -79,7 +79,9 @@ internal CosmosQueryResponseMessageHeaders CloneKnownProperties( internal static CosmosQueryResponseMessageHeaders ConvertToQueryHeaders( Headers sourceHeaders, ResourceType resourceType, - string containerRid) + string containerRid, + int? substatusCode = null, + string activityId = null) { if (sourceHeaders == null) { @@ -98,11 +100,11 @@ internal static CosmosQueryResponseMessageHeaders ConvertToQueryHeaders( { RequestCharge = sourceHeaders.RequestCharge, ContentLength = sourceHeaders.ContentLength, - ActivityId = sourceHeaders.ActivityId, + ActivityId = sourceHeaders.ActivityId ?? activityId, ETag = sourceHeaders.ETag, Location = sourceHeaders.Location, RetryAfterLiteral = sourceHeaders.RetryAfterLiteral, - SubStatusCodeLiteral = sourceHeaders.SubStatusCodeLiteral, + SubStatusCodeLiteral = sourceHeaders.SubStatusCodeLiteral ?? (substatusCode.HasValue ? substatusCode.Value.ToString() : null), ContentType = sourceHeaders.ContentType, QueryMetricsText = sourceHeaders.QueryMetricsText }; diff --git a/Microsoft.Azure.Cosmos/src/Pagination/NetworkAttachedDocumentContainer.cs b/Microsoft.Azure.Cosmos/src/Pagination/NetworkAttachedDocumentContainer.cs index 9b8b6019a4..382e641e14 100644 --- a/Microsoft.Azure.Cosmos/src/Pagination/NetworkAttachedDocumentContainer.cs +++ b/Microsoft.Azure.Cosmos/src/Pagination/NetworkAttachedDocumentContainer.cs @@ -6,6 +6,7 @@ namespace Microsoft.Azure.Cosmos.Pagination { using System; using System.Collections.Generic; + using System.ComponentModel; using System.Linq; using System.Net; using System.Threading; @@ -16,6 +17,7 @@ namespace Microsoft.Azure.Cosmos.Pagination using Microsoft.Azure.Cosmos.Query.Core.Pipeline; using Microsoft.Azure.Cosmos.Query.Core.Pipeline.CrossPartition; using Microsoft.Azure.Cosmos.Query.Core.QueryClient; + using Microsoft.Azure.Cosmos.Resource.CosmosExceptions; using Microsoft.Azure.Documents; internal sealed class NetworkAttachedDocumentContainer : IMonadicDocumentContainer @@ -82,13 +84,23 @@ public Task> MonadicReadItemAsync( public Task>> MonadicGetFeedRangesAsync( CancellationToken cancellationToken) => this.MonadicGetChildRangeAsync(FullRange, cancellationToken); - public Task>> MonadicGetChildRangeAsync( + public async Task>> MonadicGetChildRangeAsync( PartitionKeyRange partitionKeyRange, - CancellationToken cancellationToken) => this.cosmosQueryContext.QueryClient.TryGetOverlappingRangesAsync( + CancellationToken cancellationToken) + { + IReadOnlyList overlappingRanges = await this.cosmosQueryContext.QueryClient.TryGetOverlappingRangesAsync( this.cosmosQueryContext.ContainerResourceId, partitionKeyRange.ToRange(), - forceRefresh: true) - .ContinueWith(antecedent => TryCatch>.FromResult(antecedent.Result.ToList())); + forceRefresh: true); + if (overlappingRanges == null) + { + return TryCatch>.FromException( + CosmosExceptionFactory.CreateNotFoundException( + message: $"Container with resource id: {this.cosmosQueryContext.ContainerResourceId} was not found")); + } + + return TryCatch>.FromResult(overlappingRanges.ToList()); + } public Task> MonadicReadFeedAsync( int partitionKeyRangeId, diff --git a/Microsoft.Azure.Cosmos/src/Query/v3Query/QueryIterator.cs b/Microsoft.Azure.Cosmos/src/Query/v3Query/QueryIterator.cs index 6a6cb6390a..a67c8a8549 100644 --- a/Microsoft.Azure.Cosmos/src/Query/v3Query/QueryIterator.cs +++ b/Microsoft.Azure.Cosmos/src/Query/v3Query/QueryIterator.cs @@ -224,7 +224,9 @@ public override async Task ReadNextAsync(CancellationToken canc responseHeaders: CosmosQueryResponseMessageHeaders.ConvertToQueryHeaders( cosmosException.Headers, this.cosmosQueryContext.ResourceTypeEnum, - this.cosmosQueryContext.ContainerResourceId)); + this.cosmosQueryContext.ContainerResourceId, + cosmosException.SubStatusCode, + cosmosException.ActivityId)); } } diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CosmosNotFoundTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CosmosNotFoundTests.cs index 2d1999fe77..a136cfe4b3 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CosmosNotFoundTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CosmosNotFoundTests.cs @@ -166,7 +166,6 @@ private async Task VerifyQueryNotFoundResponse(FeedIterator iterator) ResponseMessage response = await iterator.ReadNextAsync(); Assert.IsNotNull(response); Assert.AreEqual(HttpStatusCode.NotFound, response.StatusCode); - Assert.IsFalse(iterator.HasMoreResults); } private void VerifyNotFoundResponse(ResponseMessage response) diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/BaselineTest/TestBaseline/DistinctHashBaselineTests.ElementsHash.xml b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/BaselineTest/TestBaseline/DistinctHashBaselineTests.ElementsHash.xml index 425589020e..47c7348fbe 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/BaselineTest/TestBaseline/DistinctHashBaselineTests.ElementsHash.xml +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/BaselineTest/TestBaseline/DistinctHashBaselineTests.ElementsHash.xml @@ -4,7 +4,7 @@ - + @@ -12,7 +12,7 @@ - + @@ -20,7 +20,7 @@ - + @@ -28,7 +28,7 @@ - + @@ -36,7 +36,7 @@ - + @@ -44,7 +44,7 @@ - + @@ -52,7 +52,7 @@ - + @@ -60,7 +60,7 @@ - + @@ -68,7 +68,7 @@ - + \ No newline at end of file diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/BaselineTest/TestBaseline/DistinctHashBaselineTests.ExtendedTypesHash.xml b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/BaselineTest/TestBaseline/DistinctHashBaselineTests.ExtendedTypesHash.xml index 2f8eec79a5..4ba1549f86 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/BaselineTest/TestBaseline/DistinctHashBaselineTests.ExtendedTypesHash.xml +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/BaselineTest/TestBaseline/DistinctHashBaselineTests.ExtendedTypesHash.xml @@ -4,7 +4,7 @@ - + @@ -12,7 +12,7 @@ - + @@ -20,7 +20,7 @@ - + \ No newline at end of file diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/BaselineTest/TestBaseline/DistinctHashBaselineTests.NumbersHash.xml b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/BaselineTest/TestBaseline/DistinctHashBaselineTests.NumbersHash.xml index d4aee5b523..7ae9a6f0e1 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/BaselineTest/TestBaseline/DistinctHashBaselineTests.NumbersHash.xml +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/BaselineTest/TestBaseline/DistinctHashBaselineTests.NumbersHash.xml @@ -4,7 +4,7 @@ - + @@ -12,7 +12,7 @@ - + @@ -20,7 +20,7 @@ - + @@ -28,7 +28,7 @@ - + @@ -36,7 +36,7 @@ - + @@ -44,7 +44,7 @@ - + @@ -52,7 +52,7 @@ - + @@ -60,7 +60,7 @@ - + \ No newline at end of file diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/BaselineTest/TestBaseline/DistinctHashBaselineTests.WrappedElementsHash.xml b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/BaselineTest/TestBaseline/DistinctHashBaselineTests.WrappedElementsHash.xml index 559f6d4a1b..3e78d1df21 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/BaselineTest/TestBaseline/DistinctHashBaselineTests.WrappedElementsHash.xml +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/BaselineTest/TestBaseline/DistinctHashBaselineTests.WrappedElementsHash.xml @@ -4,7 +4,7 @@ - + @@ -12,7 +12,7 @@ - + @@ -20,7 +20,7 @@ - + @@ -28,7 +28,7 @@ - + @@ -36,7 +36,7 @@ - + @@ -44,7 +44,7 @@ - + @@ -52,7 +52,7 @@ - + @@ -60,7 +60,7 @@ - + @@ -68,7 +68,7 @@ - + @@ -76,7 +76,7 @@ - + @@ -84,7 +84,7 @@ - + @@ -92,7 +92,7 @@ - + @@ -100,7 +100,7 @@ - + @@ -108,7 +108,7 @@ - + @@ -116,7 +116,7 @@ - + @@ -124,7 +124,7 @@ - + @@ -132,7 +132,7 @@ - + @@ -140,7 +140,7 @@ - + @@ -148,7 +148,7 @@ - + @@ -156,7 +156,7 @@ - + @@ -164,7 +164,7 @@ - + @@ -172,7 +172,7 @@ - + @@ -180,7 +180,7 @@ - + @@ -188,7 +188,7 @@ - + @@ -196,7 +196,7 @@ - + @@ -204,7 +204,7 @@ - + @@ -212,7 +212,7 @@ - + @@ -220,7 +220,7 @@ - + @@ -228,7 +228,7 @@ - + @@ -236,7 +236,7 @@ - + @@ -244,7 +244,7 @@ - + @@ -252,7 +252,7 @@ - + @@ -260,7 +260,7 @@ - + @@ -268,7 +268,7 @@ - + @@ -276,7 +276,7 @@ - + @@ -284,7 +284,7 @@ - + @@ -292,7 +292,7 @@ - + @@ -300,7 +300,7 @@ - + @@ -308,7 +308,7 @@ - + @@ -316,7 +316,7 @@ - + @@ -324,7 +324,7 @@ - + @@ -332,7 +332,7 @@ - + @@ -340,7 +340,7 @@ - + @@ -348,7 +348,7 @@ - + @@ -356,7 +356,7 @@ - + @@ -364,7 +364,7 @@ - + @@ -372,7 +372,7 @@ - + @@ -380,7 +380,7 @@ - + @@ -388,7 +388,7 @@ - + @@ -396,7 +396,7 @@ - + @@ -404,7 +404,7 @@ - + @@ -412,7 +412,7 @@ - + @@ -420,7 +420,7 @@ - + @@ -428,7 +428,7 @@ - + @@ -436,7 +436,7 @@ - + @@ -444,7 +444,7 @@ - + @@ -452,7 +452,7 @@ - + @@ -460,7 +460,7 @@ - + @@ -468,7 +468,7 @@ - + @@ -476,7 +476,7 @@ - + \ No newline at end of file diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/BaselineTest/TestBaseline/PartitionKeyHashBaselineTest.Numbers.xml b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/BaselineTest/TestBaseline/PartitionKeyHashBaselineTest.Numbers.xml index dfbf1d6ca6..de536f8710 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/BaselineTest/TestBaseline/PartitionKeyHashBaselineTest.Numbers.xml +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/BaselineTest/TestBaseline/PartitionKeyHashBaselineTest.Numbers.xml @@ -5,8 +5,8 @@ 0 - A6-80-DA-1E-00-00-00-00-00-00-00-00-00-00-00-00 - 58-DE-DD-B0-9B-DC-1C-EC-E9-B1-C4-DA-BE-95-5B-55 + 00-00-00-00-00-00-00-00-00-00-00-00-1E-DA-80-A6 + 55-5B-95-BE-DA-C4-B1-E9-EC-1C-DC-9B-B0-DD-DE-58 @@ -15,8 +15,8 @@ 0 - 50-A6-3D-B0-00-00-00-00-00-00-00-00-00-00-00-00 - 6E-40-10-FD-81-6E-76-35-72-B2-3D-DF-B1-8B-01-3D + 00-00-00-00-00-00-00-00-00-00-00-00-B0-3D-A6-50 + 3D-01-8B-B1-DF-3D-B2-72-35-76-6E-81-FD-10-40-6E @@ -25,8 +25,8 @@ 1 - 93-5A-65-E1-00-00-00-00-00-00-00-00-00-00-00-00 - B0-70-70-B8-53-69-CF-D0-A5-78-BA-39-B3-98-CD-A0 + 00-00-00-00-00-00-00-00-00-00-00-00-E1-65-5A-93 + A0-CD-98-B3-39-BA-78-A5-D0-CF-69-53-B8-70-70-B0 @@ -35,8 +35,8 @@ -1 - E7-DB-2B-45-00-00-00-00-00-00-00-00-00-00-00-00 - 39-68-C1-BB-42-E8-3A-9E-5B-1C-6C-93-7A-8E-93-59 + 00-00-00-00-00-00-00-00-00-00-00-00-45-2B-DB-E7 + 59-93-8E-7A-93-6C-1C-5B-9E-3A-E8-42-BB-C1-68-39 @@ -45,8 +45,8 @@ 4.94065645841247E-324 - 18-D3-83-AC-00-00-00-00-00-00-00-00-00-00-00-00 - 39-01-80-65-F8-DE-85-E4-7D-92-80-A2-63-BA-6C-0E + 00-00-00-00-00-00-00-00-00-00-00-00-AC-83-D3-18 + 0E-6C-BA-63-A2-80-92-7D-E4-85-DE-F8-65-80-01-39 @@ -55,8 +55,8 @@ 1.7976931348623157E+308 - 46-3C-8B-37-00-00-00-00-00-00-00-00-00-00-00-00 - 4D-CC-DB-45-F2-1F-59-34-26-10-57-64-99-4D-42-71 + 00-00-00-00-00-00-00-00-00-00-00-00-37-8B-3C-46 + 71-42-4D-99-64-57-10-26-34-59-1F-F2-45-DB-CC-4D @@ -65,8 +65,8 @@ -1.7976931348623157E+308 - A1-D8-8D-E3-00-00-00-00-00-00-00-00-00-00-00-00 - 25-DF-01-38-3D-28-E1-AC-9D-FD-0C-91-C4-0C-05-27 + 00-00-00-00-00-00-00-00-00-00-00-00-E3-8D-D8-A1 + 27-05-0C-C4-91-0C-FD-9D-AC-E1-28-3D-38-01-DF-25 @@ -75,8 +75,8 @@ "NaN" - 37-DA-04-AF-00-00-00-00-00-00-00-00-00-00-00-00 - B1-FC-ED-A8-8C-2D-92-A0-CA-0E-D7-83-25-22-A4-5A + 00-00-00-00-00-00-00-00-00-00-00-00-AF-04-DA-37 + 5A-A4-22-25-83-D7-0E-CA-A0-92-2D-8C-A8-ED-FC-B1 @@ -85,8 +85,8 @@ "-Infinity" - C5-6C-5A-0F-00-00-00-00-00-00-00-00-00-00-00-00 - E8-CD-93-E8-9F-C7-3B-CD-B6-B3-6D-26-FB-B3-63-5B + 00-00-00-00-00-00-00-00-00-00-00-00-0F-5A-6C-C5 + 5B-63-B3-FB-26-6D-B3-B6-CD-3B-C7-9F-E8-93-CD-E8 @@ -95,8 +95,8 @@ "Infinity" - 93-4F-F8-D2-00-00-00-00-00-00-00-00-00-00-00-00 - FF-B4-44-19-BB-9B-82-C9-A5-E2-8C-8A-F0-F7-89-38 + 00-00-00-00-00-00-00-00-00-00-00-00-D2-F8-4F-93 + 38-89-F7-F0-8A-8C-E2-A5-C9-82-9B-BB-19-44-B4-FF @@ -105,8 +105,8 @@ 9223372036854775807 - 6D-A9-32-17-00-00-00-00-00-00-00-00-00-00-00-00 - 9B-62-D1-84-93-F8-83-89-A1-CC-DF-78-91-95-DB-6E + 00-00-00-00-00-00-00-00-00-00-00-00-17-32-A9-6D + 6E-DB-95-91-78-DF-CC-A1-89-83-F8-93-84-D1-62-9B @@ -115,8 +115,8 @@ 9223372036854775806 - 6D-A9-32-17-00-00-00-00-00-00-00-00-00-00-00-00 - 9B-62-D1-84-93-F8-83-89-A1-CC-DF-78-91-95-DB-6E + 00-00-00-00-00-00-00-00-00-00-00-00-17-32-A9-6D + 6E-DB-95-91-78-DF-CC-A1-89-83-F8-93-84-D1-62-9B @@ -125,8 +125,8 @@ -9223372036854775808 - 87-27-AF-69-00-00-00-00-00-00-00-00-00-00-00-00 - BB-D2-8A-32-15-AD-AD-AF-FE-BD-12-55-39-C6-D5-63 + 00-00-00-00-00-00-00-00-00-00-00-00-69-AF-27-87 + 63-D5-C6-39-55-12-BD-FE-AF-AD-AD-15-32-8A-D2-BB @@ -135,8 +135,8 @@ 9007199254740991 - 8A-9B-F6-13-00-00-00-00-00-00-00-00-00-00-00-00 - CF-E6-28-14-CE-93-66-9E-1F-ED-2C-03-C7-21-C2-12 + 00-00-00-00-00-00-00-00-00-00-00-00-13-F6-9B-8A + 12-C2-21-C7-03-2C-ED-1F-9E-66-93-CE-14-28-E6-CF @@ -145,8 +145,8 @@ 9007199254740990 - 66-43-74-E5-00-00-00-00-00-00-00-00-00-00-00-00 - 9B-FD-C1-48-B2-F0-8E-62-E4-A6-4A-D9-67-34-EA-8F + 00-00-00-00-00-00-00-00-00-00-00-00-E5-74-43-66 + 8F-EA-34-67-D9-4A-A6-E4-62-8E-F0-B2-48-C1-FD-9B @@ -155,8 +155,8 @@ 9007199254740992 - 12-CE-0C-A3-00-00-00-00-00-00-00-00-00-00-00-00 - 5D-06-9B-56-DE-A5-9D-D9-E3-22-11-D0-10-59-33-55 + 00-00-00-00-00-00-00-00-00-00-00-00-A3-0C-CE-12 + 55-33-59-10-D0-11-22-E3-D9-9D-A5-DE-56-9B-06-5D \ No newline at end of file diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/BaselineTest/TestBaseline/PartitionKeyHashBaselineTest.Singletons.xml b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/BaselineTest/TestBaseline/PartitionKeyHashBaselineTest.Singletons.xml index 72c825e47b..735b049deb 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/BaselineTest/TestBaseline/PartitionKeyHashBaselineTest.Singletons.xml +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/BaselineTest/TestBaseline/PartitionKeyHashBaselineTest.Singletons.xml @@ -5,8 +5,8 @@ UNDEFINED - B7-28-4E-51-00-00-00-00-00-00-00-00-00-00-00-00 - B5-5C-FF-6E-E5-AB-10-46-83-35-F8-78-AA-2D-62-51 + 00-00-00-00-00-00-00-00-00-00-00-00-51-4E-28-B7 + 51-62-2D-AA-78-F8-35-83-46-10-AB-E5-6E-FF-5C-B5 @@ -15,8 +15,8 @@ null - AB-D1-5A-E4-00-00-00-00-00-00-00-00-00-00-00-00 - 16-FE-74-83-90-5C-CE-7A-85-67-0E-43-E4-67-88-77 + 00-00-00-00-00-00-00-00-00-00-00-00-E4-5A-D1-AB + 77-88-67-E4-43-0E-67-85-7A-CE-5C-90-83-74-FE-16 @@ -25,8 +25,8 @@ true - 76-80-2A-5E-00-00-00-00-00-00-00-00-00-00-00-00 - 59-3E-6A-30-DD-C6-6A-72-E4-A8-B5-C5-27-11-71-4E + 00-00-00-00-00-00-00-00-00-00-00-00-5E-2A-80-76 + 4E-71-11-27-C5-B5-A8-E4-72-6A-C6-DD-30-6A-3E-59 @@ -35,8 +35,8 @@ false - 9F-D0-27-6C-00-00-00-00-00-00-00-00-00-00-00-00 - F2-1E-36-37-9E-0E-5E-63-39-34-0A-E9-91-BE-E1-EF + 00-00-00-00-00-00-00-00-00-00-00-00-6C-27-D0-9F + EF-E1-BE-91-E9-0A-34-39-63-5E-0E-9E-37-36-1E-F2 \ No newline at end of file diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/BaselineTest/TestBaseline/PartitionKeyHashBaselineTest.Strings.xml b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/BaselineTest/TestBaseline/PartitionKeyHashBaselineTest.Strings.xml index f7130be38b..4cb88af72d 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/BaselineTest/TestBaseline/PartitionKeyHashBaselineTest.Strings.xml +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/BaselineTest/TestBaseline/PartitionKeyHashBaselineTest.Strings.xml @@ -5,8 +5,8 @@ "" - C7-BC-36-20-00-00-00-00-00-00-00-00-00-00-00-00 - 7E-E0-F6-1A-B7-A1-A3-DF-DB-2B-E5-33-CF-D2-63-6A + 00-00-00-00-00-00-00-00-00-00-00-00-20-36-BC-C7 + 6A-63-D2-CF-33-E5-2B-DB-DF-A3-A1-B7-1A-F6-E0-7E @@ -15,8 +15,8 @@ "asdf" - 90-E7-36-FA-00-00-00-00-00-00-00-00-00-00-00-00 - 08-39-D8-49-B1-6E-EE-05-8D-06-4B-39-FE-89-8D-49 + 00-00-00-00-00-00-00-00-00-00-00-00-FA-36-E7-90 + 49-8D-89-FE-39-4B-06-8D-05-EE-6E-B1-49-D8-39-08 @@ -25,8 +25,8 @@ "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa" - FE-6B-6F-A3-00-00-00-00-00-00-00-00-00-00-00-00 - 67-7A-49-B8-DA-D5-97-08-0B-2A-27-5D-26-59-A5-83 + 00-00-00-00-00-00-00-00-00-00-00-00-A3-6F-6B-FE + 83-A5-59-26-5D-27-2A-0B-08-97-D5-DA-B8-49-7A-67 @@ -35,8 +35,8 @@ "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa" - 25-E3-64-0D-00-00-00-00-00-00-00-00-00-00-00-00 - 4E-5B-50-72-C2-98-ED-26-CE-36-29-BD-FB-8C-28-A9 + 00-00-00-00-00-00-00-00-00-00-00-00-0D-64-E3-25 + A9-28-8C-FB-BD-29-36-CE-26-ED-98-C2-72-50-5B-4E @@ -45,8 +45,8 @@ "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa" - 25-E3-64-0D-00-00-00-00-00-00-00-00-00-00-00-00 - AE-8E-39-1D-65-72-E0-D8-81-E7-0B-F5-E3-A8-46-43 + 00-00-00-00-00-00-00-00-00-00-00-00-0D-64-E3-25 + 43-46-A8-E3-F5-0B-E7-81-D8-E0-72-65-1D-39-8E-AE @@ -55,8 +55,8 @@ "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa" - 25-E3-64-0D-00-00-00-00-00-00-00-00-00-00-00-00 - D6-6C-FB-1C-FA-0A-08-42-61-0F-22-2B-5F-B7-CA-CE + 00-00-00-00-00-00-00-00-00-00-00-00-0D-64-E3-25 + CE-CA-B7-5F-2B-22-0F-61-42-08-0A-FA-1C-FB-6C-D6 \ No newline at end of file From f6799712fa127639c179cb0bb8cc5659d081f407 Mon Sep 17 00:00:00 2001 From: Brandon Chong Date: Sun, 4 Oct 2020 15:34:25 -0700 Subject: [PATCH 83/85] setting feed iterator flag for non retriable exception --- .../src/Query/v3Query/QueryIterator.cs | 13 ++++++++++++- 1 file changed, 12 insertions(+), 1 deletion(-) diff --git a/Microsoft.Azure.Cosmos/src/Query/v3Query/QueryIterator.cs b/Microsoft.Azure.Cosmos/src/Query/v3Query/QueryIterator.cs index a67c8a8549..3abb82c41d 100644 --- a/Microsoft.Azure.Cosmos/src/Query/v3Query/QueryIterator.cs +++ b/Microsoft.Azure.Cosmos/src/Query/v3Query/QueryIterator.cs @@ -215,6 +215,10 @@ public override async Task ReadNextAsync(CancellationToken canc } CosmosException cosmosException = ExceptionToCosmosException.CreateFromException(tryGetQueryPage.Exception); + if (!IsRetriableException(cosmosException)) + { + this.hasMoreResults = false; + } return QueryResponse.CreateFailure( statusCode: cosmosException.StatusCode, @@ -232,10 +236,17 @@ public override async Task ReadNextAsync(CancellationToken canc public override CosmosElement GetCosmosElementContinuationToken() => this.queryPipelineStage.Current.Result.State?.Value; + private static bool IsRetriableException(CosmosException cosmosException) + { + return ((int)cosmosException.StatusCode == 429) + || (cosmosException.StatusCode == System.Net.HttpStatusCode.RequestTimeout) + || (cosmosException.StatusCode == System.Net.HttpStatusCode.ServiceUnavailable); + } + protected override void Dispose(bool disposing) { this.queryPipelineStage.DisposeAsync(); base.Dispose(disposing); } } -} +} \ No newline at end of file From 1596a713518bb0f306abde4476830a623c8795d6 Mon Sep 17 00:00:00 2001 From: Brandon Chong Date: Sun, 4 Oct 2020 15:54:45 -0700 Subject: [PATCH 84/85] apparently we don't actually handle name cache is stale exceptions --- .../NetworkAttachedDocumentContainer.cs | 19 ++++++++++--------- 1 file changed, 10 insertions(+), 9 deletions(-) diff --git a/Microsoft.Azure.Cosmos/src/Pagination/NetworkAttachedDocumentContainer.cs b/Microsoft.Azure.Cosmos/src/Pagination/NetworkAttachedDocumentContainer.cs index 382e641e14..2bc941d889 100644 --- a/Microsoft.Azure.Cosmos/src/Pagination/NetworkAttachedDocumentContainer.cs +++ b/Microsoft.Azure.Cosmos/src/Pagination/NetworkAttachedDocumentContainer.cs @@ -11,6 +11,7 @@ namespace Microsoft.Azure.Cosmos.Pagination using System.Net; using System.Threading; using System.Threading.Tasks; + using Microsoft.Azure.Cosmos.Common; using Microsoft.Azure.Cosmos.CosmosElements; using Microsoft.Azure.Cosmos.Query.Core; using Microsoft.Azure.Cosmos.Query.Core.Monads; @@ -88,18 +89,18 @@ public async Task>> MonadicGetChildRangeAsync( PartitionKeyRange partitionKeyRange, CancellationToken cancellationToken) { - IReadOnlyList overlappingRanges = await this.cosmosQueryContext.QueryClient.TryGetOverlappingRangesAsync( + try + { + List overlappingRanges = await this.cosmosQueryContext.QueryClient.GetTargetPartitionKeyRangesAsync( + this.cosmosQueryContext.ResourceLink, this.cosmosQueryContext.ContainerResourceId, - partitionKeyRange.ToRange(), - forceRefresh: true); - if (overlappingRanges == null) + new List>() { partitionKeyRange.ToRange() }); + return TryCatch>.FromResult(overlappingRanges); + } + catch (Exception ex) { - return TryCatch>.FromException( - CosmosExceptionFactory.CreateNotFoundException( - message: $"Container with resource id: {this.cosmosQueryContext.ContainerResourceId} was not found")); + return TryCatch>.FromException(ex); } - - return TryCatch>.FromResult(overlappingRanges.ToList()); } public Task> MonadicReadFeedAsync( From f069c81961c97b4c439efd4b30ba414b11286235 Mon Sep 17 00:00:00 2001 From: Brandon Chong Date: Mon, 5 Oct 2020 15:47:13 -0700 Subject: [PATCH 85/85] dummy commit --- .../Parallel/ParallelCrossPartitionQueryPipelineStage.cs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 e3b507bcd7..c220c10c95 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 @@ -73,7 +73,7 @@ public async ValueTask MoveNextAsync() } else { - // left most and any non null continuations + // Left most and any non null continuations List<(PartitionKeyRange, QueryState)> rangesAndStates = crossPartitionState.Value.OrderBy(tuple => tuple.Item1, PartitionKeyRangeComparer.Singleton).ToList(); List activeParallelContinuationTokens = new List(); for (int i = 0; i < rangesAndStates.Count; i++)