com.azure.cosmos.implementation.query.OrderByDocumentQueryExecutionContext Maven / Gradle / Ivy
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
package com.azure.cosmos.implementation.query;
import com.azure.cosmos.implementation.query.orderbyquery.OrderByRowResult;
import com.azure.cosmos.implementation.query.orderbyquery.OrderbyRowComparer;
import com.azure.cosmos.implementation.routing.Range;
import com.azure.cosmos.BridgeInternal;
import com.azure.cosmos.CosmosClientException;
import com.azure.cosmos.FeedOptions;
import com.azure.cosmos.FeedResponse;
import com.azure.cosmos.Resource;
import com.azure.cosmos.SqlQuerySpec;
import com.azure.cosmos.implementation.HttpConstants;
import com.azure.cosmos.implementation.IDocumentClientRetryPolicy;
import com.azure.cosmos.implementation.PartitionKeyRange;
import com.azure.cosmos.implementation.QueryMetrics;
import com.azure.cosmos.implementation.RequestChargeTracker;
import com.azure.cosmos.implementation.ResourceType;
import com.azure.cosmos.implementation.RxDocumentServiceRequest;
import com.azure.cosmos.implementation.Utils;
import com.azure.cosmos.implementation.Utils.ValueHolder;
import org.apache.commons.lang3.NotImplementedException;
import org.apache.commons.lang3.tuple.ImmutablePair;
import reactor.core.publisher.Flux;
import java.util.ArrayList;
import java.util.Collection;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.UUID;
import java.util.concurrent.Callable;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
import java.util.function.Function;
/**
* While this class is public, but it is not part of our published public APIs.
* This is meant to be internally used only by our sdk.
*/
public class OrderByDocumentQueryExecutionContext
extends ParallelDocumentQueryExecutionContextBase {
private final String FormatPlaceHolder = "{documentdb-formattableorderbyquery-filter}";
private final String True = "true";
private final String collectionRid;
private final OrderbyRowComparer consumeComparer;
private final RequestChargeTracker tracker;
private final ConcurrentMap queryMetricMap;
private Flux> orderByObservable;
private final Map targetRangeToOrderByContinuationTokenMap;
private OrderByDocumentQueryExecutionContext(
IDocumentQueryClient client,
List partitionKeyRanges,
ResourceType resourceTypeEnum,
Class klass,
SqlQuerySpec query,
FeedOptions feedOptions,
String resourceLink,
String rewrittenQuery,
boolean isContinuationExpected,
boolean getLazyFeedResponse,
OrderbyRowComparer consumeComparer,
String collectionRid,
UUID correlatedActivityId) {
super(client, partitionKeyRanges, resourceTypeEnum, klass, query, feedOptions, resourceLink, rewrittenQuery,
isContinuationExpected, getLazyFeedResponse, correlatedActivityId);
this.collectionRid = collectionRid;
this.consumeComparer = consumeComparer;
this.tracker = new RequestChargeTracker();
this.queryMetricMap = new ConcurrentHashMap<>();
targetRangeToOrderByContinuationTokenMap = new HashMap<>();
}
public static Flux> createAsync(
IDocumentQueryClient client,
ResourceType resourceTypeEnum,
Class resourceType,
SqlQuerySpec expression,
FeedOptions feedOptions,
String resourceLink,
String collectionRid,
PartitionedQueryExecutionInfo partitionedQueryExecutionInfo,
List partitionKeyRanges,
int initialPageSize,
boolean isContinuationExpected,
boolean getLazyFeedResponse,
UUID correlatedActivityId) {
OrderByDocumentQueryExecutionContext context = new OrderByDocumentQueryExecutionContext(client,
partitionKeyRanges,
resourceTypeEnum,
resourceType,
expression,
feedOptions,
resourceLink,
partitionedQueryExecutionInfo.getQueryInfo().getRewrittenQuery(),
isContinuationExpected,
getLazyFeedResponse,
new OrderbyRowComparer(partitionedQueryExecutionInfo.getQueryInfo().getOrderBy()),
collectionRid,
correlatedActivityId);
try {
context.initialize(partitionKeyRanges,
partitionedQueryExecutionInfo.getQueryInfo().getOrderBy(),
partitionedQueryExecutionInfo.getQueryInfo().getOrderByExpressions(),
initialPageSize,
feedOptions.requestContinuation());
return Flux.just(context);
} catch (CosmosClientException dce) {
return Flux.error(dce);
}
}
private void initialize(
List partitionKeyRanges,
List sortOrders,
Collection orderByExpressions,
int initialPageSize,
String continuationToken) throws CosmosClientException {
if (continuationToken == null) {
// First iteration so use null continuation tokens and "true" filters
Map partitionKeyRangeToContinuationToken = new HashMap();
for (PartitionKeyRange partitionKeyRange : partitionKeyRanges) {
partitionKeyRangeToContinuationToken.put(partitionKeyRange,
null);
}
super.initialize(collectionRid,
partitionKeyRangeToContinuationToken,
initialPageSize,
new SqlQuerySpec(querySpec.getQueryText().replace(FormatPlaceHolder,
True),
querySpec.getParameters()));
} else {
// Check to see if order by continuation token is a valid JSON.
OrderByContinuationToken orderByContinuationToken;
ValueHolder outOrderByContinuationToken = new ValueHolder();
if (!OrderByContinuationToken.tryParse(continuationToken,
outOrderByContinuationToken)) {
String message = String.format("INVALID JSON in continuation token %s for OrderBy~Context",
continuationToken);
throw BridgeInternal.createCosmosClientException(HttpConstants.StatusCodes.BADREQUEST,
message);
}
orderByContinuationToken = outOrderByContinuationToken.v;
CompositeContinuationToken compositeContinuationToken = orderByContinuationToken
.getCompositeContinuationToken();
// Check to see if the ranges inside are valid
if (compositeContinuationToken.getRange().isEmpty()) {
String message = String.format("INVALID RANGE in the continuation token %s for OrderBy~Context.",
continuationToken);
throw BridgeInternal.createCosmosClientException(HttpConstants.StatusCodes.BADREQUEST,
message);
}
// At this point the token is valid.
ImmutablePair targetIndexAndFilters = this.GetFiltersForPartitions(
orderByContinuationToken,
partitionKeyRanges,
sortOrders,
orderByExpressions);
int targetIndex = targetIndexAndFilters.left;
targetRangeToOrderByContinuationTokenMap.put(String.valueOf(targetIndex), orderByContinuationToken);
FormattedFilterInfo formattedFilterInfo = targetIndexAndFilters.right;
// Left
String filterForRangesLeftOfTheTargetRange = formattedFilterInfo.getFilterForRangesLeftOfTheTargetRange();
this.initializeRangeWithContinuationTokenAndFilter(partitionKeyRanges,
/* startInclusive */ 0,
/* endExclusive */ targetIndex,
/* continuationToken */ null,
filterForRangesLeftOfTheTargetRange,
initialPageSize);
// Target
String filterForTargetRange = formattedFilterInfo.getFilterForTargetRange();
this.initializeRangeWithContinuationTokenAndFilter(partitionKeyRanges,
/* startInclusive */ targetIndex,
/* endExclusive */ targetIndex + 1,
null,
filterForTargetRange,
initialPageSize);
// Right
String filterForRangesRightOfTheTargetRange = formattedFilterInfo.getFilterForRangesRightOfTheTargetRange();
this.initializeRangeWithContinuationTokenAndFilter(partitionKeyRanges,
/* startInclusive */ targetIndex + 1,
/* endExclusive */ partitionKeyRanges.size(),
/* continuationToken */ null,
filterForRangesRightOfTheTargetRange,
initialPageSize);
}
orderByObservable = OrderByUtils.orderedMerge(resourceType,
consumeComparer,
tracker,
documentProducers,
queryMetricMap,
targetRangeToOrderByContinuationTokenMap);
}
private void initializeRangeWithContinuationTokenAndFilter(
List partitionKeyRanges,
int startInclusive,
int endExclusive,
String continuationToken,
String filter,
int initialPageSize) {
Map partitionKeyRangeToContinuationToken = new HashMap();
for (int i = startInclusive; i < endExclusive; i++) {
PartitionKeyRange partitionKeyRange = partitionKeyRanges.get(i);
partitionKeyRangeToContinuationToken.put(partitionKeyRange,
continuationToken);
}
super.initialize(collectionRid,
partitionKeyRangeToContinuationToken,
initialPageSize,
new SqlQuerySpec(querySpec.getQueryText().replace(FormatPlaceHolder,
filter),
querySpec.getParameters()));
}
private ImmutablePair GetFiltersForPartitions(
OrderByContinuationToken orderByContinuationToken,
List partitionKeyRanges,
List sortOrders,
Collection orderByExpressions) throws CosmosClientException {
// Find the partition key range we left off on
int startIndex = this.FindTargetRangeAndExtractContinuationTokens(partitionKeyRanges,
orderByContinuationToken.getCompositeContinuationToken().getRange());
// Get the filters.
FormattedFilterInfo formattedFilterInfo = this.GetFormattedFilters(orderByExpressions,
orderByContinuationToken.getOrderByItems(),
sortOrders,
orderByContinuationToken.getInclusive());
return new ImmutablePair(startIndex,
formattedFilterInfo);
}
private OrderByDocumentQueryExecutionContext.FormattedFilterInfo GetFormattedFilters(
Collection orderByExpressionCollection,
QueryItem[] orderByItems,
Collection sortOrderCollection,
boolean inclusive) {
// Convert to arrays
SortOrder[] sortOrders = new SortOrder[sortOrderCollection.size()];
sortOrderCollection.toArray(sortOrders);
String[] expressions = new String[orderByExpressionCollection.size()];
orderByExpressionCollection.toArray(expressions);
// Validate the inputs
if (expressions.length != sortOrders.length) {
throw new IllegalArgumentException("expressions.size() != sortOrders.size()");
}
if (expressions.length != orderByItems.length) {
throw new IllegalArgumentException("expressions.size() != orderByItems.length");
}
// 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 partitions.
// 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;
boolean isSingleOrderBy = numOrderByItems == 1;
StringBuilder left = new StringBuilder();
StringBuilder target = new StringBuilder();
StringBuilder right = new StringBuilder();
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[0];
SortOrder sortOrder = sortOrders[0];
QueryItem orderByItem = orderByItems[0];
Object rawItem = orderByItem.getItem();
String orderByItemToString;
if (rawItem instanceof String) {
orderByItemToString = "\"" + rawItem.toString().replaceAll("\"",
"\\\"") + "\"";
} else {
orderByItemToString = rawItem.toString();
}
left.append(String.format("%s %s %s",
expression,
(sortOrder == SortOrder.Descending ? "<" : ">"),
orderByItemToString));
if (inclusive) {
target.append(String.format("%s %s %s",
expression,
(sortOrder == SortOrder.Descending ? "<=" : ">="),
orderByItemToString));
} else {
target.append(String.format("%s %s %s",
expression,
(sortOrder == SortOrder.Descending ? "<" : ">"),
orderByItemToString));
}
right.append(String.format("%s %s %s",
expression,
(sortOrder == SortOrder.Descending ? "<=" : ">="),
orderByItemToString));
} else {
// This code path needs to be implemented, but it's error prone and needs
// testing.
// You can port the implementation from the .net SDK and it should work if
// ported right.
throw new NotImplementedException(
"Resuming a multi order by query from a continuation token is not supported yet.");
}
return new FormattedFilterInfo(left.toString(),
target.toString(),
right.toString());
}
protected OrderByDocumentProducer createDocumentProducer(
String collectionRid,
PartitionKeyRange targetRange,
String continuationToken,
int initialPageSize,
FeedOptions feedOptions,
SqlQuerySpec querySpecForInit,
Map commonRequestHeaders,
TriFunction createRequestFunc,
Function>> executeFunc,
Callable createRetryPolicyFunc) {
return new OrderByDocumentProducer(consumeComparer,
client,
collectionRid,
feedOptions,
createRequestFunc,
executeFunc,
targetRange,
collectionRid,
() -> client.getResetSessionTokenRetryPolicy().getRequestPolicy(),
resourceType,
correlatedActivityId,
initialPageSize,
continuationToken,
top,
this.targetRangeToOrderByContinuationTokenMap);
}
private static class ItemToPageTransformer
implements Function>, Flux>> {
private final static int DEFAULT_PAGE_SIZE = 100;
private final RequestChargeTracker tracker;
private final int maxPageSize;
private final ConcurrentMap queryMetricMap;
private final Function, String> orderByContinuationTokenCallback;
private volatile FeedResponse> previousPage;
public ItemToPageTransformer(
RequestChargeTracker tracker,
int maxPageSize,
ConcurrentMap queryMetricsMap,
Function, String> orderByContinuationTokenCallback) {
this.tracker = tracker;
this.maxPageSize = maxPageSize > 0 ? maxPageSize : DEFAULT_PAGE_SIZE;
this.queryMetricMap = queryMetricsMap;
this.orderByContinuationTokenCallback = orderByContinuationTokenCallback;
this.previousPage = null;
}
private static Map headerResponse(
double requestCharge) {
return Utils.immutableMapOf(HttpConstants.HttpHeaders.REQUEST_CHARGE,
String.valueOf(requestCharge));
}
private FeedResponse> addOrderByContinuationToken(
FeedResponse> page,
String orderByContinuationToken) {
Map headers = new HashMap<>(page.getResponseHeaders());
headers.put(HttpConstants.HttpHeaders.CONTINUATION,
orderByContinuationToken);
return BridgeInternal.createFeedResponseWithQueryMetrics(page.getResults(),
headers,
BridgeInternal.queryMetricsFromFeedResponse(page));
}
@Override
public Flux> apply(Flux> source) {
return source
// .windows: creates an observable of observable where inner observable
// emits max maxPageSize elements
.window(maxPageSize).map(Flux::collectList)
// flattens the observable>>> to
// Observable>>
.flatMap(resultListObs -> resultListObs,
1)
// translates Observable>> to
// Observable>>>
.map(orderByRowResults -> {
// construct a page from result with request charge
FeedResponse> feedResponse = BridgeInternal.createFeedResponse(
orderByRowResults,
headerResponse(tracker.getAndResetCharge()));
if (!queryMetricMap.isEmpty()) {
for (String key : queryMetricMap.keySet()) {
BridgeInternal.putQueryMetricsIntoMap(feedResponse,
key,
queryMetricMap.get(key));
}
}
return feedResponse;
})
// Emit an empty page so the downstream observables know when there are no more
// results.
.concatWith(Flux.defer(() -> {
return Flux.just(BridgeInternal.createFeedResponse(Utils.immutableListOf(),
null));
}))
// CREATE pairs from the stream to allow the observables downstream to "peek"
// 1, 2, 3, null -> (null, 1), (1, 2), (2, 3), (3, null)
.map(orderByRowResults -> {
ImmutablePair>, FeedResponse>> previousCurrent = new ImmutablePair>, FeedResponse>>(
this.previousPage,
orderByRowResults);
this.previousPage = orderByRowResults;
return previousCurrent;
})
// remove the (null, 1)
.skip(1)
// Add the continuation token based on the current and next page.
.map(currentNext -> {
FeedResponse> current = currentNext.left;
FeedResponse> next = currentNext.right;
FeedResponse> page;
if (next.getResults().size() == 0) {
// No more pages no send current page with null continuation token
page = current;
page = this.addOrderByContinuationToken(page,
null);
} else {
// Give the first page but use the first value in the next page to generate the
// continuation token
page = current;
List> results = next.getResults();
OrderByRowResult firstElementInNextPage = results.get(0);
String orderByContinuationToken = this.orderByContinuationTokenCallback
.apply(firstElementInNextPage);
page = this.addOrderByContinuationToken(page,
orderByContinuationToken);
}
return page;
}).map(feedOfOrderByRowResults -> {
// FeedResponse> to FeedResponse
List unwrappedResults = new ArrayList();
for (OrderByRowResult orderByRowResult : feedOfOrderByRowResults.getResults()) {
unwrappedResults.add(orderByRowResult.getPayload());
}
return BridgeInternal.createFeedResponseWithQueryMetrics(unwrappedResults,
feedOfOrderByRowResults.getResponseHeaders(),
BridgeInternal.queryMetricsFromFeedResponse(feedOfOrderByRowResults));
}).switchIfEmpty(Flux.defer(() -> {
// create an empty page if there is no result
return Flux.just(BridgeInternal.createFeedResponse(Utils.immutableListOf(),
headerResponse(tracker.getAndResetCharge())));
}));
}
}
@Override
public Flux> drainAsync(
int maxPageSize) {
//// 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
return this.orderByObservable.compose(new ItemToPageTransformer(tracker,
maxPageSize,
this.queryMetricMap,
this::getContinuationToken));
}
@Override
public Flux> executeAsync() {
return drainAsync(feedOptions.maxItemCount());
}
private String getContinuationToken(
OrderByRowResult orderByRowResult) {
// rid
String rid = orderByRowResult.getResourceId();
// CompositeContinuationToken
String backendContinuationToken = orderByRowResult.getSourceBackendContinuationToken();
Range range = orderByRowResult.getSourcePartitionKeyRange().toRange();
boolean inclusive = true;
CompositeContinuationToken compositeContinuationToken = new CompositeContinuationToken(backendContinuationToken,
range);
// OrderByItems
QueryItem[] orderByItems = new QueryItem[orderByRowResult.getOrderByItems().size()];
orderByRowResult.getOrderByItems().toArray(orderByItems);
return new OrderByContinuationToken(compositeContinuationToken,
orderByItems,
rid,
inclusive).toJson();
}
private final class FormattedFilterInfo {
private final String filterForRangesLeftOfTheTargetRange;
private final String filterForTargetRange;
private final String filterForRangesRightOfTheTargetRange;
public FormattedFilterInfo(
String filterForRangesLeftOfTheTargetRange,
String filterForTargetRange,
String filterForRangesRightOfTheTargetRange) {
if (filterForRangesLeftOfTheTargetRange == null) {
throw new IllegalArgumentException("filterForRangesLeftOfTheTargetRange must not be null.");
}
if (filterForTargetRange == null) {
throw new IllegalArgumentException("filterForTargetRange must not be null.");
}
if (filterForRangesRightOfTheTargetRange == null) {
throw new IllegalArgumentException("filterForRangesRightOfTheTargetRange must not be null.");
}
this.filterForRangesLeftOfTheTargetRange = filterForRangesLeftOfTheTargetRange;
this.filterForTargetRange = filterForTargetRange;
this.filterForRangesRightOfTheTargetRange = filterForRangesRightOfTheTargetRange;
}
/**
* @return the filterForRangesLeftOfTheTargetRange
*/
public String getFilterForRangesLeftOfTheTargetRange() {
return filterForRangesLeftOfTheTargetRange;
}
/**
* @return the filterForTargetRange
*/
public String getFilterForTargetRange() {
return filterForTargetRange;
}
/**
* @return the filterForRangesRightOfTheTargetRange
*/
public String getFilterForRangesRightOfTheTargetRange() {
return filterForRangesRightOfTheTargetRange;
}
}
}
© 2015 - 2025 Weber Informatics LLC | Privacy Policy