Java Code Examples for io.prestosql.spi.type.Type#writeObject()

@TypeParameter("T")
@SqlType("array(T)")
public static Block repeat(
        @TypeParameter("T") Type type,
        @SqlNullable @SqlType("T") Object element,
        @SqlType(StandardTypes.INTEGER) long count)
{
    BlockBuilder blockBuilder = createBlockBuilder(type, count);
    if (element == null) {
        return repeatNullValues(blockBuilder, count);
    }
    if (count > 0) {
        type.writeObject(blockBuilder, element);
        checkMaxSize(blockBuilder.getSizeInBytes(), count);
    }
    for (int i = 1; i < count; i++) {
        type.writeObject(blockBuilder, element);
    }
    return blockBuilder.build();
}
 

private static Block createBlock(Type type, Object value)
{
    BlockBuilder blockBuilder = type.createBlockBuilder(null, 1);

    Class<?> javaType = type.getJavaType();
    if (value == null) {
        blockBuilder.appendNull();
    }
    else if (javaType == boolean.class) {
        type.writeBoolean(blockBuilder, (Boolean) value);
    }
    else if (javaType == long.class) {
        type.writeLong(blockBuilder, (Long) value);
    }
    else if (javaType == double.class) {
        type.writeDouble(blockBuilder, (Double) value);
    }
    else if (javaType == Slice.class) {
        Slice slice = (Slice) value;
        type.writeSlice(blockBuilder, slice, 0, slice.length());
    }
    else {
        type.writeObject(blockBuilder, value);
    }
    return blockBuilder.build();
}
 

private static Block convertColumnToBlock(RecordSet recordSet, int columnIndex, int positions)
{
    Type type = recordSet.getColumnTypes().get(columnIndex);
    BlockBuilder output = type.createBlockBuilder(null, positions);
    Class<?> javaType = type.getJavaType();
    RecordCursor cursor = recordSet.cursor();
    for (int position = 0; position < positions; position++) {
        checkState(cursor.advanceNextPosition(), "cursor has less values than expected");
        if (cursor.isNull(columnIndex)) {
            output.appendNull();
        }
        else {
            if (javaType == boolean.class) {
                type.writeBoolean(output, cursor.getBoolean(columnIndex));
            }
            else if (javaType == long.class) {
                type.writeLong(output, cursor.getLong(columnIndex));
            }
            else if (javaType == double.class) {
                type.writeDouble(output, cursor.getDouble(columnIndex));
            }
            else if (javaType == Slice.class) {
                Slice slice = cursor.getSlice(columnIndex);
                type.writeSlice(output, slice, 0, slice.length());
            }
            else {
                type.writeObject(output, cursor.getObject(columnIndex));
            }
        }
    }
    checkState(!cursor.advanceNextPosition(), "cursor has more values than expected");
    return output.build();
}
 

public static Block createTestBlock(Type arrayType)
{
    BlockBuilder blockBuilder = arrayType.createBlockBuilder(null, 4);
    arrayType.writeObject(blockBuilder, arrayBlockOf(INTEGER, 1, 2));
    arrayType.writeObject(blockBuilder, arrayBlockOf(INTEGER, 1, 2, 3));
    arrayType.writeObject(blockBuilder, arrayBlockOf(INTEGER, 1, 2, 3));
    arrayType.writeObject(blockBuilder, arrayBlockOf(INTEGER, 100, 200, 300));
    return blockBuilder.build();
}
 

public static Block createTestBlock(Type mapType)
{
    BlockBuilder blockBuilder = mapType.createBlockBuilder(null, 2);
    mapType.writeObject(blockBuilder, mapBlockOf(INTEGER, VARCHAR, ImmutableMap.of(1, "hi")));
    mapType.writeObject(blockBuilder, mapBlockOf(INTEGER, VARCHAR, ImmutableMap.of(1, "2", 2, "hello")));
    return blockBuilder.build();
}
 

public static Block createTestBlock(Type mapType)
{
    BlockBuilder blockBuilder = mapType.createBlockBuilder(null, 2);
    mapType.writeObject(blockBuilder, mapBlockOf(BIGINT, VARCHAR, ImmutableMap.of(1, "hi")));
    mapType.writeObject(blockBuilder, mapBlockOf(BIGINT, VARCHAR, ImmutableMap.of(1, "2", 2, "hello")));
    return blockBuilder.build();
}
 

