Examples with AggregationGroupByTrimmingService com.linkedin.pinot.core.query.aggregation.groupby.AggregationGroupByTrimmingService used on opensource projects

Example 1 with AggregationGroupByTrimmingService

use of com.linkedin.pinot.core.query.aggregation.groupby.AggregationGroupByTrimmingService in project pinot by linkedin.

the class MCombineGroupByOperator method combineBlocks.

/**
   * This method combines the result blocks from underlying operators and builds a
   * merged, sorted and trimmed result block.
   * 1. Result blocks from underlying operators are merged concurrently into a
   *   HashMap, with appropriate synchronizations. Result blocks themselves are stored
   *   in the specified blocks[].
   *   - The key in this concurrent map is the group-by key, and value is an array of
   *     Objects (one for each aggregation function).
   *   - Synchronization is provided by locking the group-key that is to be modified.
   *
   * 2. The result of the concurrent map is then translated into what is expected by
   *    the broker (List<Map<String, Object>>).
   *
   * 3. This result is then sorted and then trimmed as per 'TOP N' in the brokerRequest.
   *
   * @return IntermediateResultBlock containing the final results from combine operation.
   */
private IntermediateResultsBlock combineBlocks() throws InterruptedException {
    int numOperators = _operators.size();
    final CountDownLatch operatorLatch = new CountDownLatch(numOperators);
    final Map<String, Object[]> resultsMap = new ConcurrentHashMap<>();
    final ConcurrentLinkedQueue<ProcessingException> mergedProcessingExceptions = new ConcurrentLinkedQueue<>();
    List<AggregationInfo> aggregationInfos = _brokerRequest.getAggregationsInfo();
    final AggregationFunctionContext[] aggregationFunctionContexts = AggregationFunctionUtils.getAggregationFunctionContexts(aggregationInfos, null);
    final int numAggregationFunctions = aggregationFunctionContexts.length;
    for (int i = 0; i < numOperators; i++) {
        final int index = i;
        _executorService.execute(new TraceRunnable() {

            @SuppressWarnings("unchecked")
            @Override
            public void runJob() {
                AggregationGroupByResult aggregationGroupByResult;
                try {
                    IntermediateResultsBlock intermediateResultsBlock = (IntermediateResultsBlock) _operators.get(index).nextBlock();
                    // Merge processing exceptions.
                    List<ProcessingException> processingExceptionsToMerge = intermediateResultsBlock.getProcessingExceptions();
                    if (processingExceptionsToMerge != null) {
                        mergedProcessingExceptions.addAll(processingExceptionsToMerge);
                    }
                    // Merge aggregation group-by result.
                    aggregationGroupByResult = intermediateResultsBlock.getAggregationGroupByResult();
                    if (aggregationGroupByResult != null) {
                        // Iterate over the group-by keys, for each key, update the group-by result in the resultsMap.
                        Iterator<GroupKeyGenerator.GroupKey> groupKeyIterator = aggregationGroupByResult.getGroupKeyIterator();
                        while (groupKeyIterator.hasNext()) {
                            GroupKeyGenerator.GroupKey groupKey = groupKeyIterator.next();
                            String groupKeyString = groupKey.getStringKey();
                            // HashCode method might return negative value, make it non-negative
                            int lockIndex = (groupKeyString.hashCode() & Integer.MAX_VALUE) % NUM_LOCKS;
                            synchronized (LOCKS[lockIndex]) {
                                Object[] results = resultsMap.get(groupKeyString);
                                if (results == null) {
                                    results = new Object[numAggregationFunctions];
                                    for (int j = 0; j < numAggregationFunctions; j++) {
                                        results[j] = aggregationGroupByResult.getResultForKey(groupKey, j);
                                    }
                                    resultsMap.put(groupKeyString, results);
                                } else {
                                    for (int j = 0; j < numAggregationFunctions; j++) {
                                        results[j] = aggregationFunctionContexts[j].getAggregationFunction().merge(results[j], aggregationGroupByResult.getResultForKey(groupKey, j));
                                    }
                                }
                            }
                        }
                    }
                } catch (Exception e) {
                    LOGGER.error("Exception processing CombineGroupBy for index {}, operator {}", index, _operators.get(index).getClass().getName(), e);
                    mergedProcessingExceptions.add(QueryException.getException(QueryException.QUERY_EXECUTION_ERROR, e));
                }
                operatorLatch.countDown();
            }
        });
    }
    boolean opCompleted = operatorLatch.await(_timeOutMs, TimeUnit.MILLISECONDS);
    if (!opCompleted) {
        // If this happens, the broker side should already timed out, just log the error in server side.
        LOGGER.error("Timed out while combining group-by results, after {}ms.", _timeOutMs);
        return new IntermediateResultsBlock(new TimeoutException("CombineGroupBy timed out."));
    }
    // Trim the results map.
    AggregationGroupByTrimmingService aggregationGroupByTrimmingService = new AggregationGroupByTrimmingService(aggregationFunctionContexts, (int) _brokerRequest.getGroupBy().getTopN());
    List<Map<String, Object>> trimmedResults = aggregationGroupByTrimmingService.trimIntermediateResultsMap(resultsMap);
    IntermediateResultsBlock mergedBlock = new IntermediateResultsBlock(aggregationFunctionContexts, trimmedResults, true);
    // Set the processing exceptions.
    if (!mergedProcessingExceptions.isEmpty()) {
        mergedBlock.setProcessingExceptions(new ArrayList<>(mergedProcessingExceptions));
    }
    // Set the execution statistics.
    ExecutionStatistics executionStatistics = new ExecutionStatistics();
    for (Operator operator : _operators) {
        ExecutionStatistics executionStatisticsToMerge = operator.getExecutionStatistics();
        if (executionStatisticsToMerge != null) {
            executionStatistics.merge(executionStatisticsToMerge);
        }
    }
    mergedBlock.setNumDocsScanned(executionStatistics.getNumDocsScanned());
    mergedBlock.setNumEntriesScannedInFilter(executionStatistics.getNumEntriesScannedInFilter());
    mergedBlock.setNumEntriesScannedPostFilter(executionStatistics.getNumEntriesScannedPostFilter());
    mergedBlock.setNumTotalRawDocs(executionStatistics.getNumTotalRawDocs());
    return mergedBlock;
}

