Examples with DataTableBuilder com.linkedin.pinot.core.common.datatable.DataTableBuilder used on opensource projects

Example 1 with DataTableBuilder

use of com.linkedin.pinot.core.common.datatable.DataTableBuilder in project pinot by linkedin.

the class IntermediateResultsBlock method getAggregationResultDataTable.

@Nonnull
private DataTable getAggregationResultDataTable() throws Exception {
    // Extract each aggregation column name and type from aggregation function context.
    int numAggregationFunctions = _aggregationFunctionContexts.length;
    String[] columnNames = new String[numAggregationFunctions];
    FieldSpec.DataType[] columnTypes = new FieldSpec.DataType[numAggregationFunctions];
    for (int i = 0; i < numAggregationFunctions; i++) {
        AggregationFunctionContext aggregationFunctionContext = _aggregationFunctionContexts[i];
        columnNames[i] = aggregationFunctionContext.getAggregationColumnName();
        columnTypes[i] = aggregationFunctionContext.getAggregationFunction().getIntermediateResultDataType();
    }
    // Build the data table.
    DataTableBuilder dataTableBuilder = new DataTableBuilder(new DataSchema(columnNames, columnTypes));
    dataTableBuilder.startRow();
    for (int i = 0; i < numAggregationFunctions; i++) {
        switch(columnTypes[i]) {
            case LONG:
                dataTableBuilder.setColumn(i, ((Number) _aggregationResult.get(i)).longValue());
                break;
            case DOUBLE:
                dataTableBuilder.setColumn(i, ((Double) _aggregationResult.get(i)).doubleValue());
                break;
            case OBJECT:
                dataTableBuilder.setColumn(i, _aggregationResult.get(i));
                break;
            default:
                throw new UnsupportedOperationException("Unsupported aggregation column data type: " + columnTypes[i] + " for column: " + columnNames[i]);
        }
    }
    dataTableBuilder.finishRow();
    DataTable dataTable = dataTableBuilder.build();
    return attachMetadataToDataTable(dataTable);
}

Example 2 with DataTableBuilder

use of com.linkedin.pinot.core.common.datatable.DataTableBuilder in project pinot by linkedin.

the class SelectionOperatorUtils method getDataTableFromRows.

/**
   * Build a {@link DataTable} from a {@link Collection} of selection rows with {@link DataSchema}. (Server side)
   * <p>The passed in data schema stored the column data type that can cover all actual data types for that column.
   * <p>The actual data types for each column in rows can be different but must be compatible with each other.
   * <p>Before write each row into the data table, first convert it to match the data types in data schema.
   *
   * @param rows {@link Collection} of selection rows.
   * @param dataSchema data schema.
   * @return data table.
   * @throws Exception
   */
