com.azure.cosmos.implementation.query.OrderByUtils Maven / Gradle / Ivy
Go to download
Show more of this group Show more artifacts with this name
Show all versions of azure-cosmos Show documentation
Show all versions of azure-cosmos Show documentation
This Package contains Microsoft Azure Cosmos SDK (with Reactive Extension Reactor support) for Azure Cosmos DB SQL API
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
package com.azure.cosmos.implementation.query;
import com.azure.cosmos.implementation.ClientSideRequestStatistics;
import com.azure.cosmos.implementation.feedranges.FeedRangeEpkImpl;
import com.azure.cosmos.implementation.query.orderbyquery.OrderByRowResult;
import com.azure.cosmos.implementation.query.orderbyquery.OrderbyRowComparer;
import com.azure.cosmos.implementation.BadRequestException;
import com.azure.cosmos.BridgeInternal;
import com.azure.cosmos.models.ModelBridgeInternal;
import com.azure.cosmos.implementation.Resource;
import com.azure.cosmos.implementation.QueryMetrics;
import com.azure.cosmos.implementation.RequestChargeTracker;
import com.azure.cosmos.implementation.ResourceId;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.node.ArrayNode;
import com.azure.cosmos.implementation.apachecommons.lang.tuple.Pair;
import reactor.core.publisher.Flux;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.function.Function;
import java.util.stream.Collectors;
class OrderByUtils {
public static Flux> orderedMerge(Class klass,
OrderbyRowComparer consumeComparer,
RequestChargeTracker tracker,
List> documentProducers,
Map queryMetricsMap,
Map targetRangeToOrderByContinuationTokenMap,
List clientSideRequestStatisticsList) {
@SuppressWarnings("unchecked")
Flux>[] fluxes = documentProducers
.subList(0, documentProducers.size())
.stream()
.map(producer ->
toOrderByQueryResultObservable(klass, producer, tracker, queryMetricsMap,
targetRangeToOrderByContinuationTokenMap,
consumeComparer.getSortOrders(), clientSideRequestStatisticsList))
.toArray(Flux[]::new);
return Flux.mergeOrdered(consumeComparer, fluxes);
}
private static Flux> toOrderByQueryResultObservable(Class klass,
DocumentProducer producer,
RequestChargeTracker tracker,
Map queryMetricsMap,
Map targetRangeToOrderByContinuationTokenMap,
List sortOrders,
List clientSideRequestStatisticsList) {
return producer
.produceAsync()
.transformDeferred(new OrderByUtils.PageToItemTransformer(klass, tracker, queryMetricsMap,
targetRangeToOrderByContinuationTokenMap,
sortOrders, clientSideRequestStatisticsList));
}
private static class PageToItemTransformer implements
Function.DocumentProducerFeedResponse>, Flux>> {
private final RequestChargeTracker tracker;
private final Class klass;
private final Map queryMetricsMap;
private final Map targetRangeToOrderByContinuationTokenMap;
private final List sortOrders;
private final List clientSideRequestStatisticsList;
public PageToItemTransformer(
Class klass, RequestChargeTracker tracker, Map queryMetricsMap,
Map targetRangeToOrderByContinuationTokenMap,
List sortOrders, List clientSideRequestStatisticsList) {
this.klass = klass;
this.tracker = tracker;
this.queryMetricsMap = queryMetricsMap;
this.targetRangeToOrderByContinuationTokenMap = targetRangeToOrderByContinuationTokenMap;
this.sortOrders = sortOrders;
this.clientSideRequestStatisticsList = clientSideRequestStatisticsList;
}
@Override
public Flux> apply(Flux.DocumentProducerFeedResponse> source) {
return source.flatMap(documentProducerFeedResponse -> {
clientSideRequestStatisticsList.addAll(
BridgeInternal.getClientSideRequestStatisticsList(documentProducerFeedResponse
.pageResult.getCosmosDiagnostics()));
for (String key : BridgeInternal.queryMetricsFromFeedResponse(documentProducerFeedResponse.pageResult)
.keySet()) {
if (queryMetricsMap.containsKey(key)) {
QueryMetrics qm = BridgeInternal.queryMetricsFromFeedResponse(documentProducerFeedResponse.pageResult).get(key);
queryMetricsMap.get(key).add(qm);
} else {
queryMetricsMap.put(key, BridgeInternal.queryMetricsFromFeedResponse(documentProducerFeedResponse.pageResult).get(key));
}
}
List results = documentProducerFeedResponse.pageResult.getResults();
OrderByContinuationToken orderByContinuationToken =
targetRangeToOrderByContinuationTokenMap.get(documentProducerFeedResponse.sourceFeedRange);
if (orderByContinuationToken != null) {
Pair booleanResourceIdPair = ResourceId.tryParse(orderByContinuationToken.getRid());
if (!booleanResourceIdPair.getLeft()) {
return Flux.error(new BadRequestException(String.format("INVALID Rid in the continuation token %s for OrderBy~Context.",
orderByContinuationToken.getCompositeContinuationToken().getToken())));
}
ResourceId continuationTokenRid = booleanResourceIdPair.getRight();
results = results.stream()
.filter(tOrderByRowResult -> {
// 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.
List queryItems = new ArrayList();
ArrayNode arrayNode = (ArrayNode)ModelBridgeInternal.getObjectFromJsonSerializable(tOrderByRowResult, "orderByItems");
for (JsonNode jsonNode : arrayNode) {
QueryItem queryItem = new QueryItem(jsonNode.toString());
queryItems.add(queryItem);
}
// Check if its the same orderby item from the token
long cmp = 0;
for (int i = 0; i < sortOrders.size(); i++) {
cmp = ItemComparator.getInstance().compare(orderByContinuationToken.getOrderByItems()[i].getItem(),
queryItems.get(i).getItem());
if (cmp != 0) {
cmp = sortOrders.get(i).equals(SortOrder.Descending) ? -cmp : cmp;
break;
}
}
if (cmp == 0) {
// Once the item matches the order by items from the continuation tokens
// We still need to remove all the documents that have a lower rid in the rid sort order.
// If there is a tie in the sort order the documents should be in _rid order in the same direction as the 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 = (continuationTokenRid.getDocument() - ResourceId.tryParse(tOrderByRowResult.getResourceId()).getRight().getDocument());
if (sortOrders.iterator().next().equals(SortOrder.Descending)) {
cmp = -cmp;
}
return (cmp <= 0);
}
return true;
})
.collect(Collectors.toList());
}
tracker.addCharge(documentProducerFeedResponse.pageResult.getRequestCharge());
Flux x = Flux.fromIterable(results);
return x.map(r -> new OrderByRowResult(
klass,
ModelBridgeInternal.toJsonFromJsonSerializable(r),
documentProducerFeedResponse.sourceFeedRange,
documentProducerFeedResponse.pageResult.getContinuationToken()));
}, 1);
}
}
}