Examples with AggregationInfo com.linkedin.pinot.common.request.AggregationInfo used on opensource projects

Example 51 with AggregationInfo

use of com.linkedin.pinot.common.request.AggregationInfo 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();
    }
}