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

@Override
protected void readValue(BlockBuilder blockBuilder, Type type)
{
    if (definitionLevel == columnDescriptor.getMaxDefinitionLevel()) {
        Binary binary = valuesReader.readBytes();
        Slice value;
        if (binary.length() == 0) {
            value = EMPTY_SLICE;
        }
        else {
            value = wrappedBuffer(binary.getBytes());
        }
        if (isVarcharType(type)) {
            value = truncateToLength(value, type);
        }
        if (isCharType(type)) {
            value = truncateToLengthAndTrimSpaces(value, type);
        }
        type.writeSlice(blockBuilder, value);
    }
    else if (isValueNull()) {
        blockBuilder.appendNull();
    }
}
 

private static void writeSlice(BlockBuilder output, Type type, Object value)
{
    if (type instanceof VarcharType) {
        type.writeSlice(output, utf8Slice(((Utf8) value).toString()));
    }
    else if (type instanceof DecimalType) {
        BigDecimal bdValue = DECIMAL_CONVERTER.convert(value);
        type.writeSlice(output, Decimals.encodeScaledValue(bdValue, NUMERIC_DATA_TYPE_SCALE));
    }
    else if (type instanceof VarbinaryType) {
        if (value instanceof ByteBuffer) {
            type.writeSlice(output, Slices.wrappedBuffer((ByteBuffer) value));
        }
        else {
            output.appendNull();
        }
    }
    else {
        throw new PrestoException(GENERIC_INTERNAL_ERROR, "Unhandled type for Slice: " + type.getTypeSignature());
    }
}
 

private void writeSlice(BlockBuilder output, Type type, Object value)
{
    if (type instanceof VarcharType) {
        type.writeSlice(output, utf8Slice(toVarcharValue(value)));
    }
    else if (type instanceof CharType) {
        type.writeSlice(output, truncateToLengthAndTrimSpaces(utf8Slice((String) value), ((CharType) type)));
    }
    else if (type.equals(OBJECT_ID)) {
        type.writeSlice(output, wrappedBuffer(((ObjectId) value).toByteArray()));
    }
    else if (type instanceof VarbinaryType) {
        if (value instanceof Binary) {
            type.writeSlice(output, wrappedBuffer(((Binary) value).getData()));
        }
        else {
            output.appendNull();
        }
    }
    else if (type instanceof DecimalType) {
        type.writeSlice(output, encodeScaledValue(((Decimal128) value).bigDecimalValue(), ((DecimalType) type).getScale()));
    }
    else {
        throw new PrestoException(GENERIC_INTERNAL_ERROR, "Unhandled type for Slice: " + type.getTypeSignature());
    }
}
 

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 createChannel(int arraySize, Type elementType)
{
    BlockBuilder blockBuilder = elementType.createBlockBuilder(null, arraySize);
    for (int i = 0; i < arraySize; i++) {
        if (elementType.getJavaType() == long.class) {
            elementType.writeLong(blockBuilder, (long) i);
        }
        else if (elementType.getJavaType() == double.class) {
            elementType.writeDouble(blockBuilder, ThreadLocalRandom.current().nextDouble());
        }
        else if (elementType.getJavaType() == boolean.class) {
            elementType.writeBoolean(blockBuilder, ThreadLocalRandom.current().nextBoolean());
        }
        else if (elementType.equals(VARCHAR)) {
            // make sure the size of a varchar is rather small; otherwise the aggregated slice may overflow
            elementType.writeSlice(blockBuilder, Slices.utf8Slice(Long.toString(ThreadLocalRandom.current().nextLong() % 100)));
        }
        else {
            throw new UnsupportedOperationException();
        }
    }
    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();
}
 

