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

Example 6 with AggregationFunctionContext

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

the class IntermediateResultsBlock method getAggregationGroupByResultDataTable.

@Nonnull
private DataTable getAggregationGroupByResultDataTable() throws Exception {
    String[] columnNames = new String[] { "functionName", "GroupByResultMap" };
    FieldSpec.DataType[] columnTypes = new FieldSpec.DataType[] { FieldSpec.DataType.STRING, FieldSpec.DataType.OBJECT };
    // Build the data table.
    DataTableBuilder dataTableBuilder = new DataTableBuilder(new DataSchema(columnNames, columnTypes));
    int numAggregationFunctions = _aggregationFunctionContexts.length;
    for (int i = 0; i < numAggregationFunctions; i++) {
        dataTableBuilder.startRow();
        AggregationFunctionContext aggregationFunctionContext = _aggregationFunctionContexts[i];
        dataTableBuilder.setColumn(0, aggregationFunctionContext.getAggregationColumnName());
        dataTableBuilder.setColumn(1, _combinedAggregationGroupByResult.get(i));
        dataTableBuilder.finishRow();
    }
    DataTable dataTable = dataTableBuilder.build();
    return attachMetadataToDataTable(dataTable);
}

Example 7 with AggregationFunctionContext

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

the class TransformGroupByTest method executeGroupByQuery.

/**
   * Helper method that executes the group by query on the index and returns the group by result.
   *
   * @param query Query to execute
   * @return Group by result
   */
private AggregationGroupByResult executeGroupByQuery(IndexSegment indexSegment, String query) {
    Operator filterOperator = new MatchEntireSegmentOperator(indexSegment.getSegmentMetadata().getTotalDocs());
    final BReusableFilteredDocIdSetOperator docIdSetOperator = new BReusableFilteredDocIdSetOperator(filterOperator, indexSegment.getSegmentMetadata().getTotalDocs(), NUM_ROWS);
    final Map<String, BaseOperator> dataSourceMap = buildDataSourceMap(indexSegment.getSegmentMetadata().getSchema());
    final MProjectionOperator projectionOperator = new MProjectionOperator(dataSourceMap, docIdSetOperator);
    Pql2Compiler compiler = new Pql2Compiler();
    BrokerRequest brokerRequest = compiler.compileToBrokerRequest(query);
    List<AggregationInfo> aggregationsInfo = brokerRequest.getAggregationsInfo();
    int numAggFunctions = aggregationsInfo.size();
    AggregationFunctionContext[] aggrFuncContextArray = new AggregationFunctionContext[numAggFunctions];
    AggregationFunctionInitializer aggFuncInitializer = new AggregationFunctionInitializer(indexSegment.getSegmentMetadata());
    for (int i = 0; i < numAggFunctions; i++) {
        AggregationInfo aggregationInfo = aggregationsInfo.get(i);
        aggrFuncContextArray[i] = AggregationFunctionContext.instantiate(aggregationInfo);
        aggrFuncContextArray[i].getAggregationFunction().accept(aggFuncInitializer);
    }
    GroupBy groupBy = brokerRequest.getGroupBy();
    Set<String> expressions = new HashSet<>(groupBy.getExpressions());
    TransformExpressionOperator transformOperator = new TransformExpressionOperator(projectionOperator, TransformPlanNode.buildTransformExpressionTrees(expressions));
    AggregationGroupByOperator groupByOperator = new AggregationGroupByOperator(aggrFuncContextArray, groupBy, Integer.MAX_VALUE, transformOperator, NUM_ROWS);
    IntermediateResultsBlock block = (IntermediateResultsBlock) groupByOperator.nextBlock();
    return block.getAggregationGroupByResult();
}

Example 8 with AggregationFunctionContext

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

the class DataTableBuilder method buildEmptyDataTable.

/**
   * Build an empty data table based on the broker request.
   */
public static DataTable buildEmptyDataTable(BrokerRequest brokerRequest) throws IOException {
    // Selection query.
    if (brokerRequest.isSetSelections()) {
        Selection selection = brokerRequest.getSelections();
        List<String> selectionColumns = selection.getSelectionColumns();
        int numSelectionColumns = selectionColumns.size();
        FieldSpec.DataType[] dataTypes = new FieldSpec.DataType[numSelectionColumns];
        // Use STRING data type as default for selection query.
        Arrays.fill(dataTypes, FieldSpec.DataType.STRING);
        DataSchema dataSchema = new DataSchema(selectionColumns.toArray(new String[numSelectionColumns]), dataTypes);
        return new DataTableBuilder(dataSchema).build();
    }
    // Aggregation query.
    List<AggregationInfo> aggregationsInfo = brokerRequest.getAggregationsInfo();
    int numAggregations = aggregationsInfo.size();
    AggregationFunctionContext[] aggregationFunctionContexts = new AggregationFunctionContext[numAggregations];
    for (int i = 0; i < numAggregations; i++) {
        aggregationFunctionContexts[i] = AggregationFunctionContext.instantiate(aggregationsInfo.get(i));
    }
    if (brokerRequest.isSetGroupBy()) {
        // Aggregation group-by query.
        String[] columnNames = new String[] { "functionName", "GroupByResultMap" };
        FieldSpec.DataType[] columnTypes = new FieldSpec.DataType[] { FieldSpec.DataType.STRING, FieldSpec.DataType.OBJECT };
        // Build the data table.
        DataTableBuilder dataTableBuilder = new DataTableBuilder(new DataSchema(columnNames, columnTypes));
        for (int i = 0; i < numAggregations; i++) {
            dataTableBuilder.startRow();
            dataTableBuilder.setColumn(0, aggregationFunctionContexts[i].getAggregationColumnName());
            dataTableBuilder.setColumn(1, new HashMap<String, Object>());
            dataTableBuilder.finishRow();
        }
        return dataTableBuilder.build();
    } else {
        // Aggregation only query.
        String[] aggregationColumnNames = new String[numAggregations];
        FieldSpec.DataType[] dataTypes = new FieldSpec.DataType[numAggregations];
        Object[] aggregationResults = new Object[numAggregations];
        for (int i = 0; i < numAggregations; i++) {
            AggregationFunctionContext aggregationFunctionContext = aggregationFunctionContexts[i];
            aggregationColumnNames[i] = aggregationFunctionContext.getAggregationColumnName();
            AggregationFunction aggregationFunction = aggregationFunctionContext.getAggregationFunction();
            dataTypes[i] = aggregationFunction.getIntermediateResultDataType();
            aggregationResults[i] = aggregationFunction.extractAggregationResult(aggregationFunction.createAggregationResultHolder());
        }
        // Build the data table.
        DataTableBuilder dataTableBuilder = new DataTableBuilder(new DataSchema(aggregationColumnNames, dataTypes));
        dataTableBuilder.startRow();
        for (int i = 0; i < numAggregations; i++) {
            switch(dataTypes[i]) {
                case LONG:
                    dataTableBuilder.setColumn(i, ((Number) aggregationResults[i]).longValue());
                    break;
                case DOUBLE:
                    dataTableBuilder.setColumn(i, ((Double) aggregationResults[i]).doubleValue());
                    break;
                case OBJECT:
                    dataTableBuilder.setColumn(i, aggregationResults[i]);
                    break;
                default:
                    throw new UnsupportedOperationException("Unsupported aggregation column data type: " + dataTypes[i] + " for column: " + aggregationColumnNames[i]);
            }
        }
        dataTableBuilder.finishRow();
        return dataTableBuilder.build();
    }
}