Examples with DataTable org.apache.pinot.common.utils.DataTable used on opensource projects

Example 1 with DataTable

use of org.apache.pinot.common.utils.DataTable 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 DataTable

use of org.apache.pinot.common.utils.DataTable in project presto by prestodb.

the class TestPinotSegmentPageSource method testAllDataTypes.

@Test
public void testAllDataTypes() {
    PinotSessionProperties pinotSessionProperties = new PinotSessionProperties(pinotConfig);
    ConnectorSession session = new TestingConnectorSession(pinotSessionProperties.getSessionProperties());
    List<DataTable> dataTables = IntStream.range(0, 3).mapToObj(i -> createDataTableWithAllTypes()).collect(toImmutableList());
    List<PinotColumnHandle> pinotColumnHandles = createPinotColumnHandlesWithAllTypes();
    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, dataTables, mockPinotSplit, pinotColumnHandles);
    for (int i = 0; i < dataTables.size(); ++i) {
        Page page = requireNonNull(pinotSegmentPageSource.getNextPage(), "Expected a valid page");
        for (int j = 0; j < ALL_TYPES.size(); ++j) {
            Block block = page.getBlock(j);
            Type type = PinotColumnUtils.getPrestoTypeFromPinotType(getFieldSpec("dontcare", ALL_TYPES.get(j)), false, false);
            long maxHashCode = Long.MIN_VALUE;
            for (int k = 0; k < NUM_ROWS; ++k) {
                maxHashCode = Math.max(type.hash(block, k), maxHashCode);
            }
            Assert.assertTrue(maxHashCode != 0, "Not all column values can have hash code 0");
        }
    }
}

Example 3 with DataTable

use of org.apache.pinot.common.utils.DataTable in project presto by prestodb.

the class PinotSegmentStreamingPageSource method getNextPage.

/**
 * @return constructed page for pinot data.
 */
@Override
public Page getNextPage() {
    if (closed) {
        return null;
    }
    if (serverResponseIterator == null) {
        serverResponseIterator = queryPinot(split);
    }
    ByteBuffer byteBuffer = null;
    try {
        // So we need to check ResponseType of each ServerResponse.
        if (serverResponseIterator.hasNext()) {
            long startTimeNanos = System.nanoTime();
            ServerResponse serverResponse = serverResponseIterator.next();
            readTimeNanos += System.nanoTime() - startTimeNanos;
            final String responseType = serverResponse.getMetadataOrThrow("responseType");
            switch(responseType) {
                case CommonConstants.Query.Response.ResponseType.DATA:
                    estimatedMemoryUsageInBytes = serverResponse.getSerializedSize();
                    // Store each dataTable which will later be constructed into Pages.
                    try {
                        byteBuffer = serverResponse.getPayload().asReadOnlyByteBuffer();
                        DataTable dataTable = DataTableFactory.getDataTable(byteBuffer);
                        checkExceptions(dataTable, split, PinotSessionProperties.isMarkDataFetchExceptionsAsRetriable(session));
                        currentDataTable = new PinotSegmentPageSource.PinotDataTableWithSize(dataTable, serverResponse.getSerializedSize());
                    } catch (IOException e) {
                        throw new PinotException(PINOT_DATA_FETCH_EXCEPTION, split.getSegmentPinotQuery(), String.format("Encountered Pinot exceptions when fetching data table from Split: < %s >", split), e);
                    }
                    break;
                case CommonConstants.Query.Response.ResponseType.METADATA:
                    // The last part of the response is Metadata
                    currentDataTable = null;
                    serverResponseIterator = null;
                    close();
                    return null;
                default:
                    throw new PinotException(PINOT_UNEXPECTED_RESPONSE, split.getSegmentPinotQuery(), String.format("Encountered Pinot exceptions, unknown response type - %s", responseType));
            }
        }
        Page page = fillNextPage();
        completedPositions += currentDataTable.getDataTable().getNumberOfRows();
        return page;
    } finally {
        if (byteBuffer != null) {
            byteBuffer.clear();
        }
    }
}

Example 4 with DataTable

use of org.apache.pinot.common.utils.DataTable in project presto by prestodb.

the class TestPinotSegmentPageSource method testPrunedColumns.

@Test
public void testPrunedColumns() {
    PinotSessionProperties pinotSessionProperties = new PinotSessionProperties(pinotConfig);
    ConnectorSession session = new TestingConnectorSession(pinotSessionProperties.getSessionProperties());
    List<DataTable> dataTables = IntStream.range(0, 3).mapToObj(i -> createDataTableWithAllTypes()).collect(toImmutableList());
    List<PinotColumnHandle> expectedColumnHandles = createPinotColumnHandlesWithAllTypes();
    PinotSplit mockPinotSplit = new PinotSplit(pinotConnectorId.toString(), PinotSplit.SplitType.SEGMENT, expectedColumnHandles, Optional.empty(), Optional.of("blah"), ImmutableList.of("seg"), Optional.of("host"), getGrpcPort());
    ImmutableList.Builder<Integer> columnsSurvivingBuilder = ImmutableList.builder();
    for (int i = expectedColumnHandles.size() - 1; i >= 0; i--) {
        if (i % 2 == 0) {
            columnsSurvivingBuilder.add(i);
        }
    }
    List<Integer> columnsSurviving = columnsSurvivingBuilder.build();
    List<PinotColumnHandle> handlesSurviving = columnsSurviving.stream().map(expectedColumnHandles::get).collect(toImmutableList());
    PinotSegmentPageSource pinotSegmentPageSource = getPinotSegmentPageSource(session, dataTables, mockPinotSplit, handlesSurviving);
    for (int i = 0; i < dataTables.size(); ++i) {
        Page page = requireNonNull(pinotSegmentPageSource.getNextPage(), "Expected a valid page");
        Assert.assertEquals(page.getChannelCount(), columnsSurviving.size());
        for (int j = 0; j < columnsSurviving.size(); ++j) {
            Block block = page.getBlock(j);
            int originalColumnIndex = columnsSurviving.get(j);
            Type type = PinotColumnUtils.getPrestoTypeFromPinotType(getFieldSpec("dontcare", ALL_TYPES.get(originalColumnIndex)), false, false);
            long maxHashCode = Long.MIN_VALUE;
            for (int k = 0; k < NUM_ROWS; k++) {
                maxHashCode = Math.max(type.hash(block, k), maxHashCode);
            }
            Assert.assertTrue(maxHashCode != 0, "Not all column values can have hash code 0");
        }
    }
}