private static Block createChannel(int positionCount, int arraySize, Type elementType)
{
    ArrayType arrayType = new ArrayType(elementType);
    BlockBuilder blockBuilder = arrayType.createBlockBuilder(null, positionCount);
    for (int position = 0; position < positionCount; position++) {
        BlockBuilder entryBuilder = blockBuilder.beginBlockEntry();
        for (int i = 0; i < arraySize; i++) {
            if (elementType.getJavaType() == long.class) {
                elementType.writeLong(entryBuilder, ThreadLocalRandom.current().nextLong() % arraySize);
            }
            else if (elementType.getJavaType() == double.class) {
                elementType.writeDouble(entryBuilder, ThreadLocalRandom.current().nextDouble() % arraySize);
            }
            else if (elementType.getJavaType() == boolean.class) {
                elementType.writeBoolean(entryBuilder, ThreadLocalRandom.current().nextBoolean());
            }
            else if (elementType.equals(VARCHAR)) {
                // make sure the size of a varchar is rather small; otherwise the aggregated slice may overflow
                elementType.writeSlice(entryBuilder, Slices.utf8Slice(Long.toString(ThreadLocalRandom.current().nextLong() % arraySize)));
            }
            else {
                throw new UnsupportedOperationException();
            }
        }
        blockBuilder.closeEntry();
    }
    return blockBuilder.build();
}
 

private static Block createVarWidthValueBlock(int positionCount, int mapSize)
{
    Type valueType = createUnboundedVarcharType();
    BlockBuilder valueBlockBuilder = valueType.createBlockBuilder(null, positionCount * mapSize);
    for (int i = 0; i < positionCount * mapSize; i++) {
        int wordLength = ThreadLocalRandom.current().nextInt(5, 10);
        valueType.writeSlice(valueBlockBuilder, utf8Slice(randomString(wordLength)));
    }
    return valueBlockBuilder.build();
}
 

private static Block createVarWidthValueBlock(int positionCount, int mapSize)
{
    Type valueType = createUnboundedVarcharType();
    BlockBuilder valueBlockBuilder = valueType.createBlockBuilder(null, positionCount * mapSize);
    for (int i = 0; i < positionCount * mapSize; i++) {
        int wordLength = ThreadLocalRandom.current().nextInt(5, 10);
        valueType.writeSlice(valueBlockBuilder, utf8Slice(randomString(wordLength)));
    }
    return valueBlockBuilder.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();
}
 

@Override
protected Block[] getSequenceBlocks(int start, int length)
{
    Type type = metadata.getType(new TypeSignature(QDIGEST.getName(), TypeSignatureParameter.typeParameter(DOUBLE.getTypeSignature())));
    BlockBuilder blockBuilder = type.createBlockBuilder(null, length);
    for (int i = start; i < start + length; i++) {
        QuantileDigest qdigest = new QuantileDigest(0.0);
        qdigest.add(i);
        type.writeSlice(blockBuilder, qdigest.serialize());
    }
    return new Block[] {blockBuilder.build()};
}
 

private static void serializePrimitive(BlockBuilder blockBuilder, Object value, Type type, String columnName)
{
    requireNonNull(blockBuilder, "parent blockBuilder is null");

    if (value == null) {
        blockBuilder.appendNull();
        return;
    }

    if (type instanceof BooleanType) {
        type.writeBoolean(blockBuilder, (Boolean) value);
        return;
    }

    if ((value instanceof Integer || value instanceof Long) && (type instanceof BigintType || type instanceof IntegerType || type instanceof SmallintType || type instanceof TinyintType)) {
        type.writeLong(blockBuilder, ((Number) value).longValue());
        return;
    }

    if (type instanceof DoubleType) {
        type.writeDouble(blockBuilder, (Double) value);
        return;
    }

    if (type instanceof RealType) {
        type.writeLong(blockBuilder, floatToIntBits((Float) value));
        return;
    }

    if (type instanceof VarcharType || type instanceof VarbinaryType) {
        type.writeSlice(blockBuilder, getSlice(value, type, columnName));
        return;
    }

    throw new PrestoException(DECODER_CONVERSION_NOT_SUPPORTED, format("cannot decode object of '%s' as '%s' for column '%s'", value.getClass(), type, columnName));
}
 

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

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

    return blockBuilder.build();
}
 

