Examples with VariableWidthType com.facebook.presto.common.type.VariableWidthType used on opensource projects

Example 1 with VariableWidthType

use of com.facebook.presto.common.type.VariableWidthType 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");
        }
    }
}