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

Show more of this group Show more artifacts with this name
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

There is a newer version: 4.61.1

// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
package com.azure.cosmos.implementation.query;

import com.azure.cosmos.implementation.Document;
import com.azure.cosmos.implementation.query.aggregation.AggregateOperator;
import com.azure.cosmos.implementation.routing.UInt128;

import java.io.IOException;
import java.util.ArrayList;
import java.util.Collection;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;

public class GroupingTable {
    private static final List EMPTY_AGGREGATE_OPERATORS = new ArrayList<>();

    private final Map table;
    private final Map groupByAliasToAggregateType;
    private final List orderedAliases;
    private final boolean hasSelectValue;

    GroupingTable(Map groupByAliasToAggregateType, List orderedAliases,
        boolean hasSelectValue) {
        if (groupByAliasToAggregateType == null) {
            throw new IllegalArgumentException("groupByAliasToAggregateType cannot be null");
        }
        this.table = new HashMap<>();
        this.groupByAliasToAggregateType = groupByAliasToAggregateType;
        this.orderedAliases = orderedAliases;
        this.hasSelectValue = hasSelectValue;
    }

    public void addPayLoad(GroupByDocumentQueryExecutionContext.RewrittenGroupByProjection rewrittenGroupByProjection) {
        try {
            final UInt128 groupByKeysHash = DistinctHash.getHash(rewrittenGroupByProjection.getGroupByItems());
            SingleGroupAggregator singleGroupAggregator;
            if (!this.table.containsKey(groupByKeysHash)) {
                singleGroupAggregator = SingleGroupAggregator.create(EMPTY_AGGREGATE_OPERATORS,
                                                                     this.groupByAliasToAggregateType,
                                                                     this.orderedAliases,
                                                                     this.hasSelectValue,
                                                                      /*continuationtoken*/ null);
                this.table.put(groupByKeysHash, singleGroupAggregator);
            } else {
                singleGroupAggregator = table.get(groupByKeysHash);
            }

            singleGroupAggregator.addValues(rewrittenGroupByProjection.getPayload());

        } catch (IOException e) {
            throw new IllegalStateException("Failed to add payload to groupby projection", e);
        }
    }

    public List drain(int maxItemCount) {
        Collection keys = this.table.keySet().stream().limit(maxItemCount).collect(Collectors.toList());
        List singleGroupAggregators = new ArrayList<>(keys.size());
        for (UInt128 key : keys) {
            singleGroupAggregators.add(this.table.get(key));
            this.table.remove(key);
        }
        List results = new ArrayList<>();
        for (SingleGroupAggregator singleGroupAggregator : singleGroupAggregators) {
            results.add(singleGroupAggregator.getResult());
        }

        return results;
    }
}