@UsedByGeneratedCode
public static Block appendElement(Type elementType, Block block, Slice value)
{
    BlockBuilder blockBuilder = elementType.createBlockBuilder(null, block.getPositionCount() + 1);
    for (int i = 0; i < block.getPositionCount(); i++) {
        elementType.appendTo(block, i, blockBuilder);
    }

    elementType.writeSlice(blockBuilder, value);

    return blockBuilder.build();
}
 

private void writeSliceBlock(BlockBuilder blockBuilder, Type columnType, int columnIndex)
{
    for (int i = 0; i < currentDataTable.getDataTable().getNumberOfRows(); i++) {
        Slice slice = getSlice(i, columnIndex);
        columnType.writeSlice(blockBuilder, slice, 0, slice.length());
        completedBytes += slice.getBytes().length;
    }
}
 

@Override
protected void readValue(BlockBuilder blockBuilder, Type prestoType)
{
    if (definitionLevel == columnDescriptor.getMaxDefinitionLevel()) {
        if (!((prestoType instanceof DecimalType) || isIntegerType(prestoType))) {
            throw new ParquetDecodingException(format("Unsupported Presto column type (%s) for Parquet column (%s)", prestoType, columnDescriptor));
        }

        long value;

        // When decimals are encoded with primitive types Parquet stores unscaled values
        if (columnDescriptor.getPrimitiveType().getPrimitiveTypeName() == INT32) {
            value = valuesReader.readInteger();
        }
        else if (columnDescriptor.getPrimitiveType().getPrimitiveTypeName() == INT64) {
            value = valuesReader.readLong();
        }
        else {
            value = getShortDecimalValue(valuesReader.readBytes().getBytes());
        }

        if (prestoType instanceof DecimalType) {
            DecimalType prestoDecimalType = (DecimalType) prestoType;

            if (isShortDecimal(prestoDecimalType)) {
                long rescale = longTenToNth(Math.abs(prestoDecimalType.getScale() - parquetDecimalType.getScale()));
                long convertedValue = shortToShortCast(
                        value,
                        parquetDecimalType.getPrecision(),
                        parquetDecimalType.getScale(),
                        prestoDecimalType.getPrecision(),
                        prestoDecimalType.getScale(),
                        rescale,
                        rescale / 2);

                prestoType.writeLong(blockBuilder, convertedValue);
            }
            else if (isLongDecimal(prestoDecimalType)) {
                prestoType.writeSlice(blockBuilder, shortToLongCast(
                        value,
                        parquetDecimalType.getPrecision(),
                        parquetDecimalType.getScale(),
                        prestoDecimalType.getPrecision(),
                        prestoDecimalType.getScale()));
            }
        }
        else {
            if (parquetDecimalType.getScale() != 0) {
                throw new PrestoException(NOT_SUPPORTED, format("Unsupported Presto column type (%s) for Parquet column (%s)", prestoType, columnDescriptor));
            }

            if (!isInValidNumberRange(prestoType, value)) {
                throw new PrestoException(NOT_SUPPORTED, format("Could not coerce from %s to %s: %s", parquetDecimalType, prestoType, value));
            }
            prestoType.writeLong(blockBuilder, value);
        }
    }
    else if (isValueNull()) {
        blockBuilder.appendNull();
    }
}
 

@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;
}
 

