com.azure.cosmos.implementation.query.AggregateDocumentQueryExecutionContext 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.BridgeInternal;
import com.azure.cosmos.implementation.ClientSideRequestStatistics;
import com.azure.cosmos.implementation.Document;
import com.azure.cosmos.implementation.HttpConstants;
import com.azure.cosmos.implementation.QueryMetrics;
import com.azure.cosmos.implementation.Resource;
import com.azure.cosmos.implementation.query.aggregation.AggregateOperator;
import com.azure.cosmos.models.FeedResponse;
import com.fasterxml.jackson.databind.node.ObjectNode;
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.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
import java.util.function.BiFunction;
public class AggregateDocumentQueryExecutionContext implements IDocumentQueryExecutionComponent{
public static final String PAYLOAD_PROPERTY_NAME = "payload";
private final boolean isValueAggregateQuery;
private IDocumentQueryExecutionComponent component;
private ConcurrentMap queryMetricsMap = new ConcurrentHashMap<>();
private SingleGroupAggregator singleGroupAggregator;
//QueryInfo class used in PipelinedDocumentQueryExecutionContext returns a Collection of AggregateOperators
public AggregateDocumentQueryExecutionContext(IDocumentQueryExecutionComponent component,
List aggregateOperators,
Map groupByAliasToAggregateType,
List orderedAliases,
boolean hasSelectValue,
String continuationToken) {
this.component = component;
this.isValueAggregateQuery = hasSelectValue;
this.singleGroupAggregator = SingleGroupAggregator.create(aggregateOperators,
groupByAliasToAggregateType,
orderedAliases,
hasSelectValue,
continuationToken);
}
@SuppressWarnings("unchecked")
@Override
public Flux> drainAsync(int maxPageSize) {
return this.component.drainAsync(maxPageSize)
.collectList()
.map( superList -> {
double requestCharge = 0;
List aggregateResults = new ArrayList<>();
HashMap headers = new HashMap<>();
List diagnosticsList = new ArrayList<>();
for(FeedResponse page : superList) {
diagnosticsList.addAll(BridgeInternal
.getClientSideRequestStatisticsList(page.getCosmosDiagnostics()));
if (page.getResults().size() == 0) {
headers.put(HttpConstants.HttpHeaders.REQUEST_CHARGE, Double.toString(requestCharge));
FeedResponse frp = BridgeInternal.createFeedResponse(aggregateResults, headers);
BridgeInternal.addClientSideDiagnosticsToFeed(frp.getCosmosDiagnostics(), diagnosticsList);
return (FeedResponse) frp;
}
requestCharge += page.getRequestCharge();
for (T d : page.getResults()) {
RewrittenAggregateProjections rewrittenAggregateProjections =
new RewrittenAggregateProjections(this.isValueAggregateQuery,
(Document)d); //d is always a Document
this.singleGroupAggregator.addValues(rewrittenAggregateProjections.getPayload());
}
for(String key : BridgeInternal.queryMetricsFromFeedResponse(page).keySet()) {
if (queryMetricsMap.containsKey(key)) {
QueryMetrics qm = BridgeInternal.queryMetricsFromFeedResponse(page).get(key);
queryMetricsMap.get(key).add(qm);
} else {
queryMetricsMap.put(key, BridgeInternal.queryMetricsFromFeedResponse(page).get(key));
}
}
}
Document aggregateDocument = this.singleGroupAggregator.getResult();
if (aggregateDocument != null) {
aggregateResults.add(aggregateDocument);
}
headers.put(HttpConstants.HttpHeaders.REQUEST_CHARGE, Double.toString(requestCharge));
FeedResponse frp = BridgeInternal.createFeedResponse(aggregateResults, headers);
if(!queryMetricsMap.isEmpty()) {
for(Map.Entry entry: queryMetricsMap.entrySet()) {
BridgeInternal.putQueryMetricsIntoMap(frp, entry.getKey(), entry.getValue());
}
}
BridgeInternal.addClientSideDiagnosticsToFeed(frp.getCosmosDiagnostics(), diagnosticsList);
return (FeedResponse) frp;
}).flux();
}
public static Flux> createAsync(
BiFunction, Flux>> createSourceComponentFunction,
Collection aggregates,
Map groupByAliasToAggregateType,
List groupByAliases,
boolean hasSelectValue,
String continuationToken,
PipelinedDocumentQueryParams documentQueryParams) {
return createSourceComponentFunction
.apply(continuationToken, documentQueryParams)
.map(component -> new AggregateDocumentQueryExecutionContext(component,
new ArrayList<>(aggregates),
groupByAliasToAggregateType,
groupByAliases,
hasSelectValue,
continuationToken));
}
public IDocumentQueryExecutionComponent getComponent() {
return this.component;
}
class RewrittenAggregateProjections {
private Document payload;
public RewrittenAggregateProjections(boolean isValueAggregateQuery, Document document) {
if (document == null) {
throw new IllegalArgumentException("document cannot be null");
}
if (isValueAggregateQuery) {
this.payload = new Document(document.getPropertyBag());
} else {
if (!document.has(PAYLOAD_PROPERTY_NAME)) {
throw new IllegalStateException("Underlying object does not have an 'payload' field.");
}
if (document.get(PAYLOAD_PROPERTY_NAME) instanceof ObjectNode) {
this.payload = new Document((ObjectNode) document.get(PAYLOAD_PROPERTY_NAME));
}
}
}
public Document getPayload() {
return payload;
}
}
}