public static Block createTestBlock(Type arrayType)
{
    BlockBuilder blockBuilder = arrayType.createBlockBuilder(null, 4);
    arrayType.writeObject(blockBuilder, arrayBlockOf(COLOR, 1, 2));
    arrayType.writeObject(blockBuilder, arrayBlockOf(COLOR, 1, 2, 3));
    arrayType.writeObject(blockBuilder, arrayBlockOf(COLOR, 1, 2, 3));
    arrayType.writeObject(blockBuilder, arrayBlockOf(COLOR, 100, 200, 300));
    return blockBuilder.build();
}
 

public static Block createTestBlock(Type arrayType)
{
    BlockBuilder blockBuilder = arrayType.createBlockBuilder(null, 4);
    arrayType.writeObject(blockBuilder, arrayBlockOf(SMALLINT, 1, 2));
    arrayType.writeObject(blockBuilder, arrayBlockOf(SMALLINT, 1, 2, 3));
    arrayType.writeObject(blockBuilder, arrayBlockOf(SMALLINT, 1, 2, 3));
    arrayType.writeObject(blockBuilder, arrayBlockOf(SMALLINT, 100, 200, 300));
    return blockBuilder.build();
}
 

public static Block createTestBlock(Type arrayType)
{
    BlockBuilder blockBuilder = arrayType.createBlockBuilder(null, 4);
    arrayType.writeObject(blockBuilder, arrayBlockOf(BIGINT, 1, 2));
    arrayType.writeObject(blockBuilder, arrayBlockOf(BIGINT, 1, 2, 3));
    arrayType.writeObject(blockBuilder, arrayBlockOf(BIGINT, 1, 2, 3));
    arrayType.writeObject(blockBuilder, arrayBlockOf(BIGINT, 100, 200, 300));
    return blockBuilder.build();
}
 

public static Block createTestBlock(Type mapType)
{
    BlockBuilder blockBuilder = mapType.createBlockBuilder(null, 2);
    mapType.writeObject(blockBuilder, mapBlockOf(TINYINT, VARCHAR, ImmutableMap.of(1, "hi")));
    mapType.writeObject(blockBuilder, mapBlockOf(TINYINT, VARCHAR, ImmutableMap.of(1, "2", 2, "hello")));
    return blockBuilder.build();
}
 

public static Block createTestBlock(Type arrayType)
{
    BlockBuilder blockBuilder = arrayType.createBlockBuilder(null, 4);
    arrayType.writeObject(blockBuilder, arrayBlockOf(TINYINT, 1, 2));
    arrayType.writeObject(blockBuilder, arrayBlockOf(TINYINT, 1, 2, 3));
    arrayType.writeObject(blockBuilder, arrayBlockOf(TINYINT, 1, 2, 3));
    arrayType.writeObject(blockBuilder, arrayBlockOf(TINYINT, 100, 110, 127));
    return blockBuilder.build();
}
 

public static Block createTestBlock(Type mapType)
{
    BlockBuilder blockBuilder = mapType.createBlockBuilder(null, 2);
    mapType.writeObject(blockBuilder, mapBlockOf(SMALLINT, VARCHAR, ImmutableMap.of(1, "hi")));
    mapType.writeObject(blockBuilder, mapBlockOf(SMALLINT, VARCHAR, ImmutableMap.of(1, "2", 2, "hello")));
    return blockBuilder.build();
}
 

/**
 * Produce a block with the given values in the last field.
 */
private static Block createBlock(Type type, List<?> values)
{
    BlockBuilder blockBuilder = type.createBlockBuilder(null, values.size());

    for (Object value : values) {
        Class<?> javaType = type.getJavaType();
        if (value == null) {
            blockBuilder.appendNull();
        }
        else if (javaType == boolean.class) {
            type.writeBoolean(blockBuilder, (Boolean) value);
        }
        else if (javaType == long.class) {
            type.writeLong(blockBuilder, (Long) value);
        }
        else if (javaType == double.class) {
            type.writeDouble(blockBuilder, (Double) value);
        }
        else if (javaType == Slice.class) {
            Slice slice = (Slice) value;
            type.writeSlice(blockBuilder, slice, 0, slice.length());
        }
        else {
            type.writeObject(blockBuilder, value);
        }
    }

    return blockBuilder.build();
}
 