Example 2 with AggregationGroupByTrimmingService

use of com.linkedin.pinot.core.query.aggregation.groupby.AggregationGroupByTrimmingService in project pinot by linkedin.

the class BrokerReduceService method setGroupByResults.

/**
   * Reduce group-by results from multiple servers and set them into BrokerResponseNative passed in.
   *
   * @param brokerResponseNative broker response.
   * @param aggregationFunctions array of aggregation functions.
   * @param groupBy group-by information.
   * @param dataTableMap map from server to data table.
   */
@SuppressWarnings("unchecked")
private void setGroupByResults(@Nonnull BrokerResponseNative brokerResponseNative, @Nonnull AggregationFunction[] aggregationFunctions, @Nonnull GroupBy groupBy, @Nonnull Map<ServerInstance, DataTable> dataTableMap) {
    int numAggregationFunctions = aggregationFunctions.length;
    // Merge results from all data tables.
    String[] columnNames = new String[numAggregationFunctions];
    Map<String, Object>[] intermediateResultMaps = new Map[numAggregationFunctions];
    for (DataTable dataTable : dataTableMap.values()) {
        for (int i = 0; i < numAggregationFunctions; i++) {
            if (columnNames[i] == null) {
                columnNames[i] = dataTable.getString(i, 0);
                intermediateResultMaps[i] = dataTable.getObject(i, 1);
            } else {
                Map<String, Object> mergedIntermediateResultMap = intermediateResultMaps[i];
                Map<String, Object> intermediateResultMapToMerge = dataTable.getObject(i, 1);
                for (Map.Entry<String, Object> entry : intermediateResultMapToMerge.entrySet()) {
                    String groupKey = entry.getKey();
                    Object intermediateResultToMerge = entry.getValue();
                    if (mergedIntermediateResultMap.containsKey(groupKey)) {
                        Object mergedIntermediateResult = mergedIntermediateResultMap.get(groupKey);
                        mergedIntermediateResultMap.put(groupKey, aggregationFunctions[i].merge(mergedIntermediateResult, intermediateResultToMerge));
                    } else {
                        mergedIntermediateResultMap.put(groupKey, intermediateResultToMerge);
                    }
                }
            }
        }
    }
    // Extract final result maps from the merged intermediate result maps.
    Map<String, Comparable>[] finalResultMaps = new Map[numAggregationFunctions];
    for (int i = 0; i < numAggregationFunctions; i++) {
        Map<String, Object> intermediateResultMap = intermediateResultMaps[i];
        Map<String, Comparable> finalResultMap = new HashMap<>();
        for (String groupKey : intermediateResultMap.keySet()) {
            Object intermediateResult = intermediateResultMap.get(groupKey);
            finalResultMap.put(groupKey, aggregationFunctions[i].extractFinalResult(intermediateResult));
        }
        finalResultMaps[i] = finalResultMap;
    }
    // Trim the final result maps to topN and set them into the broker response.
    AggregationGroupByTrimmingService aggregationGroupByTrimmingService = new AggregationGroupByTrimmingService(aggregationFunctions, (int) groupBy.getTopN());
    List<GroupByResult>[] groupByResultLists = aggregationGroupByTrimmingService.trimFinalResults(finalResultMaps);
    List<AggregationResult> aggregationResults = new ArrayList<>(numAggregationFunctions);
    for (int i = 0; i < numAggregationFunctions; i++) {
        List<GroupByResult> groupByResultList = groupByResultLists[i];
        List<String> groupByColumns = groupBy.getExpressions();
        if (groupByColumns == null) {
            groupByColumns = groupBy.getColumns();
        }
        aggregationResults.add(new AggregationResult(groupByResultList, groupByColumns, columnNames[i]));
    }
    brokerResponseNative.setAggregationResults(aggregationResults);
}

