Examples with DataTableBuilder org.apache.pinot.core.common.datatable.DataTableBuilder used on opensource projects

Example 1 with DataTableBuilder

use of org.apache.pinot.core.common.datatable.DataTableBuilder in project presto by prestodb.

the class TestPinotSegmentPageSource method testMultivaluedType.

@Test
public void testMultivaluedType() throws IOException {
    String[] columnNames = { "col1", "col2" };
    DataSchema.ColumnDataType[] columnDataTypes = { DataSchema.ColumnDataType.INT_ARRAY, DataSchema.ColumnDataType.STRING_ARRAY };
    DataSchema dataSchema = new DataSchema(columnNames, columnDataTypes);
    String[] stringArray = { "stringVal1", "stringVal2" };
    int[] intArray = { 10, 34, 67 };
    DataTableBuilder dataTableBuilder = new DataTableBuilder(dataSchema);
    dataTableBuilder.startRow();
    dataTableBuilder.setColumn(0, intArray);
    dataTableBuilder.setColumn(1, stringArray);
    dataTableBuilder.finishRow();
    DataTable dataTable = dataTableBuilder.build();
    PinotSessionProperties pinotSessionProperties = new PinotSessionProperties(pinotConfig);
    ConnectorSession session = new TestingConnectorSession(pinotSessionProperties.getSessionProperties());
    List<PinotColumnHandle> pinotColumnHandles = ImmutableList.of(new PinotColumnHandle(columnNames[0], PinotColumnUtils.getPrestoTypeFromPinotType(getFieldSpec(columnNames[0], columnDataTypes[0]), false, false), PinotColumnHandle.PinotColumnType.REGULAR), new PinotColumnHandle(columnNames[1], PinotColumnUtils.getPrestoTypeFromPinotType(getFieldSpec(columnNames[1], columnDataTypes[1]), false, false), PinotColumnHandle.PinotColumnType.REGULAR));
    PinotSplit mockPinotSplit = new PinotSplit(pinotConnectorId.toString(), PinotSplit.SplitType.SEGMENT, pinotColumnHandles, Optional.empty(), Optional.of("blah"), ImmutableList.of("seg"), Optional.of("host"), getGrpcPort());
    PinotSegmentPageSource pinotSegmentPageSource = getPinotSegmentPageSource(session, ImmutableList.of(dataTable), mockPinotSplit, pinotColumnHandles);
    Page page = requireNonNull(pinotSegmentPageSource.getNextPage(), "Expected a valid page");
    for (int i = 0; i < columnDataTypes.length; i++) {
        Block block = page.getBlock(i);
        Type type = PinotColumnUtils.getPrestoTypeFromPinotType(getFieldSpec(columnNames[i], columnDataTypes[i]), false, false);
        Assert.assertTrue(type instanceof ArrayType, "presto type should be array");
        if (((ArrayType) type).getElementType() instanceof IntegerType) {
            Assert.assertTrue(block.getBlock(0).getInt(0) == 10, "Array element not matching");
            Assert.assertTrue(block.getBlock(0).getInt(1) == 34, "Array element not matching");
            Assert.assertTrue(block.getBlock(0).getInt(2) == 67, "Array element not matching");
        } else if (((ArrayType) type).getElementType() instanceof VariableWidthType) {
            Type type1 = ((ArrayType) type).getElementType();
            Assert.assertTrue(block.getBlock(0) instanceof VariableWidthBlock);
            VariableWidthBlock variableWidthBlock = (VariableWidthBlock) block.getBlock(0);
            Assert.assertTrue("stringVal1".equals(new String(variableWidthBlock.getSlice(0, 0, variableWidthBlock.getSliceLength(0)).getBytes())), "Array element not matching");
            Assert.assertTrue("stringVal2".equals(new String(variableWidthBlock.getSlice(1, 0, variableWidthBlock.getSliceLength(1)).getBytes())), "Array element not matching");
        }
    }
}

Example 2 with DataTableBuilder

use of org.apache.pinot.core.common.datatable.DataTableBuilder in project presto by prestodb.

the class TestPinotSegmentPageSource method createDataTableWithAllTypes.

