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

Example 6 with Selection

use of com.linkedin.pinot.common.request.Selection in project pinot by linkedin.

the class CombineService method mergeTwoBlocks.

public static void mergeTwoBlocks(@Nonnull BrokerRequest brokerRequest, @Nonnull IntermediateResultsBlock mergedBlock, @Nonnull IntermediateResultsBlock blockToMerge) {
    // Combine processing exceptions.
    List<ProcessingException> mergedProcessingExceptions = mergedBlock.getProcessingExceptions();
    List<ProcessingException> processingExceptionsToMerge = blockToMerge.getProcessingExceptions();
    if (mergedProcessingExceptions == null) {
        mergedBlock.setProcessingExceptions(processingExceptionsToMerge);
    } else if (processingExceptionsToMerge != null) {
        mergedProcessingExceptions.addAll(processingExceptionsToMerge);
    }
    // Combine result.
    if (brokerRequest.isSetAggregationsInfo()) {
        if (!brokerRequest.isSetGroupBy()) {
            // Combine aggregation only result.
            // Might be null if caught exception during query execution.
            List<Object> aggregationResultToMerge = blockToMerge.getAggregationResult();
            if (aggregationResultToMerge == null) {
                // No data in block to merge.
                return;
            }
            AggregationFunctionContext[] mergedAggregationFunctionContexts = mergedBlock.getAggregationFunctionContexts();
            if (mergedAggregationFunctionContexts == null) {
                // No data in merged block.
                mergedBlock.setAggregationFunctionContexts(blockToMerge.getAggregationFunctionContexts());
                mergedBlock.setAggregationResults(aggregationResultToMerge);
            }
            // Merge two block.
            List<Object> mergedAggregationResult = mergedBlock.getAggregationResult();
            int numAggregationFunctions = mergedAggregationFunctionContexts.length;
            for (int i = 0; i < numAggregationFunctions; i++) {
                mergedAggregationResult.set(i, mergedAggregationFunctionContexts[i].getAggregationFunction().merge(mergedAggregationResult.get(i), aggregationResultToMerge.get(i)));
            }
        } else {
            // Combine aggregation group-by result, which should not come into CombineService.
            throw new UnsupportedOperationException();
        }
    } else {
        // Combine selection result.
        // Data schema will be null if exceptions caught during query processing.
        // Result set size will be zero if no row matches the predicate.
        DataSchema mergedBlockSchema = mergedBlock.getSelectionDataSchema();
        DataSchema blockToMergeSchema = blockToMerge.getSelectionDataSchema();
        Collection<Serializable[]> mergedBlockResultSet = mergedBlock.getSelectionResult();
        Collection<Serializable[]> blockToMergeResultSet = blockToMerge.getSelectionResult();
        if (mergedBlockSchema == null || mergedBlockResultSet.size() == 0) {
            // If block to merge schema is not null, set its data schema and result to the merged block.
            if (blockToMergeSchema != null) {
                mergedBlock.setSelectionDataSchema(blockToMergeSchema);
                mergedBlock.setSelectionResult(blockToMergeResultSet);
            }
        } else {
            // Some data in merged block.
            Selection selection = brokerRequest.getSelections();
            boolean isSelectionOrderBy = selection.isSetSelectionSortSequence();
            int selectionSize = selection.getSize();
            // No need to merge if already got enough rows for selection only.
            if (!isSelectionOrderBy && (mergedBlockResultSet.size() == selectionSize)) {
                return;
            }
            // Merge only if there are data in block to merge.
            if (blockToMergeSchema != null && blockToMergeResultSet.size() > 0) {
                if (mergedBlockSchema.isTypeCompatibleWith(blockToMergeSchema)) {
                    // Two blocks are mergeable.
                    // Upgrade the merged block schema if necessary.
                    mergedBlockSchema.upgradeToCover(blockToMergeSchema);
                    // Merge two blocks.
                    if (isSelectionOrderBy) {
                        // Combine selection order-by.
                        SelectionOperatorUtils.mergeWithOrdering((PriorityQueue<Serializable[]>) mergedBlockResultSet, blockToMergeResultSet, selection.getOffset() + selectionSize);
                    } else {
                        // Combine selection only.
                        SelectionOperatorUtils.mergeWithoutOrdering(mergedBlockResultSet, blockToMergeResultSet, selectionSize);
                    }
                    mergedBlock.setSelectionResult(mergedBlockResultSet);
                } else {
                    // Two blocks are not mergeable.
                    throw new RuntimeException("Data schema inconsistency between merged block schema: " + mergedBlockSchema + " and block to merge schema: " + blockToMergeSchema + ", drop block to merge.");
                }
            }
        }
    }
}

Example 7 with Selection

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