Example 3 with AggregationGroupByTrimmingService

use of com.linkedin.pinot.core.query.aggregation.groupby.AggregationGroupByTrimmingService in project pinot by linkedin.

the class AggregationGroupByTrimmingServiceTest method setUp.

@BeforeClass
public void setUp() {
    // Generate a list of random groups.
    Set<String> groupSet = new HashSet<>(NUM_GROUPS);
    while (groupSet.size() < NUM_GROUPS) {
        String group = "";
        for (int i = 0; i < NUM_GROUP_KEYS; i++) {
            if (i != 0) {
                group += '\t';
            }
            // Random generate group key without '\t'.
            String groupKey = RandomStringUtils.random(RANDOM.nextInt(10));
            while (groupKey.contains("\t")) {
                groupKey = RandomStringUtils.random(RANDOM.nextInt(10));
            }
            group += groupKey;
        }
        groupSet.add(group);
    }
    _groups = new ArrayList<>(groupSet);
    // Explicitly set an empty group.
    String emptyGroup = "";
    for (int i = 1; i < NUM_GROUP_KEYS; i++) {
        emptyGroup += '\t';
    }
    _groups.set(NUM_GROUPS - 1, emptyGroup);
    _serverTrimmingService = new AggregationGroupByTrimmingService(AGGREGATION_FUNCTION_CONTEXTS, GROUP_BY_TOP_N);
    _brokerTrimmingService = new AggregationGroupByTrimmingService(AGGREGATION_FUNCTIONS, GROUP_BY_TOP_N);
}