@UsedByGeneratedCode
public static Block prependElement(Type elementType, Object value, Block block)
{
    BlockBuilder blockBuilder = elementType.createBlockBuilder(null, block.getPositionCount() + 1);

    elementType.writeObject(blockBuilder, value);
    for (int i = 0; i < block.getPositionCount(); i++) {
        elementType.appendTo(block, i, blockBuilder);
    }

    return blockBuilder.build();
}
 

static void write(Type type, BlockState state, BlockBuilder out)
{
    if (state.getBlock() == null) {
        out.appendNull();
    }
    else {
        type.writeObject(out, state.getBlock());
    }
}
 

public static void output(Type keyType, Type valueType, MultimapAggregationState state, BlockBuilder out)
{
    if (state.isEmpty()) {
        out.appendNull();
    }
    else {
        // TODO: Avoid copy value block associated with the same key by using strategy similar to multimap_from_entries
        ObjectBigArray<BlockBuilder> valueArrayBlockBuilders = new ObjectBigArray<>();
        valueArrayBlockBuilders.ensureCapacity(state.getEntryCount());
        BlockBuilder distinctKeyBlockBuilder = keyType.createBlockBuilder(null, state.getEntryCount(), expectedValueSize(keyType, 100));
        TypedSet keySet = new TypedSet(keyType, state.getEntryCount(), MultimapAggregationFunction.NAME);

        state.forEach((key, value, keyValueIndex) -> {
            // Merge values of the same key into an array
            if (!keySet.contains(key, keyValueIndex)) {
                keySet.add(key, keyValueIndex);
                keyType.appendTo(key, keyValueIndex, distinctKeyBlockBuilder);
                BlockBuilder valueArrayBuilder = valueType.createBlockBuilder(null, 10, expectedValueSize(valueType, EXPECTED_ENTRY_SIZE));
                valueArrayBlockBuilders.set(keySet.positionOf(key, keyValueIndex), valueArrayBuilder);
            }
            valueType.appendTo(value, keyValueIndex, valueArrayBlockBuilders.get(keySet.positionOf(key, keyValueIndex)));
        });

        // Write keys and value arrays into one Block
        Type valueArrayType = new ArrayType(valueType);
        BlockBuilder multimapBlockBuilder = out.beginBlockEntry();
        for (int i = 0; i < distinctKeyBlockBuilder.getPositionCount(); i++) {
            keyType.appendTo(distinctKeyBlockBuilder, i, multimapBlockBuilder);
            valueArrayType.writeObject(multimapBlockBuilder, valueArrayBlockBuilders.get(i).build());
        }
        out.closeEntry();
    }
}
 

private static SliceReadFunction arrayAsJsonReadFunction(ConnectorSession session, ColumnMapping baseElementMapping)
{
    return (resultSet, columnIndex) -> {
        // resolve array type
        Object jdbcArray = resultSet.getArray(columnIndex).getArray();
        int arrayDimensions = arrayDepth(jdbcArray);

        ReadFunction readFunction = baseElementMapping.getReadFunction();
        Type type = baseElementMapping.getType();
        for (int i = 0; i < arrayDimensions; i++) {
            readFunction = arrayReadFunction(type, readFunction);
            type = new ArrayType(type);
        }

        // read array into a block
        Block block = (Block) ((ObjectReadFunction) readFunction).readObject(resultSet, columnIndex);

        // convert block to JSON slice
        BlockBuilder builder = type.createBlockBuilder(null, 1);
        type.writeObject(builder, block);
        Object value = type.getObjectValue(session, builder.build(), 0);

        try {
            return Slices.wrappedBuffer(SORTED_MAPPER.writeValueAsBytes(value));
        }
        catch (JsonProcessingException e) {
            throw new PrestoException(JDBC_ERROR, "Conversion to JSON failed for  " + type.getDisplayName(), e);
        }
    };
}
 