protected static DataTable createDataTableWithAllTypes() {
    try {
        int numColumns = ALL_TYPES.size();
        String[] columnNames = new String[numColumns];
        for (int i = 0; i < numColumns; i++) {
            columnNames[i] = ALL_TYPES.get(i).name();
        }
        DataSchema.ColumnDataType[] columnDataTypes = ALL_TYPES_ARRAY;
        DataSchema dataSchema = new DataSchema(columnNames, columnDataTypes);
        DataTableBuilder dataTableBuilder = new DataTableBuilder(dataSchema);
        for (int rowId = 0; rowId < NUM_ROWS; rowId++) {
            dataTableBuilder.startRow();
            for (int colId = 0; colId < numColumns; colId++) {
                switch(columnDataTypes[colId]) {
                    case BOOLEAN:
                        dataTableBuilder.setColumn(colId, RANDOM.nextBoolean());
                        break;
                    case INT:
                        dataTableBuilder.setColumn(colId, RANDOM.nextInt());
                        break;
                    case LONG:
                    case TIMESTAMP:
                        dataTableBuilder.setColumn(colId, RANDOM.nextLong());
                        break;
                    case FLOAT:
                        dataTableBuilder.setColumn(colId, RANDOM.nextFloat());
                        break;
                    case DOUBLE:
                        dataTableBuilder.setColumn(colId, RANDOM.nextDouble());
                        break;
                    case STRING:
                        dataTableBuilder.setColumn(colId, generateRandomStringWithLength(RANDOM.nextInt(20)));
                        break;
                    case OBJECT:
                        dataTableBuilder.setColumn(colId, (Object) RANDOM.nextDouble());
                        break;
                    case BOOLEAN_ARRAY:
                        int length = RANDOM.nextInt(20);
                        int[] booleanArray = new int[length];
                        for (int i = 0; i < length; i++) {
                            booleanArray[i] = RANDOM.nextInt(2);
                        }
                        dataTableBuilder.setColumn(colId, booleanArray);
                        break;
                    case INT_ARRAY:
                        length = RANDOM.nextInt(20);
                        int[] intArray = new int[length];
                        for (int i = 0; i < length; i++) {
                            intArray[i] = RANDOM.nextInt();
                        }
                        dataTableBuilder.setColumn(colId, intArray);
                        break;
                    case LONG_ARRAY:
                    case TIMESTAMP_ARRAY:
                        length = RANDOM.nextInt(20);
                        long[] longArray = new long[length];
                        for (int i = 0; i < length; i++) {
                            longArray[i] = RANDOM.nextLong();
                        }
                        dataTableBuilder.setColumn(colId, longArray);
                        break;
                    case FLOAT_ARRAY:
                        length = RANDOM.nextInt(20);
                        float[] floatArray = new float[length];
                        for (int i = 0; i < length; i++) {
                            floatArray[i] = RANDOM.nextFloat();
                        }
                        dataTableBuilder.setColumn(colId, floatArray);
                        break;
                    case DOUBLE_ARRAY:
                        length = RANDOM.nextInt(20);
                        double[] doubleArray = new double[length];
                        for (int i = 0; i < length; i++) {
                            doubleArray[i] = RANDOM.nextDouble();
                        }
                        dataTableBuilder.setColumn(colId, doubleArray);
                        break;
                    case STRING_ARRAY:
                    case BYTES_ARRAY:
                        length = RANDOM.nextInt(20);
                        String[] stringArray = new String[length];
                        for (int i = 0; i < length; i++) {
                            stringArray[i] = generateRandomStringWithLength(RANDOM.nextInt(20));
                        }
                        dataTableBuilder.setColumn(colId, stringArray);
                        break;
                    case JSON:
                        dataTableBuilder.setColumn(colId, "{ " + generateRandomStringWithLength(RANDOM.nextInt(5)) + " : " + generateRandomStringWithLength(RANDOM.nextInt(10)) + " }");
                        break;
                    default:
                        throw new RuntimeException("Unsupported type - " + columnDataTypes[colId]);
                }
            }
            dataTableBuilder.finishRow();
        }
        return dataTableBuilder.build();
    } catch (IOException e) {
        Assert.fail("Failed to create Pinot DataTable with all types", e);
        throw new RuntimeException(e);
    }
}