From a001801b58723c6f05cda0ac13208e444748613b Mon Sep 17 00:00:00 2001 From: Brandon Chong Date: Wed, 8 Jan 2020 14:03:51 -0800 Subject: [PATCH 01/28] realized that this will be a lot of work --- .../CompositeContinuationTokenRefStruct.cs | 40 +++++++++ .../OrderByContinuationToken.cs | 4 +- .../OrderByContinuationTokenRefStruct.cs | 83 +++++++++++++++++++ .../Core/ContinuationTokens/RangeRefStruct.cs | 49 +++++++++++ .../DocumentQueryExecutionComponentBase.cs | 3 + ...DocumentQueryExecutionComponent.Compute.cs | 33 +++++++- .../IDocumentQueryExecutionComponent.cs | 3 + ...smosCrossPartitionQueryExecutionContext.cs | 3 + .../CosmosOrderByItemQueryExecutionContext.cs | 39 ++++++++- .../OrderBy/OrderByConsumeComparer.cs | 2 +- .../OrderBy/OrderByQueryResult.cs | 2 +- ...CosmosParallelItemQueryExecutionContext.cs | 31 +++++++ .../CosmosQueryClientCoreTest.cs | 2 +- 13 files changed, 287 insertions(+), 7 deletions(-) create mode 100644 Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/CompositeContinuationTokenRefStruct.cs create mode 100644 Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/OrderByContinuationTokenRefStruct.cs create mode 100644 Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/RangeRefStruct.cs diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/CompositeContinuationTokenRefStruct.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/CompositeContinuationTokenRefStruct.cs new file mode 100644 index 0000000000..f61f3f2a0a --- /dev/null +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/CompositeContinuationTokenRefStruct.cs @@ -0,0 +1,40 @@ +// ------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +// ------------------------------------------------------------ + +namespace Microsoft.Azure.Cosmos.Query.Core.ContinuationTokens +{ + using System; + using Microsoft.Azure.Cosmos.Json; + + internal ref struct CompositeContinuationTokenRefStruct + { + private const string TokenProperytName = "token"; + private const string RangePropertyName = "range"; + + public CompositeContinuationTokenRefStruct(string backendContinuationToken, RangeRefStruct range) + { + this.BackendContinuationToken = backendContinuationToken; + this.Range = range; + } + + public string BackendContinuationToken { get; } + + public RangeRefStruct Range { get; } + + public void WriteTo(IJsonWriter jsonWriter) + { + if (jsonWriter == null) + { + throw new ArgumentNullException(nameof(jsonWriter)); + } + + jsonWriter.WriteObjectStart(); + jsonWriter.WriteFieldName(CompositeContinuationTokenRefStruct.TokenProperytName); + jsonWriter.WriteStringValue(this.BackendContinuationToken); + jsonWriter.WriteFieldName(CompositeContinuationTokenRefStruct.RangePropertyName); + this.Range.WriteTo(jsonWriter); + jsonWriter.WriteObjectEnd(); + } + } +} diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/OrderByContinuationToken.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/OrderByContinuationToken.cs index d15f0f0509..93c8e4ef6c 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/OrderByContinuationToken.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/OrderByContinuationToken.cs @@ -62,7 +62,7 @@ internal sealed class OrderByContinuationToken /// The filter (refer to property documentation). public OrderByContinuationToken( CompositeContinuationToken compositeContinuationToken, - IList orderByItems, + IReadOnlyList orderByItems, string rid, int skipCount, string filter) @@ -118,7 +118,7 @@ public CompositeContinuationToken CompositeContinuationToken /// Right now, we don't support orderBy by multiple fields, so orderByItems is an array of one element. /// > [JsonProperty("orderByItems")] - public IList OrderByItems + public IReadOnlyList OrderByItems { get; } diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/OrderByContinuationTokenRefStruct.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/OrderByContinuationTokenRefStruct.cs new file mode 100644 index 0000000000..448f84a141 --- /dev/null +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/OrderByContinuationTokenRefStruct.cs @@ -0,0 +1,83 @@ +// ------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +// ------------------------------------------------------------ + +namespace Microsoft.Azure.Cosmos.Query.Core.ContinuationTokens +{ + using System; + using System.Collections.Generic; + using Microsoft.Azure.Cosmos.Json; + using Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.OrderBy; + + internal ref struct OrderByContinuationTokenRefStruct + { + private const string CompositeTokenPropertyName = "compositeToken"; + private const string OrderByItemsPropertyName = "orderByItems"; + private const string RidPropertyName = "rid"; + private const string SkipCountPropertyName = "skipCount"; + private const string FilterPropertyName = "filter"; + private const string ItemPropertyName = "item"; + + public OrderByContinuationTokenRefStruct( + CompositeContinuationTokenRefStruct compositeContinuationTokenRefStruct, + IReadOnlyList orderByItems, + string rid, + int skipCount, + string filter) + { + this.CompositeContinuationToken = compositeContinuationTokenRefStruct; + this.OrderByItems = orderByItems; + this.Rid = rid; + this.SkipCount = skipCount; + this.Filter = filter; + } + + public CompositeContinuationTokenRefStruct CompositeContinuationToken { get; } + public IReadOnlyList OrderByItems { get; } + public string Rid { get; } + public int SkipCount { get; } + public string Filter { get; } + + public void WriteTo(IJsonWriter jsonWriter) + { + if (jsonWriter == null) + { + throw new ArgumentNullException(nameof(jsonWriter)); + } + + jsonWriter.WriteObjectStart(); + jsonWriter.WriteFieldName(OrderByContinuationTokenRefStruct.CompositeTokenPropertyName); + this.CompositeContinuationToken.WriteTo(jsonWriter); + jsonWriter.WriteFieldName(OrderByContinuationTokenRefStruct.OrderByItemsPropertyName); + jsonWriter.WriteArrayStart(); + + foreach (OrderByItem orderByItem in this.OrderByItems) + { + jsonWriter.WriteObjectStart(); + + if (orderByItem.Item != null) + { + jsonWriter.WriteFieldName(OrderByContinuationTokenRefStruct.ItemPropertyName); + orderByItem.Item.WriteTo(jsonWriter); + } + + jsonWriter.WriteObjectEnd(); + } + + jsonWriter.WriteArrayEnd(); + jsonWriter.WriteFieldName(OrderByContinuationTokenRefStruct.RidPropertyName); + jsonWriter.WriteStringValue(this.Rid); + jsonWriter.WriteFieldName(OrderByContinuationTokenRefStruct.SkipCountPropertyName); + jsonWriter.WriteInt32Value(this.SkipCount); + jsonWriter.WriteFieldName(OrderByContinuationTokenRefStruct.FilterPropertyName); + if (this.Filter != null) + { + jsonWriter.WriteStringValue(this.Filter); + } + else + { + jsonWriter.WriteNullValue(); + } + } + } +} diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/RangeRefStruct.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/RangeRefStruct.cs new file mode 100644 index 0000000000..e4801e343e --- /dev/null +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/RangeRefStruct.cs @@ -0,0 +1,49 @@ +// ------------------------------------------------------------ +// 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/ExecutionComponent/DocumentQueryExecutionComponentBase.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/DocumentQueryExecutionComponentBase.cs index 475320bc0d..64360225b4 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/DocumentQueryExecutionComponentBase.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/DocumentQueryExecutionComponentBase.cs @@ -6,6 +6,7 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent using System; using System.Threading; using System.Threading.Tasks; + using Microsoft.Azure.Cosmos.Json; using Microsoft.Azure.Cosmos.Query.Core.QueryClient; /// @@ -67,5 +68,7 @@ public void Stop() } public abstract bool TryGetContinuationToken(out string state); + + public abstract void SerializeState(IJsonWriter jsonWriter); } } 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 cb15f58cdf..982487257c 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 @@ -8,6 +8,7 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.GroupBy using System.Threading; using System.Threading.Tasks; using Microsoft.Azure.Cosmos.CosmosElements; + using Microsoft.Azure.Cosmos.Json; using Microsoft.Azure.Cosmos.Query.Core.Exceptions; using Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.Aggregate; using Microsoft.Azure.Cosmos.Query.Core.Metrics; @@ -137,6 +138,26 @@ public override async Task DrainAsync( return response; } + public override void SerializeState(IJsonWriter jsonWriter) + { + if (jsonWriter == null) + { + throw new ArgumentNullException(nameof(jsonWriter)); + } + + if (this.IsDone) + { + jsonWriter.WriteNullValue(); + } + else + { + if (this.Source.IsDone) + { + + } + } + } + public override bool TryGetContinuationToken(out string continuationToken) { if (this.IsDone) @@ -168,7 +189,7 @@ public override bool TryGetContinuationToken(out string continuationToken) return true; } - private sealed class GroupByContinuationToken + private readonly struct GroupByContinuationToken { public GroupByContinuationToken( string groupingTableContinuationToken, @@ -248,6 +269,16 @@ public IReadOnlyDictionary GetQueryMetrics() throw new NotImplementedException(); } + public void SerializeState(IJsonWriter jsonWriter) + { + if (jsonWriter == null) + { + throw new ArgumentNullException(nameof(jsonWriter)); + } + + jsonWriter.WriteNullValue(); + } + public void Stop() { } diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/IDocumentQueryExecutionComponent.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/IDocumentQueryExecutionComponent.cs index 8203682f45..21dfefb31c 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/IDocumentQueryExecutionComponent.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/IDocumentQueryExecutionComponent.cs @@ -6,6 +6,7 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent using System; using System.Threading; using System.Threading.Tasks; + using Microsoft.Azure.Cosmos.Json; using Microsoft.Azure.Cosmos.Query.Core.QueryClient; /// @@ -32,5 +33,7 @@ internal interface IDocumentQueryExecutionComponent : IDisposable void Stop(); bool TryGetContinuationToken(out string state); + + void SerializeState(IJsonWriter jsonWriter); } } diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosCrossPartitionQueryExecutionContext.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosCrossPartitionQueryExecutionContext.cs index 8a2809e9bf..b0f442e1e5 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosCrossPartitionQueryExecutionContext.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosCrossPartitionQueryExecutionContext.cs @@ -11,6 +11,7 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionContext using System.Threading.Tasks; using Core.ExecutionComponent; using Microsoft.Azure.Cosmos.CosmosElements; + using Microsoft.Azure.Cosmos.Json; using Microsoft.Azure.Cosmos.Query.Core; using Microsoft.Azure.Cosmos.Query.Core.Collections; using Microsoft.Azure.Cosmos.Query.Core.ComparableTask; @@ -644,6 +645,8 @@ public bool TryGetContinuationToken(out string state) return true; } + public abstract void SerializeState(IJsonWriter jsonWriter); + public readonly struct InitInfo { public InitInfo(int targetIndex, IReadOnlyDictionary continuationTokens) 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 d135be3367..b19f8dbf9b 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/OrderBy/CosmosOrderByItemQueryExecutionContext.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/OrderBy/CosmosOrderByItemQueryExecutionContext.cs @@ -11,6 +11,7 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.OrderBy using System.Threading; using System.Threading.Tasks; using Microsoft.Azure.Cosmos.CosmosElements; + using Microsoft.Azure.Cosmos.Json; using Microsoft.Azure.Cosmos.Query.Core; using Microsoft.Azure.Cosmos.Query.Core.Collections; using Microsoft.Azure.Cosmos.Query.Core.ContinuationTokens; @@ -67,7 +68,7 @@ internal sealed class CosmosOrderByItemQueryExecutionContext : CosmosCrossPartit /// private string previousRid; - private IList previousOrderByItems; + private IReadOnlyList previousOrderByItems; /// /// Initializes a new instance of the CosmosOrderByItemQueryExecutionContext class. @@ -159,6 +160,42 @@ protected override string ContinuationToken } } + public override void SerializeState(IJsonWriter jsonWriter) + { + if (jsonWriter == null) + { + throw new ArgumentNullException(nameof(jsonWriter)); + } + + IEnumerable activeItemProducers = this.GetActiveItemProducers(); + if (activeItemProducers.Any()) + { + jsonWriter.WriteArrayStart(); + + foreach (ItemProducer activeItemProducer in activeItemProducers) + { + OrderByQueryResult orderByQueryResult = new OrderByQueryResult(activeItemProducer.Current); + OrderByContinuationTokenRefStruct orderByContinuationToken = new OrderByContinuationTokenRefStruct( + compositeContinuationTokenRefStruct: new CompositeContinuationTokenRefStruct( + backendContinuationToken: activeItemProducer.PreviousContinuationToken, + range: new RangeRefStruct( + min: activeItemProducer.PartitionKeyRange.MinInclusive, + max: activeItemProducer.PartitionKeyRange.MaxExclusive)), + orderByItems: orderByQueryResult.OrderByItems, + rid: orderByQueryResult.Rid, + skipCount: this.ShouldIncrementSkipCount(activeItemProducer) ? this.skipCount + 1 : 0, + filter: activeItemProducer.Filter); + orderByContinuationToken.WriteTo(jsonWriter); + } + + jsonWriter.WriteArrayEnd(); + } + else + { + jsonWriter.WriteNullValue(); + } + } + public static async Task> TryCreateAsync( CosmosQueryContext queryContext, CosmosCrossPartitionQueryExecutionContext.CrossPartitionInitParams initParams, diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/OrderBy/OrderByConsumeComparer.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/OrderBy/OrderByConsumeComparer.cs index e9cfb663bb..8bf9a296c5 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/OrderBy/OrderByConsumeComparer.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/OrderBy/OrderByConsumeComparer.cs @@ -101,7 +101,7 @@ public int Compare(ItemProducerTree producer1, ItemProducerTree producer2) /// 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(IList items1, IList items2) + public int CompareOrderByItems(IReadOnlyList items1, IReadOnlyList items2) { if (object.ReferenceEquals(items1, items2)) { diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/OrderBy/OrderByQueryResult.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/OrderBy/OrderByQueryResult.cs index c39ff13f3d..0e1b289563 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/OrderBy/OrderByQueryResult.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/OrderBy/OrderByQueryResult.cs @@ -66,7 +66,7 @@ public string Rid /// /// Gets the order by items from the document. /// - public IList OrderByItems + public IReadOnlyList OrderByItems { get { 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 cf3b7787ae..f4531ce260 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/Parallel/CosmosParallelItemQueryExecutionContext.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/Parallel/CosmosParallelItemQueryExecutionContext.cs @@ -10,6 +10,7 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.Parallel using System.Threading; using System.Threading.Tasks; using Microsoft.Azure.Cosmos.CosmosElements; + using Microsoft.Azure.Cosmos.Json; using Microsoft.Azure.Cosmos.Query.Core; using Microsoft.Azure.Cosmos.Query.Core.Collections; using Microsoft.Azure.Cosmos.Query.Core.ContinuationTokens; @@ -103,6 +104,36 @@ protected override string ContinuationToken } } + public override void SerializeState(IJsonWriter jsonWriter) + { + if (jsonWriter == null) + { + throw new ArgumentNullException(nameof(jsonWriter)); + } + + IEnumerable activeItemProducers = this.GetActiveItemProducers(); + if (activeItemProducers.Any()) + { + jsonWriter.WriteArrayStart(); + + foreach (ItemProducer activeItemProducer in activeItemProducers) + { + CompositeContinuationTokenRefStruct compositeToken = new CompositeContinuationTokenRefStruct( + backendContinuationToken: activeItemProducer.BackendContinuationToken, + range: new RangeRefStruct( + min: activeItemProducer.PartitionKeyRange.MinInclusive, + max: activeItemProducer.PartitionKeyRange.MaxExclusive)); + compositeToken.WriteTo(jsonWriter); + } + + jsonWriter.WriteArrayEnd(); + } + else + { + jsonWriter.WriteNullValue(); + } + } + public static async Task> TryCreateAsync( CosmosQueryContext queryContext, CosmosCrossPartitionQueryExecutionContext.CrossPartitionInitParams initParams, diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CosmosQueryClientCoreTest.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CosmosQueryClientCoreTest.cs index 5dc92d4007..fc80280b5a 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CosmosQueryClientCoreTest.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CosmosQueryClientCoreTest.cs @@ -46,7 +46,7 @@ public async Task TryGetOverlappingRangesAsyncTest() ContainerQueryProperties containerProperties = await this.queryClientCore.GetCachedContainerQueryPropertiesAsync( containerLink: this.Container.LinkUri, partitionKey: new PartitionKey("Test"), - cancellationToken: default); + cancellationToken: default(CancellationToken)); Assert.IsNotNull(containerProperties); Assert.IsNotNull(containerProperties.ResourceId); From f8359f2dfcdf9329fd22e710bd6e18eca54061a4 Mon Sep 17 00:00:00 2001 From: Brandon Chong Date: Tue, 14 Jan 2020 18:01:06 -0800 Subject: [PATCH 02/28] wired up continuation token on the write path --- ...DocumentQueryExecutionComponent.Compute.cs | 26 +++++++++ .../Aggregators/AverageAggregator.cs | 28 +++++++++- .../Aggregate/Aggregators/CountAggregator.cs | 11 ++++ .../Aggregate/Aggregators/IAggregator.cs | 3 + .../Aggregate/Aggregators/MinMaxAggregator.cs | 27 ++++++--- .../Aggregators/SingleGroupAggregator.cs | 56 ++++++++++++++++--- .../Aggregate/Aggregators/SumAggregator.cs | 11 ++++ ...tDocumentQueryExecutionComponent.Client.cs | 29 +--------- ...DocumentQueryExecutionComponent.Compute.cs | 20 ------- ...DistinctDocumentQueryExecutionComponent.cs | 46 +++++++++++++++ .../DistinctMap.OrderedDistinctMap.cs | 11 ++++ .../DistinctMap.UnorderedDistinctMap.cs | 26 ++++++--- .../Distinct/DistinctMap.cs | 3 + .../DocumentQueryExecutionComponentBase.cs | 1 + ...yDocumentQueryExecutionComponent.Client.cs | 6 ++ ...DocumentQueryExecutionComponent.Compute.cs | 21 +++++-- .../GroupByDocumentQueryExecutionComponent.cs | 15 ++++- .../IDocumentQueryExecutionComponent.cs | 1 + .../SkipDocumentQueryExecutionComponent.cs | 26 +++++++++ .../TakeDocumentQueryExecutionComponent.cs | 26 +++++++++ 20 files changed, 315 insertions(+), 78 deletions(-) diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/AggregateDocumentQueryExecutionComponent.Compute.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/AggregateDocumentQueryExecutionComponent.Compute.cs index a53c141c78..267ae2929e 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/AggregateDocumentQueryExecutionComponent.Compute.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/AggregateDocumentQueryExecutionComponent.Compute.cs @@ -8,6 +8,7 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.Aggregate using System.Threading; using System.Threading.Tasks; using Microsoft.Azure.Cosmos.CosmosElements; + using Microsoft.Azure.Cosmos.Json; using Microsoft.Azure.Cosmos.Query.Core.Exceptions; using Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.Aggregate.Aggregators; using Microsoft.Azure.Cosmos.Query.Core.Monads; @@ -15,6 +16,9 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.Aggregate internal abstract partial class AggregateDocumentQueryExecutionComponent : DocumentQueryExecutionComponentBase { + private const string SourceTokenName = "SourceToken"; + private const string AggregationTokenName = "AggregationToken"; + private static readonly IReadOnlyList EmptyResults = new List().AsReadOnly(); private sealed class ComputeAggregateDocumentQueryExecutionComponent : AggregateDocumentQueryExecutionComponent @@ -186,6 +190,28 @@ public override bool TryGetContinuationToken(out string state) return true; } + public override void SerializeState(IJsonWriter jsonWriter) + { + if (jsonWriter == null) + { + throw new ArgumentNullException(nameof(jsonWriter)); + } + + if (this.IsDone) + { + jsonWriter.WriteNullValue(); + } + else + { + jsonWriter.WriteObjectStart(); + jsonWriter.WriteFieldName(ComputeAggregateDocumentQueryExecutionComponent.SourceTokenName); + this.Source.SerializeState(jsonWriter); + jsonWriter.WriteFieldName(ComputeAggregateDocumentQueryExecutionComponent.AggregationTokenName); + this.singleGroupAggregator.SerializeState(jsonWriter); + jsonWriter.WriteObjectEnd(); + } + } + private struct AggregateContinuationToken { private const string SingleGroupAggregatorContinuationTokenName = "SingleGroupAggregatorContinuationToken"; diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/AverageAggregator.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/AverageAggregator.cs index bb69dc2223..8c73eef441 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/AverageAggregator.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/AverageAggregator.cs @@ -5,6 +5,7 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.Aggregate.Aggrega { using System; using Microsoft.Azure.Cosmos.CosmosElements; + using Microsoft.Azure.Cosmos.Json; using Microsoft.Azure.Cosmos.Query.Core.Exceptions; using Microsoft.Azure.Cosmos.Query.Core.Monads; @@ -52,6 +53,16 @@ public string GetContinuationToken() return this.globalAverage.ToString(); } + public void SerializeState(IJsonWriter jsonWriter) + { + if (jsonWriter == null) + { + throw new ArgumentNullException(nameof(jsonWriter)); + } + + this.globalAverage.SerializeState(jsonWriter); + } + public static TryCatch TryCreate(string continuationToken) { AverageInfo averageInfo; @@ -74,7 +85,7 @@ public static TryCatch TryCreate(string continuationToken) /// /// Struct that stores a weighted average as a sum and count so they that average across different partitions with different numbers of documents can be taken. /// - private struct AverageInfo + private readonly struct AverageInfo { private const string SumName = "sum"; private const string CountName = "count"; @@ -198,6 +209,21 @@ public CosmosNumber GetAverage() return CosmosNumber64.Create(this.Sum.Value / this.Count); } + public void SerializeState(IJsonWriter jsonWriter) + { + if (jsonWriter == null) + { + throw new ArgumentNullException(nameof(jsonWriter)); + } + + jsonWriter.WriteObjectStart(); + jsonWriter.WriteFieldName(AverageInfo.SumName); + jsonWriter.WriteFloat64Value(this.Sum.Value); + jsonWriter.WriteFieldName(AverageInfo.CountName); + jsonWriter.WriteInt64Value(this.Count); + jsonWriter.WriteObjectEnd(); + } + public override string ToString() { return $@"{{ diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/CountAggregator.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/CountAggregator.cs index cc85a874dc..f1e7bf2f86 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/CountAggregator.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/CountAggregator.cs @@ -6,6 +6,7 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.Aggregate.Aggrega using System; using System.Globalization; using Microsoft.Azure.Cosmos.CosmosElements; + using Microsoft.Azure.Cosmos.Json; using Microsoft.Azure.Cosmos.Query.Core.Exceptions; using Microsoft.Azure.Cosmos.Query.Core.Monads; @@ -66,6 +67,16 @@ public string GetContinuationToken() return this.globalCount.ToString(CultureInfo.InvariantCulture); } + public void SerializeState(IJsonWriter jsonWriter) + { + if (jsonWriter == null) + { + throw new ArgumentNullException(nameof(jsonWriter)); + } + + jsonWriter.WriteInt64Value(this.globalCount); + } + public static TryCatch TryCreate(string continuationToken) { long partialCount; diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/IAggregator.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/IAggregator.cs index 3e5d7373b4..fd7493f5bf 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/IAggregator.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/IAggregator.cs @@ -4,6 +4,7 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.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. @@ -27,5 +28,7 @@ internal interface IAggregator /// /// A continuation token that stores the partial aggregate up till this point. string GetContinuationToken(); + + void SerializeState(IJsonWriter jsonWriter); } } diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/MinMaxAggregator.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/MinMaxAggregator.cs index 3423e6c401..b9d612c02c 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/MinMaxAggregator.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/MinMaxAggregator.cs @@ -5,7 +5,7 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.Aggregate.Aggrega { using System; using Microsoft.Azure.Cosmos.CosmosElements; - using Microsoft.Azure.Cosmos.Query.Core; + 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; @@ -85,7 +85,7 @@ public void Aggregate(CosmosElement localMinMax) { max = null; } - + if (min != null) { localMinMax = min; @@ -144,25 +144,34 @@ public CosmosElement GetResult() public string GetContinuationToken() { - string continuationToken; + IJsonWriter jsonWriter = JsonWriter.Create(JsonSerializationFormat.Text); + this.SerializeState(jsonWriter); + return Utf8StringHelpers.ToString(jsonWriter.GetResult()); + } + + public void SerializeState(IJsonWriter jsonWriter) + { + if (jsonWriter == null) + { + throw new ArgumentNullException(nameof(jsonWriter)); + } + if (this.globalMinMax == ItemComparer.MinValue) { - continuationToken = MinMaxAggregator.MinValueContinuationToken; + jsonWriter.WriteStringValue(MinMaxAggregator.MinValueContinuationToken); } else if (this.globalMinMax == ItemComparer.MaxValue) { - continuationToken = MinMaxAggregator.MaxValueContinuationToken; + jsonWriter.WriteStringValue(MinMaxAggregator.MaxValueContinuationToken); } else if (this.globalMinMax == Undefined) { - continuationToken = MinMaxAggregator.UndefinedContinuationToken; + jsonWriter.WriteStringValue(MinMaxAggregator.UndefinedContinuationToken); } else { - continuationToken = this.globalMinMax.ToString(); + this.globalMinMax.WriteTo(jsonWriter); } - - return continuationToken; } public static TryCatch TryCreateMinAggregator(string continuationToken) diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/SingleGroupAggregator.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/SingleGroupAggregator.cs index ca190e08c8..120a8fa217 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/SingleGroupAggregator.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/SingleGroupAggregator.cs @@ -29,6 +29,8 @@ internal abstract class SingleGroupAggregator public abstract string GetContinuationToken(); + public abstract void SerializeState(IJsonWriter jsonWriter); + public static TryCatch TryCreate( AggregateOperator[] aggregates, IReadOnlyDictionary aggregateAliasToAggregateType, @@ -109,6 +111,16 @@ public override string GetContinuationToken() return this.aggregateValue.GetContinuationToken(); } + public override void SerializeState(IJsonWriter jsonWriter) + { + if (jsonWriter == null) + { + throw new ArgumentNullException(nameof(jsonWriter)); + } + + this.aggregateValue.SerializeState(jsonWriter); + } + public override string ToString() { return this.aggregateValue.ToString(); @@ -151,14 +163,27 @@ public override CosmosElement GetResult() public override string GetContinuationToken() { - Dictionary aliasToContinuationToken = new Dictionary(); + IJsonWriter jsonWriter = JsonWriter.Create(JsonSerializationFormat.Text); + this.SerializeState(jsonWriter); + return Utf8StringHelpers.ToString(jsonWriter.GetResult()); + } + + public override void SerializeState(IJsonWriter jsonWriter) + { + if (jsonWriter == null) + { + throw new ArgumentNullException(nameof(jsonWriter)); + } + + jsonWriter.WriteObjectStart(); + foreach (KeyValuePair kvp in this.aliasToValue) { - aliasToContinuationToken[kvp.Key] = CosmosString.Create(kvp.Value.GetContinuationToken()); + jsonWriter.WriteFieldName(kvp.Key); + kvp.Value.SerializeState(jsonWriter); } - CosmosObject cosmosObject = CosmosObject.Create(aliasToContinuationToken); - return cosmosObject.ToString(); + jsonWriter.WriteObjectEnd(); } public static TryCatch TryCreate( @@ -258,6 +283,8 @@ private abstract class AggregateValue public abstract string GetContinuationToken(); + public abstract void SerializeState(IJsonWriter jsonWriter); + public override string ToString() { return this.Result.ToString(); @@ -300,6 +327,11 @@ public override string GetContinuationToken() return this.aggregator.GetContinuationToken(); } + public override void SerializeState(IJsonWriter jsonWriter) + { + this.aggregator.SerializeState(jsonWriter); + } + public static TryCatch TryCreate( AggregateOperator aggregateOperator, string continuationToken) @@ -362,6 +394,19 @@ public override CosmosElement Result public override string GetContinuationToken() { IJsonWriter jsonWriter = JsonWriter.Create(JsonSerializationFormat.Text); + this.SerializeState(jsonWriter); + + string continuationToken = Utf8StringHelpers.ToString(jsonWriter.GetResult()); + return continuationToken; + } + + public override void SerializeState(IJsonWriter jsonWriter) + { + if (jsonWriter == null) + { + throw new ArgumentNullException(nameof(jsonWriter)); + } + jsonWriter.WriteObjectStart(); jsonWriter.WriteFieldName(nameof(this.initialized)); jsonWriter.WriteBoolValue(this.initialized); @@ -371,9 +416,6 @@ public override string GetContinuationToken() this.value.WriteTo(jsonWriter); } jsonWriter.WriteObjectEnd(); - - string continuationToken = Utf8StringHelpers.ToString(jsonWriter.GetResult()); - return continuationToken; } public static TryCatch TryCreate(string continuationToken) diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/SumAggregator.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/SumAggregator.cs index 5efc37e7da..15497b3cd9 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/SumAggregator.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/SumAggregator.cs @@ -6,6 +6,7 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.Aggregate.Aggrega using System; using System.Globalization; using Microsoft.Azure.Cosmos.CosmosElements; + using Microsoft.Azure.Cosmos.Json; using Microsoft.Azure.Cosmos.Query.Core.Exceptions; using Microsoft.Azure.Cosmos.Query.Core.Monads; @@ -74,6 +75,16 @@ public string GetContinuationToken() return this.globalSum.ToString("G17", CultureInfo.InvariantCulture); } + public void SerializeState(IJsonWriter jsonWriter) + { + if (jsonWriter == null) + { + throw new ArgumentNullException(nameof(jsonWriter)); + } + + jsonWriter.WriteFloat64Value(this.globalSum); + } + public static TryCatch TryCreate(string continuationToken) { double partialSum; diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctDocumentQueryExecutionComponent.Client.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctDocumentQueryExecutionComponent.Client.cs index 64cbe0ef53..5bc1229fea 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctDocumentQueryExecutionComponent.Client.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctDocumentQueryExecutionComponent.Client.cs @@ -8,6 +8,7 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.Distinct using System.Threading; using System.Threading.Tasks; 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.Query.Core.QueryClient; @@ -86,6 +87,8 @@ public static async Task> TryCreateAs /// A page of distinct results. public override async Task DrainAsync(int maxElements, CancellationToken cancellationToken) { + cancellationToken.ThrowIfCancellationRequested(); + List distinctResults = new List(); QueryResponseCore sourceResponse = await base.DrainAsync(maxElements, cancellationToken); @@ -129,32 +132,6 @@ public override async Task DrainAsync(int maxElements, Cancel return queryResponseCore; } - - public override bool TryGetContinuationToken(out string continuationToken) - { - if (this.distinctQueryType != DistinctQueryType.Ordered) - { - continuationToken = null; - return false; - } - - if (this.IsDone) - { - continuationToken = null; - return true; - } - - if (!this.Source.TryGetContinuationToken(out string sourceContinuationToken)) - { - continuationToken = default; - return false; - } - - continuationToken = new DistinctContinuationToken( - sourceContinuationToken, - this.distinctMap.GetContinuationToken()).ToString(); - return true; - } } } } diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctDocumentQueryExecutionComponent.Compute.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctDocumentQueryExecutionComponent.Compute.cs index 15aa1b476f..1579ecce74 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctDocumentQueryExecutionComponent.Compute.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctDocumentQueryExecutionComponent.Compute.cs @@ -103,26 +103,6 @@ public override async Task DrainAsync(int maxElements, Cancel diagnostics: sourceResponse.Diagnostics, responseLengthBytes: sourceResponse.ResponseLengthBytes); } - - public override bool TryGetContinuationToken(out string continuationToken) - { - if (this.IsDone) - { - continuationToken = null; - return true; - } - - if (!this.Source.TryGetContinuationToken(out string sourceContinuationToken)) - { - continuationToken = default; - return false; - } - - continuationToken = new DistinctContinuationToken( - sourceContinuationToken, - this.distinctMap.GetContinuationToken()).ToString(); - return true; - } } } } 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 e9ab5739c7..fa46a19e55 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctDocumentQueryExecutionComponent.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctDocumentQueryExecutionComponent.cs @@ -5,6 +5,8 @@ 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.ExecutionContext; using Microsoft.Azure.Cosmos.Query.Core.Monads; using Newtonsoft.Json; @@ -22,6 +24,8 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.Distinct /// internal abstract partial class DistinctDocumentQueryExecutionComponent : DocumentQueryExecutionComponentBase { + private const string SourceTokenName = "SourceToken"; + private const string DistinctMapTokenName = "DistinctMapToken"; /// /// An DistinctMap that efficiently stores the documents that we have already seen. /// @@ -70,6 +74,48 @@ public static async Task> TryCreateAs return tryCreateDistinctDocumentQueryExecutionComponent; } + public override bool TryGetContinuationToken(out string continuationToken) + { + if (this.IsDone) + { + continuationToken = null; + return true; + } + + if (!this.Source.TryGetContinuationToken(out string sourceContinuationToken)) + { + continuationToken = default; + return false; + } + + continuationToken = new DistinctContinuationToken( + sourceContinuationToken, + this.distinctMap.GetContinuationToken()).ToString(); + return true; + } + + public override void SerializeState(IJsonWriter jsonWriter) + { + if (jsonWriter == null) + { + throw new ArgumentNullException(nameof(jsonWriter)); + } + + if (this.IsDone) + { + jsonWriter.WriteNullValue(); + } + else + { + jsonWriter.WriteObjectStart(); + jsonWriter.WriteFieldName(DistinctDocumentQueryExecutionComponent.SourceTokenName); + this.Source.SerializeState(jsonWriter); + jsonWriter.WriteFieldName(DistinctDocumentQueryExecutionComponent.DistinctMapTokenName); + this.distinctMap.SerializeState(jsonWriter); + jsonWriter.WriteObjectEnd(); + } + } + /// /// Continuation token for distinct queries. /// diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctMap.OrderedDistinctMap.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctMap.OrderedDistinctMap.cs index cb6b356eed..832ef85f34 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctMap.OrderedDistinctMap.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctMap.OrderedDistinctMap.cs @@ -5,6 +5,7 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.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.Exceptions; using Microsoft.Azure.Cosmos.Query.Core.Monads; @@ -72,6 +73,16 @@ public override string GetContinuationToken() return this.lastHash.ToString(); } + public override void SerializeState(IJsonWriter jsonWriter) + { + if (jsonWriter == null) + { + throw new ArgumentNullException(nameof(jsonWriter)); + } + + jsonWriter.WriteBinaryValue(UInt128.ToByteArray(this.lastHash)); + } + public static TryCatch TryCreate(string continuationToken) { UInt128 lastHash; diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctMap.UnorderedDistinctMap.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctMap.UnorderedDistinctMap.cs index 07ba695367..9d0cb0c3e4 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctMap.UnorderedDistinctMap.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctMap.UnorderedDistinctMap.cs @@ -230,6 +230,24 @@ public override bool Add(CosmosElement cosmosElement, out UInt128 hash) public override string GetContinuationToken() { IJsonWriter jsonWriter = JsonWriter.Create(JsonSerializationFormat.Binary); + this.SerializeState(jsonWriter); + + ReadOnlyMemory memory = jsonWriter.GetResult(); + if (!MemoryMarshal.TryGetArray(memory, out ArraySegment buffer)) + { + buffer = new ArraySegment(memory.ToArray()); + } + + return Convert.ToBase64String(buffer.Array, buffer.Offset, buffer.Count); + } + + public override void SerializeState(IJsonWriter jsonWriter) + { + if (jsonWriter == null) + { + throw new ArgumentNullException(nameof(jsonWriter)); + } + jsonWriter.WriteObjectStart(); jsonWriter.WriteFieldName(UnorderdDistinctMap.NumbersName); @@ -292,14 +310,6 @@ public override string GetContinuationToken() jsonWriter.WriteStringValue(this.simpleValues.ToString()); jsonWriter.WriteObjectEnd(); - - ReadOnlyMemory memory = jsonWriter.GetResult(); - if (!MemoryMarshal.TryGetArray(memory, out ArraySegment buffer)) - { - buffer = new ArraySegment(memory.ToArray()); - } - - return Convert.ToBase64String(buffer.Array, buffer.Offset, buffer.Count); } /// diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctMap.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctMap.cs index eb01d763aa..fe016415ab 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctMap.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctMap.cs @@ -5,6 +5,7 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.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.Monads; @@ -59,5 +60,7 @@ public static TryCatch TryCreate( public abstract bool Add(CosmosElement cosmosElement, out UInt128 hash); public abstract string GetContinuationToken(); + + public abstract void SerializeState(IJsonWriter jsonWriter); } } diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/DocumentQueryExecutionComponentBase.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/DocumentQueryExecutionComponentBase.cs index 64360225b4..1ab3a9d5b1 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/DocumentQueryExecutionComponentBase.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/DocumentQueryExecutionComponentBase.cs @@ -6,6 +6,7 @@ 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; 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 a60dbba29a..b065da4136 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 @@ -8,6 +8,7 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.GroupBy using System.Threading; using System.Threading.Tasks; using Microsoft.Azure.Cosmos.CosmosElements; + using Microsoft.Azure.Cosmos.Json; using Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.Aggregate; using Microsoft.Azure.Cosmos.Query.Core.Metrics; using Microsoft.Azure.Cosmos.Query.Core.Monads; @@ -106,6 +107,11 @@ public override bool TryGetContinuationToken(out string state) state = default; return false; } + + public override void SerializeState(IJsonWriter jsonWriter) + { + throw new NotSupportedException(ContinuationTokenNotSupportedWithGroupBy); + } } } } \ No newline at end of file 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 982487257c..b5984cdd9a 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 @@ -19,11 +19,13 @@ internal abstract partial class GroupByDocumentQueryExecutionComponent : Documen { private sealed class ComputeGroupByDocumentQueryExecutionComponent : GroupByDocumentQueryExecutionComponent { + private const string SourceTokenName = "SourceToken"; + private const string GroupingTableContinuationTokenName = "GroupingTableContinuationToken"; + private const string DoneReadingGroupingsContinuationToken = "DONE"; + private const string UseTryGetContinuationTokenInstead = "Use TryGetContinuationTokenInstead"; + private static readonly IReadOnlyList EmptyResults = new List().AsReadOnly(); private static readonly IReadOnlyDictionary EmptyQueryMetrics = new Dictionary(); - private static readonly string DoneReadingGroupingsContinuationToken = "DONE"; - - private static readonly string UseTryGetContinuationTokenInstead = "Use TryGetContinuationTokenInstead"; private ComputeGroupByDocumentQueryExecutionComponent( IDocumentQueryExecutionComponent source, @@ -151,10 +153,21 @@ public override void SerializeState(IJsonWriter jsonWriter) } else { + jsonWriter.WriteObjectStart(); + + jsonWriter.WriteFieldName(GroupingTableContinuationTokenName); + this.groupingTable.SerializeState(jsonWriter); + jsonWriter.WriteFieldName(SourceTokenName); if (this.Source.IsDone) { - + jsonWriter.WriteStringValue(DoneReadingGroupingsContinuationToken); } + else + { + this.Source.SerializeState(jsonWriter); + } + + jsonWriter.WriteObjectEnd(); } } 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 f83fe3f831..dd493c2857 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/GroupBy/GroupByDocumentQueryExecutionComponent.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/GroupBy/GroupByDocumentQueryExecutionComponent.cs @@ -247,6 +247,18 @@ public IReadOnlyList Drain(int maxItemCount) public string GetContinuationToken() { IJsonWriter jsonWriter = JsonWriter.Create(JsonSerializationFormat.Text); + this.SerializeState(jsonWriter); + string result = Utf8StringHelpers.ToString(jsonWriter.GetResult()); + return result; + } + + public void SerializeState(IJsonWriter jsonWriter) + { + if (jsonWriter == null) + { + throw new ArgumentNullException(nameof(jsonWriter)); + } + jsonWriter.WriteObjectStart(); foreach (KeyValuePair kvp in this.table) { @@ -254,9 +266,6 @@ public string GetContinuationToken() jsonWriter.WriteStringValue(kvp.Value.GetContinuationToken()); } jsonWriter.WriteObjectEnd(); - - string result = Utf8StringHelpers.ToString(jsonWriter.GetResult()); - return result; } public IEnumerator> GetEnumerator => this.table.GetEnumerator(); diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/IDocumentQueryExecutionComponent.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/IDocumentQueryExecutionComponent.cs index 21dfefb31c..812ca08141 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/IDocumentQueryExecutionComponent.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/IDocumentQueryExecutionComponent.cs @@ -6,6 +6,7 @@ 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; 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 071784db64..a9839f1e6e 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/SkipDocumentQueryExecutionComponent.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/SkipDocumentQueryExecutionComponent.cs @@ -11,6 +11,7 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.SkipTake using System.Threading.Tasks; using Microsoft.Azure.Cosmos.Core.Trace; 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.Query.Core.QueryClient; @@ -18,6 +19,9 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.SkipTake internal sealed class SkipDocumentQueryExecutionComponent : DocumentQueryExecutionComponentBase { + private const string SourceTokenName = "SourceToken"; + private const string SkipCountName = "SkipCount"; + private int skipCount; private SkipDocumentQueryExecutionComponent(IDocumentQueryExecutionComponent source, int skipCount) @@ -128,6 +132,28 @@ public override bool TryGetContinuationToken(out string state) } } + public override void SerializeState(IJsonWriter jsonWriter) + { + if (jsonWriter == null) + { + throw new ArgumentNullException(nameof(jsonWriter)); + } + + if (this.IsDone) + { + jsonWriter.WriteNullValue(); + } + else + { + jsonWriter.WriteObjectStart(); + jsonWriter.WriteFieldName(SourceTokenName); + this.Source.SerializeState(jsonWriter); + jsonWriter.WriteFieldName(SkipCountName); + jsonWriter.WriteInt64Value(this.skipCount); + jsonWriter.WriteObjectEnd(); + } + } + /// /// A OffsetContinuationToken is a composition of a source continuation token and how many items to skip from that source. /// 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 9f474fa94b..1f6a56044a 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/TakeDocumentQueryExecutionComponent.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/TakeDocumentQueryExecutionComponent.cs @@ -11,6 +11,7 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.SkipTake using System.Threading.Tasks; using Microsoft.Azure.Cosmos.Core.Trace; 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.Query.Core.QueryClient; @@ -18,6 +19,9 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.SkipTake internal sealed class TakeDocumentQueryExecutionComponent : DocumentQueryExecutionComponentBase { + private const string SourceTokenName = "SourceToken"; + private const string TakeCountName = "TakeCount"; + private readonly TakeEnum takeEnum; private int takeCount; @@ -197,6 +201,28 @@ public override bool TryGetContinuationToken(out string state) } } + public override void SerializeState(IJsonWriter jsonWriter) + { + if (jsonWriter == null) + { + throw new ArgumentNullException(nameof(jsonWriter)); + } + + if (this.IsDone) + { + jsonWriter.WriteNullValue(); + } + else + { + jsonWriter.WriteObjectStart(); + jsonWriter.WriteFieldName(SourceTokenName); + this.Source.SerializeState(jsonWriter); + jsonWriter.WriteFieldName(TakeCountName); + jsonWriter.WriteInt64Value(this.takeCount); + jsonWriter.WriteObjectEnd(); + } + } + private enum TakeEnum { Limit, From ab6bcc124d2fd60ecb08644a6482d98cd41f0bfd Mon Sep 17 00:00:00 2001 From: Brandon Chong Date: Fri, 17 Jan 2020 14:36:52 -0800 Subject: [PATCH 03/28] wiring continuation new continuation token in the read path --- .../src/CosmosElements/CosmosElement.cs | 29 ++- .../src/FeedIteratorInteralOfT.cs | 3 + .../src/Json/IJsonReader.cs | 7 + .../CompositeContinuationToken.cs | 1 + .../CompositeContinuationTokenRefStruct.cs | 2 + .../CosmosElementRequestContinuationToken.cs | 26 +++ .../RequestContinuationToken.cs | 32 +++ .../StringRequestContinuationToken.cs | 25 +++ ...eDocumentQueryExecutionComponent.Client.cs | 5 +- ...DocumentQueryExecutionComponent.Compute.cs | 133 ++++-------- ...ggregateDocumentQueryExecutionComponent.cs | 5 +- .../Aggregators/SingleGroupAggregator.cs | 37 ++-- .../DocumentQueryExecutionComponentBase.cs | 1 - ...pDocumentQueryExecutionComponent.Client.cs | 151 ++++++++++++++ ...DocumentQueryExecutionComponent.Compute.cs | 159 +++++++++++++++ .../SkipDocumentQueryExecutionComponent.cs | 189 +++--------------- .../CatchAllCosmosQueryExecutionContext.cs | 11 + .../CosmosQueryExecutionContext.cs | 3 + ...ExecutionContextWithNameCacheStaleRetry.cs | 11 + .../LazyCosmosQueryExecutionContext.cs | 22 ++ .../CosmosOrderByItemQueryExecutionContext.cs | 34 +++- ...CosmosParallelItemQueryExecutionContext.cs | 70 +++++-- .../PipelinedDocumentQueryExecutionContext.cs | 32 ++- .../src/RequestOptions/QueryRequestOptions.cs | 3 + .../CrossPartitionQueryTests.cs | 48 +++++ .../CosmosQueryUnitTests.cs | 8 +- 26 files changed, 728 insertions(+), 319 deletions(-) create mode 100644 Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/CosmosElementRequestContinuationToken.cs create mode 100644 Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/RequestContinuationToken.cs create mode 100644 Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/StringRequestContinuationToken.cs create mode 100644 Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/SkipDocumentQueryExecutionComponent.Client.cs create mode 100644 Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/SkipDocumentQueryExecutionComponent.Compute.cs diff --git a/Microsoft.Azure.Cosmos/src/CosmosElements/CosmosElement.cs b/Microsoft.Azure.Cosmos/src/CosmosElements/CosmosElement.cs index 10f7947771..5ee94cd9be 100644 --- a/Microsoft.Azure.Cosmos/src/CosmosElements/CosmosElement.cs +++ b/Microsoft.Azure.Cosmos/src/CosmosElements/CosmosElement.cs @@ -37,6 +37,30 @@ public override string ToString() public abstract void WriteTo(IJsonWriter jsonWriter); + public static bool TryCreateFromBuffer(ReadOnlyMemory buffer, out TCosmosElement cosmosElement) + where TCosmosElement : CosmosElement + { + CosmosElement unTypedCosmosElement; + try + { + unTypedCosmosElement = CreateFromBuffer(buffer); + } + catch (JsonParseException) + { + cosmosElement = default; + return false; + } + + if (!(unTypedCosmosElement is TCosmosElement typedCosmosElement)) + { + cosmosElement = default; + return false; + } + + cosmosElement = typedCosmosElement; + return true; + } + public static CosmosElement CreateFromBuffer(ReadOnlyMemory buffer) { IJsonNavigator jsonNavigator = JsonNavigator.Create(buffer); @@ -129,9 +153,10 @@ public static bool TryParse( string serializedCosmosElement, out CosmosElement cosmosElement) { - if (serializedCosmosElement == null) + if (string.IsNullOrWhiteSpace(serializedCosmosElement)) { - throw new ArgumentNullException(nameof(serializedCosmosElement)); + cosmosElement = default; + return false; } try diff --git a/Microsoft.Azure.Cosmos/src/FeedIteratorInteralOfT.cs b/Microsoft.Azure.Cosmos/src/FeedIteratorInteralOfT.cs index 78cba950b3..68bf458404 100644 --- a/Microsoft.Azure.Cosmos/src/FeedIteratorInteralOfT.cs +++ b/Microsoft.Azure.Cosmos/src/FeedIteratorInteralOfT.cs @@ -4,6 +4,8 @@ namespace Microsoft.Azure.Cosmos { + using Microsoft.Azure.Cosmos.Json; + /// /// Internal API for FeedIterator for inheritance and mocking purposes. /// @@ -12,5 +14,6 @@ internal abstract class FeedIteratorInternal : FeedIterator #pragma warning restore SA1649 { public abstract bool TryGetContinuationToken(out string continuationToken); + public abstract void SerializeState(IJsonWriter jsonWriter); } } \ No newline at end of file diff --git a/Microsoft.Azure.Cosmos/src/Json/IJsonReader.cs b/Microsoft.Azure.Cosmos/src/Json/IJsonReader.cs index 1722eaab54..44bda39f20 100644 --- a/Microsoft.Azure.Cosmos/src/Json/IJsonReader.cs +++ b/Microsoft.Azure.Cosmos/src/Json/IJsonReader.cs @@ -49,6 +49,13 @@ interface IJsonReader /// The next JSON token from the JsonReader as a string. string GetStringValue(); + /// + /// Tries to get the buffered UTF-8 string value from the JsonReader. + /// + /// The buffered UTF-8 string value if set. + /// True if the buffered UTF-8 value is found else False. + bool TryGetBufferedUtf8StringValue(out ReadOnlyMemory bufferedUtf8StringValue); + /// /// Gets current JSON token from the JsonReader as a raw series of bytes that is buffered. /// diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/CompositeContinuationToken.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/CompositeContinuationToken.cs index 4bf81b10fa..fad272c177 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/CompositeContinuationToken.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/CompositeContinuationToken.cs @@ -6,6 +6,7 @@ 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; diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/CompositeContinuationTokenRefStruct.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/CompositeContinuationTokenRefStruct.cs index f61f3f2a0a..3fb3e60764 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/CompositeContinuationTokenRefStruct.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/CompositeContinuationTokenRefStruct.cs @@ -6,6 +6,8 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ContinuationTokens { using System; using Microsoft.Azure.Cosmos.Json; + using Microsoft.Azure.Cosmos.Query.Core.Exceptions; + using Microsoft.Azure.Cosmos.Query.Core.Monads; internal ref struct CompositeContinuationTokenRefStruct { diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/CosmosElementRequestContinuationToken.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/CosmosElementRequestContinuationToken.cs new file mode 100644 index 0000000000..8b8d85711a --- /dev/null +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/CosmosElementRequestContinuationToken.cs @@ -0,0 +1,26 @@ +// ------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +// ------------------------------------------------------------ + +namespace Microsoft.Azure.Cosmos.Query.Core.ContinuationTokens +{ + using Microsoft.Azure.Cosmos.CosmosElements; + + internal sealed class CosmosElementRequestContinuationToken : RequestContinuationToken + { + public CosmosElementRequestContinuationToken(CosmosElement continuationToken) + { + this.Value = continuationToken; + } + + public CosmosElement Value { get; } + + public override bool IsNull => this.Value == null; + + public override bool TryConvertToCosmosElement(out TCosmosElement cosmosElement) + { + cosmosElement = (TCosmosElement)this.Value; + return true; + } + } +} diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/RequestContinuationToken.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/RequestContinuationToken.cs new file mode 100644 index 0000000000..c18e9d6edf --- /dev/null +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/RequestContinuationToken.cs @@ -0,0 +1,32 @@ +// ------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +// ------------------------------------------------------------ + +namespace Microsoft.Azure.Cosmos.Query.Core.ContinuationTokens +{ + using System; + using Microsoft.Azure.Cosmos.CosmosElements; + using Microsoft.Azure.Cosmos.Json; + + /// + /// Base class that all continuation tokens will follow. + /// This serves as an adapter pattern, so that all different types of continuation tokens can have a common interface. + /// + internal abstract class RequestContinuationToken + { + public static RequestContinuationToken Create(string continuationToken) + { + return new StringRequestContinuationToken(continuationToken); + } + + public static RequestContinuationToken Create(CosmosElement continuationToken) + { + return new CosmosElementRequestContinuationToken(continuationToken); + } + + public abstract bool IsNull { get; } + + public abstract bool TryConvertToCosmosElement(out TCosmosElement cosmosElement) + where TCosmosElement : CosmosElement; + } +} diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/StringRequestContinuationToken.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/StringRequestContinuationToken.cs new file mode 100644 index 0000000000..d099384ecd --- /dev/null +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/StringRequestContinuationToken.cs @@ -0,0 +1,25 @@ +// ------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +// ------------------------------------------------------------ + +namespace Microsoft.Azure.Cosmos.Query.Core.ContinuationTokens +{ + using Microsoft.Azure.Cosmos.CosmosElements; + + internal sealed class StringRequestContinuationToken : RequestContinuationToken + { + public StringRequestContinuationToken(string continuationToken) + { + this.Value = continuationToken; + } + + public string Value { get; } + + public override bool IsNull => string.IsNullOrWhiteSpace(this.Value); + + public override bool TryConvertToCosmosElement(out TCosmosElement cosmosElement) + { + return CosmosElement.TryParse(this.Value, out cosmosElement); + } + } +} diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/AggregateDocumentQueryExecutionComponent.Client.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/AggregateDocumentQueryExecutionComponent.Client.cs index 428a527344..a124bf0ce4 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/AggregateDocumentQueryExecutionComponent.Client.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/AggregateDocumentQueryExecutionComponent.Client.cs @@ -8,6 +8,7 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.Aggregate 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.ExecutionComponent.Aggregate.Aggregators; using Microsoft.Azure.Cosmos.Query.Core.Metrics; using Microsoft.Azure.Cosmos.Query.Core.Monads; @@ -31,8 +32,8 @@ public static async Task> TryCreateAs IReadOnlyDictionary aliasToAggregateType, IReadOnlyList orderedAliases, bool hasSelectValue, - string continuationToken, - Func>> tryCreateSourceAsync) + RequestContinuationToken continuationToken, + Func>> tryCreateSourceAsync) { if (tryCreateSourceAsync == null) { diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/AggregateDocumentQueryExecutionComponent.Compute.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/AggregateDocumentQueryExecutionComponent.Compute.cs index 267ae2929e..42c6cdfda1 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/AggregateDocumentQueryExecutionComponent.Compute.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/AggregateDocumentQueryExecutionComponent.Compute.cs @@ -9,6 +9,7 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.Aggregate 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.Aggregators; using Microsoft.Azure.Cosmos.Query.Core.Monads; @@ -37,26 +38,21 @@ public static async Task> TryCreateAs IReadOnlyDictionary aliasToAggregateType, IReadOnlyList orderedAliases, bool hasSelectValue, - string requestContinuation, - Func>> tryCreateSourceAsync) + RequestContinuationToken requestContinuation, + Func>> tryCreateSourceAsync) { - string sourceContinuationToken; - string singleGroupAggregatorContinuationToken; + AggregateContinuationToken aggregateContinuationToken; if (requestContinuation != null) { - if (!AggregateContinuationToken.TryParse(requestContinuation, out AggregateContinuationToken aggregateContinuationToken)) + if (!AggregateContinuationToken.TryParse(requestContinuation, out aggregateContinuationToken)) { return TryCatch.FromException( new MalformedContinuationTokenException($"Malfomed {nameof(AggregateContinuationToken)}: '{requestContinuation}'")); } - - sourceContinuationToken = aggregateContinuationToken.SourceContinuationToken; - singleGroupAggregatorContinuationToken = aggregateContinuationToken.SingleGroupAggregatorContinuationToken; } else { - sourceContinuationToken = null; - singleGroupAggregatorContinuationToken = null; + aggregateContinuationToken = new AggregateContinuationToken(null, null); } TryCatch tryCreateSingleGroupAggregator = SingleGroupAggregator.TryCreate( @@ -64,7 +60,7 @@ public static async Task> TryCreateAs aliasToAggregateType, orderedAliases, hasSelectValue, - singleGroupAggregatorContinuationToken); + RequestContinuationToken.Create(aggregateContinuationToken.SingleGroupAggregatorContinuationToken)); if (!tryCreateSingleGroupAggregator.Succeeded) { @@ -72,12 +68,12 @@ public static async Task> TryCreateAs tryCreateSingleGroupAggregator.Exception); } - return (await tryCreateSourceAsync(sourceContinuationToken)).Try((source) => + return (await tryCreateSourceAsync(RequestContinuationToken.Create(aggregateContinuationToken.SourceContinuationToken))).Try((source) => { return new ComputeAggregateDocumentQueryExecutionComponent( - source, - tryCreateSingleGroupAggregator.Result, - hasSelectValue); + source, + tryCreateSingleGroupAggregator.Result, + hasSelectValue); }); } @@ -183,10 +179,9 @@ public override bool TryGetContinuationToken(out string state) return false; } - AggregateContinuationToken aggregateContinuationToken = AggregateContinuationToken.Create( - this.singleGroupAggregator.GetContinuationToken(), - sourceState); - state = aggregateContinuationToken.ToString(); + IJsonWriter jsonWriter = JsonWriter.Create(JsonSerializationFormat.Text); + this.SerializeState(jsonWriter); + state = Utf8StringHelpers.ToString(jsonWriter.GetResult()); return true; } @@ -197,11 +192,7 @@ public override void SerializeState(IJsonWriter jsonWriter) throw new ArgumentNullException(nameof(jsonWriter)); } - if (this.IsDone) - { - jsonWriter.WriteNullValue(); - } - else + if (!this.IsDone) { jsonWriter.WriteObjectStart(); jsonWriter.WriteFieldName(ComputeAggregateDocumentQueryExecutionComponent.SourceTokenName); @@ -212,69 +203,38 @@ public override void SerializeState(IJsonWriter jsonWriter) } } - private struct AggregateContinuationToken + private readonly struct AggregateContinuationToken { private const string SingleGroupAggregatorContinuationTokenName = "SingleGroupAggregatorContinuationToken"; private const string SourceContinuationTokenName = "SourceContinuationToken"; - private readonly CosmosObject rawCosmosObject; - - private AggregateContinuationToken(CosmosObject rawCosmosObject) + public AggregateContinuationToken( + CosmosElement singleGroupAggregatorContinuationToken, + CosmosElement sourceContinuationToken) { - if (rawCosmosObject == null) - { - throw new ArgumentNullException(nameof(rawCosmosObject)); - } - - CosmosElement rawSingleGroupAggregatorContinuationToken = rawCosmosObject[AggregateContinuationToken.SingleGroupAggregatorContinuationTokenName]; - if (!(rawSingleGroupAggregatorContinuationToken is CosmosString singleGroupAggregatorContinuationToken)) - { - throw new ArgumentException($"{nameof(rawCosmosObject)} had a property that was not a string."); - } - - CosmosElement rawSourceContinuationToken = rawCosmosObject[AggregateContinuationToken.SourceContinuationTokenName]; - if (!(rawSourceContinuationToken is CosmosString sourceContinuationToken)) - { - throw new ArgumentException($"{nameof(rawCosmosObject)} had a property that was not a string."); - } - - this.rawCosmosObject = rawCosmosObject; + this.SingleGroupAggregatorContinuationToken = singleGroupAggregatorContinuationToken; + this.SourceContinuationToken = sourceContinuationToken; } - public static AggregateContinuationToken Create( - string singleGroupAggregatorContinuationToken, - string sourceContinuationToken) - { - if (singleGroupAggregatorContinuationToken == null) - { - throw new ArgumentNullException(nameof(singleGroupAggregatorContinuationToken)); - } + public CosmosElement SingleGroupAggregatorContinuationToken { get; } - if (sourceContinuationToken == null) - { - throw new ArgumentNullException(nameof(sourceContinuationToken)); - } - - Dictionary dictionary = new Dictionary - { - [AggregateContinuationToken.SingleGroupAggregatorContinuationTokenName] = CosmosString.Create(singleGroupAggregatorContinuationToken), - [AggregateContinuationToken.SourceContinuationTokenName] = CosmosString.Create(sourceContinuationToken) - }; - - CosmosObject rawCosmosObject = CosmosObject.Create(dictionary); - return new AggregateContinuationToken(rawCosmosObject); - } + public CosmosElement SourceContinuationToken { get; } public static bool TryParse( - string serializedContinuationToken, + RequestContinuationToken continuationToken, out AggregateContinuationToken aggregateContinuationToken) { - if (serializedContinuationToken == null) + if (continuationToken == null) + { + throw new ArgumentNullException(nameof(continuationToken)); + } + + if (!(continuationToken is CosmosElementRequestContinuationToken cosmosElementRequestContinuationToken)) { - throw new ArgumentNullException(nameof(serializedContinuationToken)); + throw new ArgumentException($"Expected {nameof(CosmosElementRequestContinuationToken)} instead of: {continuationToken.GetType()}"); } - if (!CosmosElement.TryParse(serializedContinuationToken, out CosmosObject rawAggregateContinuationToken)) + if (!(cosmosElementRequestContinuationToken.Value is CosmosObject rawAggregateContinuationToken)) { aggregateContinuationToken = default; return false; @@ -282,7 +242,7 @@ public static bool TryParse( if (!rawAggregateContinuationToken.TryGetValue( AggregateContinuationToken.SingleGroupAggregatorContinuationTokenName, - out CosmosString singleGroupAggregatorContinuationToken)) + out CosmosElement singleGroupAggregatorContinuationToken)) { aggregateContinuationToken = default; return false; @@ -290,36 +250,17 @@ public static bool TryParse( if (!rawAggregateContinuationToken.TryGetValue( AggregateContinuationToken.SourceContinuationTokenName, - out CosmosString sourceContinuationToken)) + out CosmosElement sourceContinuationToken)) { aggregateContinuationToken = default; return false; } - aggregateContinuationToken = new AggregateContinuationToken(rawAggregateContinuationToken); + aggregateContinuationToken = new AggregateContinuationToken( + singleGroupAggregatorContinuationToken: singleGroupAggregatorContinuationToken, + sourceContinuationToken: sourceContinuationToken); return true; } - - public override string ToString() - { - return this.rawCosmosObject.ToString(); - } - - public string SingleGroupAggregatorContinuationToken - { - get - { - return (this.rawCosmosObject[AggregateContinuationToken.SingleGroupAggregatorContinuationTokenName] as CosmosString).Value; - } - } - - public string SourceContinuationToken - { - get - { - return (this.rawCosmosObject[AggregateContinuationToken.SourceContinuationTokenName] as CosmosString).Value; - } - } } } } 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 f0fb85c984..ab8fa68cda 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/AggregateDocumentQueryExecutionComponent.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/AggregateDocumentQueryExecutionComponent.cs @@ -7,6 +7,7 @@ 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; @@ -61,8 +62,8 @@ public static async Task> TryCreateAs IReadOnlyDictionary aliasToAggregateType, IReadOnlyList orderedAliases, bool hasSelectValue, - string continuationToken, - Func>> tryCreateSourceAsync) + RequestContinuationToken continuationToken, + Func>> tryCreateSourceAsync) { if (tryCreateSourceAsync == null) { diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/SingleGroupAggregator.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/SingleGroupAggregator.cs index 120a8fa217..501d07e99c 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/SingleGroupAggregator.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/SingleGroupAggregator.cs @@ -8,6 +8,7 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.Aggregate.Aggrega 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; @@ -36,7 +37,7 @@ public static TryCatch TryCreate( IReadOnlyDictionary aggregateAliasToAggregateType, IReadOnlyList orderedAliases, bool hasSelectValue, - string continuationToken) + RequestContinuationToken continuationToken) { if (aggregates == null) { @@ -90,7 +91,7 @@ private SelectValueAggregateValues(AggregateValue aggregateValue) this.aggregateValue = aggregateValue ?? throw new ArgumentNullException(nameof(AggregateValue)); } - public static TryCatch TryCreate(AggregateOperator? aggregateOperator, string continuationToken) + public static TryCatch TryCreate(AggregateOperator? aggregateOperator, RequestContinuationToken continuationToken) { return AggregateValue.TryCreate(aggregateOperator, continuationToken) .Try((aggregateValue) => (SingleGroupAggregator)new SelectValueAggregateValues(aggregateValue)); @@ -189,12 +190,12 @@ public override void SerializeState(IJsonWriter jsonWriter) public static TryCatch TryCreate( IReadOnlyDictionary aggregateAliasToAggregateType, IReadOnlyList orderedAliases, - string continuationToken) + RequestContinuationToken continuationToken) { CosmosObject aliasToContinuationToken; if (continuationToken != null) { - if (!CosmosElement.TryParse(continuationToken, out aliasToContinuationToken)) + if (!continuationToken.TryConvertToCosmosElement(out aliasToContinuationToken)) { return TryCatch.FromException( new MalformedContinuationTokenException( @@ -211,21 +212,14 @@ public static TryCatch TryCreate( { string alias = aliasToAggregate.Key; AggregateOperator? aggregateOperator = aliasToAggregate.Value; - string aliasContinuationToken; + RequestContinuationToken aliasContinuationToken; if (aliasToContinuationToken != null) { - if (!(aliasToContinuationToken[alias] is CosmosString parsedAliasContinuationToken)) - { - return TryCatch.FromException( - new MalformedContinuationTokenException( - $"{nameof(SelectListAggregateValues)} continuation token is malformed: {continuationToken}.")); - } - - aliasContinuationToken = parsedAliasContinuationToken.Value; + aliasContinuationToken = RequestContinuationToken.Create(aliasToContinuationToken[alias]); } else { - aliasContinuationToken = null; + aliasContinuationToken = RequestContinuationToken.Create((CosmosElement)null); } TryCatch tryCreateAggregateValue = AggregateValue.TryCreate( @@ -290,7 +284,7 @@ public override string ToString() return this.Result.ToString(); } - public static TryCatch TryCreate(AggregateOperator? aggregateOperator, string continuationToken) + public static TryCatch TryCreate(AggregateOperator? aggregateOperator, RequestContinuationToken continuationToken) { TryCatch value; if (aggregateOperator.HasValue) @@ -418,15 +412,18 @@ public override void SerializeState(IJsonWriter jsonWriter) jsonWriter.WriteObjectEnd(); } - public static TryCatch TryCreate(string continuationToken) + public static TryCatch TryCreate(RequestContinuationToken continuationToken) { + if (continuationToken == null) + { + throw new ArgumentNullException(nameof(continuationToken)); + } + CosmosElement value; bool initialized; - if (continuationToken != null) + if (continuationToken.IsNull) { - if (!CosmosElement.TryParse( - continuationToken, - out CosmosObject rawContinuationToken)) + if (!continuationToken.TryConvertToCosmosElement(out CosmosObject rawContinuationToken)) { return TryCatch.FromException( new MalformedContinuationTokenException($"Invalid {nameof(ScalarAggregateValue)}: {continuationToken}")); diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/DocumentQueryExecutionComponentBase.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/DocumentQueryExecutionComponentBase.cs index 1ab3a9d5b1..64360225b4 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/DocumentQueryExecutionComponentBase.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/DocumentQueryExecutionComponentBase.cs @@ -6,7 +6,6 @@ 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; diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/SkipDocumentQueryExecutionComponent.Client.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/SkipDocumentQueryExecutionComponent.Client.cs new file mode 100644 index 0000000000..a970e76519 --- /dev/null +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/SkipDocumentQueryExecutionComponent.Client.cs @@ -0,0 +1,151 @@ +//------------------------------------------------------------ +// 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.Json; + using Microsoft.Azure.Cosmos.Query.Core.ContinuationTokens; + using Microsoft.Azure.Cosmos.Query.Core.Exceptions; + using Microsoft.Azure.Cosmos.Query.Core.Monads; + using Newtonsoft.Json; + + internal abstract partial class SkipDocumentQueryExecutionComponent : DocumentQueryExecutionComponentBase + { + private sealed class ClientSkipDocumentQueryExecutionComponent : SkipDocumentQueryExecutionComponent + { + public ClientSkipDocumentQueryExecutionComponent(IDocumentQueryExecutionComponent source, long skipCount) + : base(source, skipCount) + { + // Work is done in base constructor. + } + + public static async Task> TryCreateAsync( + int offsetCount, + RequestContinuationToken continuationToken, + Func>> tryCreateSourceAsync) + { + if (tryCreateSourceAsync == null) + { + throw new ArgumentNullException(nameof(tryCreateSourceAsync)); + } + + if (!(continuationToken is StringRequestContinuationToken stringRequestContinuationToken)) + { + return TryCatch.FromException( + new ArgumentException($"Expected {nameof(RequestContinuationToken)} to be a {nameof(StringRequestContinuationToken)}")); + } + + OffsetContinuationToken offsetContinuationToken; + if (continuationToken != null) + { + if (!OffsetContinuationToken.TryParse(stringRequestContinuationToken.Value, out offsetContinuationToken)) + { + return TryCatch.FromException( + new MalformedContinuationTokenException($"Invalid {nameof(SkipDocumentQueryExecutionComponent)}: {continuationToken}.")); + } + } + else + { + offsetContinuationToken = new OffsetContinuationToken(offsetCount, null); + } + + 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 (await tryCreateSourceAsync(RequestContinuationToken.Create(offsetContinuationToken.SourceToken))) + .Try((source) => new ClientSkipDocumentQueryExecutionComponent( + source, + offsetContinuationToken.Offset)); + } + + public override void SerializeState(IJsonWriter jsonWriter) + { + throw new NotImplementedException(); + } + + public override bool TryGetContinuationToken(out string state) + { + throw new NotImplementedException(); + } + + /// + /// A OffsetContinuationToken is a composition of a source continuation token and how many items to skip from that source. + /// + private readonly struct OffsetContinuationToken + { + /// + /// Initializes a new instance of the OffsetContinuationToken struct. + /// + /// The number of items to skip in the query. + /// The continuation token for the source component of the query. + public OffsetContinuationToken(int offset, string sourceToken) + { + if (offset < 0) + { + throw new ArgumentException($"{nameof(offset)} must be a non negative number."); + } + + this.Offset = offset; + this.SourceToken = sourceToken; + } + + /// + /// The number of items to skip in the query. + /// + [JsonProperty("offset")] + public int Offset + { + get; + } + + /// + /// Gets the continuation token for the source component of the query. + /// + [JsonProperty("sourceToken")] + public string SourceToken + { + get; + } + + /// + /// Tries to parse out the OffsetContinuationToken. + /// + /// The value to parse from. + /// The result of parsing out the token. + /// Whether or not the LimitContinuationToken was successfully parsed out. + public static bool TryParse(string value, out OffsetContinuationToken offsetContinuationToken) + { + offsetContinuationToken = default; + if (string.IsNullOrWhiteSpace(value)) + { + return false; + } + + try + { + offsetContinuationToken = JsonConvert.DeserializeObject(value); + return true; + } + catch (JsonException) + { + return false; + } + } + + /// + /// Gets the string version of the continuation token that can be passed in a response header. + /// + /// The string version of the continuation token that can be passed in a response header. + public override string ToString() + { + return JsonConvert.SerializeObject(this); + } + } + } + } +} \ 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/ExecutionComponent/SkipTake/SkipDocumentQueryExecutionComponent.Compute.cs new file mode 100644 index 0000000000..f1d662167d --- /dev/null +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/SkipDocumentQueryExecutionComponent.Compute.cs @@ -0,0 +1,159 @@ +//------------------------------------------------------------ +// 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.Json; + using Microsoft.Azure.Cosmos.Query.Core.ContinuationTokens; + using Microsoft.Azure.Cosmos.Query.Core.Exceptions; + using Microsoft.Azure.Cosmos.Query.Core.Monads; + + internal abstract partial class SkipDocumentQueryExecutionComponent : DocumentQueryExecutionComponentBase + { + private sealed class ComputeSkipDocumentQueryExecutionComponent : SkipDocumentQueryExecutionComponent + { + private const string SkipCountPropertyName = "SkipCount"; + private const string SourceTokenPropertyName = "SourceToken"; + + public ComputeSkipDocumentQueryExecutionComponent(IDocumentQueryExecutionComponent source, long skipCount) + : base(source, skipCount) + { + // Work is done in base constructor. + } + + public static async Task> TryCreateComputeAsync( + int offsetCount, + RequestContinuationToken continuationToken, + Func>> tryCreateSourceAsync) + { + if (tryCreateSourceAsync == null) + { + throw new ArgumentNullException(nameof(tryCreateSourceAsync)); + } + + if (!(continuationToken is CosmosElementRequestContinuationToken cosmosElementRequestContinuationToken)) + { + return TryCatch.FromException( + new ArgumentException($"Expected {nameof(RequestContinuationToken)} to be a {nameof(StringRequestContinuationToken)}")); + } + + OffsetContinuationToken offsetContinuationToken; + if (continuationToken != null) + { + (bool parsed, OffsetContinuationToken parsedToken) = OffsetContinuationToken.TryParse(cosmosElementRequestContinuationToken.Value); + if (!parsed) + { + return TryCatch.FromException( + new MalformedContinuationTokenException($"Invalid {nameof(SkipDocumentQueryExecutionComponent)}: {continuationToken}.")); + } + + offsetContinuationToken = parsedToken; + } + else + { + offsetContinuationToken = new OffsetContinuationToken(offsetCount, null); + } + + 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 (await tryCreateSourceAsync(RequestContinuationToken.Create(offsetContinuationToken.SourceToken))) + .Try((source) => new ClientSkipDocumentQueryExecutionComponent( + source, + offsetContinuationToken.Offset)); + } + + public override bool TryGetContinuationToken(out string state) + { + throw new NotImplementedException(); + } + + public override void SerializeState(IJsonWriter jsonWriter) + { + if (jsonWriter == null) + { + throw new ArgumentNullException(nameof(jsonWriter)); + } + + if (!this.IsDone) + { + jsonWriter.WriteObjectStart(); + jsonWriter.WriteFieldName(SourceTokenPropertyName); + this.Source.SerializeState(jsonWriter); + jsonWriter.WriteFieldName(SkipCountPropertyName); + jsonWriter.WriteInt64Value(this.skipCount); + jsonWriter.WriteObjectEnd(); + } + } + + /// + /// A OffsetContinuationToken is a composition of a source continuation token and how many items to skip from that source. + /// + private readonly struct OffsetContinuationToken + { + /// + /// Initializes a new instance of the OffsetContinuationToken struct. + /// + /// The number of items to skip in the query. + /// The continuation token for the source component of the query. + public OffsetContinuationToken(long offset, CosmosElement sourceToken) + { + if ((offset < 0) || (offset > int.MaxValue)) + { + throw new ArgumentOutOfRangeException(nameof(offset)); + } + + this.Offset = (int)offset; + this.SourceToken = sourceToken; + } + + /// + /// The number of items to skip in the query. + /// + public int Offset + { + get; + } + + /// + /// Gets the continuation token for the source component of the query. + /// + public CosmosElement SourceToken + { + get; + } + + public static (bool parsed, OffsetContinuationToken offsetContinuationToken) TryParse(CosmosElement value) + { + if (value == null) + { + return (false, default); + } + + if (!(value is CosmosObject cosmosObject)) + { + return (false, default); + } + + if (!cosmosObject.TryGetValue(SkipCountPropertyName, out CosmosNumber offset)) + { + return (false, default); + } + + if (!cosmosObject.TryGetValue(SourceTokenPropertyName, out CosmosElement sourceToken)) + { + return (false, default); + } + + return (true, new OffsetContinuationToken(offset.AsInteger().Value, sourceToken)); + } + } + } + } +} \ No newline at end of file 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 a9839f1e6e..3382cca14e 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/SkipDocumentQueryExecutionComponent.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/SkipDocumentQueryExecutionComponent.cs @@ -5,65 +5,59 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.SkipTake { using System; using System.Collections.Generic; - using System.Globalization; using System.Linq; using System.Threading; using System.Threading.Tasks; - using Microsoft.Azure.Cosmos.Core.Trace; using Microsoft.Azure.Cosmos.CosmosElements; using Microsoft.Azure.Cosmos.Json; - using Microsoft.Azure.Cosmos.Query.Core.Exceptions; + 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.QueryClient; - using Newtonsoft.Json; - internal sealed class SkipDocumentQueryExecutionComponent : DocumentQueryExecutionComponentBase + internal abstract partial class SkipDocumentQueryExecutionComponent : DocumentQueryExecutionComponentBase { - private const string SourceTokenName = "SourceToken"; - private const string SkipCountName = "SkipCount"; - private int skipCount; - private SkipDocumentQueryExecutionComponent(IDocumentQueryExecutionComponent source, int skipCount) + protected SkipDocumentQueryExecutionComponent(IDocumentQueryExecutionComponent source, long skipCount) : base(source) { - this.skipCount = skipCount; + if (skipCount > int.MaxValue) + { + throw new ArgumentOutOfRangeException(nameof(skipCount)); + } + + this.skipCount = (int)skipCount; } - public static async Task> TryCreateAsync( + public static Task> TryCreateAsync( + ExecutionEnvironment executionEnvironment, int offsetCount, - string continuationToken, - Func>> tryCreateSourceAsync) + RequestContinuationToken continuationToken, + Func>> tryCreateSourceAsync) { - if (tryCreateSourceAsync == null) + Task> tryCreate; + switch (executionEnvironment) { - throw new ArgumentNullException(nameof(tryCreateSourceAsync)); - } + case ExecutionEnvironment.Client: + tryCreate = ClientSkipDocumentQueryExecutionComponent.TryCreateAsync( + offsetCount, + continuationToken, + tryCreateSourceAsync); + break; - OffsetContinuationToken offsetContinuationToken; - if (continuationToken != null) - { - if (!OffsetContinuationToken.TryParse(continuationToken, out offsetContinuationToken)) - { - return TryCatch.FromException( - new MalformedContinuationTokenException($"Invalid {nameof(SkipDocumentQueryExecutionComponent)}: {continuationToken}.")); - } - } - else - { - offsetContinuationToken = new OffsetContinuationToken(offsetCount, null); - } + case ExecutionEnvironment.Compute: + tryCreate = ComputeSkipDocumentQueryExecutionComponent.TryCreateComputeAsync( + offsetCount, + continuationToken, + tryCreateSourceAsync); + break; - if (offsetContinuationToken.Offset > offsetCount) - { - return TryCatch.FromException( - new MalformedContinuationTokenException("offset count in continuation token can not be greater than the offsetcount in the query.")); + default: + throw new ArgumentException($"Unknown {nameof(ExecutionEnvironment)}: {executionEnvironment}"); } - return (await tryCreateSourceAsync(offsetContinuationToken.SourceToken)) - .Try((source) => new SkipDocumentQueryExecutionComponent( - source, - offsetContinuationToken.Offset)); + return tryCreate; } public override bool IsDone @@ -107,126 +101,5 @@ public override async Task DrainAsync(int maxElements, Cancel diagnostics: sourcePage.Diagnostics, responseLengthBytes: sourcePage.ResponseLengthBytes); } - - public override bool TryGetContinuationToken(out string state) - { - if (!this.IsDone) - { - if (this.Source.TryGetContinuationToken(out string sourceState)) - { - state = new OffsetContinuationToken( - this.skipCount, - sourceState).ToString(); - return true; - } - else - { - state = null; - return false; - } - } - else - { - state = null; - return true; - } - } - - public override void SerializeState(IJsonWriter jsonWriter) - { - if (jsonWriter == null) - { - throw new ArgumentNullException(nameof(jsonWriter)); - } - - if (this.IsDone) - { - jsonWriter.WriteNullValue(); - } - else - { - jsonWriter.WriteObjectStart(); - jsonWriter.WriteFieldName(SourceTokenName); - this.Source.SerializeState(jsonWriter); - jsonWriter.WriteFieldName(SkipCountName); - jsonWriter.WriteInt64Value(this.skipCount); - jsonWriter.WriteObjectEnd(); - } - } - - /// - /// A OffsetContinuationToken is a composition of a source continuation token and how many items to skip from that source. - /// - private struct OffsetContinuationToken - { - /// - /// Initializes a new instance of the OffsetContinuationToken struct. - /// - /// The number of items to skip in the query. - /// The continuation token for the source component of the query. - public OffsetContinuationToken(int offset, string sourceToken) - { - if (offset < 0) - { - throw new ArgumentException($"{nameof(offset)} must be a non negative number."); - } - - this.Offset = offset; - this.SourceToken = sourceToken; - } - - /// - /// The number of items to skip in the query. - /// - [JsonProperty("offset")] - public int Offset - { - get; - } - - /// - /// Gets the continuation token for the source component of the query. - /// - [JsonProperty("sourceToken")] - public string SourceToken - { - get; - } - - /// - /// Tries to parse out the OffsetContinuationToken. - /// - /// The value to parse from. - /// The result of parsing out the token. - /// Whether or not the LimitContinuationToken was successfully parsed out. - public static bool TryParse(string value, out OffsetContinuationToken offsetContinuationToken) - { - offsetContinuationToken = default; - if (string.IsNullOrWhiteSpace(value)) - { - return false; - } - - try - { - offsetContinuationToken = JsonConvert.DeserializeObject(value); - return true; - } - catch (JsonException ex) - { - DefaultTrace.TraceWarning($"{DateTime.UtcNow.ToString("o", CultureInfo.InvariantCulture)} Invalid continuation token {value} for offset~Component: {ex}"); - return false; - } - } - - /// - /// Gets the string version of the continuation token that can be passed in a response header. - /// - /// The string version of the continuation token that can be passed in a response header. - public override string ToString() - { - return JsonConvert.SerializeObject(this); - } - } } } \ No newline at end of file diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CatchAllCosmosQueryExecutionContext.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CatchAllCosmosQueryExecutionContext.cs index 5b6cd4a3ae..e99f49b09a 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CatchAllCosmosQueryExecutionContext.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CatchAllCosmosQueryExecutionContext.cs @@ -7,6 +7,7 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionContext using System; using System.Threading; using System.Threading.Tasks; + using Microsoft.Azure.Cosmos.Json; using Microsoft.Azure.Cosmos.Query.Core.QueryClient; internal sealed class CatchAllCosmosQueryExecutionContext : CosmosQueryExecutionContext @@ -59,5 +60,15 @@ public override bool TryGetContinuationToken(out string continuationToken) { return this.cosmosQueryExecutionContext.TryGetContinuationToken(out continuationToken); } + + public override void SerializeState(IJsonWriter jsonWriter) + { + if (jsonWriter == null) + { + throw new ArgumentNullException(nameof(jsonWriter)); + } + + this.cosmosQueryExecutionContext.SerializeState(jsonWriter); + } } } diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosQueryExecutionContext.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosQueryExecutionContext.cs index 2ac8b061e5..872eaf90f7 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosQueryExecutionContext.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosQueryExecutionContext.cs @@ -6,6 +6,7 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionContext using System; using System.Threading; using System.Threading.Tasks; + using Microsoft.Azure.Cosmos.Json; using Microsoft.Azure.Cosmos.Query.Core.QueryClient; /// @@ -31,5 +32,7 @@ public abstract bool IsDone public abstract Task ExecuteNextAsync(CancellationToken cancellationToken); public abstract bool TryGetContinuationToken(out string continuationToken); + + public abstract void SerializeState(IJsonWriter jsonWriter); } } diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosQueryExecutionContextWithNameCacheStaleRetry.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosQueryExecutionContextWithNameCacheStaleRetry.cs index 2e8fa6d4a8..01c8df5e72 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosQueryExecutionContextWithNameCacheStaleRetry.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosQueryExecutionContextWithNameCacheStaleRetry.cs @@ -7,6 +7,7 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionContext using System; using System.Threading; using System.Threading.Tasks; + using Microsoft.Azure.Cosmos.Json; using Microsoft.Azure.Cosmos.Query.Core.QueryClient; internal sealed class CosmosQueryExecutionContextWithNameCacheStaleRetry : CosmosQueryExecutionContext @@ -63,5 +64,15 @@ public override bool TryGetContinuationToken(out string continuationToken) { return this.currentCosmosQueryExecutionContext.TryGetContinuationToken(out continuationToken); } + + public override void SerializeState(IJsonWriter jsonWriter) + { + if (jsonWriter == null) + { + throw new ArgumentNullException(nameof(jsonWriter)); + } + + this.currentCosmosQueryExecutionContext.SerializeState(jsonWriter); + } } } diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/LazyCosmosQueryExecutionContext.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/LazyCosmosQueryExecutionContext.cs index 0f04fd4937..4f6f6f04c8 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/LazyCosmosQueryExecutionContext.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/LazyCosmosQueryExecutionContext.cs @@ -7,6 +7,7 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionContext using System; using System.Threading; using System.Threading.Tasks; + using Microsoft.Azure.Cosmos.Json; using Microsoft.Azure.Cosmos.Query.Core.Monads; using Microsoft.Azure.Cosmos.Query.Core.QueryClient; @@ -92,5 +93,26 @@ public override bool TryGetContinuationToken(out string state) return tryCreateCosmosQueryExecutionContext.Result.TryGetContinuationToken(out state); } + + public override void SerializeState(IJsonWriter jsonWriter) + { + if (jsonWriter == null) + { + throw new ArgumentNullException(nameof(jsonWriter)); + } + + if (!this.lazyTryCreateCosmosQueryExecutionContext.ValueInitialized) + { + throw new InvalidOperationException(); + } + + TryCatch tryCreateCosmosQueryExecutionContext = this.lazyTryCreateCosmosQueryExecutionContext.Result; + if (!tryCreateCosmosQueryExecutionContext.Succeeded) + { + throw tryCreateCosmosQueryExecutionContext.Exception; + } + + tryCreateCosmosQueryExecutionContext.Result.SerializeState(jsonWriter); + } } } 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 9a02673eca..7116260816 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/OrderBy/CosmosOrderByItemQueryExecutionContext.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/OrderBy/CosmosOrderByItemQueryExecutionContext.cs @@ -193,7 +193,7 @@ public override void SerializeState(IJsonWriter jsonWriter) public static async Task> TryCreateAsync( CosmosQueryContext queryContext, CosmosCrossPartitionQueryExecutionContext.CrossPartitionInitParams initParams, - string requestContinuationToken, + RequestContinuationToken requestContinuationToken, CancellationToken cancellationToken) { Debug.Assert( @@ -357,7 +357,7 @@ private bool ShouldIncrementSkipCount(ItemProducer currentItemProducer) private async Task TryInitializeAsync( SqlQuerySpec sqlQuerySpec, - string requestContinuation, + RequestContinuationToken requestContinuation, string collectionRid, List partitionKeyRanges, int initialPageSize, @@ -503,7 +503,7 @@ private async Task TryInitializeAsync( } private static TryCatch TryExtractContinuationTokens( - string requestContinuation, + RequestContinuationToken requestContinuation, SortOrder[] sortOrders, string[] orderByExpressions) { @@ -515,12 +515,36 @@ private static TryCatch TryExtractContinuationTokens || orderByExpressions.Length != sortOrders.Length), "Partitioned QueryExecutionInfo returned bogus results."); - if (string.IsNullOrWhiteSpace(requestContinuation)) + if (requestContinuation == null) + { + throw new ArgumentNullException("continuation can not be null or empty."); + } + + if (requestContinuation.IsNull) { throw new ArgumentNullException("continuation can not be null or empty."); } - if (!CosmosArray.TryParse(requestContinuation, out CosmosArray cosmosArray)) + CosmosElement cosmosElement; + switch (requestContinuation) + { + case StringRequestContinuationToken stringRequestContinuationToken: + if (!CosmosElement.TryParse(stringRequestContinuationToken.Value, out cosmosElement)) + { + return TryCatch.FromException( + new MalformedContinuationTokenException($"Order by continuation token is malformed: {stringRequestContinuationToken.Value}.")); + } + break; + + case CosmosElementRequestContinuationToken cosmosElementRequestContinuation: + cosmosElement = cosmosElementRequestContinuation.Value; + break; + + default: + throw new ArgumentException($"Unknown {nameof(RequestContinuationToken)} type: {requestContinuation.GetType()}."); + } + + if (!(cosmosElement is CosmosArray cosmosArray)) { return TryCatch.FromException( new MalformedContinuationTokenException($"Order by continuation token must be an array: {requestContinuation}.")); 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 f4531ce260..efdafebfca 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/Parallel/CosmosParallelItemQueryExecutionContext.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/Parallel/CosmosParallelItemQueryExecutionContext.cs @@ -137,7 +137,7 @@ public override void SerializeState(IJsonWriter jsonWriter) public static async Task> TryCreateAsync( CosmosQueryContext queryContext, CosmosCrossPartitionQueryExecutionContext.CrossPartitionInitParams initParams, - string requestContinuationToken, + RequestContinuationToken requestContinuationToken, CancellationToken cancellationToken) { Debug.Assert( @@ -226,7 +226,7 @@ private async Task> TryInitial string collectionRid, List partitionKeyRanges, int initialPageSize, - string requestContinuation, + RequestContinuationToken requestContinuation, CancellationToken cancellationToken) { cancellationToken.ThrowIfCancellationRequested(); @@ -270,9 +270,9 @@ private async Task> TryInitial /// The subset of partition to actually target and continuation tokens. private static TryCatch TryGetInitializationInfoFromContinuationToken( List partitionKeyRanges, - string continuationToken) + RequestContinuationToken continuationToken) { - if (continuationToken == null) + if (continuationToken.IsNull) { return TryCatch.FromResult( new ParallelInitInfo( @@ -281,7 +281,26 @@ private static TryCatch TryGetInitializationInfoFromContinuati } else { - if (!TryParseContinuationToken(continuationToken, out CompositeContinuationToken[] tokens)) + CosmosElement cosmosElement; + switch (continuationToken) + { + case StringRequestContinuationToken stringRequestContinuationToken: + if (!CosmosElement.TryParse(stringRequestContinuationToken.Value, out cosmosElement)) + { + return TryCatch.FromException( + new MalformedContinuationTokenException($"Order by continuation token is malformed: {stringRequestContinuationToken.Value}.")); + } + break; + + case CosmosElementRequestContinuationToken cosmosElementRequestContinuation: + cosmosElement = cosmosElementRequestContinuation.Value; + break; + + default: + throw new ArgumentException($"Unknown {nameof(RequestContinuationToken)} type: {continuationToken.GetType()}."); + } + + if (!TryParseContinuationToken(cosmosElement, out List tokens)) { return TryCatch.FromException( new MalformedContinuationTokenException($"Invalid format for continuation token {continuationToken} for {nameof(CosmosParallelItemQueryExecutionContext)}")); @@ -308,39 +327,50 @@ private static TryCatch TryGetInitializationInfoFromContinuati } } - private static bool TryParseContinuationToken(string continuationToken, out CompositeContinuationToken[] tokens) + private static bool TryParseContinuationToken(CosmosElement continuationToken, out List tokens) { if (continuationToken == null) { throw new ArgumentNullException(nameof(continuationToken)); } - try + if (!(continuationToken is CosmosArray cosmosArray)) { - tokens = JsonConvert.DeserializeObject(continuationToken, DefaultJsonSerializationSettings.Value); + tokens = default; + return false; + } - if (tokens.Length == 0) + List compositeContinuationTokens = new List(); + foreach (CosmosElement item in cosmosArray) + { + TryCatch tryCreateCompositeContinuationToken = CompositeContinuationToken.TryCreateFromCosmosElement(item); + if (!tryCreateCompositeContinuationToken.Succeeded) { tokens = default; return false; } - foreach (CompositeContinuationToken token in tokens) - { - if ((token.Range == null) || token.Range.IsEmpty) - { - tokens = default; - return false; - } - } - - return true; + compositeContinuationTokens.Add(tryCreateCompositeContinuationToken.Result); } - catch (JsonException) + + tokens = compositeContinuationTokens; + + if (tokens.Count == 0) { tokens = default; return false; } + + foreach (CompositeContinuationToken token in tokens) + { + if ((token.Range == null) || token.Range.IsEmpty) + { + tokens = default; + return false; + } + } + + return true; } private readonly struct ParallelInitInfo diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/PipelinedDocumentQueryExecutionContext.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/PipelinedDocumentQueryExecutionContext.cs index 98742bf069..bb2f3e6a6e 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/PipelinedDocumentQueryExecutionContext.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/PipelinedDocumentQueryExecutionContext.cs @@ -9,6 +9,7 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionContext using System.Threading.Tasks; using Microsoft.Azure.Cosmos; using Microsoft.Azure.Cosmos.CosmosElements; + using Microsoft.Azure.Cosmos.Json; using Microsoft.Azure.Cosmos.Query.Core.ContinuationTokens; using Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent; using Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.Aggregate; @@ -124,11 +125,21 @@ public override bool TryGetContinuationToken(out string state) return this.component.TryGetContinuationToken(out state); } + public override void SerializeState(IJsonWriter jsonWriter) + { + if (jsonWriter == null) + { + throw new ArgumentNullException(nameof(jsonWriter)); + } + + this.component.SerializeState(jsonWriter); + } + public static async Task> TryCreateAsync( ExecutionEnvironment executionEnvironment, CosmosQueryContext queryContext, CosmosCrossPartitionQueryExecutionContext.CrossPartitionInitParams initParams, - string requestContinuationToken, + RequestContinuationToken requestContinuationToken, CancellationToken cancellationToken) { if (queryContext == null) @@ -156,7 +167,7 @@ public static async Task> TryCreateAsync( testSettings: initParams.TestSettings); } - async Task> tryCreateOrderByComponentAsync(string continuationToken) + async Task> tryCreateOrderByComponentAsync(RequestContinuationToken continuationToken) { return (await CosmosOrderByItemQueryExecutionContext.TryCreateAsync( queryContext, @@ -165,7 +176,7 @@ async Task> tryCreateOrderByComponent cancellationToken)).Try(component => component); } - async Task> tryCreateParallelComponentAsync(string continuationToken) + async Task> tryCreateParallelComponentAsync(RequestContinuationToken continuationToken) { return (await CosmosParallelItemQueryExecutionContext.TryCreateAsync( queryContext, @@ -174,7 +185,7 @@ async Task> tryCreateParallelComponen cancellationToken)).Try(component => component); } - Func>> tryCreatePipelineAsync; + Func>> tryCreatePipelineAsync; if (queryInfo.HasOrderBy) { tryCreatePipelineAsync = tryCreateOrderByComponentAsync; @@ -186,7 +197,7 @@ async Task> tryCreateParallelComponen if (queryInfo.HasAggregates && !queryInfo.HasGroupBy) { - Func>> tryCreateSourceAsync = tryCreatePipelineAsync; + Func>> tryCreateSourceAsync = tryCreatePipelineAsync; tryCreatePipelineAsync = async (continuationToken) => { return await AggregateDocumentQueryExecutionComponent.TryCreateAsync( @@ -202,7 +213,7 @@ async Task> tryCreateParallelComponen if (queryInfo.HasDistinct) { - Func>> tryCreateSourceAsync = tryCreatePipelineAsync; + Func>> tryCreateSourceAsync = tryCreatePipelineAsync; tryCreatePipelineAsync = async (continuationToken) => { return await DistinctDocumentQueryExecutionComponent.TryCreateAsync( @@ -215,7 +226,7 @@ async Task> tryCreateParallelComponen if (queryInfo.HasGroupBy) { - Func>> tryCreateSourceAsync = tryCreatePipelineAsync; + Func>> tryCreateSourceAsync = tryCreatePipelineAsync; tryCreatePipelineAsync = async (continuationToken) => { return await GroupByDocumentQueryExecutionComponent.TryCreateAsync( @@ -230,10 +241,11 @@ async Task> tryCreateParallelComponen if (queryInfo.HasOffset) { - Func>> tryCreateSourceAsync = tryCreatePipelineAsync; + Func>> tryCreateSourceAsync = tryCreatePipelineAsync; tryCreatePipelineAsync = async (continuationToken) => { return await SkipDocumentQueryExecutionComponent.TryCreateAsync( + executionEnvironment, queryInfo.Offset.Value, continuationToken, tryCreateSourceAsync); @@ -242,7 +254,7 @@ async Task> tryCreateParallelComponen if (queryInfo.HasLimit) { - Func>> tryCreateSourceAsync = tryCreatePipelineAsync; + Func>> tryCreateSourceAsync = tryCreatePipelineAsync; tryCreatePipelineAsync = async (continuationToken) => { return await TakeDocumentQueryExecutionComponent.TryCreateLimitDocumentQueryExecutionComponentAsync( @@ -254,7 +266,7 @@ async Task> tryCreateParallelComponen if (queryInfo.HasTop) { - Func>> tryCreateSourceAsync = tryCreatePipelineAsync; + Func>> tryCreateSourceAsync = tryCreatePipelineAsync; tryCreatePipelineAsync = async (continuationToken) => { return await TakeDocumentQueryExecutionComponent.TryCreateTopDocumentQueryExecutionComponentAsync( diff --git a/Microsoft.Azure.Cosmos/src/RequestOptions/QueryRequestOptions.cs b/Microsoft.Azure.Cosmos/src/RequestOptions/QueryRequestOptions.cs index 7597336173..ffa06514a7 100644 --- a/Microsoft.Azure.Cosmos/src/RequestOptions/QueryRequestOptions.cs +++ b/Microsoft.Azure.Cosmos/src/RequestOptions/QueryRequestOptions.cs @@ -141,6 +141,8 @@ public ConsistencyLevel? ConsistencyLevel /// public string SessionToken { get; set; } + internal ReadOnlyMemory BinaryContinuationToken { get; set; } + internal CosmosSerializationFormatOptions CosmosSerializationFormatOptions { get; set; } internal ExecutionEnvironment? ExecutionEnvironment { get; set; } @@ -220,6 +222,7 @@ internal QueryRequestOptions Clone() CosmosSerializationFormatOptions = this.CosmosSerializationFormatOptions, Properties = this.Properties, IsEffectivePartitionKeyRouting = this.IsEffectivePartitionKeyRouting, + BinaryContinuationToken = this.BinaryContinuationToken, }; return queryRequestOptions; diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CrossPartitionQueryTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CrossPartitionQueryTests.cs index d3d45c1d4d..5c0ec66775 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CrossPartitionQueryTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CrossPartitionQueryTests.cs @@ -650,6 +650,54 @@ private static async Task> QueryWithTryGetContinuationTokens( return resultsFromTryGetContinuationToken; } + private static async Task> QueryWithSerializeState( + Container container, + string query, + QueryRequestOptions queryRequestOptions = null) + { + if (queryRequestOptions == null) + { + queryRequestOptions = new QueryRequestOptions(); + } + + List resultsFromSerializeState = new List(); + ReadOnlyMemory continuationToken; + do + { + QueryRequestOptions computeRequestOptions = queryRequestOptions.Clone(); + computeRequestOptions.ExecutionEnvironment = Cosmos.Query.Core.ExecutionContext.ExecutionEnvironment.Compute; + + FeedIteratorInternal itemQuery = (FeedIteratorInternal)container.GetItemQueryIterator( + queryText: query, + requestOptions: computeRequestOptions, + continuationToken: continuationToken); + try + { + FeedResponse cosmosQueryResponse = await itemQuery.ReadNextAsync(); + if (queryRequestOptions.MaxItemCount.HasValue) + { + Assert.IsTrue( + cosmosQueryResponse.Count <= queryRequestOptions.MaxItemCount.Value, + "Max Item Count is not being honored"); + } + + resultsFromSerializeState.AddRange(cosmosQueryResponse); + Json.IJsonWriter jsonWriter = Json.JsonWriter.Create(Json.JsonSerializationFormat.Binary); + itemQuery.SerializeState(jsonWriter); + continuationToken = jsonWriter.GetResult(); + } + catch (CosmosException cosmosException) when (cosmosException.StatusCode == (HttpStatusCode)429) + { + itemQuery = (FeedIteratorInternal)container.GetItemQueryIterator( + queryText: query, + requestOptions: queryRequestOptions, + continuationToken: continuationToken); + } + } while (!continuationToken.IsEmpty); + + return resultsFromSerializeState; + } + private static async Task> QueryWithContinuationTokens( Container container, string query, 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 52470424d4..a9f3fc6a98 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/CosmosQueryUnitTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/CosmosQueryUnitTests.cs @@ -14,6 +14,7 @@ namespace Microsoft.Azure.Cosmos.Tests using Microsoft.Azure.Cosmos.CosmosElements; 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; @@ -287,7 +288,7 @@ public async Task TestCosmosQueryPartitionKeyDefinition() private async Task<(IList components, QueryResponseCore response)> GetAllExecutionComponents() { - (Func>> func, QueryResponseCore response) = this.SetupBaseContextToVerifyFailureScenario(); + (Func>> func, QueryResponseCore response) = this.SetupBaseContextToVerifyFailureScenario(); List components = new List(); List operators = new List() @@ -318,6 +319,7 @@ public async Task TestCosmosQueryPartitionKeyDefinition() DistinctQueryType.Ordered)).Result); components.Add((await SkipDocumentQueryExecutionComponent.TryCreateAsync( + ExecutionEnvironment.Client, 5, null, func)).Result); @@ -335,7 +337,7 @@ public async Task TestCosmosQueryPartitionKeyDefinition() return (components, response); } - private (Func>>, QueryResponseCore) SetupBaseContextToVerifyFailureScenario() + private (Func>>, QueryResponseCore) SetupBaseContextToVerifyFailureScenario() { IReadOnlyCollection diagnostics = new List() { @@ -367,7 +369,7 @@ public async Task TestCosmosQueryPartitionKeyDefinition() Mock baseContext = new Mock(); baseContext.Setup(x => x.DrainAsync(It.IsAny(), It.IsAny())).Returns(Task.FromResult(failure)); - Task> callBack(string x) => Task.FromResult>(TryCatch.FromResult(baseContext.Object)); + Task> callBack(RequestContinuationToken x) => Task.FromResult>(TryCatch.FromResult(baseContext.Object)); return (callBack, failure); } } From a5fa91daa6ab36207e564c379b1664f368b2213e Mon Sep 17 00:00:00 2001 From: Brandon Chong Date: Fri, 17 Jan 2020 17:49:08 -0800 Subject: [PATCH 04/28] updated distinct and group by --- .../src/Json/IJsonReader.cs | 7 -- ...tDocumentQueryExecutionComponent.Client.cs | 85 ++++++++++++- ...DocumentQueryExecutionComponent.Compute.cs | 89 ++++++++++++- ...DistinctDocumentQueryExecutionComponent.cs | 99 +-------------- ...yDocumentQueryExecutionComponent.Client.cs | 5 +- ...DocumentQueryExecutionComponent.Compute.cs | 66 ++++------ .../GroupByDocumentQueryExecutionComponent.cs | 17 ++- ...eDocumentQueryExecutionComponent.Client.cs | 118 ++++++++++++++++++ .../TakeDocumentQueryExecutionComponent.cs | 13 +- 9 files changed, 332 insertions(+), 167 deletions(-) create mode 100644 Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/TakeDocumentQueryExecutionComponent.Client.cs diff --git a/Microsoft.Azure.Cosmos/src/Json/IJsonReader.cs b/Microsoft.Azure.Cosmos/src/Json/IJsonReader.cs index 44bda39f20..1722eaab54 100644 --- a/Microsoft.Azure.Cosmos/src/Json/IJsonReader.cs +++ b/Microsoft.Azure.Cosmos/src/Json/IJsonReader.cs @@ -49,13 +49,6 @@ interface IJsonReader /// The next JSON token from the JsonReader as a string. string GetStringValue(); - /// - /// Tries to get the buffered UTF-8 string value from the JsonReader. - /// - /// The buffered UTF-8 string value if set. - /// True if the buffered UTF-8 value is found else False. - bool TryGetBufferedUtf8StringValue(out ReadOnlyMemory bufferedUtf8StringValue); - /// /// Gets current JSON token from the JsonReader as a raw series of bytes that is buffered. /// diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctDocumentQueryExecutionComponent.Client.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctDocumentQueryExecutionComponent.Client.cs index 5bc1229fea..d8d44a1c87 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctDocumentQueryExecutionComponent.Client.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctDocumentQueryExecutionComponent.Client.cs @@ -9,9 +9,11 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.Distinct 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 Newtonsoft.Json; internal abstract partial class DistinctDocumentQueryExecutionComponent : DocumentQueryExecutionComponentBase { @@ -40,17 +42,22 @@ private ClientDistinctDocumentQueryExecutionComponent( } public static async Task> TryCreateAsync( - string requestContinuation, - Func>> tryCreateSourceAsync, + RequestContinuationToken requestContinuation, + Func>> tryCreateSourceAsync, DistinctQueryType distinctQueryType) { + if (!(requestContinuation is StringRequestContinuationToken stringRequestContinuationToken)) + { + throw new ArgumentException($"Unknown {nameof(RequestContinuationToken)} type: {requestContinuation.GetType()}"); + } + DistinctContinuationToken distinctContinuationToken; if (requestContinuation != null) { - if (!DistinctContinuationToken.TryParse(requestContinuation, out distinctContinuationToken)) + if (!DistinctContinuationToken.TryParse(stringRequestContinuationToken, out distinctContinuationToken)) { return TryCatch.FromException( - new MalformedContinuationTokenException($"Invalid {nameof(DistinctContinuationToken)}: {requestContinuation}")); + new MalformedContinuationTokenException($"Invalid {nameof(DistinctContinuationToken)}: {stringRequestContinuationToken}")); } } else @@ -66,7 +73,7 @@ public static async Task> TryCreateAs return TryCatch.FromException(tryCreateDistinctMap.Exception); } - TryCatch tryCreateSource = await tryCreateSourceAsync(distinctContinuationToken.SourceToken); + TryCatch tryCreateSource = await tryCreateSourceAsync(RequestContinuationToken.Create(distinctContinuationToken.SourceToken)); if (!tryCreateSource.Succeeded) { return TryCatch.FromException(tryCreateSource.Exception); @@ -132,6 +139,74 @@ public override async Task DrainAsync(int maxElements, Cancel return queryResponseCore; } + + public override void SerializeState(IJsonWriter jsonWriter) + { + throw new NotImplementedException(); + } + + public override bool TryGetContinuationToken(out string state) + { + throw new NotImplementedException(); + } + + /// + /// Continuation token for distinct queries. + /// + private sealed class DistinctContinuationToken + { + public DistinctContinuationToken(string sourceToken, string distinctMapToken) + { + this.SourceToken = sourceToken; + this.DistinctMapToken = distinctMapToken; + } + + public string SourceToken { get; } + + public string DistinctMapToken { get; } + + /// + /// Tries to parse a DistinctContinuationToken from a string. + /// + /// The value to parse. + /// The output DistinctContinuationToken. + /// True if we successfully parsed the DistinctContinuationToken, else false. + public static bool TryParse( + StringRequestContinuationToken stringRequestContinuationToken, + out DistinctContinuationToken distinctContinuationToken) + { + if (stringRequestContinuationToken == null) + { + throw new ArgumentNullException(nameof(stringRequestContinuationToken)); + } + + if (stringRequestContinuationToken.IsNull) + { + distinctContinuationToken = default; + return false; + } + + try + { + distinctContinuationToken = JsonConvert.DeserializeObject(stringRequestContinuationToken.Value); + return true; + } + catch (JsonException) + { + distinctContinuationToken = default; + return false; + } + } + + /// + /// Gets the serialized form of DistinctContinuationToken + /// + /// The serialized form of DistinctContinuationToken + public override string ToString() + { + return JsonConvert.SerializeObject(this); + } + } } } } diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctDocumentQueryExecutionComponent.Compute.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctDocumentQueryExecutionComponent.Compute.cs index 1579ecce74..f63d214e60 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctDocumentQueryExecutionComponent.Compute.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctDocumentQueryExecutionComponent.Compute.cs @@ -8,6 +8,8 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.Distinct 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; @@ -31,8 +33,8 @@ private ComputeDistinctDocumentQueryExecutionComponent( } public static async Task> TryCreateAsync( - string requestContinuation, - Func>> tryCreateSourceAsync, + RequestContinuationToken requestContinuation, + Func>> tryCreateSourceAsync, DistinctQueryType distinctQueryType) { DistinctContinuationToken distinctContinuationToken; @@ -103,6 +105,89 @@ public override async Task DrainAsync(int maxElements, Cancel diagnostics: sourceResponse.Diagnostics, responseLengthBytes: sourceResponse.ResponseLengthBytes); } + + public override bool TryGetContinuationToken(out string continuationToken) + { + if (this.IsDone) + { + continuationToken = null; + return true; + } + + if (!this.Source.TryGetContinuationToken(out string sourceContinuationToken)) + { + continuationToken = default; + return false; + } + + IJsonWriter jsonWriter = JsonWriter.Create(JsonSerializationFormat.Text); + this.SerializeState(jsonWriter); + continuationToken = Utf8StringHelpers.ToString(jsonWriter.GetResult()); + return true; + } + + public override void SerializeState(IJsonWriter jsonWriter) + { + if (jsonWriter == null) + { + throw new ArgumentNullException(nameof(jsonWriter)); + } + + if (!this.IsDone) + { + jsonWriter.WriteObjectStart(); + jsonWriter.WriteFieldName(DistinctDocumentQueryExecutionComponent.SourceTokenName); + this.Source.SerializeState(jsonWriter); + jsonWriter.WriteFieldName(DistinctDocumentQueryExecutionComponent.DistinctMapTokenName); + this.distinctMap.SerializeState(jsonWriter); + jsonWriter.WriteObjectEnd(); + } + } + + private readonly struct DistinctContinuationToken + { + public DistinctContinuationToken(CosmosElement sourceToken, CosmosElement distinctMapToken) + { + this.SourceToken = sourceToken; + this.DistinctMapToken = distinctMapToken; + } + + public CosmosElement SourceToken { get; } + + public CosmosElement DistinctMapToken { get; } + + public static bool TryParse( + CosmosElement requestContinuationToken, + out DistinctContinuationToken distinctContinuationToken) + { + if (requestContinuationToken == null) + { + distinctContinuationToken = default; + return false; + } + + if (!(requestContinuationToken is CosmosObject rawObject)) + { + distinctContinuationToken = default; + return false; + } + + if (!rawObject.TryGetValue(SourceTokenName, out CosmosElement sourceToken)) + { + distinctContinuationToken = default; + return false; + } + + if (!rawObject.TryGetValue(DistinctMapTokenName, out CosmosElement distinctMapToken)) + { + distinctContinuationToken = default; + return false; + } + + distinctContinuationToken = new DistinctContinuationToken(sourceToken: sourceToken, distinctMapToken: distinctMapToken); + return true; + } + } } } } 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 fa46a19e55..3a535af119 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctDocumentQueryExecutionComponent.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctDocumentQueryExecutionComponent.cs @@ -7,6 +7,7 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.Distinct 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; @@ -41,8 +42,8 @@ protected DistinctDocumentQueryExecutionComponent( public static async Task> TryCreateAsync( ExecutionEnvironment executionEnvironment, - string requestContinuation, - Func>> tryCreateSourceAsync, + RequestContinuationToken requestContinuation, + Func>> tryCreateSourceAsync, DistinctQueryType distinctQueryType) { if (tryCreateSourceAsync == null) @@ -73,99 +74,5 @@ public static async Task> TryCreateAs return tryCreateDistinctDocumentQueryExecutionComponent; } - - public override bool TryGetContinuationToken(out string continuationToken) - { - if (this.IsDone) - { - continuationToken = null; - return true; - } - - if (!this.Source.TryGetContinuationToken(out string sourceContinuationToken)) - { - continuationToken = default; - return false; - } - - continuationToken = new DistinctContinuationToken( - sourceContinuationToken, - this.distinctMap.GetContinuationToken()).ToString(); - return true; - } - - public override void SerializeState(IJsonWriter jsonWriter) - { - if (jsonWriter == null) - { - throw new ArgumentNullException(nameof(jsonWriter)); - } - - if (this.IsDone) - { - jsonWriter.WriteNullValue(); - } - else - { - jsonWriter.WriteObjectStart(); - jsonWriter.WriteFieldName(DistinctDocumentQueryExecutionComponent.SourceTokenName); - this.Source.SerializeState(jsonWriter); - jsonWriter.WriteFieldName(DistinctDocumentQueryExecutionComponent.DistinctMapTokenName); - this.distinctMap.SerializeState(jsonWriter); - jsonWriter.WriteObjectEnd(); - } - } - - /// - /// Continuation token for distinct queries. - /// - private sealed class DistinctContinuationToken - { - public DistinctContinuationToken(string sourceToken, string distinctMapToken) - { - this.SourceToken = sourceToken; - this.DistinctMapToken = distinctMapToken; - } - - public string SourceToken { get; } - - public string DistinctMapToken { get; } - - /// - /// Tries to parse a DistinctContinuationToken from a string. - /// - /// The value to parse. - /// The output DistinctContinuationToken. - /// True if we successfully parsed the DistinctContinuationToken, else false. - public static bool TryParse( - string value, - out DistinctContinuationToken distinctContinuationToken) - { - distinctContinuationToken = default; - if (string.IsNullOrWhiteSpace(value)) - { - return false; - } - - try - { - distinctContinuationToken = JsonConvert.DeserializeObject(value); - return true; - } - catch (JsonException) - { - return false; - } - } - - /// - /// Gets the serialized form of DistinctContinuationToken - /// - /// The serialized form of DistinctContinuationToken - public override string ToString() - { - return JsonConvert.SerializeObject(this); - } - } } } 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 b065da4136..e70321e712 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 @@ -9,6 +9,7 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.GroupBy 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.ExecutionComponent.Aggregate; using Microsoft.Azure.Cosmos.Query.Core.Metrics; using Microsoft.Azure.Cosmos.Query.Core.Monads; @@ -30,8 +31,8 @@ private ClientGroupByDocumentQueryExecutionComponent( } public static async Task> TryCreateAsync( - string requestContinuation, - Func>> tryCreateSource, + RequestContinuationToken requestContinuation, + Func>> tryCreateSource, IReadOnlyDictionary groupByAliasToAggregateType, IReadOnlyList orderedAliases, bool hasSelectValue) 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 b5984cdd9a..1836b1e371 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 @@ -9,6 +9,7 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.GroupBy 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; @@ -37,16 +38,21 @@ private ComputeGroupByDocumentQueryExecutionComponent( } public static async Task> TryCreateAsync( - string requestContinuation, - Func>> tryCreateSourceAsync, + RequestContinuationToken requestContinuation, + Func>> tryCreateSourceAsync, IReadOnlyDictionary groupByAliasToAggregateType, IReadOnlyList orderedAliases, bool hasSelectValue) { + if (!(requestContinuation is CosmosElementRequestContinuationToken cosmosElementRequestContinuationToken)) + { + throw new ArgumentException($"Unknown {nameof(RequestContinuationToken)} type: {requestContinuation.GetType()}."); + } + GroupByContinuationToken groupByContinuationToken; if (requestContinuation != null) { - if (!GroupByContinuationToken.TryParse(requestContinuation, out groupByContinuationToken)) + if (!GroupByContinuationToken.TryParse(cosmosElementRequestContinuationToken.Value, out groupByContinuationToken)) { return TryCatch.FromException( new MalformedContinuationTokenException($"Invalid {nameof(GroupByContinuationToken)}: '{requestContinuation}'")); @@ -60,13 +66,15 @@ public static async Task> TryCreateAs } TryCatch tryCreateSource; - if (groupByContinuationToken.SourceContinuationToken == ComputeGroupByDocumentQueryExecutionComponent.DoneReadingGroupingsContinuationToken) + + if ((groupByContinuationToken.SourceContinuationToken is CosmosString sourceContinuationToken) + && (sourceContinuationToken.Value == ComputeGroupByDocumentQueryExecutionComponent.DoneReadingGroupingsContinuationToken)) { tryCreateSource = TryCatch.FromResult(DoneDocumentQueryExecutionComponent.Value); } else { - tryCreateSource = await tryCreateSourceAsync(groupByContinuationToken.SourceContinuationToken); + tryCreateSource = await tryCreateSourceAsync(RequestContinuationToken.Create(groupByContinuationToken.SourceContinuationToken)); } if (!tryCreateSource.Succeeded) @@ -147,11 +155,7 @@ public override void SerializeState(IJsonWriter jsonWriter) throw new ArgumentNullException(nameof(jsonWriter)); } - if (this.IsDone) - { - jsonWriter.WriteNullValue(); - } - else + if (!this.IsDone) { jsonWriter.WriteObjectStart(); @@ -185,49 +189,29 @@ public override bool TryGetContinuationToken(out string continuationToken) return false; } - if (this.Source.IsDone) - { - continuationToken = new GroupByContinuationToken( - this.groupingTable.GetContinuationToken(), - ComputeGroupByDocumentQueryExecutionComponent.DoneReadingGroupingsContinuationToken).ToString(); - } - else - { - // Still need to drain the source. - continuationToken = new GroupByContinuationToken( - this.groupingTable.GetContinuationToken(), - sourceContinuationToken).ToString(); - } - + IJsonWriter jsonWriter = JsonWriter.Create(JsonSerializationFormat.Text); + this.SerializeState(jsonWriter); + continuationToken = Utf8StringHelpers.ToString(jsonWriter.GetResult()); return true; } private readonly struct GroupByContinuationToken { public GroupByContinuationToken( - string groupingTableContinuationToken, - string sourceContinuationToken) + CosmosElement groupingTableContinuationToken, + CosmosElement sourceContinuationToken) { this.GroupingTableContinuationToken = groupingTableContinuationToken; this.SourceContinuationToken = sourceContinuationToken; } - public string GroupingTableContinuationToken { get; } + public CosmosElement GroupingTableContinuationToken { get; } - public string SourceContinuationToken { get; } - - public override string ToString() - { - return CosmosObject.Create(new Dictionary() - { - { nameof(this.GroupingTableContinuationToken), CosmosString.Create(this.GroupingTableContinuationToken) }, - { nameof(this.SourceContinuationToken), CosmosString.Create(this.SourceContinuationToken) } - }).ToString(); - } + public CosmosElement SourceContinuationToken { get; } - public static bool TryParse(string value, out GroupByContinuationToken groupByContinuationToken) + public static bool TryParse(CosmosElement value, out GroupByContinuationToken groupByContinuationToken) { - if (!CosmosElement.TryParse(value, out CosmosObject groupByContinuationTokenObject)) + if (!(value is CosmosObject groupByContinuationTokenObject)) { groupByContinuationToken = default; return false; @@ -250,8 +234,8 @@ public static bool TryParse(string value, out GroupByContinuationToken groupByCo } groupByContinuationToken = new GroupByContinuationToken( - groupingTableContinuationToken.Value, - sourceContinuationToken.Value); + groupingTableContinuationToken: groupingTableContinuationToken, + sourceContinuationToken: sourceContinuationToken); return true; } } 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 dd493c2857..8537a8e184 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/GroupBy/GroupByDocumentQueryExecutionComponent.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/GroupBy/GroupByDocumentQueryExecutionComponent.cs @@ -10,6 +10,7 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.GroupBy 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; @@ -60,8 +61,8 @@ protected GroupByDocumentQueryExecutionComponent( public static async Task> TryCreateAsync( ExecutionEnvironment executionEnvironment, - string continuationToken, - Func>> tryCreateSourceAsync, + RequestContinuationToken continuationToken, + Func>> tryCreateSourceAsync, IReadOnlyDictionary groupByAliasToAggregateType, IReadOnlyList orderedAliases, bool hasSelectValue) @@ -274,24 +275,22 @@ public static TryCatch TryCreateFromContinuationToken( IReadOnlyDictionary groupByAliasToAggregateType, IReadOnlyList orderedAliases, bool hasSelectValue, - string groupingTableContinuationToken) + CosmosElement continuationToken) { GroupingTable groupingTable = new GroupingTable( groupByAliasToAggregateType, orderedAliases, hasSelectValue); - if (groupingTableContinuationToken != null) + if (continuationToken != null) { - if (!CosmosElement.TryParse( - groupingTableContinuationToken, - out CosmosObject parsedGroupingTableContinuations)) + if (!(continuationToken is CosmosObject groupingTableContinuationToken)) { return TryCatch.FromException( new MalformedContinuationTokenException($"Invalid GroupingTableContinuationToken")); } - foreach (KeyValuePair kvp in parsedGroupingTableContinuations) + foreach (KeyValuePair kvp in groupingTableContinuationToken) { string key = kvp.Key; CosmosElement value = kvp.Value; @@ -313,7 +312,7 @@ public static TryCatch TryCreateFromContinuationToken( groupByAliasToAggregateType, orderedAliases, hasSelectValue, - singleGroupAggregatorContinuationToken.Value); + RequestContinuationToken.Create(singleGroupAggregatorContinuationToken.Value)); if (tryCreateSingleGroupAggregator.Succeeded) { diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/TakeDocumentQueryExecutionComponent.Client.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/TakeDocumentQueryExecutionComponent.Client.cs new file mode 100644 index 0000000000..3c0d7be0ff --- /dev/null +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/TakeDocumentQueryExecutionComponent.Client.cs @@ -0,0 +1,118 @@ +//------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +//------------------------------------------------------------ +namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.SkipTake +{ + using System; + using System.Collections.Generic; + using System.Globalization; + using System.Linq; + using System.Threading; + using System.Threading.Tasks; + using Microsoft.Azure.Cosmos.Core.Trace; + 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 Newtonsoft.Json; + + internal abstract partial class TakeDocumentQueryExecutionComponent : DocumentQueryExecutionComponentBase + { + private sealed class ClientTakeDocumentQueryExecutionComponent : TakeDocumentQueryExecutionComponent + { + private ClientTakeDocumentQueryExecutionComponent(IDocumentQueryExecutionComponent source, int takeCount, TakeEnum takeEnum) + : base(source, takeCount, takeEnum) + { + // Work is done in the base class. + } + + public static async Task> TryCreateLimitDocumentQueryExecutionComponentAsync( + int limitCount, + RequestContinuationToken requestContinuationToken, + Func>> tryCreateSourceAsync) + { + if (tryCreateSourceAsync == null) + { + throw new ArgumentNullException(nameof(tryCreateSourceAsync)); + } + + LimitContinuationToken limitContinuationToken; + if (requestContinuationToken.IsNull) + { + if (!LimitContinuationToken.TryParse(requestContinuationToken, out limitContinuationToken)) + { + return TryCatch.FromException( + new MalformedContinuationTokenException($"Malformed {nameof(LimitContinuationToken)}: {requestContinuationToken}.")); + } + } + else + { + limitContinuationToken = new LimitContinuationToken(limitCount, null); + } + + if (limitContinuationToken.Limit > limitCount) + { + return TryCatch.FromException( + new MalformedContinuationTokenException($"{nameof(LimitContinuationToken.Limit)} in {nameof(LimitContinuationToken)}: {continuationToken}: {limitContinuationToken.Limit} can not be greater than the limit count in the query: {limitCount}.")); + } + + if (limitCount < 0) + { + return TryCatch.FromException( + new MalformedContinuationTokenException($"{nameof(limitCount)}: {limitCount} must be a non negative number.")); + } + + return (await tryCreateSourceAsync(limitContinuationToken.SourceToken)) + .Try((source) => new TakeDocumentQueryExecutionComponent( + source, + limitContinuationToken.Limit, + TakeEnum.Limit)); + } + + public static async Task> TryCreateTopDocumentQueryExecutionComponentAsync( + int topCount, + string continuationToken, + Func>> tryCreateSourceAsync) + { + if (tryCreateSourceAsync == null) + { + throw new ArgumentNullException(nameof(tryCreateSourceAsync)); + } + + TopContinuationToken topContinuationToken; + if (continuationToken != null) + { + if (!TopContinuationToken.TryParse(continuationToken, out topContinuationToken)) + { + return TryCatch.FromException( + new MalformedContinuationTokenException($"Malformed {nameof(LimitContinuationToken)}: {continuationToken}.")); + } + } + else + { + topContinuationToken = new TopContinuationToken(topCount, null); + } + + if (topContinuationToken.Top > topCount) + { + return TryCatch.FromException( + new MalformedContinuationTokenException($"{nameof(TopContinuationToken.Top)} in {nameof(TopContinuationToken)}: {continuationToken}: {topContinuationToken.Top} can not be greater than the top count in the query: {topCount}.")); + } + + if (topCount < 0) + { + return TryCatch.FromException( + new MalformedContinuationTokenException($"{nameof(topCount)}: {topCount} must be a non negative number.")); + } + + return (await tryCreateSourceAsync(topContinuationToken.SourceToken)) + .Try((source) => new TakeDocumentQueryExecutionComponent( + source, + topContinuationToken.Top, + TakeEnum.Top)); + } + } + } +} \ 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 index 1f6a56044a..a45f83a970 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/TakeDocumentQueryExecutionComponent.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/TakeDocumentQueryExecutionComponent.cs @@ -12,12 +12,14 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.SkipTake using Microsoft.Azure.Cosmos.Core.Trace; 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.ExecutionContext; using Microsoft.Azure.Cosmos.Query.Core.Monads; using Microsoft.Azure.Cosmos.Query.Core.QueryClient; using Newtonsoft.Json; - internal sealed class TakeDocumentQueryExecutionComponent : DocumentQueryExecutionComponentBase + internal abstract partial class TakeDocumentQueryExecutionComponent : DocumentQueryExecutionComponentBase { private const string SourceTokenName = "SourceToken"; private const string TakeCountName = "TakeCount"; @@ -25,7 +27,7 @@ internal sealed class TakeDocumentQueryExecutionComponent : DocumentQueryExecuti private readonly TakeEnum takeEnum; private int takeCount; - private TakeDocumentQueryExecutionComponent( + protected TakeDocumentQueryExecutionComponent( IDocumentQueryExecutionComponent source, int takeCount, TakeEnum takeEnum) @@ -36,9 +38,10 @@ private TakeDocumentQueryExecutionComponent( } public static async Task> TryCreateLimitDocumentQueryExecutionComponentAsync( + ExecutionEnvironment executionEnvironment, int limitCount, - string continuationToken, - Func>> tryCreateSourceAsync) + RequestContinuationToken continuationToken, + Func>> tryCreateSourceAsync) { if (tryCreateSourceAsync == null) { @@ -46,7 +49,7 @@ public static async Task> TryCreateLi } LimitContinuationToken limitContinuationToken; - if (continuationToken != null) + if (continuationToken.IsNull) { if (!LimitContinuationToken.TryParse(continuationToken, out limitContinuationToken)) { From b50ae6dbf3eaa7194440135aa9212fc4accbcff7 Mon Sep 17 00:00:00 2001 From: Brandon Chong Date: Mon, 20 Jan 2020 16:00:36 -0800 Subject: [PATCH 05/28] drafted out code --- .../src/EncodingExtensions.cs | 75 ++++ .../src/FeedIteratorCore.cs | 13 +- .../src/FeedIteratorInternal.cs | 4 +- ...eDocumentQueryExecutionComponent.Client.cs | 7 - ...DocumentQueryExecutionComponent.Compute.cs | 20 - ...tDocumentQueryExecutionComponent.Client.cs | 17 +- ...DocumentQueryExecutionComponent.Compute.cs | 44 +-- .../DistinctMap.OrderedDistinctMap.cs | 43 ++- .../DistinctMap.UnorderedDistinctMap.cs | 37 +- .../Distinct/DistinctMap.cs | 3 +- .../DocumentQueryExecutionComponentBase.cs | 8 +- ...yDocumentQueryExecutionComponent.Client.cs | 8 +- ...DocumentQueryExecutionComponent.Compute.cs | 20 - ...pDocumentQueryExecutionComponent.Client.cs | 5 - ...DocumentQueryExecutionComponent.Compute.cs | 5 - ...eDocumentQueryExecutionComponent.Client.cs | 239 ++++++++++-- ...DocumentQueryExecutionComponent.Compute.cs | 144 ++++++++ .../TakeDocumentQueryExecutionComponent.cs | 344 +++--------------- .../CosmosQueryExecutionContextFactory.cs | 20 +- .../PipelinedDocumentQueryExecutionContext.cs | 2 + .../src/Query/Core/UInt128.cs | 28 +- .../Query/v3Query/FeedIteratorInlineCore.cs | 16 +- .../src/Query/v3Query/QueryIterator.cs | 33 +- ...geFeedPartitionKeyResultSetIteratorCore.cs | 6 + .../ChangeFeedResultSetIteratorCore.cs | 7 + .../CrossPartitionQueryTests.cs | 9 +- .../CosmosQueryUnitTests.cs | 2 + .../Query/QueryPipelineMockTests.cs | 16 +- 28 files changed, 684 insertions(+), 491 deletions(-) create mode 100644 Microsoft.Azure.Cosmos/src/EncodingExtensions.cs create mode 100644 Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/TakeDocumentQueryExecutionComponent.Compute.cs diff --git a/Microsoft.Azure.Cosmos/src/EncodingExtensions.cs b/Microsoft.Azure.Cosmos/src/EncodingExtensions.cs new file mode 100644 index 0000000000..c53cdde99c --- /dev/null +++ b/Microsoft.Azure.Cosmos/src/EncodingExtensions.cs @@ -0,0 +1,75 @@ +// ------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +// ------------------------------------------------------------ + +namespace Microsoft.Azure.Cosmos +{ + using System; + using System.Text; + + internal static class EncodingExtensions + { + public static int GetByteCount(this Encoding encoding, string chars) + { + return encoding.GetByteCount(chars.AsSpan()); + } + + public static unsafe int GetByteCount(this Encoding encoding, ReadOnlySpan chars) + { + if (chars.IsEmpty) + { + return 0; + } + + fixed (char* charPointer = &chars.GetPinnableReference()) + { + return encoding.GetByteCount(charPointer, chars.Length); + } + } + + public static int GetBytes(this Encoding encoding, string src, Span dst) + { + return encoding.GetBytes(src.AsSpan(), dst); + } + + public static unsafe int GetBytes(this Encoding encoding, ReadOnlySpan src, Span dst) + { + if (src.IsEmpty) + { + return 0; + } + + fixed (char* chars = &src.GetPinnableReference()) + fixed (byte* bytes = &dst.GetPinnableReference()) + { + return encoding.GetBytes(chars, src.Length, bytes, dst.Length); + } + } + + public static unsafe int GetCharCount(this Encoding encoding, ReadOnlySpan src) + { + if (src.IsEmpty) + { + return 0; + } + + fixed (byte* bytes = &src.GetPinnableReference()) + { + return encoding.GetCharCount(bytes, src.Length); + } + } + + public static unsafe string GetString(this Encoding encoding, ReadOnlySpan src) + { + if (src.IsEmpty) + { + return string.Empty; + } + + fixed (byte* bytes = &src.GetPinnableReference()) + { + return encoding.GetString(bytes, src.Length); + } + } + } +} diff --git a/Microsoft.Azure.Cosmos/src/FeedIteratorCore.cs b/Microsoft.Azure.Cosmos/src/FeedIteratorCore.cs index d8552a5c9e..5e6fc65832 100644 --- a/Microsoft.Azure.Cosmos/src/FeedIteratorCore.cs +++ b/Microsoft.Azure.Cosmos/src/FeedIteratorCore.cs @@ -7,9 +7,10 @@ namespace Microsoft.Azure.Cosmos using System; using System.IO; using System.Net; + using System.Text; using System.Threading; using System.Threading.Tasks; - using Microsoft.Azure.Cosmos.Query; + using Microsoft.Azure.Cosmos.Json; using Microsoft.Azure.Cosmos.Query.Core; using Microsoft.Azure.Documents; using static Microsoft.Azure.Documents.RuntimeConstants; @@ -114,6 +115,11 @@ internal static bool GetHasMoreResults(string continuationToken, HttpStatusCode // the enumeration is done for now. return continuationToken != null && statusCode != HttpStatusCode.NotModified; } + + public override void SerializeState(IJsonWriter jsonWriter) + { + jsonWriter.WriteJsonFragment(Encoding.UTF8.GetBytes(this.continuationToken)); + } } /// @@ -152,5 +158,10 @@ public override bool TryGetContinuationToken(out string continuationToken) { return this.feedIterator.TryGetContinuationToken(out continuationToken); } + + public override void SerializeState(IJsonWriter jsonWriter) + { + this.feedIterator.SerializeState(jsonWriter); + } } } \ No newline at end of file diff --git a/Microsoft.Azure.Cosmos/src/FeedIteratorInternal.cs b/Microsoft.Azure.Cosmos/src/FeedIteratorInternal.cs index 9a318f3f2a..ac3dfedad5 100644 --- a/Microsoft.Azure.Cosmos/src/FeedIteratorInternal.cs +++ b/Microsoft.Azure.Cosmos/src/FeedIteratorInternal.cs @@ -4,8 +4,7 @@ namespace Microsoft.Azure.Cosmos { - using System.Threading; - using System.Threading.Tasks; + using Microsoft.Azure.Cosmos.Json; /// /// Internal feed iterator API for casting and mocking purposes. @@ -13,5 +12,6 @@ namespace Microsoft.Azure.Cosmos internal abstract class FeedIteratorInternal : FeedIterator { public abstract bool TryGetContinuationToken(out string continuationToken); + public abstract void SerializeState(IJsonWriter jsonWriter); } } \ No newline at end of file diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/AggregateDocumentQueryExecutionComponent.Client.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/AggregateDocumentQueryExecutionComponent.Client.cs index a124bf0ce4..dede7b8010 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/AggregateDocumentQueryExecutionComponent.Client.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/AggregateDocumentQueryExecutionComponent.Client.cs @@ -115,13 +115,6 @@ public override async Task DrainAsync( diagnostics: diagnosticsPages, responseLengthBytes: responseLengthBytes); } - - public override bool TryGetContinuationToken(out string state) - { - // Since we block until we get the final result the continuation token is always null. - state = null; - return true; - } } } } diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/AggregateDocumentQueryExecutionComponent.Compute.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/AggregateDocumentQueryExecutionComponent.Compute.cs index 42c6cdfda1..9a4a9bbd3f 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/AggregateDocumentQueryExecutionComponent.Compute.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/AggregateDocumentQueryExecutionComponent.Compute.cs @@ -165,26 +165,6 @@ private QueryResponseCore GetEmptyPage(QueryResponseCore sourceResponse) return response; } - public override bool TryGetContinuationToken(out string state) - { - if (this.IsDone) - { - state = null; - return true; - } - - if (!this.Source.TryGetContinuationToken(out string sourceState)) - { - state = null; - return false; - } - - IJsonWriter jsonWriter = JsonWriter.Create(JsonSerializationFormat.Text); - this.SerializeState(jsonWriter); - state = Utf8StringHelpers.ToString(jsonWriter.GetResult()); - return true; - } - public override void SerializeState(IJsonWriter jsonWriter) { if (jsonWriter == null) diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctDocumentQueryExecutionComponent.Client.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctDocumentQueryExecutionComponent.Client.cs index d8d44a1c87..564a46c7a0 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctDocumentQueryExecutionComponent.Client.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctDocumentQueryExecutionComponent.Client.cs @@ -46,11 +46,21 @@ public static async Task> TryCreateAs Func>> tryCreateSourceAsync, DistinctQueryType distinctQueryType) { + if (requestContinuation == null) + { + throw new ArgumentNullException(nameof(requestContinuation)); + } + if (!(requestContinuation is StringRequestContinuationToken stringRequestContinuationToken)) { throw new ArgumentException($"Unknown {nameof(RequestContinuationToken)} type: {requestContinuation.GetType()}"); } + if (tryCreateSourceAsync == null) + { + throw new ArgumentNullException(nameof(tryCreateSourceAsync)); + } + DistinctContinuationToken distinctContinuationToken; if (requestContinuation != null) { @@ -67,7 +77,7 @@ public static async Task> TryCreateAs TryCatch tryCreateDistinctMap = DistinctMap.TryCreate( distinctQueryType, - distinctContinuationToken.DistinctMapToken); + RequestContinuationToken.Create(distinctContinuationToken.DistinctMapToken)); if (!tryCreateDistinctMap.Succeeded) { return TryCatch.FromException(tryCreateDistinctMap.Exception); @@ -145,11 +155,6 @@ public override void SerializeState(IJsonWriter jsonWriter) throw new NotImplementedException(); } - public override bool TryGetContinuationToken(out string state) - { - throw new NotImplementedException(); - } - /// /// Continuation token for distinct queries. /// diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctDocumentQueryExecutionComponent.Compute.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctDocumentQueryExecutionComponent.Compute.cs index f63d214e60..209d7c9450 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctDocumentQueryExecutionComponent.Compute.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctDocumentQueryExecutionComponent.Compute.cs @@ -37,13 +37,28 @@ public static async Task> TryCreateAs Func>> tryCreateSourceAsync, DistinctQueryType distinctQueryType) { + if (requestContinuation == null) + { + throw new ArgumentNullException(nameof(requestContinuation)); + } + + if (!(requestContinuation is CosmosElementRequestContinuationToken cosmosElementRequestContinuationToken)) + { + throw new ArgumentException($"Unknown {nameof(RequestContinuationToken)} type: {requestContinuation.GetType()}"); + } + + if (tryCreateSourceAsync == null) + { + throw new ArgumentNullException(nameof(tryCreateSourceAsync)); + } + DistinctContinuationToken distinctContinuationToken; if (requestContinuation != null) { - if (!DistinctContinuationToken.TryParse(requestContinuation, out distinctContinuationToken)) + if (!DistinctContinuationToken.TryParse(cosmosElementRequestContinuationToken.Value, out distinctContinuationToken)) { return TryCatch.FromException( - new MalformedContinuationTokenException($"Invalid {nameof(DistinctContinuationToken)}: {requestContinuation}")); + new MalformedContinuationTokenException($"Invalid {nameof(DistinctContinuationToken)}: {cosmosElementRequestContinuationToken.Value}")); } } else @@ -53,13 +68,14 @@ public static async Task> TryCreateAs TryCatch tryCreateDistinctMap = DistinctMap.TryCreate( distinctQueryType, - distinctContinuationToken.DistinctMapToken); + RequestContinuationToken.Create(distinctContinuationToken.DistinctMapToken)); if (!tryCreateDistinctMap.Succeeded) { return TryCatch.FromException(tryCreateDistinctMap.Exception); } - TryCatch tryCreateSource = await tryCreateSourceAsync(distinctContinuationToken.SourceToken); + TryCatch tryCreateSource = await tryCreateSourceAsync( + RequestContinuationToken.Create(distinctContinuationToken.SourceToken)); if (!tryCreateSource.Succeeded) { return TryCatch.FromException(tryCreateSource.Exception); @@ -106,26 +122,6 @@ public override async Task DrainAsync(int maxElements, Cancel responseLengthBytes: sourceResponse.ResponseLengthBytes); } - public override bool TryGetContinuationToken(out string continuationToken) - { - if (this.IsDone) - { - continuationToken = null; - return true; - } - - if (!this.Source.TryGetContinuationToken(out string sourceContinuationToken)) - { - continuationToken = default; - return false; - } - - IJsonWriter jsonWriter = JsonWriter.Create(JsonSerializationFormat.Text); - this.SerializeState(jsonWriter); - continuationToken = Utf8StringHelpers.ToString(jsonWriter.GetResult()); - return true; - } - public override void SerializeState(IJsonWriter jsonWriter) { if (jsonWriter == null) diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctMap.OrderedDistinctMap.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctMap.OrderedDistinctMap.cs index 832ef85f34..264c21067b 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctMap.OrderedDistinctMap.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctMap.OrderedDistinctMap.cs @@ -7,6 +7,7 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.Distinct 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.Exceptions; using Microsoft.Azure.Cosmos.Query.Core.Monads; @@ -83,17 +84,47 @@ public override void SerializeState(IJsonWriter jsonWriter) jsonWriter.WriteBinaryValue(UInt128.ToByteArray(this.lastHash)); } - public static TryCatch TryCreate(string continuationToken) + public static TryCatch TryCreate(RequestContinuationToken requestContinuationToken) { + if (requestContinuationToken == null) + { + throw new ArgumentNullException(nameof(requestContinuationToken)); + } + UInt128 lastHash; - if (continuationToken != null) + if (!requestContinuationToken.IsNull) { - if (!UInt128.TryParse(continuationToken, out lastHash)) + switch (requestContinuationToken) { - return TryCatch.FromException( - new MalformedContinuationTokenException( - $"Malformed {nameof(OrderedDistinctMap)} continuation token: {continuationToken}.")); + case StringRequestContinuationToken stringRequestContinuationToken: + if (!UInt128.TryParse(stringRequestContinuationToken.Value, out lastHash)) + { + return TryCatch.FromException( + new MalformedContinuationTokenException( + $"Malformed {nameof(OrderedDistinctMap)} continuation token: {stringRequestContinuationToken.Value}.")); + } + break; + + case CosmosElementRequestContinuationToken cosmosElementRequestContinuationToken: + if (!(cosmosElementRequestContinuationToken.Value is CosmosBinary cosmosBinary)) + { + return TryCatch.FromException( + new MalformedContinuationTokenException( + $"Malformed {nameof(OrderedDistinctMap)} continuation token: {cosmosElementRequestContinuationToken.Value}.")); + } + + if (!UInt128.TryParse(cosmosBinary.Value.Span, out lastHash)) + { + return TryCatch.FromException( + new MalformedContinuationTokenException( + $"Malformed {nameof(OrderedDistinctMap)} continuation token: {cosmosElementRequestContinuationToken.Value}.")); + } + break; + + default: + throw new ArgumentOutOfRangeException($"Unknown {nameof(RequestContinuationToken)} type. {requestContinuationToken.GetType()}."); } + } else { diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctMap.UnorderedDistinctMap.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctMap.UnorderedDistinctMap.cs index 9d0cb0c3e4..0bcc1b3266 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctMap.UnorderedDistinctMap.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctMap.UnorderedDistinctMap.cs @@ -10,6 +10,7 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.Distinct 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.Exceptions; using Microsoft.Azure.Cosmos.Query.Core.Monads; @@ -97,11 +98,6 @@ private sealed class UnorderdDistinctMap : DistinctMap private const string ObjectName = "Object"; private const string SimpleValuesName = "SimpleValues"; - /// - /// Buffer that gets reused to convert a .net string (utf-16) to a (utf-8) byte array. - /// - private readonly byte[] utf8Buffer; - /// /// HashSet for all numbers seen. /// This takes less space than a 24 byte hash and has full fidelity. @@ -159,7 +155,6 @@ private UnorderdDistinctMap( HashSet objects, SimpleValues simpleValues) { - this.utf8Buffer = new byte[UnorderdDistinctMap.UInt128Length]; this.numbers = numbers ?? throw new ArgumentNullException(nameof(numbers)); this.stringsLength4 = stringsLength4 ?? throw new ArgumentNullException(nameof(stringsLength4)); this.stringsLength8 = stringsLength8 ?? throw new ArgumentNullException(nameof(stringsLength8)); @@ -351,27 +346,25 @@ private bool AddStringValue(string value) // If you can fit the string with full fidelity in 16 bytes, then you might as well just hash the string itself. if (utf8Length <= UnorderdDistinctMap.UInt128Length) { - // Zero out the array since you want all trailing bytes to be 0 for the conversions that happen next. - Array.Clear(this.utf8Buffer, 0, this.utf8Buffer.Length); - Encoding.UTF8.GetBytes(value, 0, value.Length, this.utf8Buffer, 0); - + Span utf8Buffer = stackalloc byte[UInt128Length]; + Encoding.UTF8.GetBytes(value, utf8Buffer); if (utf8Length == 0) { added = this.AddSimpleValue(SimpleValues.EmptyString); } else if (utf8Length <= UnorderdDistinctMap.UIntLength) { - uint uintValue = BitConverter.ToUInt32(this.utf8Buffer, 0); + uint uintValue = MemoryMarshal.Read(utf8Buffer); added = this.stringsLength4.Add(uintValue); } else if (utf8Length <= UnorderdDistinctMap.ULongLength) { - ulong uLongValue = BitConverter.ToUInt64(this.utf8Buffer, 0); + ulong uLongValue = MemoryMarshal.Read(utf8Buffer); added = this.stringsLength8.Add(uLongValue); } else { - UInt128 uInt128Value = UInt128.FromByteArray(this.utf8Buffer); + UInt128 uInt128Value = UInt128.FromByteArray(utf8Buffer); added = this.stringsLength16.Add(uInt128Value); } } @@ -380,7 +373,6 @@ private bool AddStringValue(string value) // Else the string is too large and we will just store the hash. UInt128 uint128Value = DistinctHash.GetHash(CosmosString.Create(value)); added = this.stringsLength16Plus.Add(uint128Value); - } return added; @@ -408,8 +400,13 @@ private bool AddObjectValue(CosmosObject cosmosObject) return this.objects.Add(hash); } - public static TryCatch TryCreate(string continuationToken) + public static TryCatch TryCreate(RequestContinuationToken continuationToken) { + if (continuationToken == null) + { + throw new ArgumentNullException(nameof(continuationToken)); + } + HashSet numbers = new HashSet(); HashSet stringsLength4 = new HashSet(); HashSet stringsLength8 = new HashSet(); @@ -419,10 +416,14 @@ public static TryCatch TryCreate(string continuationToken) HashSet objects = new HashSet(); SimpleValues simpleValues = SimpleValues.None; - if (continuationToken != null) + if (!continuationToken.IsNull) { - byte[] binaryBuffer = Convert.FromBase64String(continuationToken); - CosmosElement cosmosElement = CosmosElement.CreateFromBuffer(binaryBuffer); + if (!(continuationToken is CosmosElementRequestContinuationToken cosmosElementRequestContinuation)) + { + throw new ArgumentException($"Unknown {nameof(RequestContinuationToken)} type: {continuationToken.GetType()}"); + } + + CosmosElement cosmosElement = cosmosElementRequestContinuation.Value; if (!(cosmosElement is CosmosObject hashDictionary)) { return TryCatch.FromException( diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctMap.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctMap.cs index fe016415ab..19e7e53dff 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctMap.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctMap.cs @@ -7,6 +7,7 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.Distinct 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; /// @@ -28,7 +29,7 @@ internal abstract partial class DistinctMap /// The appropriate IDistinctMap. public static TryCatch TryCreate( DistinctQueryType distinctQueryType, - string distinctMapContinuationToken) + RequestContinuationToken distinctMapContinuationToken) { TryCatch tryCreateDistinctMap; switch (distinctQueryType) diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/DocumentQueryExecutionComponentBase.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/DocumentQueryExecutionComponentBase.cs index 64360225b4..441d113be7 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/DocumentQueryExecutionComponentBase.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/DocumentQueryExecutionComponentBase.cs @@ -67,7 +67,13 @@ public void Stop() this.Source.Stop(); } - public abstract bool TryGetContinuationToken(out string state); + public bool TryGetContinuationToken(out string continuationToken) + { + IJsonWriter jsonWriter = JsonWriter.Create(JsonSerializationFormat.Text); + this.SerializeState(jsonWriter); + continuationToken = Utf8StringHelpers.ToString(jsonWriter.GetResult()); + return false; + } public abstract void SerializeState(IJsonWriter jsonWriter); } 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 e70321e712..3ce683257f 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 @@ -41,7 +41,7 @@ public static async Task> TryCreateAs groupByAliasToAggregateType, orderedAliases, hasSelectValue, - groupingTableContinuationToken: null); + continuationToken: null); if (!tryCreateGroupingTable.Succeeded) { @@ -103,12 +103,6 @@ public override async Task DrainAsync( return response; } - public override bool TryGetContinuationToken(out string state) - { - state = default; - return false; - } - public override void SerializeState(IJsonWriter jsonWriter) { throw new NotSupportedException(ContinuationTokenNotSupportedWithGroupBy); 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 1836b1e371..4da5259b0b 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 @@ -175,26 +175,6 @@ public override void SerializeState(IJsonWriter jsonWriter) } } - public override bool TryGetContinuationToken(out string continuationToken) - { - if (this.IsDone) - { - continuationToken = null; - return true; - } - - if (!this.Source.TryGetContinuationToken(out string sourceContinuationToken)) - { - continuationToken = default; - return false; - } - - IJsonWriter jsonWriter = JsonWriter.Create(JsonSerializationFormat.Text); - this.SerializeState(jsonWriter); - continuationToken = Utf8StringHelpers.ToString(jsonWriter.GetResult()); - return true; - } - private readonly struct GroupByContinuationToken { public GroupByContinuationToken( diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/SkipDocumentQueryExecutionComponent.Client.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/SkipDocumentQueryExecutionComponent.Client.cs index a970e76519..8db7ab7d92 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/SkipDocumentQueryExecutionComponent.Client.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/SkipDocumentQueryExecutionComponent.Client.cs @@ -68,11 +68,6 @@ public override void SerializeState(IJsonWriter jsonWriter) throw new NotImplementedException(); } - public override bool TryGetContinuationToken(out string state) - { - throw new NotImplementedException(); - } - /// /// A OffsetContinuationToken is a composition of a source continuation token and how many items to skip from that source. /// diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/SkipDocumentQueryExecutionComponent.Compute.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/SkipDocumentQueryExecutionComponent.Compute.cs index f1d662167d..1f611e2376 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/SkipDocumentQueryExecutionComponent.Compute.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/SkipDocumentQueryExecutionComponent.Compute.cs @@ -69,11 +69,6 @@ public static async Task> TryCreateCo offsetContinuationToken.Offset)); } - public override bool TryGetContinuationToken(out string state) - { - throw new NotImplementedException(); - } - public override void SerializeState(IJsonWriter jsonWriter) { if (jsonWriter == null) diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/TakeDocumentQueryExecutionComponent.Client.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/TakeDocumentQueryExecutionComponent.Client.cs index 3c0d7be0ff..27ecbe16c5 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/TakeDocumentQueryExecutionComponent.Client.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/TakeDocumentQueryExecutionComponent.Client.cs @@ -4,28 +4,23 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.SkipTake { using System; - using System.Collections.Generic; - using System.Globalization; - using System.Linq; - using System.Threading; using System.Threading.Tasks; - using Microsoft.Azure.Cosmos.Core.Trace; - 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 Newtonsoft.Json; internal abstract partial class TakeDocumentQueryExecutionComponent : DocumentQueryExecutionComponentBase { private sealed class ClientTakeDocumentQueryExecutionComponent : TakeDocumentQueryExecutionComponent { + private readonly TakeEnum takeEnum; + private ClientTakeDocumentQueryExecutionComponent(IDocumentQueryExecutionComponent source, int takeCount, TakeEnum takeEnum) - : base(source, takeCount, takeEnum) + : base(source, takeCount) { - // Work is done in the base class. + this.takeEnum = takeEnum; } public static async Task> TryCreateLimitDocumentQueryExecutionComponentAsync( @@ -33,15 +28,30 @@ public static async Task> TryCreateLi RequestContinuationToken requestContinuationToken, Func>> tryCreateSourceAsync) { + if (limitCount < 0) + { + throw new ArgumentException($"{nameof(limitCount)}: {limitCount} must be a non negative number."); + } + if (tryCreateSourceAsync == null) { throw new ArgumentNullException(nameof(tryCreateSourceAsync)); } + if (requestContinuationToken == null) + { + throw new ArgumentNullException(nameof(requestContinuationToken)); + } + + if (!(requestContinuationToken is StringRequestContinuationToken stringRequestContinuationToken)) + { + throw new ArgumentException($"Unknown {nameof(StringRequestContinuationToken)} type: {requestContinuationToken.GetType()}."); + } + LimitContinuationToken limitContinuationToken; - if (requestContinuationToken.IsNull) + if (!requestContinuationToken.IsNull) { - if (!LimitContinuationToken.TryParse(requestContinuationToken, out limitContinuationToken)) + if (!LimitContinuationToken.TryParse(stringRequestContinuationToken.Value, out limitContinuationToken)) { return TryCatch.FromException( new MalformedContinuationTokenException($"Malformed {nameof(LimitContinuationToken)}: {requestContinuationToken}.")); @@ -55,17 +65,11 @@ public static async Task> TryCreateLi if (limitContinuationToken.Limit > limitCount) { return TryCatch.FromException( - new MalformedContinuationTokenException($"{nameof(LimitContinuationToken.Limit)} in {nameof(LimitContinuationToken)}: {continuationToken}: {limitContinuationToken.Limit} can not be greater than the limit count in the query: {limitCount}.")); - } - - if (limitCount < 0) - { - return TryCatch.FromException( - new MalformedContinuationTokenException($"{nameof(limitCount)}: {limitCount} must be a non negative number.")); + new MalformedContinuationTokenException($"{nameof(LimitContinuationToken.Limit)} in {nameof(LimitContinuationToken)}: {stringRequestContinuationToken.Value}: {limitContinuationToken.Limit} can not be greater than the limit count in the query: {limitCount}.")); } - return (await tryCreateSourceAsync(limitContinuationToken.SourceToken)) - .Try((source) => new TakeDocumentQueryExecutionComponent( + return (await tryCreateSourceAsync(RequestContinuationToken.Create(limitContinuationToken.SourceToken))) + .Try((source) => new ClientTakeDocumentQueryExecutionComponent( source, limitContinuationToken.Limit, TakeEnum.Limit)); @@ -73,21 +77,36 @@ public static async Task> TryCreateLi public static async Task> TryCreateTopDocumentQueryExecutionComponentAsync( int topCount, - string continuationToken, - Func>> tryCreateSourceAsync) + RequestContinuationToken requestContinuationToken, + Func>> tryCreateSourceAsync) { + if (topCount < 0) + { + throw new ArgumentException($"{nameof(topCount)}: {topCount} must be a non negative number."); + } + if (tryCreateSourceAsync == null) { throw new ArgumentNullException(nameof(tryCreateSourceAsync)); } + if (requestContinuationToken == null) + { + throw new ArgumentNullException(nameof(requestContinuationToken)); + } + + if (!(requestContinuationToken is StringRequestContinuationToken stringRequestContinuationToken)) + { + throw new ArgumentException($"Unknown {nameof(StringRequestContinuationToken)} type: {requestContinuationToken.GetType()}."); + } + TopContinuationToken topContinuationToken; - if (continuationToken != null) + if (!requestContinuationToken.IsNull) { - if (!TopContinuationToken.TryParse(continuationToken, out topContinuationToken)) + if (!TopContinuationToken.TryParse(stringRequestContinuationToken.Value, out topContinuationToken)) { return TryCatch.FromException( - new MalformedContinuationTokenException($"Malformed {nameof(LimitContinuationToken)}: {continuationToken}.")); + new MalformedContinuationTokenException($"Malformed {nameof(LimitContinuationToken)}: {stringRequestContinuationToken.Value}.")); } } else @@ -98,21 +117,173 @@ public static async Task> TryCreateTo if (topContinuationToken.Top > topCount) { return TryCatch.FromException( - new MalformedContinuationTokenException($"{nameof(TopContinuationToken.Top)} in {nameof(TopContinuationToken)}: {continuationToken}: {topContinuationToken.Top} can not be greater than the top count in the query: {topCount}.")); - } - - if (topCount < 0) - { - return TryCatch.FromException( - new MalformedContinuationTokenException($"{nameof(topCount)}: {topCount} must be a non negative number.")); + new MalformedContinuationTokenException($"{nameof(TopContinuationToken.Top)} in {nameof(TopContinuationToken)}: {stringRequestContinuationToken.Value}: {topContinuationToken.Top} can not be greater than the top count in the query: {topCount}.")); } - return (await tryCreateSourceAsync(topContinuationToken.SourceToken)) - .Try((source) => new TakeDocumentQueryExecutionComponent( + return (await tryCreateSourceAsync(RequestContinuationToken.Create(topContinuationToken.SourceToken))) + .Try((source) => new ClientTakeDocumentQueryExecutionComponent( source, topContinuationToken.Top, TakeEnum.Top)); } + + public override void SerializeState(IJsonWriter jsonWriter) + { + throw new NotImplementedException(); + } + + private enum TakeEnum + { + Limit, + Top + } + + private abstract class TakeContinuationToken + { + } + + /// + /// A LimitContinuationToken is a composition of a source continuation token and how many items we have left to drain from that source. + /// + private sealed class LimitContinuationToken : TakeContinuationToken + { + /// + /// Initializes a new instance of the LimitContinuationToken struct. + /// + /// The limit to the number of document drained for the remainder of the query. + /// The continuation token for the source component of the query. + public LimitContinuationToken(int limit, string sourceToken) + { + if (limit < 0) + { + throw new ArgumentException($"{nameof(limit)} must be a non negative number."); + } + + this.Limit = limit; + this.SourceToken = sourceToken; + } + + /// + /// Gets the limit to the number of document drained for the remainder of the query. + /// + [JsonProperty("limit")] + public int Limit + { + get; + } + + /// + /// Gets the continuation token for the source component of the query. + /// + [JsonProperty("sourceToken")] + public string SourceToken + { + get; + } + + /// + /// Tries to parse out the LimitContinuationToken. + /// + /// The value to parse from. + /// The result of parsing out the token. + /// Whether or not the LimitContinuationToken was successfully parsed out. + public static bool TryParse(string value, out LimitContinuationToken limitContinuationToken) + { + limitContinuationToken = default; + if (string.IsNullOrWhiteSpace(value)) + { + return false; + } + + try + { + limitContinuationToken = JsonConvert.DeserializeObject(value); + return true; + } + catch (JsonException) + { + return false; + } + } + + /// + /// Gets the string version of the continuation token that can be passed in a response header. + /// + /// The string version of the continuation token that can be passed in a response header. + public override string ToString() + { + return JsonConvert.SerializeObject(this); + } + } + + /// + /// A TopContinuationToken is a composition of a source continuation token and how many items we have left to drain from that source. + /// + private sealed class TopContinuationToken : TakeContinuationToken + { + /// + /// Initializes a new instance of the TopContinuationToken struct. + /// + /// The limit to the number of document drained for the remainder of the query. + /// The continuation token for the source component of the query. + public TopContinuationToken(int top, string sourceToken) + { + this.Top = top; + this.SourceToken = sourceToken; + } + + /// + /// Gets the limit to the number of document drained for the remainder of the query. + /// + [JsonProperty("top")] + public int Top + { + get; + } + + /// + /// Gets the continuation token for the source component of the query. + /// + [JsonProperty("sourceToken")] + public string SourceToken + { + get; + } + + /// + /// Tries to parse out the TopContinuationToken. + /// + /// The value to parse from. + /// The result of parsing out the token. + /// Whether or not the TopContinuationToken was successfully parsed out. + public static bool TryParse(string value, out TopContinuationToken topContinuationToken) + { + topContinuationToken = default; + if (string.IsNullOrWhiteSpace(value)) + { + return false; + } + + try + { + topContinuationToken = JsonConvert.DeserializeObject(value); + return true; + } + catch (JsonException) + { + return false; + } + } + + /// + /// Gets the string version of the continuation token that can be passed in a response header. + /// + /// The string version of the continuation token that can be passed in a response header. + public override string ToString() + { + return JsonConvert.SerializeObject(this); + } + } } } } \ 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/ExecutionComponent/SkipTake/TakeDocumentQueryExecutionComponent.Compute.cs new file mode 100644 index 0000000000..6420372f9a --- /dev/null +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/TakeDocumentQueryExecutionComponent.Compute.cs @@ -0,0 +1,144 @@ +//------------------------------------------------------------ +// 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.Json; + using Microsoft.Azure.Cosmos.Query.Core.ContinuationTokens; + using Microsoft.Azure.Cosmos.Query.Core.Exceptions; + using Microsoft.Azure.Cosmos.Query.Core.Monads; + + internal abstract partial class TakeDocumentQueryExecutionComponent : DocumentQueryExecutionComponentBase + { + private sealed class ComputeTakeDocumentQueryExecutionComponent : TakeDocumentQueryExecutionComponent + { + private const string SourceTokenName = "SourceToken"; + private const string TakeCountName = "TakeCount"; + + private ComputeTakeDocumentQueryExecutionComponent(IDocumentQueryExecutionComponent source, int takeCount) + : base(source, takeCount) + { + // Work is done in the base class. + } + + public static async Task> TryCreateAsync( + int takeCount, + RequestContinuationToken requestContinuationToken, + Func>> tryCreateSourceAsync) + { + if (takeCount < 0) + { + throw new ArgumentException($"{nameof(takeCount)}: {takeCount} must be a non negative number."); + } + + if (tryCreateSourceAsync == null) + { + throw new ArgumentNullException(nameof(tryCreateSourceAsync)); + } + + if (requestContinuationToken == null) + { + throw new ArgumentNullException(nameof(requestContinuationToken)); + } + + if (!(requestContinuationToken is CosmosElementRequestContinuationToken cosmosElementRequestContinuationToken)) + { + throw new ArgumentException($"Unknown {nameof(RequestContinuationToken)} type: {requestContinuationToken.GetType()}."); + } + + TakeContinuationToken takeContinuationToken; + if (!requestContinuationToken.IsNull) + { + if (!TakeContinuationToken.TryParse(cosmosElementRequestContinuationToken.Value, out takeContinuationToken)) + { + return TryCatch.FromException( + new MalformedContinuationTokenException($"Malformed {nameof(TakeContinuationToken)}: {cosmosElementRequestContinuationToken.Value}.")); + } + } + else + { + takeContinuationToken = new TakeContinuationToken(takeCount, sourceToken: null); + } + + if (takeContinuationToken.TakeCount > takeCount) + { + return TryCatch.FromException( + new MalformedContinuationTokenException($"{nameof(TakeContinuationToken.TakeCount)} in {nameof(TakeContinuationToken)}: {cosmosElementRequestContinuationToken.Value}: {takeContinuationToken.TakeCount} can not be greater than the limit count in the query: {takeCount}.")); + } + + return (await tryCreateSourceAsync(RequestContinuationToken.Create(takeContinuationToken.SourceToken))) + .Try((source) => new ComputeTakeDocumentQueryExecutionComponent( + source, + takeContinuationToken.TakeCount)); + } + + public override void SerializeState(IJsonWriter jsonWriter) + { + if (jsonWriter == null) + { + throw new ArgumentNullException(nameof(jsonWriter)); + } + + if (!this.IsDone) + { + jsonWriter.WriteObjectStart(); + jsonWriter.WriteFieldName(SourceTokenName); + this.Source.SerializeState(jsonWriter); + jsonWriter.WriteFieldName(TakeCountName); + jsonWriter.WriteInt64Value(this.takeCount); + jsonWriter.WriteObjectEnd(); + } + } + + private readonly struct TakeContinuationToken + { + public TakeContinuationToken(long takeCount, CosmosElement sourceToken) + { + if ((takeCount < 0) || (takeCount > int.MaxValue)) + { + throw new ArgumentException($"{nameof(takeCount)} must be a non negative number."); + } + + this.TakeCount = (int)takeCount; + this.SourceToken = sourceToken; + } + + public int TakeCount { get; } + + public CosmosElement SourceToken { get; } + + public static bool TryParse(CosmosElement value, out TakeContinuationToken takeContinuationToken) + { + if (value == null) + { + throw new ArgumentNullException(nameof(value)); + } + + if (!(value is CosmosObject continuationToken)) + { + takeContinuationToken = default; + return false; + } + + if (!continuationToken.TryGetValue(TakeCountName, out CosmosNumber takeCount)) + { + takeContinuationToken = default; + return false; + } + + if (!continuationToken.TryGetValue(SourceTokenName, out CosmosElement sourceToken)) + { + takeContinuationToken = default; + return false; + } + + takeContinuationToken = new TakeContinuationToken(takeCount.AsInteger().Value, sourceToken); + return true; + } + } + } + } +} \ 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 index a45f83a970..608656b841 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/TakeDocumentQueryExecutionComponent.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/TakeDocumentQueryExecutionComponent.cs @@ -21,107 +21,74 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.SkipTake internal abstract partial class TakeDocumentQueryExecutionComponent : DocumentQueryExecutionComponentBase { - private const string SourceTokenName = "SourceToken"; - private const string TakeCountName = "TakeCount"; - - private readonly TakeEnum takeEnum; private int takeCount; protected TakeDocumentQueryExecutionComponent( IDocumentQueryExecutionComponent source, - int takeCount, - TakeEnum takeEnum) + int takeCount) : base(source) { this.takeCount = takeCount; - this.takeEnum = takeEnum; } - public static async Task> TryCreateLimitDocumentQueryExecutionComponentAsync( + public static Task> TryCreateLimitDocumentQueryExecutionComponentAsync( ExecutionEnvironment executionEnvironment, int limitCount, - RequestContinuationToken continuationToken, + RequestContinuationToken requestContinuationToken, Func>> tryCreateSourceAsync) { - if (tryCreateSourceAsync == null) + Task> tryCreateComponentAsync; + switch (executionEnvironment) { - throw new ArgumentNullException(nameof(tryCreateSourceAsync)); - } + case ExecutionEnvironment.Client: + tryCreateComponentAsync = ClientTakeDocumentQueryExecutionComponent.TryCreateLimitDocumentQueryExecutionComponentAsync( + limitCount, + requestContinuationToken, + tryCreateSourceAsync); + break; - LimitContinuationToken limitContinuationToken; - if (continuationToken.IsNull) - { - if (!LimitContinuationToken.TryParse(continuationToken, out limitContinuationToken)) - { - return TryCatch.FromException( - new MalformedContinuationTokenException($"Malformed {nameof(LimitContinuationToken)}: {continuationToken}.")); - } - } - else - { - limitContinuationToken = new LimitContinuationToken(limitCount, null); - } + case ExecutionEnvironment.Compute: + tryCreateComponentAsync = ComputeTakeDocumentQueryExecutionComponent.TryCreateAsync( + limitCount, + requestContinuationToken, + tryCreateSourceAsync); + break; - if (limitContinuationToken.Limit > limitCount) - { - return TryCatch.FromException( - new MalformedContinuationTokenException($"{nameof(LimitContinuationToken.Limit)} in {nameof(LimitContinuationToken)}: {continuationToken}: {limitContinuationToken.Limit} can not be greater than the limit count in the query: {limitCount}.")); + default: + throw new ArgumentOutOfRangeException($"Unknown {nameof(ExecutionEnvironment)}: {executionEnvironment}."); } - if (limitCount < 0) - { - return TryCatch.FromException( - new MalformedContinuationTokenException($"{nameof(limitCount)}: {limitCount} must be a non negative number.")); - } - - return (await tryCreateSourceAsync(limitContinuationToken.SourceToken)) - .Try((source) => new TakeDocumentQueryExecutionComponent( - source, - limitContinuationToken.Limit, - TakeEnum.Limit)); + return tryCreateComponentAsync; } - public static async Task> TryCreateTopDocumentQueryExecutionComponentAsync( + public static Task> TryCreateTopDocumentQueryExecutionComponentAsync( + ExecutionEnvironment executionEnvironment, int topCount, - string continuationToken, - Func>> tryCreateSourceAsync) + RequestContinuationToken requestContinuationToken, + Func>> tryCreateSourceAsync) { - if (tryCreateSourceAsync == null) - { - throw new ArgumentNullException(nameof(tryCreateSourceAsync)); - } - - TopContinuationToken topContinuationToken; - if (continuationToken != null) - { - if (!TopContinuationToken.TryParse(continuationToken, out topContinuationToken)) - { - return TryCatch.FromException( - new MalformedContinuationTokenException($"Malformed {nameof(LimitContinuationToken)}: {continuationToken}.")); - } - } - else + Task> tryCreateComponentAsync; + switch (executionEnvironment) { - topContinuationToken = new TopContinuationToken(topCount, null); - } + case ExecutionEnvironment.Client: + tryCreateComponentAsync = ClientTakeDocumentQueryExecutionComponent.TryCreateTopDocumentQueryExecutionComponentAsync( + topCount, + requestContinuationToken, + tryCreateSourceAsync); + break; - if (topContinuationToken.Top > topCount) - { - return TryCatch.FromException( - new MalformedContinuationTokenException($"{nameof(TopContinuationToken.Top)} in {nameof(TopContinuationToken)}: {continuationToken}: {topContinuationToken.Top} can not be greater than the top count in the query: {topCount}.")); - } + case ExecutionEnvironment.Compute: + tryCreateComponentAsync = ComputeTakeDocumentQueryExecutionComponent.TryCreateAsync( + topCount, + requestContinuationToken, + tryCreateSourceAsync); + break; - if (topCount < 0) - { - return TryCatch.FromException( - new MalformedContinuationTokenException($"{nameof(topCount)}: {topCount} must be a non negative number.")); + default: + throw new ArgumentOutOfRangeException($"Unknown {nameof(ExecutionEnvironment)}: {executionEnvironment}."); } - return (await tryCreateSourceAsync(topContinuationToken.SourceToken)) - .Try((source) => new TakeDocumentQueryExecutionComponent( - source, - topContinuationToken.Top, - TakeEnum.Top)); + return tryCreateComponentAsync; } public override bool IsDone @@ -147,10 +114,9 @@ public override async Task DrainAsync(int maxElements, Cancel string updatedContinuationToken = null; if (results.DisallowContinuationTokenMessage == null) { - if (!this.TryGetContinuationToken(out updatedContinuationToken)) - { - throw new InvalidOperationException($"Failed to get state for {nameof(TakeDocumentQueryExecutionComponent)}."); - } + IJsonWriter jsonWriter = Json.JsonWriter.Create(JsonSerializationFormat.Text); + this.SerializeState(jsonWriter); + updatedContinuationToken = Utf8StringHelpers.ToString(jsonWriter.GetResult()); } return QueryResponseCore.CreateSuccess( @@ -162,227 +128,5 @@ public override async Task DrainAsync(int maxElements, Cancel diagnostics: results.Diagnostics, responseLengthBytes: results.ResponseLengthBytes); } - - public override bool TryGetContinuationToken(out string state) - { - if (!this.IsDone) - { - if (this.Source.TryGetContinuationToken(out string sourceState)) - { - TakeContinuationToken takeContinuationToken; - switch (this.takeEnum) - { - case TakeEnum.Limit: - takeContinuationToken = new LimitContinuationToken( - this.takeCount, - sourceState); - break; - - case TakeEnum.Top: - takeContinuationToken = new TopContinuationToken( - this.takeCount, - sourceState); - break; - - default: - throw new ArgumentException($"Unknown {nameof(TakeEnum)}: {this.takeEnum}"); - } - - state = takeContinuationToken.ToString(); - return true; - } - else - { - state = default; - return false; - } - } - else - { - state = default; - return true; - } - } - - public override void SerializeState(IJsonWriter jsonWriter) - { - if (jsonWriter == null) - { - throw new ArgumentNullException(nameof(jsonWriter)); - } - - if (this.IsDone) - { - jsonWriter.WriteNullValue(); - } - else - { - jsonWriter.WriteObjectStart(); - jsonWriter.WriteFieldName(SourceTokenName); - this.Source.SerializeState(jsonWriter); - jsonWriter.WriteFieldName(TakeCountName); - jsonWriter.WriteInt64Value(this.takeCount); - jsonWriter.WriteObjectEnd(); - } - } - - private enum TakeEnum - { - Limit, - Top - } - - private abstract class TakeContinuationToken - { - } - - /// - /// A LimitContinuationToken is a composition of a source continuation token and how many items we have left to drain from that source. - /// - private sealed class LimitContinuationToken : TakeContinuationToken - { - /// - /// Initializes a new instance of the LimitContinuationToken struct. - /// - /// The limit to the number of document drained for the remainder of the query. - /// The continuation token for the source component of the query. - public LimitContinuationToken(int limit, string sourceToken) - { - if (limit < 0) - { - throw new ArgumentException($"{nameof(limit)} must be a non negative number."); - } - - this.Limit = limit; - this.SourceToken = sourceToken; - } - - /// - /// Gets the limit to the number of document drained for the remainder of the query. - /// - [JsonProperty("limit")] - public int Limit - { - get; - } - - /// - /// Gets the continuation token for the source component of the query. - /// - [JsonProperty("sourceToken")] - public string SourceToken - { - get; - } - - /// - /// Tries to parse out the LimitContinuationToken. - /// - /// The value to parse from. - /// The result of parsing out the token. - /// Whether or not the LimitContinuationToken was successfully parsed out. - public static bool TryParse(string value, out LimitContinuationToken limitContinuationToken) - { - limitContinuationToken = default; - if (string.IsNullOrWhiteSpace(value)) - { - return false; - } - - try - { - limitContinuationToken = JsonConvert.DeserializeObject(value); - return true; - } - catch (JsonException) - { - return false; - } - } - - /// - /// Gets the string version of the continuation token that can be passed in a response header. - /// - /// The string version of the continuation token that can be passed in a response header. - public override string ToString() - { - return JsonConvert.SerializeObject(this); - } - } - - /// - /// A TopContinuationToken is a composition of a source continuation token and how many items we have left to drain from that source. - /// - private sealed class TopContinuationToken : TakeContinuationToken - { - /// - /// Initializes a new instance of the TopContinuationToken struct. - /// - /// The limit to the number of document drained for the remainder of the query. - /// The continuation token for the source component of the query. - public TopContinuationToken(int top, string sourceToken) - { - this.Top = top; - this.SourceToken = sourceToken; - } - - /// - /// Gets the limit to the number of document drained for the remainder of the query. - /// - [JsonProperty("top")] - public int Top - { - get; - } - - /// - /// Gets the continuation token for the source component of the query. - /// - [JsonProperty("sourceToken")] - public string SourceToken - { - get; - } - - /// - /// Tries to parse out the TopContinuationToken. - /// - /// The value to parse from. - /// The result of parsing out the token. - /// Whether or not the TopContinuationToken was successfully parsed out. - public static bool TryParse(string value, out TopContinuationToken topContinuationToken) - { - topContinuationToken = default; - if (string.IsNullOrWhiteSpace(value)) - { - return false; - } - - try - { - topContinuationToken = JsonConvert.DeserializeObject(value); - return true; - } - catch (JsonException ex) - { - DefaultTrace.TraceWarning(string.Format( - CultureInfo.InvariantCulture, - "{0} Invalid continuation token {1} for Top~Component: {2}", - DateTime.UtcNow.ToString("o", CultureInfo.InvariantCulture), - value, - ex.Message)); - return false; - } - } - - /// - /// Gets the string version of the continuation token that can be passed in a response header. - /// - /// The string version of the continuation token that can be passed in a response header. - public override string ToString() - { - return JsonConvert.SerializeObject(this); - } - } } } \ No newline at end of file diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosQueryExecutionContextFactory.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosQueryExecutionContextFactory.cs index 45a7bd9fbe..9130a80795 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosQueryExecutionContextFactory.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosQueryExecutionContextFactory.cs @@ -9,6 +9,7 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionContext 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.ContinuationTokens; using Microsoft.Azure.Cosmos.Query.Core.Exceptions; @@ -63,12 +64,19 @@ private static async Task> TryCreateCoreCo CancellationToken cancellationToken) { // Try to parse the continuation token. - string continuationToken = inputParameters.InitialUserContinuationToken; + RequestContinuationToken continuationToken = inputParameters.InitialUserContinuationToken; PartitionedQueryExecutionInfo queryPlanFromContinuationToken = inputParameters.PartitionedQueryExecutionInfo; - if (continuationToken != null) + if (!continuationToken.IsNull) { + if (!continuationToken.TryConvertToCosmosElement(out CosmosElement cosmosElement)) + { + return TryCatch.FromException( + new MalformedContinuationTokenException( + $"Malformed {nameof(PipelineContinuationToken)}.")); + } + if (!PipelineContinuationToken.TryParse( - continuationToken, + cosmosElement.ToString(), out PipelineContinuationToken pipelineContinuationToken)) { return TryCatch.FromException( @@ -94,7 +102,7 @@ private static async Task> TryCreateCoreCo $"{nameof(PipelineContinuationToken)}: '{continuationToken}' is no longer supported.")); } - continuationToken = latestVersionPipelineContinuationToken.SourceContinuationToken; + continuationToken = RequestContinuationToken.Create(latestVersionPipelineContinuationToken.SourceContinuationToken); if (latestVersionPipelineContinuationToken.QueryPlan != null) { queryPlanFromContinuationToken = latestVersionPipelineContinuationToken.QueryPlan; @@ -342,7 +350,7 @@ public sealed class InputParameters public InputParameters( SqlQuerySpec sqlQuerySpec, - string initialUserContinuationToken, + RequestContinuationToken initialUserContinuationToken, int? maxConcurrency, int? maxItemCount, int? maxBufferedItemCount, @@ -386,7 +394,7 @@ public InputParameters( } public SqlQuerySpec SqlQuerySpec { get; } - public string InitialUserContinuationToken { get; } + public RequestContinuationToken InitialUserContinuationToken { get; } public int MaxConcurrency { get; } public int MaxItemCount { get; } public int MaxBufferedItemCount { get; } diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/PipelinedDocumentQueryExecutionContext.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/PipelinedDocumentQueryExecutionContext.cs index bb2f3e6a6e..77a8c1bec2 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/PipelinedDocumentQueryExecutionContext.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/PipelinedDocumentQueryExecutionContext.cs @@ -258,6 +258,7 @@ async Task> tryCreateParallelComponen tryCreatePipelineAsync = async (continuationToken) => { return await TakeDocumentQueryExecutionComponent.TryCreateLimitDocumentQueryExecutionComponentAsync( + executionEnvironment, queryInfo.Limit.Value, continuationToken, tryCreateSourceAsync); @@ -270,6 +271,7 @@ async Task> tryCreateParallelComponen tryCreatePipelineAsync = async (continuationToken) => { return await TakeDocumentQueryExecutionComponent.TryCreateTopDocumentQueryExecutionComponentAsync( + executionEnvironment, queryInfo.Top.Value, continuationToken, tryCreateSourceAsync); diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/UInt128.cs b/Microsoft.Azure.Cosmos/src/Query/Core/UInt128.cs index 5ddc369f41..eff9ed0399 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/UInt128.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/UInt128.cs @@ -233,17 +233,14 @@ public static UInt128 Create(ulong low, ulong high) return new UInt128(low, high); } - /// - /// Creates a UInt128 from a byte array. - /// - /// The bytes. - /// The UInt128 from the byte array. - public static UInt128 FromByteArray(ReadOnlySpan bytes) + public static UInt128 FromByteArray(ReadOnlySpan buffer) { - ulong low = MemoryMarshal.Read(bytes); - ulong high = MemoryMarshal.Read(bytes.Slice(sizeof(ulong))); + if (!UInt128.TryParse(buffer, out UInt128 value)) + { + throw new FormatException($"Malformed buffer"); + } - return new UInt128(low, high); + return value; } /// @@ -402,5 +399,18 @@ public static bool TryParse(string value, out UInt128 uInt128) uInt128 = UInt128.FromByteArray(bytes); return true; } + + public static bool TryParse(ReadOnlySpan buffer, out UInt128 value) + { + if (buffer.Length < UInt128.Length) + { + value = default; + return false; + } + + ReadOnlySpan bufferAsULongs = MemoryMarshal.Cast(buffer); + value = new UInt128(low: bufferAsULongs[0], high: bufferAsULongs[1]); + return true; + } } } diff --git a/Microsoft.Azure.Cosmos/src/Query/v3Query/FeedIteratorInlineCore.cs b/Microsoft.Azure.Cosmos/src/Query/v3Query/FeedIteratorInlineCore.cs index 55ca75e868..2abb2a39e8 100644 --- a/Microsoft.Azure.Cosmos/src/Query/v3Query/FeedIteratorInlineCore.cs +++ b/Microsoft.Azure.Cosmos/src/Query/v3Query/FeedIteratorInlineCore.cs @@ -7,11 +7,7 @@ namespace Microsoft.Azure.Cosmos using System; using System.Threading; using System.Threading.Tasks; - 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.QueryClient; - using Microsoft.Azure.Cosmos.Query.Core.QueryPlan; + using Microsoft.Azure.Cosmos.Json; internal class FeedIteratorInlineCore : FeedIteratorInternal { @@ -43,6 +39,11 @@ public override Task ReadNextAsync(CancellationToken cancellati return TaskHelper.RunInlineIfNeededAsync(() => this.feedIteratorInternal.ReadNextAsync(cancellationToken)); } + public override void SerializeState(IJsonWriter jsonWriter) + { + this.feedIteratorInternal.SerializeState(jsonWriter); + } + public override bool TryGetContinuationToken(out string continuationToken) { return this.feedIteratorInternal.TryGetContinuationToken(out continuationToken); @@ -83,5 +84,10 @@ public override bool TryGetContinuationToken(out string continuationToken) { return this.feedIteratorInternal.TryGetContinuationToken(out continuationToken); } + + public override void SerializeState(IJsonWriter jsonWriter) + { + this.feedIteratorInternal.SerializeState(jsonWriter); + } } } diff --git a/Microsoft.Azure.Cosmos/src/Query/v3Query/QueryIterator.cs b/Microsoft.Azure.Cosmos/src/Query/v3Query/QueryIterator.cs index 0051e6da7d..74997f13d6 100644 --- a/Microsoft.Azure.Cosmos/src/Query/v3Query/QueryIterator.cs +++ b/Microsoft.Azure.Cosmos/src/Query/v3Query/QueryIterator.cs @@ -7,9 +7,11 @@ namespace Microsoft.Azure.Cosmos.Query 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; + 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.QueryClient; using Microsoft.Azure.Cosmos.Query.Core.QueryPlan; @@ -55,9 +57,31 @@ public static QueryIterator Create( allowNonValueAggregateQuery: allowNonValueAggregateQuery, correlatedActivityId: Guid.NewGuid()); + RequestContinuationToken requestContinuationToken; + if (queryRequestOptions.ExecutionEnvironment.HasValue) + { + switch (queryRequestOptions.ExecutionEnvironment.Value) + { + case ExecutionEnvironment.Client: + requestContinuationToken = RequestContinuationToken.Create(continuationToken); + break; + + case ExecutionEnvironment.Compute: + requestContinuationToken = RequestContinuationToken.Create(CosmosElement.CreateFromBuffer(queryRequestOptions.BinaryContinuationToken)); + break; + + default: + throw new ArgumentOutOfRangeException($"Unknown {nameof(ExecutionEnvironment)}: {queryRequestOptions.ExecutionEnvironment.Value}."); + } + } + else + { + requestContinuationToken = RequestContinuationToken.Create(continuationToken); + } + CosmosQueryExecutionContextFactory.InputParameters inputParameters = new CosmosQueryExecutionContextFactory.InputParameters( sqlQuerySpec: sqlQuerySpec, - initialUserContinuationToken: continuationToken, + initialUserContinuationToken: requestContinuationToken, maxConcurrency: queryRequestOptions.MaxConcurrency, maxItemCount: queryRequestOptions.MaxItemCount, maxBufferedItemCount: queryRequestOptions.MaxBufferedItemCount, @@ -128,5 +152,10 @@ public override bool TryGetContinuationToken(out string continuationToken) { return this.cosmosQueryExecutionContext.TryGetContinuationToken(out continuationToken); } + + public override void SerializeState(IJsonWriter jsonWriter) + { + this.cosmosQueryExecutionContext.SerializeState(jsonWriter); + } } } diff --git a/Microsoft.Azure.Cosmos/src/Resource/QueryResponses/ChangeFeedPartitionKeyResultSetIteratorCore.cs b/Microsoft.Azure.Cosmos/src/Resource/QueryResponses/ChangeFeedPartitionKeyResultSetIteratorCore.cs index 2a6fd32234..de3b29651f 100644 --- a/Microsoft.Azure.Cosmos/src/Resource/QueryResponses/ChangeFeedPartitionKeyResultSetIteratorCore.cs +++ b/Microsoft.Azure.Cosmos/src/Resource/QueryResponses/ChangeFeedPartitionKeyResultSetIteratorCore.cs @@ -8,6 +8,7 @@ namespace Microsoft.Azure.Cosmos using System.Net; using System.Threading; using System.Threading.Tasks; + using Microsoft.Azure.Cosmos.Json; /// /// Cosmos Change Feed Iterator for a particular Partition Key Range @@ -75,6 +76,11 @@ internal ChangeFeedPartitionKeyResultSetIteratorCore( }, cancellationToken); } + public override void SerializeState(IJsonWriter jsonWriter) + { + throw new NotImplementedException(); + } + public override bool TryGetContinuationToken(out string continuationToken) { continuationToken = this.continuationToken; diff --git a/Microsoft.Azure.Cosmos/src/Resource/QueryResponses/ChangeFeedResultSetIteratorCore.cs b/Microsoft.Azure.Cosmos/src/Resource/QueryResponses/ChangeFeedResultSetIteratorCore.cs index 1d2fa2b98b..b31461cea3 100644 --- a/Microsoft.Azure.Cosmos/src/Resource/QueryResponses/ChangeFeedResultSetIteratorCore.cs +++ b/Microsoft.Azure.Cosmos/src/Resource/QueryResponses/ChangeFeedResultSetIteratorCore.cs @@ -6,8 +6,10 @@ namespace Microsoft.Azure.Cosmos { using System; using System.Net; + using System.Text; using System.Threading; using System.Threading.Tasks; + using Microsoft.Azure.Cosmos.Json; using Microsoft.Azure.Cosmos.Query; using Microsoft.Azure.Cosmos.Query.Core.ContinuationTokens; using Microsoft.Azure.Cosmos.Routing; @@ -170,5 +172,10 @@ internal virtual Task NextResultSetDelegateAsync( streamPayload: null, cancellationToken: cancellationToken); } + + public override void SerializeState(IJsonWriter jsonWriter) + { + throw new NotImplementedException(); + } } } \ No newline at end of file diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CrossPartitionQueryTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CrossPartitionQueryTests.cs index 5c0ec66775..013a127ed2 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CrossPartitionQueryTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CrossPartitionQueryTests.cs @@ -661,16 +661,16 @@ private static async Task> QueryWithSerializeState( } List resultsFromSerializeState = new List(); - ReadOnlyMemory continuationToken; + ReadOnlyMemory continuationToken = default(ReadOnlyMemory); do { QueryRequestOptions computeRequestOptions = queryRequestOptions.Clone(); computeRequestOptions.ExecutionEnvironment = Cosmos.Query.Core.ExecutionContext.ExecutionEnvironment.Compute; + computeRequestOptions.BinaryContinuationToken = continuationToken; FeedIteratorInternal itemQuery = (FeedIteratorInternal)container.GetItemQueryIterator( queryText: query, - requestOptions: computeRequestOptions, - continuationToken: continuationToken); + requestOptions: computeRequestOptions); try { FeedResponse cosmosQueryResponse = await itemQuery.ReadNextAsync(); @@ -690,8 +690,7 @@ private static async Task> QueryWithSerializeState( { itemQuery = (FeedIteratorInternal)container.GetItemQueryIterator( queryText: query, - requestOptions: queryRequestOptions, - continuationToken: continuationToken); + requestOptions: queryRequestOptions); } } while (!continuationToken.IsEmpty); 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 a9f3fc6a98..34ea6de660 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/CosmosQueryUnitTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/CosmosQueryUnitTests.cs @@ -325,11 +325,13 @@ public async Task TestCosmosQueryPartitionKeyDefinition() func)).Result); components.Add((await TakeDocumentQueryExecutionComponent.TryCreateLimitDocumentQueryExecutionComponentAsync( + ExecutionEnvironment.Client, 5, null, func)).Result); components.Add((await TakeDocumentQueryExecutionComponent.TryCreateTopDocumentQueryExecutionComponentAsync( + ExecutionEnvironment.Client, 5, null, func)).Result); diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/QueryPipelineMockTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/QueryPipelineMockTests.cs index c0de6c65d1..be05bfa018 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/QueryPipelineMockTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/QueryPipelineMockTests.cs @@ -89,7 +89,7 @@ public async Task TestCosmosCrossPartitionQueryExecutionContextWithEmptyPagesAnd CosmosParallelItemQueryExecutionContext executionContext = (await CosmosParallelItemQueryExecutionContext.TryCreateAsync( context, initParams, - fullConitnuationToken, + RequestContinuationToken.Create(fullConitnuationToken), this.cancellationToken)).Result; // Read all the pages from both splits @@ -173,7 +173,7 @@ public async Task TestCosmosCrossPartitionQueryExecutionContextWithFailuresAsync CosmosParallelItemQueryExecutionContext executionContext = (await CosmosParallelItemQueryExecutionContext.TryCreateAsync( context, initParams, - fullConitnuationToken, + RequestContinuationToken.Create(fullConitnuationToken), this.cancellationToken)).Result; // Read all the pages from both splits @@ -301,7 +301,7 @@ public async Task TestCosmosOrderByQueryExecutionContextWithEmptyPagesAndSplitAs CosmosOrderByItemQueryExecutionContext executionContext = (await CosmosOrderByItemQueryExecutionContext.TryCreateAsync( context, initParams, - fullConitnuationToken, + RequestContinuationToken.Create(fullConitnuationToken), this.cancellationToken)).Result; // For order by it will drain all the pages till it gets a value. @@ -424,7 +424,7 @@ public async Task TestCosmosOrderByQueryExecutionContextWithFailurePageAsync(boo TryCatch tryCreate = await CosmosOrderByItemQueryExecutionContext.TryCreateAsync( context, initParams, - fullConitnuationToken, + RequestContinuationToken.Create(fullConitnuationToken), this.cancellationToken); if (tryCreate.Succeeded) @@ -513,19 +513,21 @@ public async Task TestNegativeDistinctComponentCreation() TryCatch tryCreateWhenInvalidContinuationToken = await DistinctDocumentQueryExecutionComponent.TryCreateAsync( ExecutionEnvironment.Client, - "This is not a valid continuation token", + RequestContinuationToken.Create("This is not a valid continuation token"), CreateSource, DistinctQueryType.Unordered); Assert.IsFalse(tryCreateWhenInvalidContinuationToken.Succeeded); } - private static Task> FailToCreateSource(string continuationToken) + // ADD MORE TESTS HERE + + private static Task> FailToCreateSource(RequestContinuationToken continuationToken) { return Task.FromResult(TryCatch.FromException(new Exception())); } - private static Task> CreateSource(string continuationToken) + private static Task> CreateSource(RequestContinuationToken continuationToken) { return Task.FromResult(TryCatch.FromResult(new Mock().Object)); } From 81418d0ca50def2975c96f0d91b64ba370fdce8f Mon Sep 17 00:00:00 2001 From: Brandon Chong Date: Sat, 25 Jan 2020 09:34:37 -0800 Subject: [PATCH 06/28] wired in serialize state --- .../CrossPartitionQueryTests.cs | 13 ++++++++++++- 1 file changed, 12 insertions(+), 1 deletion(-) diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CrossPartitionQueryTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CrossPartitionQueryTests.cs index b9f09f89b3..4945aa5216 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CrossPartitionQueryTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CrossPartitionQueryTests.cs @@ -4684,7 +4684,7 @@ private static async Task> RunQuery( container, query, queryRequestOptions, - QueryDrainingMode.ContinuationToken | QueryDrainingMode.HoldState | QueryDrainingMode.TryGetContinuationTokens); + QueryDrainingMode.ContinuationToken | QueryDrainingMode.HoldState | QueryDrainingMode.TryGetContinuationTokens | QueryDrainingMode.SerializeState); } [Flags] @@ -4694,6 +4694,7 @@ public enum QueryDrainingMode HoldState = 1, ContinuationToken = 2, TryGetContinuationTokens = 4, + SerializeState = 8, } private static async Task> RunQueryCombinations( @@ -4739,6 +4740,16 @@ private static async Task> RunQueryCombinations( queryExecutionResults[QueryDrainingMode.TryGetContinuationTokens] = queryResultsWithTryGetContinuationToken; } + if (queryDrainingMode.HasFlag(QueryDrainingMode.SerializeState)) + { + List queryResultsWithSerializeState = await QueryWithSerializeState( + container, + query, + queryRequestOptions); + + queryExecutionResults[QueryDrainingMode.SerializeState] = queryResultsWithSerializeState; + } + foreach (QueryDrainingMode queryDrainingMode1 in queryExecutionResults.Keys) { foreach (QueryDrainingMode queryDrainingMode2 in queryExecutionResults.Keys) From f4e689efc2ce494f3f8fb3a3de5cd082e1d9522e Mon Sep 17 00:00:00 2001 From: Brandon Chong Date: Wed, 29 Jan 2020 12:57:19 -0800 Subject: [PATCH 07/28] fixed test cases --- ...DocumentQueryExecutionComponent.Compute.cs | 2 +- ...tDocumentQueryExecutionComponent.Client.cs | 9 ++-- ...DocumentQueryExecutionComponent.Compute.cs | 2 +- .../DocumentQueryExecutionComponentBase.cs | 4 +- ...DocumentQueryExecutionComponent.Compute.cs | 7 ++- ...pDocumentQueryExecutionComponent.Client.cs | 44 ++++++++++++++++++- ...DocumentQueryExecutionComponent.Compute.cs | 29 ++++++++++++ .../SkipDocumentQueryExecutionComponent.cs | 34 -------------- ...eDocumentQueryExecutionComponent.Client.cs | 39 ++++++++++++++++ ...DocumentQueryExecutionComponent.Compute.cs | 26 +++++++++++ .../TakeDocumentQueryExecutionComponent.cs | 30 ------------- .../CosmosOrderByItemQueryExecutionContext.cs | 7 ++- 12 files changed, 157 insertions(+), 76 deletions(-) diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/AggregateDocumentQueryExecutionComponent.Compute.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/AggregateDocumentQueryExecutionComponent.Compute.cs index 9a4a9bbd3f..1ab29645e0 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/AggregateDocumentQueryExecutionComponent.Compute.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/AggregateDocumentQueryExecutionComponent.Compute.cs @@ -42,7 +42,7 @@ public static async Task> TryCreateAs Func>> tryCreateSourceAsync) { AggregateContinuationToken aggregateContinuationToken; - if (requestContinuation != null) + if (!requestContinuation.IsNull) { if (!AggregateContinuationToken.TryParse(requestContinuation, out aggregateContinuationToken)) { diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctDocumentQueryExecutionComponent.Client.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctDocumentQueryExecutionComponent.Client.cs index 564a46c7a0..a2c29e3941 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctDocumentQueryExecutionComponent.Client.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctDocumentQueryExecutionComponent.Client.cs @@ -62,7 +62,7 @@ public static async Task> TryCreateAs } DistinctContinuationToken distinctContinuationToken; - if (requestContinuation != null) + if (!requestContinuation.IsNull) { if (!DistinctContinuationToken.TryParse(stringRequestContinuationToken, out distinctContinuationToken)) { @@ -124,11 +124,14 @@ public override async Task DrainAsync(int maxElements, Cancel // For clients we write out the continuation token if it's a streaming query. QueryResponseCore queryResponseCore; - if (this.TryGetContinuationToken(out string continuationToken)) + if (this.distinctQueryType == DistinctQueryType.Ordered) { + DistinctContinuationToken distinctContinuationToken = new DistinctContinuationToken( + sourceToken: sourceResponse.ContinuationToken, + distinctMapToken: this.distinctMap.GetContinuationToken()); queryResponseCore = QueryResponseCore.CreateSuccess( result: distinctResults, - continuationToken: continuationToken, + continuationToken: distinctContinuationToken.ToString(), disallowContinuationTokenMessage: null, activityId: sourceResponse.ActivityId, requestCharge: sourceResponse.RequestCharge, diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctDocumentQueryExecutionComponent.Compute.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctDocumentQueryExecutionComponent.Compute.cs index 209d7c9450..d93429d30b 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctDocumentQueryExecutionComponent.Compute.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctDocumentQueryExecutionComponent.Compute.cs @@ -53,7 +53,7 @@ public static async Task> TryCreateAs } DistinctContinuationToken distinctContinuationToken; - if (requestContinuation != null) + if (!requestContinuation.IsNull) { if (!DistinctContinuationToken.TryParse(cosmosElementRequestContinuationToken.Value, out distinctContinuationToken)) { diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/DocumentQueryExecutionComponentBase.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/DocumentQueryExecutionComponentBase.cs index 441d113be7..37e7a4c52b 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/DocumentQueryExecutionComponentBase.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/DocumentQueryExecutionComponentBase.cs @@ -14,6 +14,8 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent /// internal abstract class DocumentQueryExecutionComponentBase : IDocumentQueryExecutionComponent { + public static readonly string UseSerializeStateInstead = $"Use Serialize State instead"; + /// /// Source DocumentQueryExecutionComponent that this component will drain from. /// @@ -72,7 +74,7 @@ public bool TryGetContinuationToken(out string continuationToken) IJsonWriter jsonWriter = JsonWriter.Create(JsonSerializationFormat.Text); this.SerializeState(jsonWriter); continuationToken = Utf8StringHelpers.ToString(jsonWriter.GetResult()); - return false; + return true; } public abstract void SerializeState(IJsonWriter jsonWriter); 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 4da5259b0b..4da479037d 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 @@ -23,7 +23,6 @@ private sealed class ComputeGroupByDocumentQueryExecutionComponent : GroupByDocu private const string SourceTokenName = "SourceToken"; private const string GroupingTableContinuationTokenName = "GroupingTableContinuationToken"; private const string DoneReadingGroupingsContinuationToken = "DONE"; - private const string UseTryGetContinuationTokenInstead = "Use TryGetContinuationTokenInstead"; private static readonly IReadOnlyList EmptyResults = new List().AsReadOnly(); private static readonly IReadOnlyDictionary EmptyQueryMetrics = new Dictionary(); @@ -50,7 +49,7 @@ public static async Task> TryCreateAs } GroupByContinuationToken groupByContinuationToken; - if (requestContinuation != null) + if (!requestContinuation.IsNull) { if (!GroupByContinuationToken.TryParse(cosmosElementRequestContinuationToken.Value, out groupByContinuationToken)) { @@ -123,7 +122,7 @@ public override async Task DrainAsync( response = QueryResponseCore.CreateSuccess( result: EmptyResults, continuationToken: null, - disallowContinuationTokenMessage: UseTryGetContinuationTokenInstead, + disallowContinuationTokenMessage: DocumentQueryExecutionComponentBase.UseSerializeStateInstead, activityId: sourceResponse.ActivityId, requestCharge: sourceResponse.RequestCharge, diagnostics: sourceResponse.Diagnostics, @@ -138,7 +137,7 @@ public override async Task DrainAsync( response = QueryResponseCore.CreateSuccess( result: results, continuationToken: null, - disallowContinuationTokenMessage: UseTryGetContinuationTokenInstead, + disallowContinuationTokenMessage: DocumentQueryExecutionComponentBase.UseSerializeStateInstead, activityId: null, requestCharge: 0, diagnostics: QueryResponseCore.EmptyDiagnostics, diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/SkipDocumentQueryExecutionComponent.Client.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/SkipDocumentQueryExecutionComponent.Client.cs index 8db7ab7d92..f83c6d456e 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/SkipDocumentQueryExecutionComponent.Client.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/SkipDocumentQueryExecutionComponent.Client.cs @@ -4,11 +4,16 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.SkipTake { 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.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 Newtonsoft.Json; internal abstract partial class SkipDocumentQueryExecutionComponent : DocumentQueryExecutionComponentBase @@ -38,7 +43,7 @@ public static async Task> TryCreateAs } OffsetContinuationToken offsetContinuationToken; - if (continuationToken != null) + if (!continuationToken.IsNull) { if (!OffsetContinuationToken.TryParse(stringRequestContinuationToken.Value, out offsetContinuationToken)) { @@ -63,6 +68,43 @@ public static async Task> TryCreateAs offsetContinuationToken.Offset)); } + public override async Task DrainAsync(int maxElements, CancellationToken token) + { + token.ThrowIfCancellationRequested(); + QueryResponseCore sourcePage = await base.DrainAsync(maxElements, token); + if (!sourcePage.IsSuccess) + { + return sourcePage; + } + + // skip the documents but keep all the other headers + IReadOnlyList documentsAfterSkip = sourcePage.CosmosElements.Skip(this.skipCount).ToList(); + + int numberOfDocumentsSkipped = sourcePage.CosmosElements.Count() - documentsAfterSkip.Count(); + this.skipCount -= numberOfDocumentsSkipped; + + string updatedContinuationToken; + if (sourcePage.DisallowContinuationTokenMessage == null) + { + updatedContinuationToken = new OffsetContinuationToken( + offset: this.skipCount, + sourceToken: sourcePage.ContinuationToken).ToString(); + } + else + { + updatedContinuationToken = null; + } + + return QueryResponseCore.CreateSuccess( + result: documentsAfterSkip, + continuationToken: updatedContinuationToken, + disallowContinuationTokenMessage: sourcePage.DisallowContinuationTokenMessage, + activityId: sourcePage.ActivityId, + requestCharge: sourcePage.RequestCharge, + diagnostics: sourcePage.Diagnostics, + responseLengthBytes: sourcePage.ResponseLengthBytes); + } + public override void SerializeState(IJsonWriter jsonWriter) { throw new NotImplementedException(); diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/SkipDocumentQueryExecutionComponent.Compute.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/SkipDocumentQueryExecutionComponent.Compute.cs index 1010e70f87..29f79dcddd 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/SkipDocumentQueryExecutionComponent.Compute.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/SkipDocumentQueryExecutionComponent.Compute.cs @@ -4,12 +4,16 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.SkipTake { 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.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; internal abstract partial class SkipDocumentQueryExecutionComponent : DocumentQueryExecutionComponentBase { @@ -69,6 +73,31 @@ public static async Task> TryCreateCo offsetContinuationToken.Offset)); } + public override async Task DrainAsync(int maxElements, CancellationToken token) + { + token.ThrowIfCancellationRequested(); + QueryResponseCore sourcePage = await base.DrainAsync(maxElements, token); + if (!sourcePage.IsSuccess) + { + return sourcePage; + } + + // skip the documents but keep all the other headers + IReadOnlyList documentsAfterSkip = sourcePage.CosmosElements.Skip(this.skipCount).ToList(); + + int numberOfDocumentsSkipped = sourcePage.CosmosElements.Count() - documentsAfterSkip.Count(); + this.skipCount -= numberOfDocumentsSkipped; + + return QueryResponseCore.CreateSuccess( + result: documentsAfterSkip, + continuationToken: null, + disallowContinuationTokenMessage: DocumentQueryExecutionComponentBase.UseSerializeStateInstead, + activityId: sourcePage.ActivityId, + requestCharge: sourcePage.RequestCharge, + diagnostics: sourcePage.Diagnostics, + responseLengthBytes: sourcePage.ResponseLengthBytes); + } + public override void SerializeState(IJsonWriter jsonWriter) { if (jsonWriter == null) 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 3382cca14e..a9ae9c9680 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/SkipDocumentQueryExecutionComponent.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/SkipDocumentQueryExecutionComponent.cs @@ -67,39 +67,5 @@ public override bool IsDone return this.Source.IsDone; } } - - public override async Task DrainAsync(int maxElements, CancellationToken token) - { - token.ThrowIfCancellationRequested(); - QueryResponseCore sourcePage = await base.DrainAsync(maxElements, token); - if (!sourcePage.IsSuccess) - { - return sourcePage; - } - - // skip the documents but keep all the other headers - IReadOnlyList documentsAfterSkip = sourcePage.CosmosElements.Skip(this.skipCount).ToList(); - - int numberOfDocumentsSkipped = sourcePage.CosmosElements.Count() - documentsAfterSkip.Count(); - this.skipCount -= numberOfDocumentsSkipped; - string updatedContinuationToken = null; - - if (sourcePage.DisallowContinuationTokenMessage == null) - { - if (!this.TryGetContinuationToken(out updatedContinuationToken)) - { - throw new InvalidOperationException($"Failed to get state for {nameof(SkipDocumentQueryExecutionComponent)}."); - } - } - - return QueryResponseCore.CreateSuccess( - result: documentsAfterSkip, - continuationToken: updatedContinuationToken, - disallowContinuationTokenMessage: sourcePage.DisallowContinuationTokenMessage, - activityId: sourcePage.ActivityId, - requestCharge: sourcePage.RequestCharge, - diagnostics: sourcePage.Diagnostics, - responseLengthBytes: sourcePage.ResponseLengthBytes); - } } } \ No newline at end of file diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/TakeDocumentQueryExecutionComponent.Client.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/TakeDocumentQueryExecutionComponent.Client.cs index 27ecbe16c5..539a8c499a 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/TakeDocumentQueryExecutionComponent.Client.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/TakeDocumentQueryExecutionComponent.Client.cs @@ -4,11 +4,16 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.SkipTake { 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.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 Newtonsoft.Json; internal abstract partial class TakeDocumentQueryExecutionComponent : DocumentQueryExecutionComponentBase @@ -127,6 +132,40 @@ public static async Task> TryCreateTo TakeEnum.Top)); } + public override async Task DrainAsync(int maxElements, CancellationToken token) + { + token.ThrowIfCancellationRequested(); + QueryResponseCore sourcePage = await base.DrainAsync(maxElements, token); + if (!sourcePage.IsSuccess) + { + return sourcePage; + } + + List takedDocuments = sourcePage.CosmosElements.Take(this.takeCount).ToList(); + this.takeCount -= takedDocuments.Count; + + string updatedContinuationToken; + if (sourcePage.DisallowContinuationTokenMessage == null) + { + updatedContinuationToken = new LimitContinuationToken( + limit: this.takeCount, + sourceToken: sourcePage.ContinuationToken).ToString(); + } + else + { + updatedContinuationToken = null; + } + + return QueryResponseCore.CreateSuccess( + result: takedDocuments, + continuationToken: updatedContinuationToken, + disallowContinuationTokenMessage: sourcePage.DisallowContinuationTokenMessage, + activityId: sourcePage.ActivityId, + requestCharge: sourcePage.RequestCharge, + diagnostics: sourcePage.Diagnostics, + responseLengthBytes: sourcePage.ResponseLengthBytes); + } + public override void SerializeState(IJsonWriter jsonWriter) { throw new NotImplementedException(); diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/TakeDocumentQueryExecutionComponent.Compute.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/TakeDocumentQueryExecutionComponent.Compute.cs index 96ff3e1b91..d99239aaab 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/TakeDocumentQueryExecutionComponent.Compute.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/TakeDocumentQueryExecutionComponent.Compute.cs @@ -4,12 +4,16 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.SkipTake { 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.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; internal abstract partial class TakeDocumentQueryExecutionComponent : DocumentQueryExecutionComponentBase { @@ -75,6 +79,28 @@ public static async Task> TryCreateAs takeContinuationToken.TakeCount)); } + public override async Task DrainAsync(int maxElements, CancellationToken token) + { + token.ThrowIfCancellationRequested(); + QueryResponseCore sourcePage = await base.DrainAsync(maxElements, token); + if (!sourcePage.IsSuccess) + { + return sourcePage; + } + + List takedDocuments = sourcePage.CosmosElements.Take(this.takeCount).ToList(); + this.takeCount -= takedDocuments.Count; + + return QueryResponseCore.CreateSuccess( + result: takedDocuments, + continuationToken: null, + disallowContinuationTokenMessage: DocumentQueryExecutionComponentBase.UseSerializeStateInstead, + activityId: sourcePage.ActivityId, + requestCharge: sourcePage.RequestCharge, + diagnostics: sourcePage.Diagnostics, + responseLengthBytes: sourcePage.ResponseLengthBytes); + } + public override void SerializeState(IJsonWriter jsonWriter) { if (jsonWriter == null) 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 608656b841..51655109cb 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/TakeDocumentQueryExecutionComponent.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/TakeDocumentQueryExecutionComponent.cs @@ -98,35 +98,5 @@ public override bool IsDone return this.Source.IsDone || this.takeCount <= 0; } } - - public override async Task DrainAsync(int maxElements, CancellationToken token) - { - token.ThrowIfCancellationRequested(); - QueryResponseCore results = await base.DrainAsync(maxElements, token); - if (!results.IsSuccess) - { - return results; - } - - List takedDocuments = results.CosmosElements.Take(this.takeCount).ToList(); - this.takeCount -= takedDocuments.Count; - - string updatedContinuationToken = null; - if (results.DisallowContinuationTokenMessage == null) - { - IJsonWriter jsonWriter = Json.JsonWriter.Create(JsonSerializationFormat.Text); - this.SerializeState(jsonWriter); - updatedContinuationToken = Utf8StringHelpers.ToString(jsonWriter.GetResult()); - } - - return QueryResponseCore.CreateSuccess( - result: takedDocuments, - continuationToken: updatedContinuationToken, - disallowContinuationTokenMessage: results.DisallowContinuationTokenMessage, - activityId: results.ActivityId, - requestCharge: results.RequestCharge, - diagnostics: results.Diagnostics, - responseLengthBytes: results.ResponseLengthBytes); - } } } \ No newline at end of file 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 a843273c55..586097aad5 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/OrderBy/CosmosOrderByItemQueryExecutionContext.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/OrderBy/CosmosOrderByItemQueryExecutionContext.cs @@ -374,6 +374,11 @@ private async Task TryInitializeAsync( throw new ArgumentNullException(nameof(sqlQuerySpec)); } + if (requestContinuation == null) + { + throw new ArgumentNullException(nameof(requestContinuation)); + } + if (collectionRid == null) { throw new ArgumentNullException(nameof(collectionRid)); @@ -396,7 +401,7 @@ private async Task TryInitializeAsync( cancellationToken.ThrowIfCancellationRequested(); - if (requestContinuation == null) + if (requestContinuation.IsNull) { SqlQuerySpec sqlQuerySpecForInit = new SqlQuerySpec( sqlQuerySpec.QueryText.Replace(oldValue: FormatPlaceHolder, newValue: True), From 3e470e0fd520155a452d65ce9456cf5e68e249db Mon Sep 17 00:00:00 2001 From: Brandon Chong Date: Wed, 29 Jan 2020 14:54:33 -0800 Subject: [PATCH 08/28] I give up --- .../StringRequestContinuationToken.cs | 2 ++ ...ateDocumentQueryExecutionComponent.Client.cs | 10 +++++----- .../Aggregators/SingleGroupAggregator.cs | 17 ++++++++++++++++- .../GroupByDocumentQueryExecutionComponent.cs | 2 +- ...ipDocumentQueryExecutionComponent.Compute.cs | 2 +- ...keDocumentQueryExecutionComponent.Compute.cs | 14 +++++++------- .../CrossPartitionQueryTests.cs | 6 +++--- 7 files changed, 35 insertions(+), 18 deletions(-) diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/StringRequestContinuationToken.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/StringRequestContinuationToken.cs index d099384ecd..7a3588fb86 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/StringRequestContinuationToken.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/StringRequestContinuationToken.cs @@ -8,6 +8,8 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ContinuationTokens internal sealed class StringRequestContinuationToken : RequestContinuationToken { + public static readonly StringRequestContinuationToken Null = new StringRequestContinuationToken(continuationToken: null); + public StringRequestContinuationToken(string continuationToken) { this.Value = continuationToken; diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/AggregateDocumentQueryExecutionComponent.Client.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/AggregateDocumentQueryExecutionComponent.Client.cs index 253782746b..c4538041db 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/AggregateDocumentQueryExecutionComponent.Client.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/AggregateDocumentQueryExecutionComponent.Client.cs @@ -41,11 +41,11 @@ public static async Task> TryCreateAs } TryCatch tryCreateSingleGroupAggregator = SingleGroupAggregator.TryCreate( - aggregates, - aliasToAggregateType, - orderedAliases, - hasSelectValue, - continuationToken: null); + aggregates, + aliasToAggregateType, + orderedAliases, + hasSelectValue, + continuationToken: StringRequestContinuationToken.Null); if (!tryCreateSingleGroupAggregator.Succeeded) { diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/SingleGroupAggregator.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/SingleGroupAggregator.cs index 501d07e99c..bb0ef8fdb0 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/SingleGroupAggregator.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/SingleGroupAggregator.cs @@ -192,8 +192,23 @@ public static TryCatch TryCreate( IReadOnlyList orderedAliases, RequestContinuationToken continuationToken) { + if (aggregateAliasToAggregateType == null) + { + throw new ArgumentNullException(nameof(aggregateAliasToAggregateType)); + } + + if (orderedAliases == null) + { + throw new ArgumentNullException(nameof(orderedAliases)); + } + + if (continuationToken == null) + { + throw new ArgumentNullException(nameof(continuationToken)); + } + CosmosObject aliasToContinuationToken; - if (continuationToken != null) + if (!continuationToken.IsNull) { if (!continuationToken.TryConvertToCosmosElement(out aliasToContinuationToken)) { 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 8537a8e184..77f9f4ca01 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/GroupBy/GroupByDocumentQueryExecutionComponent.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/GroupBy/GroupByDocumentQueryExecutionComponent.cs @@ -208,7 +208,7 @@ public void AddPayload(RewrittenGroupByProjection rewrittenGroupByProjection) this.groupByAliasToAggregateType, this.orderedAliases, this.hasSelectValue, - continuationToken: null).Result; + continuationToken: StringRequestContinuationToken.Null).Result; this.table[groupByKeysHash] = singleGroupAggregator; } diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/SkipDocumentQueryExecutionComponent.Compute.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/SkipDocumentQueryExecutionComponent.Compute.cs index 29f79dcddd..2d51a4a36a 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/SkipDocumentQueryExecutionComponent.Compute.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/SkipDocumentQueryExecutionComponent.Compute.cs @@ -45,7 +45,7 @@ public static async Task> TryCreateCo } OffsetContinuationToken offsetContinuationToken; - if (continuationToken != null) + if (!continuationToken.IsNull) { (bool parsed, OffsetContinuationToken parsedToken) = OffsetContinuationToken.TryParse(cosmosElementRequestContinuationToken.Value); if (!parsed) diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/TakeDocumentQueryExecutionComponent.Compute.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/TakeDocumentQueryExecutionComponent.Compute.cs index d99239aaab..c32fb08b8e 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/TakeDocumentQueryExecutionComponent.Compute.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/TakeDocumentQueryExecutionComponent.Compute.cs @@ -92,13 +92,13 @@ public override async Task DrainAsync(int maxElements, Cancel this.takeCount -= takedDocuments.Count; return QueryResponseCore.CreateSuccess( - result: takedDocuments, - continuationToken: null, - disallowContinuationTokenMessage: DocumentQueryExecutionComponentBase.UseSerializeStateInstead, - activityId: sourcePage.ActivityId, - requestCharge: sourcePage.RequestCharge, - diagnostics: sourcePage.Diagnostics, - responseLengthBytes: sourcePage.ResponseLengthBytes); + result: takedDocuments, + continuationToken: null, + disallowContinuationTokenMessage: DocumentQueryExecutionComponentBase.UseSerializeStateInstead, + activityId: sourcePage.ActivityId, + requestCharge: sourcePage.RequestCharge, + diagnostics: sourcePage.Diagnostics, + responseLengthBytes: sourcePage.ResponseLengthBytes); } public override void SerializeState(IJsonWriter jsonWriter) diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CrossPartitionQueryTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CrossPartitionQueryTests.cs index 4945aa5216..27e5982ce3 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CrossPartitionQueryTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CrossPartitionQueryTests.cs @@ -2272,7 +2272,7 @@ public async Task TestNonValueAggregates() await this.CreateIngestQueryDelete( ConnectionModes.Direct, - CollectionTypes.SinglePartition | CollectionTypes.MultiPartition, + /*CollectionTypes.SinglePartition |*/ CollectionTypes.MultiPartition, documents, this.TestNonValueAggregates); } @@ -2733,7 +2733,7 @@ private async Task TestQueryDistinct(Container container, IEnumerable MaxConcurrency = 10, MaxItemCount = 100, }, - QueryDrainingMode.TryGetContinuationTokens | QueryDrainingMode.HoldState); + QueryDrainingMode.TryGetContinuationTokens | QueryDrainingMode.HoldState | QueryDrainingMode.SerializeState); documentsFromWithoutDistinct = documentsFromWithoutDistinct .Where(document => documentsSeen.Add(document, out UInt128 hash)) .ToList(); @@ -2749,7 +2749,7 @@ private async Task TestQueryDistinct(Container container, IEnumerable MaxConcurrency = 10, MaxItemCount = pageSize }, - QueryDrainingMode.TryGetContinuationTokens | QueryDrainingMode.HoldState); + QueryDrainingMode.TryGetContinuationTokens | QueryDrainingMode.HoldState | QueryDrainingMode.SerializeState); Assert.IsTrue( documentsFromWithDistinct.SequenceEqual(documentsFromWithoutDistinct, JsonTokenEqualityComparer.Value), From d5ff581069f16d8456e3a205a025a247e62cd252 Mon Sep 17 00:00:00 2001 From: Brandon Chong Date: Thu, 30 Jan 2020 08:29:07 -0800 Subject: [PATCH 09/28] can't do it --- .../CompositeContinuationTokenRefStruct.cs | 2 -- .../Aggregators/AverageAggregator.cs | 18 ++++++++++----- .../Aggregate/Aggregators/CountAggregator.cs | 12 ++++++++-- .../Aggregate/Aggregators/MinMaxAggregator.cs | 22 ++++++++++++++----- .../Aggregators/SingleGroupAggregator.cs | 2 +- 5 files changed, 39 insertions(+), 17 deletions(-) diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/CompositeContinuationTokenRefStruct.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/CompositeContinuationTokenRefStruct.cs index 3fb3e60764..f61f3f2a0a 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/CompositeContinuationTokenRefStruct.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/CompositeContinuationTokenRefStruct.cs @@ -6,8 +6,6 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ContinuationTokens { using System; using Microsoft.Azure.Cosmos.Json; - using Microsoft.Azure.Cosmos.Query.Core.Exceptions; - using Microsoft.Azure.Cosmos.Query.Core.Monads; internal ref struct CompositeContinuationTokenRefStruct { diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/AverageAggregator.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/AverageAggregator.cs index 99f8c9b6f9..f1f081838d 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/AverageAggregator.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/AverageAggregator.cs @@ -7,6 +7,7 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.Aggregate.Aggrega 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; @@ -64,10 +65,15 @@ public void SerializeState(IJsonWriter jsonWriter) this.globalAverage.SerializeState(jsonWriter); } - public static TryCatch TryCreate(string continuationToken) + public static TryCatch TryCreate(RequestContinuationToken continuationToken) { + if (continuationToken == null) + { + throw new ArgumentNullException(nameof(continuationToken)); + } + AverageInfo averageInfo; - if (continuationToken != null) + if (!continuationToken.IsNull) { if (!AverageInfo.TryParse(continuationToken, out averageInfo)) { @@ -213,14 +219,14 @@ public override string ToString() }}"; } - public static bool TryParse(string serializedAverageInfo, out AverageInfo averageInfo) + public static bool TryParse(RequestContinuationToken requestContinuationToken, out AverageInfo averageInfo) { - if (serializedAverageInfo == null) + if (requestContinuationToken == null) { - throw new ArgumentNullException(nameof(serializedAverageInfo)); + throw new ArgumentNullException(nameof(requestContinuationToken)); } - if (!CosmosElement.TryParse(serializedAverageInfo, out CosmosElement cosmosElementAverageInfo)) + if (!requestContinuationToken.TryConvertToCosmosElement(out CosmosElement cosmosElementAverageInfo)) { averageInfo = default; return false; diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/CountAggregator.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/CountAggregator.cs index dc2c1acbd2..383b044c8a 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/CountAggregator.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/CountAggregator.cs @@ -8,6 +8,7 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.Aggregate.Aggrega 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; @@ -71,16 +72,23 @@ public void SerializeState(IJsonWriter jsonWriter) jsonWriter.WriteInt64Value(this.globalCount); } - public static TryCatch TryCreate(string continuationToken) + public static TryCatch TryCreate(RequestContinuationToken continuationToken) { + if (continuationToken == null) + { + throw new ArgumentNullException(nameof(continuationToken)); + } + long partialCount; if (continuationToken != null) { - if (!long.TryParse(continuationToken, out partialCount)) + if (!continuationToken.TryConvertToCosmosElement(out CosmosInt64 cosmosPartialCount)) { return TryCatch.FromException( new MalformedContinuationTokenException($@"Invalid count continuation token: ""{continuationToken}"".")); } + + partialCount = cosmosPartialCount.GetValue(); } else { diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/MinMaxAggregator.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/MinMaxAggregator.cs index 479477dfd8..17b4047d4b 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/MinMaxAggregator.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/MinMaxAggregator.cs @@ -6,6 +6,7 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.Aggregate.Aggrega using System; 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.ExecutionContext.ItemProducers; using Microsoft.Azure.Cosmos.Query.Core.Monads; @@ -165,21 +166,31 @@ public void SerializeState(IJsonWriter jsonWriter) } } - public static TryCatch TryCreateMinAggregator(string continuationToken) + public static TryCatch TryCreateMinAggregator(RequestContinuationToken continuationToken) { return MinMaxAggregator.TryCreate(isMinAggregation: true, continuationToken: continuationToken); } - public static TryCatch TryCreateMaxAggregator(string continuationToken) + public static TryCatch TryCreateMaxAggregator(RequestContinuationToken continuationToken) { return MinMaxAggregator.TryCreate(isMinAggregation: false, continuationToken: continuationToken); } - private static TryCatch TryCreate(bool isMinAggregation, string continuationToken) + private static TryCatch TryCreate(bool isMinAggregation, RequestContinuationToken continuationToken) { + if (continuationToken == null) + { + throw new ArgumentNullException(nameof(continuationToken)); + } + CosmosElement globalMinMax; - if (continuationToken != null) + if (!continuationToken.IsNull) { + if (!continuationToken.TryConvertToCosmosElement(out CosmosElement globalMinMax)) + { + return TryCatch.FromException( + new MalformedContinuationTokenException($"Malformed continuation token: {continuationToken}")); + } if (continuationToken == MinMaxAggregator.MaxValueContinuationToken) { globalMinMax = ItemComparer.MaxValue; @@ -196,8 +207,7 @@ private static TryCatch TryCreate(bool isMinAggregation, string con { if (!CosmosElement.TryParse(continuationToken, out globalMinMax)) { - return TryCatch.FromException( - new MalformedContinuationTokenException($"Malformed continuation token: {continuationToken}")); + } } } diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/SingleGroupAggregator.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/SingleGroupAggregator.cs index bb0ef8fdb0..0c10128397 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/SingleGroupAggregator.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/SingleGroupAggregator.cs @@ -343,7 +343,7 @@ public override void SerializeState(IJsonWriter jsonWriter) public static TryCatch TryCreate( AggregateOperator aggregateOperator, - string continuationToken) + RequestContinuationToken continuationToken) { TryCatch tryCreateAggregator; switch (aggregateOperator) From 7ebd1589d6230c0f2cc4b3bbc35f5d9bb4021e63 Mon Sep 17 00:00:00 2001 From: Brandon Chong Date: Wed, 12 Feb 2020 18:03:51 -0800 Subject: [PATCH 10/28] added min max continuation token --- .../Aggregate/Aggregators/MinMaxAggregator.cs | 103 ++++------- .../Aggregators/MinMaxContinuationToken.cs | 165 ++++++++++++++++++ 2 files changed, 203 insertions(+), 65 deletions(-) create mode 100644 Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/MinMaxContinuationToken.cs diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/MinMaxAggregator.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/MinMaxAggregator.cs index 17b4047d4b..a29e0b84e9 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/MinMaxAggregator.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/MinMaxAggregator.cs @@ -4,6 +4,8 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.Aggregate.Aggregators { using System; + using System.Collections.Generic; + using System.Runtime.CompilerServices; using Microsoft.Azure.Cosmos.CosmosElements; using Microsoft.Azure.Cosmos.Json; using Microsoft.Azure.Cosmos.Query.Core.ContinuationTokens; @@ -18,10 +20,6 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.Aggregate.Aggrega /// internal sealed class MinMaxAggregator : IAggregator { - private const string MinValueContinuationToken = "MIN_VALUE"; - private const string MaxValueContinuationToken = "MAX_VALUE"; - private const string UndefinedContinuationToken = "UNDEFINED"; - private static readonly CosmosElement Undefined = null; /// /// Whether or not the aggregation is a min or a max. @@ -121,7 +119,7 @@ public void Aggregate(CosmosElement localMinMax) public CosmosElement GetResult() { CosmosElement result; - if (this.globalMinMax == ItemComparer.MinValue || this.globalMinMax == ItemComparer.MaxValue) + if ((this.globalMinMax == ItemComparer.MinValue) || (this.globalMinMax == ItemComparer.MaxValue)) { // The filter did not match any documents. result = Undefined; @@ -148,22 +146,27 @@ public void SerializeState(IJsonWriter jsonWriter) throw new ArgumentNullException(nameof(jsonWriter)); } + jsonWriter.WriteObjectStart(); + + MinMaxContinuationToken minMaxContinuationToken; if (this.globalMinMax == ItemComparer.MinValue) { - jsonWriter.WriteStringValue(MinMaxAggregator.MinValueContinuationToken); + minMaxContinuationToken = MinMaxContinuationToken.CreateMinValueContinuationToken(); } else if (this.globalMinMax == ItemComparer.MaxValue) { - jsonWriter.WriteStringValue(MinMaxAggregator.MaxValueContinuationToken); + minMaxContinuationToken = MinMaxContinuationToken.CreateMaxValueContinuationToken(); } else if (this.globalMinMax == Undefined) { - jsonWriter.WriteStringValue(MinMaxAggregator.UndefinedContinuationToken); + minMaxContinuationToken = MinMaxContinuationToken.CreateUndefinedValueContinuationToken(); } else { - this.globalMinMax.WriteTo(jsonWriter); + minMaxContinuationToken = MinMaxContinuationToken.CreateValueContinuationToken(this.globalMinMax); } + + MinMaxContinuationToken.ToCosmosElement(minMaxContinuationToken).WriteTo(jsonWriter); } public static TryCatch TryCreateMinAggregator(RequestContinuationToken continuationToken) @@ -186,29 +189,38 @@ private static TryCatch TryCreate(bool isMinAggregation, RequestCon CosmosElement globalMinMax; if (!continuationToken.IsNull) { - if (!continuationToken.TryConvertToCosmosElement(out CosmosElement globalMinMax)) + if (!continuationToken.TryConvertToCosmosElement(out CosmosElement cosmosElementContinuationToken)) { return TryCatch.FromException( - new MalformedContinuationTokenException($"Malformed continuation token: {continuationToken}")); - } - if (continuationToken == MinMaxAggregator.MaxValueContinuationToken) - { - globalMinMax = ItemComparer.MaxValue; - } - else if (continuationToken == MinMaxAggregator.MinValueContinuationToken) - { - globalMinMax = ItemComparer.MinValue; + new MalformedContinuationTokenException($"Malformed continuation token: {continuationToken}")); } - else if (continuationToken == MinMaxAggregator.UndefinedContinuationToken) + + TryCatch tryCreateMinMaxContinuationToken = MinMaxContinuationToken.TryCreateFromCosmosElement(cosmosElementContinuationToken); + if (!tryCreateMinMaxContinuationToken.Succeeded) { - globalMinMax = MinMaxAggregator.Undefined; + return TryCatch.FromException(tryCreateMinMaxContinuationToken.Exception); } - else + + switch (tryCreateMinMaxContinuationToken.Result.Type) { - if (!CosmosElement.TryParse(continuationToken, out globalMinMax)) - { - - } + case MinMaxContinuationToken.MinMaxContinuationTokenType.MinValue: + globalMinMax = ItemComparer.MinValue; + break; + + case MinMaxContinuationToken.MinMaxContinuationTokenType.MaxValue: + globalMinMax = ItemComparer.MaxValue; + break; + + case MinMaxContinuationToken.MinMaxContinuationTokenType.Undefined: + globalMinMax = MinMaxAggregator.Undefined; + break; + + case MinMaxContinuationToken.MinMaxContinuationTokenType.Value: + globalMinMax = tryCreateMinMaxContinuationToken.Result.Value; + break; + + default: + throw new ArgumentOutOfRangeException($"Unknown {nameof(MinMaxContinuationToken.MinMaxContinuationTokenType)}: {tryCreateMinMaxContinuationToken.Result.Type}"); } } else @@ -219,44 +231,5 @@ private static TryCatch TryCreate(bool isMinAggregation, RequestCon return TryCatch.FromResult( new MinMaxAggregator(isMinAggregation: isMinAggregation, globalMinMax: globalMinMax)); } - - private static bool CosmosElementIsPrimitive(CosmosElement cosmosElement) - { - if (cosmosElement == null) - { - return false; - } - - CosmosElementType cosmosElementType = cosmosElement.Type; - switch (cosmosElementType) - { - case CosmosElementType.Array: - return false; - - case CosmosElementType.Boolean: - return true; - - case CosmosElementType.Null: - return true; - - case CosmosElementType.Number: - return true; - - case CosmosElementType.Object: - return false; - - case CosmosElementType.String: - return true; - - case CosmosElementType.Guid: - return true; - - case CosmosElementType.Binary: - return true; - - default: - throw new ArgumentException($"Unknown {nameof(CosmosElementType)} : {cosmosElementType}."); - } - } } } diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/MinMaxContinuationToken.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/MinMaxContinuationToken.cs new file mode 100644 index 0000000000..df8c199c39 --- /dev/null +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/MinMaxContinuationToken.cs @@ -0,0 +1,165 @@ +// ------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +// ------------------------------------------------------------ + +namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.Aggregate.Aggregators +{ + 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; + + internal sealed class MinMaxContinuationToken + { + private const string TypeName = "type"; + private const string ValueName = "value"; + + private MinMaxContinuationToken( + MinMaxContinuationTokenType type, + CosmosElement value) + { + switch (type) + { + case MinMaxContinuationTokenType.MinValue: + case MinMaxContinuationTokenType.MaxValue: + case MinMaxContinuationTokenType.Undefined: + if (value != null) + { + throw new ArgumentException($"{nameof(value)} must be set with type: {type}."); + } + break; + + case MinMaxContinuationTokenType.Value: + if (value == null) + { + throw new ArgumentException($"{nameof(value)} must not be set with type: {type}."); + } + break; + + default: + throw new ArgumentOutOfRangeException($"Unknown {nameof(type)}: {type}."); + } + + this.Type = type; + this.Value = value; + } + + public MinMaxContinuationTokenType Type { get; } + public CosmosElement Value { get; } + + public static MinMaxContinuationToken CreateMinValueContinuationToken() + { + return new MinMaxContinuationToken(type: MinMaxContinuationTokenType.MinValue, value: null); + } + + public static MinMaxContinuationToken CreateMaxValueContinuationToken() + { + return new MinMaxContinuationToken(type: MinMaxContinuationTokenType.MaxValue, value: null); + } + + public static MinMaxContinuationToken CreateUndefinedValueContinuationToken() + { + return new MinMaxContinuationToken(type: MinMaxContinuationTokenType.MaxValue, value: null); + } + + public static MinMaxContinuationToken CreateValueContinuationToken(CosmosElement value) + { + return new MinMaxContinuationToken(type: MinMaxContinuationTokenType.Value, value: value); + } + + public static CosmosElement ToCosmosElement(MinMaxContinuationToken minMaxContinuationToken) + { + if (minMaxContinuationToken == null) + { + throw new ArgumentNullException(nameof(minMaxContinuationToken)); + } + + Dictionary dictionary = new Dictionary(); + dictionary.Add( + MinMaxContinuationToken.TypeName, + EnumToCosmosString.ConvertEnumToCosmosString(minMaxContinuationToken.Type)); + if (minMaxContinuationToken.Value != null) + { + dictionary.Add(MinMaxContinuationToken.ValueName, minMaxContinuationToken.Value); + } + + return CosmosObject.Create(dictionary); + } + + public static TryCatch TryCreateFromCosmosElement(CosmosElement cosmosElement) + { + if (!(cosmosElement is CosmosObject cosmosObject)) + { + return TryCatch.FromException( + new MalformedContinuationTokenException($"{nameof(MinMaxContinuationToken)} was not an object.")); + } + + if (!cosmosObject.TryGetValue(MinMaxContinuationToken.TypeName, out CosmosString typeValue)) + { + return TryCatch.FromException( + new MalformedContinuationTokenException($"{nameof(MinMaxContinuationToken)} is missing property: {MinMaxContinuationToken.TypeName}.")); + } + + if (!Enum.TryParse(typeValue.Value, out MinMaxContinuationTokenType minMaxContinuationTokenType)) + { + return TryCatch.FromException( + new MalformedContinuationTokenException($"{nameof(MinMaxContinuationToken)} has malformed '{MinMaxContinuationToken.TypeName}': {typeValue.Value}.")); + } + + CosmosElement value; + if (minMaxContinuationTokenType == MinMaxContinuationTokenType.Value) + { + if (!cosmosObject.TryGetValue(MinMaxContinuationToken.ValueName, out value)) + { + return TryCatch.FromException( + new MalformedContinuationTokenException($"{nameof(MinMaxContinuationToken)} is missing property: {MinMaxContinuationToken.ValueName}.")); + } + } + else + { + value = null; + } + + return TryCatch.FromResult( + new MinMaxContinuationToken(minMaxContinuationTokenType, value)); + } + + private static class EnumToCosmosString + { + private static readonly CosmosString MinValueCosmosString = CosmosString.Create(MinMaxContinuationTokenType.MinValue.ToString()); + private static readonly CosmosString MaxValueCosmosString = CosmosString.Create(MinMaxContinuationTokenType.MaxValue.ToString()); + private static readonly CosmosString UndefinedCosmosString = CosmosString.Create(MinMaxContinuationTokenType.Undefined.ToString()); + private static readonly CosmosString ValueCosmosString = CosmosString.Create(MinMaxContinuationTokenType.Value.ToString()); + + public static CosmosString ConvertEnumToCosmosString(MinMaxContinuationTokenType type) + { + switch (type) + { + case MinMaxContinuationTokenType.MinValue: + return EnumToCosmosString.MinValueCosmosString; + + case MinMaxContinuationTokenType.MaxValue: + return EnumToCosmosString.MaxValueCosmosString; + + case MinMaxContinuationTokenType.Undefined: + return EnumToCosmosString.UndefinedCosmosString; + + case MinMaxContinuationTokenType.Value: + return EnumToCosmosString.ValueCosmosString; + + default: + throw new ArgumentOutOfRangeException($"Unknown {nameof(MinMaxContinuationTokenType)}: {type}"); + } + } + } + + public enum MinMaxContinuationTokenType + { + MinValue, + MaxValue, + Undefined, + Value, + } + } +} From 20000cb02138f4d29d539d7629360c146792d295 Mon Sep 17 00:00:00 2001 From: Brandon Chong Date: Wed, 12 Feb 2020 18:44:17 -0800 Subject: [PATCH 11/28] removed TryGetContinuationToken --- .../src/FeedIteratorCore.cs | 11 --- .../src/FeedIteratorInternal.cs | 1 - .../src/FeedIteratorInternal{T}.cs | 1 - .../src/Json/JsonNavigator.cs | 5 ++ ...DocumentQueryExecutionComponent.Compute.cs | 7 +- .../Aggregate/Aggregators/MinMaxAggregator.cs | 62 +++++++++++++++- .../Aggregate/Aggregators/SumAggregator.cs | 11 +-- ...DocumentQueryExecutionComponent.Compute.cs | 2 +- .../DocumentQueryExecutionComponentBase.cs | 8 --- ...DocumentQueryExecutionComponent.Compute.cs | 6 -- .../IDocumentQueryExecutionComponent.cs | 2 - .../CatchAllCosmosQueryExecutionContext.cs | 5 -- ...smosCrossPartitionQueryExecutionContext.cs | 6 -- .../CosmosQueryExecutionContext.cs | 2 - ...ExecutionContextWithNameCacheStaleRetry.cs | 5 -- .../LazyCosmosQueryExecutionContext.cs | 18 ----- .../PipelinedDocumentQueryExecutionContext.cs | 5 -- .../Query/v3Query/FeedIteratorInlineCore.cs | 10 --- .../src/Query/v3Query/QueryIterator.cs | 5 -- ...geFeedPartitionKeyResultSetIteratorCore.cs | 6 -- .../ChangeFeedResultSetIteratorCore.cs | 6 -- .../CrossPartitionQueryTests.cs | 71 ++----------------- .../CosmosQueryUnitTests.cs | 2 +- 23 files changed, 81 insertions(+), 176 deletions(-) diff --git a/Microsoft.Azure.Cosmos/src/FeedIteratorCore.cs b/Microsoft.Azure.Cosmos/src/FeedIteratorCore.cs index 37ff23e5c9..60b809f716 100644 --- a/Microsoft.Azure.Cosmos/src/FeedIteratorCore.cs +++ b/Microsoft.Azure.Cosmos/src/FeedIteratorCore.cs @@ -98,12 +98,6 @@ public override async Task ReadNextAsync(CancellationToken canc return response; } - public override bool TryGetContinuationToken(out string continuationToken) - { - continuationToken = this.continuationToken; - return true; - } - internal static string GetContinuationToken(ResponseMessage httpResponseMessage) { return httpResponseMessage.Headers.ContinuationToken; @@ -155,11 +149,6 @@ public override async Task> ReadNextAsync(CancellationToken canc return this.responseCreator(response); } - public override bool TryGetContinuationToken(out string continuationToken) - { - return this.feedIterator.TryGetContinuationToken(out continuationToken); - } - public override void SerializeState(IJsonWriter jsonWriter) { this.feedIterator.SerializeState(jsonWriter); diff --git a/Microsoft.Azure.Cosmos/src/FeedIteratorInternal.cs b/Microsoft.Azure.Cosmos/src/FeedIteratorInternal.cs index 75ae9f944b..b4da79e520 100644 --- a/Microsoft.Azure.Cosmos/src/FeedIteratorInternal.cs +++ b/Microsoft.Azure.Cosmos/src/FeedIteratorInternal.cs @@ -18,7 +18,6 @@ namespace Microsoft.Azure.Cosmos #endif abstract class FeedIteratorInternal : FeedIterator { - public abstract bool TryGetContinuationToken(out string continuationToken); public abstract void SerializeState(IJsonWriter jsonWriter); } } \ No newline at end of file diff --git a/Microsoft.Azure.Cosmos/src/FeedIteratorInternal{T}.cs b/Microsoft.Azure.Cosmos/src/FeedIteratorInternal{T}.cs index a468f8d7bc..b8732678a9 100644 --- a/Microsoft.Azure.Cosmos/src/FeedIteratorInternal{T}.cs +++ b/Microsoft.Azure.Cosmos/src/FeedIteratorInternal{T}.cs @@ -18,7 +18,6 @@ namespace Microsoft.Azure.Cosmos #endif abstract class FeedIteratorInternal : FeedIterator { - public abstract bool TryGetContinuationToken(out string continuationToken); public abstract void SerializeState(IJsonWriter jsonWriter); } } \ No newline at end of file diff --git a/Microsoft.Azure.Cosmos/src/Json/JsonNavigator.cs b/Microsoft.Azure.Cosmos/src/Json/JsonNavigator.cs index 9cf693cd8f..754cbf8f5d 100644 --- a/Microsoft.Azure.Cosmos/src/Json/JsonNavigator.cs +++ b/Microsoft.Azure.Cosmos/src/Json/JsonNavigator.cs @@ -39,6 +39,11 @@ public static IJsonNavigator Create( JsonStringDictionary jsonStringDictionary = null, bool skipValidation = false) { + if (buffer.IsEmpty) + { + throw new ArgumentOutOfRangeException($"{nameof(buffer)} buffer can not be empty."); + } + // Examine the first buffer byte to determine the serialization format byte firstByte = buffer.Span[0]; diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/AggregateDocumentQueryExecutionComponent.Compute.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/AggregateDocumentQueryExecutionComponent.Compute.cs index 1ab29645e0..229cb3d2cc 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/AggregateDocumentQueryExecutionComponent.Compute.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/AggregateDocumentQueryExecutionComponent.Compute.cs @@ -147,11 +147,6 @@ private QueryResponseCore GetFinalResponse() private QueryResponseCore GetEmptyPage(QueryResponseCore sourceResponse) { - if (!this.TryGetContinuationToken(out string updatedContinuationToken)) - { - throw new InvalidOperationException("Failed to get source continuation token."); - } - // We need to give empty pages until the results are fully drained. QueryResponseCore response = QueryResponseCore.CreateSuccess( result: EmptyResults, @@ -159,7 +154,7 @@ private QueryResponseCore GetEmptyPage(QueryResponseCore sourceResponse) activityId: sourceResponse.ActivityId, responseLengthBytes: sourceResponse.ResponseLengthBytes, disallowContinuationTokenMessage: null, - continuationToken: updatedContinuationToken, + continuationToken: sourceResponse.ContinuationToken, diagnostics: sourceResponse.Diagnostics); return response; diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/MinMaxAggregator.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/MinMaxAggregator.cs index a29e0b84e9..856fe6ac9e 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/MinMaxAggregator.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/MinMaxAggregator.cs @@ -91,8 +91,7 @@ public void Aggregate(CosmosElement localMinMax) } } - if (!ItemComparer.IsMinOrMax(this.globalMinMax) - && (!CosmosElementIsPrimitive(localMinMax) || !CosmosElementIsPrimitive(this.globalMinMax))) + if (!ItemComparer.IsMinOrMax(this.globalMinMax) && (!IsPrimitve(localMinMax) || !IsPrimitve(this.globalMinMax))) { // This means we are comparing non primitives which is undefined this.globalMinMax = Undefined; @@ -231,5 +230,64 @@ private static TryCatch TryCreate(bool isMinAggregation, RequestCon return TryCatch.FromResult( new MinMaxAggregator(isMinAggregation: isMinAggregation, globalMinMax: globalMinMax)); } + + private static bool IsPrimitve(CosmosElement cosmosElement) + { + if (cosmosElement == Undefined) + { + return false; + } + + return cosmosElement.Accept(IsPrimitiveCosmosElementVisitor.Singleton); + } + + private sealed class IsPrimitiveCosmosElementVisitor : ICosmosElementVisitor + { + public static readonly IsPrimitiveCosmosElementVisitor Singleton = new IsPrimitiveCosmosElementVisitor(); + + private IsPrimitiveCosmosElementVisitor() + { + } + + public bool Visit(CosmosArray cosmosArray) + { + return false; + } + + public bool Visit(CosmosBinary cosmosBinary) + { + return true; + } + + public bool Visit(CosmosBoolean cosmosBoolean) + { + return true; + } + + public bool Visit(CosmosGuid cosmosGuid) + { + return true; + } + + public bool Visit(CosmosNull cosmosNull) + { + return true; + } + + public bool Visit(CosmosNumber cosmosNumber) + { + return true; + } + + public bool Visit(CosmosObject cosmosObject) + { + return false; + } + + public bool Visit(CosmosString cosmosString) + { + return true; + } + } } } diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/SumAggregator.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/SumAggregator.cs index 1e16fbd565..36b6d4ab54 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/SumAggregator.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/SumAggregator.cs @@ -8,6 +8,7 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.Aggregate.Aggrega 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; @@ -79,16 +80,18 @@ public void SerializeState(IJsonWriter jsonWriter) jsonWriter.WriteFloat64Value(this.globalSum); } - public static TryCatch TryCreate(string continuationToken) + public static TryCatch TryCreate(RequestContinuationToken requestContinuationToken) { double partialSum; - if (continuationToken != null) + if (!requestContinuationToken.IsNull) { - if (!double.TryParse(continuationToken, out partialSum)) + if (!requestContinuationToken.TryConvertToCosmosElement(out CosmosFloat64 cosmosFloat64)) { return TryCatch.FromException( - new MalformedContinuationTokenException($"Malformed {nameof(SumAggregator)} continuation token: {continuationToken}")); + new MalformedContinuationTokenException($"Malformed {nameof(SumAggregator)} continuation token: {requestContinuationToken}")); } + + partialSum = cosmosFloat64.GetValue(); } else { diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctDocumentQueryExecutionComponent.Compute.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctDocumentQueryExecutionComponent.Compute.cs index d93429d30b..0ed0af3f52 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctDocumentQueryExecutionComponent.Compute.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctDocumentQueryExecutionComponent.Compute.cs @@ -22,7 +22,7 @@ internal abstract partial class DistinctDocumentQueryExecutionComponent : Docume /// private sealed class ComputeDistinctDocumentQueryExecutionComponent : DistinctDocumentQueryExecutionComponent { - private static readonly string UseTryGetContinuationTokenMessage = $"Use {nameof(ComputeDistinctDocumentQueryExecutionComponent.TryGetContinuationToken)}"; + private static readonly string UseTryGetContinuationTokenMessage = $"Use TryGetContinuationToken"; private ComputeDistinctDocumentQueryExecutionComponent( DistinctQueryType distinctQueryType, diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/DocumentQueryExecutionComponentBase.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/DocumentQueryExecutionComponentBase.cs index 37e7a4c52b..81ec8b548b 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/DocumentQueryExecutionComponentBase.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/DocumentQueryExecutionComponentBase.cs @@ -69,14 +69,6 @@ public void Stop() this.Source.Stop(); } - public bool TryGetContinuationToken(out string continuationToken) - { - IJsonWriter jsonWriter = JsonWriter.Create(JsonSerializationFormat.Text); - this.SerializeState(jsonWriter); - continuationToken = Utf8StringHelpers.ToString(jsonWriter.GetResult()); - return true; - } - public abstract void SerializeState(IJsonWriter jsonWriter); } } 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 4da479037d..f1df7f4b01 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 @@ -258,12 +258,6 @@ public void SerializeState(IJsonWriter jsonWriter) public void Stop() { } - - public bool TryGetContinuationToken(out string state) - { - state = null; - return true; - } } } } diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/IDocumentQueryExecutionComponent.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/IDocumentQueryExecutionComponent.cs index 812ca08141..5cb50ad609 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/IDocumentQueryExecutionComponent.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/IDocumentQueryExecutionComponent.cs @@ -33,8 +33,6 @@ internal interface IDocumentQueryExecutionComponent : IDisposable /// void Stop(); - bool TryGetContinuationToken(out string state); - void SerializeState(IJsonWriter jsonWriter); } } diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CatchAllCosmosQueryExecutionContext.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CatchAllCosmosQueryExecutionContext.cs index e99f49b09a..68edc998f8 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CatchAllCosmosQueryExecutionContext.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CatchAllCosmosQueryExecutionContext.cs @@ -56,11 +56,6 @@ public override async Task ExecuteNextAsync(CancellationToken return queryResponseCore; } - public override bool TryGetContinuationToken(out string continuationToken) - { - return this.cosmosQueryExecutionContext.TryGetContinuationToken(out continuationToken); - } - public override void SerializeState(IJsonWriter jsonWriter) { if (jsonWriter == null) diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosCrossPartitionQueryExecutionContext.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosCrossPartitionQueryExecutionContext.cs index 7ad91b1830..c95c048310 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosCrossPartitionQueryExecutionContext.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosCrossPartitionQueryExecutionContext.cs @@ -662,12 +662,6 @@ private void OnItemProducerTreeCompleteFetching( } } - public bool TryGetContinuationToken(out string state) - { - state = this.ContinuationToken; - return true; - } - public abstract void SerializeState(IJsonWriter jsonWriter); public readonly struct InitInfo diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosQueryExecutionContext.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosQueryExecutionContext.cs index 872eaf90f7..7d234a1ba0 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosQueryExecutionContext.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosQueryExecutionContext.cs @@ -31,8 +31,6 @@ public abstract bool IsDone /// A task to await on, which in return provides a DoucmentFeedResponse of documents. public abstract Task ExecuteNextAsync(CancellationToken cancellationToken); - public abstract bool TryGetContinuationToken(out string continuationToken); - public abstract void SerializeState(IJsonWriter jsonWriter); } } diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosQueryExecutionContextWithNameCacheStaleRetry.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosQueryExecutionContextWithNameCacheStaleRetry.cs index 01c8df5e72..f25a11d684 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosQueryExecutionContextWithNameCacheStaleRetry.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosQueryExecutionContextWithNameCacheStaleRetry.cs @@ -60,11 +60,6 @@ await this.cosmosQueryContext.QueryClient.ForceRefreshCollectionCacheAsync( return queryResponse; } - public override bool TryGetContinuationToken(out string continuationToken) - { - return this.currentCosmosQueryExecutionContext.TryGetContinuationToken(out continuationToken); - } - public override void SerializeState(IJsonWriter jsonWriter) { if (jsonWriter == null) diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/LazyCosmosQueryExecutionContext.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/LazyCosmosQueryExecutionContext.cs index 4f6f6f04c8..a48fa2d30a 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/LazyCosmosQueryExecutionContext.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/LazyCosmosQueryExecutionContext.cs @@ -76,24 +76,6 @@ public override async Task ExecuteNextAsync(CancellationToken return queryResponseCore; } - public override bool TryGetContinuationToken(out string state) - { - if (!this.lazyTryCreateCosmosQueryExecutionContext.ValueInitialized) - { - state = null; - return false; - } - - TryCatch tryCreateCosmosQueryExecutionContext = this.lazyTryCreateCosmosQueryExecutionContext.Result; - if (!tryCreateCosmosQueryExecutionContext.Succeeded) - { - state = null; - return false; - } - - return tryCreateCosmosQueryExecutionContext.Result.TryGetContinuationToken(out state); - } - public override void SerializeState(IJsonWriter jsonWriter) { if (jsonWriter == null) diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/PipelinedDocumentQueryExecutionContext.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/PipelinedDocumentQueryExecutionContext.cs index 91cf0fd65d..3f02a9df1e 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/PipelinedDocumentQueryExecutionContext.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/PipelinedDocumentQueryExecutionContext.cs @@ -120,11 +120,6 @@ public override bool IsDone } } - public override bool TryGetContinuationToken(out string state) - { - return this.component.TryGetContinuationToken(out state); - } - public override void SerializeState(IJsonWriter jsonWriter) { if (jsonWriter == null) diff --git a/Microsoft.Azure.Cosmos/src/Query/v3Query/FeedIteratorInlineCore.cs b/Microsoft.Azure.Cosmos/src/Query/v3Query/FeedIteratorInlineCore.cs index 2abb2a39e8..86e98941a5 100644 --- a/Microsoft.Azure.Cosmos/src/Query/v3Query/FeedIteratorInlineCore.cs +++ b/Microsoft.Azure.Cosmos/src/Query/v3Query/FeedIteratorInlineCore.cs @@ -43,11 +43,6 @@ public override void SerializeState(IJsonWriter jsonWriter) { this.feedIteratorInternal.SerializeState(jsonWriter); } - - public override bool TryGetContinuationToken(out string continuationToken) - { - return this.feedIteratorInternal.TryGetContinuationToken(out continuationToken); - } } internal class FeedIteratorInlineCore : FeedIteratorInternal @@ -80,11 +75,6 @@ public override Task> ReadNextAsync(CancellationToken cancellati return TaskHelper.RunInlineIfNeededAsync(() => this.feedIteratorInternal.ReadNextAsync(cancellationToken)); } - public override bool TryGetContinuationToken(out string continuationToken) - { - return this.feedIteratorInternal.TryGetContinuationToken(out continuationToken); - } - public override void SerializeState(IJsonWriter jsonWriter) { this.feedIteratorInternal.SerializeState(jsonWriter); diff --git a/Microsoft.Azure.Cosmos/src/Query/v3Query/QueryIterator.cs b/Microsoft.Azure.Cosmos/src/Query/v3Query/QueryIterator.cs index fa1cca7ade..99293b92ff 100644 --- a/Microsoft.Azure.Cosmos/src/Query/v3Query/QueryIterator.cs +++ b/Microsoft.Azure.Cosmos/src/Query/v3Query/QueryIterator.cs @@ -164,11 +164,6 @@ public override async Task ReadNextAsync(CancellationToken canc } } - public override bool TryGetContinuationToken(out string continuationToken) - { - return this.cosmosQueryExecutionContext.TryGetContinuationToken(out continuationToken); - } - public override void SerializeState(IJsonWriter jsonWriter) { this.cosmosQueryExecutionContext.SerializeState(jsonWriter); diff --git a/Microsoft.Azure.Cosmos/src/Resource/QueryResponses/ChangeFeedPartitionKeyResultSetIteratorCore.cs b/Microsoft.Azure.Cosmos/src/Resource/QueryResponses/ChangeFeedPartitionKeyResultSetIteratorCore.cs index 8626a50514..30e12ca6c0 100644 --- a/Microsoft.Azure.Cosmos/src/Resource/QueryResponses/ChangeFeedPartitionKeyResultSetIteratorCore.cs +++ b/Microsoft.Azure.Cosmos/src/Resource/QueryResponses/ChangeFeedPartitionKeyResultSetIteratorCore.cs @@ -81,12 +81,6 @@ public override void SerializeState(IJsonWriter jsonWriter) throw new NotImplementedException(); } - public override bool TryGetContinuationToken(out string continuationToken) - { - continuationToken = this.continuationToken; - return true; - } - private Task NextResultSetDelegateAsync( string continuationToken, string partitionKeyRangeId, diff --git a/Microsoft.Azure.Cosmos/src/Resource/QueryResponses/ChangeFeedResultSetIteratorCore.cs b/Microsoft.Azure.Cosmos/src/Resource/QueryResponses/ChangeFeedResultSetIteratorCore.cs index 488d92a977..0219ec88ef 100644 --- a/Microsoft.Azure.Cosmos/src/Resource/QueryResponses/ChangeFeedResultSetIteratorCore.cs +++ b/Microsoft.Azure.Cosmos/src/Resource/QueryResponses/ChangeFeedResultSetIteratorCore.cs @@ -87,12 +87,6 @@ internal ChangeFeedResultSetIteratorCore( return response; } - public override bool TryGetContinuationToken(out string state) - { - state = this.continuationToken; - return true; - } - internal async Task> ReadNextInternalAsync(CancellationToken cancellationToken) { cancellationToken.ThrowIfCancellationRequested(); diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CrossPartitionQueryTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CrossPartitionQueryTests.cs index 27e5982ce3..ed9810bf7a 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CrossPartitionQueryTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CrossPartitionQueryTests.cs @@ -603,54 +603,6 @@ private CosmosClient CreateNewCosmosClient(ConnectionMode connectionMode) } } - private static async Task> QueryWithTryGetContinuationTokens( - Container container, - string query, - QueryRequestOptions queryRequestOptions = null) - { - if (queryRequestOptions == null) - { - queryRequestOptions = new QueryRequestOptions(); - } - - List resultsFromTryGetContinuationToken = new List(); - string continuationToken = null; - do - { - QueryRequestOptions computeRequestOptions = queryRequestOptions.Clone(); - computeRequestOptions.ExecutionEnvironment = Cosmos.Query.Core.ExecutionContext.ExecutionEnvironment.Compute; - - FeedIteratorInternal itemQuery = (FeedIteratorInternal)container.GetItemQueryIterator( - queryText: query, - requestOptions: computeRequestOptions, - continuationToken: continuationToken); - try - { - FeedResponse cosmosQueryResponse = await itemQuery.ReadNextAsync(); - if (queryRequestOptions.MaxItemCount.HasValue) - { - Assert.IsTrue( - cosmosQueryResponse.Count <= queryRequestOptions.MaxItemCount.Value, - "Max Item Count is not being honored"); - } - - resultsFromTryGetContinuationToken.AddRange(cosmosQueryResponse); - Assert.IsTrue( - itemQuery.TryGetContinuationToken(out continuationToken), - "Failed to get state for query"); - } - catch (CosmosException cosmosException) when (cosmosException.StatusCode == (HttpStatusCode)429) - { - itemQuery = (FeedIteratorInternal)container.GetItemQueryIterator( - queryText: query, - requestOptions: queryRequestOptions, - continuationToken: continuationToken); - } - } while (continuationToken != null); - - return resultsFromTryGetContinuationToken; - } - private static async Task> QueryWithSerializeState( Container container, string query, @@ -1460,7 +1412,7 @@ async Task ValidateNonDeterministicQuery(Func, useOrderBy); await ValidateNonDeterministicQuery(CrossPartitionQueryTests.QueryWithContinuationTokens, useOrderBy); - await ValidateNonDeterministicQuery(CrossPartitionQueryTests.QueryWithTryGetContinuationTokens, useOrderBy); + await ValidateNonDeterministicQuery(CrossPartitionQueryTests.QueryWithSerializeState, useOrderBy); } } } @@ -2733,7 +2685,7 @@ private async Task TestQueryDistinct(Container container, IEnumerable MaxConcurrency = 10, MaxItemCount = 100, }, - QueryDrainingMode.TryGetContinuationTokens | QueryDrainingMode.HoldState | QueryDrainingMode.SerializeState); + QueryDrainingMode.HoldState | QueryDrainingMode.SerializeState); documentsFromWithoutDistinct = documentsFromWithoutDistinct .Where(document => documentsSeen.Add(document, out UInt128 hash)) .ToList(); @@ -2749,7 +2701,7 @@ private async Task TestQueryDistinct(Container container, IEnumerable MaxConcurrency = 10, MaxItemCount = pageSize }, - QueryDrainingMode.TryGetContinuationTokens | QueryDrainingMode.HoldState | QueryDrainingMode.SerializeState); + QueryDrainingMode.HoldState | QueryDrainingMode.SerializeState); Assert.IsTrue( documentsFromWithDistinct.SequenceEqual(documentsFromWithoutDistinct, JsonTokenEqualityComparer.Value), @@ -4341,7 +4293,7 @@ FROM c }); HashSet actualWithoutContinuationTokensSet = new HashSet(actualWithoutContinuationTokens, JsonTokenEqualityComparer.Value); - List actualWithTryGetContinuationTokens = await CrossPartitionQueryTests.QueryWithTryGetContinuationTokens( + List actualWithTryGetContinuationTokens = await CrossPartitionQueryTests.QueryWithSerializeState( container, query, new QueryRequestOptions() @@ -4684,7 +4636,7 @@ private static async Task> RunQuery( container, query, queryRequestOptions, - QueryDrainingMode.ContinuationToken | QueryDrainingMode.HoldState | QueryDrainingMode.TryGetContinuationTokens | QueryDrainingMode.SerializeState); + QueryDrainingMode.ContinuationToken | QueryDrainingMode.HoldState | QueryDrainingMode.SerializeState); } [Flags] @@ -4693,8 +4645,7 @@ public enum QueryDrainingMode None = 0, HoldState = 1, ContinuationToken = 2, - TryGetContinuationTokens = 4, - SerializeState = 8, + SerializeState = 4, } private static async Task> RunQueryCombinations( @@ -4730,16 +4681,6 @@ private static async Task> RunQueryCombinations( queryExecutionResults[QueryDrainingMode.ContinuationToken] = queryResultsWithContinuationTokens; } - if (queryDrainingMode.HasFlag(QueryDrainingMode.TryGetContinuationTokens)) - { - List queryResultsWithTryGetContinuationToken = await QueryWithTryGetContinuationTokens( - container, - query, - queryRequestOptions); - - queryExecutionResults[QueryDrainingMode.TryGetContinuationTokens] = queryResultsWithTryGetContinuationToken; - } - if (queryDrainingMode.HasFlag(QueryDrainingMode.SerializeState)) { List queryResultsWithSerializeState = await QueryWithSerializeState( 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 e1ef51f7b1..7c2d86b8ad 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/CosmosQueryUnitTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/CosmosQueryUnitTests.cs @@ -342,7 +342,7 @@ public async Task TestCosmosQueryPartitionKeyDefinition() private (Func>>, QueryResponseCore) SetupBaseContextToVerifyFailureScenario() { - CosmosDiagnosticsContext diagnosticsContext = new CosmosDiagnosticsContext(); + CosmosDiagnosticsContext diagnosticsContext = CosmosDiagnosticsContext.Create(); diagnosticsContext.AddDiagnosticsInternal( new PointOperationStatistics( Guid.NewGuid().ToString(), System.Net.HttpStatusCode.Unauthorized, From e76d0531b7d4f66290b7cacfe1f25e81d5bee910 Mon Sep 17 00:00:00 2001 From: Brandon Chong Date: Fri, 14 Feb 2020 00:36:08 -0800 Subject: [PATCH 12/28] fixed random bugs --- .../CosmosElements/CosmosElementSerializer.cs | 2 +- .../src/Json/JsonNavigator.cs | 2 +- Microsoft.Azure.Cosmos/src/Json/JsonReader.cs | 5 + .../src/Json/JsonWriter.JsonBinaryWriter.cs | 6 + .../CompositeContinuationTokenRefStruct.cs | 12 +- .../CosmosElementRequestContinuationToken.cs | 2 + .../OrderByContinuationTokenRefStruct.cs | 9 +- .../Core/ContinuationTokens/RangeRefStruct.cs | 3 + ...eDocumentQueryExecutionComponent.Client.cs | 6 + ...DocumentQueryExecutionComponent.Compute.cs | 187 +++++++++--------- .../Aggregators/AverageAggregator.cs | 10 +- .../Aggregate/Aggregators/CountAggregator.cs | 2 +- .../Aggregate/Aggregators/MinMaxAggregator.cs | 5 +- .../Aggregators/MinMaxContinuationToken.cs | 2 +- .../Aggregators/SingleGroupAggregator.cs | 7 +- ...DocumentQueryExecutionComponent.Compute.cs | 8 +- .../GroupByDocumentQueryExecutionComponent.cs | 12 +- ...pDocumentQueryExecutionComponent.Client.cs | 2 +- ...DocumentQueryExecutionComponent.Compute.cs | 6 +- .../SkipDocumentQueryExecutionComponent.cs | 9 +- ...eDocumentQueryExecutionComponent.Client.cs | 36 ++-- .../CosmosOrderByItemQueryExecutionContext.cs | 4 - ...CosmosParallelItemQueryExecutionContext.cs | 6 +- .../src/Query/v3Query/QueryIterator.cs | 12 +- .../src/Serializer/CosmosSerializerCore.cs | 7 +- .../CrossPartitionQueryTests.cs | 162 +++++++-------- 26 files changed, 283 insertions(+), 241 deletions(-) diff --git a/Microsoft.Azure.Cosmos/src/CosmosElements/CosmosElementSerializer.cs b/Microsoft.Azure.Cosmos/src/CosmosElements/CosmosElementSerializer.cs index bcaf594471..917b04ecd4 100644 --- a/Microsoft.Azure.Cosmos/src/CosmosElements/CosmosElementSerializer.cs +++ b/Microsoft.Azure.Cosmos/src/CosmosElements/CosmosElementSerializer.cs @@ -254,7 +254,7 @@ internal static IEnumerable GetResources( return new List(); } - if (typeof(T) == typeof(CosmosElement)) + if (typeof(CosmosElement).IsAssignableFrom(typeof(T))) { return cosmosArray.Cast(); } diff --git a/Microsoft.Azure.Cosmos/src/Json/JsonNavigator.cs b/Microsoft.Azure.Cosmos/src/Json/JsonNavigator.cs index 754cbf8f5d..c3d411312c 100644 --- a/Microsoft.Azure.Cosmos/src/Json/JsonNavigator.cs +++ b/Microsoft.Azure.Cosmos/src/Json/JsonNavigator.cs @@ -41,7 +41,7 @@ public static IJsonNavigator Create( { if (buffer.IsEmpty) { - throw new ArgumentOutOfRangeException($"{nameof(buffer)} buffer can not be empty."); + throw new ArgumentOutOfRangeException($"{nameof(buffer)} can not be empty."); } // Examine the first buffer byte to determine the serialization format diff --git a/Microsoft.Azure.Cosmos/src/Json/JsonReader.cs b/Microsoft.Azure.Cosmos/src/Json/JsonReader.cs index 6a00ac92ba..70d41f8727 100644 --- a/Microsoft.Azure.Cosmos/src/Json/JsonReader.cs +++ b/Microsoft.Azure.Cosmos/src/Json/JsonReader.cs @@ -67,6 +67,11 @@ public JsonTokenType CurrentTokenType /// A concrete JsonReader that can read the supplied byte array. public static IJsonReader Create(ReadOnlyMemory buffer, JsonStringDictionary jsonStringDictionary = null, bool skipValidation = false) { + if (buffer.IsEmpty) + { + throw new ArgumentOutOfRangeException($"{nameof(buffer)} can not be empty."); + } + byte firstByte = buffer.Span[0]; // Explicitly pick from the set of supported formats, or otherwise assume text format diff --git a/Microsoft.Azure.Cosmos/src/Json/JsonWriter.JsonBinaryWriter.cs b/Microsoft.Azure.Cosmos/src/Json/JsonWriter.JsonBinaryWriter.cs index fbca600a5a..04208bb4cf 100644 --- a/Microsoft.Azure.Cosmos/src/Json/JsonWriter.JsonBinaryWriter.cs +++ b/Microsoft.Azure.Cosmos/src/Json/JsonWriter.JsonBinaryWriter.cs @@ -301,6 +301,12 @@ public override ReadOnlyMemory GetResult() throw new JsonNotCompleteException(); } + if (this.binaryWriter.Position == 1) + { + // We haven't written anything but the type marker, so just return an empty buffer. + return ReadOnlyMemory.Empty; + } + return this.binaryWriter.Buffer.Slice( 0, this.binaryWriter.Position); diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/CompositeContinuationTokenRefStruct.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/CompositeContinuationTokenRefStruct.cs index f61f3f2a0a..2d86a99b90 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/CompositeContinuationTokenRefStruct.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/CompositeContinuationTokenRefStruct.cs @@ -30,10 +30,20 @@ public void WriteTo(IJsonWriter jsonWriter) } jsonWriter.WriteObjectStart(); + jsonWriter.WriteFieldName(CompositeContinuationTokenRefStruct.TokenProperytName); - jsonWriter.WriteStringValue(this.BackendContinuationToken); + if (this.BackendContinuationToken != null) + { + jsonWriter.WriteStringValue(this.BackendContinuationToken); + } + else + { + jsonWriter.WriteNullValue(); + } + jsonWriter.WriteFieldName(CompositeContinuationTokenRefStruct.RangePropertyName); this.Range.WriteTo(jsonWriter); + jsonWriter.WriteObjectEnd(); } } diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/CosmosElementRequestContinuationToken.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/CosmosElementRequestContinuationToken.cs index 8b8d85711a..9f6a74ce37 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/CosmosElementRequestContinuationToken.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/CosmosElementRequestContinuationToken.cs @@ -8,6 +8,8 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ContinuationTokens internal sealed class CosmosElementRequestContinuationToken : RequestContinuationToken { + public static readonly CosmosElementRequestContinuationToken Null = new CosmosElementRequestContinuationToken(null); + public CosmosElementRequestContinuationToken(CosmosElement continuationToken) { this.Value = continuationToken; diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/OrderByContinuationTokenRefStruct.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/OrderByContinuationTokenRefStruct.cs index 448f84a141..1a74d4d420 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/OrderByContinuationTokenRefStruct.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/OrderByContinuationTokenRefStruct.cs @@ -46,8 +46,10 @@ public void WriteTo(IJsonWriter jsonWriter) } jsonWriter.WriteObjectStart(); + jsonWriter.WriteFieldName(OrderByContinuationTokenRefStruct.CompositeTokenPropertyName); this.CompositeContinuationToken.WriteTo(jsonWriter); + jsonWriter.WriteFieldName(OrderByContinuationTokenRefStruct.OrderByItemsPropertyName); jsonWriter.WriteArrayStart(); @@ -65,10 +67,13 @@ public void WriteTo(IJsonWriter jsonWriter) } jsonWriter.WriteArrayEnd(); + jsonWriter.WriteFieldName(OrderByContinuationTokenRefStruct.RidPropertyName); jsonWriter.WriteStringValue(this.Rid); + jsonWriter.WriteFieldName(OrderByContinuationTokenRefStruct.SkipCountPropertyName); - jsonWriter.WriteInt32Value(this.SkipCount); + jsonWriter.WriteNumberValue(this.SkipCount); + jsonWriter.WriteFieldName(OrderByContinuationTokenRefStruct.FilterPropertyName); if (this.Filter != null) { @@ -78,6 +83,8 @@ public void WriteTo(IJsonWriter jsonWriter) { jsonWriter.WriteNullValue(); } + + jsonWriter.WriteObjectEnd(); } } } diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/RangeRefStruct.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/RangeRefStruct.cs index e4801e343e..bd9856605a 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/RangeRefStruct.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/RangeRefStruct.cs @@ -39,10 +39,13 @@ public void WriteTo(IJsonWriter 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/ExecutionComponent/Aggregate/AggregateDocumentQueryExecutionComponent.Client.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/AggregateDocumentQueryExecutionComponent.Client.cs index c4538041db..3180fe52b1 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/AggregateDocumentQueryExecutionComponent.Client.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/AggregateDocumentQueryExecutionComponent.Client.cs @@ -9,6 +9,7 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.Aggregate 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; @@ -115,6 +116,11 @@ public override async Task DrainAsync( diagnostics: diagnosticsPages, responseLengthBytes: responseLengthBytes); } + + public override void SerializeState(IJsonWriter jsonWriter) + { + throw new NotImplementedException(); + } } } } diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/AggregateDocumentQueryExecutionComponent.Compute.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/AggregateDocumentQueryExecutionComponent.Compute.cs index 229cb3d2cc..66b32ff367 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/AggregateDocumentQueryExecutionComponent.Compute.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/AggregateDocumentQueryExecutionComponent.Compute.cs @@ -17,13 +17,13 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.Aggregate internal abstract partial class AggregateDocumentQueryExecutionComponent : DocumentQueryExecutionComponentBase { - private const string SourceTokenName = "SourceToken"; - private const string AggregationTokenName = "AggregationToken"; - - private static readonly IReadOnlyList EmptyResults = new List().AsReadOnly(); - private sealed class ComputeAggregateDocumentQueryExecutionComponent : AggregateDocumentQueryExecutionComponent { + private const string SourceTokenName = "SourceToken"; + private const string AggregationTokenName = "AggregationToken"; + + private static readonly IReadOnlyList EmptyResults = new List().AsReadOnly(); + private ComputeAggregateDocumentQueryExecutionComponent( IDocumentQueryExecutionComponent source, SingleGroupAggregator singleGroupAggregator, @@ -122,119 +122,116 @@ public override async Task DrainAsync( return response; } - } - private QueryResponseCore GetFinalResponse() - { - List finalResult = new List(); - CosmosElement aggregationResult = this.singleGroupAggregator.GetResult(); - if (aggregationResult != null) + private QueryResponseCore GetFinalResponse() { - finalResult.Add(aggregationResult); - } - - QueryResponseCore response = QueryResponseCore.CreateSuccess( - result: finalResult, - requestCharge: 0, - activityId: null, - responseLengthBytes: 0, - disallowContinuationTokenMessage: null, - continuationToken: null, - diagnostics: QueryResponseCore.EmptyDiagnostics); + List finalResult = new List(); + CosmosElement aggregationResult = this.singleGroupAggregator.GetResult(); + if (aggregationResult != null) + { + finalResult.Add(aggregationResult); + } - return response; - } + QueryResponseCore response = QueryResponseCore.CreateSuccess( + result: finalResult, + requestCharge: 0, + activityId: null, + responseLengthBytes: 0, + disallowContinuationTokenMessage: null, + continuationToken: null, + diagnostics: QueryResponseCore.EmptyDiagnostics); - 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, - diagnostics: sourceResponse.Diagnostics); - - return response; - } - - public override void SerializeState(IJsonWriter jsonWriter) - { - if (jsonWriter == null) - { - throw new ArgumentNullException(nameof(jsonWriter)); + return response; } - if (!this.IsDone) + private QueryResponseCore GetEmptyPage(QueryResponseCore sourceResponse) { - jsonWriter.WriteObjectStart(); - jsonWriter.WriteFieldName(ComputeAggregateDocumentQueryExecutionComponent.SourceTokenName); - this.Source.SerializeState(jsonWriter); - jsonWriter.WriteFieldName(ComputeAggregateDocumentQueryExecutionComponent.AggregationTokenName); - this.singleGroupAggregator.SerializeState(jsonWriter); - jsonWriter.WriteObjectEnd(); - } - } - - private readonly struct AggregateContinuationToken - { - private const string SingleGroupAggregatorContinuationTokenName = "SingleGroupAggregatorContinuationToken"; - private const string SourceContinuationTokenName = "SourceContinuationToken"; + // 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, + diagnostics: sourceResponse.Diagnostics); - public AggregateContinuationToken( - CosmosElement singleGroupAggregatorContinuationToken, - CosmosElement sourceContinuationToken) - { - this.SingleGroupAggregatorContinuationToken = singleGroupAggregatorContinuationToken; - this.SourceContinuationToken = sourceContinuationToken; + return response; } - public CosmosElement SingleGroupAggregatorContinuationToken { get; } - - public CosmosElement SourceContinuationToken { get; } - - public static bool TryParse( - RequestContinuationToken continuationToken, - out AggregateContinuationToken aggregateContinuationToken) + public override void SerializeState(IJsonWriter jsonWriter) { - if (continuationToken == null) + if (jsonWriter == null) { - throw new ArgumentNullException(nameof(continuationToken)); + throw new ArgumentNullException(nameof(jsonWriter)); } - if (!(continuationToken is CosmosElementRequestContinuationToken cosmosElementRequestContinuationToken)) + if (!this.IsDone) { - throw new ArgumentException($"Expected {nameof(CosmosElementRequestContinuationToken)} instead of: {continuationToken.GetType()}"); + jsonWriter.WriteObjectStart(); + jsonWriter.WriteFieldName(ComputeAggregateDocumentQueryExecutionComponent.SourceTokenName); + this.Source.SerializeState(jsonWriter); + jsonWriter.WriteFieldName(ComputeAggregateDocumentQueryExecutionComponent.AggregationTokenName); + this.singleGroupAggregator.SerializeState(jsonWriter); + jsonWriter.WriteObjectEnd(); } + } - if (!(cosmosElementRequestContinuationToken.Value is CosmosObject rawAggregateContinuationToken)) + private readonly struct AggregateContinuationToken + { + public AggregateContinuationToken( + CosmosElement singleGroupAggregatorContinuationToken, + CosmosElement sourceContinuationToken) { - aggregateContinuationToken = default; - return false; + this.SingleGroupAggregatorContinuationToken = singleGroupAggregatorContinuationToken; + this.SourceContinuationToken = sourceContinuationToken; } - if (!rawAggregateContinuationToken.TryGetValue( - AggregateContinuationToken.SingleGroupAggregatorContinuationTokenName, - out CosmosElement singleGroupAggregatorContinuationToken)) - { - aggregateContinuationToken = default; - return false; - } + public CosmosElement SingleGroupAggregatorContinuationToken { get; } - if (!rawAggregateContinuationToken.TryGetValue( - AggregateContinuationToken.SourceContinuationTokenName, - out CosmosElement sourceContinuationToken)) + public CosmosElement SourceContinuationToken { get; } + + public static bool TryParse( + RequestContinuationToken continuationToken, + out AggregateContinuationToken aggregateContinuationToken) { - aggregateContinuationToken = default; - return false; - } + if (continuationToken == null) + { + throw new ArgumentNullException(nameof(continuationToken)); + } - aggregateContinuationToken = new AggregateContinuationToken( - singleGroupAggregatorContinuationToken: singleGroupAggregatorContinuationToken, - sourceContinuationToken: sourceContinuationToken); - return true; + if (!(continuationToken is CosmosElementRequestContinuationToken cosmosElementRequestContinuationToken)) + { + throw new ArgumentException($"Expected {nameof(CosmosElementRequestContinuationToken)} instead of: {continuationToken.GetType()}"); + } + + if (!(cosmosElementRequestContinuationToken.Value is CosmosObject rawAggregateContinuationToken)) + { + aggregateContinuationToken = default; + return false; + } + + if (!rawAggregateContinuationToken.TryGetValue( + ComputeAggregateDocumentQueryExecutionComponent.AggregationTokenName, + out CosmosElement singleGroupAggregatorContinuationToken)) + { + aggregateContinuationToken = default; + return false; + } + + if (!rawAggregateContinuationToken.TryGetValue( + ComputeAggregateDocumentQueryExecutionComponent.SourceTokenName, + out CosmosElement sourceContinuationToken)) + { + aggregateContinuationToken = default; + return false; + } + + aggregateContinuationToken = new AggregateContinuationToken( + singleGroupAggregatorContinuationToken: singleGroupAggregatorContinuationToken, + sourceContinuationToken: sourceContinuationToken); + return true; + } } } } diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/AverageAggregator.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/AverageAggregator.cs index f1f081838d..09d0a7c681 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/AverageAggregator.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/AverageAggregator.cs @@ -204,10 +204,16 @@ public void SerializeState(IJsonWriter jsonWriter) } jsonWriter.WriteObjectStart(); - jsonWriter.WriteFieldName(AverageInfo.SumName); - jsonWriter.WriteFloat64Value(this.Sum.Value); + + if (this.Sum.HasValue) + { + jsonWriter.WriteFieldName(AverageInfo.SumName); + jsonWriter.WriteFloat64Value(this.Sum.Value); + } + jsonWriter.WriteFieldName(AverageInfo.CountName); jsonWriter.WriteInt64Value(this.Count); + jsonWriter.WriteObjectEnd(); } diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/CountAggregator.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/CountAggregator.cs index 383b044c8a..9b05acf3f8 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/CountAggregator.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/CountAggregator.cs @@ -80,7 +80,7 @@ public static TryCatch TryCreate(RequestContinuationToken continuat } long partialCount; - if (continuationToken != null) + if (!continuationToken.IsNull) { if (!continuationToken.TryConvertToCosmosElement(out CosmosInt64 cosmosPartialCount)) { diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/MinMaxAggregator.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/MinMaxAggregator.cs index 856fe6ac9e..4fcb08cec2 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/MinMaxAggregator.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/MinMaxAggregator.cs @@ -145,8 +145,6 @@ public void SerializeState(IJsonWriter jsonWriter) throw new ArgumentNullException(nameof(jsonWriter)); } - jsonWriter.WriteObjectStart(); - MinMaxContinuationToken minMaxContinuationToken; if (this.globalMinMax == ItemComparer.MinValue) { @@ -165,7 +163,8 @@ public void SerializeState(IJsonWriter jsonWriter) minMaxContinuationToken = MinMaxContinuationToken.CreateValueContinuationToken(this.globalMinMax); } - MinMaxContinuationToken.ToCosmosElement(minMaxContinuationToken).WriteTo(jsonWriter); + CosmosElement minMaxContinuationTokenAsCosmosElement = MinMaxContinuationToken.ToCosmosElement(minMaxContinuationToken); + minMaxContinuationTokenAsCosmosElement.WriteTo(jsonWriter); } public static TryCatch TryCreateMinAggregator(RequestContinuationToken continuationToken) diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/MinMaxContinuationToken.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/MinMaxContinuationToken.cs index df8c199c39..65d003eb99 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/MinMaxContinuationToken.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/MinMaxContinuationToken.cs @@ -60,7 +60,7 @@ public static MinMaxContinuationToken CreateMaxValueContinuationToken() public static MinMaxContinuationToken CreateUndefinedValueContinuationToken() { - return new MinMaxContinuationToken(type: MinMaxContinuationTokenType.MaxValue, value: null); + return new MinMaxContinuationToken(type: MinMaxContinuationTokenType.Undefined, value: null); } public static MinMaxContinuationToken CreateValueContinuationToken(CosmosElement value) diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/SingleGroupAggregator.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/SingleGroupAggregator.cs index 0c10128397..29a7fbfb4d 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/SingleGroupAggregator.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/SingleGroupAggregator.cs @@ -234,7 +234,7 @@ public static TryCatch TryCreate( } else { - aliasContinuationToken = RequestContinuationToken.Create((CosmosElement)null); + aliasContinuationToken = CosmosElementRequestContinuationToken.Null; } TryCatch tryCreateAggregateValue = AggregateValue.TryCreate( @@ -417,13 +417,16 @@ public override void SerializeState(IJsonWriter jsonWriter) } jsonWriter.WriteObjectStart(); + jsonWriter.WriteFieldName(nameof(this.initialized)); jsonWriter.WriteBoolValue(this.initialized); + if (this.value != null) { jsonWriter.WriteFieldName(nameof(this.value)); this.value.WriteTo(jsonWriter); } + jsonWriter.WriteObjectEnd(); } @@ -436,7 +439,7 @@ public static TryCatch TryCreate(RequestContinuationToken contin CosmosElement value; bool initialized; - if (continuationToken.IsNull) + if (!continuationToken.IsNull) { if (!continuationToken.TryConvertToCosmosElement(out CosmosObject rawContinuationToken)) { 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 f1df7f4b01..47f3c3ce08 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 @@ -197,16 +197,16 @@ public static bool TryParse(CosmosElement value, out GroupByContinuationToken gr } if (!groupByContinuationTokenObject.TryGetValue( - nameof(GroupByContinuationToken.GroupingTableContinuationToken), - out CosmosString groupingTableContinuationToken)) + ComputeGroupByDocumentQueryExecutionComponent.GroupingTableContinuationTokenName, + out CosmosElement groupingTableContinuationToken)) { groupByContinuationToken = default; return false; } if (!groupByContinuationTokenObject.TryGetValue( - nameof(GroupByContinuationToken.SourceContinuationToken), - out CosmosString sourceContinuationToken)) + ComputeGroupByDocumentQueryExecutionComponent.SourceTokenName, + out CosmosElement sourceContinuationToken)) { groupByContinuationToken = default; return false; 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 77f9f4ca01..e5cac473ee 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/GroupBy/GroupByDocumentQueryExecutionComponent.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/GroupBy/GroupByDocumentQueryExecutionComponent.cs @@ -208,7 +208,7 @@ public void AddPayload(RewrittenGroupByProjection rewrittenGroupByProjection) this.groupByAliasToAggregateType, this.orderedAliases, this.hasSelectValue, - continuationToken: StringRequestContinuationToken.Null).Result; + continuationToken: CosmosElementRequestContinuationToken.Null).Result; this.table[groupByKeysHash] = singleGroupAggregator; } @@ -264,7 +264,7 @@ public void SerializeState(IJsonWriter jsonWriter) foreach (KeyValuePair kvp in this.table) { jsonWriter.WriteFieldName(kvp.Key.ToString()); - jsonWriter.WriteStringValue(kvp.Value.GetContinuationToken()); + kvp.Value.SerializeState(jsonWriter); } jsonWriter.WriteObjectEnd(); } @@ -301,18 +301,12 @@ public static TryCatch TryCreateFromContinuationToken( new MalformedContinuationTokenException($"Invalid GroupingTableContinuationToken")); } - if (!(value is CosmosString singleGroupAggregatorContinuationToken)) - { - return TryCatch.FromException( - new MalformedContinuationTokenException($"Invalid GroupingTableContinuationToken")); - } - TryCatch tryCreateSingleGroupAggregator = SingleGroupAggregator.TryCreate( EmptyAggregateOperators, groupByAliasToAggregateType, orderedAliases, hasSelectValue, - RequestContinuationToken.Create(singleGroupAggregatorContinuationToken.Value)); + RequestContinuationToken.Create(value)); if (tryCreateSingleGroupAggregator.Succeeded) { diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/SkipDocumentQueryExecutionComponent.Client.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/SkipDocumentQueryExecutionComponent.Client.cs index f83c6d456e..51fb9e1198 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/SkipDocumentQueryExecutionComponent.Client.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/SkipDocumentQueryExecutionComponent.Client.cs @@ -20,7 +20,7 @@ internal abstract partial class SkipDocumentQueryExecutionComponent : DocumentQu { private sealed class ClientSkipDocumentQueryExecutionComponent : SkipDocumentQueryExecutionComponent { - public ClientSkipDocumentQueryExecutionComponent(IDocumentQueryExecutionComponent source, long skipCount) + private ClientSkipDocumentQueryExecutionComponent(IDocumentQueryExecutionComponent source, long skipCount) : base(source, skipCount) { // Work is done in base constructor. diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/SkipDocumentQueryExecutionComponent.Compute.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/SkipDocumentQueryExecutionComponent.Compute.cs index 2d51a4a36a..8e27da989d 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/SkipDocumentQueryExecutionComponent.Compute.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/SkipDocumentQueryExecutionComponent.Compute.cs @@ -22,13 +22,13 @@ private sealed class ComputeSkipDocumentQueryExecutionComponent : SkipDocumentQu private const string SkipCountPropertyName = "SkipCount"; private const string SourceTokenPropertyName = "SourceToken"; - public ComputeSkipDocumentQueryExecutionComponent(IDocumentQueryExecutionComponent source, long skipCount) + private ComputeSkipDocumentQueryExecutionComponent(IDocumentQueryExecutionComponent source, long skipCount) : base(source, skipCount) { // Work is done in base constructor. } - public static async Task> TryCreateComputeAsync( + public static async Task> TryCreateAsync( int offsetCount, RequestContinuationToken continuationToken, Func>> tryCreateSourceAsync) @@ -68,7 +68,7 @@ public static async Task> TryCreateCo } return (await tryCreateSourceAsync(RequestContinuationToken.Create(offsetContinuationToken.SourceToken))) - .Try((source) => new ClientSkipDocumentQueryExecutionComponent( + .Try((source) => new ComputeSkipDocumentQueryExecutionComponent( source, offsetContinuationToken.Offset)); } 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 a9ae9c9680..331d27fc42 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/SkipDocumentQueryExecutionComponent.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/SkipDocumentQueryExecutionComponent.cs @@ -1,19 +1,14 @@ //------------------------------------------------------------ // Copyright (c) Microsoft Corporation. All rights reserved. //------------------------------------------------------------ + namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.SkipTake { 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.Json; 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.QueryClient; internal abstract partial class SkipDocumentQueryExecutionComponent : DocumentQueryExecutionComponentBase { @@ -47,7 +42,7 @@ public static Task> TryCreateAsync( break; case ExecutionEnvironment.Compute: - tryCreate = ComputeSkipDocumentQueryExecutionComponent.TryCreateComputeAsync( + tryCreate = ComputeSkipDocumentQueryExecutionComponent.TryCreateAsync( offsetCount, continuationToken, tryCreateSourceAsync); diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/TakeDocumentQueryExecutionComponent.Client.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/TakeDocumentQueryExecutionComponent.Client.cs index 539a8c499a..8c7bb9fdf0 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/TakeDocumentQueryExecutionComponent.Client.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/TakeDocumentQueryExecutionComponent.Client.cs @@ -145,11 +145,25 @@ public override async Task DrainAsync(int maxElements, Cancel this.takeCount -= takedDocuments.Count; string updatedContinuationToken; - if (sourcePage.DisallowContinuationTokenMessage == null) + if (!this.IsDone && (sourcePage.DisallowContinuationTokenMessage == null)) { - updatedContinuationToken = new LimitContinuationToken( - limit: this.takeCount, - sourceToken: sourcePage.ContinuationToken).ToString(); + switch (this.takeEnum) + { + 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}."); + } } else { @@ -157,13 +171,13 @@ public override async Task DrainAsync(int maxElements, Cancel } return QueryResponseCore.CreateSuccess( - result: takedDocuments, - continuationToken: updatedContinuationToken, - disallowContinuationTokenMessage: sourcePage.DisallowContinuationTokenMessage, - activityId: sourcePage.ActivityId, - requestCharge: sourcePage.RequestCharge, - diagnostics: sourcePage.Diagnostics, - responseLengthBytes: sourcePage.ResponseLengthBytes); + result: takedDocuments, + continuationToken: updatedContinuationToken, + disallowContinuationTokenMessage: sourcePage.DisallowContinuationTokenMessage, + activityId: sourcePage.ActivityId, + requestCharge: sourcePage.RequestCharge, + diagnostics: sourcePage.Diagnostics, + responseLengthBytes: sourcePage.ResponseLengthBytes); } public override void SerializeState(IJsonWriter jsonWriter) 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 586097aad5..291ed6a83b 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/OrderBy/CosmosOrderByItemQueryExecutionContext.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/OrderBy/CosmosOrderByItemQueryExecutionContext.cs @@ -184,10 +184,6 @@ public override void SerializeState(IJsonWriter jsonWriter) jsonWriter.WriteArrayEnd(); } - else - { - jsonWriter.WriteNullValue(); - } } public static async Task> TryCreateAsync( 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 058c7591fd..41bfb144d9 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/Parallel/CosmosParallelItemQueryExecutionContext.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/Parallel/CosmosParallelItemQueryExecutionContext.cs @@ -123,7 +123,7 @@ public override void SerializeState(IJsonWriter jsonWriter) foreach (ItemProducer activeItemProducer in activeItemProducers) { CompositeContinuationTokenRefStruct compositeToken = new CompositeContinuationTokenRefStruct( - backendContinuationToken: activeItemProducer.BackendContinuationToken, + backendContinuationToken: activeItemProducer.CurrentContinuationToken, range: new RangeRefStruct( min: activeItemProducer.PartitionKeyRange.MinInclusive, max: activeItemProducer.PartitionKeyRange.MaxExclusive)); @@ -132,10 +132,6 @@ public override void SerializeState(IJsonWriter jsonWriter) jsonWriter.WriteArrayEnd(); } - else - { - jsonWriter.WriteNullValue(); - } } public static async Task> TryCreateAsync( diff --git a/Microsoft.Azure.Cosmos/src/Query/v3Query/QueryIterator.cs b/Microsoft.Azure.Cosmos/src/Query/v3Query/QueryIterator.cs index 99293b92ff..231de1aa4f 100644 --- a/Microsoft.Azure.Cosmos/src/Query/v3Query/QueryIterator.cs +++ b/Microsoft.Azure.Cosmos/src/Query/v3Query/QueryIterator.cs @@ -72,7 +72,17 @@ public static QueryIterator Create( break; case ExecutionEnvironment.Compute: - requestContinuationToken = RequestContinuationToken.Create(CosmosElement.CreateFromBuffer(queryRequestOptions.BinaryContinuationToken)); + CosmosElement cosmosElement; + if (queryRequestOptions.BinaryContinuationToken.IsEmpty) + { + cosmosElement = null; + } + else + { + cosmosElement = CosmosElement.CreateFromBuffer(queryRequestOptions.BinaryContinuationToken); + } + + requestContinuationToken = RequestContinuationToken.Create(cosmosElement); break; default: diff --git a/Microsoft.Azure.Cosmos/src/Serializer/CosmosSerializerCore.cs b/Microsoft.Azure.Cosmos/src/Serializer/CosmosSerializerCore.cs index b898af2a9e..bc769e8616 100644 --- a/Microsoft.Azure.Cosmos/src/Serializer/CosmosSerializerCore.cs +++ b/Microsoft.Azure.Cosmos/src/Serializer/CosmosSerializerCore.cs @@ -135,11 +135,14 @@ private CosmosSerializer GetSerializer() } #if DEBUG + // This check is used to stop internal developers from deserializing an internal with the user's serialier that doesn't know how to materialize said type. string clientAssemblyName = typeof(DatabaseProperties).Assembly.GetName().Name; string directAssemblyName = typeof(Documents.PartitionKeyRange).Assembly.GetName().Name; string inputAssemblyName = inputType.Assembly.GetName().Name; - if (string.Equals(inputAssemblyName, clientAssemblyName) || - string.Equals(inputAssemblyName, directAssemblyName)) + bool inputIsClientOrDirect = string.Equals(inputAssemblyName, clientAssemblyName) || string.Equals(inputAssemblyName, directAssemblyName); + bool typeIsWhiteListed = inputType == typeof(Document); + + if (!typeIsWhiteListed && inputIsClientOrDirect) { throw new ArgumentException($"User serializer is being used for internal type:{inputType.FullName}."); } diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CrossPartitionQueryTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CrossPartitionQueryTests.cs index ed9810bf7a..328cac2a25 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CrossPartitionQueryTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CrossPartitionQueryTests.cs @@ -13,6 +13,7 @@ namespace Microsoft.Azure.Cosmos.SDK.EmulatorTests using System.IO; using System.Linq; using System.Net; + using System.Runtime.ExceptionServices; using System.Text; using System.Text.RegularExpressions; using System.Threading; @@ -21,9 +22,7 @@ namespace Microsoft.Azure.Cosmos.SDK.EmulatorTests using Microsoft.Azure.Cosmos.CosmosElements; using Microsoft.Azure.Cosmos.CosmosElements.Numbers; using Microsoft.Azure.Cosmos.Query.Core; - using Microsoft.Azure.Cosmos.Query.Core.ContinuationTokens; using Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.ItemProducers; - using Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.OrderBy; using Microsoft.Azure.Cosmos.Query.Core.Metrics; using Microsoft.Azure.Cosmos.Query.Core.QueryClient; using Microsoft.Azure.Cosmos.Query.Core.QueryPlan; @@ -31,12 +30,10 @@ namespace Microsoft.Azure.Cosmos.SDK.EmulatorTests using Microsoft.Azure.Documents; using Microsoft.Azure.Documents.Routing; using Microsoft.VisualStudio.TestTools.UnitTesting; - using Moq; using Newtonsoft.Json; using Newtonsoft.Json.Converters; using Newtonsoft.Json.Linq; - using Query; - using UInt128 = Query.Core.UInt128; + using UInt128 = Microsoft.Azure.Cosmos.Query.Core.UInt128; /// /// Tests for CrossPartitionQueryTests. @@ -457,103 +454,90 @@ private async Task CreateIngestQueryDelete( string partitionKey = "/id", Cosmos.IndexingPolicy indexingPolicy = null) { - int retryCount = 1; - AggregateException exceptionHistory = new AggregateException(); - while (retryCount-- > 0) + try { - try + List>> collectionsAndDocuments = new List>>(); + foreach (CollectionTypes collectionType in Enum.GetValues(collectionTypes.GetType()).Cast().Where(collectionTypes.HasFlag)) { - List>> collectionsAndDocuments = new List>>(); - foreach (CollectionTypes collectionType in Enum.GetValues(collectionTypes.GetType()).Cast().Where(collectionTypes.HasFlag)) + if (collectionType == CollectionTypes.None) { - if (collectionType == CollectionTypes.None) - { - continue; - } - - Task>> createContainerTask; - switch (collectionType) - { - case CollectionTypes.NonPartitioned: - createContainerTask = this.CreateNonPartitionedContainerAndIngestDocuments( - documents, - indexingPolicy); - break; - - case CollectionTypes.SinglePartition: - createContainerTask = this.CreateSinglePartitionContainerAndIngestDocuments( - documents, - partitionKey, - indexingPolicy); - break; - - case CollectionTypes.MultiPartition: - createContainerTask = this.CreateMultiPartitionContainerAndIngestDocuments( - documents, - partitionKey, - indexingPolicy); - break; - - default: - throw new ArgumentException($"Unknown {nameof(CollectionTypes)} : {collectionType}"); - } - - collectionsAndDocuments.Add(await createContainerTask); + continue; } - List cosmosClients = new List(); - foreach (ConnectionModes connectionMode in Enum.GetValues(connectionModes.GetType()).Cast().Where(connectionModes.HasFlag)) + Task>> createContainerTask; + switch (collectionType) { - if (connectionMode == ConnectionModes.None) - { - continue; - } - - ConnectionMode targetConnectionMode = GetTargetConnectionMode(connectionMode); - CosmosClient cosmosClient = cosmosClientFactory(targetConnectionMode); - - Assert.AreEqual( - targetConnectionMode, - cosmosClient.ClientOptions.ConnectionMode, - "Test setup: Invalid connection policy applied to CosmosClient"); - cosmosClients.Add(cosmosClient); + case CollectionTypes.NonPartitioned: + createContainerTask = this.CreateNonPartitionedContainerAndIngestDocuments( + documents, + indexingPolicy); + break; + + case CollectionTypes.SinglePartition: + createContainerTask = this.CreateSinglePartitionContainerAndIngestDocuments( + documents, + partitionKey, + indexingPolicy); + break; + + case CollectionTypes.MultiPartition: + createContainerTask = this.CreateMultiPartitionContainerAndIngestDocuments( + documents, + partitionKey, + indexingPolicy); + break; + + default: + throw new ArgumentException($"Unknown {nameof(CollectionTypes)} : {collectionType}"); } - List queryTasks = new List(); - foreach (CosmosClient cosmosClient in cosmosClients) + collectionsAndDocuments.Add(await createContainerTask); + } + + List cosmosClients = new List(); + foreach (ConnectionModes connectionMode in Enum.GetValues(connectionModes.GetType()).Cast().Where(connectionModes.HasFlag)) + { + if (connectionMode == ConnectionModes.None) { - foreach (Tuple> containerAndDocuments in collectionsAndDocuments) - { - Container container = cosmosClient.GetContainer(((ContainerCore)(ContainerInlineCore)containerAndDocuments.Item1).Database.Id, containerAndDocuments.Item1.Id); - Task queryTask = Task.Run(() => query(container, containerAndDocuments.Item2, testArgs)); - queryTasks.Add(queryTask); - } + continue; } - await Task.WhenAll(queryTasks); + ConnectionMode targetConnectionMode = GetTargetConnectionMode(connectionMode); + CosmosClient cosmosClient = cosmosClientFactory(targetConnectionMode); - List> deleteContainerTasks = new List>(); - foreach (Container container in collectionsAndDocuments.Select(tuple => tuple.Item1)) + Assert.AreEqual( + targetConnectionMode, + cosmosClient.ClientOptions.ConnectionMode, + "Test setup: Invalid connection policy applied to CosmosClient"); + cosmosClients.Add(cosmosClient); + } + + List queryTasks = new List(); + foreach (CosmosClient cosmosClient in cosmosClients) + { + foreach (Tuple> containerAndDocuments in collectionsAndDocuments) { - deleteContainerTasks.Add(container.DeleteContainerAsync()); + Container container = cosmosClient.GetContainer(((ContainerCore)(ContainerInlineCore)containerAndDocuments.Item1).Database.Id, containerAndDocuments.Item1.Id); + Task queryTask = Task.Run(() => query(container, containerAndDocuments.Item2, testArgs)); + queryTasks.Add(queryTask); } + } - await Task.WhenAll(deleteContainerTasks); + await Task.WhenAll(queryTasks); - // If you made it here then it's all good - break; - } - catch (Exception ex) when (ex.GetType() != typeof(AssertFailedException)) + List> deleteContainerTasks = new List>(); + foreach (Container container in collectionsAndDocuments.Select(tuple => tuple.Item1)) { - List previousExceptions = exceptionHistory.InnerExceptions.ToList(); - previousExceptions.Add(ex); - exceptionHistory = new AggregateException(previousExceptions); + deleteContainerTasks.Add(container.DeleteContainerAsync()); } - } - if (exceptionHistory.InnerExceptions.Count > 0) + await Task.WhenAll(deleteContainerTasks); + } + catch (Exception ex) when (ex.GetType() != typeof(AssertFailedException)) { - throw exceptionHistory; + while (ex.InnerException != null) ex = ex.InnerException; + + ExceptionDispatchInfo.Capture(ex).Throw(); } } @@ -638,6 +622,11 @@ private static async Task> QueryWithSerializeState( Json.IJsonWriter jsonWriter = Json.JsonWriter.Create(Json.JsonSerializationFormat.Binary); itemQuery.SerializeState(jsonWriter); continuationToken = jsonWriter.GetResult(); + + if (!continuationToken.IsEmpty) + { + string stringContinuationToken = CosmosElement.CreateFromBuffer(continuationToken).ToString(); + } } catch (CosmosException cosmosException) when (cosmosException.StatusCode == (HttpStatusCode)429) { @@ -1933,7 +1922,6 @@ private async Task TestQueryCrossPartitionAggregateFunctionsEmptyPartitionsHelpe } } - [TestCategory("Quarantine")] [TestMethod] public async Task TestQueryCrossPartitionAggregateFunctionsWithMixedTypes() { @@ -2100,15 +2088,17 @@ FROM c } string filename = $"CrossPartitionQueryTests.AggregateMixedTypes"; - string outputPath = $"{filename}_output.xml"; string baselinePath = $"{filename}_baseline.xml"; + XmlWriterSettings settings = new XmlWriterSettings() { OmitXmlDeclaration = true, Indent = true, NewLineOnAttributes = true, }; - using (XmlWriter writer = XmlWriter.Create(outputPath, settings)) + + StringBuilder builder = new StringBuilder(); + using (XmlWriter writer = XmlWriter.Create(builder, settings)) { writer.WriteStartDocument(); writer.WriteStartElement("Results"); @@ -2147,7 +2137,7 @@ FROM c Regex r = new Regex(">\\s+"); string normalizedBaseline = r.Replace(File.ReadAllText(baselinePath), ">"); - string normalizedOutput = r.Replace(File.ReadAllText(outputPath), ">"); + string normalizedOutput = r.Replace(builder.ToString(), ">"); Assert.AreEqual(normalizedBaseline, normalizedOutput); } From d9455b14cef0a33cdf9040cdc3b0230eb07bbc06 Mon Sep 17 00:00:00 2001 From: Brandon Chong Date: Fri, 14 Feb 2020 13:08:03 -0800 Subject: [PATCH 13/28] fixed more bugs --- ...nctDocumentQueryExecutionComponent.Client.cs | 17 +++++++++++++---- .../CosmosQueryUnitTests.cs | 12 ++++++------ .../Query/QueryPipelineMockTests.cs | 2 +- 3 files changed, 20 insertions(+), 11 deletions(-) diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctDocumentQueryExecutionComponent.Client.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctDocumentQueryExecutionComponent.Client.cs index a2c29e3941..605a928ff7 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctDocumentQueryExecutionComponent.Client.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctDocumentQueryExecutionComponent.Client.cs @@ -126,12 +126,21 @@ public override async Task DrainAsync(int maxElements, Cancel QueryResponseCore queryResponseCore; if (this.distinctQueryType == DistinctQueryType.Ordered) { - DistinctContinuationToken distinctContinuationToken = new DistinctContinuationToken( - sourceToken: sourceResponse.ContinuationToken, - distinctMapToken: this.distinctMap.GetContinuationToken()); + string updatedContinuationToken; + if (this.IsDone) + { + updatedContinuationToken = null; + } + else + { + updatedContinuationToken = new DistinctContinuationToken( + sourceToken: sourceResponse.ContinuationToken, + distinctMapToken: this.distinctMap.GetContinuationToken()).ToString(); + } + queryResponseCore = QueryResponseCore.CreateSuccess( result: distinctResults, - continuationToken: distinctContinuationToken.ToString(), + continuationToken: updatedContinuationToken, disallowContinuationTokenMessage: null, activityId: sourceResponse.ActivityId, requestCharge: sourceResponse.RequestCharge, 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 7c2d86b8ad..1783bcb0d5 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/CosmosQueryUnitTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/CosmosQueryUnitTests.cs @@ -256,7 +256,7 @@ public async Task TestCosmosQueryPartitionKeyDefinition() CosmosQueryExecutionContextFactory.InputParameters inputParameters = new CosmosQueryExecutionContextFactory.InputParameters( sqlQuerySpec: sqlQuerySpec, - initialUserContinuationToken: null, + initialUserContinuationToken: StringRequestContinuationToken.Null, maxConcurrency: queryRequestOptions?.MaxConcurrency, maxItemCount: queryRequestOptions?.MaxItemCount, maxBufferedItemCount: queryRequestOptions?.MaxBufferedItemCount, @@ -310,31 +310,31 @@ public async Task TestCosmosQueryPartitionKeyDefinition() }, new List() { "test" }, false, - null, + StringRequestContinuationToken.Null, func)).Result); components.Add((await DistinctDocumentQueryExecutionComponent.TryCreateAsync( ExecutionEnvironment.Client, - null, + StringRequestContinuationToken.Null, func, DistinctQueryType.Ordered)).Result); components.Add((await SkipDocumentQueryExecutionComponent.TryCreateAsync( ExecutionEnvironment.Client, 5, - null, + StringRequestContinuationToken.Null, func)).Result); components.Add((await TakeDocumentQueryExecutionComponent.TryCreateLimitDocumentQueryExecutionComponentAsync( ExecutionEnvironment.Client, 5, - null, + StringRequestContinuationToken.Null, func)).Result); components.Add((await TakeDocumentQueryExecutionComponent.TryCreateTopDocumentQueryExecutionComponentAsync( ExecutionEnvironment.Client, 5, - null, + StringRequestContinuationToken.Null, func)).Result); return (components, response); diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/QueryPipelineMockTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/QueryPipelineMockTests.cs index 4c5b3f0eab..0e872d7dfc 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/QueryPipelineMockTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/QueryPipelineMockTests.cs @@ -523,7 +523,7 @@ public async Task TestNegativeDistinctComponentCreation() { TryCatch tryCreateWhenSourceFails = await DistinctDocumentQueryExecutionComponent.TryCreateAsync( ExecutionEnvironment.Client, - null, + StringRequestContinuationToken.Null, FailToCreateSource, DistinctQueryType.Ordered); From ca3d0cf6f0536d5c58b6190ed50957fb87427b82 Mon Sep 17 00:00:00 2001 From: Brandon Chong Date: Wed, 26 Feb 2020 20:56:13 -0800 Subject: [PATCH 14/28] about to gut out RequestContinuationToken and just wire through CosmosElement --- .../src/FeedIteratorCore.cs | 25 +++--- .../CompositeContinuationTokenRefStruct.cs | 50 ----------- .../CosmosElementRequestContinuationToken.cs | 28 ------ .../OrderByContinuationTokenRefStruct.cs | 90 ------------------- .../RequestContinuationToken.cs | 32 ------- .../StringRequestContinuationToken.cs | 27 ------ .../DistinctMap.OrderedDistinctMap.cs | 2 +- .../CosmosOrderByItemQueryExecutionContext.cs | 17 ++-- ...CosmosParallelItemQueryExecutionContext.cs | 15 ++-- .../src/Query/Core/UInt128.cs | 4 +- .../src/Resource/Offer/CosmosOffers.cs | 3 +- .../CrossPartitionQueryTests.cs | 11 ++- 12 files changed, 42 insertions(+), 262 deletions(-) delete mode 100644 Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/CompositeContinuationTokenRefStruct.cs delete mode 100644 Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/CosmosElementRequestContinuationToken.cs delete mode 100644 Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/OrderByContinuationTokenRefStruct.cs delete mode 100644 Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/RequestContinuationToken.cs delete mode 100644 Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/StringRequestContinuationToken.cs diff --git a/Microsoft.Azure.Cosmos/src/FeedIteratorCore.cs b/Microsoft.Azure.Cosmos/src/FeedIteratorCore.cs index 60b809f716..8fc5491b14 100644 --- a/Microsoft.Azure.Cosmos/src/FeedIteratorCore.cs +++ b/Microsoft.Azure.Cosmos/src/FeedIteratorCore.cs @@ -18,31 +18,28 @@ namespace Microsoft.Azure.Cosmos /// /// Cosmos feed stream iterator. This is used to get the query responses with a Stream content /// - internal class FeedIteratorCore : FeedIteratorInternal + internal sealed class FeedIteratorCore : FeedIteratorInternal { private readonly CosmosClientContext clientContext; private readonly Uri resourceLink; private readonly ResourceType resourceType; private readonly SqlQuerySpec querySpec; - private readonly bool usePropertySerializer; private bool hasMoreResultsInternal; - internal FeedIteratorCore( + public FeedIteratorCore( CosmosClientContext clientContext, Uri resourceLink, ResourceType resourceType, QueryDefinition queryDefinition, string continuationToken, - QueryRequestOptions options, - bool usePropertySerializer = false) + QueryRequestOptions options) { this.resourceLink = resourceLink; this.clientContext = clientContext; this.resourceType = resourceType; this.querySpec = queryDefinition?.ToSqlQuerySpec(); - this.continuationToken = continuationToken; + this.ContinuationToken = continuationToken; this.requestOptions = options; - this.usePropertySerializer = usePropertySerializer; this.hasMoreResultsInternal = true; } @@ -51,12 +48,12 @@ internal FeedIteratorCore( /// /// The query options for the result set /// - protected QueryRequestOptions requestOptions { get; } + public QueryRequestOptions requestOptions { get; } /// /// The Continuation Token /// - protected string continuationToken { get; set; } + public string ContinuationToken { get; set; } /// /// Get the next set of results from the cosmos service @@ -65,6 +62,8 @@ internal FeedIteratorCore( /// A query response from cosmos service public override async Task ReadNextAsync(CancellationToken cancellationToken = default) { + cancellationToken.ThrowIfCancellationRequested(); + Stream stream = null; OperationType operation = OperationType.ReadFeed; if (this.querySpec != null) @@ -83,7 +82,7 @@ public override async Task ReadNextAsync(CancellationToken canc streamPayload: stream, requestEnricher: request => { - QueryRequestOptions.FillContinuationToken(request, this.continuationToken); + QueryRequestOptions.FillContinuationToken(request, this.ContinuationToken); if (this.querySpec != null) { request.Headers.Add(HttpConstants.HttpHeaders.ContentType, MediaTypes.QueryJson); @@ -93,8 +92,8 @@ public override async Task ReadNextAsync(CancellationToken canc diagnosticsScope: null, cancellationToken: cancellationToken); - this.continuationToken = response.Headers.ContinuationToken; - this.hasMoreResultsInternal = GetHasMoreResults(this.continuationToken, response.StatusCode); + this.ContinuationToken = response.Headers.ContinuationToken; + this.hasMoreResultsInternal = GetHasMoreResults(this.ContinuationToken, response.StatusCode); return response; } @@ -113,7 +112,7 @@ internal static bool GetHasMoreResults(string continuationToken, HttpStatusCode public override void SerializeState(IJsonWriter jsonWriter) { - jsonWriter.WriteJsonFragment(Encoding.UTF8.GetBytes(this.continuationToken)); + throw new NotImplementedException(); } } diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/CompositeContinuationTokenRefStruct.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/CompositeContinuationTokenRefStruct.cs deleted file mode 100644 index 2d86a99b90..0000000000 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/CompositeContinuationTokenRefStruct.cs +++ /dev/null @@ -1,50 +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 CompositeContinuationTokenRefStruct - { - private const string TokenProperytName = "token"; - private const string RangePropertyName = "range"; - - public CompositeContinuationTokenRefStruct(string backendContinuationToken, RangeRefStruct range) - { - this.BackendContinuationToken = backendContinuationToken; - this.Range = range; - } - - public string BackendContinuationToken { get; } - - public RangeRefStruct Range { get; } - - public void WriteTo(IJsonWriter jsonWriter) - { - if (jsonWriter == null) - { - throw new ArgumentNullException(nameof(jsonWriter)); - } - - jsonWriter.WriteObjectStart(); - - jsonWriter.WriteFieldName(CompositeContinuationTokenRefStruct.TokenProperytName); - if (this.BackendContinuationToken != null) - { - jsonWriter.WriteStringValue(this.BackendContinuationToken); - } - else - { - jsonWriter.WriteNullValue(); - } - - jsonWriter.WriteFieldName(CompositeContinuationTokenRefStruct.RangePropertyName); - this.Range.WriteTo(jsonWriter); - - jsonWriter.WriteObjectEnd(); - } - } -} diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/CosmosElementRequestContinuationToken.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/CosmosElementRequestContinuationToken.cs deleted file mode 100644 index 9f6a74ce37..0000000000 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/CosmosElementRequestContinuationToken.cs +++ /dev/null @@ -1,28 +0,0 @@ -// ------------------------------------------------------------ -// Copyright (c) Microsoft Corporation. All rights reserved. -// ------------------------------------------------------------ - -namespace Microsoft.Azure.Cosmos.Query.Core.ContinuationTokens -{ - using Microsoft.Azure.Cosmos.CosmosElements; - - internal sealed class CosmosElementRequestContinuationToken : RequestContinuationToken - { - public static readonly CosmosElementRequestContinuationToken Null = new CosmosElementRequestContinuationToken(null); - - public CosmosElementRequestContinuationToken(CosmosElement continuationToken) - { - this.Value = continuationToken; - } - - public CosmosElement Value { get; } - - public override bool IsNull => this.Value == null; - - public override bool TryConvertToCosmosElement(out TCosmosElement cosmosElement) - { - cosmosElement = (TCosmosElement)this.Value; - return true; - } - } -} diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/OrderByContinuationTokenRefStruct.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/OrderByContinuationTokenRefStruct.cs deleted file mode 100644 index 1a74d4d420..0000000000 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/OrderByContinuationTokenRefStruct.cs +++ /dev/null @@ -1,90 +0,0 @@ -// ------------------------------------------------------------ -// Copyright (c) Microsoft Corporation. All rights reserved. -// ------------------------------------------------------------ - -namespace Microsoft.Azure.Cosmos.Query.Core.ContinuationTokens -{ - using System; - using System.Collections.Generic; - using Microsoft.Azure.Cosmos.Json; - using Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.OrderBy; - - internal ref struct OrderByContinuationTokenRefStruct - { - private const string CompositeTokenPropertyName = "compositeToken"; - private const string OrderByItemsPropertyName = "orderByItems"; - private const string RidPropertyName = "rid"; - private const string SkipCountPropertyName = "skipCount"; - private const string FilterPropertyName = "filter"; - private const string ItemPropertyName = "item"; - - public OrderByContinuationTokenRefStruct( - CompositeContinuationTokenRefStruct compositeContinuationTokenRefStruct, - IReadOnlyList orderByItems, - string rid, - int skipCount, - string filter) - { - this.CompositeContinuationToken = compositeContinuationTokenRefStruct; - this.OrderByItems = orderByItems; - this.Rid = rid; - this.SkipCount = skipCount; - this.Filter = filter; - } - - public CompositeContinuationTokenRefStruct CompositeContinuationToken { get; } - public IReadOnlyList OrderByItems { get; } - public string Rid { get; } - public int SkipCount { get; } - public string Filter { get; } - - public void WriteTo(IJsonWriter jsonWriter) - { - if (jsonWriter == null) - { - throw new ArgumentNullException(nameof(jsonWriter)); - } - - jsonWriter.WriteObjectStart(); - - jsonWriter.WriteFieldName(OrderByContinuationTokenRefStruct.CompositeTokenPropertyName); - this.CompositeContinuationToken.WriteTo(jsonWriter); - - jsonWriter.WriteFieldName(OrderByContinuationTokenRefStruct.OrderByItemsPropertyName); - jsonWriter.WriteArrayStart(); - - foreach (OrderByItem orderByItem in this.OrderByItems) - { - jsonWriter.WriteObjectStart(); - - if (orderByItem.Item != null) - { - jsonWriter.WriteFieldName(OrderByContinuationTokenRefStruct.ItemPropertyName); - orderByItem.Item.WriteTo(jsonWriter); - } - - jsonWriter.WriteObjectEnd(); - } - - jsonWriter.WriteArrayEnd(); - - jsonWriter.WriteFieldName(OrderByContinuationTokenRefStruct.RidPropertyName); - jsonWriter.WriteStringValue(this.Rid); - - jsonWriter.WriteFieldName(OrderByContinuationTokenRefStruct.SkipCountPropertyName); - jsonWriter.WriteNumberValue(this.SkipCount); - - jsonWriter.WriteFieldName(OrderByContinuationTokenRefStruct.FilterPropertyName); - if (this.Filter != null) - { - jsonWriter.WriteStringValue(this.Filter); - } - else - { - jsonWriter.WriteNullValue(); - } - - jsonWriter.WriteObjectEnd(); - } - } -} diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/RequestContinuationToken.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/RequestContinuationToken.cs deleted file mode 100644 index c18e9d6edf..0000000000 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/RequestContinuationToken.cs +++ /dev/null @@ -1,32 +0,0 @@ -// ------------------------------------------------------------ -// Copyright (c) Microsoft Corporation. All rights reserved. -// ------------------------------------------------------------ - -namespace Microsoft.Azure.Cosmos.Query.Core.ContinuationTokens -{ - using System; - using Microsoft.Azure.Cosmos.CosmosElements; - using Microsoft.Azure.Cosmos.Json; - - /// - /// Base class that all continuation tokens will follow. - /// This serves as an adapter pattern, so that all different types of continuation tokens can have a common interface. - /// - internal abstract class RequestContinuationToken - { - public static RequestContinuationToken Create(string continuationToken) - { - return new StringRequestContinuationToken(continuationToken); - } - - public static RequestContinuationToken Create(CosmosElement continuationToken) - { - return new CosmosElementRequestContinuationToken(continuationToken); - } - - public abstract bool IsNull { get; } - - public abstract bool TryConvertToCosmosElement(out TCosmosElement cosmosElement) - where TCosmosElement : CosmosElement; - } -} diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/StringRequestContinuationToken.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/StringRequestContinuationToken.cs deleted file mode 100644 index 7a3588fb86..0000000000 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/StringRequestContinuationToken.cs +++ /dev/null @@ -1,27 +0,0 @@ -// ------------------------------------------------------------ -// Copyright (c) Microsoft Corporation. All rights reserved. -// ------------------------------------------------------------ - -namespace Microsoft.Azure.Cosmos.Query.Core.ContinuationTokens -{ - using Microsoft.Azure.Cosmos.CosmosElements; - - internal sealed class StringRequestContinuationToken : RequestContinuationToken - { - public static readonly StringRequestContinuationToken Null = new StringRequestContinuationToken(continuationToken: null); - - public StringRequestContinuationToken(string continuationToken) - { - this.Value = continuationToken; - } - - public string Value { get; } - - public override bool IsNull => string.IsNullOrWhiteSpace(this.Value); - - public override bool TryConvertToCosmosElement(out TCosmosElement cosmosElement) - { - return CosmosElement.TryParse(this.Value, out cosmosElement); - } - } -} diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctMap.OrderedDistinctMap.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctMap.OrderedDistinctMap.cs index 264c21067b..5da0801c67 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctMap.OrderedDistinctMap.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctMap.OrderedDistinctMap.cs @@ -113,7 +113,7 @@ public static TryCatch TryCreate(RequestContinuationToken requestCo $"Malformed {nameof(OrderedDistinctMap)} continuation token: {cosmosElementRequestContinuationToken.Value}.")); } - if (!UInt128.TryParse(cosmosBinary.Value.Span, out lastHash)) + if (!UInt128.TryCreateFromByteArray(cosmosBinary.Value.Span, out lastHash)) { return TryCatch.FromException( new MalformedContinuationTokenException( 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 291ed6a83b..52b514a56c 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/OrderBy/CosmosOrderByItemQueryExecutionContext.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/OrderBy/CosmosOrderByItemQueryExecutionContext.cs @@ -169,17 +169,22 @@ public override void SerializeState(IJsonWriter jsonWriter) foreach (ItemProducer activeItemProducer in activeItemProducers) { OrderByQueryResult orderByQueryResult = new OrderByQueryResult(activeItemProducer.Current); - OrderByContinuationTokenRefStruct orderByContinuationToken = new OrderByContinuationTokenRefStruct( - compositeContinuationTokenRefStruct: new CompositeContinuationTokenRefStruct( - backendContinuationToken: activeItemProducer.PreviousContinuationToken, - range: new RangeRefStruct( + OrderByContinuationToken orderByContinuationToken = new OrderByContinuationToken( + compositeContinuationToken: new CompositeContinuationToken() + { + Token = activeItemProducer.PreviousContinuationToken, + Range = new Documents.Routing.Range( min: activeItemProducer.PartitionKeyRange.MinInclusive, - max: activeItemProducer.PartitionKeyRange.MaxExclusive)), + 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); - orderByContinuationToken.WriteTo(jsonWriter); + + OrderByContinuationToken.ToCosmosElement(orderByContinuationToken).WriteTo(jsonWriter); } jsonWriter.WriteArrayEnd(); 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 41bfb144d9..b558ae2be9 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/Parallel/CosmosParallelItemQueryExecutionContext.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/Parallel/CosmosParallelItemQueryExecutionContext.cs @@ -122,12 +122,17 @@ public override void SerializeState(IJsonWriter jsonWriter) foreach (ItemProducer activeItemProducer in activeItemProducers) { - CompositeContinuationTokenRefStruct compositeToken = new CompositeContinuationTokenRefStruct( - backendContinuationToken: activeItemProducer.CurrentContinuationToken, - range: new RangeRefStruct( + CompositeContinuationToken compositeToken = new CompositeContinuationToken() + { + Token = activeItemProducer.CurrentContinuationToken, + Range = new Documents.Routing.Range( min: activeItemProducer.PartitionKeyRange.MinInclusive, - max: activeItemProducer.PartitionKeyRange.MaxExclusive)); - compositeToken.WriteTo(jsonWriter); + max: activeItemProducer.PartitionKeyRange.MaxExclusive, + isMinInclusive: false, + isMaxInclusive: true) + }; + + CompositeContinuationToken.ToCosmosElement(compositeToken).WriteTo(jsonWriter); } jsonWriter.WriteArrayEnd(); diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/UInt128.cs b/Microsoft.Azure.Cosmos/src/Query/Core/UInt128.cs index eff9ed0399..c4dfc86b5c 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/UInt128.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/UInt128.cs @@ -235,7 +235,7 @@ public static UInt128 Create(ulong low, ulong high) public static UInt128 FromByteArray(ReadOnlySpan buffer) { - if (!UInt128.TryParse(buffer, out UInt128 value)) + if (!UInt128.TryCreateFromByteArray(buffer, out UInt128 value)) { throw new FormatException($"Malformed buffer"); } @@ -400,7 +400,7 @@ public static bool TryParse(string value, out UInt128 uInt128) return true; } - public static bool TryParse(ReadOnlySpan buffer, out UInt128 value) + public static bool TryCreateFromByteArray(ReadOnlySpan buffer, out UInt128 value) { if (buffer.Length < UInt128.Length) { diff --git a/Microsoft.Azure.Cosmos/src/Resource/Offer/CosmosOffers.cs b/Microsoft.Azure.Cosmos/src/Resource/Offer/CosmosOffers.cs index 34bc33629a..aec29dc45f 100644 --- a/Microsoft.Azure.Cosmos/src/Resource/Offer/CosmosOffers.cs +++ b/Microsoft.Azure.Cosmos/src/Resource/Offer/CosmosOffers.cs @@ -185,8 +185,7 @@ internal virtual FeedIterator GetOfferQueryStreamIterator( resourceType: ResourceType.Offer, queryDefinition: queryDefinition, continuationToken: continuationToken, - options: requestOptions, - usePropertySerializer: true); + options: requestOptions); } private async Task SingleOrDefaultAsync( diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CrossPartitionQueryTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CrossPartitionQueryTests.cs index 328cac2a25..0fa3639cc7 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CrossPartitionQueryTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CrossPartitionQueryTests.cs @@ -33,7 +33,6 @@ namespace Microsoft.Azure.Cosmos.SDK.EmulatorTests using Newtonsoft.Json; using Newtonsoft.Json.Converters; using Newtonsoft.Json.Linq; - using UInt128 = Microsoft.Azure.Cosmos.Query.Core.UInt128; /// /// Tests for CrossPartitionQueryTests. @@ -2571,7 +2570,7 @@ private async Task TestQueryDistinct(Container container, IEnumerable FeedResponse cosmosQueryResponse = await documentQueryWithoutDistinct.ReadNextAsync(); foreach (JToken document in cosmosQueryResponse) { - if (documentsSeen.Add(document, out UInt128 hash)) + if (documentsSeen.Add(document, out Cosmos.Query.Core.UInt128 hash)) { documentsFromWithoutDistinct.Add(document); } @@ -2629,7 +2628,7 @@ private async Task TestQueryDistinct(Container container, IEnumerable }, QueryDrainingMode.ContinuationToken | QueryDrainingMode.HoldState); documentsFromWithoutDistinct = documentsFromWithoutDistinct - .Where(document => documentsSeen.Add(document, out UInt128 hash)) + .Where(document => documentsSeen.Add(document, out Cosmos.Query.Core.UInt128 hash)) .ToList(); foreach (int pageSize in new int[] { 1, 10, 100 }) @@ -2677,7 +2676,7 @@ private async Task TestQueryDistinct(Container container, IEnumerable }, QueryDrainingMode.HoldState | QueryDrainingMode.SerializeState); documentsFromWithoutDistinct = documentsFromWithoutDistinct - .Where(document => documentsSeen.Add(document, out UInt128 hash)) + .Where(document => documentsSeen.Add(document, out Cosmos.Query.Core.UInt128 hash)) .ToList(); foreach (int pageSize in new int[] { 1, 10, 100 }) @@ -4742,9 +4741,9 @@ internal sealed class MockDistinctMap // >> True private readonly HashSet jTokenSet = new HashSet(JsonTokenEqualityComparer.Value); - public bool Add(JToken jToken, out UInt128 hash) + public bool Add(JToken jToken, out Cosmos.Query.Core.UInt128 hash) { - hash = default(UInt128); + hash = default(Cosmos.Query.Core.UInt128); return this.jTokenSet.Add(jToken); } } From fb031f3537778a8e4670a45b82f68d0201cbbd01 Mon Sep 17 00:00:00 2001 From: Brandon Chong Date: Thu, 27 Feb 2020 17:47:03 -0800 Subject: [PATCH 15/28] made input a cosmos element --- .../PipelineContinuationToken.cs | 38 +++---- .../PipelineContinuationTokenV0.cs | 19 ++-- .../PipelineContinuationTokenV1.cs | 35 ++---- .../PipelineContinuationTokenV1_1.cs | 39 ++----- ...eDocumentQueryExecutionComponent.Client.cs | 6 +- ...DocumentQueryExecutionComponent.Compute.cs | 36 +++--- ...ggregateDocumentQueryExecutionComponent.cs | 4 +- .../Aggregators/AverageAggregator.cs | 60 +++++----- .../Aggregate/Aggregators/CountAggregator.cs | 13 +-- .../Aggregate/Aggregators/MinMaxAggregator.cs | 21 +--- .../Aggregators/SingleGroupAggregator.cs | 38 +++---- .../Aggregate/Aggregators/SumAggregator.cs | 8 +- ...tDocumentQueryExecutionComponent.Client.cs | 69 ++++++----- ...DocumentQueryExecutionComponent.Compute.cs | 24 ++-- ...DistinctDocumentQueryExecutionComponent.cs | 4 +- .../DistinctMap.OrderedDistinctMap.cs | 28 ++--- .../DistinctMap.UnorderedDistinctMap.cs | 17 +-- .../Distinct/DistinctMap.cs | 2 +- ...yDocumentQueryExecutionComponent.Client.cs | 4 +- ...DocumentQueryExecutionComponent.Compute.cs | 15 +-- .../GroupByDocumentQueryExecutionComponent.cs | 8 +- ...pDocumentQueryExecutionComponent.Client.cs | 32 ++++-- ...DocumentQueryExecutionComponent.Compute.cs | 16 +-- .../SkipDocumentQueryExecutionComponent.cs | 5 +- ...eDocumentQueryExecutionComponent.Client.cs | 74 ++++++------ ...DocumentQueryExecutionComponent.Compute.cs | 25 ++-- .../TakeDocumentQueryExecutionComponent.cs | 8 +- .../CosmosQueryExecutionContextFactory.cs | 21 ++-- .../CosmosOrderByItemQueryExecutionContext.cs | 39 +------ ...CosmosParallelItemQueryExecutionContext.cs | 107 +++++------------- .../PipelinedDocumentQueryExecutionContext.cs | 22 ++-- .../QueryExecutionContextWithException.cs | 47 ++++++++ .../src/Query/v3Query/QueryIterator.cs | 56 +++++---- .../src/RequestOptions/QueryRequestOptions.cs | 5 +- .../CrossPartitionQueryTests.cs | 16 ++- .../CosmosQueryUnitTests.cs | 18 +-- .../PipelineContinuationTokenTests.cs | 11 +- .../Query/QueryPipelineMockTests.cs | 16 +-- 38 files changed, 426 insertions(+), 580 deletions(-) create mode 100644 Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/QueryExecutionContextWithException.cs diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/PipelineContinuationToken.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/PipelineContinuationToken.cs index 7254dcb97e..b587288192 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/PipelineContinuationToken.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/PipelineContinuationToken.cs @@ -20,34 +20,24 @@ protected PipelineContinuationToken(Version version) public Version Version { get; } - public static bool TryParse( - string rawContinuationToken, + public static bool TryCreateFromCosmosElement( + CosmosElement cosmosElement, out PipelineContinuationToken pipelineContinuationToken) { - if (rawContinuationToken == null) + if (cosmosElement == null) { - throw new ArgumentNullException(nameof(rawContinuationToken)); + throw new ArgumentNullException(nameof(cosmosElement)); } - if (!CosmosElement.TryParse( - rawContinuationToken, - out CosmosObject parsedContinuationToken)) + if (!(cosmosElement is CosmosObject cosmosObject)) { - // Failed to parse so we need to assume it's a V0 token - if (!PipelineContinuationTokenV0.TryParse( - rawContinuationToken, - out PipelineContinuationTokenV0 pipelineContinuationTokenV0)) - { - pipelineContinuationToken = default; - return false; - } - - pipelineContinuationToken = pipelineContinuationTokenV0; + // Not a CosmosObject, so it doesn't have a version number, so + pipelineContinuationToken = new PipelineContinuationTokenV0(cosmosElement); return true; } if (!PipelineContinuationToken.TryParseVersion( - parsedContinuationToken, + cosmosObject, out Version version)) { pipelineContinuationToken = default; @@ -56,8 +46,8 @@ public static bool TryParse( if (version == PipelineContinuationTokenV0.VersionNumber) { - if (!PipelineContinuationTokenV0.TryParse( - rawContinuationToken, + if (!PipelineContinuationTokenV0.TryCreateFromCosmosElement( + cosmosElement, out PipelineContinuationTokenV0 pipelineContinuationTokenV0)) { pipelineContinuationToken = default; @@ -68,8 +58,8 @@ public static bool TryParse( } else if (version == PipelineContinuationTokenV1.VersionNumber) { - if (!PipelineContinuationTokenV1.TryParse( - parsedContinuationToken, + if (!PipelineContinuationTokenV1.TryCreateFromCosmosElement( + cosmosObject, out PipelineContinuationTokenV1 pipelineContinuationTokenV1)) { pipelineContinuationToken = default; @@ -80,8 +70,8 @@ public static bool TryParse( } else if (version == PipelineContinuationTokenV1_1.VersionNumber) { - if (!PipelineContinuationTokenV1_1.TryParse( - parsedContinuationToken, + if (!PipelineContinuationTokenV1_1.TryCreateFromCosmosElement( + cosmosObject, out PipelineContinuationTokenV1_1 pipelineContinuationTokenV1_1)) { pipelineContinuationToken = default; diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/PipelineContinuationTokenV0.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/PipelineContinuationTokenV0.cs index 5a1ce26202..39c318e667 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/PipelineContinuationTokenV0.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/PipelineContinuationTokenV0.cs @@ -5,6 +5,7 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ContinuationTokens { using System; + using Microsoft.Azure.Cosmos.CosmosElements; /// /// Pipelined continuation token before we started versioning continuation tokens. @@ -13,29 +14,29 @@ internal sealed class PipelineContinuationTokenV0 : PipelineContinuationToken { public static readonly Version VersionNumber = new Version(major: 0, minor: 0); - public PipelineContinuationTokenV0(string sourceContinuationToken) + public PipelineContinuationTokenV0(CosmosElement sourceContinuationToken) : base(PipelineContinuationTokenV0.VersionNumber) { this.SourceContinuationToken = sourceContinuationToken ?? throw new ArgumentNullException(nameof(sourceContinuationToken)); } - public string SourceContinuationToken { get; } + public CosmosElement SourceContinuationToken { get; } public override string ToString() { - return this.SourceContinuationToken; + return this.SourceContinuationToken.ToString(); } - public static bool TryParse( - string rawContinuationToken, - out PipelineContinuationTokenV0 pipelinedContinuationTokenV0) + public static bool TryCreateFromCosmosElement( + CosmosElement cosmosElement, + out PipelineContinuationTokenV0 pipelineContinuationTokenV0) { - if (rawContinuationToken == null) + if (cosmosElement == null) { - throw new ArgumentNullException(nameof(rawContinuationToken)); + throw new ArgumentNullException(nameof(cosmosElement)); } - pipelinedContinuationTokenV0 = new PipelineContinuationTokenV0(rawContinuationToken); + pipelineContinuationTokenV0 = new PipelineContinuationTokenV0(cosmosElement); return true; } } diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/PipelineContinuationTokenV1.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/PipelineContinuationTokenV1.cs index ab4cce7020..9594d6f5dc 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/PipelineContinuationTokenV1.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/PipelineContinuationTokenV1.cs @@ -17,13 +17,13 @@ internal sealed class PipelineContinuationTokenV1 : PipelineContinuationToken private static readonly string SourceContinuationTokenPropertyName = "SourceContinuationToken"; - public PipelineContinuationTokenV1(string sourceContinuationToken) + public PipelineContinuationTokenV1(CosmosElement sourceContinuationToken) : base(PipelineContinuationTokenV1.VersionNumber) { this.SourceContinuationToken = sourceContinuationToken ?? throw new ArgumentNullException(nameof(sourceContinuationToken)); } - public string SourceContinuationToken { get; } + public CosmosElement SourceContinuationToken { get; } public override string ToString() { @@ -35,12 +35,12 @@ public override string ToString() }, { PipelineContinuationTokenV1.SourceContinuationTokenPropertyName, - CosmosString.Create(this.SourceContinuationToken) + this.SourceContinuationToken }, }).ToString(); } - public static bool TryParse( + public static bool TryCreateFromCosmosElement( CosmosObject parsedContinuationToken, out PipelineContinuationTokenV1 pipelinedContinuationTokenV1) { @@ -63,9 +63,9 @@ public static bool TryParse( return false; } - if (!PipelineContinuationTokenV1.TryParseSourceContinuationToken( - parsedContinuationToken, - out string sourceContinuationToken)) + if (!parsedContinuationToken.TryGetValue( + SourceContinuationTokenPropertyName, + out CosmosElement sourceContinuationToken)) { pipelinedContinuationTokenV1 = default; return false; @@ -85,26 +85,5 @@ public static PipelineContinuationTokenV1 CreateFromV0Token( return new PipelineContinuationTokenV1(pipelinedContinuationTokenV0.SourceContinuationToken); } - - private static bool TryParseSourceContinuationToken( - CosmosObject parsedContinuationToken, - out string sourceContinuationToken) - { - if (parsedContinuationToken == null) - { - throw new ArgumentNullException(nameof(parsedContinuationToken)); - } - - if (!parsedContinuationToken.TryGetValue( - PipelineContinuationTokenV1.SourceContinuationTokenPropertyName, - out CosmosString parsedSourceContinuationToken)) - { - sourceContinuationToken = default; - return false; - } - - sourceContinuationToken = parsedSourceContinuationToken.Value; - return true; - } } } diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/PipelineContinuationTokenV1_1.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/PipelineContinuationTokenV1_1.cs index c5c18a8d4f..c2ab949206 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/PipelineContinuationTokenV1_1.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/PipelineContinuationTokenV1_1.cs @@ -21,14 +21,14 @@ internal sealed class PipelineContinuationTokenV1_1 : PipelineContinuationToken public PipelineContinuationTokenV1_1( PartitionedQueryExecutionInfo queryPlan, - string sourceContinuationToken) + CosmosElement sourceContinuationToken) : base(PipelineContinuationTokenV1_1.VersionNumber) { this.QueryPlan = queryPlan; this.SourceContinuationToken = sourceContinuationToken ?? throw new ArgumentNullException(nameof(sourceContinuationToken)); } - public string SourceContinuationToken { get; } + public CosmosElement SourceContinuationToken { get; } public PartitionedQueryExecutionInfo QueryPlan { get; } @@ -47,7 +47,7 @@ public string ToString(int lengthLimitInBytes) } else { - shouldSerializeQueryPlan = (queryPlanString.Length + this.SourceContinuationToken.Length) < lengthLimitInBytes; + shouldSerializeQueryPlan = (queryPlanString.Length + this.SourceContinuationToken.ToString().Length) < lengthLimitInBytes; } return CosmosObject.Create(new Dictionary() @@ -58,16 +58,16 @@ public string ToString(int lengthLimitInBytes) }, { PipelineContinuationTokenV1_1.QueryPlanPropertyName, - shouldSerializeQueryPlan ? (CosmosElement)CosmosString.Create(queryPlanString) : (CosmosElement)CosmosNull.Create() + shouldSerializeQueryPlan ? (CosmosElement)CosmosString.Create(queryPlanString) : (CosmosElement)CosmosNull.Create() }, { PipelineContinuationTokenV1_1.SourceContinuationTokenPropertyName, - CosmosString.Create(this.SourceContinuationToken) + this.SourceContinuationToken }, }).ToString(); } - public static bool TryParse( + public static bool TryCreateFromCosmosElement( CosmosObject parsedContinuationToken, out PipelineContinuationTokenV1_1 pipelinedContinuationToken) { @@ -98,9 +98,9 @@ public static bool TryParse( return false; } - if (!PipelineContinuationTokenV1_1.TryParseSourceContinuationToken( - parsedContinuationToken, - out string sourceContinuationToken)) + if (!parsedContinuationToken.TryGetValue( + SourceContinuationTokenPropertyName, + out CosmosElement sourceContinuationToken)) { pipelinedContinuationToken = default; return false; @@ -148,26 +148,5 @@ private static bool TryParseQueryPlan( return true; } - - private static bool TryParseSourceContinuationToken( - CosmosObject parsedContinuationToken, - out string sourceContinuationToken) - { - if (parsedContinuationToken == null) - { - throw new ArgumentNullException(nameof(parsedContinuationToken)); - } - - if (!parsedContinuationToken.TryGetValue( - PipelineContinuationTokenV1_1.SourceContinuationTokenPropertyName, - out CosmosString parsedSourceContinuationToken)) - { - sourceContinuationToken = default; - return false; - } - - sourceContinuationToken = parsedSourceContinuationToken.Value; - return true; - } } } diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/AggregateDocumentQueryExecutionComponent.Client.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/AggregateDocumentQueryExecutionComponent.Client.cs index 3180fe52b1..a90e6edf33 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/AggregateDocumentQueryExecutionComponent.Client.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/AggregateDocumentQueryExecutionComponent.Client.cs @@ -33,8 +33,8 @@ public static async Task> TryCreateAs IReadOnlyDictionary aliasToAggregateType, IReadOnlyList orderedAliases, bool hasSelectValue, - RequestContinuationToken continuationToken, - Func>> tryCreateSourceAsync) + CosmosElement continuationToken, + Func>> tryCreateSourceAsync) { if (tryCreateSourceAsync == null) { @@ -46,7 +46,7 @@ public static async Task> TryCreateAs aliasToAggregateType, orderedAliases, hasSelectValue, - continuationToken: StringRequestContinuationToken.Null); + continuationToken: null); if (!tryCreateSingleGroupAggregator.Succeeded) { diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/AggregateDocumentQueryExecutionComponent.Compute.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/AggregateDocumentQueryExecutionComponent.Compute.cs index 66b32ff367..e794215e36 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/AggregateDocumentQueryExecutionComponent.Compute.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/AggregateDocumentQueryExecutionComponent.Compute.cs @@ -38,13 +38,13 @@ public static async Task> TryCreateAs IReadOnlyDictionary aliasToAggregateType, IReadOnlyList orderedAliases, bool hasSelectValue, - RequestContinuationToken requestContinuation, - Func>> tryCreateSourceAsync) + CosmosElement requestContinuation, + Func>> tryCreateSourceAsync) { AggregateContinuationToken aggregateContinuationToken; - if (!requestContinuation.IsNull) + if (requestContinuation != null) { - if (!AggregateContinuationToken.TryParse(requestContinuation, out aggregateContinuationToken)) + if (!AggregateContinuationToken.TryCreateFromCosmosElement(requestContinuation, out aggregateContinuationToken)) { return TryCatch.FromException( new MalformedContinuationTokenException($"Malfomed {nameof(AggregateContinuationToken)}: '{requestContinuation}'")); @@ -60,7 +60,7 @@ public static async Task> TryCreateAs aliasToAggregateType, orderedAliases, hasSelectValue, - RequestContinuationToken.Create(aggregateContinuationToken.SingleGroupAggregatorContinuationToken)); + aggregateContinuationToken.SingleGroupAggregatorContinuationToken); if (!tryCreateSingleGroupAggregator.Succeeded) { @@ -68,13 +68,14 @@ public static async Task> TryCreateAs tryCreateSingleGroupAggregator.Exception); } - return (await tryCreateSourceAsync(RequestContinuationToken.Create(aggregateContinuationToken.SourceContinuationToken))).Try((source) => - { - return new ComputeAggregateDocumentQueryExecutionComponent( - source, - tryCreateSingleGroupAggregator.Result, - hasSelectValue); - }); + return (await tryCreateSourceAsync(aggregateContinuationToken.SourceContinuationToken)) + .Try((source) => + { + return new ComputeAggregateDocumentQueryExecutionComponent( + source, + tryCreateSingleGroupAggregator.Result, + hasSelectValue); + }); } public override async Task DrainAsync( @@ -191,8 +192,8 @@ public AggregateContinuationToken( public CosmosElement SourceContinuationToken { get; } - public static bool TryParse( - RequestContinuationToken continuationToken, + public static bool TryCreateFromCosmosElement( + CosmosElement continuationToken, out AggregateContinuationToken aggregateContinuationToken) { if (continuationToken == null) @@ -200,12 +201,7 @@ public static bool TryParse( throw new ArgumentNullException(nameof(continuationToken)); } - if (!(continuationToken is CosmosElementRequestContinuationToken cosmosElementRequestContinuationToken)) - { - throw new ArgumentException($"Expected {nameof(CosmosElementRequestContinuationToken)} instead of: {continuationToken.GetType()}"); - } - - if (!(cosmosElementRequestContinuationToken.Value is CosmosObject rawAggregateContinuationToken)) + if (!(continuationToken is CosmosObject rawAggregateContinuationToken)) { aggregateContinuationToken = default; return false; 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 ab8fa68cda..36e26b1f4a 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/AggregateDocumentQueryExecutionComponent.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/AggregateDocumentQueryExecutionComponent.cs @@ -62,8 +62,8 @@ public static async Task> TryCreateAs IReadOnlyDictionary aliasToAggregateType, IReadOnlyList orderedAliases, bool hasSelectValue, - RequestContinuationToken continuationToken, - Func>> tryCreateSourceAsync) + CosmosElement continuationToken, + Func>> tryCreateSourceAsync) { if (tryCreateSourceAsync == null) { diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/AverageAggregator.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/AverageAggregator.cs index 09d0a7c681..c5c3316765 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/AverageAggregator.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/AverageAggregator.cs @@ -36,8 +36,14 @@ private AverageAggregator(AverageInfo globalAverage) /// The local average to add to the global average. public void Aggregate(CosmosElement localAverage) { - // item is a JObject of the form : { "sum": , "count": } - AverageInfo newInfo = AverageInfo.Create(localAverage); + // item is a JObject of the form : { "sum": , "count": } + TryCatch tryCreateAverageInfo = AverageInfo.TryCreateFromCosmosElement(localAverage); + if (!tryCreateAverageInfo.Succeeded) + { + throw tryCreateAverageInfo.Exception; + } + + AverageInfo newInfo = tryCreateAverageInfo.Result; this.globalAverage += newInfo; } @@ -65,21 +71,21 @@ public void SerializeState(IJsonWriter jsonWriter) this.globalAverage.SerializeState(jsonWriter); } - public static TryCatch TryCreate(RequestContinuationToken continuationToken) + public static TryCatch TryCreate(CosmosElement continuationToken) { - if (continuationToken == null) - { - throw new ArgumentNullException(nameof(continuationToken)); - } - AverageInfo averageInfo; - if (!continuationToken.IsNull) + if (continuationToken != null) { - if (!AverageInfo.TryParse(continuationToken, out averageInfo)) + TryCatch tryCreateAverageInfo = AverageInfo.TryCreateFromCosmosElement(continuationToken); + if (!tryCreateAverageInfo.Succeeded) { return TryCatch.FromException( - new MalformedContinuationTokenException($"Invalid continuation token: {continuationToken}")); + new MalformedContinuationTokenException( + message: $"Invalid continuation token: {continuationToken}", + innerException: tryCreateAverageInfo.Exception)); } + + averageInfo = tryCreateAverageInfo.Result; } else { @@ -111,7 +117,7 @@ public AverageInfo(double? sum, long count) /// /// Initializes a new instance of the AverageInfo class. /// - public static AverageInfo Create(CosmosElement cosmosElement) + public static TryCatch TryCreateFromCosmosElement(CosmosElement cosmosElement) { if (cosmosElement == null) { @@ -120,7 +126,8 @@ public static AverageInfo Create(CosmosElement cosmosElement) if (!(cosmosElement is CosmosObject cosmosObject)) { - throw new ArgumentException($"{nameof(cosmosElement)} must not be an object."); + return TryCatch.FromException( + new ArgumentException($"{nameof(cosmosElement)} must be an object.")); } double? sum; @@ -128,7 +135,8 @@ public static AverageInfo Create(CosmosElement cosmosElement) { if (!(sumPropertyValue is CosmosNumber cosmosSum)) { - throw new ArgumentException($"value for the {SumName} field was not a number"); + return TryCatch.FromException( + new ArgumentException($"value for the {SumName} field was not a number")); } sum = Number64.ToDouble(cosmosSum.Value); @@ -140,12 +148,13 @@ public static AverageInfo Create(CosmosElement cosmosElement) if (!cosmosObject.TryGetValue(CountName, out CosmosNumber cosmosCount)) { - throw new ArgumentException($"value for the {CountName} field was not a number"); + return TryCatch.FromException( + new ArgumentException($"value for the {CountName} field was not a number")); } long count = Number64.ToLong(cosmosCount.Value); - - return new AverageInfo(sum, count); + + return TryCatch.FromResult(new AverageInfo(sum, count)); } /// @@ -224,23 +233,6 @@ public override string ToString() ""{CountName}"" : {this.Count} }}"; } - - public static bool TryParse(RequestContinuationToken requestContinuationToken, out AverageInfo averageInfo) - { - if (requestContinuationToken == null) - { - throw new ArgumentNullException(nameof(requestContinuationToken)); - } - - if (!requestContinuationToken.TryConvertToCosmosElement(out CosmosElement cosmosElementAverageInfo)) - { - averageInfo = default; - return false; - } - - averageInfo = AverageInfo.Create(cosmosElementAverageInfo); - return true; - } } } } \ No newline at end of file diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/CountAggregator.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/CountAggregator.cs index 9b05acf3f8..abaa599f38 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/CountAggregator.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/CountAggregator.cs @@ -72,23 +72,18 @@ public void SerializeState(IJsonWriter jsonWriter) jsonWriter.WriteInt64Value(this.globalCount); } - public static TryCatch TryCreate(RequestContinuationToken continuationToken) + public static TryCatch TryCreate(CosmosElement continuationToken) { - if (continuationToken == null) - { - throw new ArgumentNullException(nameof(continuationToken)); - } - long partialCount; - if (!continuationToken.IsNull) + if (continuationToken != null) { - if (!continuationToken.TryConvertToCosmosElement(out CosmosInt64 cosmosPartialCount)) + if (!(continuationToken is CosmosNumber cosmosPartialCount)) { return TryCatch.FromException( new MalformedContinuationTokenException($@"Invalid count continuation token: ""{continuationToken}"".")); } - partialCount = cosmosPartialCount.GetValue(); + partialCount = Number64.ToLong(cosmosPartialCount.Value); } else { diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/MinMaxAggregator.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/MinMaxAggregator.cs index 4fcb08cec2..af88b822cd 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/MinMaxAggregator.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/MinMaxAggregator.cs @@ -167,33 +167,22 @@ public void SerializeState(IJsonWriter jsonWriter) minMaxContinuationTokenAsCosmosElement.WriteTo(jsonWriter); } - public static TryCatch TryCreateMinAggregator(RequestContinuationToken continuationToken) + public static TryCatch TryCreateMinAggregator(CosmosElement continuationToken) { return MinMaxAggregator.TryCreate(isMinAggregation: true, continuationToken: continuationToken); } - public static TryCatch TryCreateMaxAggregator(RequestContinuationToken continuationToken) + public static TryCatch TryCreateMaxAggregator(CosmosElement continuationToken) { return MinMaxAggregator.TryCreate(isMinAggregation: false, continuationToken: continuationToken); } - private static TryCatch TryCreate(bool isMinAggregation, RequestContinuationToken continuationToken) + private static TryCatch TryCreate(bool isMinAggregation, CosmosElement continuationToken) { - if (continuationToken == null) - { - throw new ArgumentNullException(nameof(continuationToken)); - } - CosmosElement globalMinMax; - if (!continuationToken.IsNull) + if (continuationToken != null) { - if (!continuationToken.TryConvertToCosmosElement(out CosmosElement cosmosElementContinuationToken)) - { - return TryCatch.FromException( - new MalformedContinuationTokenException($"Malformed continuation token: {continuationToken}")); - } - - TryCatch tryCreateMinMaxContinuationToken = MinMaxContinuationToken.TryCreateFromCosmosElement(cosmosElementContinuationToken); + TryCatch tryCreateMinMaxContinuationToken = MinMaxContinuationToken.TryCreateFromCosmosElement(continuationToken); if (!tryCreateMinMaxContinuationToken.Succeeded) { return TryCatch.FromException(tryCreateMinMaxContinuationToken.Exception); diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/SingleGroupAggregator.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/SingleGroupAggregator.cs index 29a7fbfb4d..188d74c0f4 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/SingleGroupAggregator.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/SingleGroupAggregator.cs @@ -37,7 +37,7 @@ public static TryCatch TryCreate( IReadOnlyDictionary aggregateAliasToAggregateType, IReadOnlyList orderedAliases, bool hasSelectValue, - RequestContinuationToken continuationToken) + CosmosElement continuationToken) { if (aggregates == null) { @@ -91,7 +91,7 @@ private SelectValueAggregateValues(AggregateValue aggregateValue) this.aggregateValue = aggregateValue ?? throw new ArgumentNullException(nameof(AggregateValue)); } - public static TryCatch TryCreate(AggregateOperator? aggregateOperator, RequestContinuationToken continuationToken) + public static TryCatch TryCreate(AggregateOperator? aggregateOperator, CosmosElement continuationToken) { return AggregateValue.TryCreate(aggregateOperator, continuationToken) .Try((aggregateValue) => (SingleGroupAggregator)new SelectValueAggregateValues(aggregateValue)); @@ -190,7 +190,7 @@ public override void SerializeState(IJsonWriter jsonWriter) public static TryCatch TryCreate( IReadOnlyDictionary aggregateAliasToAggregateType, IReadOnlyList orderedAliases, - RequestContinuationToken continuationToken) + CosmosElement continuationToken) { if (aggregateAliasToAggregateType == null) { @@ -202,20 +202,17 @@ public static TryCatch TryCreate( throw new ArgumentNullException(nameof(orderedAliases)); } - if (continuationToken == null) - { - throw new ArgumentNullException(nameof(continuationToken)); - } - CosmosObject aliasToContinuationToken; - if (!continuationToken.IsNull) + if (continuationToken != null) { - if (!continuationToken.TryConvertToCosmosElement(out aliasToContinuationToken)) + if (!(continuationToken is CosmosObject cosmosObject)) { return TryCatch.FromException( new MalformedContinuationTokenException( $"{nameof(SelectListAggregateValues)} continuation token is malformed: {continuationToken}.")); } + + aliasToContinuationToken = cosmosObject; } else { @@ -227,14 +224,14 @@ public static TryCatch TryCreate( { string alias = aliasToAggregate.Key; AggregateOperator? aggregateOperator = aliasToAggregate.Value; - RequestContinuationToken aliasContinuationToken; + CosmosElement aliasContinuationToken; if (aliasToContinuationToken != null) { - aliasContinuationToken = RequestContinuationToken.Create(aliasToContinuationToken[alias]); + aliasContinuationToken = aliasToContinuationToken[alias]; } else { - aliasContinuationToken = CosmosElementRequestContinuationToken.Null; + aliasContinuationToken = null; } TryCatch tryCreateAggregateValue = AggregateValue.TryCreate( @@ -299,7 +296,7 @@ public override string ToString() return this.Result.ToString(); } - public static TryCatch TryCreate(AggregateOperator? aggregateOperator, RequestContinuationToken continuationToken) + public static TryCatch TryCreate(AggregateOperator? aggregateOperator, CosmosElement continuationToken) { TryCatch value; if (aggregateOperator.HasValue) @@ -343,7 +340,7 @@ public override void SerializeState(IJsonWriter jsonWriter) public static TryCatch TryCreate( AggregateOperator aggregateOperator, - RequestContinuationToken continuationToken) + CosmosElement continuationToken) { TryCatch tryCreateAggregator; switch (aggregateOperator) @@ -430,18 +427,13 @@ public override void SerializeState(IJsonWriter jsonWriter) jsonWriter.WriteObjectEnd(); } - public static TryCatch TryCreate(RequestContinuationToken continuationToken) + public static TryCatch TryCreate(CosmosElement continuationToken) { - if (continuationToken == null) - { - throw new ArgumentNullException(nameof(continuationToken)); - } - CosmosElement value; bool initialized; - if (!continuationToken.IsNull) + if (continuationToken != null) { - if (!continuationToken.TryConvertToCosmosElement(out CosmosObject rawContinuationToken)) + if (!(continuationToken is CosmosObject rawContinuationToken)) { return TryCatch.FromException( new MalformedContinuationTokenException($"Invalid {nameof(ScalarAggregateValue)}: {continuationToken}")); diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/SumAggregator.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/SumAggregator.cs index 36b6d4ab54..546a885540 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/SumAggregator.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/SumAggregator.cs @@ -80,18 +80,18 @@ public void SerializeState(IJsonWriter jsonWriter) jsonWriter.WriteFloat64Value(this.globalSum); } - public static TryCatch TryCreate(RequestContinuationToken requestContinuationToken) + public static TryCatch TryCreate(CosmosElement requestContinuationToken) { double partialSum; - if (!requestContinuationToken.IsNull) + if (requestContinuationToken != null) { - if (!requestContinuationToken.TryConvertToCosmosElement(out CosmosFloat64 cosmosFloat64)) + if (!(requestContinuationToken is CosmosNumber cosmosNumber)) { return TryCatch.FromException( new MalformedContinuationTokenException($"Malformed {nameof(SumAggregator)} continuation token: {requestContinuationToken}")); } - partialSum = cosmosFloat64.GetValue(); + partialSum = Number64.ToLong(cosmosNumber.Value); } else { diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctDocumentQueryExecutionComponent.Client.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctDocumentQueryExecutionComponent.Client.cs index 605a928ff7..7a0f9256d5 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctDocumentQueryExecutionComponent.Client.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctDocumentQueryExecutionComponent.Client.cs @@ -42,48 +42,66 @@ private ClientDistinctDocumentQueryExecutionComponent( } public static async Task> TryCreateAsync( - RequestContinuationToken requestContinuation, - Func>> tryCreateSourceAsync, + CosmosElement requestContinuation, + Func>> tryCreateSourceAsync, DistinctQueryType distinctQueryType) { - if (requestContinuation == null) - { - throw new ArgumentNullException(nameof(requestContinuation)); - } - - if (!(requestContinuation is StringRequestContinuationToken stringRequestContinuationToken)) - { - throw new ArgumentException($"Unknown {nameof(RequestContinuationToken)} type: {requestContinuation.GetType()}"); - } - if (tryCreateSourceAsync == null) { throw new ArgumentNullException(nameof(tryCreateSourceAsync)); } DistinctContinuationToken distinctContinuationToken; - if (!requestContinuation.IsNull) + if (requestContinuation != null) { - if (!DistinctContinuationToken.TryParse(stringRequestContinuationToken, out distinctContinuationToken)) + if (!DistinctContinuationToken.TryParse(requestContinuation.ToString(), out distinctContinuationToken)) { return TryCatch.FromException( - new MalformedContinuationTokenException($"Invalid {nameof(DistinctContinuationToken)}: {stringRequestContinuationToken}")); + new MalformedContinuationTokenException( + $"Invalid {nameof(DistinctContinuationToken)}: {requestContinuation}")); } } else { - distinctContinuationToken = new DistinctContinuationToken(sourceToken: null, distinctMapToken: null); + distinctContinuationToken = new DistinctContinuationToken( + sourceToken: null, + distinctMapToken: null); + } + + CosmosElement distinctMapToken; + if (distinctContinuationToken.DistinctMapToken != null) + { + distinctMapToken = CosmosString.Create(distinctContinuationToken.DistinctMapToken); + } + else + { + distinctMapToken = null; } TryCatch tryCreateDistinctMap = DistinctMap.TryCreate( distinctQueryType, - RequestContinuationToken.Create(distinctContinuationToken.DistinctMapToken)); + distinctMapToken); if (!tryCreateDistinctMap.Succeeded) { return TryCatch.FromException(tryCreateDistinctMap.Exception); } - TryCatch tryCreateSource = await tryCreateSourceAsync(RequestContinuationToken.Create(distinctContinuationToken.SourceToken)); + CosmosElement sourceToken; + if (distinctContinuationToken.SourceToken != null) + { + if (!CosmosElement.TryParse(distinctContinuationToken.SourceToken, out sourceToken)) + { + return TryCatch.FromException( + new MalformedContinuationTokenException( + $"Invalid Source Token: {distinctContinuationToken.SourceToken}")); + } + } + else + { + sourceToken = null; + } + + TryCatch tryCreateSource = await tryCreateSourceAsync(sourceToken); if (!tryCreateSource.Succeeded) { return TryCatch.FromException(tryCreateSource.Exception); @@ -189,23 +207,12 @@ public DistinctContinuationToken(string sourceToken, string distinctMapToken) /// The output DistinctContinuationToken. /// True if we successfully parsed the DistinctContinuationToken, else false. public static bool TryParse( - StringRequestContinuationToken stringRequestContinuationToken, + string stringRequestContinuationToken, out DistinctContinuationToken distinctContinuationToken) { - if (stringRequestContinuationToken == null) - { - throw new ArgumentNullException(nameof(stringRequestContinuationToken)); - } - - if (stringRequestContinuationToken.IsNull) - { - distinctContinuationToken = default; - return false; - } - try { - distinctContinuationToken = JsonConvert.DeserializeObject(stringRequestContinuationToken.Value); + distinctContinuationToken = JsonConvert.DeserializeObject(stringRequestContinuationToken); return true; } catch (JsonException) diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctDocumentQueryExecutionComponent.Compute.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctDocumentQueryExecutionComponent.Compute.cs index 0ed0af3f52..99e0604c21 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctDocumentQueryExecutionComponent.Compute.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctDocumentQueryExecutionComponent.Compute.cs @@ -33,32 +33,22 @@ private ComputeDistinctDocumentQueryExecutionComponent( } public static async Task> TryCreateAsync( - RequestContinuationToken requestContinuation, - Func>> tryCreateSourceAsync, + CosmosElement requestContinuation, + Func>> tryCreateSourceAsync, DistinctQueryType distinctQueryType) { - if (requestContinuation == null) - { - throw new ArgumentNullException(nameof(requestContinuation)); - } - - if (!(requestContinuation is CosmosElementRequestContinuationToken cosmosElementRequestContinuationToken)) - { - throw new ArgumentException($"Unknown {nameof(RequestContinuationToken)} type: {requestContinuation.GetType()}"); - } - if (tryCreateSourceAsync == null) { throw new ArgumentNullException(nameof(tryCreateSourceAsync)); } DistinctContinuationToken distinctContinuationToken; - if (!requestContinuation.IsNull) + if (requestContinuation != null) { - if (!DistinctContinuationToken.TryParse(cosmosElementRequestContinuationToken.Value, out distinctContinuationToken)) + if (!DistinctContinuationToken.TryParse(requestContinuation, out distinctContinuationToken)) { return TryCatch.FromException( - new MalformedContinuationTokenException($"Invalid {nameof(DistinctContinuationToken)}: {cosmosElementRequestContinuationToken.Value}")); + new MalformedContinuationTokenException($"Invalid {nameof(DistinctContinuationToken)}: {requestContinuation}")); } } else @@ -68,14 +58,14 @@ public static async Task> TryCreateAs TryCatch tryCreateDistinctMap = DistinctMap.TryCreate( distinctQueryType, - RequestContinuationToken.Create(distinctContinuationToken.DistinctMapToken)); + distinctContinuationToken.DistinctMapToken); if (!tryCreateDistinctMap.Succeeded) { return TryCatch.FromException(tryCreateDistinctMap.Exception); } TryCatch tryCreateSource = await tryCreateSourceAsync( - RequestContinuationToken.Create(distinctContinuationToken.SourceToken)); + distinctContinuationToken.SourceToken); if (!tryCreateSource.Succeeded) { return TryCatch.FromException(tryCreateSource.Exception); 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 3a535af119..a213306321 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctDocumentQueryExecutionComponent.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctDocumentQueryExecutionComponent.cs @@ -42,8 +42,8 @@ protected DistinctDocumentQueryExecutionComponent( public static async Task> TryCreateAsync( ExecutionEnvironment executionEnvironment, - RequestContinuationToken requestContinuation, - Func>> tryCreateSourceAsync, + CosmosElement requestContinuation, + Func>> tryCreateSourceAsync, DistinctQueryType distinctQueryType) { if (tryCreateSourceAsync == null) diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctMap.OrderedDistinctMap.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctMap.OrderedDistinctMap.cs index 5da0801c67..e09cec1333 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctMap.OrderedDistinctMap.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctMap.OrderedDistinctMap.cs @@ -84,47 +84,35 @@ public override void SerializeState(IJsonWriter jsonWriter) jsonWriter.WriteBinaryValue(UInt128.ToByteArray(this.lastHash)); } - public static TryCatch TryCreate(RequestContinuationToken requestContinuationToken) + public static TryCatch TryCreate(CosmosElement requestContinuationToken) { - if (requestContinuationToken == null) - { - throw new ArgumentNullException(nameof(requestContinuationToken)); - } - UInt128 lastHash; - if (!requestContinuationToken.IsNull) + if (requestContinuationToken != null) { switch (requestContinuationToken) { - case StringRequestContinuationToken stringRequestContinuationToken: - if (!UInt128.TryParse(stringRequestContinuationToken.Value, out lastHash)) + case CosmosString cosmosString: + if (!UInt128.TryParse(cosmosString.Value, out lastHash)) { return TryCatch.FromException( new MalformedContinuationTokenException( - $"Malformed {nameof(OrderedDistinctMap)} continuation token: {stringRequestContinuationToken.Value}.")); + $"Malformed {nameof(OrderedDistinctMap)} continuation token: {requestContinuationToken}.")); } break; - case CosmosElementRequestContinuationToken cosmosElementRequestContinuationToken: - if (!(cosmosElementRequestContinuationToken.Value is CosmosBinary cosmosBinary)) - { - return TryCatch.FromException( - new MalformedContinuationTokenException( - $"Malformed {nameof(OrderedDistinctMap)} continuation token: {cosmosElementRequestContinuationToken.Value}.")); - } + case CosmosBinary cosmosBinary: if (!UInt128.TryCreateFromByteArray(cosmosBinary.Value.Span, out lastHash)) { return TryCatch.FromException( new MalformedContinuationTokenException( - $"Malformed {nameof(OrderedDistinctMap)} continuation token: {cosmosElementRequestContinuationToken.Value}.")); + $"Malformed {nameof(OrderedDistinctMap)} continuation token: {requestContinuationToken}.")); } break; default: - throw new ArgumentOutOfRangeException($"Unknown {nameof(RequestContinuationToken)} type. {requestContinuationToken.GetType()}."); + throw new ArgumentOutOfRangeException($"Unknown {nameof(requestContinuationToken)} type. {requestContinuationToken.GetType()}."); } - } else { diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctMap.UnorderedDistinctMap.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctMap.UnorderedDistinctMap.cs index ddbfb5c8af..a86709fbe5 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctMap.UnorderedDistinctMap.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctMap.UnorderedDistinctMap.cs @@ -391,13 +391,8 @@ private bool AddObjectValue(CosmosObject cosmosObject) return this.objects.Add(hash); } - public static TryCatch TryCreate(RequestContinuationToken continuationToken) + public static TryCatch TryCreate(CosmosElement continuationToken) { - if (continuationToken == null) - { - throw new ArgumentNullException(nameof(continuationToken)); - } - HashSet numbers = new HashSet(); HashSet stringsLength4 = new HashSet(); HashSet stringsLength8 = new HashSet(); @@ -407,15 +402,9 @@ public static TryCatch TryCreate(RequestContinuationToken continuat HashSet objects = new HashSet(); SimpleValues simpleValues = SimpleValues.None; - if (!continuationToken.IsNull) + if (continuationToken != null) { - if (!(continuationToken is CosmosElementRequestContinuationToken cosmosElementRequestContinuation)) - { - throw new ArgumentException($"Unknown {nameof(RequestContinuationToken)} type: {continuationToken.GetType()}"); - } - - CosmosElement cosmosElement = cosmosElementRequestContinuation.Value; - if (!(cosmosElement is CosmosObject hashDictionary)) + if (!(continuationToken is CosmosObject hashDictionary)) { return TryCatch.FromException( new MalformedContinuationTokenException( diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctMap.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctMap.cs index 19e7e53dff..dae3f7e60a 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctMap.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctMap.cs @@ -29,7 +29,7 @@ internal abstract partial class DistinctMap /// The appropriate IDistinctMap. public static TryCatch TryCreate( DistinctQueryType distinctQueryType, - RequestContinuationToken distinctMapContinuationToken) + CosmosElement distinctMapContinuationToken) { TryCatch tryCreateDistinctMap; switch (distinctQueryType) 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 e19dc54f14..2f5a2b59ac 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 @@ -31,8 +31,8 @@ private ClientGroupByDocumentQueryExecutionComponent( } public static async Task> TryCreateAsync( - RequestContinuationToken requestContinuation, - Func>> tryCreateSource, + CosmosElement requestContinuation, + Func>> tryCreateSource, IReadOnlyDictionary groupByAliasToAggregateType, IReadOnlyList orderedAliases, bool hasSelectValue) 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 47f3c3ce08..e0f48f9212 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 @@ -37,21 +37,16 @@ private ComputeGroupByDocumentQueryExecutionComponent( } public static async Task> TryCreateAsync( - RequestContinuationToken requestContinuation, - Func>> tryCreateSourceAsync, + CosmosElement requestContinuation, + Func>> tryCreateSourceAsync, IReadOnlyDictionary groupByAliasToAggregateType, IReadOnlyList orderedAliases, bool hasSelectValue) { - if (!(requestContinuation is CosmosElementRequestContinuationToken cosmosElementRequestContinuationToken)) - { - throw new ArgumentException($"Unknown {nameof(RequestContinuationToken)} type: {requestContinuation.GetType()}."); - } - GroupByContinuationToken groupByContinuationToken; - if (!requestContinuation.IsNull) + if (requestContinuation != null) { - if (!GroupByContinuationToken.TryParse(cosmosElementRequestContinuationToken.Value, out groupByContinuationToken)) + if (!GroupByContinuationToken.TryParse(requestContinuation, out groupByContinuationToken)) { return TryCatch.FromException( new MalformedContinuationTokenException($"Invalid {nameof(GroupByContinuationToken)}: '{requestContinuation}'")); @@ -73,7 +68,7 @@ public static async Task> TryCreateAs } else { - tryCreateSource = await tryCreateSourceAsync(RequestContinuationToken.Create(groupByContinuationToken.SourceContinuationToken)); + tryCreateSource = await tryCreateSourceAsync(groupByContinuationToken.SourceContinuationToken); } if (!tryCreateSource.Succeeded) 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 e5cac473ee..f6e869e597 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/GroupBy/GroupByDocumentQueryExecutionComponent.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/GroupBy/GroupByDocumentQueryExecutionComponent.cs @@ -61,8 +61,8 @@ protected GroupByDocumentQueryExecutionComponent( public static async Task> TryCreateAsync( ExecutionEnvironment executionEnvironment, - RequestContinuationToken continuationToken, - Func>> tryCreateSourceAsync, + CosmosElement continuationToken, + Func>> tryCreateSourceAsync, IReadOnlyDictionary groupByAliasToAggregateType, IReadOnlyList orderedAliases, bool hasSelectValue) @@ -208,7 +208,7 @@ public void AddPayload(RewrittenGroupByProjection rewrittenGroupByProjection) this.groupByAliasToAggregateType, this.orderedAliases, this.hasSelectValue, - continuationToken: CosmosElementRequestContinuationToken.Null).Result; + continuationToken: null).Result; this.table[groupByKeysHash] = singleGroupAggregator; } @@ -306,7 +306,7 @@ public static TryCatch TryCreateFromContinuationToken( groupByAliasToAggregateType, orderedAliases, hasSelectValue, - RequestContinuationToken.Create(value)); + value); if (tryCreateSingleGroupAggregator.Succeeded) { diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/SkipDocumentQueryExecutionComponent.Client.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/SkipDocumentQueryExecutionComponent.Client.cs index 51fb9e1198..491920365c 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/SkipDocumentQueryExecutionComponent.Client.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/SkipDocumentQueryExecutionComponent.Client.cs @@ -28,24 +28,18 @@ private ClientSkipDocumentQueryExecutionComponent(IDocumentQueryExecutionCompone public static async Task> TryCreateAsync( int offsetCount, - RequestContinuationToken continuationToken, - Func>> tryCreateSourceAsync) + CosmosElement continuationToken, + Func>> tryCreateSourceAsync) { if (tryCreateSourceAsync == null) { throw new ArgumentNullException(nameof(tryCreateSourceAsync)); } - if (!(continuationToken is StringRequestContinuationToken stringRequestContinuationToken)) - { - return TryCatch.FromException( - new ArgumentException($"Expected {nameof(RequestContinuationToken)} to be a {nameof(StringRequestContinuationToken)}")); - } - OffsetContinuationToken offsetContinuationToken; - if (!continuationToken.IsNull) + if (continuationToken != null) { - if (!OffsetContinuationToken.TryParse(stringRequestContinuationToken.Value, out offsetContinuationToken)) + if (!OffsetContinuationToken.TryParse(continuationToken.ToString(), out offsetContinuationToken)) { return TryCatch.FromException( new MalformedContinuationTokenException($"Invalid {nameof(SkipDocumentQueryExecutionComponent)}: {continuationToken}.")); @@ -59,10 +53,24 @@ 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.")); + new MalformedContinuationTokenException("offset count in continuation token can not be greater than the offsetcount in the query.")); + } + + CosmosElement sourceToken; + if (offsetContinuationToken.SourceToken != null) + { + if (!CosmosElement.TryParse(offsetContinuationToken.SourceToken, out sourceToken)) + { + return TryCatch.FromException( + new MalformedContinuationTokenException("source token is not valid.")); + } + } + else + { + sourceToken = null; } - return (await tryCreateSourceAsync(RequestContinuationToken.Create(offsetContinuationToken.SourceToken))) + return (await tryCreateSourceAsync(sourceToken)) .Try((source) => new ClientSkipDocumentQueryExecutionComponent( source, offsetContinuationToken.Offset)); diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/SkipDocumentQueryExecutionComponent.Compute.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/SkipDocumentQueryExecutionComponent.Compute.cs index 8e27da989d..add3d09845 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/SkipDocumentQueryExecutionComponent.Compute.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/SkipDocumentQueryExecutionComponent.Compute.cs @@ -30,24 +30,18 @@ private ComputeSkipDocumentQueryExecutionComponent(IDocumentQueryExecutionCompon public static async Task> TryCreateAsync( int offsetCount, - RequestContinuationToken continuationToken, - Func>> tryCreateSourceAsync) + CosmosElement continuationToken, + Func>> tryCreateSourceAsync) { if (tryCreateSourceAsync == null) { throw new ArgumentNullException(nameof(tryCreateSourceAsync)); } - if (!(continuationToken is CosmosElementRequestContinuationToken cosmosElementRequestContinuationToken)) - { - return TryCatch.FromException( - new ArgumentException($"Expected {nameof(RequestContinuationToken)} to be a {nameof(StringRequestContinuationToken)}")); - } - OffsetContinuationToken offsetContinuationToken; - if (!continuationToken.IsNull) + if (continuationToken != null) { - (bool parsed, OffsetContinuationToken parsedToken) = OffsetContinuationToken.TryParse(cosmosElementRequestContinuationToken.Value); + (bool parsed, OffsetContinuationToken parsedToken) = OffsetContinuationToken.TryParse(continuationToken); if (!parsed) { return TryCatch.FromException( @@ -67,7 +61,7 @@ public static async Task> TryCreateAs new MalformedContinuationTokenException("offset count in continuation token can not be greater than the offsetcount in the query.")); } - return (await tryCreateSourceAsync(RequestContinuationToken.Create(offsetContinuationToken.SourceToken))) + return (await tryCreateSourceAsync(offsetContinuationToken.SourceToken)) .Try((source) => new ComputeSkipDocumentQueryExecutionComponent( source, offsetContinuationToken.Offset)); 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 331d27fc42..0153f6a1e6 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/SkipDocumentQueryExecutionComponent.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/SkipDocumentQueryExecutionComponent.cs @@ -6,6 +6,7 @@ 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; @@ -28,8 +29,8 @@ protected SkipDocumentQueryExecutionComponent(IDocumentQueryExecutionComponent s public static Task> TryCreateAsync( ExecutionEnvironment executionEnvironment, int offsetCount, - RequestContinuationToken continuationToken, - Func>> tryCreateSourceAsync) + CosmosElement continuationToken, + Func>> tryCreateSourceAsync) { Task> tryCreate; switch (executionEnvironment) diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/TakeDocumentQueryExecutionComponent.Client.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/TakeDocumentQueryExecutionComponent.Client.cs index 8c7bb9fdf0..8058ab2ec2 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/TakeDocumentQueryExecutionComponent.Client.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/TakeDocumentQueryExecutionComponent.Client.cs @@ -30,8 +30,8 @@ private ClientTakeDocumentQueryExecutionComponent(IDocumentQueryExecutionCompone public static async Task> TryCreateLimitDocumentQueryExecutionComponentAsync( int limitCount, - RequestContinuationToken requestContinuationToken, - Func>> tryCreateSourceAsync) + CosmosElement requestContinuationToken, + Func>> tryCreateSourceAsync) { if (limitCount < 0) { @@ -43,20 +43,10 @@ public static async Task> TryCreateLi throw new ArgumentNullException(nameof(tryCreateSourceAsync)); } - if (requestContinuationToken == null) - { - throw new ArgumentNullException(nameof(requestContinuationToken)); - } - - if (!(requestContinuationToken is StringRequestContinuationToken stringRequestContinuationToken)) - { - throw new ArgumentException($"Unknown {nameof(StringRequestContinuationToken)} type: {requestContinuationToken.GetType()}."); - } - LimitContinuationToken limitContinuationToken; - if (!requestContinuationToken.IsNull) + if (requestContinuationToken != null) { - if (!LimitContinuationToken.TryParse(stringRequestContinuationToken.Value, out limitContinuationToken)) + if (!LimitContinuationToken.TryParse(requestContinuationToken.ToString(), out limitContinuationToken)) { return TryCatch.FromException( new MalformedContinuationTokenException($"Malformed {nameof(LimitContinuationToken)}: {requestContinuationToken}.")); @@ -70,10 +60,24 @@ public static async Task> TryCreateLi if (limitContinuationToken.Limit > limitCount) { return TryCatch.FromException( - new MalformedContinuationTokenException($"{nameof(LimitContinuationToken.Limit)} in {nameof(LimitContinuationToken)}: {stringRequestContinuationToken.Value}: {limitContinuationToken.Limit} can not be greater than the limit count in the query: {limitCount}.")); + 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; + if (limitContinuationToken.SourceToken != null) + { + if (!CosmosElement.TryParse(limitContinuationToken.SourceToken, out sourceToken)) + { + return TryCatch.FromException( + new MalformedContinuationTokenException($"Malformed {nameof(LimitContinuationToken)}: {requestContinuationToken}.")); + } + } + else + { + sourceToken = null; } - return (await tryCreateSourceAsync(RequestContinuationToken.Create(limitContinuationToken.SourceToken))) + return (await tryCreateSourceAsync(sourceToken)) .Try((source) => new ClientTakeDocumentQueryExecutionComponent( source, limitContinuationToken.Limit, @@ -82,8 +86,8 @@ public static async Task> TryCreateLi public static async Task> TryCreateTopDocumentQueryExecutionComponentAsync( int topCount, - RequestContinuationToken requestContinuationToken, - Func>> tryCreateSourceAsync) + CosmosElement requestContinuationToken, + Func>> tryCreateSourceAsync) { if (topCount < 0) { @@ -95,23 +99,13 @@ public static async Task> TryCreateTo throw new ArgumentNullException(nameof(tryCreateSourceAsync)); } - if (requestContinuationToken == null) - { - throw new ArgumentNullException(nameof(requestContinuationToken)); - } - - if (!(requestContinuationToken is StringRequestContinuationToken stringRequestContinuationToken)) - { - throw new ArgumentException($"Unknown {nameof(StringRequestContinuationToken)} type: {requestContinuationToken.GetType()}."); - } - TopContinuationToken topContinuationToken; - if (!requestContinuationToken.IsNull) + if (requestContinuationToken != null) { - if (!TopContinuationToken.TryParse(stringRequestContinuationToken.Value, out topContinuationToken)) + if (!TopContinuationToken.TryParse(requestContinuationToken.ToString(), out topContinuationToken)) { return TryCatch.FromException( - new MalformedContinuationTokenException($"Malformed {nameof(LimitContinuationToken)}: {stringRequestContinuationToken.Value}.")); + new MalformedContinuationTokenException($"Malformed {nameof(LimitContinuationToken)}: {requestContinuationToken}.")); } } else @@ -122,10 +116,24 @@ public static async Task> TryCreateTo if (topContinuationToken.Top > topCount) { return TryCatch.FromException( - new MalformedContinuationTokenException($"{nameof(TopContinuationToken.Top)} in {nameof(TopContinuationToken)}: {stringRequestContinuationToken.Value}: {topContinuationToken.Top} can not be greater than the top count in the query: {topCount}.")); + 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; + if (topContinuationToken.SourceToken != null) + { + if (!CosmosElement.TryParse(topContinuationToken.SourceToken, out sourceToken)) + { + return TryCatch.FromException( + new MalformedContinuationTokenException($"{nameof(TopContinuationToken.SourceToken)} in {nameof(TopContinuationToken)}: {requestContinuationToken}: {topContinuationToken.SourceToken} was malformed.")); + } + } + else + { + sourceToken = null; } - return (await tryCreateSourceAsync(RequestContinuationToken.Create(topContinuationToken.SourceToken))) + return (await tryCreateSourceAsync(sourceToken)) .Try((source) => new ClientTakeDocumentQueryExecutionComponent( source, topContinuationToken.Top, diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/TakeDocumentQueryExecutionComponent.Compute.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/TakeDocumentQueryExecutionComponent.Compute.cs index c32fb08b8e..40996a41de 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/TakeDocumentQueryExecutionComponent.Compute.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/TakeDocumentQueryExecutionComponent.Compute.cs @@ -10,7 +10,6 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.SkipTake 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; @@ -30,8 +29,8 @@ private ComputeTakeDocumentQueryExecutionComponent(IDocumentQueryExecutionCompon public static async Task> TryCreateAsync( int takeCount, - RequestContinuationToken requestContinuationToken, - Func>> tryCreateSourceAsync) + CosmosElement requestContinuationToken, + Func>> tryCreateSourceAsync) { if (takeCount < 0) { @@ -43,23 +42,13 @@ public static async Task> TryCreateAs throw new ArgumentNullException(nameof(tryCreateSourceAsync)); } - if (requestContinuationToken == null) - { - throw new ArgumentNullException(nameof(requestContinuationToken)); - } - - if (!(requestContinuationToken is CosmosElementRequestContinuationToken cosmosElementRequestContinuationToken)) - { - throw new ArgumentException($"Unknown {nameof(RequestContinuationToken)} type: {requestContinuationToken.GetType()}."); - } - TakeContinuationToken takeContinuationToken; - if (!requestContinuationToken.IsNull) + if (requestContinuationToken != null) { - if (!TakeContinuationToken.TryParse(cosmosElementRequestContinuationToken.Value, out takeContinuationToken)) + if (!TakeContinuationToken.TryParse(requestContinuationToken, out takeContinuationToken)) { return TryCatch.FromException( - new MalformedContinuationTokenException($"Malformed {nameof(TakeContinuationToken)}: {cosmosElementRequestContinuationToken.Value}.")); + new MalformedContinuationTokenException($"Malformed {nameof(TakeContinuationToken)}: {requestContinuationToken}.")); } } else @@ -70,10 +59,10 @@ public static async Task> TryCreateAs if (takeContinuationToken.TakeCount > takeCount) { return TryCatch.FromException( - new MalformedContinuationTokenException($"{nameof(TakeContinuationToken.TakeCount)} in {nameof(TakeContinuationToken)}: {cosmosElementRequestContinuationToken.Value}: {takeContinuationToken.TakeCount} can not be greater than the limit count in the query: {takeCount}.")); + new MalformedContinuationTokenException($"{nameof(TakeContinuationToken.TakeCount)} in {nameof(TakeContinuationToken)}: {requestContinuationToken}: {takeContinuationToken.TakeCount} can not be greater than the limit count in the query: {takeCount}.")); } - return (await tryCreateSourceAsync(RequestContinuationToken.Create(takeContinuationToken.SourceToken))) + return (await tryCreateSourceAsync(takeContinuationToken.SourceToken)) .Try((source) => new ComputeTakeDocumentQueryExecutionComponent( source, takeContinuationToken.TakeCount)); 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 51655109cb..2aa61c01db 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/TakeDocumentQueryExecutionComponent.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/TakeDocumentQueryExecutionComponent.cs @@ -34,8 +34,8 @@ protected TakeDocumentQueryExecutionComponent( public static Task> TryCreateLimitDocumentQueryExecutionComponentAsync( ExecutionEnvironment executionEnvironment, int limitCount, - RequestContinuationToken requestContinuationToken, - Func>> tryCreateSourceAsync) + CosmosElement requestContinuationToken, + Func>> tryCreateSourceAsync) { Task> tryCreateComponentAsync; switch (executionEnvironment) @@ -64,8 +64,8 @@ public static Task> TryCreateLimitDoc public static Task> TryCreateTopDocumentQueryExecutionComponentAsync( ExecutionEnvironment executionEnvironment, int topCount, - RequestContinuationToken requestContinuationToken, - Func>> tryCreateSourceAsync) + CosmosElement requestContinuationToken, + Func>> tryCreateSourceAsync) { Task> tryCreateComponentAsync; switch (executionEnvironment) diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosQueryExecutionContextFactory.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosQueryExecutionContextFactory.cs index 1b0ed4e861..3bf8ee4dd1 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosQueryExecutionContextFactory.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosQueryExecutionContextFactory.cs @@ -64,19 +64,12 @@ private static async Task> TryCreateCoreCo CancellationToken cancellationToken) { // Try to parse the continuation token. - RequestContinuationToken continuationToken = inputParameters.InitialUserContinuationToken; + CosmosElement continuationToken = inputParameters.InitialUserContinuationToken; PartitionedQueryExecutionInfo queryPlanFromContinuationToken = inputParameters.PartitionedQueryExecutionInfo; - if (!continuationToken.IsNull) + if (continuationToken != null) { - if (!continuationToken.TryConvertToCosmosElement(out CosmosElement cosmosElement)) - { - return TryCatch.FromException( - new MalformedContinuationTokenException( - $"Malformed {nameof(PipelineContinuationToken)}.")); - } - - if (!PipelineContinuationToken.TryParse( - cosmosElement.ToString(), + if (!PipelineContinuationToken.TryCreateFromCosmosElement( + continuationToken, out PipelineContinuationToken pipelineContinuationToken)) { return TryCatch.FromException( @@ -102,7 +95,7 @@ private static async Task> TryCreateCoreCo $"{nameof(PipelineContinuationToken)}: '{continuationToken}' is no longer supported.")); } - continuationToken = RequestContinuationToken.Create(latestVersionPipelineContinuationToken.SourceContinuationToken); + continuationToken = latestVersionPipelineContinuationToken.SourceContinuationToken; if (latestVersionPipelineContinuationToken.QueryPlan != null) { queryPlanFromContinuationToken = latestVersionPipelineContinuationToken.QueryPlan; @@ -354,7 +347,7 @@ public sealed class InputParameters public InputParameters( SqlQuerySpec sqlQuerySpec, - RequestContinuationToken initialUserContinuationToken, + CosmosElement initialUserContinuationToken, int? maxConcurrency, int? maxItemCount, int? maxBufferedItemCount, @@ -398,7 +391,7 @@ public InputParameters( } public SqlQuerySpec SqlQuerySpec { get; } - public RequestContinuationToken InitialUserContinuationToken { get; } + public CosmosElement InitialUserContinuationToken { get; } public int MaxConcurrency { get; } public int MaxItemCount { get; } public int MaxBufferedItemCount { 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 52b514a56c..7b13b99405 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/OrderBy/CosmosOrderByItemQueryExecutionContext.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/OrderBy/CosmosOrderByItemQueryExecutionContext.cs @@ -194,7 +194,7 @@ public override void SerializeState(IJsonWriter jsonWriter) public static async Task> TryCreateAsync( CosmosQueryContext queryContext, CosmosCrossPartitionQueryExecutionContext.CrossPartitionInitParams initParams, - RequestContinuationToken requestContinuationToken, + CosmosElement requestContinuationToken, CancellationToken cancellationToken) { Debug.Assert( @@ -362,7 +362,7 @@ private bool ShouldIncrementSkipCount(ItemProducer currentItemProducer) private async Task TryInitializeAsync( SqlQuerySpec sqlQuerySpec, - RequestContinuationToken requestContinuation, + CosmosElement requestContinuation, string collectionRid, List partitionKeyRanges, int initialPageSize, @@ -375,11 +375,6 @@ private async Task TryInitializeAsync( throw new ArgumentNullException(nameof(sqlQuerySpec)); } - if (requestContinuation == null) - { - throw new ArgumentNullException(nameof(requestContinuation)); - } - if (collectionRid == null) { throw new ArgumentNullException(nameof(collectionRid)); @@ -402,7 +397,7 @@ private async Task TryInitializeAsync( cancellationToken.ThrowIfCancellationRequested(); - if (requestContinuation.IsNull) + if (requestContinuation == null) { SqlQuerySpec sqlQuerySpecForInit = new SqlQuerySpec( sqlQuerySpec.QueryText.Replace(oldValue: FormatPlaceHolder, newValue: True), @@ -513,7 +508,7 @@ private async Task TryInitializeAsync( } private static TryCatch TryExtractContinuationTokens( - RequestContinuationToken requestContinuation, + CosmosElement requestContinuation, SortOrder[] sortOrders, string[] orderByExpressions) { @@ -530,31 +525,7 @@ private static TryCatch TryExtractContinuationTokens throw new ArgumentNullException("continuation can not be null or empty."); } - if (requestContinuation.IsNull) - { - throw new ArgumentNullException("continuation can not be null or empty."); - } - - CosmosElement cosmosElement; - switch (requestContinuation) - { - case StringRequestContinuationToken stringRequestContinuationToken: - if (!CosmosElement.TryParse(stringRequestContinuationToken.Value, out cosmosElement)) - { - return TryCatch.FromException( - new MalformedContinuationTokenException($"Order by continuation token is malformed: {stringRequestContinuationToken.Value}.")); - } - break; - - case CosmosElementRequestContinuationToken cosmosElementRequestContinuation: - cosmosElement = cosmosElementRequestContinuation.Value; - break; - - default: - throw new ArgumentException($"Unknown {nameof(RequestContinuationToken)} type: {requestContinuation.GetType()}."); - } - - if (!(cosmosElement is CosmosArray cosmosArray)) + if (!(requestContinuation is CosmosArray cosmosArray)) { return TryCatch.FromException( new MalformedContinuationTokenException($"Order by continuation token must be an array: {requestContinuation}.")); 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 b558ae2be9..83a9ba193c 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/Parallel/CosmosParallelItemQueryExecutionContext.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/Parallel/CosmosParallelItemQueryExecutionContext.cs @@ -142,7 +142,7 @@ public override void SerializeState(IJsonWriter jsonWriter) public static async Task> TryCreateAsync( CosmosQueryContext queryContext, CosmosCrossPartitionQueryExecutionContext.CrossPartitionInitParams initParams, - RequestContinuationToken requestContinuationToken, + CosmosElement requestContinuationToken, CancellationToken cancellationToken) { Debug.Assert( @@ -243,7 +243,7 @@ private async Task> TryInitial string collectionRid, List partitionKeyRanges, int initialPageSize, - RequestContinuationToken requestContinuation, + CosmosElement requestContinuation, CancellationToken cancellationToken) { cancellationToken.ThrowIfCancellationRequested(); @@ -287,107 +287,52 @@ private async Task> TryInitial /// The subset of partition to actually target and continuation tokens. private static TryCatch TryGetInitializationInfoFromContinuationToken( List partitionKeyRanges, - RequestContinuationToken continuationToken) + CosmosElement continuationToken) { - if (continuationToken.IsNull) + if (continuationToken == null) { return TryCatch.FromResult( new ParallelInitInfo( partitionKeyRanges, null)); } - else - { - CosmosElement cosmosElement; - switch (continuationToken) - { - case StringRequestContinuationToken stringRequestContinuationToken: - if (!CosmosElement.TryParse(stringRequestContinuationToken.Value, out cosmosElement)) - { - return TryCatch.FromException( - new MalformedContinuationTokenException($"Order by continuation token is malformed: {stringRequestContinuationToken.Value}.")); - } - break; - - case CosmosElementRequestContinuationToken cosmosElementRequestContinuation: - cosmosElement = cosmosElementRequestContinuation.Value; - break; - - default: - throw new ArgumentException($"Unknown {nameof(RequestContinuationToken)} type: {continuationToken.GetType()}."); - } - - if (!TryParseContinuationToken(cosmosElement, out List tokens)) - { - return TryCatch.FromException( - new MalformedContinuationTokenException($"Invalid format for continuation token {continuationToken} for {nameof(CosmosParallelItemQueryExecutionContext)}")); - } - return CosmosCrossPartitionQueryExecutionContext.TryFindTargetRangeAndExtractContinuationTokens( - partitionKeyRanges, - tokens.Select(token => Tuple.Create(token, token.Range))) - .Try((indexAndTokens) => - { - int minIndex = indexAndTokens.TargetIndex; - IReadOnlyDictionary rangeToToken = indexAndTokens.ContinuationTokens; - - // We know that all partitions to the left of the continuation token are fully drained so we can filter them out - IReadOnlyList filteredRanges = new PartialReadOnlyList( - partitionKeyRanges, - minIndex, - partitionKeyRanges.Count - minIndex); - - return new ParallelInitInfo( - filteredRanges, - rangeToToken); - }); - } - } - - private static bool TryParseContinuationToken(CosmosElement continuationToken, out List tokens) - { - if (continuationToken == null) - { - throw new ArgumentNullException(nameof(continuationToken)); - } - - if (!(continuationToken is CosmosArray cosmosArray)) + if (!(continuationToken is CosmosArray compositeContinuationTokenListRaw)) { - tokens = default; - return false; + return TryCatch.FromException( + new MalformedContinuationTokenException($"Invalid format for continuation token {continuationToken} for {nameof(CosmosParallelItemQueryExecutionContext)}")); } List compositeContinuationTokens = new List(); - foreach (CosmosElement item in cosmosArray) + foreach (CosmosElement compositeContinuationTokenRaw in compositeContinuationTokenListRaw) { - TryCatch tryCreateCompositeContinuationToken = CompositeContinuationToken.TryCreateFromCosmosElement(item); + TryCatch tryCreateCompositeContinuationToken = CompositeContinuationToken.TryCreateFromCosmosElement(compositeContinuationTokenRaw); if (!tryCreateCompositeContinuationToken.Succeeded) { - tokens = default; - return false; + return TryCatch.FromException(tryCreateCompositeContinuationToken.Exception); } compositeContinuationTokens.Add(tryCreateCompositeContinuationToken.Result); } - tokens = compositeContinuationTokens; - - if (tokens.Count == 0) - { - tokens = default; - return false; - } - - foreach (CompositeContinuationToken token in tokens) - { - if ((token.Range == null) || token.Range.IsEmpty) + return CosmosCrossPartitionQueryExecutionContext.TryFindTargetRangeAndExtractContinuationTokens( + partitionKeyRanges, + compositeContinuationTokens.Select(token => Tuple.Create(token, token.Range))) + .Try((indexAndTokens) => { - tokens = default; - return false; - } - } + int minIndex = indexAndTokens.TargetIndex; + IReadOnlyDictionary rangeToToken = indexAndTokens.ContinuationTokens; + + // We know that all partitions to the left of the continuation token are fully drained so we can filter them out + IReadOnlyList filteredRanges = new PartialReadOnlyList( + partitionKeyRanges, + minIndex, + partitionKeyRanges.Count - minIndex); - return true; + return new ParallelInitInfo( + filteredRanges, + rangeToToken); + }); } private readonly struct ParallelInitInfo diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/PipelinedDocumentQueryExecutionContext.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/PipelinedDocumentQueryExecutionContext.cs index 3f02a9df1e..72a351baf0 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/PipelinedDocumentQueryExecutionContext.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/PipelinedDocumentQueryExecutionContext.cs @@ -134,7 +134,7 @@ public static async Task> TryCreateAsync( ExecutionEnvironment executionEnvironment, CosmosQueryContext queryContext, CosmosCrossPartitionQueryExecutionContext.CrossPartitionInitParams initParams, - RequestContinuationToken requestContinuationToken, + CosmosElement requestContinuationToken, CancellationToken cancellationToken) { if (queryContext == null) @@ -163,7 +163,7 @@ public static async Task> TryCreateAsync( testSettings: initParams.TestSettings); } - Task> tryCreateOrderByComponentAsync(RequestContinuationToken continuationToken) + Task> tryCreateOrderByComponentAsync(CosmosElement continuationToken) { return CosmosOrderByItemQueryExecutionContext.TryCreateAsync( queryContext, @@ -172,7 +172,7 @@ Task> tryCreateOrderByComponentAsync( cancellationToken); } - Task> tryCreateParallelComponentAsync(RequestContinuationToken continuationToken) + Task> tryCreateParallelComponentAsync(CosmosElement continuationToken) { return CosmosParallelItemQueryExecutionContext.TryCreateAsync( queryContext, @@ -181,7 +181,7 @@ Task> tryCreateParallelComponentAsync cancellationToken); } - Func>> tryCreatePipelineAsync; + Func>> tryCreatePipelineAsync; if (queryInfo.HasOrderBy) { tryCreatePipelineAsync = tryCreateOrderByComponentAsync; @@ -193,7 +193,7 @@ Task> tryCreateParallelComponentAsync if (queryInfo.HasAggregates && !queryInfo.HasGroupBy) { - Func>> tryCreateSourceAsync = tryCreatePipelineAsync; + Func>> tryCreateSourceAsync = tryCreatePipelineAsync; tryCreatePipelineAsync = async (continuationToken) => { return await AggregateDocumentQueryExecutionComponent.TryCreateAsync( @@ -209,7 +209,7 @@ Task> tryCreateParallelComponentAsync if (queryInfo.HasDistinct) { - Func>> tryCreateSourceAsync = tryCreatePipelineAsync; + Func>> tryCreateSourceAsync = tryCreatePipelineAsync; tryCreatePipelineAsync = async (continuationToken) => { return await DistinctDocumentQueryExecutionComponent.TryCreateAsync( @@ -222,7 +222,7 @@ Task> tryCreateParallelComponentAsync if (queryInfo.HasGroupBy) { - Func>> tryCreateSourceAsync = tryCreatePipelineAsync; + Func>> tryCreateSourceAsync = tryCreatePipelineAsync; tryCreatePipelineAsync = async (continuationToken) => { return await GroupByDocumentQueryExecutionComponent.TryCreateAsync( @@ -237,7 +237,7 @@ Task> tryCreateParallelComponentAsync if (queryInfo.HasOffset) { - Func>> tryCreateSourceAsync = tryCreatePipelineAsync; + Func>> tryCreateSourceAsync = tryCreatePipelineAsync; tryCreatePipelineAsync = async (continuationToken) => { return await SkipDocumentQueryExecutionComponent.TryCreateAsync( @@ -250,7 +250,7 @@ Task> tryCreateParallelComponentAsync if (queryInfo.HasLimit) { - Func>> tryCreateSourceAsync = tryCreatePipelineAsync; + Func>> tryCreateSourceAsync = tryCreatePipelineAsync; tryCreatePipelineAsync = async (continuationToken) => { return await TakeDocumentQueryExecutionComponent.TryCreateLimitDocumentQueryExecutionComponentAsync( @@ -263,7 +263,7 @@ Task> tryCreateParallelComponentAsync if (queryInfo.HasTop) { - Func>> tryCreateSourceAsync = tryCreatePipelineAsync; + Func>> tryCreateSourceAsync = tryCreatePipelineAsync; tryCreatePipelineAsync = async (continuationToken) => { return await TakeDocumentQueryExecutionComponent.TryCreateTopDocumentQueryExecutionComponentAsync( @@ -326,7 +326,7 @@ public override async Task ExecuteNextAsync(CancellationToken { if (queryResponse.ContinuationToken != null) { - updatedContinuationToken = new PipelineContinuationTokenV0(queryResponse.ContinuationToken).ToString(); + updatedContinuationToken = new PipelineContinuationTokenV0(CosmosElement.Parse(queryResponse.ContinuationToken)).ToString(); } else { diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/QueryExecutionContextWithException.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/QueryExecutionContextWithException.cs new file mode 100644 index 0000000000..520cea9d85 --- /dev/null +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/QueryExecutionContextWithException.cs @@ -0,0 +1,47 @@ +// ------------------------------------------------------------ +// 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.Json; + using Microsoft.Azure.Cosmos.Query.Core; + using Microsoft.Azure.Cosmos.Query.Core.ExecutionContext; + 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 void SerializeState(IJsonWriter jsonWriter) + { + throw new NotImplementedException(); + } + } +} diff --git a/Microsoft.Azure.Cosmos/src/Query/v3Query/QueryIterator.cs b/Microsoft.Azure.Cosmos/src/Query/v3Query/QueryIterator.cs index 231de1aa4f..3bf46e23c9 100644 --- a/Microsoft.Azure.Cosmos/src/Query/v3Query/QueryIterator.cs +++ b/Microsoft.Azure.Cosmos/src/Query/v3Query/QueryIterator.cs @@ -11,9 +11,8 @@ namespace Microsoft.Azure.Cosmos.Query using Microsoft.Azure.Cosmos.Diagnostics; using Microsoft.Azure.Cosmos.Json; 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; - using Microsoft.Azure.Cosmos.Query.Core.Metrics; using Microsoft.Azure.Cosmos.Query.Core.QueryClient; using Microsoft.Azure.Cosmos.Query.Core.QueryPlan; @@ -62,36 +61,35 @@ public static QueryIterator Create( allowNonValueAggregateQuery: allowNonValueAggregateQuery, correlatedActivityId: Guid.NewGuid()); - RequestContinuationToken requestContinuationToken; - if (queryRequestOptions.ExecutionEnvironment.HasValue) + CosmosElement requestContinuationToken; + switch (queryRequestOptions.ExecutionEnvironment.GetValueOrDefault(ExecutionEnvironment.Client)) { - switch (queryRequestOptions.ExecutionEnvironment.Value) - { - case ExecutionEnvironment.Client: - requestContinuationToken = RequestContinuationToken.Create(continuationToken); - break; - - case ExecutionEnvironment.Compute: - CosmosElement cosmosElement; - if (queryRequestOptions.BinaryContinuationToken.IsEmpty) + case ExecutionEnvironment.Client: + if (continuationToken != null) + { + if (!CosmosElement.TryParse(continuationToken, out requestContinuationToken)) { - cosmosElement = null; + return new QueryIterator( + cosmosQueryContext, + new QueryExecutionContextWithException( + new MalformedContinuationTokenException( + $"Malformed Continuation Token: {requestContinuationToken}")), + queryRequestOptions.CosmosSerializationFormatOptions, + queryRequestOptions); } - else - { - cosmosElement = CosmosElement.CreateFromBuffer(queryRequestOptions.BinaryContinuationToken); - } - - requestContinuationToken = RequestContinuationToken.Create(cosmosElement); - break; - - default: - throw new ArgumentOutOfRangeException($"Unknown {nameof(ExecutionEnvironment)}: {queryRequestOptions.ExecutionEnvironment.Value}."); - } - } - else - { - requestContinuationToken = RequestContinuationToken.Create(continuationToken); + } + else + { + requestContinuationToken = null; + } + break; + + case ExecutionEnvironment.Compute: + requestContinuationToken = queryRequestOptions.CosmosElementContinuationToken; + break; + + default: + throw new ArgumentOutOfRangeException($"Unknown {nameof(ExecutionEnvironment)}: {queryRequestOptions.ExecutionEnvironment.Value}."); } CosmosQueryExecutionContextFactory.InputParameters inputParameters = new CosmosQueryExecutionContextFactory.InputParameters( diff --git a/Microsoft.Azure.Cosmos/src/RequestOptions/QueryRequestOptions.cs b/Microsoft.Azure.Cosmos/src/RequestOptions/QueryRequestOptions.cs index dd4196353f..1345dc51e8 100644 --- a/Microsoft.Azure.Cosmos/src/RequestOptions/QueryRequestOptions.cs +++ b/Microsoft.Azure.Cosmos/src/RequestOptions/QueryRequestOptions.cs @@ -6,6 +6,7 @@ namespace Microsoft.Azure.Cosmos { using System; using System.Globalization; + using Microsoft.Azure.Cosmos.CosmosElements; using Microsoft.Azure.Cosmos.Query.Core; using Microsoft.Azure.Cosmos.Query.Core.ExecutionContext; using Microsoft.Azure.Documents; @@ -141,7 +142,7 @@ public ConsistencyLevel? ConsistencyLevel /// public string SessionToken { get; set; } - internal ReadOnlyMemory BinaryContinuationToken { get; set; } + internal CosmosElement CosmosElementContinuationToken { get; set; } internal CosmosSerializationFormatOptions CosmosSerializationFormatOptions { get; set; } @@ -224,7 +225,7 @@ internal QueryRequestOptions Clone() CosmosSerializationFormatOptions = this.CosmosSerializationFormatOptions, Properties = this.Properties, IsEffectivePartitionKeyRouting = this.IsEffectivePartitionKeyRouting, - BinaryContinuationToken = this.BinaryContinuationToken, + CosmosElementContinuationToken = this.CosmosElementContinuationToken, }; return queryRequestOptions; diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CrossPartitionQueryTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CrossPartitionQueryTests.cs index 0fa3639cc7..9408e99f71 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CrossPartitionQueryTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CrossPartitionQueryTests.cs @@ -597,12 +597,12 @@ private static async Task> QueryWithSerializeState( } List resultsFromSerializeState = new List(); - ReadOnlyMemory continuationToken = default(ReadOnlyMemory); + CosmosElement continuationToken = null; do { QueryRequestOptions computeRequestOptions = queryRequestOptions.Clone(); computeRequestOptions.ExecutionEnvironment = Cosmos.Query.Core.ExecutionContext.ExecutionEnvironment.Compute; - computeRequestOptions.BinaryContinuationToken = continuationToken; + computeRequestOptions.CosmosElementContinuationToken = continuationToken; FeedIteratorInternal itemQuery = (FeedIteratorInternal)container.GetItemQueryIterator( queryText: query, @@ -620,11 +620,15 @@ private static async Task> QueryWithSerializeState( resultsFromSerializeState.AddRange(cosmosQueryResponse); Json.IJsonWriter jsonWriter = Json.JsonWriter.Create(Json.JsonSerializationFormat.Binary); itemQuery.SerializeState(jsonWriter); - continuationToken = jsonWriter.GetResult(); + ReadOnlyMemory continuationTokenBuffer = jsonWriter.GetResult(); - if (!continuationToken.IsEmpty) + if (!continuationTokenBuffer.IsEmpty) { - string stringContinuationToken = CosmosElement.CreateFromBuffer(continuationToken).ToString(); + continuationToken = CosmosElement.CreateFromBuffer(continuationTokenBuffer); + } + else + { + continuationToken = null; } } catch (CosmosException cosmosException) when (cosmosException.StatusCode == (HttpStatusCode)429) @@ -633,7 +637,7 @@ private static async Task> QueryWithSerializeState( queryText: query, requestOptions: queryRequestOptions); } - } while (!continuationToken.IsEmpty); + } while (continuationToken != null); return resultsFromSerializeState; } 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 1783bcb0d5..4c19344c7c 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/CosmosQueryUnitTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/CosmosQueryUnitTests.cs @@ -256,7 +256,7 @@ public async Task TestCosmosQueryPartitionKeyDefinition() CosmosQueryExecutionContextFactory.InputParameters inputParameters = new CosmosQueryExecutionContextFactory.InputParameters( sqlQuerySpec: sqlQuerySpec, - initialUserContinuationToken: StringRequestContinuationToken.Null, + initialUserContinuationToken: null, maxConcurrency: queryRequestOptions?.MaxConcurrency, maxItemCount: queryRequestOptions?.MaxItemCount, maxBufferedItemCount: queryRequestOptions?.MaxBufferedItemCount, @@ -289,7 +289,7 @@ public async Task TestCosmosQueryPartitionKeyDefinition() private async Task<(IList components, QueryResponseCore response)> GetAllExecutionComponents() { - (Func>> func, QueryResponseCore response) = this.SetupBaseContextToVerifyFailureScenario(); + (Func>> func, QueryResponseCore response) = this.SetupBaseContextToVerifyFailureScenario(); List components = new List(); List operators = new List() @@ -310,37 +310,37 @@ public async Task TestCosmosQueryPartitionKeyDefinition() }, new List() { "test" }, false, - StringRequestContinuationToken.Null, + null, func)).Result); components.Add((await DistinctDocumentQueryExecutionComponent.TryCreateAsync( ExecutionEnvironment.Client, - StringRequestContinuationToken.Null, + null, func, DistinctQueryType.Ordered)).Result); components.Add((await SkipDocumentQueryExecutionComponent.TryCreateAsync( ExecutionEnvironment.Client, 5, - StringRequestContinuationToken.Null, + null, func)).Result); components.Add((await TakeDocumentQueryExecutionComponent.TryCreateLimitDocumentQueryExecutionComponentAsync( ExecutionEnvironment.Client, 5, - StringRequestContinuationToken.Null, + null, func)).Result); components.Add((await TakeDocumentQueryExecutionComponent.TryCreateTopDocumentQueryExecutionComponentAsync( ExecutionEnvironment.Client, 5, - StringRequestContinuationToken.Null, + null, func)).Result); return (components, response); } - private (Func>>, QueryResponseCore) SetupBaseContextToVerifyFailureScenario() + private (Func>>, QueryResponseCore) SetupBaseContextToVerifyFailureScenario() { CosmosDiagnosticsContext diagnosticsContext = CosmosDiagnosticsContext.Create(); diagnosticsContext.AddDiagnosticsInternal( new PointOperationStatistics( @@ -374,7 +374,7 @@ public async Task TestCosmosQueryPartitionKeyDefinition() Mock baseContext = new Mock(); baseContext.Setup(x => x.DrainAsync(It.IsAny(), It.IsAny())).Returns(Task.FromResult(failure)); - Task> callBack(RequestContinuationToken x) => Task.FromResult>(TryCatch.FromResult(baseContext.Object)); + Task> callBack(CosmosElement x) => Task.FromResult>(TryCatch.FromResult(baseContext.Object)); return (callBack, failure); } } 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 7da61c738e..b8f5efbd7d 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 @@ -6,6 +6,7 @@ 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.Test.BaselineTest; using VisualStudio.TestTools.UnitTesting; @@ -72,13 +73,17 @@ public void Tests() public override PipelineContinuationTokenTestsOutput ExecuteTest( PipelineContinuationTokenTestsInput input) { - if (!PipelineContinuationToken.TryParse( - input.ContinuationToken, - out PipelineContinuationToken pipelineContinuationToken)) + if (!CosmosElement.TryParse(input.ContinuationToken, out CosmosElement cosmosElementContinuationToken)) { return new PipelineContinuationTokenTestsOutputNegative("Failed to parse token."); } + if (!PipelineContinuationToken.TryCreateFromCosmosElement( + cosmosElementContinuationToken, + out PipelineContinuationToken pipelineContinuationToken)) + { + return new PipelineContinuationTokenTestsOutputNegative("Failed to parse token."); + } if (!PipelineContinuationToken.TryConvertToLatest( pipelineContinuationToken, diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/QueryPipelineMockTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/QueryPipelineMockTests.cs index 0e872d7dfc..566d8473bc 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/QueryPipelineMockTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/QueryPipelineMockTests.cs @@ -94,7 +94,7 @@ public async Task TestCosmosCrossPartitionQueryExecutionContextWithEmptyPagesAnd IDocumentQueryExecutionComponent executionContext = (await CosmosParallelItemQueryExecutionContext.TryCreateAsync( context, initParams, - RequestContinuationToken.Create(fullConitnuationToken), + CosmosElement.Parse(fullConitnuationToken), this.cancellationToken)).Result; // Read all the pages from both splits @@ -182,7 +182,7 @@ public async Task TestCosmosCrossPartitionQueryExecutionContextWithFailuresAsync IDocumentQueryExecutionComponent executionContext = (await CosmosParallelItemQueryExecutionContext.TryCreateAsync( context, initParams, - RequestContinuationToken.Create(fullConitnuationToken), + CosmosElement.Parse(fullConitnuationToken), this.cancellationToken)).Result; // Read all the pages from both splits @@ -315,7 +315,7 @@ public async Task TestCosmosOrderByQueryExecutionContextWithEmptyPagesAndSplitAs IDocumentQueryExecutionComponent executionContext = (await CosmosOrderByItemQueryExecutionContext.TryCreateAsync( context, initParams, - RequestContinuationToken.Create(fullConitnuationToken), + CosmosElement.Parse(fullConitnuationToken), this.cancellationToken)).Result; // For order by it will drain all the pages till it gets a value. @@ -442,7 +442,7 @@ public async Task TestCosmosOrderByQueryExecutionContextWithFailurePageAsync(boo TryCatch tryCreate = await CosmosOrderByItemQueryExecutionContext.TryCreateAsync( context, initParams, - RequestContinuationToken.Create(fullContinuationToken), + CosmosElement.Parse(fullContinuationToken), this.cancellationToken); if (tryCreate.Succeeded) @@ -523,7 +523,7 @@ public async Task TestNegativeDistinctComponentCreation() { TryCatch tryCreateWhenSourceFails = await DistinctDocumentQueryExecutionComponent.TryCreateAsync( ExecutionEnvironment.Client, - StringRequestContinuationToken.Null, + null, FailToCreateSource, DistinctQueryType.Ordered); @@ -531,7 +531,7 @@ public async Task TestNegativeDistinctComponentCreation() TryCatch tryCreateWhenInvalidContinuationToken = await DistinctDocumentQueryExecutionComponent.TryCreateAsync( ExecutionEnvironment.Client, - RequestContinuationToken.Create("This is not a valid continuation token"), + CosmosElement.Parse("This is not a valid continuation token"), CreateSource, DistinctQueryType.Unordered); @@ -540,12 +540,12 @@ public async Task TestNegativeDistinctComponentCreation() // ADD MORE TESTS HERE - private static Task> FailToCreateSource(RequestContinuationToken continuationToken) + private static Task> FailToCreateSource(CosmosElement continuationToken) { return Task.FromResult(TryCatch.FromException(new Exception())); } - private static Task> CreateSource(RequestContinuationToken continuationToken) + private static Task> CreateSource(CosmosElement continuationToken) { return Task.FromResult(TryCatch.FromResult(new Mock().Object)); } From 8d8279fc2a4676aa77c3c87ce8be643c1ee04cdb Mon Sep 17 00:00:00 2001 From: Brandon Chong Date: Fri, 28 Feb 2020 09:12:31 -0800 Subject: [PATCH 16/28] returning continuation token instead --- .../src/FeedIteratorCore.cs | 14 +- .../src/FeedIteratorInternal.cs | 3 +- .../src/FeedIteratorInternal{T}.cs | 3 +- ...eDocumentQueryExecutionComponent.Client.cs | 2 +- ...DocumentQueryExecutionComponent.Compute.cs | 47 ++-- .../Aggregators/AverageAggregator.cs | 20 +- .../Aggregate/Aggregators/CountAggregator.cs | 15 +- .../Aggregate/Aggregators/IAggregator.cs | 8 +- .../Aggregate/Aggregators/MinMaxAggregator.cs | 215 ++++++++++++++---- .../Aggregators/MinMaxContinuationToken.cs | 165 -------------- .../Aggregators/SingleGroupAggregator.cs | 80 ++----- .../Aggregate/Aggregators/SumAggregator.cs | 5 +- ...tDocumentQueryExecutionComponent.Client.cs | 2 +- ...DocumentQueryExecutionComponent.Compute.cs | 39 +++- ...DistinctDocumentQueryExecutionComponent.cs | 2 - .../DistinctMap.OrderedDistinctMap.cs | 9 +- .../DistinctMap.UnorderedDistinctMap.cs | 115 ++++------ .../Distinct/DistinctMap.cs | 2 +- .../DocumentQueryExecutionComponentBase.cs | 3 +- ...yDocumentQueryExecutionComponent.Client.cs | 5 +- ...DocumentQueryExecutionComponent.Compute.cs | 73 +++--- .../GroupByDocumentQueryExecutionComponent.cs | 23 +- .../IDocumentQueryExecutionComponent.cs | 2 +- ...pDocumentQueryExecutionComponent.Client.cs | 3 +- ...DocumentQueryExecutionComponent.Compute.cs | 49 ++-- ...eDocumentQueryExecutionComponent.Client.cs | 4 +- ...DocumentQueryExecutionComponent.Compute.cs | 44 ++-- .../CatchAllCosmosQueryExecutionContext.cs | 10 +- ...smosCrossPartitionQueryExecutionContext.cs | 2 +- .../CosmosQueryExecutionContext.cs | 3 +- ...ExecutionContextWithNameCacheStaleRetry.cs | 10 +- .../LazyCosmosQueryExecutionContext.cs | 10 +- .../CosmosOrderByItemQueryExecutionContext.cs | 71 +++--- ...CosmosParallelItemQueryExecutionContext.cs | 59 +++-- .../PipelinedDocumentQueryExecutionContext.cs | 15 +- .../QueryExecutionContextWithException.cs | 5 +- .../Query/v3Query/FeedIteratorInlineCore.cs | 17 +- .../src/Query/v3Query/QueryIterator.cs | 4 +- ...geFeedPartitionKeyResultSetIteratorCore.cs | 11 +- .../ChangeFeedResultSetIteratorCore.cs | 4 +- .../CrossPartitionQueryTests.cs | 13 +- 41 files changed, 546 insertions(+), 640 deletions(-) delete mode 100644 Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/MinMaxContinuationToken.cs diff --git a/Microsoft.Azure.Cosmos/src/FeedIteratorCore.cs b/Microsoft.Azure.Cosmos/src/FeedIteratorCore.cs index 8fc5491b14..9e169ec05f 100644 --- a/Microsoft.Azure.Cosmos/src/FeedIteratorCore.cs +++ b/Microsoft.Azure.Cosmos/src/FeedIteratorCore.cs @@ -10,7 +10,7 @@ namespace Microsoft.Azure.Cosmos using System.Text; using System.Threading; using System.Threading.Tasks; - using Microsoft.Azure.Cosmos.Json; + using Microsoft.Azure.Cosmos.CosmosElements; using Microsoft.Azure.Cosmos.Query.Core; using Microsoft.Azure.Documents; using static Microsoft.Azure.Documents.RuntimeConstants; @@ -110,7 +110,7 @@ internal static bool GetHasMoreResults(string continuationToken, HttpStatusCode return continuationToken != null && statusCode != HttpStatusCode.NotModified; } - public override void SerializeState(IJsonWriter jsonWriter) + public override CosmosElement GetCosmsoElementContinuationToken() { throw new NotImplementedException(); } @@ -135,6 +135,11 @@ internal FeedIteratorCore( public override bool HasMoreResults => this.feedIterator.HasMoreResults; + public override CosmosElement GetCosmosElementContinuationToken() + { + return this.feedIterator.GetCosmsoElementContinuationToken(); + } + /// /// Get the next set of results from the cosmos service /// @@ -147,10 +152,5 @@ public override async Task> ReadNextAsync(CancellationToken canc ResponseMessage response = await this.feedIterator.ReadNextAsync(cancellationToken); return this.responseCreator(response); } - - public override void SerializeState(IJsonWriter jsonWriter) - { - this.feedIterator.SerializeState(jsonWriter); - } } } \ No newline at end of file diff --git a/Microsoft.Azure.Cosmos/src/FeedIteratorInternal.cs b/Microsoft.Azure.Cosmos/src/FeedIteratorInternal.cs index b4da79e520..6f3aaea1f4 100644 --- a/Microsoft.Azure.Cosmos/src/FeedIteratorInternal.cs +++ b/Microsoft.Azure.Cosmos/src/FeedIteratorInternal.cs @@ -4,6 +4,7 @@ namespace Microsoft.Azure.Cosmos { + using Microsoft.Azure.Cosmos.CosmosElements; using Microsoft.Azure.Cosmos.Json; /// @@ -18,6 +19,6 @@ namespace Microsoft.Azure.Cosmos #endif abstract class FeedIteratorInternal : FeedIterator { - public abstract void SerializeState(IJsonWriter jsonWriter); + public abstract CosmosElement GetCosmsoElementContinuationToken(); } } \ No newline at end of file diff --git a/Microsoft.Azure.Cosmos/src/FeedIteratorInternal{T}.cs b/Microsoft.Azure.Cosmos/src/FeedIteratorInternal{T}.cs index b8732678a9..d26df5d5f0 100644 --- a/Microsoft.Azure.Cosmos/src/FeedIteratorInternal{T}.cs +++ b/Microsoft.Azure.Cosmos/src/FeedIteratorInternal{T}.cs @@ -4,6 +4,7 @@ namespace Microsoft.Azure.Cosmos { + using Microsoft.Azure.Cosmos.CosmosElements; using Microsoft.Azure.Cosmos.Json; /// @@ -18,6 +19,6 @@ namespace Microsoft.Azure.Cosmos #endif abstract class FeedIteratorInternal : FeedIterator { - public abstract void SerializeState(IJsonWriter jsonWriter); + public abstract CosmosElement GetCosmosElementContinuationToken(); } } \ No newline at end of file diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/AggregateDocumentQueryExecutionComponent.Client.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/AggregateDocumentQueryExecutionComponent.Client.cs index a90e6edf33..9b6621c308 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/AggregateDocumentQueryExecutionComponent.Client.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/AggregateDocumentQueryExecutionComponent.Client.cs @@ -117,7 +117,7 @@ public override async Task DrainAsync( responseLengthBytes: responseLengthBytes); } - public override void SerializeState(IJsonWriter jsonWriter) + public override CosmosElement GetCosmosElementContinuationToken() { throw new NotImplementedException(); } diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/AggregateDocumentQueryExecutionComponent.Compute.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/AggregateDocumentQueryExecutionComponent.Compute.cs index e794215e36..8b509d0bbe 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/AggregateDocumentQueryExecutionComponent.Compute.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/AggregateDocumentQueryExecutionComponent.Compute.cs @@ -5,6 +5,7 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.Aggregate { using System; using System.Collections.Generic; + using System.Runtime.CompilerServices; using System.Threading; using System.Threading.Tasks; using Microsoft.Azure.Cosmos.CosmosElements; @@ -19,9 +20,6 @@ internal abstract partial class AggregateDocumentQueryExecutionComponent : Docum { private sealed class ComputeAggregateDocumentQueryExecutionComponent : AggregateDocumentQueryExecutionComponent { - private const string SourceTokenName = "SourceToken"; - private const string AggregationTokenName = "AggregationToken"; - private static readonly IReadOnlyList EmptyResults = new List().AsReadOnly(); private ComputeAggregateDocumentQueryExecutionComponent( @@ -160,26 +158,24 @@ private QueryResponseCore GetEmptyPage(QueryResponseCore sourceResponse) return response; } - public override void SerializeState(IJsonWriter jsonWriter) + public override CosmosElement GetCosmosElementContinuationToken() { - if (jsonWriter == null) + if (this.IsDone) { - throw new ArgumentNullException(nameof(jsonWriter)); + return default; } - if (!this.IsDone) - { - jsonWriter.WriteObjectStart(); - jsonWriter.WriteFieldName(ComputeAggregateDocumentQueryExecutionComponent.SourceTokenName); - this.Source.SerializeState(jsonWriter); - jsonWriter.WriteFieldName(ComputeAggregateDocumentQueryExecutionComponent.AggregationTokenName); - this.singleGroupAggregator.SerializeState(jsonWriter); - jsonWriter.WriteObjectEnd(); - } + AggregateContinuationToken aggregateContinuationToken = new AggregateContinuationToken( + singleGroupAggregatorContinuationToken: this.singleGroupAggregator.GetCosmosElementContinuationToken(), + sourceContinuationToken: this.Source.GetCosmosElementContinuationToken()); + return AggregateContinuationToken.ToCosmosElement(aggregateContinuationToken); } private readonly struct AggregateContinuationToken { + private const string SourceTokenName = "SourceToken"; + private const string AggregationTokenName = "AggregationToken"; + public AggregateContinuationToken( CosmosElement singleGroupAggregatorContinuationToken, CosmosElement sourceContinuationToken) @@ -192,6 +188,23 @@ public AggregateContinuationToken( public CosmosElement SourceContinuationToken { get; } + public static CosmosElement ToCosmosElement(AggregateContinuationToken aggregateContinuationToken) + { + Dictionary dictionary = new Dictionary() + { + { + AggregateContinuationToken.SourceTokenName, + aggregateContinuationToken.SourceContinuationToken + }, + { + AggregateContinuationToken.AggregationTokenName, + aggregateContinuationToken.SingleGroupAggregatorContinuationToken + } + }; + + return CosmosObject.Create(dictionary); + } + public static bool TryCreateFromCosmosElement( CosmosElement continuationToken, out AggregateContinuationToken aggregateContinuationToken) @@ -208,7 +221,7 @@ public static bool TryCreateFromCosmosElement( } if (!rawAggregateContinuationToken.TryGetValue( - ComputeAggregateDocumentQueryExecutionComponent.AggregationTokenName, + AggregateContinuationToken.AggregationTokenName, out CosmosElement singleGroupAggregatorContinuationToken)) { aggregateContinuationToken = default; @@ -216,7 +229,7 @@ public static bool TryCreateFromCosmosElement( } if (!rawAggregateContinuationToken.TryGetValue( - ComputeAggregateDocumentQueryExecutionComponent.SourceTokenName, + AggregateContinuationToken.SourceTokenName, out CosmosElement sourceContinuationToken)) { aggregateContinuationToken = default; diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/AverageAggregator.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/AverageAggregator.cs index c5c3316765..4ead8af534 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/AverageAggregator.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/AverageAggregator.cs @@ -4,10 +4,10 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.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.ContinuationTokens; using Microsoft.Azure.Cosmos.Query.Core.Exceptions; using Microsoft.Azure.Cosmos.Query.Core.Monads; @@ -95,6 +95,11 @@ public static TryCatch TryCreate(CosmosElement continuationToken) return TryCatch.FromResult(new AverageAggregator(averageInfo)); } + public CosmosElement GetCosmosElementContinuationToken() + { + return AverageInfo.ToCosmosElement(this.globalAverage); + } + /// /// Struct that stores a weighted average as a sum and count so they that average across different partitions with different numbers of documents can be taken. /// @@ -114,6 +119,19 @@ public AverageInfo(double? sum, long count) this.Count = count; } + public static CosmosElement ToCosmosElement(AverageInfo averageInfo) + { + Dictionary dictionary = new Dictionary(); + if (averageInfo.Sum.HasValue) + { + dictionary.Add(AverageInfo.SumName, CosmosNumber64.Create(averageInfo.Sum.Value)); + } + + dictionary.Add(AverageInfo.CountName, CosmosNumber64.Create(averageInfo.Count)); + + return CosmosObject.Create(dictionary); + } + /// /// Initializes a new instance of the AverageInfo class. /// diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/CountAggregator.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/CountAggregator.cs index abaa599f38..84851aaf66 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/CountAggregator.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/CountAggregator.cs @@ -62,16 +62,6 @@ public string GetContinuationToken() return this.globalCount.ToString(CultureInfo.InvariantCulture); } - public void SerializeState(IJsonWriter jsonWriter) - { - if (jsonWriter == null) - { - throw new ArgumentNullException(nameof(jsonWriter)); - } - - jsonWriter.WriteInt64Value(this.globalCount); - } - public static TryCatch TryCreate(CosmosElement continuationToken) { long partialCount; @@ -93,5 +83,10 @@ public static TryCatch TryCreate(CosmosElement continuationToken) return TryCatch.FromResult( new CountAggregator(initialCount: partialCount)); } + + public CosmosElement GetCosmosElementContinuationToken() + { + return CosmosNumber64.Create(this.globalCount); + } } } \ No newline at end of file diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/IAggregator.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/IAggregator.cs index fd7493f5bf..dace3bf6a7 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/IAggregator.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/IAggregator.cs @@ -23,12 +23,6 @@ internal interface IAggregator /// The result of the aggregation. CosmosElement GetResult(); - /// - /// Gets a continuation token that stores the partial aggregate up till this point. - /// - /// A continuation token that stores the partial aggregate up till this point. - string GetContinuationToken(); - - void SerializeState(IJsonWriter jsonWriter); + CosmosElement GetCosmosElementContinuationToken(); } } diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/MinMaxAggregator.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/MinMaxAggregator.cs index af88b822cd..acf167809d 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/MinMaxAggregator.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/MinMaxAggregator.cs @@ -5,10 +5,8 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.Aggregate.Aggrega { using System; using System.Collections.Generic; - using System.Runtime.CompilerServices; 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.ExecutionContext.ItemProducers; using Microsoft.Azure.Cosmos.Query.Core.Monads; @@ -131,42 +129,6 @@ public CosmosElement GetResult() return result; } - public string GetContinuationToken() - { - IJsonWriter jsonWriter = JsonWriter.Create(JsonSerializationFormat.Text); - this.SerializeState(jsonWriter); - return Utf8StringHelpers.ToString(jsonWriter.GetResult()); - } - - public void SerializeState(IJsonWriter jsonWriter) - { - if (jsonWriter == null) - { - throw new ArgumentNullException(nameof(jsonWriter)); - } - - MinMaxContinuationToken minMaxContinuationToken; - if (this.globalMinMax == ItemComparer.MinValue) - { - minMaxContinuationToken = MinMaxContinuationToken.CreateMinValueContinuationToken(); - } - else if (this.globalMinMax == ItemComparer.MaxValue) - { - minMaxContinuationToken = MinMaxContinuationToken.CreateMaxValueContinuationToken(); - } - else if (this.globalMinMax == Undefined) - { - minMaxContinuationToken = MinMaxContinuationToken.CreateUndefinedValueContinuationToken(); - } - else - { - minMaxContinuationToken = MinMaxContinuationToken.CreateValueContinuationToken(this.globalMinMax); - } - - CosmosElement minMaxContinuationTokenAsCosmosElement = MinMaxContinuationToken.ToCosmosElement(minMaxContinuationToken); - minMaxContinuationTokenAsCosmosElement.WriteTo(jsonWriter); - } - public static TryCatch TryCreateMinAggregator(CosmosElement continuationToken) { return MinMaxAggregator.TryCreate(isMinAggregation: true, continuationToken: continuationToken); @@ -229,6 +191,30 @@ private static bool IsPrimitve(CosmosElement cosmosElement) return cosmosElement.Accept(IsPrimitiveCosmosElementVisitor.Singleton); } + public CosmosElement GetCosmosElementContinuationToken() + { + MinMaxContinuationToken minMaxContinuationToken; + if (this.globalMinMax == ItemComparer.MinValue) + { + minMaxContinuationToken = MinMaxContinuationToken.CreateMinValueContinuationToken(); + } + else if (this.globalMinMax == ItemComparer.MaxValue) + { + minMaxContinuationToken = MinMaxContinuationToken.CreateMaxValueContinuationToken(); + } + else if (this.globalMinMax == Undefined) + { + minMaxContinuationToken = MinMaxContinuationToken.CreateUndefinedValueContinuationToken(); + } + else + { + minMaxContinuationToken = MinMaxContinuationToken.CreateValueContinuationToken(this.globalMinMax); + } + + CosmosElement minMaxContinuationTokenAsCosmosElement = MinMaxContinuationToken.ToCosmosElement(minMaxContinuationToken); + return minMaxContinuationTokenAsCosmosElement; + } + private sealed class IsPrimitiveCosmosElementVisitor : ICosmosElementVisitor { public static readonly IsPrimitiveCosmosElementVisitor Singleton = new IsPrimitiveCosmosElementVisitor(); @@ -277,5 +263,158 @@ public bool Visit(CosmosString cosmosString) return true; } } + + private sealed class MinMaxContinuationToken + { + private const string TypeName = "type"; + private const string ValueName = "value"; + + private MinMaxContinuationToken( + MinMaxContinuationTokenType type, + CosmosElement value) + { + switch (type) + { + case MinMaxContinuationTokenType.MinValue: + case MinMaxContinuationTokenType.MaxValue: + case MinMaxContinuationTokenType.Undefined: + if (value != null) + { + throw new ArgumentException($"{nameof(value)} must be set with type: {type}."); + } + break; + + case MinMaxContinuationTokenType.Value: + if (value == null) + { + throw new ArgumentException($"{nameof(value)} must not be set with type: {type}."); + } + break; + + default: + throw new ArgumentOutOfRangeException($"Unknown {nameof(type)}: {type}."); + } + + this.Type = type; + this.Value = value; + } + + public MinMaxContinuationTokenType Type { get; } + public CosmosElement Value { get; } + + public static MinMaxContinuationToken CreateMinValueContinuationToken() + { + return new MinMaxContinuationToken(type: MinMaxContinuationTokenType.MinValue, value: null); + } + + public static MinMaxContinuationToken CreateMaxValueContinuationToken() + { + return new MinMaxContinuationToken(type: MinMaxContinuationTokenType.MaxValue, value: null); + } + + public static MinMaxContinuationToken CreateUndefinedValueContinuationToken() + { + return new MinMaxContinuationToken(type: MinMaxContinuationTokenType.Undefined, value: null); + } + + public static MinMaxContinuationToken CreateValueContinuationToken(CosmosElement value) + { + return new MinMaxContinuationToken(type: MinMaxContinuationTokenType.Value, value: value); + } + + public static CosmosElement ToCosmosElement(MinMaxContinuationToken minMaxContinuationToken) + { + if (minMaxContinuationToken == null) + { + throw new ArgumentNullException(nameof(minMaxContinuationToken)); + } + + Dictionary dictionary = new Dictionary(); + dictionary.Add( + MinMaxContinuationToken.TypeName, + EnumToCosmosString.ConvertEnumToCosmosString(minMaxContinuationToken.Type)); + if (minMaxContinuationToken.Value != null) + { + dictionary.Add(MinMaxContinuationToken.ValueName, minMaxContinuationToken.Value); + } + + return CosmosObject.Create(dictionary); + } + + public static TryCatch TryCreateFromCosmosElement(CosmosElement cosmosElement) + { + if (!(cosmosElement is CosmosObject cosmosObject)) + { + return TryCatch.FromException( + new MalformedContinuationTokenException($"{nameof(MinMaxContinuationToken)} was not an object.")); + } + + if (!cosmosObject.TryGetValue(MinMaxContinuationToken.TypeName, out CosmosString typeValue)) + { + return TryCatch.FromException( + new MalformedContinuationTokenException($"{nameof(MinMaxContinuationToken)} is missing property: {MinMaxContinuationToken.TypeName}.")); + } + + if (!Enum.TryParse(typeValue.Value, out MinMaxContinuationTokenType minMaxContinuationTokenType)) + { + return TryCatch.FromException( + new MalformedContinuationTokenException($"{nameof(MinMaxContinuationToken)} has malformed '{MinMaxContinuationToken.TypeName}': {typeValue.Value}.")); + } + + CosmosElement value; + if (minMaxContinuationTokenType == MinMaxContinuationTokenType.Value) + { + if (!cosmosObject.TryGetValue(MinMaxContinuationToken.ValueName, out value)) + { + return TryCatch.FromException( + new MalformedContinuationTokenException($"{nameof(MinMaxContinuationToken)} is missing property: {MinMaxContinuationToken.ValueName}.")); + } + } + else + { + value = null; + } + + return TryCatch.FromResult( + new MinMaxContinuationToken(minMaxContinuationTokenType, value)); + } + + private static class EnumToCosmosString + { + private static readonly CosmosString MinValueCosmosString = CosmosString.Create(MinMaxContinuationTokenType.MinValue.ToString()); + private static readonly CosmosString MaxValueCosmosString = CosmosString.Create(MinMaxContinuationTokenType.MaxValue.ToString()); + private static readonly CosmosString UndefinedCosmosString = CosmosString.Create(MinMaxContinuationTokenType.Undefined.ToString()); + private static readonly CosmosString ValueCosmosString = CosmosString.Create(MinMaxContinuationTokenType.Value.ToString()); + + public static CosmosString ConvertEnumToCosmosString(MinMaxContinuationTokenType type) + { + switch (type) + { + case MinMaxContinuationTokenType.MinValue: + return EnumToCosmosString.MinValueCosmosString; + + case MinMaxContinuationTokenType.MaxValue: + return EnumToCosmosString.MaxValueCosmosString; + + case MinMaxContinuationTokenType.Undefined: + return EnumToCosmosString.UndefinedCosmosString; + + case MinMaxContinuationTokenType.Value: + return EnumToCosmosString.ValueCosmosString; + + default: + throw new ArgumentOutOfRangeException($"Unknown {nameof(MinMaxContinuationTokenType)}: {type}"); + } + } + } + + public enum MinMaxContinuationTokenType + { + MinValue, + MaxValue, + Undefined, + Value, + } + } } } diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/MinMaxContinuationToken.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/MinMaxContinuationToken.cs deleted file mode 100644 index 65d003eb99..0000000000 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/MinMaxContinuationToken.cs +++ /dev/null @@ -1,165 +0,0 @@ -// ------------------------------------------------------------ -// Copyright (c) Microsoft Corporation. All rights reserved. -// ------------------------------------------------------------ - -namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.Aggregate.Aggregators -{ - 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; - - internal sealed class MinMaxContinuationToken - { - private const string TypeName = "type"; - private const string ValueName = "value"; - - private MinMaxContinuationToken( - MinMaxContinuationTokenType type, - CosmosElement value) - { - switch (type) - { - case MinMaxContinuationTokenType.MinValue: - case MinMaxContinuationTokenType.MaxValue: - case MinMaxContinuationTokenType.Undefined: - if (value != null) - { - throw new ArgumentException($"{nameof(value)} must be set with type: {type}."); - } - break; - - case MinMaxContinuationTokenType.Value: - if (value == null) - { - throw new ArgumentException($"{nameof(value)} must not be set with type: {type}."); - } - break; - - default: - throw new ArgumentOutOfRangeException($"Unknown {nameof(type)}: {type}."); - } - - this.Type = type; - this.Value = value; - } - - public MinMaxContinuationTokenType Type { get; } - public CosmosElement Value { get; } - - public static MinMaxContinuationToken CreateMinValueContinuationToken() - { - return new MinMaxContinuationToken(type: MinMaxContinuationTokenType.MinValue, value: null); - } - - public static MinMaxContinuationToken CreateMaxValueContinuationToken() - { - return new MinMaxContinuationToken(type: MinMaxContinuationTokenType.MaxValue, value: null); - } - - public static MinMaxContinuationToken CreateUndefinedValueContinuationToken() - { - return new MinMaxContinuationToken(type: MinMaxContinuationTokenType.Undefined, value: null); - } - - public static MinMaxContinuationToken CreateValueContinuationToken(CosmosElement value) - { - return new MinMaxContinuationToken(type: MinMaxContinuationTokenType.Value, value: value); - } - - public static CosmosElement ToCosmosElement(MinMaxContinuationToken minMaxContinuationToken) - { - if (minMaxContinuationToken == null) - { - throw new ArgumentNullException(nameof(minMaxContinuationToken)); - } - - Dictionary dictionary = new Dictionary(); - dictionary.Add( - MinMaxContinuationToken.TypeName, - EnumToCosmosString.ConvertEnumToCosmosString(minMaxContinuationToken.Type)); - if (minMaxContinuationToken.Value != null) - { - dictionary.Add(MinMaxContinuationToken.ValueName, minMaxContinuationToken.Value); - } - - return CosmosObject.Create(dictionary); - } - - public static TryCatch TryCreateFromCosmosElement(CosmosElement cosmosElement) - { - if (!(cosmosElement is CosmosObject cosmosObject)) - { - return TryCatch.FromException( - new MalformedContinuationTokenException($"{nameof(MinMaxContinuationToken)} was not an object.")); - } - - if (!cosmosObject.TryGetValue(MinMaxContinuationToken.TypeName, out CosmosString typeValue)) - { - return TryCatch.FromException( - new MalformedContinuationTokenException($"{nameof(MinMaxContinuationToken)} is missing property: {MinMaxContinuationToken.TypeName}.")); - } - - if (!Enum.TryParse(typeValue.Value, out MinMaxContinuationTokenType minMaxContinuationTokenType)) - { - return TryCatch.FromException( - new MalformedContinuationTokenException($"{nameof(MinMaxContinuationToken)} has malformed '{MinMaxContinuationToken.TypeName}': {typeValue.Value}.")); - } - - CosmosElement value; - if (minMaxContinuationTokenType == MinMaxContinuationTokenType.Value) - { - if (!cosmosObject.TryGetValue(MinMaxContinuationToken.ValueName, out value)) - { - return TryCatch.FromException( - new MalformedContinuationTokenException($"{nameof(MinMaxContinuationToken)} is missing property: {MinMaxContinuationToken.ValueName}.")); - } - } - else - { - value = null; - } - - return TryCatch.FromResult( - new MinMaxContinuationToken(minMaxContinuationTokenType, value)); - } - - private static class EnumToCosmosString - { - private static readonly CosmosString MinValueCosmosString = CosmosString.Create(MinMaxContinuationTokenType.MinValue.ToString()); - private static readonly CosmosString MaxValueCosmosString = CosmosString.Create(MinMaxContinuationTokenType.MaxValue.ToString()); - private static readonly CosmosString UndefinedCosmosString = CosmosString.Create(MinMaxContinuationTokenType.Undefined.ToString()); - private static readonly CosmosString ValueCosmosString = CosmosString.Create(MinMaxContinuationTokenType.Value.ToString()); - - public static CosmosString ConvertEnumToCosmosString(MinMaxContinuationTokenType type) - { - switch (type) - { - case MinMaxContinuationTokenType.MinValue: - return EnumToCosmosString.MinValueCosmosString; - - case MinMaxContinuationTokenType.MaxValue: - return EnumToCosmosString.MaxValueCosmosString; - - case MinMaxContinuationTokenType.Undefined: - return EnumToCosmosString.UndefinedCosmosString; - - case MinMaxContinuationTokenType.Value: - return EnumToCosmosString.ValueCosmosString; - - default: - throw new ArgumentOutOfRangeException($"Unknown {nameof(MinMaxContinuationTokenType)}: {type}"); - } - } - } - - public enum MinMaxContinuationTokenType - { - MinValue, - MaxValue, - Undefined, - Value, - } - } -} diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/SingleGroupAggregator.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/SingleGroupAggregator.cs index 188d74c0f4..58da49d586 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/SingleGroupAggregator.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/SingleGroupAggregator.cs @@ -28,9 +28,7 @@ internal abstract class SingleGroupAggregator /// public abstract CosmosElement GetResult(); - public abstract string GetContinuationToken(); - - public abstract void SerializeState(IJsonWriter jsonWriter); + public abstract CosmosElement GetCosmosElementContinuationToken(); public static TryCatch TryCreate( AggregateOperator[] aggregates, @@ -107,19 +105,9 @@ public override CosmosElement GetResult() return this.aggregateValue.Result; } - public override string GetContinuationToken() - { - return this.aggregateValue.GetContinuationToken(); - } - - public override void SerializeState(IJsonWriter jsonWriter) + public override CosmosElement GetCosmosElementContinuationToken() { - if (jsonWriter == null) - { - throw new ArgumentNullException(nameof(jsonWriter)); - } - - this.aggregateValue.SerializeState(jsonWriter); + return this.aggregateValue.GetCosmosElementContinuationToken(); } public override string ToString() @@ -162,29 +150,15 @@ public override CosmosElement GetResult() return CosmosObject.Create(aliasToElement); } - public override string GetContinuationToken() + public override CosmosElement GetCosmosElementContinuationToken() { - IJsonWriter jsonWriter = JsonWriter.Create(JsonSerializationFormat.Text); - this.SerializeState(jsonWriter); - return Utf8StringHelpers.ToString(jsonWriter.GetResult()); - } - - public override void SerializeState(IJsonWriter jsonWriter) - { - if (jsonWriter == null) - { - throw new ArgumentNullException(nameof(jsonWriter)); - } - - jsonWriter.WriteObjectStart(); - + Dictionary dictionary = new Dictionary(); foreach (KeyValuePair kvp in this.aliasToValue) { - jsonWriter.WriteFieldName(kvp.Key); - kvp.Value.SerializeState(jsonWriter); + dictionary.Add(kvp.Key, kvp.Value.GetCosmosElementContinuationToken()); } - jsonWriter.WriteObjectEnd(); + return CosmosObject.Create(dictionary); } public static TryCatch TryCreate( @@ -287,9 +261,7 @@ private abstract class AggregateValue public abstract CosmosElement Result { get; } - public abstract string GetContinuationToken(); - - public abstract void SerializeState(IJsonWriter jsonWriter); + public abstract CosmosElement GetCosmosElementContinuationToken(); public override string ToString() { @@ -328,14 +300,9 @@ public override void AddValue(CosmosElement aggregateValue) this.aggregator.Aggregate(aggregateItem.Item); } - public override string GetContinuationToken() + public override CosmosElement GetCosmosElementContinuationToken() { - return this.aggregator.GetContinuationToken(); - } - - public override void SerializeState(IJsonWriter jsonWriter) - { - this.aggregator.SerializeState(jsonWriter); + return this.aggregator.GetCosmosElementContinuationToken(); } public static TryCatch TryCreate( @@ -397,34 +364,17 @@ public override CosmosElement Result } } - public override string GetContinuationToken() + public override CosmosElement GetCosmosElementContinuationToken() { - IJsonWriter jsonWriter = JsonWriter.Create(JsonSerializationFormat.Text); - this.SerializeState(jsonWriter); - - string continuationToken = Utf8StringHelpers.ToString(jsonWriter.GetResult()); - return continuationToken; - } - - public override void SerializeState(IJsonWriter jsonWriter) - { - if (jsonWriter == null) - { - throw new ArgumentNullException(nameof(jsonWriter)); - } - - jsonWriter.WriteObjectStart(); - - jsonWriter.WriteFieldName(nameof(this.initialized)); - jsonWriter.WriteBoolValue(this.initialized); + Dictionary dictionary = new Dictionary(); + dictionary.Add(nameof(this.initialized), CosmosBoolean.Create(this.initialized)); if (this.value != null) { - jsonWriter.WriteFieldName(nameof(this.value)); - this.value.WriteTo(jsonWriter); + dictionary.Add(nameof(this.value), this.value); } - jsonWriter.WriteObjectEnd(); + return CosmosObject.Create(dictionary); } public static TryCatch TryCreate(CosmosElement continuationToken) diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/SumAggregator.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/SumAggregator.cs index 546a885540..ac8801b5d3 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/SumAggregator.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/SumAggregator.cs @@ -76,8 +76,11 @@ public void SerializeState(IJsonWriter jsonWriter) { throw new ArgumentNullException(nameof(jsonWriter)); } + } - jsonWriter.WriteFloat64Value(this.globalSum); + public CosmosElement GetCosmosElementContinuationToken() + { + return CosmosNumber64.Create(this.globalSum); } public static TryCatch TryCreate(CosmosElement requestContinuationToken) diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctDocumentQueryExecutionComponent.Client.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctDocumentQueryExecutionComponent.Client.cs index 7a0f9256d5..aa88246259 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctDocumentQueryExecutionComponent.Client.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctDocumentQueryExecutionComponent.Client.cs @@ -180,7 +180,7 @@ public override async Task DrainAsync(int maxElements, Cancel return queryResponseCore; } - public override void SerializeState(IJsonWriter jsonWriter) + public override CosmosElement GetCosmosElementContinuationToken() { throw new NotImplementedException(); } diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctDocumentQueryExecutionComponent.Compute.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctDocumentQueryExecutionComponent.Compute.cs index 99e0604c21..99d3c21db3 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctDocumentQueryExecutionComponent.Compute.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctDocumentQueryExecutionComponent.Compute.cs @@ -112,26 +112,24 @@ public override async Task DrainAsync(int maxElements, Cancel responseLengthBytes: sourceResponse.ResponseLengthBytes); } - public override void SerializeState(IJsonWriter jsonWriter) + public override CosmosElement GetCosmosElementContinuationToken() { - if (jsonWriter == null) + if (this.IsDone) { - throw new ArgumentNullException(nameof(jsonWriter)); + return default; } - if (!this.IsDone) - { - jsonWriter.WriteObjectStart(); - jsonWriter.WriteFieldName(DistinctDocumentQueryExecutionComponent.SourceTokenName); - this.Source.SerializeState(jsonWriter); - jsonWriter.WriteFieldName(DistinctDocumentQueryExecutionComponent.DistinctMapTokenName); - this.distinctMap.SerializeState(jsonWriter); - jsonWriter.WriteObjectEnd(); - } + DistinctContinuationToken distinctContinuationToken = new DistinctContinuationToken( + sourceToken: this.Source.GetCosmosElementContinuationToken(), + distinctMapToken: this.distinctMap.GetCosmosElementContinuationToken()); + return DistinctContinuationToken.ToCosmosElement(distinctContinuationToken); } private readonly struct DistinctContinuationToken { + private const string SourceTokenName = "SourceToken"; + private const string DistinctMapTokenName = "DistinctMapToken"; + public DistinctContinuationToken(CosmosElement sourceToken, CosmosElement distinctMapToken) { this.SourceToken = sourceToken; @@ -142,6 +140,23 @@ public DistinctContinuationToken(CosmosElement sourceToken, CosmosElement distin public CosmosElement DistinctMapToken { get; } + public static CosmosElement ToCosmosElement(DistinctContinuationToken distinctContinuationToken) + { + Dictionary dictionary = new Dictionary() + { + { + DistinctContinuationToken.SourceTokenName, + distinctContinuationToken.SourceToken + }, + { + DistinctContinuationToken.DistinctMapTokenName, + distinctContinuationToken.DistinctMapToken + } + }; + + return CosmosObject.Create(dictionary); + } + public static bool TryParse( CosmosElement requestContinuationToken, out DistinctContinuationToken distinctContinuationToken) 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 a213306321..09e35f81c7 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctDocumentQueryExecutionComponent.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctDocumentQueryExecutionComponent.cs @@ -25,8 +25,6 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.Distinct /// internal abstract partial class DistinctDocumentQueryExecutionComponent : DocumentQueryExecutionComponentBase { - private const string SourceTokenName = "SourceToken"; - private const string DistinctMapTokenName = "DistinctMapToken"; /// /// An DistinctMap that efficiently stores the documents that we have already seen. /// diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctMap.OrderedDistinctMap.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctMap.OrderedDistinctMap.cs index e09cec1333..879bbecca8 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctMap.OrderedDistinctMap.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctMap.OrderedDistinctMap.cs @@ -74,14 +74,9 @@ public override string GetContinuationToken() return this.lastHash.ToString(); } - public override void SerializeState(IJsonWriter jsonWriter) + public override CosmosElement GetCosmosElementContinuationToken() { - if (jsonWriter == null) - { - throw new ArgumentNullException(nameof(jsonWriter)); - } - - jsonWriter.WriteBinaryValue(UInt128.ToByteArray(this.lastHash)); + return CosmosBinary.Create(UInt128.ToByteArray(this.lastHash)); } public static TryCatch TryCreate(CosmosElement requestContinuationToken) diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctMap.UnorderedDistinctMap.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctMap.UnorderedDistinctMap.cs index a86709fbe5..6f71e5dcf2 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctMap.UnorderedDistinctMap.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctMap.UnorderedDistinctMap.cs @@ -5,13 +5,13 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.Distinct { using System; using System.Collections.Generic; + using System.Linq; using System.Runtime.InteropServices; 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.ContinuationTokens; using Microsoft.Azure.Cosmos.Query.Core.Exceptions; using Microsoft.Azure.Cosmos.Query.Core.Monads; @@ -215,87 +215,48 @@ public override bool Add(CosmosElement cosmosElement, out UInt128 hash) public override string GetContinuationToken() { - IJsonWriter jsonWriter = JsonWriter.Create(JsonSerializationFormat.Binary); - this.SerializeState(jsonWriter); - - ReadOnlyMemory memory = jsonWriter.GetResult(); - if (!MemoryMarshal.TryGetArray(memory, out ArraySegment buffer)) - { - buffer = new ArraySegment(memory.ToArray()); - } - - return Convert.ToBase64String(buffer.Array, buffer.Offset, buffer.Count); + return this.GetCosmosElementContinuationToken().ToString(); } - public override void SerializeState(IJsonWriter jsonWriter) + public override CosmosElement GetCosmosElementContinuationToken() { - if (jsonWriter == null) - { - throw new ArgumentNullException(nameof(jsonWriter)); - } - - jsonWriter.WriteObjectStart(); - - jsonWriter.WriteFieldName(UnorderdDistinctMap.NumbersName); - jsonWriter.WriteArrayStart(); - foreach (Number64 number in this.numbers) + Dictionary dictionary = new Dictionary() { - jsonWriter.WriteNumberValue(number); - } - jsonWriter.WriteArrayEnd(); - - jsonWriter.WriteFieldName(UnorderdDistinctMap.StringsLength4Name); - jsonWriter.WriteArrayStart(); - foreach (uint stringLength4 in this.stringsLength4) - { - jsonWriter.WriteUInt32Value(stringLength4); - } - jsonWriter.WriteArrayEnd(); - - jsonWriter.WriteFieldName(UnorderdDistinctMap.StringsLength8Name); - jsonWriter.WriteArrayStart(); - foreach (ulong stringLength8 in this.stringsLength8) - { - jsonWriter.WriteInt64Value((long)stringLength8); - } - jsonWriter.WriteArrayEnd(); - - jsonWriter.WriteFieldName(UnorderdDistinctMap.StringsLength16Name); - jsonWriter.WriteArrayStart(); - foreach (UInt128 stringLength16 in this.stringsLength16) - { - jsonWriter.WriteBinaryValue(UInt128.ToByteArray(stringLength16)); - } - jsonWriter.WriteArrayEnd(); - - jsonWriter.WriteFieldName(UnorderdDistinctMap.StringsLength16PlusName); - jsonWriter.WriteArrayStart(); - foreach (UInt128 stringLength16Plus in this.stringsLength16Plus) - { - jsonWriter.WriteBinaryValue(UInt128.ToByteArray(stringLength16Plus)); - } - jsonWriter.WriteArrayEnd(); - - jsonWriter.WriteFieldName(UnorderdDistinctMap.ArraysName); - jsonWriter.WriteArrayStart(); - foreach (UInt128 array in this.arrays) - { - jsonWriter.WriteBinaryValue(UInt128.ToByteArray(array)); - } - jsonWriter.WriteArrayEnd(); - - jsonWriter.WriteFieldName(UnorderdDistinctMap.ObjectName); - jsonWriter.WriteArrayStart(); - foreach (UInt128 objectHash in this.objects) - { - jsonWriter.WriteBinaryValue(UInt128.ToByteArray(objectHash)); - } - jsonWriter.WriteArrayEnd(); - - jsonWriter.WriteFieldName(UnorderdDistinctMap.SimpleValuesName); - jsonWriter.WriteStringValue(this.simpleValues.ToString()); + { + UnorderdDistinctMap.NumbersName, + CosmosArray.Create(this.numbers.Select(x => CosmosNumber64.Create(x))) + }, + { + UnorderdDistinctMap.StringsLength4Name, + CosmosArray.Create(this.stringsLength4.Select(x => CosmosUInt32.Create(x))) + }, + { + UnorderdDistinctMap.StringsLength8Name, + CosmosArray.Create(this.stringsLength8.Select(x => CosmosInt64.Create((long)x))) + }, + { + UnorderdDistinctMap.StringsLength16Name, + CosmosArray.Create(this.stringsLength16.Select(x => CosmosBinary.Create(UInt128.ToByteArray(x)))) + }, + { + UnorderdDistinctMap.StringsLength16PlusName, + CosmosArray.Create(this.stringsLength16Plus.Select(x => CosmosBinary.Create(UInt128.ToByteArray(x)))) + }, + { + UnorderdDistinctMap.ArraysName, + CosmosArray.Create(this.arrays.Select(x => CosmosBinary.Create(UInt128.ToByteArray(x)))) + }, + { + UnorderdDistinctMap.ObjectName, + CosmosArray.Create(this.objects.Select(x => CosmosBinary.Create(UInt128.ToByteArray(x)))) + }, + { + UnorderdDistinctMap.SimpleValuesName, + CosmosString.Create(this.simpleValues.ToString()) + } + }; - jsonWriter.WriteObjectEnd(); + return CosmosObject.Create(dictionary); } /// diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctMap.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctMap.cs index dae3f7e60a..4267bc9fff 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctMap.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctMap.cs @@ -62,6 +62,6 @@ public static TryCatch TryCreate( public abstract string GetContinuationToken(); - public abstract void SerializeState(IJsonWriter jsonWriter); + public abstract CosmosElement GetCosmosElementContinuationToken(); } } diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/DocumentQueryExecutionComponentBase.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/DocumentQueryExecutionComponentBase.cs index 81ec8b548b..98e907f802 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/DocumentQueryExecutionComponentBase.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/DocumentQueryExecutionComponentBase.cs @@ -6,6 +6,7 @@ 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; @@ -69,6 +70,6 @@ public void Stop() this.Source.Stop(); } - public abstract void SerializeState(IJsonWriter jsonWriter); + public abstract CosmosElement GetCosmosElementContinuationToken(); } } 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 2f5a2b59ac..0b5fcc586a 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,7 +10,6 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.GroupBy 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; using Microsoft.Azure.Cosmos.Query.Core.Monads; using Microsoft.Azure.Cosmos.Query.Core.QueryClient; @@ -103,9 +102,9 @@ public override async Task DrainAsync( return response; } - public override void SerializeState(IJsonWriter jsonWriter) + public override CosmosElement GetCosmosElementContinuationToken() { - throw new NotSupportedException(ContinuationTokenNotSupportedWithGroupBy); + throw new NotImplementedException(); } } } 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 e0f48f9212..1924195ed0 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 @@ -5,6 +5,7 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.GroupBy { using System; using System.Collections.Generic; + using System.Runtime.CompilerServices; using System.Threading; using System.Threading.Tasks; using Microsoft.Azure.Cosmos.CosmosElements; @@ -20,9 +21,8 @@ internal abstract partial class GroupByDocumentQueryExecutionComponent : Documen { private sealed class ComputeGroupByDocumentQueryExecutionComponent : GroupByDocumentQueryExecutionComponent { - private const string SourceTokenName = "SourceToken"; - private const string GroupingTableContinuationTokenName = "GroupingTableContinuationToken"; private const string DoneReadingGroupingsContinuationToken = "DONE"; + private static readonly CosmosElement DoneCosmosElementToken = CosmosString.Create(DoneReadingGroupingsContinuationToken); private static readonly IReadOnlyList EmptyResults = new List().AsReadOnly(); private static readonly IReadOnlyDictionary EmptyQueryMetrics = new Dictionary(); @@ -60,7 +60,6 @@ public static async Task> TryCreateAs } TryCatch tryCreateSource; - if ((groupByContinuationToken.SourceContinuationToken is CosmosString sourceContinuationToken) && (sourceContinuationToken.Value == ComputeGroupByDocumentQueryExecutionComponent.DoneReadingGroupingsContinuationToken)) { @@ -142,35 +141,35 @@ public override async Task DrainAsync( return response; } - public override void SerializeState(IJsonWriter jsonWriter) + public override CosmosElement GetCosmosElementContinuationToken() { - if (jsonWriter == null) + if (this.IsDone) { - throw new ArgumentNullException(nameof(jsonWriter)); + return default; } - if (!this.IsDone) + CosmosElement sourceContinuationToken; + if (this.Source.IsDone) + { + sourceContinuationToken = DoneCosmosElementToken; + } + else { - jsonWriter.WriteObjectStart(); + sourceContinuationToken = this.Source.GetCosmosElementContinuationToken(); + } - jsonWriter.WriteFieldName(GroupingTableContinuationTokenName); - this.groupingTable.SerializeState(jsonWriter); - jsonWriter.WriteFieldName(SourceTokenName); - if (this.Source.IsDone) - { - jsonWriter.WriteStringValue(DoneReadingGroupingsContinuationToken); - } - else - { - this.Source.SerializeState(jsonWriter); - } + GroupByContinuationToken groupByContinuationToken = new GroupByContinuationToken( + groupingTableContinuationToken: this.groupingTable.GetCosmosElementContinuationToken(), + sourceContinuationToken: sourceContinuationToken); - jsonWriter.WriteObjectEnd(); - } + return GroupByContinuationToken.ToCosmosElement(groupByContinuationToken); } private readonly struct GroupByContinuationToken { + private const string SourceTokenName = "SourceToken"; + private const string GroupingTableContinuationTokenName = "GroupingTableContinuationToken"; + public GroupByContinuationToken( CosmosElement groupingTableContinuationToken, CosmosElement sourceContinuationToken) @@ -183,6 +182,23 @@ public GroupByContinuationToken( public CosmosElement SourceContinuationToken { get; } + public static CosmosElement ToCosmosElement(GroupByContinuationToken groupByContinuationToken) + { + Dictionary dictionary = new Dictionary() + { + { + GroupByContinuationToken.SourceTokenName, + groupByContinuationToken.SourceContinuationToken + }, + { + GroupByContinuationToken.GroupingTableContinuationTokenName, + groupByContinuationToken.GroupingTableContinuationToken + }, + }; + + return CosmosObject.Create(dictionary); + } + public static bool TryParse(CosmosElement value, out GroupByContinuationToken groupByContinuationToken) { if (!(value is CosmosObject groupByContinuationTokenObject)) @@ -192,7 +208,7 @@ public static bool TryParse(CosmosElement value, out GroupByContinuationToken gr } if (!groupByContinuationTokenObject.TryGetValue( - ComputeGroupByDocumentQueryExecutionComponent.GroupingTableContinuationTokenName, + GroupByContinuationToken.GroupingTableContinuationTokenName, out CosmosElement groupingTableContinuationToken)) { groupByContinuationToken = default; @@ -200,7 +216,7 @@ public static bool TryParse(CosmosElement value, out GroupByContinuationToken gr } if (!groupByContinuationTokenObject.TryGetValue( - ComputeGroupByDocumentQueryExecutionComponent.SourceTokenName, + GroupByContinuationToken.SourceTokenName, out CosmosElement sourceContinuationToken)) { groupByContinuationToken = default; @@ -235,19 +251,14 @@ public Task DrainAsync(int maxElements, CancellationToken tok throw new NotImplementedException(); } - public IReadOnlyDictionary GetQueryMetrics() + public CosmosElement GetCosmosElementContinuationToken() { throw new NotImplementedException(); } - public void SerializeState(IJsonWriter jsonWriter) + public IReadOnlyDictionary GetQueryMetrics() { - if (jsonWriter == null) - { - throw new ArgumentNullException(nameof(jsonWriter)); - } - - jsonWriter.WriteNullValue(); + throw new NotImplementedException(); } public void Stop() 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 f6e869e597..d135f0e2f9 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/GroupBy/GroupByDocumentQueryExecutionComponent.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/GroupBy/GroupByDocumentQueryExecutionComponent.cs @@ -245,28 +245,15 @@ public IReadOnlyList Drain(int maxItemCount) return results; } - public string GetContinuationToken() + public CosmosElement GetCosmosElementContinuationToken() { - IJsonWriter jsonWriter = JsonWriter.Create(JsonSerializationFormat.Text); - this.SerializeState(jsonWriter); - string result = Utf8StringHelpers.ToString(jsonWriter.GetResult()); - return result; - } - - public void SerializeState(IJsonWriter jsonWriter) - { - if (jsonWriter == null) - { - throw new ArgumentNullException(nameof(jsonWriter)); - } - - jsonWriter.WriteObjectStart(); + Dictionary dictionary = new Dictionary(); foreach (KeyValuePair kvp in this.table) { - jsonWriter.WriteFieldName(kvp.Key.ToString()); - kvp.Value.SerializeState(jsonWriter); + dictionary.Add(kvp.Key.ToString(), kvp.Value.GetCosmosElementContinuationToken()); } - jsonWriter.WriteObjectEnd(); + + return CosmosObject.Create(dictionary); } public IEnumerator> GetEnumerator => this.table.GetEnumerator(); diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/IDocumentQueryExecutionComponent.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/IDocumentQueryExecutionComponent.cs index 5cb50ad609..eb547dad68 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/IDocumentQueryExecutionComponent.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/IDocumentQueryExecutionComponent.cs @@ -33,6 +33,6 @@ internal interface IDocumentQueryExecutionComponent : IDisposable /// void Stop(); - void SerializeState(IJsonWriter jsonWriter); + CosmosElement GetCosmosElementContinuationToken(); } } diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/SkipDocumentQueryExecutionComponent.Client.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/SkipDocumentQueryExecutionComponent.Client.cs index 491920365c..0fdf27eefa 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/SkipDocumentQueryExecutionComponent.Client.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/SkipDocumentQueryExecutionComponent.Client.cs @@ -10,7 +10,6 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.SkipTake 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; @@ -113,7 +112,7 @@ public override async Task DrainAsync(int maxElements, Cancel responseLengthBytes: sourcePage.ResponseLengthBytes); } - public override void SerializeState(IJsonWriter jsonWriter) + public override CosmosElement GetCosmosElementContinuationToken() { throw new NotImplementedException(); } diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/SkipDocumentQueryExecutionComponent.Compute.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/SkipDocumentQueryExecutionComponent.Compute.cs index add3d09845..6f97b9fbc6 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/SkipDocumentQueryExecutionComponent.Compute.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/SkipDocumentQueryExecutionComponent.Compute.cs @@ -6,11 +6,12 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.SkipTake 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.ContinuationTokens; using Microsoft.Azure.Cosmos.Query.Core.Exceptions; using Microsoft.Azure.Cosmos.Query.Core.Monads; using Microsoft.Azure.Cosmos.Query.Core.QueryClient; @@ -19,9 +20,6 @@ internal abstract partial class SkipDocumentQueryExecutionComponent : DocumentQu { private sealed class ComputeSkipDocumentQueryExecutionComponent : SkipDocumentQueryExecutionComponent { - private const string SkipCountPropertyName = "SkipCount"; - private const string SourceTokenPropertyName = "SourceToken"; - private ComputeSkipDocumentQueryExecutionComponent(IDocumentQueryExecutionComponent source, long skipCount) : base(source, skipCount) { @@ -92,22 +90,17 @@ public override async Task DrainAsync(int maxElements, Cancel responseLengthBytes: sourcePage.ResponseLengthBytes); } - public override void SerializeState(IJsonWriter jsonWriter) + public override CosmosElement GetCosmosElementContinuationToken() { - if (jsonWriter == null) + if (this.IsDone) { - throw new ArgumentNullException(nameof(jsonWriter)); + return default; } - if (!this.IsDone) - { - jsonWriter.WriteObjectStart(); - jsonWriter.WriteFieldName(SourceTokenPropertyName); - this.Source.SerializeState(jsonWriter); - jsonWriter.WriteFieldName(SkipCountPropertyName); - jsonWriter.WriteInt64Value(this.skipCount); - jsonWriter.WriteObjectEnd(); - } + OffsetContinuationToken offsetContinuationToken = new OffsetContinuationToken( + offset: this.skipCount, + sourceToken: this.Source.GetCosmosElementContinuationToken()); + return OffsetContinuationToken.ToCosmosElement(offsetContinuationToken); } /// @@ -115,6 +108,9 @@ public override void SerializeState(IJsonWriter jsonWriter) /// private readonly struct OffsetContinuationToken { + private const string SkipCountPropertyName = "SkipCount"; + private const string SourceTokenPropertyName = "SourceToken"; + /// /// Initializes a new instance of the OffsetContinuationToken struct. /// @@ -147,6 +143,23 @@ public CosmosElement SourceToken get; } + public static CosmosElement ToCosmosElement(OffsetContinuationToken offsetContinuationToken) + { + Dictionary dictionary = new Dictionary() + { + { + OffsetContinuationToken.SkipCountPropertyName, + CosmosNumber64.Create(offsetContinuationToken.Offset) + }, + { + OffsetContinuationToken.SourceTokenPropertyName, + offsetContinuationToken.SourceToken + } + }; + + return CosmosObject.Create(dictionary); + } + public static (bool parsed, OffsetContinuationToken offsetContinuationToken) TryParse(CosmosElement value) { if (value == null) @@ -159,12 +172,12 @@ public static (bool parsed, OffsetContinuationToken offsetContinuationToken) Try return (false, default); } - if (!cosmosObject.TryGetValue(SkipCountPropertyName, out CosmosNumber offset)) + if (!cosmosObject.TryGetValue(OffsetContinuationToken.SkipCountPropertyName, out CosmosNumber offset)) { return (false, default); } - if (!cosmosObject.TryGetValue(SourceTokenPropertyName, out CosmosElement sourceToken)) + if (!cosmosObject.TryGetValue(OffsetContinuationToken.SourceTokenPropertyName, out CosmosElement sourceToken)) { return (false, default); } diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/TakeDocumentQueryExecutionComponent.Client.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/TakeDocumentQueryExecutionComponent.Client.cs index 8058ab2ec2..8625a9229d 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/TakeDocumentQueryExecutionComponent.Client.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/TakeDocumentQueryExecutionComponent.Client.cs @@ -9,8 +9,6 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.SkipTake 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; @@ -188,7 +186,7 @@ public override async Task DrainAsync(int maxElements, Cancel responseLengthBytes: sourcePage.ResponseLengthBytes); } - public override void SerializeState(IJsonWriter jsonWriter) + public override CosmosElement GetCosmosElementContinuationToken() { throw new NotImplementedException(); } diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/TakeDocumentQueryExecutionComponent.Compute.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/TakeDocumentQueryExecutionComponent.Compute.cs index 40996a41de..80cf6c1012 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/TakeDocumentQueryExecutionComponent.Compute.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/TakeDocumentQueryExecutionComponent.Compute.cs @@ -9,7 +9,7 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.SkipTake using System.Threading; using System.Threading.Tasks; using Microsoft.Azure.Cosmos.CosmosElements; - using Microsoft.Azure.Cosmos.Json; + 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; @@ -18,9 +18,6 @@ internal abstract partial class TakeDocumentQueryExecutionComponent : DocumentQu { private sealed class ComputeTakeDocumentQueryExecutionComponent : TakeDocumentQueryExecutionComponent { - private const string SourceTokenName = "SourceToken"; - private const string TakeCountName = "TakeCount"; - private ComputeTakeDocumentQueryExecutionComponent(IDocumentQueryExecutionComponent source, int takeCount) : base(source, takeCount) { @@ -90,26 +87,24 @@ public override async Task DrainAsync(int maxElements, Cancel responseLengthBytes: sourcePage.ResponseLengthBytes); } - public override void SerializeState(IJsonWriter jsonWriter) + public override CosmosElement GetCosmosElementContinuationToken() { - if (jsonWriter == null) + if (this.IsDone) { - throw new ArgumentNullException(nameof(jsonWriter)); + return default; } - if (!this.IsDone) - { - jsonWriter.WriteObjectStart(); - jsonWriter.WriteFieldName(SourceTokenName); - this.Source.SerializeState(jsonWriter); - jsonWriter.WriteFieldName(TakeCountName); - jsonWriter.WriteInt64Value(this.takeCount); - jsonWriter.WriteObjectEnd(); - } + TakeContinuationToken takeContinuationToken = new TakeContinuationToken( + takeCount: this.takeCount, + sourceToken: this.Source.GetCosmosElementContinuationToken()); + return TakeContinuationToken.ToCosmosElement(takeContinuationToken); } private readonly struct TakeContinuationToken { + private const string SourceTokenName = "SourceToken"; + private const string TakeCountName = "TakeCount"; + public TakeContinuationToken(long takeCount, CosmosElement sourceToken) { if ((takeCount < 0) || (takeCount > int.MaxValue)) @@ -125,6 +120,23 @@ public TakeContinuationToken(long takeCount, CosmosElement sourceToken) public CosmosElement SourceToken { get; } + public static CosmosElement ToCosmosElement(TakeContinuationToken takeContinuationToken) + { + Dictionary dictionary = new Dictionary() + { + { + TakeContinuationToken.SourceTokenName, + takeContinuationToken.SourceToken + }, + { + TakeContinuationToken.TakeCountName, + CosmosNumber64.Create(takeContinuationToken.TakeCount) + }, + }; + + return CosmosObject.Create(dictionary); + } + public static bool TryParse(CosmosElement value, out TakeContinuationToken takeContinuationToken) { if (value == null) diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CatchAllCosmosQueryExecutionContext.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CatchAllCosmosQueryExecutionContext.cs index 68edc998f8..43e68d2bf8 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CatchAllCosmosQueryExecutionContext.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CatchAllCosmosQueryExecutionContext.cs @@ -7,6 +7,7 @@ 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; @@ -56,14 +57,9 @@ public override async Task ExecuteNextAsync(CancellationToken return queryResponseCore; } - public override void SerializeState(IJsonWriter jsonWriter) + public override CosmosElement GetCosmosElementContinuationToken() { - if (jsonWriter == null) - { - throw new ArgumentNullException(nameof(jsonWriter)); - } - - this.cosmosQueryExecutionContext.SerializeState(jsonWriter); + return this.cosmosQueryExecutionContext.GetCosmosElementContinuationToken(); } } } diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosCrossPartitionQueryExecutionContext.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosCrossPartitionQueryExecutionContext.cs index c95c048310..b9a9af8a06 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosCrossPartitionQueryExecutionContext.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosCrossPartitionQueryExecutionContext.cs @@ -662,7 +662,7 @@ private void OnItemProducerTreeCompleteFetching( } } - public abstract void SerializeState(IJsonWriter jsonWriter); + public abstract CosmosElement GetCosmosElementContinuationToken(); public readonly struct InitInfo { diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosQueryExecutionContext.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosQueryExecutionContext.cs index 7d234a1ba0..1cc7f29ab2 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosQueryExecutionContext.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosQueryExecutionContext.cs @@ -6,6 +6,7 @@ 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; @@ -31,6 +32,6 @@ public abstract bool IsDone /// A task to await on, which in return provides a DoucmentFeedResponse of documents. public abstract Task ExecuteNextAsync(CancellationToken cancellationToken); - public abstract void SerializeState(IJsonWriter jsonWriter); + public abstract CosmosElement GetCosmosElementContinuationToken(); } } diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosQueryExecutionContextWithNameCacheStaleRetry.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosQueryExecutionContextWithNameCacheStaleRetry.cs index f25a11d684..50be1cacd6 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosQueryExecutionContextWithNameCacheStaleRetry.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosQueryExecutionContextWithNameCacheStaleRetry.cs @@ -7,6 +7,7 @@ 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; @@ -60,14 +61,9 @@ await this.cosmosQueryContext.QueryClient.ForceRefreshCollectionCacheAsync( return queryResponse; } - public override void SerializeState(IJsonWriter jsonWriter) + public override CosmosElement GetCosmosElementContinuationToken() { - if (jsonWriter == null) - { - throw new ArgumentNullException(nameof(jsonWriter)); - } - - this.currentCosmosQueryExecutionContext.SerializeState(jsonWriter); + 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 index a48fa2d30a..ce0d5eb5d3 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/LazyCosmosQueryExecutionContext.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/LazyCosmosQueryExecutionContext.cs @@ -7,6 +7,7 @@ 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.Monads; using Microsoft.Azure.Cosmos.Query.Core.QueryClient; @@ -76,13 +77,8 @@ public override async Task ExecuteNextAsync(CancellationToken return queryResponseCore; } - public override void SerializeState(IJsonWriter jsonWriter) + public override CosmosElement GetCosmosElementContinuationToken() { - if (jsonWriter == null) - { - throw new ArgumentNullException(nameof(jsonWriter)); - } - if (!this.lazyTryCreateCosmosQueryExecutionContext.ValueInitialized) { throw new InvalidOperationException(); @@ -94,7 +90,7 @@ public override void SerializeState(IJsonWriter jsonWriter) throw tryCreateCosmosQueryExecutionContext.Exception; } - tryCreateCosmosQueryExecutionContext.Result.SerializeState(jsonWriter); + return tryCreateCosmosQueryExecutionContext.Result.GetCosmosElementContinuationToken(); } } } 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 7b13b99405..6c8a791047 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/OrderBy/CosmosOrderByItemQueryExecutionContext.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/OrderBy/CosmosOrderByItemQueryExecutionContext.cs @@ -154,43 +154,6 @@ protected override string ContinuationToken } } - public override void SerializeState(IJsonWriter jsonWriter) - { - if (jsonWriter == null) - { - throw new ArgumentNullException(nameof(jsonWriter)); - } - - IEnumerable activeItemProducers = this.GetActiveItemProducers(); - if (activeItemProducers.Any()) - { - jsonWriter.WriteArrayStart(); - - 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); - - OrderByContinuationToken.ToCosmosElement(orderByContinuationToken).WriteTo(jsonWriter); - } - - jsonWriter.WriteArrayEnd(); - } - } - public static async Task> TryCreateAsync( CosmosQueryContext queryContext, CosmosCrossPartitionQueryExecutionContext.CrossPartitionInitParams initParams, @@ -940,6 +903,40 @@ private static Tuple GetFormattedFilters( return new Tuple(left.ToString(), target.ToString(), right.ToString()); } + 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); + } + private readonly struct OrderByInitInfo { public OrderByInitInfo(RangeFilterInitializationInfo[] filters, IReadOnlyDictionary continuationTokens) 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 83a9ba193c..a3d3553b83 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/Parallel/CosmosParallelItemQueryExecutionContext.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/Parallel/CosmosParallelItemQueryExecutionContext.cs @@ -108,37 +108,6 @@ protected override string ContinuationToken } } - public override void SerializeState(IJsonWriter jsonWriter) - { - if (jsonWriter == null) - { - throw new ArgumentNullException(nameof(jsonWriter)); - } - - IEnumerable activeItemProducers = this.GetActiveItemProducers(); - if (activeItemProducers.Any()) - { - jsonWriter.WriteArrayStart(); - - 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: false, - isMaxInclusive: true) - }; - - CompositeContinuationToken.ToCosmosElement(compositeToken).WriteTo(jsonWriter); - } - - jsonWriter.WriteArrayEnd(); - } - } - public static async Task> TryCreateAsync( CosmosQueryContext queryContext, CosmosCrossPartitionQueryExecutionContext.CrossPartitionInitParams initParams, @@ -335,6 +304,34 @@ private static TryCatch TryGetInitializationInfoFromContinuati }); } + 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: false, + isMaxInclusive: true) + }; + + CosmosElement compositeContinuationToken = CompositeContinuationToken.ToCosmosElement(compositeToken); + compositeContinuationTokens.Add(compositeContinuationToken); + } + + return CosmosArray.Create(compositeContinuationTokens); + } + private readonly struct ParallelInitInfo { public ParallelInitInfo(IReadOnlyList partialRanges, IReadOnlyDictionary continuationTokens) diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/PipelinedDocumentQueryExecutionContext.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/PipelinedDocumentQueryExecutionContext.cs index 72a351baf0..2867201d7f 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/PipelinedDocumentQueryExecutionContext.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/PipelinedDocumentQueryExecutionContext.cs @@ -120,16 +120,6 @@ public override bool IsDone } } - public override void SerializeState(IJsonWriter jsonWriter) - { - if (jsonWriter == null) - { - throw new ArgumentNullException(nameof(jsonWriter)); - } - - this.component.SerializeState(jsonWriter); - } - public static async Task> TryCreateAsync( ExecutionEnvironment executionEnvironment, CosmosQueryContext queryContext, @@ -353,5 +343,10 @@ public override async Task ExecuteNextAsync(CancellationToken throw; } } + + public override CosmosElement GetCosmosElementContinuationToken() + { + return this.component.GetCosmosElementContinuationToken(); + } } } diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/QueryExecutionContextWithException.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/QueryExecutionContextWithException.cs index 520cea9d85..625acc6bbf 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/QueryExecutionContextWithException.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/QueryExecutionContextWithException.cs @@ -6,9 +6,8 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionContext using System; using System.Threading; using System.Threading.Tasks; - using Microsoft.Azure.Cosmos.Json; + using Microsoft.Azure.Cosmos.CosmosElements; using Microsoft.Azure.Cosmos.Query.Core; - using Microsoft.Azure.Cosmos.Query.Core.ExecutionContext; using Microsoft.Azure.Cosmos.Query.Core.QueryClient; internal sealed class QueryExecutionContextWithException : CosmosQueryExecutionContext @@ -39,7 +38,7 @@ public override Task ExecuteNextAsync(CancellationToken cance return Task.FromResult(queryResponse); } - public override void SerializeState(IJsonWriter jsonWriter) + public override CosmosElement GetCosmosElementContinuationToken() { throw new NotImplementedException(); } diff --git a/Microsoft.Azure.Cosmos/src/Query/v3Query/FeedIteratorInlineCore.cs b/Microsoft.Azure.Cosmos/src/Query/v3Query/FeedIteratorInlineCore.cs index 86e98941a5..cdfad030d6 100644 --- a/Microsoft.Azure.Cosmos/src/Query/v3Query/FeedIteratorInlineCore.cs +++ b/Microsoft.Azure.Cosmos/src/Query/v3Query/FeedIteratorInlineCore.cs @@ -7,9 +7,10 @@ namespace Microsoft.Azure.Cosmos using System; using System.Threading; using System.Threading.Tasks; + using Microsoft.Azure.Cosmos.CosmosElements; using Microsoft.Azure.Cosmos.Json; - internal class FeedIteratorInlineCore : FeedIteratorInternal + internal sealed class FeedIteratorInlineCore : FeedIteratorInternal { private readonly FeedIteratorInternal feedIteratorInternal; @@ -34,18 +35,18 @@ internal FeedIteratorInlineCore( public override bool HasMoreResults => this.feedIteratorInternal.HasMoreResults; - public override Task ReadNextAsync(CancellationToken cancellationToken = default) + public override CosmosElement GetCosmsoElementContinuationToken() { - return TaskHelper.RunInlineIfNeededAsync(() => this.feedIteratorInternal.ReadNextAsync(cancellationToken)); + return this.feedIteratorInternal.GetCosmsoElementContinuationToken(); } - public override void SerializeState(IJsonWriter jsonWriter) + public override Task ReadNextAsync(CancellationToken cancellationToken = default) { - this.feedIteratorInternal.SerializeState(jsonWriter); + return TaskHelper.RunInlineIfNeededAsync(() => this.feedIteratorInternal.ReadNextAsync(cancellationToken)); } } - internal class FeedIteratorInlineCore : FeedIteratorInternal + internal sealed class FeedIteratorInlineCore : FeedIteratorInternal { private readonly FeedIteratorInternal feedIteratorInternal; @@ -75,9 +76,9 @@ public override Task> ReadNextAsync(CancellationToken cancellati return TaskHelper.RunInlineIfNeededAsync(() => this.feedIteratorInternal.ReadNextAsync(cancellationToken)); } - public override void SerializeState(IJsonWriter jsonWriter) + public override CosmosElement GetCosmosElementContinuationToken() { - this.feedIteratorInternal.SerializeState(jsonWriter); + return this.feedIteratorInternal.GetCosmosElementContinuationToken(); } } } diff --git a/Microsoft.Azure.Cosmos/src/Query/v3Query/QueryIterator.cs b/Microsoft.Azure.Cosmos/src/Query/v3Query/QueryIterator.cs index 3bf46e23c9..e54782ae22 100644 --- a/Microsoft.Azure.Cosmos/src/Query/v3Query/QueryIterator.cs +++ b/Microsoft.Azure.Cosmos/src/Query/v3Query/QueryIterator.cs @@ -172,9 +172,9 @@ public override async Task ReadNextAsync(CancellationToken canc } } - public override void SerializeState(IJsonWriter jsonWriter) + public override CosmosElement GetCosmsoElementContinuationToken() { - this.cosmosQueryExecutionContext.SerializeState(jsonWriter); + return this.cosmosQueryExecutionContext.GetCosmosElementContinuationToken(); } } } diff --git a/Microsoft.Azure.Cosmos/src/Resource/QueryResponses/ChangeFeedPartitionKeyResultSetIteratorCore.cs b/Microsoft.Azure.Cosmos/src/Resource/QueryResponses/ChangeFeedPartitionKeyResultSetIteratorCore.cs index 30e12ca6c0..8a408a050e 100644 --- a/Microsoft.Azure.Cosmos/src/Resource/QueryResponses/ChangeFeedPartitionKeyResultSetIteratorCore.cs +++ b/Microsoft.Azure.Cosmos/src/Resource/QueryResponses/ChangeFeedPartitionKeyResultSetIteratorCore.cs @@ -8,6 +8,7 @@ namespace Microsoft.Azure.Cosmos using System.Net; using System.Threading; using System.Threading.Tasks; + using Microsoft.Azure.Cosmos.CosmosElements; using Microsoft.Azure.Cosmos.Json; /// @@ -55,6 +56,11 @@ internal ChangeFeedPartitionKeyResultSetIteratorCore( public override bool HasMoreResults => this.hasMoreResultsInternal; + public override CosmosElement GetCosmsoElementContinuationToken() + { + throw new NotImplementedException(); + } + /// /// Get the next set of results from the cosmos service /// @@ -76,11 +82,6 @@ internal ChangeFeedPartitionKeyResultSetIteratorCore( }, cancellationToken); } - public override void SerializeState(IJsonWriter jsonWriter) - { - throw new NotImplementedException(); - } - private Task NextResultSetDelegateAsync( string continuationToken, string partitionKeyRangeId, diff --git a/Microsoft.Azure.Cosmos/src/Resource/QueryResponses/ChangeFeedResultSetIteratorCore.cs b/Microsoft.Azure.Cosmos/src/Resource/QueryResponses/ChangeFeedResultSetIteratorCore.cs index 0219ec88ef..3d57f5aaa9 100644 --- a/Microsoft.Azure.Cosmos/src/Resource/QueryResponses/ChangeFeedResultSetIteratorCore.cs +++ b/Microsoft.Azure.Cosmos/src/Resource/QueryResponses/ChangeFeedResultSetIteratorCore.cs @@ -9,7 +9,7 @@ namespace Microsoft.Azure.Cosmos using System.Text; using System.Threading; using System.Threading.Tasks; - using Microsoft.Azure.Cosmos.Json; + using Microsoft.Azure.Cosmos.CosmosElements; using Microsoft.Azure.Cosmos.Query; using Microsoft.Azure.Cosmos.Query.Core.ContinuationTokens; using Microsoft.Azure.Cosmos.Routing; @@ -168,7 +168,7 @@ internal virtual Task NextResultSetDelegateAsync( cancellationToken: cancellationToken); } - public override void SerializeState(IJsonWriter jsonWriter) + public override CosmosElement GetCosmsoElementContinuationToken() { throw new NotImplementedException(); } diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CrossPartitionQueryTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CrossPartitionQueryTests.cs index 9408e99f71..5f29ec6623 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CrossPartitionQueryTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CrossPartitionQueryTests.cs @@ -618,18 +618,7 @@ private static async Task> QueryWithSerializeState( } resultsFromSerializeState.AddRange(cosmosQueryResponse); - Json.IJsonWriter jsonWriter = Json.JsonWriter.Create(Json.JsonSerializationFormat.Binary); - itemQuery.SerializeState(jsonWriter); - ReadOnlyMemory continuationTokenBuffer = jsonWriter.GetResult(); - - if (!continuationTokenBuffer.IsEmpty) - { - continuationToken = CosmosElement.CreateFromBuffer(continuationTokenBuffer); - } - else - { - continuationToken = null; - } + continuationToken = itemQuery.GetCosmosElementContinuationToken(); } catch (CosmosException cosmosException) when (cosmosException.StatusCode == (HttpStatusCode)429) { From 4fe4933f25148d1f532ff2fd4f02d4c84ad29206 Mon Sep 17 00:00:00 2001 From: Brandon Chong Date: Fri, 28 Feb 2020 12:35:06 -0800 Subject: [PATCH 17/28] updated tests --- .../PipelineContinuationTokenTests.Tests.xml | 27 +++---------------- .../PipelineContinuationTokenTests.cs | 6 ----- .../Query/QueryPipelineMockTests.cs | 20 +++++++------- 3 files changed, 14 insertions(+), 39 deletions(-) diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/BaselineTest/TestBaseline/PipelineContinuationTokenTests.Tests.xml b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/BaselineTest/TestBaseline/PipelineContinuationTokenTests.Tests.xml index f96f2ebe67..6e2a5e6073 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/BaselineTest/TestBaseline/PipelineContinuationTokenTests.Tests.xml +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/BaselineTest/TestBaseline/PipelineContinuationTokenTests.Tests.xml @@ -5,8 +5,8 @@ - - + + @@ -35,8 +35,7 @@ - - + @@ -46,7 +45,7 @@ - + @@ -67,15 +66,6 @@ - - - V1 Source Continuation Token Is Not A String - - - - - - V1.1 No Query Plan @@ -103,13 +93,4 @@ - - - V1.1 Source Continuation Is Not A String - - - - - - \ No newline at end of file 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 b8f5efbd7d..5a56dd4d55 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 @@ -48,9 +48,6 @@ public void Tests() new PipelineContinuationTokenTestsInput( description: "V1 No Source Continuation Token", continuationToken: "{\"Version\":\"1.0\"}"), - new PipelineContinuationTokenTestsInput( - description: "V1 Source Continuation Token Is Not A String", - continuationToken: "{\"Version\":\"1.0\", \"SourceContinuationToken\": 42}"), // Version 1.1 Negative Tests new PipelineContinuationTokenTestsInput( @@ -62,9 +59,6 @@ public void Tests() new PipelineContinuationTokenTestsInput( description: "V1.1 No Source Continuation", continuationToken: "{\"Version\":\"1.1\",\"QueryPlan\": null}"), - new PipelineContinuationTokenTestsInput( - description: "V1.1 Source Continuation Is Not A String", - continuationToken: "{\"Version\":\"1.1\",\"QueryPlan\": null,\"SourceContinuationToken\": 42}"), }; this.ExecuteTestSuite(pipelineContinuationTokenTestsInputs); diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/QueryPipelineMockTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/QueryPipelineMockTests.cs index f1268b815e..a09f8bb104 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/QueryPipelineMockTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/QueryPipelineMockTests.cs @@ -44,7 +44,7 @@ public async Task TestCosmosCrossPartitionQueryExecutionContextWithEmptyPagesAnd foreach (MockPartitionResponse[] mockResponse in mockResponsesScenario) { string initialContinuationToken = null; - string fullConitnuationToken = null; + string fullContinuationToken = null; if (createInitialContinuationToken) { initialContinuationToken = " - RID:02FYAIvUH1kCAAAAAAAAAA ==#RT:1#TRC:1"; @@ -58,7 +58,7 @@ public async Task TestCosmosCrossPartitionQueryExecutionContextWithEmptyPagesAnd Token = initialContinuationToken }; - fullConitnuationToken = CosmosArray.Create( + fullContinuationToken = CosmosArray.Create( new List() { CompositeContinuationToken.ToCosmosElement(compositeContinuation) @@ -94,7 +94,7 @@ public async Task TestCosmosCrossPartitionQueryExecutionContextWithEmptyPagesAnd IDocumentQueryExecutionComponent executionContext = (await CosmosParallelItemQueryExecutionContext.TryCreateAsync( context, initParams, - CosmosElement.Parse(fullConitnuationToken), + fullContinuationToken != null ? CosmosElement.Parse(fullContinuationToken) : null, this.cancellationToken)).Result; // Read all the pages from both splits @@ -132,7 +132,7 @@ public async Task TestCosmosCrossPartitionQueryExecutionContextWithFailuresAsync foreach (MockPartitionResponse[] mockResponse in mockResponsesScenario) { string initialContinuationToken = null; - string fullConitnuationToken = null; + string fullContinuationToken = null; if (createInitialContinuationToken) { initialContinuationToken = " - RID:02FYAIvUH1kCAAAAAAAAAA ==#RT:1#TRC:1"; @@ -146,7 +146,7 @@ public async Task TestCosmosCrossPartitionQueryExecutionContextWithFailuresAsync Token = initialContinuationToken }; - fullConitnuationToken = CosmosArray.Create( + fullContinuationToken = CosmosArray.Create( new List() { CompositeContinuationToken.ToCosmosElement(compositeContinuation) @@ -182,7 +182,7 @@ public async Task TestCosmosCrossPartitionQueryExecutionContextWithFailuresAsync IDocumentQueryExecutionComponent executionContext = (await CosmosParallelItemQueryExecutionContext.TryCreateAsync( context, initParams, - CosmosElement.Parse(fullConitnuationToken), + fullContinuationToken != null ? CosmosElement.Parse(fullContinuationToken) : null, this.cancellationToken)).Result; // Read all the pages from both splits @@ -235,7 +235,7 @@ public async Task TestCosmosOrderByQueryExecutionContextWithEmptyPagesAndSplitAs foreach (MockPartitionResponse[] mockResponse in mockResponsesScenario) { string initialContinuationToken = null; - string fullConitnuationToken = null; + string fullContinuationToken = null; if (createInitialContinuationToken) { ToDoItem itemToRepresentPreviousQuery = ToDoItem.CreateItems( @@ -266,7 +266,7 @@ public async Task TestCosmosOrderByQueryExecutionContextWithEmptyPagesAndSplitAs skipCount: 0, filter: null); - fullConitnuationToken = CosmosArray.Create( + fullContinuationToken = CosmosArray.Create( new List() { OrderByContinuationToken.ToCosmosElement(orderByContinuationToken) @@ -315,7 +315,7 @@ public async Task TestCosmosOrderByQueryExecutionContextWithEmptyPagesAndSplitAs IDocumentQueryExecutionComponent executionContext = (await CosmosOrderByItemQueryExecutionContext.TryCreateAsync( context, initParams, - CosmosElement.Parse(fullConitnuationToken), + fullContinuationToken != null ? CosmosElement.Parse(fullContinuationToken) : null, this.cancellationToken)).Result; // For order by it will drain all the pages till it gets a value. @@ -442,7 +442,7 @@ public async Task TestCosmosOrderByQueryExecutionContextWithFailurePageAsync(boo TryCatch tryCreate = await CosmosOrderByItemQueryExecutionContext.TryCreateAsync( context, initParams, - CosmosElement.Parse(fullContinuationToken), + fullContinuationToken != null ? CosmosElement.Parse(fullContinuationToken) : null, this.cancellationToken); if (tryCreate.Succeeded) From abec934a836d793d92d484065b8e5dbf4d42c59a Mon Sep 17 00:00:00 2001 From: Brandon Chong Date: Mon, 2 Mar 2020 13:12:46 -0800 Subject: [PATCH 18/28] resolved iteration comments --- .../src/CosmosElements/CosmosElement.cs | 32 ++- .../CompositeContinuationToken.cs | 39 +-- .../OrderByContinuationToken.cs | 55 ++-- ...eDocumentQueryExecutionComponent.Client.cs | 2 +- ...DocumentQueryExecutionComponent.Compute.cs | 4 +- ...ggregateDocumentQueryExecutionComponent.cs | 2 +- .../Aggregators/AverageAggregator.cs | 38 +-- .../Aggregate/Aggregators/MinMaxAggregator.cs | 21 +- .../Aggregators/SingleGroupAggregator.cs | 2 +- .../Aggregate/Aggregators/SumAggregator.cs | 15 +- ...tDocumentQueryExecutionComponent.Client.cs | 34 ++- .../DistinctMap.OrderedDistinctMap.cs | 3 - .../DistinctMap.UnorderedDistinctMap.cs | 51 ++-- .../DocumentQueryExecutionComponentBase.cs | 3 +- ...DocumentQueryExecutionComponent.Compute.cs | 19 +- .../GroupByDocumentQueryExecutionComponent.cs | 2 +- ...DocumentQueryExecutionComponent.Compute.cs | 18 +- ...DocumentQueryExecutionComponent.Compute.cs | 17 +- .../TakeDocumentQueryExecutionComponent.cs | 10 - .../CosmosOrderByItemQueryExecutionContext.cs | 14 +- .../src/Query/Core/QueryPlan/QueryInfo.cs | 18 +- .../Query/Core/QueryPlan/QueryPlanHandler.cs | 4 +- .../CrossPartitionQueryTests.cs | 248 +++++++----------- .../Query/QueryPipelineMockTests.cs | 4 +- 24 files changed, 297 insertions(+), 358 deletions(-) diff --git a/Microsoft.Azure.Cosmos/src/CosmosElements/CosmosElement.cs b/Microsoft.Azure.Cosmos/src/CosmosElements/CosmosElement.cs index a760111d70..07bf9b2b4f 100644 --- a/Microsoft.Azure.Cosmos/src/CosmosElements/CosmosElement.cs +++ b/Microsoft.Azure.Cosmos/src/CosmosElements/CosmosElement.cs @@ -33,8 +33,34 @@ public override string ToString() return Utf8StringHelpers.ToString(jsonWriter.GetResult()); } + public override bool Equals(object obj) + { + if (!(obj is CosmosElement cosmosElement)) + { + return false; + } + + return this.Equals(cosmosElement); + } + + public bool Equals(CosmosElement cosmosElement) + { + return CosmosElementEqualityComparer.Value.Equals(this, cosmosElement); + } + + public override int GetHashCode() + { + return CosmosElementEqualityComparer.Value.GetHashCode(this); + } + public abstract void WriteTo(IJsonWriter jsonWriter); + public abstract void Accept(ICosmosElementVisitor cosmosElementVisitor); + + public abstract TResult Accept(ICosmosElementVisitor cosmosElementVisitor); + + public abstract TResult Accept(ICosmosElementVisitor cosmosElementVisitor, TArg input); + public static bool TryCreateFromBuffer(ReadOnlyMemory buffer, out TCosmosElement cosmosElement) where TCosmosElement : CosmosElement { @@ -59,12 +85,6 @@ public static bool TryCreateFromBuffer(ReadOnlyMemory buff return true; } - public abstract void Accept(ICosmosElementVisitor cosmosElementVisitor); - - public abstract TResult Accept(ICosmosElementVisitor cosmosElementVisitor); - - public abstract TResult Accept(ICosmosElementVisitor cosmosElementVisitor, TArg input); - public static CosmosElement CreateFromBuffer(ReadOnlyMemory buffer) { IJsonNavigator jsonNavigator = JsonNavigator.Create(buffer); diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/CompositeContinuationToken.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/CompositeContinuationToken.cs index fad272c177..78d0290bb9 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/CompositeContinuationToken.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/CompositeContinuationToken.cs @@ -17,20 +17,23 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ContinuationTokens /// internal sealed class CompositeContinuationToken { - private const string TokenName = "token"; - private const string RangeName = "range"; + private static class PropertyNames + { + public const string Token = "token"; + public const string Range = "range"; - private const string MinName = "min"; - private const string MaxName = "max"; + public const string Min = "min"; + public const string Max = "max"; + } - [JsonProperty(TokenName)] + [JsonProperty(PropertyNames.Token)] public string Token { get; set; } - [JsonProperty(RangeName)] + [JsonProperty(PropertyNames.Range)] [JsonConverter(typeof(RangeJsonConverter))] public Documents.Routing.Range Range { @@ -49,14 +52,14 @@ public static CosmosElement ToCosmosElement(CompositeContinuationToken composite return CosmosObject.Create( new Dictionary() { - { CompositeContinuationToken.TokenName, token }, + { CompositeContinuationToken.PropertyNames.Token, token }, { - CompositeContinuationToken.RangeName, + CompositeContinuationToken.PropertyNames.Range, CosmosObject.Create( new Dictionary() { - { MinName, CosmosString.Create(compositeContinuationToken.Range.Min) }, - { MaxName, CosmosString.Create(compositeContinuationToken.Range.Max) } + { PropertyNames.Min, CosmosString.Create(compositeContinuationToken.Range.Min) }, + { PropertyNames.Max, CosmosString.Create(compositeContinuationToken.Range.Max) } }) }, }); @@ -70,10 +73,10 @@ public static TryCatch TryCreateFromCosmosElement(Co new MalformedContinuationTokenException($"{nameof(CompositeContinuationToken)} is not an object: {cosmosElement}")); } - if (!cosmosObject.TryGetValue(TokenName, out CosmosElement rawToken)) + if (!cosmosObject.TryGetValue(PropertyNames.Token, out CosmosElement rawToken)) { return TryCatch.FromException( - new MalformedContinuationTokenException($"{nameof(CompositeContinuationToken)} is missing field: '{TokenName}': {cosmosElement}")); + new MalformedContinuationTokenException($"{nameof(CompositeContinuationToken)} is missing field: '{PropertyNames.Token}': {cosmosElement}")); } string token; @@ -86,24 +89,24 @@ public static TryCatch TryCreateFromCosmosElement(Co token = null; } - if (!cosmosObject.TryGetValue(RangeName, out CosmosObject rawRange)) + if (!cosmosObject.TryGetValue(PropertyNames.Range, out CosmosObject rawRange)) { return TryCatch.FromException( - new MalformedContinuationTokenException($"{nameof(CompositeContinuationToken)} is missing field: '{RangeName}': {cosmosElement}")); + new MalformedContinuationTokenException($"{nameof(CompositeContinuationToken)} is missing field: '{PropertyNames.Range}': {cosmosElement}")); } - if (!rawRange.TryGetValue(MinName, out CosmosString rawMin)) + if (!rawRange.TryGetValue(PropertyNames.Min, out CosmosString rawMin)) { return TryCatch.FromException( - new MalformedContinuationTokenException($"{nameof(CompositeContinuationToken)} is missing field: '{MinName}': {cosmosElement}")); + new MalformedContinuationTokenException($"{nameof(CompositeContinuationToken)} is missing field: '{PropertyNames.Min}': {cosmosElement}")); } string min = rawMin.Value; - if (!rawRange.TryGetValue(MaxName, out CosmosString rawMax)) + if (!rawRange.TryGetValue(PropertyNames.Max, out CosmosString rawMax)) { return TryCatch.FromException( - new MalformedContinuationTokenException($"{nameof(CompositeContinuationToken)} is missing field: '{MaxName}': {cosmosElement}")); + new MalformedContinuationTokenException($"{nameof(CompositeContinuationToken)} is missing field: '{PropertyNames.Max}': {cosmosElement}")); } string max = rawMax.Value; diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/OrderByContinuationToken.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/OrderByContinuationToken.cs index 79f735db29..6fd5929708 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/OrderByContinuationToken.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/OrderByContinuationToken.cs @@ -5,9 +5,7 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ContinuationTokens { using System; using System.Collections.Generic; - using System.Globalization; using System.Linq; - using Microsoft.Azure.Cosmos.Core.Trace; using Microsoft.Azure.Cosmos.CosmosElements; using Microsoft.Azure.Cosmos.CosmosElements.Numbers; using Microsoft.Azure.Cosmos.Query.Core.Exceptions; @@ -57,11 +55,14 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ContinuationTokens /// internal sealed class OrderByContinuationToken { - private const string CompositeTokenName = "compositeToken"; - private const string OrderByItemsName = "orderByItems"; - private const string RidName = "rid"; - private const string SkipCountName = "skipCount"; - private const string FilterName = "filter"; + private static class PropertyNames + { + public const string CompositeToken = "compositeToken"; + public const string OrderByItems = "orderByItems"; + public const string Rid = "rid"; + public const string SkipCount = "skipCount"; + public const string Filter = "filter"; + } /// /// Initializes a new instance of the OrderByContinuationToken struct. @@ -109,7 +110,7 @@ public OrderByContinuationToken( /// {"compositeToken":{"token":"+RID:OpY0AN-mFAACAAAAAAAABA==#RT:1#TRC:1#RTD:qdTAEA==","range":{"min":"05C1D9CD673398","max":"05C1E399CD6732"}} /// ]]> /// - [JsonProperty(CompositeTokenName)] + [JsonProperty(PropertyNames.CompositeToken)] public CompositeContinuationToken CompositeContinuationToken { get; @@ -129,7 +130,7 @@ public CompositeContinuationToken CompositeContinuationToken /// /// Right now, we don't support orderBy by multiple fields, so orderByItems is an array of one element. /// > - [JsonProperty("orderByItems")] + [JsonProperty(PropertyNames.OrderByItems)] public IReadOnlyList OrderByItems { get; @@ -150,7 +151,7 @@ public IReadOnlyList OrderByItems /// "rid":"OpY0AN-mFAACAAAAAAAABA==" /// ]]> /// - [JsonProperty(RidName)] + [JsonProperty(PropertyNames.Rid)] public string Rid { get; @@ -176,7 +177,7 @@ public string Rid /// /// The skip count keeps track of that information. /// - [JsonProperty(SkipCountName)] + [JsonProperty(PropertyNames.SkipCount)] public int SkipCount { get; @@ -199,7 +200,7 @@ public int SkipCount /// ]]> /// /// - [JsonProperty(FilterName)] + [JsonProperty(PropertyNames.Filter)] public string Filter { get; @@ -221,11 +222,11 @@ public static CosmosElement ToCosmosElement(OrderByContinuationToken orderByCont CosmosObject cosmosObject = CosmosObject.Create( new Dictionary() { - { CompositeTokenName, compositeContinuationToken }, - { OrderByItemsName, orderByItems }, - { RidName, CosmosString.Create(orderByContinuationToken.Rid) }, - { SkipCountName, CosmosNumber64.Create(orderByContinuationToken.SkipCount) }, - { FilterName, filter }, + { PropertyNames.CompositeToken, compositeContinuationToken }, + { PropertyNames.OrderByItems, orderByItems }, + { PropertyNames.Rid, CosmosString.Create(orderByContinuationToken.Rid) }, + { PropertyNames.SkipCount, CosmosNumber64.Create(orderByContinuationToken.SkipCount) }, + { PropertyNames.Filter, filter }, }); return cosmosObject; @@ -239,10 +240,10 @@ public static TryCatch TryCreateFromCosmosElement(Cosm new MalformedContinuationTokenException($"{nameof(OrderByContinuationToken)} is not an object: {cosmosElement}")); } - if (!cosmosObject.TryGetValue(CompositeTokenName, out CosmosElement compositeContinuationTokenElement)) + if (!cosmosObject.TryGetValue(PropertyNames.CompositeToken, out CosmosElement compositeContinuationTokenElement)) { return TryCatch.FromException( - new MalformedContinuationTokenException($"{nameof(OrderByContinuationToken)} is missing field: '{CompositeTokenName}': {cosmosElement}")); + new MalformedContinuationTokenException($"{nameof(OrderByContinuationToken)} is missing field: '{PropertyNames.CompositeToken}': {cosmosElement}")); } TryCatch tryCompositeContinuation = CompositeContinuationToken.TryCreateFromCosmosElement(compositeContinuationTokenElement); @@ -253,34 +254,34 @@ public static TryCatch TryCreateFromCosmosElement(Cosm CompositeContinuationToken compositeContinuationToken = tryCompositeContinuation.Result; - if (!cosmosObject.TryGetValue(OrderByItemsName, out CosmosArray orderByItemsRaw)) + if (!cosmosObject.TryGetValue(PropertyNames.OrderByItems, out CosmosArray orderByItemsRaw)) { return TryCatch.FromException( - new MalformedContinuationTokenException($"{nameof(OrderByContinuationToken)} is missing field: '{OrderByItemsName}': {cosmosElement}")); + new MalformedContinuationTokenException($"{nameof(OrderByContinuationToken)} is missing field: '{PropertyNames.OrderByItems}': {cosmosElement}")); } List orderByItems = orderByItemsRaw.Select(x => OrderByItem.FromCosmosElement(x)).ToList(); - if (!cosmosObject.TryGetValue(RidName, out CosmosString ridRaw)) + if (!cosmosObject.TryGetValue(PropertyNames.Rid, out CosmosString ridRaw)) { return TryCatch.FromException( - new MalformedContinuationTokenException($"{nameof(OrderByContinuationToken)} is missing field: '{RidName}': {cosmosElement}")); + new MalformedContinuationTokenException($"{nameof(OrderByContinuationToken)} is missing field: '{PropertyNames.Rid}': {cosmosElement}")); } string rid = ridRaw.Value; - if (!cosmosObject.TryGetValue(SkipCountName, out CosmosNumber64 skipCountRaw)) + if (!cosmosObject.TryGetValue(PropertyNames.SkipCount, out CosmosNumber64 skipCountRaw)) { return TryCatch.FromException( - new MalformedContinuationTokenException($"{nameof(OrderByContinuationToken)} is missing field: '{SkipCountName}': {cosmosElement}")); + new MalformedContinuationTokenException($"{nameof(OrderByContinuationToken)} is missing field: '{PropertyNames.SkipCount}': {cosmosElement}")); } int skipCount = (int)Number64.ToLong(skipCountRaw.GetValue()); - if (!cosmosObject.TryGetValue(FilterName, out CosmosElement filterRaw)) + if (!cosmosObject.TryGetValue(PropertyNames.Filter, out CosmosElement filterRaw)) { return TryCatch.FromException( - new MalformedContinuationTokenException($"{nameof(OrderByContinuationToken)} is missing field: '{FilterName}': {cosmosElement}")); + new MalformedContinuationTokenException($"{nameof(OrderByContinuationToken)} is missing field: '{PropertyNames.Filter}': {cosmosElement}")); } string filter; diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/AggregateDocumentQueryExecutionComponent.Client.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/AggregateDocumentQueryExecutionComponent.Client.cs index 9b6621c308..c39639b13d 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/AggregateDocumentQueryExecutionComponent.Client.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/AggregateDocumentQueryExecutionComponent.Client.cs @@ -29,7 +29,7 @@ private ClientAggregateDocumentQueryExecutionComponent( } public static async Task> TryCreateAsync( - AggregateOperator[] aggregates, + IReadOnlyList aggregates, IReadOnlyDictionary aliasToAggregateType, IReadOnlyList orderedAliases, bool hasSelectValue, diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/AggregateDocumentQueryExecutionComponent.Compute.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/AggregateDocumentQueryExecutionComponent.Compute.cs index 8b509d0bbe..968a5b419b 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/AggregateDocumentQueryExecutionComponent.Compute.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/AggregateDocumentQueryExecutionComponent.Compute.cs @@ -9,8 +9,6 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.Aggregate 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.Aggregators; using Microsoft.Azure.Cosmos.Query.Core.Monads; @@ -32,7 +30,7 @@ private ComputeAggregateDocumentQueryExecutionComponent( } public static async Task> TryCreateAsync( - AggregateOperator[] aggregates, + IReadOnlyList aggregates, IReadOnlyDictionary aliasToAggregateType, IReadOnlyList orderedAliases, bool hasSelectValue, 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 36e26b1f4a..a61e5946f4 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/AggregateDocumentQueryExecutionComponent.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/AggregateDocumentQueryExecutionComponent.cs @@ -58,7 +58,7 @@ protected AggregateDocumentQueryExecutionComponent( public static async Task> TryCreateAsync( ExecutionEnvironment executionEnvironment, - AggregateOperator[] aggregates, + IReadOnlyList aggregates, IReadOnlyDictionary aliasToAggregateType, IReadOnlyList orderedAliases, bool hasSelectValue, diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/AverageAggregator.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/AverageAggregator.cs index 4ead8af534..1501e35d6d 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/AverageAggregator.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/AverageAggregator.cs @@ -56,21 +56,6 @@ public CosmosElement GetResult() return this.globalAverage.GetAverage(); } - public string GetContinuationToken() - { - return this.globalAverage.ToString(); - } - - public void SerializeState(IJsonWriter jsonWriter) - { - if (jsonWriter == null) - { - throw new ArgumentNullException(nameof(jsonWriter)); - } - - this.globalAverage.SerializeState(jsonWriter); - } - public static TryCatch TryCreate(CosmosElement continuationToken) { AverageInfo averageInfo; @@ -220,28 +205,7 @@ public CosmosNumber GetAverage() return null; } - return CosmosNumber64.Create(this.Sum.Value / this.Count); - } - - public void SerializeState(IJsonWriter jsonWriter) - { - if (jsonWriter == null) - { - throw new ArgumentNullException(nameof(jsonWriter)); - } - - jsonWriter.WriteObjectStart(); - - if (this.Sum.HasValue) - { - jsonWriter.WriteFieldName(AverageInfo.SumName); - jsonWriter.WriteFloat64Value(this.Sum.Value); - } - - jsonWriter.WriteFieldName(AverageInfo.CountName); - jsonWriter.WriteInt64Value(this.Count); - - jsonWriter.WriteObjectEnd(); + return CosmosNumber64.Create(this.Sum.Value / (double)this.Count); } public override string ToString() diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/MinMaxAggregator.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/MinMaxAggregator.cs index acf167809d..5f77dbba77 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/MinMaxAggregator.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/MinMaxAggregator.cs @@ -266,8 +266,11 @@ public bool Visit(CosmosString cosmosString) private sealed class MinMaxContinuationToken { - private const string TypeName = "type"; - private const string ValueName = "value"; + private static class PropertyNames + { + public const string Type = "type"; + public const string Value = "value"; + } private MinMaxContinuationToken( MinMaxContinuationTokenType type, @@ -331,11 +334,11 @@ public static CosmosElement ToCosmosElement(MinMaxContinuationToken minMaxContin Dictionary dictionary = new Dictionary(); dictionary.Add( - MinMaxContinuationToken.TypeName, + MinMaxContinuationToken.PropertyNames.Type, EnumToCosmosString.ConvertEnumToCosmosString(minMaxContinuationToken.Type)); if (minMaxContinuationToken.Value != null) { - dictionary.Add(MinMaxContinuationToken.ValueName, minMaxContinuationToken.Value); + dictionary.Add(MinMaxContinuationToken.PropertyNames.Value, minMaxContinuationToken.Value); } return CosmosObject.Create(dictionary); @@ -349,25 +352,25 @@ public static TryCatch TryCreateFromCosmosElement(Cosmo new MalformedContinuationTokenException($"{nameof(MinMaxContinuationToken)} was not an object.")); } - if (!cosmosObject.TryGetValue(MinMaxContinuationToken.TypeName, out CosmosString typeValue)) + if (!cosmosObject.TryGetValue(MinMaxContinuationToken.PropertyNames.Type, out CosmosString typeValue)) { return TryCatch.FromException( - new MalformedContinuationTokenException($"{nameof(MinMaxContinuationToken)} is missing property: {MinMaxContinuationToken.TypeName}.")); + new MalformedContinuationTokenException($"{nameof(MinMaxContinuationToken)} is missing property: {MinMaxContinuationToken.PropertyNames.Type}.")); } if (!Enum.TryParse(typeValue.Value, out MinMaxContinuationTokenType minMaxContinuationTokenType)) { return TryCatch.FromException( - new MalformedContinuationTokenException($"{nameof(MinMaxContinuationToken)} has malformed '{MinMaxContinuationToken.TypeName}': {typeValue.Value}.")); + new MalformedContinuationTokenException($"{nameof(MinMaxContinuationToken)} has malformed '{MinMaxContinuationToken.PropertyNames.Type}': {typeValue.Value}.")); } CosmosElement value; if (minMaxContinuationTokenType == MinMaxContinuationTokenType.Value) { - if (!cosmosObject.TryGetValue(MinMaxContinuationToken.ValueName, out value)) + if (!cosmosObject.TryGetValue(MinMaxContinuationToken.PropertyNames.Value, out value)) { return TryCatch.FromException( - new MalformedContinuationTokenException($"{nameof(MinMaxContinuationToken)} is missing property: {MinMaxContinuationToken.ValueName}.")); + new MalformedContinuationTokenException($"{nameof(MinMaxContinuationToken)} is missing property: {MinMaxContinuationToken.PropertyNames.Value}.")); } } else diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/SingleGroupAggregator.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/SingleGroupAggregator.cs index 58da49d586..95a567226d 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/SingleGroupAggregator.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/SingleGroupAggregator.cs @@ -31,7 +31,7 @@ internal abstract class SingleGroupAggregator public abstract CosmosElement GetCosmosElementContinuationToken(); public static TryCatch TryCreate( - AggregateOperator[] aggregates, + IReadOnlyList aggregates, IReadOnlyDictionary aggregateAliasToAggregateType, IReadOnlyList orderedAliases, bool hasSelectValue, diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/SumAggregator.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/SumAggregator.cs index ac8801b5d3..b0c0a94b84 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/SumAggregator.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Aggregate/Aggregators/SumAggregator.cs @@ -65,19 +65,6 @@ public CosmosElement GetResult() return CosmosNumber64.Create(this.globalSum); } - public string GetContinuationToken() - { - return this.globalSum.ToString("G17", CultureInfo.InvariantCulture); - } - - public void SerializeState(IJsonWriter jsonWriter) - { - if (jsonWriter == null) - { - throw new ArgumentNullException(nameof(jsonWriter)); - } - } - public CosmosElement GetCosmosElementContinuationToken() { return CosmosNumber64.Create(this.globalSum); @@ -94,7 +81,7 @@ public static TryCatch TryCreate(CosmosElement requestContinuationT new MalformedContinuationTokenException($"Malformed {nameof(SumAggregator)} continuation token: {requestContinuationToken}")); } - partialSum = Number64.ToLong(cosmosNumber.Value); + partialSum = Number64.ToDouble(cosmosNumber.Value); } else { diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctDocumentQueryExecutionComponent.Client.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctDocumentQueryExecutionComponent.Client.cs index aa88246259..7f69a03b63 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctDocumentQueryExecutionComponent.Client.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctDocumentQueryExecutionComponent.Client.cs @@ -54,7 +54,7 @@ public static async Task> TryCreateAs DistinctContinuationToken distinctContinuationToken; if (requestContinuation != null) { - if (!DistinctContinuationToken.TryParse(requestContinuation.ToString(), out distinctContinuationToken)) + if (!DistinctContinuationToken.TryParse(requestContinuation, out distinctContinuationToken)) { return TryCatch.FromException( new MalformedContinuationTokenException( @@ -190,6 +190,12 @@ public override CosmosElement GetCosmosElementContinuationToken() /// private sealed class DistinctContinuationToken { + private static class PropertyNames + { + public const string SourceToken = "SourceToken"; + public const string DistinctMapToken = "DistinctMapToken"; + } + public DistinctContinuationToken(string sourceToken, string distinctMapToken) { this.SourceToken = sourceToken; @@ -203,23 +209,37 @@ public DistinctContinuationToken(string sourceToken, string distinctMapToken) /// /// Tries to parse a DistinctContinuationToken from a string. /// - /// The value to parse. + /// The value to parse. /// The output DistinctContinuationToken. /// True if we successfully parsed the DistinctContinuationToken, else false. public static bool TryParse( - string stringRequestContinuationToken, + CosmosElement cosmosElement, out DistinctContinuationToken distinctContinuationToken) { - try + if (!(cosmosElement is CosmosObject cosmosObject)) + { + distinctContinuationToken = default; + return false; + } + + if (!cosmosObject.TryGetValue( + DistinctContinuationToken.PropertyNames.SourceToken, + out CosmosString sourceToken)) { - distinctContinuationToken = JsonConvert.DeserializeObject(stringRequestContinuationToken); - return true; + distinctContinuationToken = default; + return false; } - catch (JsonException) + + if (!cosmosObject.TryGetValue( + DistinctContinuationToken.PropertyNames.DistinctMapToken, + out CosmosString distinctMapToken)) { distinctContinuationToken = default; return false; } + + distinctContinuationToken = new DistinctContinuationToken(sourceToken.Value, distinctMapToken.Value); + return true; } /// diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctMap.OrderedDistinctMap.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctMap.OrderedDistinctMap.cs index 879bbecca8..da6e632d4c 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctMap.OrderedDistinctMap.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctMap.OrderedDistinctMap.cs @@ -5,9 +5,7 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.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.Exceptions; using Microsoft.Azure.Cosmos.Query.Core.Monads; @@ -96,7 +94,6 @@ public static TryCatch TryCreate(CosmosElement requestContinuationT break; case CosmosBinary cosmosBinary: - if (!UInt128.TryCreateFromByteArray(cosmosBinary.Value.Span, out lastHash)) { return TryCatch.FromException( diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctMap.UnorderedDistinctMap.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctMap.UnorderedDistinctMap.cs index 6f71e5dcf2..50c19ff862 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctMap.UnorderedDistinctMap.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/Distinct/DistinctMap.UnorderedDistinctMap.cs @@ -90,14 +90,17 @@ private sealed class UnorderdDistinctMap : DistinctMap /// private const int UIntLength = 4; - private const string NumbersName = "Numbers"; - private const string StringsLength4Name = "StringsLength4"; - private const string StringsLength8Name = "StringsLength8"; - private const string StringsLength16Name = "StringsLength16"; - private const string StringsLength16PlusName = "StringsLength16+"; - private const string ArraysName = "Arrays"; - private const string ObjectName = "Object"; - private const string SimpleValuesName = "SimpleValues"; + private static class PropertyNames + { + public const string Numbers = "Numbers"; + public const string StringsLength4 = "StringsLength4"; + public const string StringsLength8 = "StringsLength8"; + public const string StringsLength16 = "StringsLength16"; + public const string StringsLength16Plus = "StringsLength16+"; + public const string Arrays = "Arrays"; + public const string Object = "Object"; + public const string SimpleValues = "SimpleValues"; + } /// /// HashSet for all numbers seen. @@ -223,35 +226,35 @@ public override CosmosElement GetCosmosElementContinuationToken() Dictionary dictionary = new Dictionary() { { - UnorderdDistinctMap.NumbersName, + UnorderdDistinctMap.PropertyNames.Numbers, CosmosArray.Create(this.numbers.Select(x => CosmosNumber64.Create(x))) }, { - UnorderdDistinctMap.StringsLength4Name, + UnorderdDistinctMap.PropertyNames.StringsLength4, CosmosArray.Create(this.stringsLength4.Select(x => CosmosUInt32.Create(x))) }, { - UnorderdDistinctMap.StringsLength8Name, + UnorderdDistinctMap.PropertyNames.StringsLength8, CosmosArray.Create(this.stringsLength8.Select(x => CosmosInt64.Create((long)x))) }, { - UnorderdDistinctMap.StringsLength16Name, + UnorderdDistinctMap.PropertyNames.StringsLength16, CosmosArray.Create(this.stringsLength16.Select(x => CosmosBinary.Create(UInt128.ToByteArray(x)))) }, { - UnorderdDistinctMap.StringsLength16PlusName, + UnorderdDistinctMap.PropertyNames.StringsLength16Plus, CosmosArray.Create(this.stringsLength16Plus.Select(x => CosmosBinary.Create(UInt128.ToByteArray(x)))) }, { - UnorderdDistinctMap.ArraysName, + UnorderdDistinctMap.PropertyNames.Arrays, CosmosArray.Create(this.arrays.Select(x => CosmosBinary.Create(UInt128.ToByteArray(x)))) }, { - UnorderdDistinctMap.ObjectName, + UnorderdDistinctMap.PropertyNames.Object, CosmosArray.Create(this.objects.Select(x => CosmosBinary.Create(UInt128.ToByteArray(x)))) }, { - UnorderdDistinctMap.SimpleValuesName, + UnorderdDistinctMap.PropertyNames.SimpleValues, CosmosString.Create(this.simpleValues.ToString()) } }; @@ -373,7 +376,7 @@ public static TryCatch TryCreate(CosmosElement continuationToken) } // Numbers - if (!hashDictionary.TryGetValue(UnorderdDistinctMap.NumbersName, out CosmosArray numbersArray)) + if (!hashDictionary.TryGetValue(UnorderdDistinctMap.PropertyNames.Numbers, out CosmosArray numbersArray)) { return TryCatch.FromException( new MalformedContinuationTokenException( @@ -393,7 +396,7 @@ public static TryCatch TryCreate(CosmosElement continuationToken) } // Strings Length 4 - if (!hashDictionary.TryGetValue(UnorderdDistinctMap.StringsLength4Name, out CosmosArray stringsLength4Array)) + if (!hashDictionary.TryGetValue(UnorderdDistinctMap.PropertyNames.StringsLength4, out CosmosArray stringsLength4Array)) { return TryCatch.FromException( new MalformedContinuationTokenException( @@ -413,7 +416,7 @@ public static TryCatch TryCreate(CosmosElement continuationToken) } // Strings Length 8 - if (!hashDictionary.TryGetValue(UnorderdDistinctMap.StringsLength8Name, out CosmosArray stringsLength8Array)) + if (!hashDictionary.TryGetValue(UnorderdDistinctMap.PropertyNames.StringsLength8, out CosmosArray stringsLength8Array)) { return TryCatch.FromException( new MalformedContinuationTokenException( @@ -433,19 +436,19 @@ public static TryCatch TryCreate(CosmosElement continuationToken) } // Strings Length 16 - stringsLength16 = Parse128BitHashes(hashDictionary, UnorderdDistinctMap.StringsLength16Name); + stringsLength16 = Parse128BitHashes(hashDictionary, UnorderdDistinctMap.PropertyNames.StringsLength16); // Strings Length 24 - stringsLength16Plus = Parse128BitHashes(hashDictionary, UnorderdDistinctMap.StringsLength16PlusName); + stringsLength16Plus = Parse128BitHashes(hashDictionary, UnorderdDistinctMap.PropertyNames.StringsLength16Plus); // Array - arrays = Parse128BitHashes(hashDictionary, UnorderdDistinctMap.ArraysName); + arrays = Parse128BitHashes(hashDictionary, UnorderdDistinctMap.PropertyNames.Arrays); // Object - objects = Parse128BitHashes(hashDictionary, UnorderdDistinctMap.ObjectName); + objects = Parse128BitHashes(hashDictionary, UnorderdDistinctMap.PropertyNames.Object); // Simple Values - CosmosElement rawSimpleValues = hashDictionary[UnorderdDistinctMap.SimpleValuesName]; + CosmosElement rawSimpleValues = hashDictionary[UnorderdDistinctMap.PropertyNames.SimpleValues]; if (!(rawSimpleValues is CosmosString simpleValuesString)) { return TryCatch.FromException( diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/DocumentQueryExecutionComponentBase.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/DocumentQueryExecutionComponentBase.cs index 98e907f802..e3ab5efd67 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/DocumentQueryExecutionComponentBase.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/DocumentQueryExecutionComponentBase.cs @@ -7,7 +7,6 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent using System.Threading; using System.Threading.Tasks; using Microsoft.Azure.Cosmos.CosmosElements; - using Microsoft.Azure.Cosmos.Json; using Microsoft.Azure.Cosmos.Query.Core.QueryClient; /// @@ -15,7 +14,7 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent /// internal abstract class DocumentQueryExecutionComponentBase : IDocumentQueryExecutionComponent { - public static readonly string UseSerializeStateInstead = $"Use Serialize State instead"; + public static readonly string UseCosmosElementContinuationTokenInstead = $"Use Cosmos Element Continuation Token instead."; /// /// Source DocumentQueryExecutionComponent that this component will drain from. 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 1924195ed0..f7265dc64e 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 @@ -116,7 +116,7 @@ public override async Task DrainAsync( response = QueryResponseCore.CreateSuccess( result: EmptyResults, continuationToken: null, - disallowContinuationTokenMessage: DocumentQueryExecutionComponentBase.UseSerializeStateInstead, + disallowContinuationTokenMessage: DocumentQueryExecutionComponentBase.UseCosmosElementContinuationTokenInstead, activityId: sourceResponse.ActivityId, requestCharge: sourceResponse.RequestCharge, diagnostics: sourceResponse.Diagnostics, @@ -131,7 +131,7 @@ public override async Task DrainAsync( response = QueryResponseCore.CreateSuccess( result: results, continuationToken: null, - disallowContinuationTokenMessage: DocumentQueryExecutionComponentBase.UseSerializeStateInstead, + disallowContinuationTokenMessage: DocumentQueryExecutionComponentBase.UseCosmosElementContinuationTokenInstead, activityId: null, requestCharge: 0, diagnostics: QueryResponseCore.EmptyDiagnostics, @@ -167,8 +167,11 @@ public override CosmosElement GetCosmosElementContinuationToken() private readonly struct GroupByContinuationToken { - private const string SourceTokenName = "SourceToken"; - private const string GroupingTableContinuationTokenName = "GroupingTableContinuationToken"; + private static class PropertyNames + { + public const string SourceToken = "SourceToken"; + public const string GroupingTableContinuationToken = "GroupingTableContinuationToken"; + } public GroupByContinuationToken( CosmosElement groupingTableContinuationToken, @@ -187,11 +190,11 @@ public static CosmosElement ToCosmosElement(GroupByContinuationToken groupByCont Dictionary dictionary = new Dictionary() { { - GroupByContinuationToken.SourceTokenName, + GroupByContinuationToken.PropertyNames.SourceToken, groupByContinuationToken.SourceContinuationToken }, { - GroupByContinuationToken.GroupingTableContinuationTokenName, + GroupByContinuationToken.PropertyNames.GroupingTableContinuationToken, groupByContinuationToken.GroupingTableContinuationToken }, }; @@ -208,7 +211,7 @@ public static bool TryParse(CosmosElement value, out GroupByContinuationToken gr } if (!groupByContinuationTokenObject.TryGetValue( - GroupByContinuationToken.GroupingTableContinuationTokenName, + GroupByContinuationToken.PropertyNames.GroupingTableContinuationToken, out CosmosElement groupingTableContinuationToken)) { groupByContinuationToken = default; @@ -216,7 +219,7 @@ public static bool TryParse(CosmosElement value, out GroupByContinuationToken gr } if (!groupByContinuationTokenObject.TryGetValue( - GroupByContinuationToken.SourceTokenName, + GroupByContinuationToken.PropertyNames.SourceToken, out CosmosElement sourceContinuationToken)) { groupByContinuationToken = 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 d135f0e2f9..a509ca4a20 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/GroupBy/GroupByDocumentQueryExecutionComponent.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/GroupBy/GroupByDocumentQueryExecutionComponent.cs @@ -175,7 +175,7 @@ public CosmosElement Payload protected sealed class GroupingTable : IEnumerable> { - private static readonly AggregateOperator[] EmptyAggregateOperators = new AggregateOperator[] { }; + private static readonly IReadOnlyList EmptyAggregateOperators = new AggregateOperator[] { }; private readonly Dictionary table; private readonly IReadOnlyDictionary groupByAliasToAggregateType; diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/SkipDocumentQueryExecutionComponent.Compute.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/SkipDocumentQueryExecutionComponent.Compute.cs index 6f97b9fbc6..ad94557b2b 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/SkipDocumentQueryExecutionComponent.Compute.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/SkipDocumentQueryExecutionComponent.Compute.cs @@ -15,6 +15,7 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.SkipTake 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 { @@ -83,7 +84,7 @@ public override async Task DrainAsync(int maxElements, Cancel return QueryResponseCore.CreateSuccess( result: documentsAfterSkip, continuationToken: null, - disallowContinuationTokenMessage: DocumentQueryExecutionComponentBase.UseSerializeStateInstead, + disallowContinuationTokenMessage: DocumentQueryExecutionComponentBase.UseCosmosElementContinuationTokenInstead, activityId: sourcePage.ActivityId, requestCharge: sourcePage.RequestCharge, diagnostics: sourcePage.Diagnostics, @@ -108,8 +109,11 @@ public override CosmosElement GetCosmosElementContinuationToken() /// private readonly struct OffsetContinuationToken { - private const string SkipCountPropertyName = "SkipCount"; - private const string SourceTokenPropertyName = "SourceToken"; + private static class ProperytNames + { + public const string SkipCountProperty = "SkipCount"; + public const string SourceTokenProperty = "SourceToken"; + } /// /// Initializes a new instance of the OffsetContinuationToken struct. @@ -148,11 +152,11 @@ public static CosmosElement ToCosmosElement(OffsetContinuationToken offsetContin Dictionary dictionary = new Dictionary() { { - OffsetContinuationToken.SkipCountPropertyName, + OffsetContinuationToken.ProperytNames.SkipCountProperty, CosmosNumber64.Create(offsetContinuationToken.Offset) }, { - OffsetContinuationToken.SourceTokenPropertyName, + OffsetContinuationToken.ProperytNames.SourceTokenProperty, offsetContinuationToken.SourceToken } }; @@ -172,12 +176,12 @@ public static (bool parsed, OffsetContinuationToken offsetContinuationToken) Try return (false, default); } - if (!cosmosObject.TryGetValue(OffsetContinuationToken.SkipCountPropertyName, out CosmosNumber offset)) + if (!cosmosObject.TryGetValue(OffsetContinuationToken.ProperytNames.SkipCountProperty, out CosmosNumber offset)) { return (false, default); } - if (!cosmosObject.TryGetValue(OffsetContinuationToken.SourceTokenPropertyName, out CosmosElement sourceToken)) + if (!cosmosObject.TryGetValue(OffsetContinuationToken.ProperytNames.SourceTokenProperty, out CosmosElement sourceToken)) { return (false, default); } diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/TakeDocumentQueryExecutionComponent.Compute.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/TakeDocumentQueryExecutionComponent.Compute.cs index 80cf6c1012..adac2c147a 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/TakeDocumentQueryExecutionComponent.Compute.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/TakeDocumentQueryExecutionComponent.Compute.cs @@ -80,7 +80,7 @@ public override async Task DrainAsync(int maxElements, Cancel return QueryResponseCore.CreateSuccess( result: takedDocuments, continuationToken: null, - disallowContinuationTokenMessage: DocumentQueryExecutionComponentBase.UseSerializeStateInstead, + disallowContinuationTokenMessage: DocumentQueryExecutionComponentBase.UseCosmosElementContinuationTokenInstead, activityId: sourcePage.ActivityId, requestCharge: sourcePage.RequestCharge, diagnostics: sourcePage.Diagnostics, @@ -102,8 +102,11 @@ public override CosmosElement GetCosmosElementContinuationToken() private readonly struct TakeContinuationToken { - private const string SourceTokenName = "SourceToken"; - private const string TakeCountName = "TakeCount"; + public static class PropertyNames + { + public const string SourceToken = "SourceToken"; + public const string TakeCount = "TakeCount"; + } public TakeContinuationToken(long takeCount, CosmosElement sourceToken) { @@ -125,11 +128,11 @@ public static CosmosElement ToCosmosElement(TakeContinuationToken takeContinuati Dictionary dictionary = new Dictionary() { { - TakeContinuationToken.SourceTokenName, + TakeContinuationToken.PropertyNames.SourceToken, takeContinuationToken.SourceToken }, { - TakeContinuationToken.TakeCountName, + TakeContinuationToken.PropertyNames.TakeCount, CosmosNumber64.Create(takeContinuationToken.TakeCount) }, }; @@ -150,13 +153,13 @@ public static bool TryParse(CosmosElement value, out TakeContinuationToken takeC return false; } - if (!continuationToken.TryGetValue(TakeCountName, out CosmosNumber takeCount)) + if (!continuationToken.TryGetValue(TakeContinuationToken.PropertyNames.TakeCount, out CosmosNumber takeCount)) { takeContinuationToken = default; return false; } - if (!continuationToken.TryGetValue(SourceTokenName, out CosmosElement sourceToken)) + if (!continuationToken.TryGetValue(TakeContinuationToken.PropertyNames.SourceToken, out CosmosElement sourceToken)) { takeContinuationToken = default; return false; 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 2aa61c01db..48a76be5f7 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/TakeDocumentQueryExecutionComponent.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionComponent/SkipTake/TakeDocumentQueryExecutionComponent.cs @@ -4,20 +4,10 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionComponent.SkipTake { using System; - using System.Collections.Generic; - using System.Globalization; - using System.Linq; - using System.Threading; using System.Threading.Tasks; - using Microsoft.Azure.Cosmos.Core.Trace; 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.ExecutionContext; using Microsoft.Azure.Cosmos.Query.Core.Monads; - using Microsoft.Azure.Cosmos.Query.Core.QueryClient; - using Newtonsoft.Json; internal abstract partial class TakeDocumentQueryExecutionComponent : DocumentQueryExecutionComponentBase { 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 02352807b6..6c534ec9bd 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/OrderBy/CosmosOrderByItemQueryExecutionContext.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/OrderBy/CosmosOrderByItemQueryExecutionContext.cs @@ -11,7 +11,6 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.OrderBy using System.Threading; using System.Threading.Tasks; using Microsoft.Azure.Cosmos.CosmosElements; - using Microsoft.Azure.Cosmos.Json; using Microsoft.Azure.Cosmos.Query.Core; using Microsoft.Azure.Cosmos.Query.Core.Collections; using Microsoft.Azure.Cosmos.Query.Core.ContinuationTokens; @@ -174,7 +173,7 @@ public static async Task> TryCreateAs // 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); + OrderByItemProducerTreeComparer orderByItemProducerTreeComparer = new OrderByItemProducerTreeComparer(initParams.PartitionedQueryExecutionInfo.QueryInfo.OrderBy.ToArray()); CosmosOrderByItemQueryExecutionContext context = new CosmosOrderByItemQueryExecutionContext( initPararms: queryContext, maxConcurrency: initParams.MaxConcurrency, @@ -189,8 +188,8 @@ public static async Task> TryCreateAs collectionRid: initParams.CollectionRid, partitionKeyRanges: initParams.PartitionKeyRanges, initialPageSize: initParams.InitialPageSize, - sortOrders: initParams.PartitionedQueryExecutionInfo.QueryInfo.OrderBy, - orderByExpressions: initParams.PartitionedQueryExecutionInfo.QueryInfo.OrderByExpressions, + sortOrders: initParams.PartitionedQueryExecutionInfo.QueryInfo.OrderBy.ToArray(), + orderByExpressions: initParams.PartitionedQueryExecutionInfo.QueryInfo.OrderByExpressions.ToArray(), cancellationToken: cancellationToken)) .Try(() => context); } @@ -595,8 +594,8 @@ private async Task TryFilterAsync( 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.")); - + new MalformedContinuationTokenException( + $"Invalid Rid in the continuation token {continuationToken.CompositeContinuationToken.Token} for OrderBy~Context~TryParse.")); } resourceIds.Add(orderByResult.Rid, rid); @@ -607,7 +606,8 @@ private async Task TryFilterAsync( 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.")); + new MalformedContinuationTokenException( + $"Invalid Rid in the continuation token {continuationToken.CompositeContinuationToken.Token} for OrderBy~Context.")); } continuationRidVerified = true; diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/QueryPlan/QueryInfo.cs b/Microsoft.Azure.Cosmos/src/Query/Core/QueryPlan/QueryInfo.cs index 5d6ee3b462..fd59b1d3f6 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/QueryPlan/QueryInfo.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/QueryPlan/QueryInfo.cs @@ -45,42 +45,42 @@ public int? Limit } [JsonProperty("orderBy", ItemConverterType = typeof(StringEnumConverter))] - public SortOrder[] OrderBy + public IReadOnlyList OrderBy { get; set; } [JsonProperty("orderByExpressions")] - public string[] OrderByExpressions + public IReadOnlyList OrderByExpressions { get; set; } [JsonProperty("groupByExpressions")] - public string[] GroupByExpressions + public IReadOnlyList GroupByExpressions { get; set; } [JsonProperty("groupByAliases")] - public string[] GroupByAliases + public IReadOnlyList GroupByAliases { get; set; } [JsonProperty("aggregates", ItemConverterType = typeof(StringEnumConverter))] - public AggregateOperator[] Aggregates + public IReadOnlyList Aggregates { get; set; } [JsonProperty("groupByAliasToAggregateType", ItemConverterType = typeof(StringEnumConverter))] - public Dictionary GroupByAliasToAggregateType + public IReadOnlyDictionary GroupByAliasToAggregateType { get; set; @@ -119,7 +119,7 @@ public bool HasAggregates { get { - bool aggregatesListNonEmpty = (this.Aggregates != null) && (this.Aggregates.Length > 0); + bool aggregatesListNonEmpty = (this.Aggregates != null) && (this.Aggregates.Count > 0); if (aggregatesListNonEmpty) { return true; @@ -137,7 +137,7 @@ public bool HasGroupBy { get { - return this.GroupByExpressions != null && this.GroupByExpressions.Length > 0; + return this.GroupByExpressions != null && this.GroupByExpressions.Count > 0; } } @@ -145,7 +145,7 @@ public bool HasOrderBy { get { - return this.OrderBy != null && this.OrderBy.Length > 0; + return this.OrderBy != null && this.OrderBy.Count > 0; } } diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/QueryPlan/QueryPlanHandler.cs b/Microsoft.Azure.Cosmos/src/Query/Core/QueryPlan/QueryPlanHandler.cs index eb7f7e9fa9..84f92a811f 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/QueryPlan/QueryPlanHandler.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/QueryPlan/QueryPlanHandler.cs @@ -146,7 +146,7 @@ private static QueryFeatures GetNeededQueryFeaturesIfAggregateQuery( QueryFeatures neededQueryFeatures = QueryFeatures.None; if (queryInfo.HasAggregates) { - bool isSingleAggregate = (queryInfo.Aggregates.Length == 1) + bool isSingleAggregate = (queryInfo.Aggregates.Count == 1) || (queryInfo.GroupByAliasToAggregateType.Values.Where(aggregateOperator => aggregateOperator.HasValue).Count() == 1); if (isSingleAggregate) { @@ -253,7 +253,7 @@ private static QueryFeatures GetNeededQueryFeaturesIfOrderByQuery( QueryFeatures neededQueryFeatures = QueryFeatures.None; if (queryInfo.HasOrderBy) { - if (queryInfo.OrderByExpressions.Length == 1) + if (queryInfo.OrderByExpressions.Count == 1) { if (!supportedQueryFeatures.HasFlag(QueryFeatures.OrderBy)) { diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CrossPartitionQueryTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CrossPartitionQueryTests.cs index 8cfec38984..195947d506 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CrossPartitionQueryTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CrossPartitionQueryTests.cs @@ -587,7 +587,7 @@ private CosmosClient CreateNewCosmosClient(ConnectionMode connectionMode) } } - private static async Task> QueryWithSerializeState( + private static async Task> QueryWithCosmosElementContinuationTokenAsync( Container container, string query, QueryRequestOptions queryRequestOptions = null) @@ -597,7 +597,7 @@ private static async Task> QueryWithSerializeState( queryRequestOptions = new QueryRequestOptions(); } - List resultsFromSerializeState = new List(); + List resultsFromCosmosElementContinuationToken = new List(); CosmosElement continuationToken = null; do { @@ -618,7 +618,7 @@ private static async Task> QueryWithSerializeState( "Max Item Count is not being honored"); } - resultsFromSerializeState.AddRange(cosmosQueryResponse); + resultsFromCosmosElementContinuationToken.AddRange(cosmosQueryResponse); continuationToken = itemQuery.GetCosmosElementContinuationToken(); } catch (CosmosException cosmosException) when (cosmosException.StatusCode == (HttpStatusCode)429) @@ -629,10 +629,10 @@ private static async Task> QueryWithSerializeState( } } while (continuationToken != null); - return resultsFromSerializeState; + return resultsFromCosmosElementContinuationToken; } - private static async Task> QueryWithContinuationTokens( + private static async Task> QueryWithContinuationTokensAsync( Container container, string query, QueryRequestOptions queryRequestOptions = null) @@ -680,7 +680,7 @@ private static async Task> QueryWithContinuationTokens( return resultsFromContinuationToken; } - private static async Task> QueryWithoutContinuationTokens( + private static async Task> QueryWithoutContinuationTokensAsync( Container container, string query, QueryRequestOptions queryRequestOptions = null) @@ -849,7 +849,7 @@ private async Task TestQueryWithPartitionKeyHelper( Container container, IEnumerable documents) { - Assert.AreEqual(0, (await CrossPartitionQueryTests.RunQuery( + Assert.AreEqual(0, (await CrossPartitionQueryTests.RunQueryAsync( container, @"SELECT * FROM Root r WHERE false", new QueryRequestOptions() @@ -1329,7 +1329,7 @@ private async Task TestBasicCrossPartitionQueryHelper( ReturnResultsInDeterministicOrder = true, }; - List queryResults = await CrossPartitionQueryTests.RunQuery( + List queryResults = await CrossPartitionQueryTests.RunQueryAsync( container, query, feedOptions); @@ -1392,9 +1392,9 @@ async Task ValidateNonDeterministicQuery(Func, useOrderBy); - await ValidateNonDeterministicQuery(CrossPartitionQueryTests.QueryWithContinuationTokens, useOrderBy); - await ValidateNonDeterministicQuery(CrossPartitionQueryTests.QueryWithSerializeState, useOrderBy); + await ValidateNonDeterministicQuery(CrossPartitionQueryTests.QueryWithoutContinuationTokensAsync, useOrderBy); + await ValidateNonDeterministicQuery(CrossPartitionQueryTests.QueryWithContinuationTokensAsync, useOrderBy); + await ValidateNonDeterministicQuery(CrossPartitionQueryTests.QueryWithCosmosElementContinuationTokenAsync, useOrderBy); } } } @@ -1438,7 +1438,7 @@ private async Task TestExceptionlessFailuresHelper( TestSettings = new TestInjections(simulate429s: true, simulateEmptyPages: false) }; - List queryResults = await CrossPartitionQueryTests.RunQuery( + List queryResults = await CrossPartitionQueryTests.RunQueryAsync( container, query, feedOptions); @@ -1482,7 +1482,7 @@ private async Task TestEmptyPagesHelper( TestSettings = new TestInjections(simulate429s: false, simulateEmptyPages: true) }; - List queryResults = await CrossPartitionQueryTests.RunQuery( + List queryResults = await CrossPartitionQueryTests.RunQueryAsync( container, query, feedOptions); @@ -1542,7 +1542,7 @@ private async Task TestQueryPlanGatewayAndServiceInteropHelper( MaxItemCount = maxItemCount, }; - List queryResults = await CrossPartitionQueryTests.RunQuery( + List queryResults = await CrossPartitionQueryTests.RunQueryAsync( containerWithForcedPlan, "SELECT * FROM c ORDER BY c._ts", feedOptions); @@ -1594,7 +1594,7 @@ private async Task TestUnsupportedQueriesHelper( { try { - await CrossPartitionQueryTests.RunQuery( + await CrossPartitionQueryTests.RunQueryAsync( container, unsupportedQuery, queryRequestOptions: feedOptions); @@ -1613,8 +1613,8 @@ public async Task TestQueryCrossPartitionAggregateFunctions() { AggregateTestArgs aggregateTestArgs = new AggregateTestArgs() { - NumberOfDocsWithSamePartitionKey = 100, - NumberOfDocumentsDifferentPartitionKey = 100, + NumberOfDocsWithSamePartitionKey = 37, + NumberOfDocumentsDifferentPartitionKey = 43, PartitionKey = "key", UniquePartitionKey = "uniquePartitionKey", Field = "field", @@ -1636,18 +1636,14 @@ public async Task TestQueryCrossPartitionAggregateFunctions() { Document doc = new Document(); doc.SetPropertyValue(aggregateTestArgs.PartitionKey, aggregateTestArgs.UniquePartitionKey); - doc.ResourceId = i.ToString(CultureInfo.InvariantCulture); - doc.SetPropertyValue(aggregateTestArgs.Field, i + 1); - doc.SetPropertyValue("id", Guid.NewGuid().ToString()); - documents.Add(doc.ToString()); } + Random random = new Random(); for (int i = 0; i < aggregateTestArgs.NumberOfDocumentsDifferentPartitionKey; ++i) { Document doc = new Document(); - doc.SetPropertyValue(aggregateTestArgs.PartitionKey, i + 1); - doc.SetPropertyValue("id", Guid.NewGuid().ToString()); + doc.SetPropertyValue(aggregateTestArgs.PartitionKey, random.NextDouble()); documents.Add(doc.ToString()); } @@ -1673,63 +1669,70 @@ private struct AggregateTestArgs private struct AggregateQueryArguments { public string AggregateOperator; - public object ExpectedValue; + public CosmosElement ExpectedValue; public string Predicate; } - private async Task TestQueryCrossPartitionAggregateFunctionsAsync(Container container, IEnumerable documents, AggregateTestArgs aggregateTestArgs) + private async Task TestQueryCrossPartitionAggregateFunctionsAsync( + Container container, + IEnumerable documents, + AggregateTestArgs aggregateTestArgs) { - int numberOfDocumentsDifferentPartitionKey = aggregateTestArgs.NumberOfDocumentsDifferentPartitionKey; - int numberOfDocumentSamePartitionKey = aggregateTestArgs.NumberOfDocsWithSamePartitionKey; - int numberOfDocuments = aggregateTestArgs.NumberOfDocumentsDifferentPartitionKey + aggregateTestArgs.NumberOfDocsWithSamePartitionKey; - object[] values = aggregateTestArgs.Values; - string partitionKey = aggregateTestArgs.PartitionKey; - - double samePartitionSum = numberOfDocumentSamePartitionKey * (numberOfDocumentSamePartitionKey + 1) / 2; - double differentPartitionSum = numberOfDocumentsDifferentPartitionKey * (numberOfDocumentsDifferentPartitionKey + 1) / 2; - double partitionSum = samePartitionSum + differentPartitionSum; + IEnumerable documentsWherePkIsANumber = documents + .Where(doc => + { + return double.TryParse( + JObject.Parse(doc.ToString())[aggregateTestArgs.PartitionKey].ToString(), + out double result); + }); + double numberSum = documentsWherePkIsANumber + .Sum(doc => + { + return double.Parse(JObject.Parse(doc.ToString())[aggregateTestArgs.PartitionKey].ToString()); + }); + double count = documentsWherePkIsANumber.Count(); AggregateQueryArguments[] aggregateQueryArgumentsList = new AggregateQueryArguments[] { new AggregateQueryArguments() { AggregateOperator = "AVG", - ExpectedValue = partitionSum / numberOfDocuments, - Predicate = $"IS_NUMBER(r.{partitionKey})", + ExpectedValue = CosmosNumber64.Create(numberSum / count), + Predicate = $"IS_NUMBER(r.{aggregateTestArgs.PartitionKey})", }, new AggregateQueryArguments() { AggregateOperator = "AVG", - ExpectedValue = Undefined.Value, + ExpectedValue = null, Predicate = "true", }, new AggregateQueryArguments() { AggregateOperator = "COUNT", - ExpectedValue = (long)numberOfDocuments + values.Length, + ExpectedValue = CosmosNumber64.Create(documents.Count()), Predicate = "true", }, new AggregateQueryArguments() { AggregateOperator = "MAX", - ExpectedValue = "xyz", + ExpectedValue = CosmosString.Create("xyz"), Predicate = "true", }, new AggregateQueryArguments() { AggregateOperator = "MIN", - ExpectedValue = false, + ExpectedValue = CosmosBoolean.Create(false), Predicate = "true", }, new AggregateQueryArguments() { AggregateOperator = "SUM", - ExpectedValue = differentPartitionSum, - Predicate = $"IS_NUMBER(r.{partitionKey})", + ExpectedValue = CosmosNumber64.Create(numberSum), + Predicate = $"IS_NUMBER(r.{aggregateTestArgs.PartitionKey})", }, new AggregateQueryArguments() { AggregateOperator = "SUM", - ExpectedValue = Undefined.Value, + ExpectedValue = null, Predicate = $"true", }, }; @@ -1746,97 +1749,47 @@ private async Task TestQueryCrossPartitionAggregateFunctionsAsync(Container cont foreach (string queryFormat in queryFormats) { - string query = string.Format(CultureInfo.InvariantCulture, queryFormat, argument.AggregateOperator, partitionKey, argument.Predicate); - string message = string.Format(CultureInfo.InvariantCulture, "query: {0}, data: {1}", query, JsonConvert.SerializeObject(argument)); + string query = string.Format( + CultureInfo.InvariantCulture, + queryFormat, + argument.AggregateOperator, + aggregateTestArgs.PartitionKey, + argument.Predicate); + string message = string.Format( + CultureInfo.InvariantCulture, + "query: {0}, data: {1}", + query, + JsonConvert.SerializeObject(argument)); - List items = await CrossPartitionQueryTests.RunQuery( + List items = await CrossPartitionQueryTests.RunQueryAsync( container, query, new QueryRequestOptions() { MaxConcurrency = maxDoP, }); - if (Undefined.Value.Equals(argument.ExpectedValue)) + + if (argument.ExpectedValue == null) { Assert.AreEqual(0, items.Count, message); } else { - object expected = argument.ExpectedValue; - object actual = items.Single(); + Assert.AreEqual(1, items.Count, message); + CosmosElement expected = argument.ExpectedValue; + CosmosElement actual = items.Single(); - if (expected is long) + if ((expected is CosmosNumber expectedNumber) && (actual is CosmosNumber actualNumber)) { - expected = (double)(long)expected; + Assert.AreEqual(Number64.ToDouble(expectedNumber.Value), Number64.ToDouble(actualNumber.Value), .01); } - - if (actual is long) + else { - actual = (double)(long)actual; + Assert.AreEqual(expected, actual, message); } - - Assert.AreEqual(expected, actual, message); } } } - - // Single partition queries - double singlePartitionSum = samePartitionSum; - Tuple[] datum = new[] - { - Tuple.Create("AVG", singlePartitionSum / numberOfDocumentSamePartitionKey), - Tuple.Create("COUNT", (long)numberOfDocumentSamePartitionKey), - Tuple.Create("MAX", (long)numberOfDocumentSamePartitionKey), - Tuple.Create("MIN", (long)1), - Tuple.Create("SUM", (long)singlePartitionSum), - }; - - string field = aggregateTestArgs.Field; - string uniquePartitionKey = aggregateTestArgs.UniquePartitionKey; - foreach (Tuple data in datum) - { - string query = $"SELECT VALUE {data.Item1}(r.{field}) FROM r WHERE r.{partitionKey} = '{uniquePartitionKey}'"; - dynamic aggregate = (await CrossPartitionQueryTests.RunQuery( - container, - query)).Single(); - object expected = data.Item2; - - if (aggregate is long) - { - aggregate = (long)aggregate; - } - - if (expected is long) - { - expected = (long)expected; - } - - Assert.AreEqual( - expected, - aggregate, - string.Format(CultureInfo.InvariantCulture, "query: {0}, data: {1}", query, JsonConvert.SerializeObject(data))); - - // V3 doesn't support an equivalent to ToList() - // Aggregate queries need to be in the form SELECT VALUE - //query = $"SELECT {data.Item1}(r.{field}) FROM r WHERE r.{partitionKey} = '{uniquePartitionKey}'"; - //try - //{ - // documentClient.CreateDocumentQuery( - // collection, - // query).ToList().Single(); - // Assert.Fail($"Expect exception query: {query}"); - //} - //catch (AggregateException ex) - //{ - // if (!(ex.InnerException is CosmosException) || ((CosmosException)ex.InnerException).StatusCode != HttpStatusCode.BadRequest) - // { - // throw; - // } - //} - - // Make sure ExecuteNextAsync works for unsupported aggregate projection - FeedResponse page = await container.GetItemQueryIterator(query, requestOptions: new QueryRequestOptions() { MaxConcurrency = 1 }).ReadNextAsync(); - } } } @@ -1898,7 +1851,7 @@ private async Task TestQueryCrossPartitionAggregateFunctionsEmptyPartitionsHelpe { try { - List items = await CrossPartitionQueryTests.RunQuery( + List items = await CrossPartitionQueryTests.RunQueryAsync( container, query, new QueryRequestOptions() @@ -2104,7 +2057,7 @@ FROM c StringSplitOptions.None) .Select(x => x.Trim())); - List items = await CrossPartitionQueryTests.RunQuery( + List items = await CrossPartitionQueryTests.RunQueryAsync( container, query, new QueryRequestOptions() @@ -2406,7 +2359,7 @@ private async Task TestNonValueAggregates( { foreach (int maxItemCount in new int[] { 1, 5, 10 }) { - List actual = await RunQuery( + List actual = await RunQueryAsync( container: container, query: query, queryRequestOptions: new QueryRequestOptions() @@ -2440,7 +2393,7 @@ private async Task TestNonValueAggregates( { try { - List actual = await QueryWithoutContinuationTokens( + List actual = await QueryWithoutContinuationTokensAsync( container: container, query: query, queryRequestOptions: new QueryRequestOptions() @@ -2612,7 +2565,7 @@ private async Task TestQueryDistinct(Container container, IEnumerable { string queryWithoutDistinct = string.Format(query, ""); MockDistinctMap documentsSeen = new MockDistinctMap(); - List documentsFromWithoutDistinct = await CrossPartitionQueryTests.RunQueryCombinations( + List documentsFromWithoutDistinct = await CrossPartitionQueryTests.RunQueryCombinationsAsync( container, queryWithoutDistinct, new QueryRequestOptions() @@ -2628,7 +2581,7 @@ private async Task TestQueryDistinct(Container container, IEnumerable foreach (int pageSize in new int[] { 1, 10, 100 }) { string queryWithDistinct = string.Format(query, "DISTINCT"); - List documentsFromWithDistinct = await CrossPartitionQueryTests.RunQueryCombinations( + List documentsFromWithDistinct = await CrossPartitionQueryTests.RunQueryCombinationsAsync( container, queryWithDistinct, new QueryRequestOptions() @@ -2660,7 +2613,7 @@ private async Task TestQueryDistinct(Container container, IEnumerable { string queryWithoutDistinct = string.Format(query, ""); MockDistinctMap documentsSeen = new MockDistinctMap(); - List documentsFromWithoutDistinct = await CrossPartitionQueryTests.RunQueryCombinations( + List documentsFromWithoutDistinct = await CrossPartitionQueryTests.RunQueryCombinationsAsync( container, queryWithoutDistinct, new QueryRequestOptions() @@ -2668,7 +2621,7 @@ private async Task TestQueryDistinct(Container container, IEnumerable MaxConcurrency = 10, MaxItemCount = 100, }, - QueryDrainingMode.HoldState | QueryDrainingMode.SerializeState); + QueryDrainingMode.HoldState | QueryDrainingMode.CosmosElementContinuationToken); documentsFromWithoutDistinct = documentsFromWithoutDistinct .Where(document => documentsSeen.Add(document, out Cosmos.Query.Core.UInt128 hash)) .ToList(); @@ -2676,7 +2629,7 @@ private async Task TestQueryDistinct(Container container, IEnumerable foreach (int pageSize in new int[] { 1, 10, 100 }) { string queryWithDistinct = string.Format(query, "DISTINCT"); - List documentsFromWithDistinct = await CrossPartitionQueryTests.RunQueryCombinations( + List documentsFromWithDistinct = await CrossPartitionQueryTests.RunQueryCombinationsAsync( container, queryWithDistinct, new QueryRequestOptions() @@ -2684,7 +2637,7 @@ private async Task TestQueryDistinct(Container container, IEnumerable MaxConcurrency = 10, MaxItemCount = pageSize }, - QueryDrainingMode.HoldState | QueryDrainingMode.SerializeState); + QueryDrainingMode.HoldState | QueryDrainingMode.CosmosElementContinuationToken); Assert.IsTrue( documentsFromWithDistinct.SequenceEqual(documentsFromWithoutDistinct, JsonTokenEqualityComparer.Value), @@ -2723,7 +2676,7 @@ private async Task TestQueryCrossPartitionTopOrderByDifferentDimensionHelper(Con await CrossPartitionQueryTests.NoOp(); string[] expected = new[] { "documentId2", "documentId5", "documentId8" }; - List query = await CrossPartitionQueryTests.RunQuery( + List query = await CrossPartitionQueryTests.RunQueryAsync( container, "SELECT r.id FROM r ORDER BY r.prop DESC", new QueryRequestOptions() @@ -2946,7 +2899,7 @@ FROM c }; List actualFromQueryWithoutContinutionTokens; - actualFromQueryWithoutContinutionTokens = await CrossPartitionQueryTests.QueryWithoutContinuationTokens( + actualFromQueryWithoutContinutionTokens = await CrossPartitionQueryTests.QueryWithoutContinuationTokensAsync( container, query, queryRequestOptions: feedOptions); @@ -3114,7 +3067,7 @@ private async Task TestQueryCrossPartitionTopOrderByHelper(Container container, List computedResults = new List(); string emptyQueryText = @"SELECT TOP 5 * FROM Root r WHERE r.partitionKey = 9991123 OR r.partitionKey = 9991124 OR r.partitionKey = 99991125"; - List queryEmptyResult = await CrossPartitionQueryTests.RunQuery( + List queryEmptyResult = await CrossPartitionQueryTests.RunQueryAsync( container, emptyQueryText); @@ -3376,7 +3329,7 @@ ORDER BY c.guid expectedResults = expectedResults.Skip(offsetCount); expectedResults = expectedResults.Take(limitCount); - List queryResults = await CrossPartitionQueryTests.RunQuery( + List queryResults = await CrossPartitionQueryTests.RunQueryAsync( container, query, queryRequestOptions: queryRequestOptions); @@ -3613,24 +3566,15 @@ await container.GetItemQueryIterator( foreach (string query in queries) { - List queryResultsWithoutContinuationTokens = await QueryWithoutContinuationTokens( - container, - query); - foreach (int pageSize in new int[] { 1, documentCount / 2, documentCount }) { - List queryResultsWithContinuationTokens = await QueryWithContinuationTokens( + await RunQueryAsync( container, query, new QueryRequestOptions() { MaxItemCount = pageSize, }); - - Assert.AreEqual( - string.Join(", ", queryResultsWithoutContinuationTokens.Select(doc => doc.GetPropertyValue(partitionKey))), - string.Join(", ", queryResultsWithContinuationTokens.Select(doc => doc.GetPropertyValue(partitionKey))), - $"query: {query}, page size: {pageSize}"); } } } @@ -3989,7 +3933,7 @@ FROM root { whereString } MaxConcurrency = 10, }; - List> actual = await CrossPartitionQueryTests.RunQuery>( + List> actual = await CrossPartitionQueryTests.RunQueryAsync>( container, query, queryRequestOptions: feedOptions); @@ -4265,7 +4209,7 @@ FROM c { foreach (int maxItemCount in new int[] { 1, 5, 10 }) { - List actualWithoutContinuationTokens = await CrossPartitionQueryTests.QueryWithoutContinuationTokens( + List actualWithoutContinuationTokens = await CrossPartitionQueryTests.QueryWithoutContinuationTokensAsync( container, query, new QueryRequestOptions() @@ -4276,7 +4220,7 @@ FROM c }); HashSet actualWithoutContinuationTokensSet = new HashSet(actualWithoutContinuationTokens, JsonTokenEqualityComparer.Value); - List actualWithTryGetContinuationTokens = await CrossPartitionQueryTests.QueryWithSerializeState( + List actualWithTryGetContinuationTokens = await CrossPartitionQueryTests.QueryWithCosmosElementContinuationTokenAsync( container, query, new QueryRequestOptions() @@ -4308,7 +4252,7 @@ FROM c { try { - List actual = await CrossPartitionQueryTests.QueryWithContinuationTokens( + List actual = await CrossPartitionQueryTests.QueryWithContinuationTokensAsync( container, "SELECT c.age FROM c GROUP BY c.age", new QueryRequestOptions() @@ -4610,16 +4554,16 @@ private static bool Equals(Headers headers1, Headers headers2) // this.responseLengthBytes.Value = null; //} - private static async Task> RunQuery( + private static async Task> RunQueryAsync( Container container, string query, QueryRequestOptions queryRequestOptions = null) { - return await RunQueryCombinations( + return await RunQueryCombinationsAsync( container, query, queryRequestOptions, - QueryDrainingMode.ContinuationToken | QueryDrainingMode.HoldState | QueryDrainingMode.SerializeState); + QueryDrainingMode.ContinuationToken | QueryDrainingMode.HoldState | QueryDrainingMode.CosmosElementContinuationToken); } [Flags] @@ -4628,10 +4572,10 @@ public enum QueryDrainingMode None = 0, HoldState = 1, ContinuationToken = 2, - SerializeState = 4, + CosmosElementContinuationToken = 4, } - private static async Task> RunQueryCombinations( + private static async Task> RunQueryCombinationsAsync( Container container, string query, QueryRequestOptions queryRequestOptions, @@ -4646,7 +4590,7 @@ private static async Task> RunQueryCombinations( if (queryDrainingMode.HasFlag(QueryDrainingMode.HoldState)) { - List queryResultsWithoutContinuationToken = await QueryWithoutContinuationTokens( + List queryResultsWithoutContinuationToken = await QueryWithoutContinuationTokensAsync( container, query, queryRequestOptions); @@ -4656,7 +4600,7 @@ private static async Task> RunQueryCombinations( if (queryDrainingMode.HasFlag(QueryDrainingMode.ContinuationToken)) { - List queryResultsWithContinuationTokens = await QueryWithContinuationTokens( + List queryResultsWithContinuationTokens = await QueryWithContinuationTokensAsync( container, query, queryRequestOptions); @@ -4664,14 +4608,14 @@ private static async Task> RunQueryCombinations( queryExecutionResults[QueryDrainingMode.ContinuationToken] = queryResultsWithContinuationTokens; } - if (queryDrainingMode.HasFlag(QueryDrainingMode.SerializeState)) + if (queryDrainingMode.HasFlag(QueryDrainingMode.CosmosElementContinuationToken)) { - List queryResultsWithSerializeState = await QueryWithSerializeState( + List queryResultsWithCosmosElementContinuationToken = await QueryWithCosmosElementContinuationTokenAsync( container, query, queryRequestOptions); - queryExecutionResults[QueryDrainingMode.SerializeState] = queryResultsWithSerializeState; + queryExecutionResults[QueryDrainingMode.CosmosElementContinuationToken] = queryResultsWithCosmosElementContinuationToken; } foreach (QueryDrainingMode queryDrainingMode1 in queryExecutionResults.Keys) diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/QueryPipelineMockTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/QueryPipelineMockTests.cs index a09f8bb104..3d54301b6a 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/QueryPipelineMockTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/QueryPipelineMockTests.cs @@ -83,7 +83,7 @@ public async Task TestCosmosCrossPartitionQueryExecutionContextWithEmptyPagesAnd sqlQuerySpec: MockQueryFactory.DefaultQuerySpec, collectionRid: MockQueryFactory.DefaultCollectionRid, partitionedQueryExecutionInfo: new PartitionedQueryExecutionInfo() { QueryInfo = new QueryInfo() }, - partitionKeyRanges: new List() { MockQueryFactory.DefaultPartitionKeyRange }, + partitionKeyRanges: new List { MockQueryFactory.DefaultPartitionKeyRange }, initialPageSize: maxPageSize, maxConcurrency: null, maxItemCount: maxPageSize, @@ -531,7 +531,7 @@ public async Task TestNegativeDistinctComponentCreation() TryCatch tryCreateWhenInvalidContinuationToken = await DistinctDocumentQueryExecutionComponent.TryCreateAsync( ExecutionEnvironment.Client, - CosmosElement.Parse("This is not a valid continuation token"), + CosmosElement.Parse("\"This is not a valid continuation token\""), CreateSource, DistinctQueryType.Unordered); From 144364f7ac183aad46245abe0e3823f63a7c4da3 Mon Sep 17 00:00:00 2001 From: Brandon Chong Date: Mon, 2 Mar 2020 14:16:20 -0800 Subject: [PATCH 19/28] broke up parallel execution context into different files based on stages --- ...QueryExecutionContext.ContinuationToken.cs | 111 ++++++ ...ParallelItemQueryExecutionContext.Drain.cs | 72 ++++ ...arallelItemQueryExecutionContext.Resume.cs | 185 ++++++++++ ...CosmosParallelItemQueryExecutionContext.cs | 323 +----------------- 4 files changed, 374 insertions(+), 317 deletions(-) create mode 100644 Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/Parallel/CosmosParallelItemQueryExecutionContext.ContinuationToken.cs create mode 100644 Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/Parallel/CosmosParallelItemQueryExecutionContext.Drain.cs create mode 100644 Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/Parallel/CosmosParallelItemQueryExecutionContext.Resume.cs 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 new file mode 100644 index 0000000000..2d5a42884a --- /dev/null +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/Parallel/CosmosParallelItemQueryExecutionContext.ContinuationToken.cs @@ -0,0 +1,111 @@ +//------------------------------------------------------------ +// 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: false, + isMaxInclusive: true) + }; + + 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 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.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 new file mode 100644 index 0000000000..45845fa4d8 --- /dev/null +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/Parallel/CosmosParallelItemQueryExecutionContext.Drain.cs @@ -0,0 +1,72 @@ +//------------------------------------------------------------ +// 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, + diagnostics: this.GetAndResetDiagnostics()); + } + } +} 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 new file mode 100644 index 0000000000..e388281bc4 --- /dev/null +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/Parallel/CosmosParallelItemQueryExecutionContext.Resume.cs @@ -0,0 +1,185 @@ +// ------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +// ------------------------------------------------------------ + +namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.Parallel +{ + using System; + using System.Collections.Generic; + using System.Diagnostics; + using System.Linq; + using System.Threading; + using System.Threading.Tasks; + using Microsoft.Azure.Cosmos.CosmosElements; + 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.ExecutionComponent; + using Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.ItemProducers; + using Microsoft.Azure.Cosmos.Query.Core.Monads; + using Microsoft.Azure.Cosmos.Query.Core.QueryClient; + using PartitionKeyRange = Documents.PartitionKeyRange; + + internal sealed partial class CosmosParallelItemQueryExecutionContext : CosmosCrossPartitionQueryExecutionContext + { + public static async Task> TryCreateAsync( + CosmosQueryContext queryContext, + CosmosCrossPartitionQueryExecutionContext.CrossPartitionInitParams initParams, + CosmosElement requestContinuationToken, + CancellationToken cancellationToken) + { + Debug.Assert( + !initParams.PartitionedQueryExecutionInfo.QueryInfo.HasOrderBy, + "Parallel~Context must not have order by query info."); + + 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, + List partitionKeyRanges, + int initialPageSize, + CosmosElement requestContinuation, + CancellationToken cancellationToken) + { + cancellationToken.ThrowIfCancellationRequested(); + + TryCatch tryGetInitInfo = TryGetInitializationInfoFromContinuationToken( + partitionKeyRanges, + requestContinuation); + if (!tryGetInitInfo.Succeeded) + { + return TryCatch.FromException(tryGetInitInfo.Exception); + } + + ParallelInitInfo initializationInfo = tryGetInitInfo.Result; + IReadOnlyList filteredPartitionKeyRanges = initializationInfo.PartialRanges; + IReadOnlyDictionary targetIndicesForFullContinuation = initializationInfo.ContinuationTokens; + TryCatch tryInitialize = await base.TryInitializeAsync( + collectionRid, + filteredPartitionKeyRanges, + initialPageSize, + sqlQuerySpec, + targetIndicesForFullContinuation?.ToDictionary(kvp => kvp.Key, kvp => kvp.Value.Token), + true, + null, + null, + cancellationToken); + if (!tryInitialize.Succeeded) + { + return TryCatch.FromException(tryInitialize.Exception); + } + + return TryCatch.FromResult(this); + } + + /// + /// Given a continuation token and a list of partitionKeyRanges this function will return a list of partition key ranges you should resume with. + /// Note that the output list is just a right hand slice of the input list, since we know that for any continuation of a parallel query it is just + /// resuming from the partition that the query left off that. + /// + /// The partition key ranges. + /// The continuation tokens that the user has supplied. + /// The subset of partition to actually target and continuation tokens. + private static TryCatch TryGetInitializationInfoFromContinuationToken( + List partitionKeyRanges, + CosmosElement continuationToken) + { + if (continuationToken == null) + { + return TryCatch.FromResult( + new ParallelInitInfo( + partitionKeyRanges, + null)); + } + + if (!(continuationToken is CosmosArray compositeContinuationTokenListRaw)) + { + return TryCatch.FromException( + new MalformedContinuationTokenException($"Invalid format for continuation token {continuationToken} 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 CosmosCrossPartitionQueryExecutionContext.TryFindTargetRangeAndExtractContinuationTokens( + partitionKeyRanges, + compositeContinuationTokens.Select(token => Tuple.Create(token, token.Range))) + .Try((indexAndTokens) => + { + int minIndex = indexAndTokens.TargetIndex; + IReadOnlyDictionary rangeToToken = indexAndTokens.ContinuationTokens; + + // We know that all partitions to the left of the continuation token are fully drained so we can filter them out + IReadOnlyList filteredRanges = new PartialReadOnlyList( + partitionKeyRanges, + minIndex, + partitionKeyRanges.Count - minIndex); + + return new ParallelInitInfo( + filteredRanges, + rangeToToken); + }); + } + + private readonly struct ParallelInitInfo + { + public ParallelInitInfo(IReadOnlyList partialRanges, IReadOnlyDictionary continuationTokens) + { + this.PartialRanges = partialRanges; + this.ContinuationTokens = continuationTokens; + } + + public IReadOnlyList PartialRanges { get; } + + public IReadOnlyDictionary ContinuationTokens { get; } + } + } +} 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 a3d3553b83..56a2f6225e 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/Parallel/CosmosParallelItemQueryExecutionContext.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/Parallel/CosmosParallelItemQueryExecutionContext.cs @@ -5,22 +5,9 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.Parallel { using System; using System.Collections.Generic; - using System.Diagnostics; - using System.Linq; - using System.Threading; - using System.Threading.Tasks; using Microsoft.Azure.Cosmos.CosmosElements; - using Microsoft.Azure.Cosmos.Json; - using Microsoft.Azure.Cosmos.Query.Core; - 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.ExecutionComponent; using Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.ItemProducers; - using Microsoft.Azure.Cosmos.Query.Core.Monads; using Microsoft.Azure.Cosmos.Query.Core.QueryClient; - using Newtonsoft.Json; - using PartitionKeyRange = Documents.PartitionKeyRange; /// /// CosmosParallelItemQueryExecutionContext is a concrete implementation for CrossPartitionQueryExecutionContext. @@ -29,20 +16,20 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.Parallel /// 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 CosmosParallelItemQueryExecutionContext : CosmosCrossPartitionQueryExecutionContext + internal sealed partial class CosmosParallelItemQueryExecutionContext : CosmosCrossPartitionQueryExecutionContext { /// - /// The function to determine which partition to fetch from first. + /// The comparer used to determine, which continuation tokens should be returned to the user. /// - private static readonly Func FetchPriorityFunction = documentProducerTree => int.Parse(documentProducerTree.PartitionKeyRange.Id); + private static readonly IEqualityComparer EqualityComparer = new ParallelEqualityComparer(); /// - /// The comparer used to determine, which continuation tokens should be returned to the user. + /// The function to determine which partition to fetch from first. /// - private static readonly IEqualityComparer EqualityComparer = new ParallelEqualityComparer(); + private static readonly Func FetchPriorityFunction = documentProducerTree => int.Parse(documentProducerTree.PartitionKeyRange.Id); private readonly bool returnResultsInDeterministicOrder; - + /// /// Initializes a new instance of the CosmosParallelItemQueryExecutionContext class. /// @@ -74,303 +61,5 @@ private CosmosParallelItemQueryExecutionContext( { this.returnResultsInDeterministicOrder = returnResultsInDeterministicOrder; } - - /// - /// 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 static async Task> TryCreateAsync( - CosmosQueryContext queryContext, - CosmosCrossPartitionQueryExecutionContext.CrossPartitionInitParams initParams, - CosmosElement requestContinuationToken, - CancellationToken cancellationToken) - { - Debug.Assert( - !initParams.PartitionedQueryExecutionInfo.QueryInfo.HasOrderBy, - "Parallel~Context must not have order by query info."); - - 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); - } - - 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, - diagnostics: this.GetAndResetDiagnostics()); - } - - /// - /// 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, - List partitionKeyRanges, - int initialPageSize, - CosmosElement requestContinuation, - CancellationToken cancellationToken) - { - cancellationToken.ThrowIfCancellationRequested(); - - TryCatch tryGetInitInfo = TryGetInitializationInfoFromContinuationToken( - partitionKeyRanges, - requestContinuation); - if (!tryGetInitInfo.Succeeded) - { - return TryCatch.FromException(tryGetInitInfo.Exception); - } - - ParallelInitInfo initializationInfo = tryGetInitInfo.Result; - IReadOnlyList filteredPartitionKeyRanges = initializationInfo.PartialRanges; - IReadOnlyDictionary targetIndicesForFullContinuation = initializationInfo.ContinuationTokens; - TryCatch tryInitialize = await base.TryInitializeAsync( - collectionRid, - filteredPartitionKeyRanges, - initialPageSize, - sqlQuerySpec, - targetIndicesForFullContinuation?.ToDictionary(kvp => kvp.Key, kvp => kvp.Value.Token), - true, - null, - null, - cancellationToken); - if (!tryInitialize.Succeeded) - { - return TryCatch.FromException(tryInitialize.Exception); - } - - return TryCatch.FromResult(this); - } - - /// - /// Given a continuation token and a list of partitionKeyRanges this function will return a list of partition key ranges you should resume with. - /// Note that the output list is just a right hand slice of the input list, since we know that for any continuation of a parallel query it is just - /// resuming from the partition that the query left off that. - /// - /// The partition key ranges. - /// The continuation tokens that the user has supplied. - /// The subset of partition to actually target and continuation tokens. - private static TryCatch TryGetInitializationInfoFromContinuationToken( - List partitionKeyRanges, - CosmosElement continuationToken) - { - if (continuationToken == null) - { - return TryCatch.FromResult( - new ParallelInitInfo( - partitionKeyRanges, - null)); - } - - if (!(continuationToken is CosmosArray compositeContinuationTokenListRaw)) - { - return TryCatch.FromException( - new MalformedContinuationTokenException($"Invalid format for continuation token {continuationToken} 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 CosmosCrossPartitionQueryExecutionContext.TryFindTargetRangeAndExtractContinuationTokens( - partitionKeyRanges, - compositeContinuationTokens.Select(token => Tuple.Create(token, token.Range))) - .Try((indexAndTokens) => - { - int minIndex = indexAndTokens.TargetIndex; - IReadOnlyDictionary rangeToToken = indexAndTokens.ContinuationTokens; - - // We know that all partitions to the left of the continuation token are fully drained so we can filter them out - IReadOnlyList filteredRanges = new PartialReadOnlyList( - partitionKeyRanges, - minIndex, - partitionKeyRanges.Count - minIndex); - - return new ParallelInitInfo( - filteredRanges, - rangeToToken); - }); - } - - 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: false, - isMaxInclusive: true) - }; - - CosmosElement compositeContinuationToken = CompositeContinuationToken.ToCosmosElement(compositeToken); - compositeContinuationTokens.Add(compositeContinuationToken); - } - - return CosmosArray.Create(compositeContinuationTokens); - } - - private readonly struct ParallelInitInfo - { - public ParallelInitInfo(IReadOnlyList partialRanges, IReadOnlyDictionary continuationTokens) - { - this.PartialRanges = partialRanges; - this.ContinuationTokens = continuationTokens; - } - - public IReadOnlyList PartialRanges { get; } - - public IReadOnlyDictionary ContinuationTokens { get; } - } - - /// - /// 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 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.GetHashCode(); - } - } } } From 26e2fc5bbabaee9e63e25480e2bca414a4170935 Mon Sep 17 00:00:00 2001 From: Brandon Chong Date: Mon, 2 Mar 2020 18:48:23 -0800 Subject: [PATCH 20/28] refactored parrallel resume code --- .../CompositeContinuationToken.cs | 5 +- .../ContinuationTokens/IPartitionedToken.cs | 11 ++ .../OrderByContinuationToken.cs | 5 +- ...smosCrossPartitionQueryExecutionContext.cs | 111 +++++----------- ...arallelItemQueryExecutionContext.Resume.cs | 122 ++++++++---------- 5 files changed, 108 insertions(+), 146 deletions(-) create mode 100644 Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/IPartitionedToken.cs diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/CompositeContinuationToken.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/CompositeContinuationToken.cs index 78d0290bb9..1f1c16f24f 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/CompositeContinuationToken.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/CompositeContinuationToken.cs @@ -10,12 +10,13 @@ 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.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 + internal sealed class CompositeContinuationToken : IPartitionedToken { private static class PropertyNames { @@ -41,6 +42,8 @@ public Documents.Routing.Range Range set; } + public Range PartitionRange => this.Range; + public object ShallowCopy() { return this.MemberwiseClone(); diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/IPartitionedToken.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/IPartitionedToken.cs new file mode 100644 index 0000000000..2e81a39fa1 --- /dev/null +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/IPartitionedToken.cs @@ -0,0 +1,11 @@ +//------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +//------------------------------------------------------------ + +namespace Microsoft.Azure.Cosmos.Query.Core.ContinuationTokens +{ + internal interface IPartitionedToken + { + Documents.Routing.Range PartitionRange { get; } + } +} diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/OrderByContinuationToken.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/OrderByContinuationToken.cs index 6fd5929708..720a3a6c13 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/OrderByContinuationToken.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/OrderByContinuationToken.cs @@ -11,6 +11,7 @@ namespace 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.Documents.Routing; using Newtonsoft.Json; /// @@ -53,7 +54,7 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ContinuationTokens /// {"compositeToken":{"token":"+RID:OpY0AN-mFAACAAAAAAAABA==#RT:1#TRC:1#RTD:qdTAEA==","range":{"min":"05C1D9CD673398","max":"05C1E399CD6732"}},"orderByItems"[{"item":2}],"rid":"OpY0AN-mFAACAAAAAAAABA==","skipCount":0,"filter":"r.key > 1"} /// ]]> /// - internal sealed class OrderByContinuationToken + internal sealed class OrderByContinuationToken : IPartitionedToken { private static class PropertyNames { @@ -206,6 +207,8 @@ public string Filter get; } + public Range PartitionRange => this.CompositeContinuationToken.Range; + public static CosmosElement ToCosmosElement(OrderByContinuationToken orderByContinuationToken) { CosmosElement compositeContinuationToken = CompositeContinuationToken.ToCosmosElement(orderByContinuationToken.CompositeContinuationToken); diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosCrossPartitionQueryExecutionContext.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosCrossPartitionQueryExecutionContext.cs index ea2964f834..490fed658f 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosCrossPartitionQueryExecutionContext.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosCrossPartitionQueryExecutionContext.cs @@ -12,16 +12,17 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionContext using Core.ExecutionComponent; using Microsoft.Azure.Cosmos.CosmosElements; using Microsoft.Azure.Cosmos.Diagnostics; - using Microsoft.Azure.Cosmos.Json; 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 Microsoft.Azure.Documents.Routing; using PartitionKeyRange = Documents.PartitionKeyRange; using RequestChargeTracker = Documents.RequestChargeTracker; using RMResources = Documents.RMResources; @@ -322,25 +323,11 @@ public override void Stop() this.comparableTaskScheduler.Stop(); } - /// - /// Initializes cross partition query execution context by initializing the necessary document producers. - /// - /// The collection to drain from. - /// The partitions to target. - /// The page size to start the document producers off with. - /// The query specification for the rewritten query. - /// Map from partition to it's corresponding continuation token. - /// Whether or not we should defer the fetch of the first page from each partition. - /// The filter to inject in the predicate. - /// The callback used to filter each partition. - /// The cancellation token. - /// A task to await on. protected async Task TryInitializeAsync( string collectionRid, - IReadOnlyList partitionKeyRanges, int initialPageSize, SqlQuerySpec querySpecForInit, - IReadOnlyDictionary targetRangeToContinuationMap, + IReadOnlyDictionary targetRangeToContinuationMap, bool deferFirstPage, string filter, Func> tryFilterAsync, @@ -438,7 +425,7 @@ protected async Task TryInitializeAsync( /// /// - /// If a query encounters split up resuming using continuation, we need to regenerate the continuation tokens. + /// If a query encounters split upon resuming using continuation, we need to regenerate the continuation tokens. /// Specifically, since after split we will have new set of ranges, we need to remove continuation token for the /// parent partition and introduce continuation token for the child partitions. /// @@ -451,90 +438,60 @@ protected async Task TryInitializeAsync( /// /// /// The partition key ranges to extract continuation tokens for. - /// The continuation token that the user supplied. - /// The type of continuation token to generate. + /// The continuation token that the user supplied. + /// The type of continuation token to generate. /// /// The code assumes that merge doesn't happen and /// /// The index of the partition whose MinInclusive is equal to the suppliedContinuationTokens along with the continuation tokens. - protected static TryCatch> TryFindTargetRangeAndExtractContinuationTokens( - List partitionKeyRanges, - IEnumerable>> suppliedContinuationTokens) + protected static TryCatch> TryMatchRangesToContinuationTokens( + IReadOnlyList partitionKeyRanges, + IReadOnlyList partitionedContinuationTokens) + where PartitionedToken : IPartitionedToken { if (partitionKeyRanges == null) { throw new ArgumentNullException(nameof(partitionKeyRanges)); } - if (partitionKeyRanges.Count < 1) + if (partitionKeyRanges.Count() < 1) { throw new ArgumentException(nameof(partitionKeyRanges)); } - foreach (PartitionKeyRange partitionKeyRange in partitionKeyRanges) - { - if (partitionKeyRange == null) - { - throw new ArgumentException(nameof(partitionKeyRanges)); - } - } - - if (suppliedContinuationTokens == null) - { - throw new ArgumentNullException(nameof(suppliedContinuationTokens)); - } - - if (suppliedContinuationTokens.Count() < 1) + if (partitionedContinuationTokens == null) { - throw new ArgumentException(nameof(suppliedContinuationTokens)); + throw new ArgumentNullException(nameof(partitionedContinuationTokens)); } - if (suppliedContinuationTokens.Count() > partitionKeyRanges.Count) + if (partitionedContinuationTokens.Count() < 1) { - throw new ArgumentException($"{nameof(suppliedContinuationTokens)} can not have more elements than {nameof(partitionKeyRanges)}."); + throw new ArgumentException(nameof(partitionedContinuationTokens)); } - Dictionary targetRangeToContinuationTokenMap = new Dictionary(); - - // Find the minimum index. - Tuple> firstContinuationTokenAndRange = suppliedContinuationTokens - .OrderBy((tuple) => tuple.Item2.Min) - .First(); - TContinuationToken firstContinuationToken = firstContinuationTokenAndRange.Item1; - PartitionKeyRange firstContinuationRange = new PartitionKeyRange + if (partitionedContinuationTokens.Count() > partitionKeyRanges.Count()) { - MinInclusive = firstContinuationTokenAndRange.Item2.Min, - MaxExclusive = firstContinuationTokenAndRange.Item2.Max - }; - - int minIndex = partitionKeyRanges.BinarySearch( - firstContinuationRange, - Comparer.Create((range1, range2) => string.CompareOrdinal(range1.MinInclusive, range2.MinInclusive))); - if (minIndex < 0) - { - return TryCatch>.FromException( - new MalformedContinuationTokenException( - $"{RMResources.InvalidContinuationToken} - Could not find continuation token: {firstContinuationToken}")); + throw new ArgumentException($"{nameof(partitionedContinuationTokens)} can not have more elements than {nameof(partitionKeyRanges)}."); } - foreach (Tuple> suppledContinuationToken in suppliedContinuationTokens) + Dictionary partitionKeyRangeToToken = new Dictionary(); + foreach (IPartitionedToken partitionedToken in partitionedContinuationTokens) { - // find what ranges make up the supplied continuation token - TContinuationToken continuationToken = suppledContinuationToken.Item1; - Documents.Routing.Range range = suppledContinuationToken.Item2; - - IEnumerable replacementRanges = partitionKeyRanges + // Find what ranges make up the supplied continuation token + Range range = partitionedToken.PartitionRange; + IReadOnlyList replacementRanges = partitionKeyRanges .Where((partitionKeyRange) => string.CompareOrdinal(range.Min, partitionKeyRange.MinInclusive) <= 0 && string.CompareOrdinal(range.Max, partitionKeyRange.MaxExclusive) >= 0) - .OrderBy((partitionKeyRange) => partitionKeyRange.MinInclusive); + .OrderBy((partitionKeyRange) => partitionKeyRange.MinInclusive) + .ToList(); // Could not find the child ranges - if (replacementRanges.Count() == 0) + if (replacementRanges.Count == 0) { - return TryCatch>.FromException( + return TryCatch>.FromException( new MalformedContinuationTokenException( - $"{RMResources.InvalidContinuationToken} - Could not find continuation token: {continuationToken}")); + $"{RMResources.InvalidContinuationToken} - Could not find continuation token: {partitionedToken}")); } // PMax = C2Max > C2Min > C1Max > C1Min = PMin. @@ -551,21 +508,19 @@ protected static TryCatch> TryFindTargetRangeAndExt string.CompareOrdinal(child1Max, child1Min) >= 0 && child1Min == parentMin)) { - return TryCatch>.FromException( + return TryCatch>.FromException( new MalformedContinuationTokenException( - $"{RMResources.InvalidContinuationToken} - PMax = C2Max > C2Min > C1Max > C1Min = PMin: {continuationToken}")); + $"{RMResources.InvalidContinuationToken} - PMax = C2Max > C2Min > C1Max > C1Min = PMin: {partitionedToken}")); } foreach (PartitionKeyRange partitionKeyRange in replacementRanges) { - targetRangeToContinuationTokenMap.Add(partitionKeyRange.Id, continuationToken); + partitionKeyRangeToToken[partitionKeyRange] = partitionedToken; } } - return TryCatch>.FromResult( - new InitInfo( - minIndex, - targetRangeToContinuationTokenMap)); + return TryCatch>.FromResult( + (IReadOnlyDictionary)partitionKeyRangeToToken); } protected virtual long GetAndResetResponseLengthBytes() @@ -677,7 +632,7 @@ public InitInfo(int targetIndex, IReadOnlyDictionary /// 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 struct CrossPartitionInitParams + public readonly struct CrossPartitionInitParams { /// /// Initializes a new instance of the InitParams struct. 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 e388281bc4..54d4a47fe5 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 @@ -6,12 +6,10 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.Parallel { using System; using System.Collections.Generic; - using System.Diagnostics; using System.Linq; using System.Threading; using System.Threading.Tasks; using Microsoft.Azure.Cosmos.CosmosElements; - 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.ExecutionComponent; @@ -28,9 +26,10 @@ public static async Task> TryCreateAs CosmosElement requestContinuationToken, CancellationToken cancellationToken) { - Debug.Assert( - !initParams.PartitionedQueryExecutionInfo.QueryInfo.HasOrderBy, - "Parallel~Context must not have order by query info."); + if (queryContext == null) + { + throw new ArgumentNullException(nameof(queryContext)); + } cancellationToken.ThrowIfCancellationRequested(); @@ -75,33 +74,32 @@ public static async Task> TryCreateAs private async Task> TryInitializeAsync( SqlQuerySpec sqlQuerySpec, string collectionRid, - List partitionKeyRanges, + IReadOnlyList partitionKeyRanges, int initialPageSize, CosmosElement requestContinuation, CancellationToken cancellationToken) { cancellationToken.ThrowIfCancellationRequested(); - TryCatch tryGetInitInfo = TryGetInitializationInfoFromContinuationToken( + TryCatch> tryGetPartitionKeyRangeToCompositeContinuationToken = CosmosParallelItemQueryExecutionContext.TryGetPartitionKeyRangeToCompositeContinuationToken( partitionKeyRanges, requestContinuation); - if (!tryGetInitInfo.Succeeded) + if (!tryGetPartitionKeyRangeToCompositeContinuationToken.Succeeded) { - return TryCatch.FromException(tryGetInitInfo.Exception); + return TryCatch.FromException(tryGetPartitionKeyRangeToCompositeContinuationToken.Exception); } - ParallelInitInfo initializationInfo = tryGetInitInfo.Result; - IReadOnlyList filteredPartitionKeyRanges = initializationInfo.PartialRanges; - IReadOnlyDictionary targetIndicesForFullContinuation = initializationInfo.ContinuationTokens; + IReadOnlyDictionary partitionKeyRangeToContinuationToken = tryGetPartitionKeyRangeToCompositeContinuationToken + .Result + .ToDictionary(kvp => kvp.Key, kvp => kvp.Value.Token); TryCatch tryInitialize = await base.TryInitializeAsync( collectionRid, - filteredPartitionKeyRanges, initialPageSize, sqlQuerySpec, - targetIndicesForFullContinuation?.ToDictionary(kvp => kvp.Key, kvp => kvp.Value.Token), - true, - null, - null, + partitionKeyRangeToContinuationToken, + deferFirstPage: true, + filter: null, + tryFilterAsync: null, cancellationToken); if (!tryInitialize.Succeeded) { @@ -111,30 +109,52 @@ private async Task> TryInitial return TryCatch.FromResult(this); } - /// - /// Given a continuation token and a list of partitionKeyRanges this function will return a list of partition key ranges you should resume with. - /// Note that the output list is just a right hand slice of the input list, since we know that for any continuation of a parallel query it is just - /// resuming from the partition that the query left off that. - /// - /// The partition key ranges. - /// The continuation tokens that the user has supplied. - /// The subset of partition to actually target and continuation tokens. - private static TryCatch TryGetInitializationInfoFromContinuationToken( - List partitionKeyRanges, + private static TryCatch> TryGetPartitionKeyRangeToCompositeContinuationToken( + IReadOnlyList partitionKeyRanges, CosmosElement continuationToken) { if (continuationToken == null) { - return TryCatch.FromResult( - new ParallelInitInfo( - partitionKeyRanges, - null)); + Dictionary dictionary = new Dictionary(); + foreach (PartitionKeyRange partitionKeyRange in partitionKeyRanges) + { + dictionary.Add(key: partitionKeyRange, value: null); + } + + return TryCatch>.FromResult(dictionary); } - if (!(continuationToken is CosmosArray compositeContinuationTokenListRaw)) + TryCatch> tryParseCompositeContinuationTokens = TryParseCompositeContinuationList(continuationToken); + if (!tryParseCompositeContinuationTokens.Succeeded) { - return TryCatch.FromException( - new MalformedContinuationTokenException($"Invalid format for continuation token {continuationToken} for {nameof(CosmosParallelItemQueryExecutionContext)}")); + return TryCatch>.FromException(tryParseCompositeContinuationTokens.Exception); + } + + TryCatch> tryMatchContinuationTokensToRanges = CosmosCrossPartitionQueryExecutionContext.TryMatchRangesToContinuationTokens( + partitionKeyRanges, + tryParseCompositeContinuationTokens.Result); + if (!tryMatchContinuationTokensToRanges.Succeeded) + { + return TryCatch>.FromException( + tryMatchContinuationTokensToRanges.Exception); + } + + return tryMatchContinuationTokensToRanges; + } + + 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(); @@ -143,43 +163,13 @@ private static TryCatch TryGetInitializationInfoFromContinuati TryCatch tryCreateCompositeContinuationToken = CompositeContinuationToken.TryCreateFromCosmosElement(compositeContinuationTokenRaw); if (!tryCreateCompositeContinuationToken.Succeeded) { - return TryCatch.FromException(tryCreateCompositeContinuationToken.Exception); + return TryCatch>.FromException(tryCreateCompositeContinuationToken.Exception); } compositeContinuationTokens.Add(tryCreateCompositeContinuationToken.Result); } - return CosmosCrossPartitionQueryExecutionContext.TryFindTargetRangeAndExtractContinuationTokens( - partitionKeyRanges, - compositeContinuationTokens.Select(token => Tuple.Create(token, token.Range))) - .Try((indexAndTokens) => - { - int minIndex = indexAndTokens.TargetIndex; - IReadOnlyDictionary rangeToToken = indexAndTokens.ContinuationTokens; - - // We know that all partitions to the left of the continuation token are fully drained so we can filter them out - IReadOnlyList filteredRanges = new PartialReadOnlyList( - partitionKeyRanges, - minIndex, - partitionKeyRanges.Count - minIndex); - - return new ParallelInitInfo( - filteredRanges, - rangeToToken); - }); - } - - private readonly struct ParallelInitInfo - { - public ParallelInitInfo(IReadOnlyList partialRanges, IReadOnlyDictionary continuationTokens) - { - this.PartialRanges = partialRanges; - this.ContinuationTokens = continuationTokens; - } - - public IReadOnlyList PartialRanges { get; } - - public IReadOnlyDictionary ContinuationTokens { get; } + return TryCatch>.FromResult(compositeContinuationTokens); } } } From 3f65c719e1dfe324321261781a9f54eff7f55eaf Mon Sep 17 00:00:00 2001 From: Brandon Chong Date: Mon, 2 Mar 2020 18:58:13 -0800 Subject: [PATCH 21/28] broke order by into separate files --- ...QueryExecutionContext.ContinuationToken.cs | 144 +++ ...sOrderByItemQueryExecutionContext.Drain.cs | 140 +++ ...OrderByItemQueryExecutionContext.Resume.cs | 691 +++++++++++++ .../CosmosOrderByItemQueryExecutionContext.cs | 945 +----------------- 4 files changed, 979 insertions(+), 941 deletions(-) create mode 100644 Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/OrderBy/CosmosOrderByItemQueryExecutionContext.ContinuationToken.cs create mode 100644 Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/OrderBy/CosmosOrderByItemQueryExecutionContext.Drain.cs create mode 100644 Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/OrderBy/CosmosOrderByItemQueryExecutionContext.Resume.cs 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 new file mode 100644 index 0000000000..dd4ddd4c19 --- /dev/null +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/OrderBy/CosmosOrderByItemQueryExecutionContext.ContinuationToken.cs @@ -0,0 +1,144 @@ +namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.OrderBy +{ + 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 new file mode 100644 index 0000000000..b5da8ae6b5 --- /dev/null +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/OrderBy/CosmosOrderByItemQueryExecutionContext.Drain.cs @@ -0,0 +1,140 @@ +namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.OrderBy +{ + 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, + diagnostics: this.GetAndResetDiagnostics()); + } + + /// + /// 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 new file mode 100644 index 0000000000..c5e0c50421 --- /dev/null +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/OrderBy/CosmosOrderByItemQueryExecutionContext.Resume.cs @@ -0,0 +1,691 @@ +//------------------------------------------------------------ +// 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.Collections; + 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 PartitionKeyRange = Documents.PartitionKeyRange; + using ResourceId = Documents.ResourceId; + + internal sealed partial class CosmosOrderByItemQueryExecutionContext : CosmosCrossPartitionQueryExecutionContext + { + public static async Task> TryCreateAsync( + 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(); + + // 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, + maxConcurrency: initParams.MaxConcurrency, + maxItemCount: initParams.MaxItemCount, + maxBufferedItemCount: initParams.MaxBufferedItemCount, + consumeComparer: orderByItemProducerTreeComparer, + testSettings: initParams.TestSettings); + + return (await context.TryInitializeAsync( + sqlQuerySpec: initParams.SqlQuerySpec, + requestContinuation: requestContinuationToken, + collectionRid: initParams.CollectionRid, + partitionKeyRanges: initParams.PartitionKeyRanges, + initialPageSize: initParams.InitialPageSize, + sortOrders: initParams.PartitionedQueryExecutionInfo.QueryInfo.OrderBy.ToArray(), + orderByExpressions: initParams.PartitionedQueryExecutionInfo.QueryInfo.OrderByExpressions.ToArray(), + cancellationToken: cancellationToken)) + .Try(() => context); + } + + private async Task TryInitializeAsync( + SqlQuerySpec sqlQuerySpec, + CosmosElement requestContinuation, + string collectionRid, + List partitionKeyRanges, + int initialPageSize, + SortOrder[] sortOrders, + string[] orderByExpressions, + 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 (sortOrders == null) + { + throw new ArgumentNullException(nameof(sortOrders)); + } + + if (orderByExpressions == null) + { + throw new ArgumentNullException(nameof(orderByExpressions)); + } + + cancellationToken.ThrowIfCancellationRequested(); + + if (requestContinuation == null) + { + SqlQuerySpec sqlQuerySpecForInit = new SqlQuerySpec( + sqlQuerySpec.QueryText.Replace(oldValue: FormatPlaceHolder, newValue: True), + sqlQuerySpec.Parameters); + + TryCatch tryInitialize = await base.TryInitializeAsync( + collectionRid, + partitionKeyRanges, + initialPageSize, + sqlQuerySpecForInit, + cancellationToken: cancellationToken, + targetRangeToContinuationMap: null, + deferFirstPage: false, + filter: null, + tryFilterAsync: null); + if (!tryInitialize.Succeeded) + { + return tryInitialize; + } + } + else + { + TryCatch tryExtractContinuationTokens = CosmosOrderByItemQueryExecutionContext.TryExtractContinuationTokens( + requestContinuation, + sortOrders, + orderByExpressions); + if (!tryExtractContinuationTokens.Succeeded) + { + return TryCatch.FromException(tryExtractContinuationTokens.Exception); + } + + TryCatch tryGetOrderByInitInfo = CosmosOrderByItemQueryExecutionContext.TryGetOrderByPartitionKeyRangesInitializationInfo( + tryExtractContinuationTokens.Result, + partitionKeyRanges, + sortOrders, + orderByExpressions); + if (!tryGetOrderByInitInfo.Succeeded) + { + return TryCatch.FromException(tryGetOrderByInitInfo.Exception); + } + + OrderByInitInfo initiaizationInfo = tryGetOrderByInitInfo.Result; + RangeFilterInitializationInfo[] orderByInfos = initiaizationInfo.Filters; + IReadOnlyDictionary targetRangeToOrderByContinuationMap = initiaizationInfo.ContinuationTokens; + Debug.Assert( + targetRangeToOrderByContinuationMap != null, + "If targetRangeToOrderByContinuationMap can't be null is valid continuation is supplied"); + + // 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) + foreach (RangeFilterInitializationInfo info in orderByInfos) + { + if (info.StartIndex > info.EndIndex) + { + continue; + } + + PartialReadOnlyList partialRanges = + new PartialReadOnlyList( + partitionKeyRanges, + info.StartIndex, + info.EndIndex - info.StartIndex + 1); + + SqlQuerySpec sqlQuerySpecForInit = new SqlQuerySpec( + sqlQuerySpec.QueryText.Replace(FormatPlaceHolder, info.Filter), + sqlQuerySpec.Parameters); + + TryCatch tryInitialize = await base.TryInitializeAsync( + collectionRid, + partialRanges, + initialPageSize, + sqlQuerySpecForInit, + targetRangeToOrderByContinuationMap.ToDictionary( + kvp => kvp.Key, + kvp => kvp.Value.CompositeContinuationToken.Token), + false, + info.Filter, + async (itemProducerTree) => + { + if (targetRangeToOrderByContinuationMap.TryGetValue( + itemProducerTree.Root.PartitionKeyRange.Id, + out OrderByContinuationToken continuationToken)) + { + TryCatch tryFilter = await this.TryFilterAsync( + itemProducerTree, + sortOrders, + continuationToken, + cancellationToken); + + if (!tryFilter.Succeeded) + { + return tryFilter; + } + } + + return TryCatch.FromResult(); + }, + cancellationToken); + if (!tryInitialize.Succeeded) + { + return tryInitialize; + } + } + } + + return TryCatch.FromResult(); + } + + private static TryCatch TryExtractContinuationTokens( + CosmosElement requestContinuation, + SortOrder[] sortOrders, + string[] orderByExpressions) + { + Debug.Assert( + !(orderByExpressions == null + || orderByExpressions.Length <= 0 + || sortOrders == null + || sortOrders.Length <= 0 + || orderByExpressions.Length != sortOrders.Length), + "Partitioned QueryExecutionInfo returned bogus results."); + + if (requestContinuation == null) + { + throw new ArgumentNullException("continuation can not be null or empty."); + } + + 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 != sortOrders.Length) + { + return TryCatch.FromException( + new MalformedContinuationTokenException($"Invalid order-by items in continuation token {requestContinuation} for OrderBy~Context.")); + } + } + + return TryCatch.FromResult(orderByContinuationTokens.ToArray()); + } + + /// + /// 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, + SortOrder[] 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.Length; ++i) + { + cmp = ItemComparer.Instance.Compare( + continuationToken.OrderByItems[i].Item, + orderByResult.OrderByItems[i].Item); + + if (cmp != 0) + { + cmp = sortOrders[i] != SortOrder.Descending ? 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 first order by field. + // So if it's ORDER BY c.age ASC, c.name DESC the _rids are ASC + // If ti's ORDER BY c.age DESC, c.name DESC the _rids are DESC + cmp = continuationRid.Document.CompareTo(rid.Document); + if (sortOrders[0] == SortOrder.Descending) + { + 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(); + } + + /// + /// Gets the filters for every partition. + /// + private static TryCatch TryGetOrderByPartitionKeyRangesInitializationInfo( + OrderByContinuationToken[] suppliedContinuationTokens, + List partitionKeyRanges, + SortOrder[] sortOrders, + string[] orderByExpressions) + { + TryCatch> tryFindRangeAndContinuationTokensMonad = CosmosCrossPartitionQueryExecutionContext.TryFindTargetRangeAndExtractContinuationTokens( + partitionKeyRanges, + suppliedContinuationTokens + .Select(token => Tuple.Create(token, token.CompositeContinuationToken.Range))); + + return tryFindRangeAndContinuationTokensMonad.Try((indexAndContinuationTokens) => + { + int minIndex = indexAndContinuationTokens.TargetIndex; + IReadOnlyDictionary partitionKeyRangeToContinuationToken = indexAndContinuationTokens.ContinuationTokens; + + FormattedFilterInfo formattedFilterInfo = CosmosOrderByItemQueryExecutionContext.GetFormattedFilters( + orderByExpressions, + suppliedContinuationTokens, + sortOrders); + + RangeFilterInitializationInfo[] filters = new RangeFilterInitializationInfo[] + { + new RangeFilterInitializationInfo( + filter: formattedFilterInfo.FilterForRangesLeftOfTargetRanges, + startIndex: 0, + endIndex: minIndex - 1), + new RangeFilterInitializationInfo( + filter: formattedFilterInfo.FiltersForTargetRange, + startIndex: minIndex, + endIndex: minIndex), + new RangeFilterInitializationInfo( + filter: formattedFilterInfo.FilterForRangesRightOfTargetRanges, + startIndex: minIndex + 1, + endIndex: partitionKeyRanges.Count - 1), + }; + + return new OrderByInitInfo( + filters, + partitionKeyRangeToContinuationToken); + }); + } + + /// + /// Gets the formatted filters for every partition. + /// + /// The filter expressions. + /// The continuation token. + /// The sort orders. + /// The formatted filters for every partition. + private static FormattedFilterInfo GetFormattedFilters( + string[] expressions, + OrderByContinuationToken[] continuationTokens, + SortOrder[] sortOrders) + { + // Validate the inputs + for (int index = 0; index < continuationTokens.Length; index++) + { + Debug.Assert(continuationTokens[index].OrderByItems.Count == sortOrders.Length, "Expect values and orders are the same size."); + Debug.Assert(expressions.Length == sortOrders.Length, "Expect expressions and orders are the same size."); + } + + Tuple filters = CosmosOrderByItemQueryExecutionContext.GetFormattedFilters( + expressions, + continuationTokens[0].OrderByItems.Select(orderByItem => orderByItem.Item).ToArray(), + sortOrders); + + return new FormattedFilterInfo(filters.Item1, filters.Item2, filters.Item3); + } + + private static void AppendToBuilders(Tuple builders, object str) + { + CosmosOrderByItemQueryExecutionContext.AppendToBuilders(builders, str, str, str); + } + + private static void AppendToBuilders(Tuple builders, object left, object target, object right) + { + builders.Item1.Append(left); + builders.Item2.Append(target); + builders.Item3.Append(right); + } + + private static Tuple GetFormattedFilters( + string[] expressions, + CosmosElement[] orderByItems, + SortOrder[] sortOrders) + { + // 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 = expressions.Length; + bool isSingleOrderBy = numOrderByItems == 1; + StringBuilder left = new StringBuilder(); + StringBuilder target = new StringBuilder(); + StringBuilder right = new StringBuilder(); + + Tuple builders = new Tuple(left, right, target); + + 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 + string expression = expressions.First(); + SortOrder sortOrder = sortOrders.First(); + CosmosElement orderByItem = orderByItems.First(); + StringBuilder sb = new StringBuilder(); + CosmosElementToQueryLiteral cosmosElementToQueryLiteral = new CosmosElementToQueryLiteral(sb); + orderByItem.Accept(cosmosElementToQueryLiteral); + string orderByItemToString = sb.ToString(); + left.Append($"{expression} {(sortOrder == SortOrder.Descending ? "<" : ">")} {orderByItemToString}"); + target.Append($"{expression} {(sortOrder == SortOrder.Descending ? "<=" : ">=")} {orderByItemToString}"); + right.Append($"{expression} {(sortOrder == SortOrder.Descending ? "<=" : ">=")} {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++) + { + ArraySegment expressionPrefix = new ArraySegment(expressions, 0, prefixLength); + ArraySegment sortOrderPrefix = new ArraySegment(sortOrders, 0, prefixLength); + ArraySegment orderByItemsPrefix = new ArraySegment(orderByItems, 0, prefixLength); + + bool lastPrefix = prefixLength == numOrderByItems; + bool firstPrefix = prefixLength == 1; + + CosmosOrderByItemQueryExecutionContext.AppendToBuilders(builders, "("); + + for (int index = 0; index < prefixLength; index++) + { + string expression = expressionPrefix.ElementAt(index); + SortOrder sortOrder = sortOrderPrefix.ElementAt(index); + CosmosElement orderByItem = orderByItemsPrefix.ElementAt(index); + bool lastItem = index == prefixLength - 1; + + // Append Expression + CosmosOrderByItemQueryExecutionContext.AppendToBuilders(builders, expression); + CosmosOrderByItemQueryExecutionContext.AppendToBuilders(builders, " "); + + // Append binary operator + if (lastItem) + { + string inequality = sortOrder == SortOrder.Descending ? "<" : ">"; + CosmosOrderByItemQueryExecutionContext.AppendToBuilders(builders, inequality); + if (lastPrefix) + { + CosmosOrderByItemQueryExecutionContext.AppendToBuilders(builders, string.Empty, "=", "="); + } + } + else + { + CosmosOrderByItemQueryExecutionContext.AppendToBuilders(builders, "="); + } + + // 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 "); + } + } + } + + return new Tuple(left.ToString(), target.ToString(), 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; } + } + + /// + /// Struct to hold all the filters for every partition. + /// + private readonly struct FormattedFilterInfo + { + /// + /// Filters for current partition. + /// + public readonly string FiltersForTargetRange; + + /// + /// Filters for partitions left of the current partition. + /// + public readonly string FilterForRangesLeftOfTargetRanges; + + /// + /// Filters for partitions right of the current partition. + /// + public readonly string FilterForRangesRightOfTargetRanges; + + /// + /// Initializes a new instance of the FormattedFilterInfo struct. + /// + /// The filters for the partitions left of the current partition. + /// The filters for the current partition. + /// The filters for the partitions right of the current partition. + public FormattedFilterInfo(string leftFilter, string targetFilter, string rightFilters) + { + this.FilterForRangesLeftOfTargetRanges = leftFilter; + this.FiltersForTargetRange = targetFilter; + this.FilterForRangesRightOfTargetRanges = rightFilters; + } + } + } +} 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 6c534ec9bd..9c536fb41d 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/OrderBy/CosmosOrderByItemQueryExecutionContext.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/OrderBy/CosmosOrderByItemQueryExecutionContext.cs @@ -1,26 +1,14 @@ -//------------------------------------------------------------ +// ------------------------------------------------------------ // 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.Collections; - 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 PartitionKeyRange = Documents.PartitionKeyRange; - using ResourceId = Documents.ResourceId; /// /// CosmosOrderByItemQueryExecutionContext is a concrete implementation for CrossPartitionQueryExecutionContext. @@ -30,7 +18,7 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.OrderBy /// 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 CosmosOrderByItemQueryExecutionContext : CosmosCrossPartitionQueryExecutionContext + internal sealed partial class CosmosOrderByItemQueryExecutionContext { /// /// Order by queries are rewritten to allow us to inject a filter. @@ -99,930 +87,5 @@ private CosmosOrderByItemQueryExecutionContext( testSettings: testSettings) { } - - /// - /// 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 static async Task> TryCreateAsync( - 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(); - - // 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, - maxConcurrency: initParams.MaxConcurrency, - maxItemCount: initParams.MaxItemCount, - maxBufferedItemCount: initParams.MaxBufferedItemCount, - consumeComparer: orderByItemProducerTreeComparer, - testSettings: initParams.TestSettings); - - return (await context.TryInitializeAsync( - sqlQuerySpec: initParams.SqlQuerySpec, - requestContinuation: requestContinuationToken, - collectionRid: initParams.CollectionRid, - partitionKeyRanges: initParams.PartitionKeyRanges, - initialPageSize: initParams.InitialPageSize, - sortOrders: initParams.PartitionedQueryExecutionInfo.QueryInfo.OrderBy.ToArray(), - orderByExpressions: initParams.PartitionedQueryExecutionInfo.QueryInfo.OrderByExpressions.ToArray(), - cancellationToken: cancellationToken)) - .Try(() => context); - } - - /// - /// 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, - diagnostics: this.GetAndResetDiagnostics()); - } - - /// - /// 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); - } - - private async Task TryInitializeAsync( - SqlQuerySpec sqlQuerySpec, - CosmosElement requestContinuation, - string collectionRid, - List partitionKeyRanges, - int initialPageSize, - SortOrder[] sortOrders, - string[] orderByExpressions, - 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 (sortOrders == null) - { - throw new ArgumentNullException(nameof(sortOrders)); - } - - if (orderByExpressions == null) - { - throw new ArgumentNullException(nameof(orderByExpressions)); - } - - cancellationToken.ThrowIfCancellationRequested(); - - if (requestContinuation == null) - { - SqlQuerySpec sqlQuerySpecForInit = new SqlQuerySpec( - sqlQuerySpec.QueryText.Replace(oldValue: FormatPlaceHolder, newValue: True), - sqlQuerySpec.Parameters); - - TryCatch tryInitialize = await base.TryInitializeAsync( - collectionRid, - partitionKeyRanges, - initialPageSize, - sqlQuerySpecForInit, - cancellationToken: cancellationToken, - targetRangeToContinuationMap: null, - deferFirstPage: false, - filter: null, - tryFilterAsync: null); - if (!tryInitialize.Succeeded) - { - return tryInitialize; - } - } - else - { - TryCatch tryExtractContinuationTokens = CosmosOrderByItemQueryExecutionContext.TryExtractContinuationTokens( - requestContinuation, - sortOrders, - orderByExpressions); - if (!tryExtractContinuationTokens.Succeeded) - { - return TryCatch.FromException(tryExtractContinuationTokens.Exception); - } - - TryCatch tryGetOrderByInitInfo = CosmosOrderByItemQueryExecutionContext.TryGetOrderByPartitionKeyRangesInitializationInfo( - tryExtractContinuationTokens.Result, - partitionKeyRanges, - sortOrders, - orderByExpressions); - if (!tryGetOrderByInitInfo.Succeeded) - { - return TryCatch.FromException(tryGetOrderByInitInfo.Exception); - } - - OrderByInitInfo initiaizationInfo = tryGetOrderByInitInfo.Result; - RangeFilterInitializationInfo[] orderByInfos = initiaizationInfo.Filters; - IReadOnlyDictionary targetRangeToOrderByContinuationMap = initiaizationInfo.ContinuationTokens; - Debug.Assert( - targetRangeToOrderByContinuationMap != null, - "If targetRangeToOrderByContinuationMap can't be null is valid continuation is supplied"); - - // 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) - foreach (RangeFilterInitializationInfo info in orderByInfos) - { - if (info.StartIndex > info.EndIndex) - { - continue; - } - - PartialReadOnlyList partialRanges = - new PartialReadOnlyList( - partitionKeyRanges, - info.StartIndex, - info.EndIndex - info.StartIndex + 1); - - SqlQuerySpec sqlQuerySpecForInit = new SqlQuerySpec( - sqlQuerySpec.QueryText.Replace(FormatPlaceHolder, info.Filter), - sqlQuerySpec.Parameters); - - TryCatch tryInitialize = await base.TryInitializeAsync( - collectionRid, - partialRanges, - initialPageSize, - sqlQuerySpecForInit, - targetRangeToOrderByContinuationMap.ToDictionary( - kvp => kvp.Key, - kvp => kvp.Value.CompositeContinuationToken.Token), - false, - info.Filter, - async (itemProducerTree) => - { - if (targetRangeToOrderByContinuationMap.TryGetValue( - itemProducerTree.Root.PartitionKeyRange.Id, - out OrderByContinuationToken continuationToken)) - { - TryCatch tryFilter = await this.TryFilterAsync( - itemProducerTree, - sortOrders, - continuationToken, - cancellationToken); - - if (!tryFilter.Succeeded) - { - return tryFilter; - } - } - - return TryCatch.FromResult(); - }, - cancellationToken); - if (!tryInitialize.Succeeded) - { - return tryInitialize; - } - } - } - - return TryCatch.FromResult(); - } - - private static TryCatch TryExtractContinuationTokens( - CosmosElement requestContinuation, - SortOrder[] sortOrders, - string[] orderByExpressions) - { - Debug.Assert( - !(orderByExpressions == null - || orderByExpressions.Length <= 0 - || sortOrders == null - || sortOrders.Length <= 0 - || orderByExpressions.Length != sortOrders.Length), - "Partitioned QueryExecutionInfo returned bogus results."); - - if (requestContinuation == null) - { - throw new ArgumentNullException("continuation can not be null or empty."); - } - - 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 != sortOrders.Length) - { - return TryCatch.FromException( - new MalformedContinuationTokenException($"Invalid order-by items in continuation token {requestContinuation} for OrderBy~Context.")); - } - } - - return TryCatch.FromResult(orderByContinuationTokens.ToArray()); - } - - /// - /// 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, - SortOrder[] 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.Length; ++i) - { - cmp = ItemComparer.Instance.Compare( - continuationToken.OrderByItems[i].Item, - orderByResult.OrderByItems[i].Item); - - if (cmp != 0) - { - cmp = sortOrders[i] != SortOrder.Descending ? 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 first order by field. - // So if it's ORDER BY c.age ASC, c.name DESC the _rids are ASC - // If ti's ORDER BY c.age DESC, c.name DESC the _rids are DESC - cmp = continuationRid.Document.CompareTo(rid.Document); - if (sortOrders[0] == SortOrder.Descending) - { - 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(); - } - - /// - /// Gets the filters for every partition. - /// - private static TryCatch TryGetOrderByPartitionKeyRangesInitializationInfo( - OrderByContinuationToken[] suppliedContinuationTokens, - List partitionKeyRanges, - SortOrder[] sortOrders, - string[] orderByExpressions) - { - TryCatch> tryFindRangeAndContinuationTokensMonad = CosmosCrossPartitionQueryExecutionContext.TryFindTargetRangeAndExtractContinuationTokens( - partitionKeyRanges, - suppliedContinuationTokens - .Select(token => Tuple.Create(token, token.CompositeContinuationToken.Range))); - - return tryFindRangeAndContinuationTokensMonad.Try((indexAndContinuationTokens) => - { - int minIndex = indexAndContinuationTokens.TargetIndex; - IReadOnlyDictionary partitionKeyRangeToContinuationToken = indexAndContinuationTokens.ContinuationTokens; - - FormattedFilterInfo formattedFilterInfo = CosmosOrderByItemQueryExecutionContext.GetFormattedFilters( - orderByExpressions, - suppliedContinuationTokens, - sortOrders); - - RangeFilterInitializationInfo[] filters = new RangeFilterInitializationInfo[] - { - new RangeFilterInitializationInfo( - filter: formattedFilterInfo.FilterForRangesLeftOfTargetRanges, - startIndex: 0, - endIndex: minIndex - 1), - new RangeFilterInitializationInfo( - filter: formattedFilterInfo.FiltersForTargetRange, - startIndex: minIndex, - endIndex: minIndex), - new RangeFilterInitializationInfo( - filter: formattedFilterInfo.FilterForRangesRightOfTargetRanges, - startIndex: minIndex + 1, - endIndex: partitionKeyRanges.Count - 1), - }; - - return new OrderByInitInfo( - filters, - partitionKeyRangeToContinuationToken); - }); - } - - /// - /// Gets the formatted filters for every partition. - /// - /// The filter expressions. - /// The continuation token. - /// The sort orders. - /// The formatted filters for every partition. - private static FormattedFilterInfo GetFormattedFilters( - string[] expressions, - OrderByContinuationToken[] continuationTokens, - SortOrder[] sortOrders) - { - // Validate the inputs - for (int index = 0; index < continuationTokens.Length; index++) - { - Debug.Assert(continuationTokens[index].OrderByItems.Count == sortOrders.Length, "Expect values and orders are the same size."); - Debug.Assert(expressions.Length == sortOrders.Length, "Expect expressions and orders are the same size."); - } - - Tuple filters = CosmosOrderByItemQueryExecutionContext.GetFormattedFilters( - expressions, - continuationTokens[0].OrderByItems.Select(orderByItem => orderByItem.Item).ToArray(), - sortOrders); - - return new FormattedFilterInfo(filters.Item1, filters.Item2, filters.Item3); - } - - private static void AppendToBuilders(Tuple builders, object str) - { - CosmosOrderByItemQueryExecutionContext.AppendToBuilders(builders, str, str, str); - } - - private static void AppendToBuilders(Tuple builders, object left, object target, object right) - { - builders.Item1.Append(left); - builders.Item2.Append(target); - builders.Item3.Append(right); - } - - private static Tuple GetFormattedFilters( - string[] expressions, - CosmosElement[] orderByItems, - SortOrder[] sortOrders) - { - // 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 = expressions.Length; - bool isSingleOrderBy = numOrderByItems == 1; - StringBuilder left = new StringBuilder(); - StringBuilder target = new StringBuilder(); - StringBuilder right = new StringBuilder(); - - Tuple builders = new Tuple(left, right, target); - - 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 - string expression = expressions.First(); - SortOrder sortOrder = sortOrders.First(); - CosmosElement orderByItem = orderByItems.First(); - StringBuilder sb = new StringBuilder(); - CosmosElementToQueryLiteral cosmosElementToQueryLiteral = new CosmosElementToQueryLiteral(sb); - orderByItem.Accept(cosmosElementToQueryLiteral); - string orderByItemToString = sb.ToString(); - left.Append($"{expression} {(sortOrder == SortOrder.Descending ? "<" : ">")} {orderByItemToString}"); - target.Append($"{expression} {(sortOrder == SortOrder.Descending ? "<=" : ">=")} {orderByItemToString}"); - right.Append($"{expression} {(sortOrder == SortOrder.Descending ? "<=" : ">=")} {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++) - { - ArraySegment expressionPrefix = new ArraySegment(expressions, 0, prefixLength); - ArraySegment sortOrderPrefix = new ArraySegment(sortOrders, 0, prefixLength); - ArraySegment orderByItemsPrefix = new ArraySegment(orderByItems, 0, prefixLength); - - bool lastPrefix = prefixLength == numOrderByItems; - bool firstPrefix = prefixLength == 1; - - CosmosOrderByItemQueryExecutionContext.AppendToBuilders(builders, "("); - - for (int index = 0; index < prefixLength; index++) - { - string expression = expressionPrefix.ElementAt(index); - SortOrder sortOrder = sortOrderPrefix.ElementAt(index); - CosmosElement orderByItem = orderByItemsPrefix.ElementAt(index); - bool lastItem = index == prefixLength - 1; - - // Append Expression - CosmosOrderByItemQueryExecutionContext.AppendToBuilders(builders, expression); - CosmosOrderByItemQueryExecutionContext.AppendToBuilders(builders, " "); - - // Append binary operator - if (lastItem) - { - string inequality = sortOrder == SortOrder.Descending ? "<" : ">"; - CosmosOrderByItemQueryExecutionContext.AppendToBuilders(builders, inequality); - if (lastPrefix) - { - CosmosOrderByItemQueryExecutionContext.AppendToBuilders(builders, string.Empty, "=", "="); - } - } - else - { - CosmosOrderByItemQueryExecutionContext.AppendToBuilders(builders, "="); - } - - // 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 "); - } - } - } - - return new Tuple(left.ToString(), target.ToString(), right.ToString()); - } - - 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); - } - - private readonly struct OrderByInitInfo - { - public OrderByInitInfo(RangeFilterInitializationInfo[] filters, IReadOnlyDictionary continuationTokens) - { - this.Filters = filters; - this.ContinuationTokens = continuationTokens; - } - - public RangeFilterInitializationInfo[] Filters { get; } - - public IReadOnlyDictionary ContinuationTokens { get; } - } - - /// - /// Struct to hold all the filters for every partition. - /// - private readonly struct FormattedFilterInfo - { - /// - /// Filters for current partition. - /// - public readonly string FiltersForTargetRange; - - /// - /// Filters for partitions left of the current partition. - /// - public readonly string FilterForRangesLeftOfTargetRanges; - - /// - /// Filters for partitions right of the current partition. - /// - public readonly string FilterForRangesRightOfTargetRanges; - - /// - /// Initializes a new instance of the FormattedFilterInfo struct. - /// - /// The filters for the partitions left of the current partition. - /// The filters for the current partition. - /// The filters for the partitions right of the current partition. - public FormattedFilterInfo(string leftFilter, string targetFilter, string rightFilters) - { - this.FilterForRangesLeftOfTargetRanges = leftFilter; - this.FiltersForTargetRange = targetFilter; - this.FilterForRangesRightOfTargetRanges = rightFilters; - } - } - - /// - /// 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; - } - } } } From f0d2874864e1744ef181a0be429601a2ad9cb1c2 Mon Sep 17 00:00:00 2001 From: Brandon Chong Date: Tue, 3 Mar 2020 16:47:52 -0800 Subject: [PATCH 22/28] refactored code --- ...smosCrossPartitionQueryExecutionContext.cs | 219 ++++++---- ...QueryExecutionContext.ContinuationToken.cs | 7 +- ...sOrderByItemQueryExecutionContext.Drain.cs | 7 +- ...OrderByItemQueryExecutionContext.Resume.cs | 399 +++++++----------- ...arallelItemQueryExecutionContext.Resume.cs | 59 +-- 5 files changed, 341 insertions(+), 350 deletions(-) diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosCrossPartitionQueryExecutionContext.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosCrossPartitionQueryExecutionContext.cs index 490fed658f..ec9c063a70 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosCrossPartitionQueryExecutionContext.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosCrossPartitionQueryExecutionContext.cs @@ -22,6 +22,7 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionContext using Microsoft.Azure.Cosmos.Query.Core.Monads; using Microsoft.Azure.Cosmos.Query.Core.QueryClient; using Microsoft.Azure.Cosmos.Query.Core.QueryPlan; + using Microsoft.Azure.Cosmos.Sql; using Microsoft.Azure.Documents.Routing; using PartitionKeyRange = Documents.PartitionKeyRange; using RequestChargeTracker = Documents.RequestChargeTracker; @@ -333,24 +334,38 @@ protected async Task TryInitializeAsync( 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)); + } + + if (tryFilterAsync == null) + { + throw new ArgumentNullException(nameof(tryFilterAsync)); + } + cancellationToken.ThrowIfCancellationRequested(); List itemProducerTrees = new List(); - foreach (PartitionKeyRange partitionKeyRange in partitionKeyRanges) + foreach (KeyValuePair rangeAndContinuationToken in targetRangeToContinuationMap) { - string initialContinuationToken; - if (targetRangeToContinuationMap != null) - { - if (!targetRangeToContinuationMap.TryGetValue(partitionKeyRange.Id, out initialContinuationToken)) - { - initialContinuationToken = null; - } - } - else - { - initialContinuationToken = null; - } - + PartitionKeyRange partitionKeyRange = rangeAndContinuationToken.Key; + string continuationToken = rangeAndContinuationToken.Value; ItemProducerTree itemProducerTree = new ItemProducerTree( this.queryContext, querySpecForInit, @@ -362,7 +377,7 @@ protected async Task TryInitializeAsync( deferFirstPage, collectionRid, initialPageSize, - initialContinuationToken) + continuationToken) { Filter = filter }; @@ -423,6 +438,84 @@ protected async Task TryInitializeAsync( return TryCatch.FromResult(); } + protected 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 - 1); + 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)); + } + /// /// /// If a query encounters split upon resuming using continuation, we need to regenerate the continuation tokens. @@ -444,83 +537,34 @@ protected async Task TryInitializeAsync( /// The code assumes that merge doesn't happen and /// /// The index of the partition whose MinInclusive is equal to the suppliedContinuationTokens along with the continuation tokens. - protected static TryCatch> TryMatchRangesToContinuationTokens( - IReadOnlyList partitionKeyRanges, + private static IReadOnlyDictionary MatchRangesToContinuationTokens( + ReadOnlyMemory partitionKeyRanges, IReadOnlyList partitionedContinuationTokens) where PartitionedToken : IPartitionedToken { - if (partitionKeyRanges == null) - { - throw new ArgumentNullException(nameof(partitionKeyRanges)); - } - - if (partitionKeyRanges.Count() < 1) - { - throw new ArgumentException(nameof(partitionKeyRanges)); - } - if (partitionedContinuationTokens == null) { throw new ArgumentNullException(nameof(partitionedContinuationTokens)); } - if (partitionedContinuationTokens.Count() < 1) - { - throw new ArgumentException(nameof(partitionedContinuationTokens)); - } - - if (partitionedContinuationTokens.Count() > partitionKeyRanges.Count()) - { - throw new ArgumentException($"{nameof(partitionedContinuationTokens)} can not have more elements than {nameof(partitionKeyRanges)}."); - } - Dictionary partitionKeyRangeToToken = new Dictionary(); - foreach (IPartitionedToken partitionedToken in partitionedContinuationTokens) - { - // Find what ranges make up the supplied continuation token - Range range = partitionedToken.PartitionRange; - IReadOnlyList replacementRanges = partitionKeyRanges - .Where((partitionKeyRange) => - string.CompareOrdinal(range.Min, partitionKeyRange.MinInclusive) <= 0 && - string.CompareOrdinal(range.Max, partitionKeyRange.MaxExclusive) >= 0) - .OrderBy((partitionKeyRange) => partitionKeyRange.MinInclusive) - .ToList(); - - // Could not find the child ranges - if (replacementRanges.Count == 0) - { - return TryCatch>.FromException( - new MalformedContinuationTokenException( - $"{RMResources.InvalidContinuationToken} - Could not find continuation token: {partitionedToken}")); - } - - // PMax = C2Max > C2Min > C1Max > C1Min = PMin. - string parentMax = range.Max; - string child2Max = replacementRanges.Last().MaxExclusive; - string child2Min = replacementRanges.Last().MinInclusive; - string child1Max = replacementRanges.First().MaxExclusive; - string child1Min = replacementRanges.First().MinInclusive; - string parentMin = range.Min; - - if (!(parentMax == child2Max && - string.CompareOrdinal(child2Max, child2Min) >= 0 && - (replacementRanges.Count() == 1 ? true : string.CompareOrdinal(child2Min, child1Max) >= 0) && - string.CompareOrdinal(child1Max, child1Min) >= 0 && - child1Min == parentMin)) - { - return TryCatch>.FromException( - new MalformedContinuationTokenException( - $"{RMResources.InvalidContinuationToken} - PMax = C2Max > C2Min > C1Max > C1Min = PMin: {partitionedToken}")); - } - - foreach (PartitionKeyRange partitionKeyRange in replacementRanges) + ReadOnlySpan partitionKeyRangeSpan = partitionKeyRanges.Span; + for (int i = 0; i < partitionKeyRangeSpan.Length; i++) + { + PartitionKeyRange partitionKeyRange = partitionKeyRangeSpan[i]; + foreach (PartitionedToken partitionedToken in partitionedContinuationTokens) { - partitionKeyRangeToToken[partitionKeyRange] = partitionedToken; + // 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; + } } } - return TryCatch>.FromResult( - (IReadOnlyDictionary)partitionKeyRangeToToken); + return (IReadOnlyDictionary)partitionKeyRangeToToken; } protected virtual long GetAndResetResponseLengthBytes() @@ -627,6 +671,23 @@ public InitInfo(int targetIndex, IReadOnlyDictionary 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.ContinuationToken.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/OrderBy/CosmosOrderByItemQueryExecutionContext.ContinuationToken.cs index dd4ddd4c19..c0df108ecd 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/OrderBy/CosmosOrderByItemQueryExecutionContext.ContinuationToken.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/OrderBy/CosmosOrderByItemQueryExecutionContext.ContinuationToken.cs @@ -1,5 +1,10 @@ -namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.OrderBy +// ------------------------------------------------------------ +// 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; 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 index b5da8ae6b5..b0c3a796b9 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/OrderBy/CosmosOrderByItemQueryExecutionContext.Drain.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/OrderBy/CosmosOrderByItemQueryExecutionContext.Drain.cs @@ -1,5 +1,10 @@ -namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.OrderBy +// ------------------------------------------------------------ +// 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; 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 c5e0c50421..47cd90b276 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 @@ -12,7 +12,6 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.OrderBy using System.Threading.Tasks; using Microsoft.Azure.Cosmos.CosmosElements; using Microsoft.Azure.Cosmos.Query.Core; - 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.ExecutionComponent; @@ -53,14 +52,24 @@ public static async Task> TryCreateAs 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(); + return (await context.TryInitializeAsync( sqlQuerySpec: initParams.SqlQuerySpec, requestContinuation: requestContinuationToken, collectionRid: initParams.CollectionRid, partitionKeyRanges: initParams.PartitionKeyRanges, initialPageSize: initParams.InitialPageSize, - sortOrders: initParams.PartitionedQueryExecutionInfo.QueryInfo.OrderBy.ToArray(), - orderByExpressions: initParams.PartitionedQueryExecutionInfo.QueryInfo.OrderByExpressions.ToArray(), + orderByColumns: columns, cancellationToken: cancellationToken)) .Try(() => context); } @@ -69,10 +78,9 @@ private async Task TryInitializeAsync( SqlQuerySpec sqlQuerySpec, CosmosElement requestContinuation, string collectionRid, - List partitionKeyRanges, + IReadOnlyList partitionKeyRanges, int initialPageSize, - SortOrder[] sortOrders, - string[] orderByExpressions, + IReadOnlyList orderByColumns, CancellationToken cancellationToken) { if (sqlQuerySpec == null) @@ -90,150 +98,149 @@ private async Task TryInitializeAsync( throw new ArgumentNullException(nameof(partitionKeyRanges)); } - if (sortOrders == null) + if (orderByColumns == null) { - throw new ArgumentNullException(nameof(sortOrders)); + throw new ArgumentNullException(nameof(orderByColumns)); } - if (orderByExpressions == null) + cancellationToken.ThrowIfCancellationRequested(); + + TryCatch> tryGetOrderByContinuationTokenMapping = TryGetOrderByContinuationTokenMapping( + partitionKeyRanges, + requestContinuation, + orderByColumns.Count()); + if (!tryGetOrderByContinuationTokenMapping.Succeeded) { - throw new ArgumentNullException(nameof(orderByExpressions)); + return TryCatch.FromException(tryGetOrderByContinuationTokenMapping.Exception); } - cancellationToken.ThrowIfCancellationRequested(); + SqlQuerySpec rewrittenQueryForOrderBy = new SqlQuerySpec( + sqlQuerySpec.QueryText.Replace(oldValue: FormatPlaceHolder, newValue: True), + sqlQuerySpec.Parameters); + + 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.")); + } - if (requestContinuation == null) + 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)>() { - SqlQuerySpec sqlQuerySpecForInit = new SqlQuerySpec( - sqlQuerySpec.QueryText.Replace(oldValue: FormatPlaceHolder, newValue: True), - sqlQuerySpec.Parameters); + { (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) + { TryCatch tryInitialize = await base.TryInitializeAsync( collectionRid, - partitionKeyRanges, initialPageSize, - sqlQuerySpecForInit, - cancellationToken: cancellationToken, - targetRangeToContinuationMap: null, + rewrittenQueryForOrderBy, + tokenMapping.ToDictionary(kvp => kvp.Key, kvp => kvp.Value?.CompositeContinuationToken.Token), deferFirstPage: false, - filter: null, - tryFilterAsync: null); + filter, + tryFilterAsync: async (itemProducerTree) => + { + if (!tokenMapping.TryGetValue( + itemProducerTree.PartitionKeyRange, + out OrderByContinuationToken continuationToken)) + { + throw new InvalidOperationException($"Failed to retrieve {nameof(OrderByContinuationToken)}."); + } + + return await this.TryFilterAsync( + itemProducerTree, + sortOrders, + continuationToken, + cancellationToken); + }, + cancellationToken); if (!tryInitialize.Succeeded) { return tryInitialize; } } - else + + return TryCatch.FromResult(); + } + + private static TryCatch> TryGetOrderByContinuationTokenMapping( + IReadOnlyList partitionKeyRanges, + CosmosElement continuationToken, + int numOrderByItems) + { + if (partitionKeyRanges == null) { - TryCatch tryExtractContinuationTokens = CosmosOrderByItemQueryExecutionContext.TryExtractContinuationTokens( - requestContinuation, - sortOrders, - orderByExpressions); - if (!tryExtractContinuationTokens.Succeeded) - { - return TryCatch.FromException(tryExtractContinuationTokens.Exception); - } + throw new ArgumentOutOfRangeException(nameof(partitionKeyRanges)); + } - TryCatch tryGetOrderByInitInfo = CosmosOrderByItemQueryExecutionContext.TryGetOrderByPartitionKeyRangesInitializationInfo( - tryExtractContinuationTokens.Result, - partitionKeyRanges, - sortOrders, - orderByExpressions); - if (!tryGetOrderByInitInfo.Succeeded) - { - return TryCatch.FromException(tryGetOrderByInitInfo.Exception); - } + if (numOrderByItems < 0) + { + throw new ArgumentOutOfRangeException(nameof(numOrderByItems)); + } - OrderByInitInfo initiaizationInfo = tryGetOrderByInitInfo.Result; - RangeFilterInitializationInfo[] orderByInfos = initiaizationInfo.Filters; - IReadOnlyDictionary targetRangeToOrderByContinuationMap = initiaizationInfo.ContinuationTokens; - Debug.Assert( - targetRangeToOrderByContinuationMap != null, - "If targetRangeToOrderByContinuationMap can't be null is valid continuation is supplied"); - - // 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) - foreach (RangeFilterInitializationInfo info in orderByInfos) + if (continuationToken == null) + { + // Create a mapping for all the ranges being right of the target partition (since they are let to be consumed). + Dictionary dictionary = new Dictionary(); + foreach (PartitionKeyRange partitionKeyRange in partitionKeyRanges) { - if (info.StartIndex > info.EndIndex) - { - continue; - } - - PartialReadOnlyList partialRanges = - new PartialReadOnlyList( - partitionKeyRanges, - info.StartIndex, - info.EndIndex - info.StartIndex + 1); - - SqlQuerySpec sqlQuerySpecForInit = new SqlQuerySpec( - sqlQuerySpec.QueryText.Replace(FormatPlaceHolder, info.Filter), - sqlQuerySpec.Parameters); - - TryCatch tryInitialize = await base.TryInitializeAsync( - collectionRid, - partialRanges, - initialPageSize, - sqlQuerySpecForInit, - targetRangeToOrderByContinuationMap.ToDictionary( - kvp => kvp.Key, - kvp => kvp.Value.CompositeContinuationToken.Token), - false, - info.Filter, - async (itemProducerTree) => - { - if (targetRangeToOrderByContinuationMap.TryGetValue( - itemProducerTree.Root.PartitionKeyRange.Id, - out OrderByContinuationToken continuationToken)) - { - TryCatch tryFilter = await this.TryFilterAsync( - itemProducerTree, - sortOrders, - continuationToken, - cancellationToken); + dictionary.Add(key: partitionKeyRange, value: null); + } - if (!tryFilter.Succeeded) - { - return tryFilter; - } - } + return TryCatch>.FromResult( + new PartitionMapping( + partitionsLeftOfTarget: new Dictionary(), + targetPartition: new Dictionary(), + partitionsRightOfTarget: dictionary)); + } - return TryCatch.FromResult(); - }, - cancellationToken); - if (!tryInitialize.Succeeded) - { - return tryInitialize; - } - } + TryCatch> tryExtractContinuationTokens = TryExtractContinuationTokens(continuationToken, numOrderByItems); + if (!tryExtractContinuationTokens.Succeeded) + { + return TryCatch>.FromException(tryExtractContinuationTokens.Exception); } - return TryCatch.FromResult(); + return TryGetInitializationInfo( + partitionKeyRanges, + tryExtractContinuationTokens.Result); } - private static TryCatch TryExtractContinuationTokens( + private static TryCatch> TryExtractContinuationTokens( CosmosElement requestContinuation, - SortOrder[] sortOrders, - string[] orderByExpressions) + int numOrderByItems) { - Debug.Assert( - !(orderByExpressions == null - || orderByExpressions.Length <= 0 - || sortOrders == null - || sortOrders.Length <= 0 - || orderByExpressions.Length != sortOrders.Length), - "Partitioned QueryExecutionInfo returned bogus results."); - 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}.")); + return TryCatch>.FromException( + new MalformedContinuationTokenException( + $"Order by continuation token must be an array: {requestContinuation}.")); } List orderByContinuationTokens = new List(); @@ -242,7 +249,7 @@ private static TryCatch TryExtractContinuationTokens TryCatch tryCreateOrderByContinuationToken = OrderByContinuationToken.TryCreateFromCosmosElement(arrayItem); if (!tryCreateOrderByContinuationToken.Succeeded) { - return TryCatch.FromException(tryCreateOrderByContinuationToken.Exception); + return TryCatch>.FromException(tryCreateOrderByContinuationToken.Exception); } orderByContinuationTokens.Add(tryCreateOrderByContinuationToken.Result); @@ -250,20 +257,22 @@ private static TryCatch TryExtractContinuationTokens if (orderByContinuationTokens.Count == 0) { - return TryCatch.FromException( - new MalformedContinuationTokenException($"Order by continuation token cannot be empty: {requestContinuation}.")); + return TryCatch>.FromException( + new MalformedContinuationTokenException( + $"Order by continuation token cannot be empty: {requestContinuation}.")); } foreach (OrderByContinuationToken suppliedOrderByContinuationToken in orderByContinuationTokens) { - if (suppliedOrderByContinuationToken.OrderByItems.Count != sortOrders.Length) + if (suppliedOrderByContinuationToken.OrderByItems.Count != numOrderByItems) { - return TryCatch.FromException( - new MalformedContinuationTokenException($"Invalid order-by items in continuation token {requestContinuation} for OrderBy~Context.")); + return TryCatch>.FromException( + new MalformedContinuationTokenException( + $"Invalid order-by items in continuation token {requestContinuation} for OrderBy~Context.")); } } - return TryCatch.FromResult(orderByContinuationTokens.ToArray()); + return TryCatch>.FromResult(orderByContinuationTokens); } /// @@ -276,7 +285,7 @@ private static TryCatch TryExtractContinuationTokens /// A task to await on. private async Task TryFilterAsync( ItemProducerTree producer, - SortOrder[] sortOrders, + IReadOnlyList sortOrders, OrderByContinuationToken continuationToken, CancellationToken cancellationToken) { @@ -293,7 +302,8 @@ private async Task TryFilterAsync( 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.")); + new MalformedContinuationTokenException( + $"Invalid Rid in the continuation token {continuationToken.CompositeContinuationToken.Token} for OrderBy~Context.")); } Dictionary resourceIds = new Dictionary(); @@ -311,7 +321,7 @@ private async Task TryFilterAsync( 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.Length; ++i) + for (int i = 0; i < sortOrders.Count(); ++i) { cmp = ItemComparer.Instance.Compare( continuationToken.OrderByItems[i].Item, @@ -408,95 +418,20 @@ private async Task TryFilterAsync( return TryCatch.FromResult(); } - /// - /// Gets the filters for every partition. - /// - private static TryCatch TryGetOrderByPartitionKeyRangesInitializationInfo( - OrderByContinuationToken[] suppliedContinuationTokens, - List partitionKeyRanges, - SortOrder[] sortOrders, - string[] orderByExpressions) - { - TryCatch> tryFindRangeAndContinuationTokensMonad = CosmosCrossPartitionQueryExecutionContext.TryFindTargetRangeAndExtractContinuationTokens( - partitionKeyRanges, - suppliedContinuationTokens - .Select(token => Tuple.Create(token, token.CompositeContinuationToken.Range))); - - return tryFindRangeAndContinuationTokensMonad.Try((indexAndContinuationTokens) => - { - int minIndex = indexAndContinuationTokens.TargetIndex; - IReadOnlyDictionary partitionKeyRangeToContinuationToken = indexAndContinuationTokens.ContinuationTokens; - - FormattedFilterInfo formattedFilterInfo = CosmosOrderByItemQueryExecutionContext.GetFormattedFilters( - orderByExpressions, - suppliedContinuationTokens, - sortOrders); - - RangeFilterInitializationInfo[] filters = new RangeFilterInitializationInfo[] - { - new RangeFilterInitializationInfo( - filter: formattedFilterInfo.FilterForRangesLeftOfTargetRanges, - startIndex: 0, - endIndex: minIndex - 1), - new RangeFilterInitializationInfo( - filter: formattedFilterInfo.FiltersForTargetRange, - startIndex: minIndex, - endIndex: minIndex), - new RangeFilterInitializationInfo( - filter: formattedFilterInfo.FilterForRangesRightOfTargetRanges, - startIndex: minIndex + 1, - endIndex: partitionKeyRanges.Count - 1), - }; - - return new OrderByInitInfo( - filters, - partitionKeyRangeToContinuationToken); - }); - } - - /// - /// Gets the formatted filters for every partition. - /// - /// The filter expressions. - /// The continuation token. - /// The sort orders. - /// The formatted filters for every partition. - private static FormattedFilterInfo GetFormattedFilters( - string[] expressions, - OrderByContinuationToken[] continuationTokens, - SortOrder[] sortOrders) - { - // Validate the inputs - for (int index = 0; index < continuationTokens.Length; index++) - { - Debug.Assert(continuationTokens[index].OrderByItems.Count == sortOrders.Length, "Expect values and orders are the same size."); - Debug.Assert(expressions.Length == sortOrders.Length, "Expect expressions and orders are the same size."); - } - - Tuple filters = CosmosOrderByItemQueryExecutionContext.GetFormattedFilters( - expressions, - continuationTokens[0].OrderByItems.Select(orderByItem => orderByItem.Item).ToArray(), - sortOrders); - - return new FormattedFilterInfo(filters.Item1, filters.Item2, filters.Item3); - } - - private static void AppendToBuilders(Tuple builders, object str) + private static void AppendToBuilders((StringBuilder, StringBuilder, StringBuilder) builders, object str) { CosmosOrderByItemQueryExecutionContext.AppendToBuilders(builders, str, str, str); } - private static void AppendToBuilders(Tuple builders, object left, object target, object right) + private static void AppendToBuilders((StringBuilder, StringBuilder, StringBuilder) builders, object left, object target, object right) { builders.Item1.Append(left); builders.Item2.Append(target); builders.Item3.Append(right); } - private static Tuple GetFormattedFilters( - string[] expressions, - CosmosElement[] orderByItems, - SortOrder[] sortOrders) + 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. @@ -506,13 +441,13 @@ private static Tuple GetFormattedFilters( // 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 = expressions.Length; + int numOrderByItems = columnAndItems.Length; bool isSingleOrderBy = numOrderByItems == 1; StringBuilder left = new StringBuilder(); StringBuilder target = new StringBuilder(); StringBuilder right = new StringBuilder(); - Tuple builders = new Tuple(left, right, target); + (StringBuilder, StringBuilder, StringBuilder) builders = (left, right, target); if (isSingleOrderBy) { @@ -536,13 +471,15 @@ private static Tuple GetFormattedFilters( // < for partitions to the left // <= for the partition we left off on // <= for the partitions to the right - string expression = expressions.First(); - SortOrder sortOrder = sortOrders.First(); - CosmosElement orderByItem = orderByItems.First(); + (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 ? "<" : ">")} {orderByItemToString}"); target.Append($"{expression} {(sortOrder == SortOrder.Descending ? "<=" : ">=")} {orderByItemToString}"); right.Append($"{expression} {(sortOrder == SortOrder.Descending ? "<=" : ">=")} {orderByItemToString}"); @@ -580,9 +517,7 @@ private static Tuple GetFormattedFilters( // And then you just OR all the possibilities together for (int prefixLength = 1; prefixLength <= numOrderByItems; prefixLength++) { - ArraySegment expressionPrefix = new ArraySegment(expressions, 0, prefixLength); - ArraySegment sortOrderPrefix = new ArraySegment(sortOrders, 0, prefixLength); - ArraySegment orderByItemsPrefix = new ArraySegment(orderByItems, 0, prefixLength); + ReadOnlySpan<(OrderByColumn orderByColumn, CosmosElement orderByItem)> columnAndItemPrefix = columnAndItems.Span.Slice(start: 0, length: prefixLength); bool lastPrefix = prefixLength == numOrderByItems; bool firstPrefix = prefixLength == 1; @@ -591,9 +526,9 @@ private static Tuple GetFormattedFilters( for (int index = 0; index < prefixLength; index++) { - string expression = expressionPrefix.ElementAt(index); - SortOrder sortOrder = sortOrderPrefix.ElementAt(index); - CosmosElement orderByItem = orderByItemsPrefix.ElementAt(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 @@ -638,7 +573,7 @@ private static Tuple GetFormattedFilters( } } - return new Tuple(left.ToString(), target.ToString(), right.ToString()); + return (left.ToString(), target.ToString(), right.ToString()); } private readonly struct OrderByInitInfo @@ -654,38 +589,16 @@ public OrderByInitInfo(RangeFilterInitializationInfo[] filters, IReadOnlyDiction public IReadOnlyDictionary ContinuationTokens { get; } } - /// - /// Struct to hold all the filters for every partition. - /// - private readonly struct FormattedFilterInfo + private readonly struct OrderByColumn { - /// - /// Filters for current partition. - /// - public readonly string FiltersForTargetRange; - - /// - /// Filters for partitions left of the current partition. - /// - public readonly string FilterForRangesLeftOfTargetRanges; - - /// - /// Filters for partitions right of the current partition. - /// - public readonly string FilterForRangesRightOfTargetRanges; - - /// - /// Initializes a new instance of the FormattedFilterInfo struct. - /// - /// The filters for the partitions left of the current partition. - /// The filters for the current partition. - /// The filters for the partitions right of the current partition. - public FormattedFilterInfo(string leftFilter, string targetFilter, string rightFilters) + public OrderByColumn(string expression, SortOrder sortOrder) { - this.FilterForRangesLeftOfTargetRanges = leftFilter; - this.FiltersForTargetRange = targetFilter; - this.FilterForRangesRightOfTargetRanges = rightFilters; + 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/Parallel/CosmosParallelItemQueryExecutionContext.Resume.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/Parallel/CosmosParallelItemQueryExecutionContext.Resume.cs index 54d4a47fe5..e4a3577f8c 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 @@ -81,7 +81,7 @@ private async Task> TryInitial { cancellationToken.ThrowIfCancellationRequested(); - TryCatch> tryGetPartitionKeyRangeToCompositeContinuationToken = CosmosParallelItemQueryExecutionContext.TryGetPartitionKeyRangeToCompositeContinuationToken( + TryCatch> tryGetPartitionKeyRangeToCompositeContinuationToken = CosmosParallelItemQueryExecutionContext.TryGetPartitionKeyRangeToCompositeContinuationToken( partitionKeyRanges, requestContinuation); if (!tryGetPartitionKeyRangeToCompositeContinuationToken.Succeeded) @@ -89,57 +89,64 @@ private async Task> TryInitial return TryCatch.FromException(tryGetPartitionKeyRangeToCompositeContinuationToken.Exception); } - IReadOnlyDictionary partitionKeyRangeToContinuationToken = tryGetPartitionKeyRangeToCompositeContinuationToken - .Result - .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) + List> rangesToInitialize = new List>() { - return TryCatch.FromException(tryInitialize.Exception); + // 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( + private static TryCatch> TryGetPartitionKeyRangeToCompositeContinuationToken( IReadOnlyList partitionKeyRanges, CosmosElement continuationToken) { if (continuationToken == null) { + // Create a mapping for all the ranges being right of the target partition (since they are let to be consumed). Dictionary dictionary = new Dictionary(); foreach (PartitionKeyRange partitionKeyRange in partitionKeyRanges) { dictionary.Add(key: partitionKeyRange, value: null); } - return TryCatch>.FromResult(dictionary); + return TryCatch>.FromResult( + new PartitionMapping( + partitionsLeftOfTarget: new Dictionary(), + targetPartition: new Dictionary(), + partitionsRightOfTarget: dictionary)); } TryCatch> tryParseCompositeContinuationTokens = TryParseCompositeContinuationList(continuationToken); if (!tryParseCompositeContinuationTokens.Succeeded) { - return TryCatch>.FromException(tryParseCompositeContinuationTokens.Exception); + return TryCatch>.FromException(tryParseCompositeContinuationTokens.Exception); } - TryCatch> tryMatchContinuationTokensToRanges = CosmosCrossPartitionQueryExecutionContext.TryMatchRangesToContinuationTokens( + return CosmosCrossPartitionQueryExecutionContext.TryGetInitializationInfo( partitionKeyRanges, tryParseCompositeContinuationTokens.Result); - if (!tryMatchContinuationTokensToRanges.Succeeded) - { - return TryCatch>.FromException( - tryMatchContinuationTokensToRanges.Exception); - } - - return tryMatchContinuationTokensToRanges; } private static TryCatch> TryParseCompositeContinuationList( From c52f60d888e57b2e53936cdf8d4a702a0527250a Mon Sep 17 00:00:00 2001 From: Brandon Chong Date: Wed, 4 Mar 2020 08:43:43 -0800 Subject: [PATCH 23/28] fixed bugs --- .../CompositeContinuationToken.cs | 1 + .../OrderByContinuationToken.cs | 1 + ...smosCrossPartitionQueryExecutionContext.cs | 15 +++--- ...OrderByItemQueryExecutionContext.Resume.cs | 50 +++++++++++++------ ...arallelItemQueryExecutionContext.Resume.cs | 5 +- .../CrossPartitionQueryTests.cs | 18 +++++-- 6 files changed, 61 insertions(+), 29 deletions(-) diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/CompositeContinuationToken.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/CompositeContinuationToken.cs index 1f1c16f24f..c8b9474562 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/CompositeContinuationToken.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/CompositeContinuationToken.cs @@ -42,6 +42,7 @@ public Documents.Routing.Range Range set; } + [JsonIgnore] public Range PartitionRange => this.Range; public object ShallowCopy() diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/OrderByContinuationToken.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/OrderByContinuationToken.cs index 720a3a6c13..74cb8fc3f6 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/OrderByContinuationToken.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ContinuationTokens/OrderByContinuationToken.cs @@ -207,6 +207,7 @@ public string Filter get; } + [JsonIgnore] public Range PartitionRange => this.CompositeContinuationToken.Range; public static CosmosElement ToCosmosElement(OrderByContinuationToken orderByContinuationToken) diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosCrossPartitionQueryExecutionContext.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosCrossPartitionQueryExecutionContext.cs index ec9c063a70..3c66dc1f13 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosCrossPartitionQueryExecutionContext.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosCrossPartitionQueryExecutionContext.cs @@ -354,11 +354,6 @@ protected async Task TryInitializeAsync( throw new ArgumentNullException(nameof(targetRangeToContinuationMap)); } - if (tryFilterAsync == null) - { - throw new ArgumentNullException(nameof(tryFilterAsync)); - } - cancellationToken.ThrowIfCancellationRequested(); List itemProducerTrees = new List(); @@ -547,7 +542,7 @@ private static IReadOnlyDictionary MatchRan throw new ArgumentNullException(nameof(partitionedContinuationTokens)); } - Dictionary partitionKeyRangeToToken = new Dictionary(); + Dictionary partitionKeyRangeToToken = new Dictionary(); ReadOnlySpan partitionKeyRangeSpan = partitionKeyRanges.Span; for (int i = 0; i < partitionKeyRangeSpan.Length; i++) { @@ -562,9 +557,15 @@ private static IReadOnlyDictionary MatchRan break; } } + + if (!partitionKeyRangeToToken.ContainsKey(partitionKeyRange)) + { + // Could not find a matching token so just set it to null + partitionKeyRangeToToken[partitionKeyRange] = default; + } } - return (IReadOnlyDictionary)partitionKeyRangeToToken; + return partitionKeyRangeToToken; } protected virtual long GetAndResetResponseLengthBytes() 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 47cd90b276..56076aab7b 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 @@ -11,6 +11,7 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.OrderBy using System.Threading; using System.Threading.Tasks; using Microsoft.Azure.Cosmos.CosmosElements; + using Microsoft.Azure.Cosmos.Linq; using Microsoft.Azure.Cosmos.Query.Core; using Microsoft.Azure.Cosmos.Query.Core.ContinuationTokens; using Microsoft.Azure.Cosmos.Query.Core.Exceptions; @@ -105,6 +106,29 @@ private async Task TryInitializeAsync( 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, @@ -114,10 +138,6 @@ private async Task TryInitializeAsync( return TryCatch.FromException(tryGetOrderByContinuationTokenMapping.Exception); } - SqlQuerySpec rewrittenQueryForOrderBy = new SqlQuerySpec( - sqlQuerySpec.QueryText.Replace(oldValue: FormatPlaceHolder, newValue: True), - sqlQuerySpec.Parameters); - IReadOnlyList orderByItems = tryGetOrderByContinuationTokenMapping .Result .TargetPartition @@ -148,6 +168,10 @@ private async Task TryInitializeAsync( 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, @@ -164,6 +188,11 @@ private async Task TryInitializeAsync( throw new InvalidOperationException($"Failed to retrieve {nameof(OrderByContinuationToken)}."); } + if (continuationToken == null) + { + return TryCatch.FromResult(); + } + return await this.TryFilterAsync( itemProducerTree, sortOrders, @@ -197,18 +226,7 @@ private static TryCatch> TryGetOrderB if (continuationToken == null) { - // Create a mapping for all the ranges being right of the target partition (since they are let to be consumed). - Dictionary dictionary = new Dictionary(); - foreach (PartitionKeyRange partitionKeyRange in partitionKeyRanges) - { - dictionary.Add(key: partitionKeyRange, value: null); - } - - return TryCatch>.FromResult( - new PartitionMapping( - partitionsLeftOfTarget: new Dictionary(), - targetPartition: new Dictionary(), - partitionsRightOfTarget: dictionary)); + throw new ArgumentNullException(nameof(continuationToken)); } TryCatch> tryExtractContinuationTokens = TryExtractContinuationTokens(continuationToken, numOrderByItems); 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 e4a3577f8c..0ad350287f 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 @@ -124,7 +124,6 @@ private static TryCatch> TryGetPart { if (continuationToken == null) { - // Create a mapping for all the ranges being right of the target partition (since they are let to be consumed). Dictionary dictionary = new Dictionary(); foreach (PartitionKeyRange partitionKeyRange in partitionKeyRanges) { @@ -134,8 +133,8 @@ private static TryCatch> TryGetPart return TryCatch>.FromResult( new PartitionMapping( partitionsLeftOfTarget: new Dictionary(), - targetPartition: new Dictionary(), - partitionsRightOfTarget: dictionary)); + targetPartition: dictionary, + partitionsRightOfTarget: new Dictionary())); } TryCatch> tryParseCompositeContinuationTokens = TryParseCompositeContinuationList(continuationToken); diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CrossPartitionQueryTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CrossPartitionQueryTests.cs index 195947d506..8b1f9aeb70 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CrossPartitionQueryTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CrossPartitionQueryTests.cs @@ -1315,11 +1315,23 @@ private async Task TestBasicCrossPartitionQueryHelper( Container container, IEnumerable documents) { - foreach (int maxDegreeOfParallelism in new int[] { 1, 100 }) + foreach (int maxDegreeOfParallelism in new int[] { - foreach (int maxItemCount in new int[] { 10, 100 }) + 1, + 100, + }) + { + foreach (int maxItemCount in new int[] { - foreach (string query in new string[] { "SELECT c.id FROM c", "SELECT c._ts, c.id FROM c ORDER BY c._ts" }) + 10, + 100, + }) + { + foreach (string query in new string[] + { + "SELECT c.id FROM c", + "SELECT c._ts, c.id FROM c ORDER BY c._ts", + }) { QueryRequestOptions feedOptions = new QueryRequestOptions { From c894ddc2f2415f36d9a36d0ce734b6d40fa31a05 Mon Sep 17 00:00:00 2001 From: Brandon Chong Date: Thu, 5 Mar 2020 16:18:43 -0800 Subject: [PATCH 24/28] resolved iteration comments --- ...smosCrossPartitionQueryExecutionContext.cs | 32 +- ...OrderByItemQueryExecutionContext.Resume.cs | 36 +- .../CosmosOrderByItemQueryExecutionContext.cs | 892 +----------------- 3 files changed, 34 insertions(+), 926 deletions(-) diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosCrossPartitionQueryExecutionContext.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosCrossPartitionQueryExecutionContext.cs index efadf464bc..fee65379e8 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosCrossPartitionQueryExecutionContext.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosCrossPartitionQueryExecutionContext.cs @@ -449,17 +449,17 @@ protected static TryCatch> TryGetInitializati throw new ArgumentNullException(nameof(partitionedContinuationTokens)); } - if (partitionKeyRanges.Count() < 1) + if (partitionKeyRanges.Count < 1) { throw new ArgumentException(nameof(partitionKeyRanges)); } - if (partitionedContinuationTokens.Count() < 1) + if (partitionedContinuationTokens.Count < 1) { throw new ArgumentException(nameof(partitionKeyRanges)); } - if (partitionedContinuationTokens.Count() > partitionKeyRanges.Count()) + if (partitionedContinuationTokens.Count > partitionKeyRanges.Count) { throw new ArgumentException($"{nameof(partitionedContinuationTokens)} can not have more elements than {nameof(partitionKeyRanges)}."); } @@ -513,26 +513,14 @@ protected static TryCatch> TryGetInitializati } /// - /// - /// If a query encounters split upon resuming using continuation, we need to regenerate the continuation tokens. - /// Specifically, since after split we will have new set of ranges, we need to remove continuation token for the - /// parent partition and introduce continuation token for the child partitions. - /// - /// - /// This function does that. Also in that process, we also check validity of the input continuation tokens. For example, - /// even after split the boundary ranges of the child partitions should match with the parent partitions. If the Min and Max - /// range of a target partition in the continuation token was Min1 and Max1. Then the Min and Max range info for the two - /// corresponding child partitions C1Min, C1Max, C2Min, and C2Max should follow the constrain below: - /// PMax = C2Max > C2Min > C1Max > C1Min = PMin. - /// + /// 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 partition key ranges to extract continuation tokens for. - /// The continuation token that the user supplied. - /// The type of continuation token to generate. - /// - /// The code assumes that merge doesn't happen and - /// - /// The index of the partition whose MinInclusive is equal to the suppliedContinuationTokens along with the continuation tokens. + /// 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. private static IReadOnlyDictionary MatchRangesToContinuationTokens( ReadOnlyMemory partitionKeyRanges, IReadOnlyList partitionedContinuationTokens) 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 56076aab7b..c39de51b99 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 @@ -11,7 +11,6 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.OrderBy using System.Threading; using System.Threading.Tasks; using Microsoft.Azure.Cosmos.CosmosElements; - using Microsoft.Azure.Cosmos.Linq; using Microsoft.Azure.Cosmos.Query.Core; using Microsoft.Azure.Cosmos.Query.Core.ContinuationTokens; using Microsoft.Azure.Cosmos.Query.Core.Exceptions; @@ -24,6 +23,15 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.OrderBy internal sealed partial class CosmosOrderByItemQueryExecutionContext : CosmosCrossPartitionQueryExecutionContext { + 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> TryCreateAsync( CosmosQueryContext queryContext, CosmosCrossPartitionQueryExecutionContext.CrossPartitionInitParams initParams, @@ -436,16 +444,16 @@ private async Task TryFilterAsync( return TryCatch.FromResult(); } - private static void AppendToBuilders((StringBuilder, StringBuilder, StringBuilder) builders, object str) + private static void AppendToBuilders((StringBuilder leftFilter, StringBuilder targetFilter, StringBuilder rightFilter) builders, object str) { CosmosOrderByItemQueryExecutionContext.AppendToBuilders(builders, str, str, str); } - private static void AppendToBuilders((StringBuilder, StringBuilder, StringBuilder) builders, object left, object target, object right) + private static void AppendToBuilders((StringBuilder leftFilter, StringBuilder targetFilter, StringBuilder rightFilter) builders, object left, object target, object right) { - builders.Item1.Append(left); - builders.Item2.Append(target); - builders.Item3.Append(right); + builders.leftFilter.Append(left); + builders.targetFilter.Append(target); + builders.rightFilter.Append(right); } private static (string leftFilter, string targetFilter, string rightFilter) GetFormattedFilters( @@ -498,9 +506,9 @@ private static (string leftFilter, string targetFilter, string rightFilter) GetF string orderByItemToString = sb.ToString(); - left.Append($"{expression} {(sortOrder == SortOrder.Descending ? "<" : ">")} {orderByItemToString}"); - target.Append($"{expression} {(sortOrder == SortOrder.Descending ? "<=" : ">=")} {orderByItemToString}"); - right.Append($"{expression} {(sortOrder == SortOrder.Descending ? "<=" : ">=")} {orderByItemToString}"); + 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 { @@ -556,16 +564,16 @@ private static (string leftFilter, string targetFilter, string rightFilter) GetF // Append binary operator if (lastItem) { - string inequality = sortOrder == SortOrder.Descending ? "<" : ">"; + string inequality = sortOrder == SortOrder.Descending ? Expressions.LessThan : Expressions.GreaterThan; CosmosOrderByItemQueryExecutionContext.AppendToBuilders(builders, inequality); if (lastPrefix) { - CosmosOrderByItemQueryExecutionContext.AppendToBuilders(builders, string.Empty, "=", "="); + CosmosOrderByItemQueryExecutionContext.AppendToBuilders(builders, string.Empty, Expressions.EqualTo, Expressions.EqualTo); } } else { - CosmosOrderByItemQueryExecutionContext.AppendToBuilders(builders, "="); + CosmosOrderByItemQueryExecutionContext.AppendToBuilders(builders, Expressions.EqualTo); } // Append SortOrder @@ -596,7 +604,9 @@ private static (string leftFilter, string targetFilter, string rightFilter) GetF private readonly struct OrderByInitInfo { - public OrderByInitInfo(RangeFilterInitializationInfo[] filters, IReadOnlyDictionary continuationTokens) + public OrderByInitInfo( + RangeFilterInitializationInfo[] filters, + IReadOnlyDictionary continuationTokens) { this.Filters = filters; this.ContinuationTokens = continuationTokens; 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 f40f51f27d..8debd89d5a 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/OrderBy/CosmosOrderByItemQueryExecutionContext.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/OrderBy/CosmosOrderByItemQueryExecutionContext.cs @@ -6,7 +6,7 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.OrderBy { using System; using System.Collections.Generic; - using System.Text; + using Microsoft.Azure.Cosmos.Query.Core; using Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.ItemProducers; using Microsoft.Azure.Cosmos.Query.Core.QueryClient; @@ -87,895 +87,5 @@ private CosmosOrderByItemQueryExecutionContext( testSettings: testSettings) { } - - /// - /// 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 static async Task> TryCreateAsync( - CosmosQueryContext queryContext, - CosmosCrossPartitionQueryExecutionContext.CrossPartitionInitParams initParams, - string 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(); - - // 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); - CosmosOrderByItemQueryExecutionContext context = new CosmosOrderByItemQueryExecutionContext( - initPararms: queryContext, - maxConcurrency: initParams.MaxConcurrency, - maxItemCount: initParams.MaxItemCount, - maxBufferedItemCount: initParams.MaxBufferedItemCount, - consumeComparer: orderByItemProducerTreeComparer, - testSettings: initParams.TestSettings); - - return (await context.TryInitializeAsync( - sqlQuerySpec: initParams.SqlQuerySpec, - requestContinuation: requestContinuationToken, - collectionRid: initParams.CollectionRid, - partitionKeyRanges: initParams.PartitionKeyRanges, - initialPageSize: initParams.InitialPageSize, - sortOrders: initParams.PartitionedQueryExecutionInfo.QueryInfo.OrderBy, - orderByExpressions: initParams.PartitionedQueryExecutionInfo.QueryInfo.OrderByExpressions, - cancellationToken: cancellationToken)) - .Try(() => context); - } - - /// - /// 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, - diagnostics: this.GetAndResetDiagnostics()); - } - - /// - /// 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); - } - - private async Task TryInitializeAsync( - SqlQuerySpec sqlQuerySpec, - string requestContinuation, - string collectionRid, - List partitionKeyRanges, - int initialPageSize, - SortOrder[] sortOrders, - string[] orderByExpressions, - 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 (sortOrders == null) - { - throw new ArgumentNullException(nameof(sortOrders)); - } - - if (orderByExpressions == null) - { - throw new ArgumentNullException(nameof(orderByExpressions)); - } - - cancellationToken.ThrowIfCancellationRequested(); - - if (requestContinuation == null) - { - SqlQuerySpec sqlQuerySpecForInit = new SqlQuerySpec( - sqlQuerySpec.QueryText.Replace(oldValue: FormatPlaceHolder, newValue: True), - sqlQuerySpec.Parameters); - - TryCatch tryInitialize = await base.TryInitializeAsync( - collectionRid, - partitionKeyRanges, - initialPageSize, - sqlQuerySpecForInit, - cancellationToken: cancellationToken, - targetRangeToContinuationMap: null, - deferFirstPage: false, - filter: null, - tryFilterAsync: null); - if (!tryInitialize.Succeeded) - { - return tryInitialize; - } - } - else - { - TryCatch tryExtractContinuationTokens = CosmosOrderByItemQueryExecutionContext.TryExtractContinuationTokens( - requestContinuation, - sortOrders, - orderByExpressions); - if (!tryExtractContinuationTokens.Succeeded) - { - return TryCatch.FromException(tryExtractContinuationTokens.Exception); - } - - TryCatch tryGetOrderByInitInfo = CosmosOrderByItemQueryExecutionContext.TryGetOrderByPartitionKeyRangesInitializationInfo( - tryExtractContinuationTokens.Result, - partitionKeyRanges, - sortOrders, - orderByExpressions); - if (!tryGetOrderByInitInfo.Succeeded) - { - return TryCatch.FromException(tryGetOrderByInitInfo.Exception); - } - - OrderByInitInfo initiaizationInfo = tryGetOrderByInitInfo.Result; - RangeFilterInitializationInfo[] orderByInfos = initiaizationInfo.Filters; - IReadOnlyDictionary targetRangeToOrderByContinuationMap = initiaizationInfo.ContinuationTokens; - Debug.Assert( - targetRangeToOrderByContinuationMap != null, - "If targetRangeToOrderByContinuationMap can't be null is valid continuation is supplied"); - - // 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) - foreach (RangeFilterInitializationInfo info in orderByInfos) - { - if (info.StartIndex > info.EndIndex) - { - continue; - } - - PartialReadOnlyList partialRanges = - new PartialReadOnlyList( - partitionKeyRanges, - info.StartIndex, - info.EndIndex - info.StartIndex + 1); - - SqlQuerySpec sqlQuerySpecForInit = new SqlQuerySpec( - sqlQuerySpec.QueryText.Replace(FormatPlaceHolder, info.Filter), - sqlQuerySpec.Parameters); - - TryCatch tryInitialize = await base.TryInitializeAsync( - collectionRid, - partialRanges, - initialPageSize, - sqlQuerySpecForInit, - targetRangeToOrderByContinuationMap.ToDictionary( - kvp => kvp.Key, - kvp => kvp.Value.CompositeContinuationToken.Token), - false, - info.Filter, - async (itemProducerTree) => - { - if (targetRangeToOrderByContinuationMap.TryGetValue( - itemProducerTree.Root.PartitionKeyRange.Id, - out OrderByContinuationToken continuationToken)) - { - TryCatch tryFilter = await this.TryFilterAsync( - itemProducerTree, - sortOrders, - continuationToken, - cancellationToken); - - if (!tryFilter.Succeeded) - { - return tryFilter; - } - } - - return TryCatch.FromResult(); - }, - cancellationToken); - if (!tryInitialize.Succeeded) - { - return tryInitialize; - } - } - } - - return TryCatch.FromResult(); - } - - private static TryCatch TryExtractContinuationTokens( - string requestContinuation, - SortOrder[] sortOrders, - string[] orderByExpressions) - { - Debug.Assert( - !(orderByExpressions == null - || orderByExpressions.Length <= 0 - || sortOrders == null - || sortOrders.Length <= 0 - || orderByExpressions.Length != sortOrders.Length), - "Partitioned QueryExecutionInfo returned bogus results."); - - if (string.IsNullOrWhiteSpace(requestContinuation)) - { - throw new ArgumentNullException("continuation can not be null or empty."); - } - - if (!CosmosArray.TryParse(requestContinuation, out 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 != sortOrders.Length) - { - return TryCatch.FromException( - new MalformedContinuationTokenException($"Invalid order-by items in continuation token {requestContinuation} for OrderBy~Context.")); - } - } - - return TryCatch.FromResult(orderByContinuationTokens.ToArray()); - } - - /// - /// 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, - SortOrder[] 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.Length; ++i) - { - cmp = ItemComparer.Instance.Compare( - continuationToken.OrderByItems[i].Item, - orderByResult.OrderByItems[i].Item); - - if (cmp != 0) - { - cmp = sortOrders[i] != SortOrder.Descending ? 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 first order by field. - // So if it's ORDER BY c.age ASC, c.name DESC the _rids are ASC - // If ti's ORDER BY c.age DESC, c.name DESC the _rids are DESC - cmp = continuationRid.Document.CompareTo(rid.Document); - if (sortOrders[0] == SortOrder.Descending) - { - 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(); - } - - /// - /// Gets the filters for every partition. - /// - private static TryCatch TryGetOrderByPartitionKeyRangesInitializationInfo( - OrderByContinuationToken[] suppliedContinuationTokens, - List partitionKeyRanges, - SortOrder[] sortOrders, - string[] orderByExpressions) - { - TryCatch> tryFindRangeAndContinuationTokensMonad = CosmosCrossPartitionQueryExecutionContext.TryFindTargetRangeAndExtractContinuationTokens( - partitionKeyRanges, - suppliedContinuationTokens - .Select(token => Tuple.Create(token, token.CompositeContinuationToken.Range))); - - return tryFindRangeAndContinuationTokensMonad.Try((indexAndContinuationTokens) => - { - int minIndex = indexAndContinuationTokens.TargetIndex; - IReadOnlyDictionary partitionKeyRangeToContinuationToken = indexAndContinuationTokens.ContinuationTokens; - - FormattedFilterInfo formattedFilterInfo = CosmosOrderByItemQueryExecutionContext.GetFormattedFilters( - orderByExpressions, - suppliedContinuationTokens, - sortOrders); - - RangeFilterInitializationInfo[] filters = new RangeFilterInitializationInfo[] - { - new RangeFilterInitializationInfo( - filter: formattedFilterInfo.FilterForRangesLeftOfTargetRanges, - startIndex: 0, - endIndex: minIndex - 1), - new RangeFilterInitializationInfo( - filter: formattedFilterInfo.FiltersForTargetRange, - startIndex: minIndex, - endIndex: minIndex), - new RangeFilterInitializationInfo( - filter: formattedFilterInfo.FilterForRangesRightOfTargetRanges, - startIndex: minIndex + 1, - endIndex: partitionKeyRanges.Count - 1), - }; - - return new OrderByInitInfo( - filters, - partitionKeyRangeToContinuationToken); - }); - } - - /// - /// Gets the formatted filters for every partition. - /// - /// The filter expressions. - /// The continuation token. - /// The sort orders. - /// The formatted filters for every partition. - private static FormattedFilterInfo GetFormattedFilters( - string[] expressions, - OrderByContinuationToken[] continuationTokens, - SortOrder[] sortOrders) - { - // Validate the inputs - for (int index = 0; index < continuationTokens.Length; index++) - { - Debug.Assert(continuationTokens[index].OrderByItems.Count == sortOrders.Length, "Expect values and orders are the same size."); - Debug.Assert(expressions.Length == sortOrders.Length, "Expect expressions and orders are the same size."); - } - - Tuple filters = CosmosOrderByItemQueryExecutionContext.GetFormattedFilters( - expressions, - continuationTokens[0].OrderByItems.Select(orderByItem => orderByItem.Item).ToArray(), - sortOrders); - - return new FormattedFilterInfo(filters.Item1, filters.Item2, filters.Item3); - } - - private static void AppendToBuilders(Tuple builders, object str) - { - CosmosOrderByItemQueryExecutionContext.AppendToBuilders(builders, str, str, str); - } - - private static void AppendToBuilders(Tuple builders, object left, object target, object right) - { - builders.Item1.Append(left); - builders.Item2.Append(target); - builders.Item3.Append(right); - } - - private static Tuple GetFormattedFilters( - string[] expressions, - CosmosElement[] orderByItems, - SortOrder[] sortOrders) - { - // 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 = expressions.Length; - bool isSingleOrderBy = numOrderByItems == 1; - StringBuilder left = new StringBuilder(); - StringBuilder target = new StringBuilder(); - StringBuilder right = new StringBuilder(); - - Tuple builders = new Tuple(left, right, target); - - 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 - string expression = expressions.First(); - SortOrder sortOrder = sortOrders.First(); - CosmosElement orderByItem = orderByItems.First(); - StringBuilder sb = new StringBuilder(); - CosmosElementToQueryLiteral cosmosElementToQueryLiteral = new CosmosElementToQueryLiteral(sb); - orderByItem.Accept(cosmosElementToQueryLiteral); - string orderByItemToString = sb.ToString(); - left.Append($"{expression} {(sortOrder == SortOrder.Descending ? "<" : ">")} {orderByItemToString}"); - target.Append($"{expression} {(sortOrder == SortOrder.Descending ? "<=" : ">=")} {orderByItemToString}"); - right.Append($"{expression} {(sortOrder == SortOrder.Descending ? "<=" : ">=")} {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++) - { - ArraySegment expressionPrefix = new ArraySegment(expressions, 0, prefixLength); - ArraySegment sortOrderPrefix = new ArraySegment(sortOrders, 0, prefixLength); - ArraySegment orderByItemsPrefix = new ArraySegment(orderByItems, 0, prefixLength); - - bool lastPrefix = prefixLength == numOrderByItems; - bool firstPrefix = prefixLength == 1; - - CosmosOrderByItemQueryExecutionContext.AppendToBuilders(builders, "("); - - for (int index = 0; index < prefixLength; index++) - { - string expression = expressionPrefix.ElementAt(index); - SortOrder sortOrder = sortOrderPrefix.ElementAt(index); - CosmosElement orderByItem = orderByItemsPrefix.ElementAt(index); - bool lastItem = index == prefixLength - 1; - - // Append Expression - CosmosOrderByItemQueryExecutionContext.AppendToBuilders(builders, expression); - CosmosOrderByItemQueryExecutionContext.AppendToBuilders(builders, " "); - - // Append binary operator - if (lastItem) - { - string inequality = sortOrder == SortOrder.Descending ? "<" : ">"; - CosmosOrderByItemQueryExecutionContext.AppendToBuilders(builders, inequality); - if (lastPrefix) - { - CosmosOrderByItemQueryExecutionContext.AppendToBuilders(builders, string.Empty, "=", "="); - } - } - else - { - CosmosOrderByItemQueryExecutionContext.AppendToBuilders(builders, "="); - } - - // 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 "); - } - } - } - - return new Tuple(left.ToString(), target.ToString(), 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; } - } - - /// - /// Struct to hold all the filters for every partition. - /// - private readonly struct FormattedFilterInfo - { - /// - /// Filters for current partition. - /// - public readonly string FiltersForTargetRange; - - /// - /// Filters for partitions left of the current partition. - /// - public readonly string FilterForRangesLeftOfTargetRanges; - - /// - /// Filters for partitions right of the current partition. - /// - public readonly string FilterForRangesRightOfTargetRanges; - - /// - /// Initializes a new instance of the FormattedFilterInfo struct. - /// - /// The filters for the partitions left of the current partition. - /// The filters for the current partition. - /// The filters for the partitions right of the current partition. - public FormattedFilterInfo(string leftFilter, string targetFilter, string rightFilters) - { - this.FilterForRangesLeftOfTargetRanges = leftFilter; - this.FiltersForTargetRange = targetFilter; - this.FilterForRangesRightOfTargetRanges = rightFilters; - } - } - - /// - /// 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; - } - } } } From 6494d65c996b7a512ccaf5044d9da4cadf2eb3fa Mon Sep 17 00:00:00 2001 From: Brandon Chong Date: Mon, 9 Mar 2020 11:37:23 -0700 Subject: [PATCH 25/28] resolved iteration comments --- .../CosmosOrderByItemQueryExecutionContext.Resume.cs | 10 +++++----- ...allelItemQueryExecutionContext.ContinuationToken.cs | 2 +- 2 files changed, 6 insertions(+), 6 deletions(-) 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 c39de51b99..13987e1b07 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 @@ -63,7 +63,7 @@ public static async Task> TryCreateAs IReadOnlyList orderByExpressions = initParams.PartitionedQueryExecutionInfo.QueryInfo.OrderByExpressions; IReadOnlyList sortOrders = initParams.PartitionedQueryExecutionInfo.QueryInfo.OrderBy; - if (orderByExpressions.Count() != sortOrders.Count()) + if (orderByExpressions.Count != sortOrders.Count) { throw new ArgumentException("order by expressions count does not match sort order"); } @@ -140,7 +140,7 @@ private async Task TryInitializeAsync( TryCatch> tryGetOrderByContinuationTokenMapping = TryGetOrderByContinuationTokenMapping( partitionKeyRanges, requestContinuation, - orderByColumns.Count()); + orderByColumns.Count); if (!tryGetOrderByContinuationTokenMapping.Succeeded) { return TryCatch.FromException(tryGetOrderByContinuationTokenMapping.Exception); @@ -154,7 +154,7 @@ private async Task TryInitializeAsync( .OrderByItems .Select(x => x.Item) .ToList(); - if (orderByItems.Count() != orderByColumns.Count()) + if (orderByItems.Count != orderByColumns.Count) { return TryCatch.FromException( new MalformedContinuationTokenException("Order By Items from continuation token did not match the query text.")); @@ -347,7 +347,7 @@ private async Task TryFilterAsync( 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) + for (int i = 0; i < sortOrders.Count; ++i) { cmp = ItemComparer.Instance.Compare( continuationToken.OrderByItems[i].Item, @@ -355,7 +355,7 @@ private async Task TryFilterAsync( if (cmp != 0) { - cmp = sortOrders[i] != SortOrder.Descending ? cmp : -cmp; + cmp = sortOrders[i] == SortOrder.Ascending ? cmp : -cmp; break; } } 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 index 2d5a42884a..8e2f3e5ca4 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/Parallel/CosmosParallelItemQueryExecutionContext.ContinuationToken.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/Parallel/CosmosParallelItemQueryExecutionContext.ContinuationToken.cs @@ -104,7 +104,7 @@ public bool Equals(CosmosElement x, CosmosElement y) /// The hash code for the object. public int GetHashCode(CosmosElement obj) { - return obj.GetHashCode(); + return obj == null ? 0 : obj.GetHashCode(); } } } From 69d098534651bf1db4b0c9510784ee6ef3991d11 Mon Sep 17 00:00:00 2001 From: Brandon Chong Date: Tue, 10 Mar 2020 21:44:36 -0700 Subject: [PATCH 26/28] Update Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/OrderBy/CosmosOrderByItemQueryExecutionContext.Resume.cs Co-Authored-By: j82w --- .../OrderBy/CosmosOrderByItemQueryExecutionContext.Resume.cs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 13987e1b07..5ae272e245 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 @@ -157,7 +157,7 @@ private async Task TryInitializeAsync( if (orderByItems.Count != orderByColumns.Count) { return TryCatch.FromException( - new MalformedContinuationTokenException("Order By Items from continuation token did not match the query text.")); + 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(); From 78339f098e1faabc8b73d1357201ad33957178f9 Mon Sep 17 00:00:00 2001 From: Brandon Chong Date: Wed, 11 Mar 2020 19:59:30 -0700 Subject: [PATCH 27/28] added tests and fixed off by one error --- ...smosCrossPartitionQueryExecutionContext.cs | 9 +- ...OrderByItemQueryExecutionContext.Resume.cs | 4 +- .../Query/ContinuationResumeLogicTests.cs | 352 ++++++++++++++++++ 3 files changed, 357 insertions(+), 8 deletions(-) create mode 100644 Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/ContinuationResumeLogicTests.cs diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosCrossPartitionQueryExecutionContext.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosCrossPartitionQueryExecutionContext.cs index fee65379e8..f9ae8f77bd 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosCrossPartitionQueryExecutionContext.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosCrossPartitionQueryExecutionContext.cs @@ -12,7 +12,6 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionContext using Core.ExecutionComponent; using Microsoft.Azure.Cosmos.CosmosElements; using Microsoft.Azure.Cosmos.Diagnostics; - using Microsoft.Azure.Cosmos.Json; using Microsoft.Azure.Cosmos.Query.Core; using Microsoft.Azure.Cosmos.Query.Core.Collections; using Microsoft.Azure.Cosmos.Query.Core.ComparableTask; @@ -23,8 +22,6 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionContext using Microsoft.Azure.Cosmos.Query.Core.Monads; using Microsoft.Azure.Cosmos.Query.Core.QueryClient; using Microsoft.Azure.Cosmos.Query.Core.QueryPlan; - using Microsoft.Azure.Cosmos.Sql; - using Microsoft.Azure.Documents.Routing; using PartitionKeyRange = Documents.PartitionKeyRange; using RequestChargeTracker = Documents.RequestChargeTracker; using RMResources = Documents.RMResources; @@ -434,7 +431,7 @@ protected async Task TryInitializeAsync( return TryCatch.FromResult(); } - protected static TryCatch> TryGetInitializationInfo( + public static TryCatch> TryGetInitializationInfo( IReadOnlyList partitionKeyRanges, IReadOnlyList partitionedContinuationTokens) where PartitionedToken : IPartitionedToken @@ -490,7 +487,7 @@ protected static TryCatch> TryGetInitializati $"{RMResources.InvalidContinuationToken} - Could not find continuation token: {firstContinuationToken}")); } - ReadOnlyMemory partitionsLeftOfTarget = matchedIndex == 0 ? ReadOnlyMemory.Empty : sortedRanges.Slice(start: 0, length: matchedIndex - 1); + 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); @@ -521,7 +518,7 @@ protected static TryCatch> TryGetInitializati /// The partition key ranges to match. /// The continuation tokens to match with. /// A dictionary of ranges matched with their continuation tokens. - private static IReadOnlyDictionary MatchRangesToContinuationTokens( + public static IReadOnlyDictionary MatchRangesToContinuationTokens( ReadOnlyMemory partitionKeyRanges, IReadOnlyList partitionedContinuationTokens) where PartitionedToken : IPartitionedToken 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 13987e1b07..657a966d5b 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 @@ -190,7 +190,7 @@ private async Task TryInitializeAsync( tryFilterAsync: async (itemProducerTree) => { if (!tokenMapping.TryGetValue( - itemProducerTree.PartitionKeyRange, + itemProducerTree.Root.PartitionKeyRange, out OrderByContinuationToken continuationToken)) { throw new InvalidOperationException($"Failed to retrieve {nameof(OrderByContinuationToken)}."); @@ -473,7 +473,7 @@ private static (string leftFilter, string targetFilter, string rightFilter) GetF StringBuilder target = new StringBuilder(); StringBuilder right = new StringBuilder(); - (StringBuilder, StringBuilder, StringBuilder) builders = (left, right, target); + (StringBuilder, StringBuilder, StringBuilder) builders = (left, target, right); if (isSingleOrderBy) { 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 new file mode 100644 index 0000000000..2bba535810 --- /dev/null +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/ContinuationResumeLogicTests.cs @@ -0,0 +1,352 @@ +namespace Microsoft.Azure.Cosmos.Query +{ + 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.Documents; + using Microsoft.VisualStudio.TestTools.UnitTesting; + using Newtonsoft.Json; + using static Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.CosmosCrossPartitionQueryExecutionContext; + + [TestClass] + public class ContinuationResumeLogicTests + { + [TestMethod] + public void TestMatchRangesTocontinuationTokens_OneToOne() + { + PartitionKeyRange partitionKeyRange = new PartitionKeyRange() + { + MinInclusive = string.Empty, + MaxExclusive = "FF", + Id = "0" + }; + + CompositeContinuationToken token = new CompositeContinuationToken() + { + Range = new Documents.Routing.Range( + min: string.Empty, + max: "FF", + isMinInclusive: true, + isMaxInclusive: false), + Token = "asdf" + }; + + IReadOnlyDictionary expectedMapping = new Dictionary() + { + { partitionKeyRange, token } + }; + + ContinuationResumeLogicTests.RunMatchRangesToContinuationTokens( + expectedMapping, + new PartitionKeyRange[] { partitionKeyRange }, + new CompositeContinuationToken[] { token }); + } + + [TestMethod] + public void TestMatchRangesTocontinuationTokens_OneToMany() + { + PartitionKeyRange partitionKeyRange1 = new PartitionKeyRange() + { + MinInclusive = string.Empty, + MaxExclusive = "A", + Id = "1" + }; + + PartitionKeyRange partitionKeyRange2 = new PartitionKeyRange() + { + MinInclusive = "A", + MaxExclusive = "B", + Id = "1" + }; + + CompositeContinuationToken token = new CompositeContinuationToken() + { + Range = new Documents.Routing.Range( + min: string.Empty, + max: "B", + isMinInclusive: true, + isMaxInclusive: false), + Token = "asdf" + }; + + IReadOnlyDictionary expectedMapping = new Dictionary() + { + { partitionKeyRange1, token }, + { partitionKeyRange2, token } + }; + + ContinuationResumeLogicTests.RunMatchRangesToContinuationTokens( + expectedMapping, + new PartitionKeyRange[] { partitionKeyRange1, partitionKeyRange2 }, + new CompositeContinuationToken[] { token }); + } + + [TestMethod] + public void TestMatchRangesTocontinuationTokens_OneToNone() + { + PartitionKeyRange partitionKeyRange = new PartitionKeyRange() + { + MinInclusive = string.Empty, + MaxExclusive = "A", + Id = "1" + }; + + CompositeContinuationToken token = new CompositeContinuationToken() + { + Range = new Documents.Routing.Range( + min: "B", + max: "C", + isMinInclusive: true, + isMaxInclusive: false), + Token = "asdf" + }; + + IReadOnlyDictionary expectedMapping = new Dictionary() + { + { partitionKeyRange, null }, + }; + + ContinuationResumeLogicTests.RunMatchRangesToContinuationTokens( + expectedMapping, + new PartitionKeyRange[] { partitionKeyRange }, + new CompositeContinuationToken[] { token }); + } + + [TestMethod] + [ExpectedException(typeof(ArgumentNullException))] + public void TestMatchRangesTocontinuationTokens_ArgumentNullException() + { + ContinuationResumeLogicTests.RunMatchRangesToContinuationTokens( + expectedMapping: null, + partitionKeyRanges: new PartitionKeyRange[] { }, + partitionedTokens: null); + } + + [TestMethod] + public void TestTryGetInitializationInfo_ResumeLeftMostPartition() + { + 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 = "C", + Id = "3" + }; + + CompositeContinuationToken token = new CompositeContinuationToken() + { + Range = new Documents.Routing.Range( + min: string.Empty, + max: "A", + isMinInclusive: true, + isMaxInclusive: false), + Token = "asdf" + }; + + IReadOnlyDictionary expectedMappingLeftPartitions = new Dictionary() + { + }; + + IReadOnlyDictionary expectedMappingTargetPartition = new Dictionary() + { + { pkRange1, token} + }; + + IReadOnlyDictionary expectedMappingRightPartitions = new Dictionary() + { + { pkRange2, null}, + { pkRange3, null}, + }; + + RunTryGetInitializationInfo( + expectedMappingLeftPartitions, + expectedMappingTargetPartition, + expectedMappingRightPartitions, + new PartitionKeyRange[] { pkRange1, pkRange2, pkRange3 }, + new IPartitionedToken[] { token }); + } + + [TestMethod] + public void TestTryGetInitializationInfo_ResumeMiddlePartition() + { + 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 = "C", + Id = "3" + }; + + CompositeContinuationToken token = new CompositeContinuationToken() + { + Range = new Documents.Routing.Range( + min: "A", + max: "B", + isMinInclusive: true, + isMaxInclusive: false), + Token = "asdf" + }; + + IReadOnlyDictionary expectedMappingLeftPartitions = new Dictionary() + { + { pkRange1, null} + }; + + IReadOnlyDictionary expectedMappingTargetPartition = new Dictionary() + { + { pkRange2, token}, + }; + + IReadOnlyDictionary expectedMappingRightPartitions = new Dictionary() + { + { pkRange3, null}, + }; + + RunTryGetInitializationInfo( + expectedMappingLeftPartitions, + expectedMappingTargetPartition, + expectedMappingRightPartitions, + new PartitionKeyRange[] { pkRange1, pkRange2, pkRange3 }, + new IPartitionedToken[] { token }); + } + + [TestMethod] + public void TestTryGetInitializationInfo_ResumeRightPartition() + { + 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 = "C", + Id = "3" + }; + + CompositeContinuationToken token = new CompositeContinuationToken() + { + Range = new Documents.Routing.Range( + min: "B", + max: "C", + isMinInclusive: true, + isMaxInclusive: false), + Token = "asdf" + }; + + IReadOnlyDictionary expectedMappingLeftPartitions = new Dictionary() + { + { pkRange1, null}, + { pkRange2, null}, + }; + + IReadOnlyDictionary expectedMappingTargetPartition = new Dictionary() + { + { pkRange3, token}, + }; + + IReadOnlyDictionary expectedMappingRightPartitions = new Dictionary() + { + }; + + RunTryGetInitializationInfo( + expectedMappingLeftPartitions, + expectedMappingTargetPartition, + expectedMappingRightPartitions, + new PartitionKeyRange[] { pkRange1, pkRange2, pkRange3 }, + new IPartitionedToken[] { token }); + } + + private static void RunMatchRangesToContinuationTokens( + IReadOnlyDictionary expectedMapping, + IEnumerable partitionKeyRanges, + IEnumerable partitionedTokens) + { + IReadOnlyDictionary actualMapping = CosmosCrossPartitionQueryExecutionContext.MatchRangesToContinuationTokens( + partitionKeyRanges.OrderBy(x => Guid.NewGuid()).ToArray(), + partitionedTokens.OrderBy(x => Guid.NewGuid()).ToList()); + + ContinuationResumeLogicTests.AssertPartitionMappingAreEqual( + expectedMapping, + actualMapping); + } + + private static void RunTryGetInitializationInfo( + IReadOnlyDictionary expectedMappingLeftPartitions, + IReadOnlyDictionary expectedMappingTargetPartition, + IReadOnlyDictionary expectedMappingRightPartitions, + IEnumerable partitionKeyRanges, + IEnumerable partitionedTokens) + { + TryCatch> tryGetInitializationInfo = CosmosCrossPartitionQueryExecutionContext.TryGetInitializationInfo( + partitionKeyRanges.OrderBy(x => Guid.NewGuid()).ToArray(), + partitionedTokens.OrderBy(x => Guid.NewGuid()).ToList()); + Assert.IsTrue(tryGetInitializationInfo.Succeeded); + PartitionMapping partitionMapping = tryGetInitializationInfo.Result; + + AssertPartitionMappingAreEqual(expectedMappingLeftPartitions, partitionMapping.PartitionsLeftOfTarget); + AssertPartitionMappingAreEqual(expectedMappingTargetPartition, partitionMapping.TargetPartition); + AssertPartitionMappingAreEqual(expectedMappingRightPartitions, partitionMapping.PartitionsRightOfTarget); + } + + private static void AssertPartitionMappingAreEqual( + IReadOnlyDictionary expectedMapping, + IReadOnlyDictionary actualMapping) + { + Assert.IsNotNull(expectedMapping); + Assert.IsNotNull(actualMapping); + + Assert.AreEqual(expected: expectedMapping.Count, actual: actualMapping.Count); + + foreach (KeyValuePair kvp in expectedMapping) + { + Assert.IsTrue( + actualMapping.TryGetValue( + kvp.Key, + out IPartitionedToken partitionedToken)); + Assert.AreEqual( + expected: JsonConvert.SerializeObject(kvp.Value), + actual: JsonConvert.SerializeObject(partitionedToken)); + } + } + } +} From 2737293af4590516169da705d393a417da8e823e Mon Sep 17 00:00:00 2001 From: Brandon Chong Date: Thu, 12 Mar 2020 09:04:08 -0700 Subject: [PATCH 28/28] fixed typo --- ...smosParallelItemQueryExecutionContext.ContinuationToken.cs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) 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 index 8e2f3e5ca4..0380c994e9 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/Parallel/CosmosParallelItemQueryExecutionContext.ContinuationToken.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/Parallel/CosmosParallelItemQueryExecutionContext.ContinuationToken.cs @@ -69,8 +69,8 @@ public override CosmosElement GetCosmosElementContinuationToken() Range = new Documents.Routing.Range( min: activeItemProducer.PartitionKeyRange.MinInclusive, max: activeItemProducer.PartitionKeyRange.MaxExclusive, - isMinInclusive: false, - isMaxInclusive: true) + isMinInclusive: true, + isMaxInclusive: false) }; CosmosElement compositeContinuationToken = CompositeContinuationToken.ToCosmosElement(compositeToken);