private static void rewriteBlock(
        Type fromType,
        Type toType,
        Block block,
        int position,
        BlockBuilder blockBuilder,
        Coercer coercer,
        BridgingRecordCursor bridgingRecordCursor)
{
    Class<?> fromJavaType = fromType.getJavaType();
    if (fromJavaType == long.class) {
        bridgingRecordCursor.setValue(fromType.getLong(block, position));
    }
    else if (fromJavaType == double.class) {
        bridgingRecordCursor.setValue(fromType.getDouble(block, position));
    }
    else if (fromJavaType == boolean.class) {
        bridgingRecordCursor.setValue(fromType.getBoolean(block, position));
    }
    else if (fromJavaType == Slice.class) {
        bridgingRecordCursor.setValue(fromType.getSlice(block, position));
    }
    else if (fromJavaType == Block.class) {
        bridgingRecordCursor.setValue(fromType.getObject(block, position));
    }
    else {
        bridgingRecordCursor.setValue(null);
    }
    coercer.reset();
    Class<?> toJaveType = toType.getJavaType();
    if (coercer.isNull(bridgingRecordCursor, 0)) {
        blockBuilder.appendNull();
    }
    else if (toJaveType == long.class) {
        toType.writeLong(blockBuilder, coercer.getLong(bridgingRecordCursor, 0));
    }
    else if (toJaveType == double.class) {
        toType.writeDouble(blockBuilder, coercer.getDouble(bridgingRecordCursor, 0));
    }
    else if (toJaveType == boolean.class) {
        toType.writeBoolean(blockBuilder, coercer.getBoolean(bridgingRecordCursor, 0));
    }
    else if (toJaveType == Slice.class) {
        toType.writeSlice(blockBuilder, coercer.getSlice(bridgingRecordCursor, 0));
    }
    else if (toJaveType == Block.class) {
        toType.writeObject(blockBuilder, coercer.getObject(bridgingRecordCursor, 0));
    }
    else {
        throw new PrestoException(NOT_SUPPORTED, format("Unsupported coercion from %s to %s", fromType.getDisplayName(), toType.getDisplayName()));
    }
    coercer.reset();
    bridgingRecordCursor.close();
}
 

@Override
public Page getNextPage()
{
    if (!closed) {
        for (int i = 0; i < ROWS_PER_REQUEST && !pageBuilder.isFull() && pageBuilder.getPositionCount() < maxRowsPerPage; i++) {
            if (!cursor.advanceNextPosition()) {
                closed = true;
                break;
            }

            pageBuilder.declarePosition();
            for (int column = 0; column < types.size(); column++) {
                BlockBuilder output = pageBuilder.getBlockBuilder(column);
                if (cursor.isNull(column)) {
                    output.appendNull();
                }
                else {
                    Type type = types.get(column);
                    Class<?> javaType = type.getJavaType();
                    if (javaType == boolean.class) {
                        type.writeBoolean(output, cursor.getBoolean(column));
                    }
                    else if (javaType == long.class) {
                        type.writeLong(output, cursor.getLong(column));
                    }
                    else if (javaType == double.class) {
                        type.writeDouble(output, cursor.getDouble(column));
                    }
                    else if (javaType == Slice.class) {
                        Slice slice = cursor.getSlice(column);
                        type.writeSlice(output, slice, 0, slice.length());
                    }
                    else {
                        type.writeObject(output, cursor.getObject(column));
                    }
                }
            }
        }
    }

    if ((closed && !pageBuilder.isEmpty()) || pageBuilder.isFull() || pageBuilder.getPositionCount() >= maxRowsPerPage) {
        Page page = pageBuilder.build();
        pageBuilder.reset();
        return lazyWrapper(page);
    }

    return null;
}
 

@Override
public Page getNextPage()
{
    if (!closed) {
        for (int i = 0; i < ROWS_PER_REQUEST && !pageBuilder.isFull(); i++) {
            if (!cursor.advanceNextPosition()) {
                closed = true;
                break;
            }

            pageBuilder.declarePosition();
            for (int column = 0; column < types.size(); column++) {
                BlockBuilder output = pageBuilder.getBlockBuilder(column);
                if (cursor.isNull(column)) {
                    output.appendNull();
                }
                else {
                    Type type = types.get(column);
                    Class<?> javaType = type.getJavaType();
                    if (javaType == boolean.class) {
                        type.writeBoolean(output, cursor.getBoolean(column));
                    }
                    else if (javaType == long.class) {
                        type.writeLong(output, cursor.getLong(column));
                    }
                    else if (javaType == double.class) {
                        type.writeDouble(output, cursor.getDouble(column));
                    }
                    else if (javaType == Slice.class) {
                        Slice slice = cursor.getSlice(column);
                        type.writeSlice(output, slice, 0, slice.length());
                    }
                    else {
                        type.writeObject(output, cursor.getObject(column));
                    }
                }
            }
        }
    }

    // only return a page if the buffer is full or we are finishing
    if ((closed && !pageBuilder.isEmpty()) || pageBuilder.isFull()) {
        Page page = pageBuilder.build();
        pageBuilder.reset();
        return page;
    }

    return null;
}