public static async Task <TryCatch <IDocumentQueryExecutionComponent> > 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);
IReadOnlyList <string> orderByExpressions = initParams.PartitionedQueryExecutionInfo.QueryInfo.OrderByExpressions;
IReadOnlyList <SortOrder> sortOrders = initParams.PartitionedQueryExecutionInfo.QueryInfo.OrderBy;
if (orderByExpressions.Count != sortOrders.Count)
{
throw new ArgumentException("order by expressions count does not match sort order");
}
IReadOnlyList <OrderByColumn> 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,
orderByColumns: columns,
cancellationToken: cancellationToken))
.Try <IDocumentQueryExecutionComponent>(() => context));
}
private async Task <TryCatch> TryInitializeAsync(
SqlQuerySpec sqlQuerySpec,
CosmosElement requestContinuation,
string collectionRid,
IReadOnlyList <PartitionKeyRange> partitionKeyRanges,
int initialPageSize,
IReadOnlyList <OrderByColumn> orderByColumns,
CancellationToken cancellationToken)
{
if (sqlQuerySpec == null)
{
throw new ArgumentNullException(nameof(sqlQuerySpec));
}
if (collectionRid == null)
{
throw new ArgumentNullException(nameof(collectionRid));
}
if (partitionKeyRanges == null)
{
throw new ArgumentNullException(nameof(partitionKeyRanges));
}
if (orderByColumns == null)
{
throw new ArgumentNullException(nameof(orderByColumns));
}
cancellationToken.ThrowIfCancellationRequested();
if (requestContinuation == null)
{
// Start off all the partition key ranges with null continuation
SqlQuerySpec rewrittenQueryForOrderBy = new SqlQuerySpec(
sqlQuerySpec.QueryText.Replace(oldValue: FormatPlaceHolder, newValue: True),
sqlQuerySpec.Parameters);
Dictionary <PartitionKeyRange, string> partitionKeyRangeToContinuationToken = new Dictionary <PartitionKeyRange, string>();
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 <PartitionMapping <OrderByContinuationToken> > tryGetOrderByContinuationTokenMapping = TryGetOrderByContinuationTokenMapping(
partitionKeyRanges,
requestContinuation,
orderByColumns.Count);
if (!tryGetOrderByContinuationTokenMapping.Succeeded)
{
return(TryCatch.FromException(tryGetOrderByContinuationTokenMapping.Exception));
}
IReadOnlyList <CosmosElement> orderByItems = tryGetOrderByContinuationTokenMapping
.Result
.TargetPartition
.Values
.First()
.OrderByItems
.Select(x => x.Item)
.ToList();
if (orderByItems.Count != orderByColumns.Count)
{
return(TryCatch.FromException(
new MalformedContinuationTokenException($"Order By Items from continuation token did not match the query text. Order by item count: {orderByItems.Count()} did not match column count {orderByColumns.Count()}. Continuation token: {requestContinuation}")));
}
ReadOnlyMemory <(OrderByColumn, CosmosElement)> columnAndItems = orderByColumns.Zip(orderByItems, (column, item) => (column, item)).ToArray();
// For ascending order-by, left of target partition has filter expression > value,
// right of target partition has filter expression >= value,
// and target partition takes the previous filter from continuation (or true if no continuation)
(string leftFilter, string targetFilter, string rightFilter) = CosmosOrderByItemQueryExecutionContext.GetFormattedFilters(columnAndItems);
List <(IReadOnlyDictionary <PartitionKeyRange, OrderByContinuationToken>, string)> tokenMappingAndFilters = new List <(IReadOnlyDictionary <PartitionKeyRange, OrderByContinuationToken>, string)>()
{
{ (tryGetOrderByContinuationTokenMapping.Result.PartitionsLeftOfTarget, leftFilter) },
{ (tryGetOrderByContinuationTokenMapping.Result.TargetPartition, targetFilter) },