private Block createZeroBlock(Type type, int rowsCount, Slice constantSlice)
{
    checkArgument(isSupportedType(type), "Unsupported type [%s]", type);

    Slice slice;
    // do not exceed varchar limit
    if (isVarcharType(type) && !((VarcharType) type).isUnbounded()) {
        slice = constantSlice.slice(0, Math.min(((VarcharType) type).getBoundedLength(), constantSlice.length()));
    }
    else if (isLongDecimal(type)) {
        slice = encodeScaledValue(ZERO);
    }
    else {
        slice = constantSlice;
    }

    BlockBuilder builder;
    if (type instanceof FixedWidthType) {
        builder = type.createBlockBuilder(null, rowsCount);
    }
    else {
        builder = type.createBlockBuilder(null, rowsCount, slice.length());
    }

    for (int i = 0; i < rowsCount; i++) {
        Class<?> javaType = type.getJavaType();
        if (javaType == boolean.class) {
            type.writeBoolean(builder, false);
        }
        else if (javaType == long.class) {
            type.writeLong(builder, 0);
        }
        else if (javaType == double.class) {
            type.writeDouble(builder, 0.0);
        }
        else if (javaType == Slice.class) {
            requireNonNull(slice, "slice is null");
            type.writeSlice(builder, slice, 0, slice.length());
        }
        else {
            throw new UnsupportedOperationException("Unknown javaType: " + javaType.getName());
        }
    }
    return builder.build();
}
 

private static void serializePrimitive(Type type, BlockBuilder builder, Object object, PrimitiveObjectInspector inspector)
{
    requireNonNull(builder, "parent builder is null");

    if (object == null) {
        builder.appendNull();
        return;
    }

    switch (inspector.getPrimitiveCategory()) {
        case BOOLEAN:
            BooleanType.BOOLEAN.writeBoolean(builder, ((BooleanObjectInspector) inspector).get(object));
            return;
        case BYTE:
            TinyintType.TINYINT.writeLong(builder, ((ByteObjectInspector) inspector).get(object));
            return;
        case SHORT:
            SmallintType.SMALLINT.writeLong(builder, ((ShortObjectInspector) inspector).get(object));
            return;
        case INT:
            IntegerType.INTEGER.writeLong(builder, ((IntObjectInspector) inspector).get(object));
            return;
        case LONG:
            BigintType.BIGINT.writeLong(builder, ((LongObjectInspector) inspector).get(object));
            return;
        case FLOAT:
            RealType.REAL.writeLong(builder, floatToRawIntBits(((FloatObjectInspector) inspector).get(object)));
            return;
        case DOUBLE:
            DoubleType.DOUBLE.writeDouble(builder, ((DoubleObjectInspector) inspector).get(object));
            return;
        case STRING:
            type.writeSlice(builder, Slices.utf8Slice(((StringObjectInspector) inspector).getPrimitiveJavaObject(object)));
            return;
        case VARCHAR:
            type.writeSlice(builder, Slices.utf8Slice(((HiveVarcharObjectInspector) inspector).getPrimitiveJavaObject(object).getValue()));
            return;
        case CHAR:
            CharType charType = (CharType) type;
            HiveChar hiveChar = ((HiveCharObjectInspector) inspector).getPrimitiveJavaObject(object);
            type.writeSlice(builder, truncateToLengthAndTrimSpaces(Slices.utf8Slice(hiveChar.getValue()), charType.getLength()));
            return;
        case DATE:
            DateType.DATE.writeLong(builder, formatDateAsLong(object, (DateObjectInspector) inspector));
            return;
        case TIMESTAMP:
            TimestampType.TIMESTAMP.writeLong(builder, formatTimestampAsLong(object, (TimestampObjectInspector) inspector));
            return;
        case BINARY:
            VARBINARY.writeSlice(builder, Slices.wrappedBuffer(((BinaryObjectInspector) inspector).getPrimitiveJavaObject(object)));
            return;
        case DECIMAL:
            DecimalType decimalType = (DecimalType) type;
            HiveDecimalWritable hiveDecimal = ((HiveDecimalObjectInspector) inspector).getPrimitiveWritableObject(object);
            if (decimalType.isShort()) {
                decimalType.writeLong(builder, DecimalUtils.getShortDecimalValue(hiveDecimal, decimalType.getScale()));
            }
            else {
                decimalType.writeSlice(builder, DecimalUtils.getLongDecimalValue(hiveDecimal, decimalType.getScale()));
            }
            return;
    }
    throw new RuntimeException("Unknown primitive type: " + inspector.getPrimitiveCategory());
}