com.azure.cosmos.implementation.query.AggregateDocumentQueryExecutionContext Maven / Gradle / Ivy

Go to download
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.

package com.azure.cosmos.implementation.query;

import com.azure.cosmos.implementation.query.aggregation.AggregateOperator;
import com.azure.cosmos.implementation.query.aggregation.Aggregator;
import com.azure.cosmos.implementation.query.aggregation.AverageAggregator;
import com.azure.cosmos.implementation.query.aggregation.CountAggregator;
import com.azure.cosmos.implementation.query.aggregation.MaxAggregator;
import com.azure.cosmos.implementation.query.aggregation.MinAggregator;
import com.azure.cosmos.implementation.query.aggregation.SumAggregator;
import com.azure.cosmos.BridgeInternal;
import com.azure.cosmos.implementation.Document;
import com.azure.cosmos.FeedResponse;
import com.azure.cosmos.Resource;
import com.azure.cosmos.implementation.Undefined;
import com.azure.cosmos.implementation.Constants;
import com.azure.cosmos.implementation.HttpConstants;
import com.azure.cosmos.implementation.QueryMetrics;
import reactor.core.publisher.Flux;

import java.util.ArrayList;
import java.util.Collection;
import java.util.HashMap;
import java.util.List;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
import java.util.function.Function;

public class AggregateDocumentQueryExecutionContext implements IDocumentQueryExecutionComponent{

    private IDocumentQueryExecutionComponent component;
    private Aggregator aggregator;
    private ConcurrentMap queryMetricsMap = new ConcurrentHashMap<>();

    //QueryInfo class used in PipelinedDocumentQueryExecutionContext returns a Collection of AggregateOperators
    //while Multiple aggregates are allowed in queries targeted at a single partition, only a single aggregate is allowed in x-partition queries (currently)
    public AggregateDocumentQueryExecutionContext (IDocumentQueryExecutionComponent component, Collection aggregateOperators) {

        this.component = component;
        AggregateOperator aggregateOperator = aggregateOperators.iterator().next();
        
        switch (aggregateOperator) {
            case Average:
                this.aggregator = new AverageAggregator();
                break;
            case Count:
                this.aggregator = new CountAggregator();
                break;
            case Max:
                this.aggregator = new MaxAggregator();
                break;
            case Min:
                this.aggregator = new MinAggregator();
                break;
            case Sum:
                this.aggregator = new SumAggregator();
                break;
            default:
                throw new IllegalStateException("Unexpected value: " + aggregateOperator.toString());
            }
        }

    @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();
                    
                    for(FeedResponse page : superList) {
                        
                        if (page.getResults().size() == 0) {
                            headers.put(HttpConstants.HttpHeaders.REQUEST_CHARGE, Double.toString(requestCharge));
                            FeedResponse frp = BridgeInternal.createFeedResponse(aggregateResults, headers);
                            return (FeedResponse) frp;
                        }
                        
                        Document doc = ((Document)page.getResults().get(0));
                        requestCharge += page.getRequestCharge();
                        QueryItem values = new QueryItem(doc.toJson());
                        this.aggregator.aggregate(values.getItem());
                        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));
                            }
                        }
                    }
                    
                    if (this.aggregator.getResult() == null || !this.aggregator.getResult().equals(Undefined.Value())) {
                        Document aggregateDocument = new Document();
                        BridgeInternal.setProperty(aggregateDocument, Constants.Properties.VALUE, this.aggregator.getResult());
                        aggregateResults.add(aggregateDocument);
                    }

                    headers.put(HttpConstants.HttpHeaders.REQUEST_CHARGE, Double.toString(requestCharge));
                    FeedResponse frp = BridgeInternal.createFeedResponse(aggregateResults, headers);
                    if(!queryMetricsMap.isEmpty()) {
                        for(String key: queryMetricsMap.keySet()) {
                            BridgeInternal.putQueryMetricsIntoMap(frp, key, queryMetricsMap.get(key));
                        }
                    }
                    return (FeedResponse) frp;
                }).flux();
    }

    public static   Flux> createAsync(
            Function>> createSourceComponentFunction,
            Collection aggregates,
            String continuationToken) {

        return createSourceComponentFunction
                .apply(continuationToken)
                .map( component -> { return new AggregateDocumentQueryExecutionContext(component, aggregates);});
    }

    public IDocumentQueryExecutionComponent getComponent() {
        return this.component;
    }

}