@Nonnull
public static DataTable getDataTableFromRows(@Nonnull Collection<Serializable[]> rows, @Nonnull DataSchema dataSchema) throws Exception {
    int numColumns = dataSchema.size();
    DataTableBuilder dataTableBuilder = new DataTableBuilder(dataSchema);
    for (Serializable[] row : rows) {
        dataTableBuilder.startRow();
        for (int i = 0; i < numColumns; i++) {
            Serializable columnValue = row[i];
            DataType columnType = dataSchema.getColumnType(i);
            switch(columnType) {
                // Single-value column.
                case INT:
                    dataTableBuilder.setColumn(i, ((Number) columnValue).intValue());
                    break;
                case LONG:
                    dataTableBuilder.setColumn(i, ((Number) columnValue).longValue());
                    break;
                case FLOAT:
                    dataTableBuilder.setColumn(i, ((Number) columnValue).floatValue());
                    break;
                case DOUBLE:
                    dataTableBuilder.setColumn(i, ((Number) columnValue).doubleValue());
                    break;
                case STRING:
                    dataTableBuilder.setColumn(i, ((String) columnValue));
                    break;
                // Multi-value column.
                case INT_ARRAY:
                    dataTableBuilder.setColumn(i, (int[]) columnValue);
                    break;
                case LONG_ARRAY:
                    // LONG_ARRAY type covers INT_ARRAY and LONG_ARRAY.
                    if (columnValue instanceof int[]) {
                        int[] ints = (int[]) columnValue;
                        int length = ints.length;
                        long[] longs = new long[length];
                        for (int j = 0; j < length; j++) {
                            longs[j] = ints[j];
                        }
                        dataTableBuilder.setColumn(i, longs);
                    } else {
                        dataTableBuilder.setColumn(i, (long[]) columnValue);
                    }
                    break;
                case FLOAT_ARRAY:
                    dataTableBuilder.setColumn(i, (float[]) columnValue);
                    break;
                case DOUBLE_ARRAY:
                    // DOUBLE_ARRAY type covers INT_ARRAY, LONG_ARRAY, FLOAT_ARRAY and DOUBLE_ARRAY.
                    if (columnValue instanceof int[]) {
                        int[] ints = (int[]) columnValue;
                        int length = ints.length;
                        double[] doubles = new double[length];
                        for (int j = 0; j < length; j++) {
                            doubles[j] = ints[j];
                        }
                        dataTableBuilder.setColumn(i, doubles);
                    } else if (columnValue instanceof long[]) {
                        long[] longs = (long[]) columnValue;
                        int length = longs.length;
                        double[] doubles = new double[length];
                        for (int j = 0; j < length; j++) {
                            doubles[j] = longs[j];
                        }
                        dataTableBuilder.setColumn(i, doubles);
                    } else if (columnValue instanceof float[]) {
                        float[] floats = (float[]) columnValue;
                        int length = floats.length;
                        double[] doubles = new double[length];
                        for (int j = 0; j < length; j++) {
                            doubles[j] = floats[j];
                        }
                        dataTableBuilder.setColumn(i, doubles);
                    } else {
                        dataTableBuilder.setColumn(i, (double[]) columnValue);
                    }
                    break;
                case STRING_ARRAY:
                    dataTableBuilder.setColumn(i, (String[]) columnValue);
                    break;
                default:
                    throw new UnsupportedOperationException("Unsupported data type: " + columnType + " for column: " + dataSchema.getColumnName(i));
            }
        }
        dataTableBuilder.finishRow();
    }
    return dataTableBuilder.build();
}

Example 3 with DataTableBuilder

use of com.linkedin.pinot.core.common.datatable.DataTableBuilder 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 4 with DataTableBuilder

use of com.linkedin.pinot.core.common.datatable.DataTableBuilder in project pinot by linkedin.

the class ScheduledRequestHandlerTest method testValidQueryResponse.

@Test
public void testValidQueryResponse() throws InterruptedException, ExecutionException, TimeoutException, IOException {
    ScheduledRequestHandler handler = new ScheduledRequestHandler(new QueryScheduler(queryExecutor) {

        @Override
        public ListenableFuture<DataTable> submit(QueryRequest queryRequest) {
            return queryRunners.submit(new Callable<DataTable>() {

                @Override
                public DataTable call() throws Exception {
                    String[] columns = new String[] { "foo", "bar" };
                    FieldSpec.DataType[] columnTypes = new FieldSpec.DataType[] { FieldSpec.DataType.STRING, FieldSpec.DataType.INT };
                    DataSchema dataSchema = new DataSchema(columns, columnTypes);
                    DataTableBuilder dtBuilder = new DataTableBuilder(dataSchema);
                    dtBuilder.startRow();
                    dtBuilder.setColumn(0, "mars");
                    dtBuilder.setColumn(1, 10);
                    dtBuilder.finishRow();
                    dtBuilder.startRow();
                    dtBuilder.setColumn(0, "jupiter");
                    dtBuilder.setColumn(1, 100);
                    dtBuilder.finishRow();
                    return dtBuilder.build();
                }
            });
        }
    }, serverMetrics);
    ByteBuf requestBuf = getSerializedInstanceRequest(getInstanceRequest());
    ListenableFuture<byte[]> responseFuture = handler.processRequest(channelHandlerContext, requestBuf);
    byte[] responseBytes = responseFuture.get(2, TimeUnit.SECONDS);
    DataTable responseDT = DataTableFactory.getDataTable(responseBytes);
    Assert.assertEquals(responseDT.getNumberOfRows(), 2);
    Assert.assertEquals(responseDT.getString(0, 0), "mars");
    Assert.assertEquals(responseDT.getInt(0, 1), 10);
    Assert.assertEquals(responseDT.getString(1, 0), "jupiter");
    Assert.assertEquals(responseDT.getInt(1, 1), 100);
}