com.facebook.presto.spi.block.Block Java Examples

/**
 * @param type Presto type
 * @param block Block to decode
 * @param position Position in the block to get
 * @return Java object from the Block
 */
private static Object readObject(Type type, Block block, int position)
{
    if (Types.isArrayType(type)) {
        Type elementType = Types.getElementType(type);
        return getArrayFromBlock(elementType, block.getObject(position, Block.class));
    }
    else if (Types.isMapType(type)) {
        return getMapFromBlock(type, block.getObject(position, Block.class));
    }
    else {
        if (type.getJavaType() == Slice.class) {
            Slice slice = (Slice) TypeUtils.readNativeValue(type, block, position);
            return type.equals(VarcharType.VARCHAR) ? slice.toStringUtf8() : slice.getBytes();
        }

        return TypeUtils.readNativeValue(type, block, position);
    }
}
 

/**
 * Encodes the given map into a Block.
 *
 * @param mapType Presto type of the map
 * @param bytes hbase byte[]
 * @return Presto Block
 */
public static Block getBlockFromMap(Type mapType, byte[] bytes)
        throws IOException
{
    Type keyType = mapType.getTypeParameters().get(0);
    Type valueType = mapType.getTypeParameters().get(1);
    Map<?, ?> map = MAPPER.readValue(bytes, new MyTypeReference(keyType.getJavaType(), valueType.getJavaType()));

    BlockBuilder mapBlockBuilder = mapType.createBlockBuilder(null, 1);
    BlockBuilder builder = mapBlockBuilder.beginBlockEntry();

    for (Map.Entry<?, ?> entry : map.entrySet()) {
        writeObject(builder, keyType, entry.getKey());
        writeObject(builder, valueType, entry.getValue());
    }

    mapBlockBuilder.closeEntry();
    return (Block) mapType.getObject(mapBlockBuilder, 0);
}
 

@Override
public void deleteRows(Block rowIds) {
    Schema schema = table.getSchema();
    KuduSession session = clientSession.newSession();
    session.setFlushMode(FlushMode.AUTO_FLUSH_BACKGROUND);
    try {
        try {
            for (int i = 0; i < rowIds.getPositionCount(); i++) {
                int len = rowIds.getSliceLength(i);
                Slice slice = rowIds.getSlice(i, 0, len);
                PartialRow row = KeyEncoderAccessor.decodePrimaryKey(schema, slice.getBytes());
                Delete delete = table.newDelete();
                RowHelper.copyPrimaryKey(schema, row, delete.getRow());
                session.apply(delete);
            }
        } finally {
            session.close();
        }
    } catch (KuduException e) {
        throw new RuntimeException(e);
    }
}
 

public Block readPrimitive(ColumnDescriptor columnDescriptor, Type type, IntList offsets)
        throws IOException
{
    ParquetColumnReader columnReader = columnReadersMap.get(columnDescriptor);
    if (columnReader.getPageReader() == null) {
        validateParquet(currentBlockMetadata.getRowCount() > 0, "Row group has 0 rows");
        ColumnChunkMetaData metadata = getColumnChunkMetaData(columnDescriptor);
        long startingPosition = metadata.getStartingPos();
        int totalSize = checkedCast(metadata.getTotalSize());
        byte[] buffer = new byte[totalSize];
        dataSource.readFully(startingPosition, buffer);
        ParquetColumnChunkDescriptor descriptor = new ParquetColumnChunkDescriptor(columnDescriptor, metadata, totalSize);
        ParquetColumnChunk columnChunk = new ParquetColumnChunk(descriptor, buffer, 0);
        columnReader.setPageReader(columnChunk.readAllPages());
    }
    return columnReader.readPrimitive(type, offsets);
}
 

@Override
public final void load(LazyBlock lazyBlock)
{
    if (loaded) {
        return;
    }

    checkState(batchId == expectedBatchId);

    try {
        Block block = parquetReader.readPrimitive(columnDescriptor, type);
        lazyBlock.setBlock(block);
    }
    catch (IOException e) {
        throw new ParaflowCursorException();
    }
    loaded = true;
}
 

@Override
public boolean equalTo(Block leftBlock, int leftPosition, Block rightBlock, int rightPosition)
{
    // Check that the position is valid
    checkArgument(leftBlock.isNull(leftPosition), "Expected NULL value for UnknownType");
    checkArgument(rightBlock.isNull(rightPosition), "Expected NULL value for UnknownType");
    return true;
}
 

@Override
public Object getObjectValue(ConnectorSession session, Block block, int position)
{
    if (block.isNull(position)) {
        return null;
    }

    return new SqlVarbinary(block.getSlice(position, 0, block.getSliceLength(position)).getBytes());
}
 

@Override
public void appendTo(Block block, int position, BlockBuilder blockBuilder)
{
    if (block.isNull(position)) {
        blockBuilder.appendNull();
    }
    else {
        block.writeBytesTo(position, 0, block.getSliceLength(position), blockBuilder);
        blockBuilder.closeEntry();
    }
}
 

@Override
public void deserialize(Block block, int index, BloomFilterState state)
{
    if (!block.isNull(index)) {
        state.setBloomFilter(BloomFilter.newInstance(BloomFilterType.BLOOM_FILTER.getSlice(block, index)));
    }
}
 

private void appendColumn(PartialRow row, Page page, int position, int channel, int destChannel) {
    Block block = page.getBlock(channel);
    Type type = columnTypes.get(destChannel);
    if (block.isNull(position)) {
        row.setNull(destChannel);
    } else if (TIMESTAMP.equals(type)) {
        row.addLong(destChannel, type.getLong(block, position) * 1000);
    } else if (REAL.equals(type)) {
        row.addFloat(destChannel, intBitsToFloat((int) type.getLong(block, position)));
    } else if (BIGINT.equals(type)) {
        row.addLong(destChannel, type.getLong(block, position));
    } else if (INTEGER.equals(type)) {
        row.addInt(destChannel, (int) type.getLong(block, position));
    } else if (SMALLINT.equals(type)) {
        row.addShort(destChannel, (short) type.getLong(block, position));
    } else if (TINYINT.equals(type)) {
        row.addByte(destChannel, (byte) type.getLong(block, position));
    } else if (BOOLEAN.equals(type)) {
        row.addBoolean(destChannel, type.getBoolean(block, position));
    } else if (DOUBLE.equals(type)) {
        row.addDouble(destChannel, type.getDouble(block, position));
    } else if (isVarcharType(type)) {
        Type originalType = originalColumnTypes.get(destChannel);
        if (DATE.equals(originalType)) {
            SqlDate date = (SqlDate) originalType.getObjectValue(connectorSession, block, position);
            LocalDateTime ldt = LocalDateTime.ofEpochSecond(TimeUnit.DAYS.toSeconds(date.getDays()), 0, ZoneOffset.UTC);
            byte[] bytes = ldt.format(DateTimeFormatter.ISO_LOCAL_DATE).getBytes(Charsets.UTF_8);
            row.addStringUtf8(destChannel, bytes);
        } else {
            row.addString(destChannel, type.getSlice(block, position).toStringUtf8());
        }
    } else if (VARBINARY.equals(type)) {
        row.addBinary(destChannel, type.getSlice(block, position).toByteBuffer());
    } else if (type instanceof DecimalType) {
        SqlDecimal sqlDecimal = (SqlDecimal) type.getObjectValue(connectorSession, block, position);
        row.addDecimal(destChannel, sqlDecimal.toBigDecimal());
    } else {
        throw new UnsupportedOperationException("Type is not supported: " + type);
    }
}
 

private static Block serializeStruct(Type type, BlockBuilder builder, Object object) {
    if (object == null) {
        requireNonNull(builder, "parent builder is null").appendNull();
        return null;
    }

    List<Type> typeParameters = type.getTypeParameters();
    EthBlock.TransactionObject structData = (EthBlock.TransactionObject) object;
    boolean builderSynthesized = false;
    if (builder == null) {
        builderSynthesized = true;
        builder = type.createBlockBuilder(new BlockBuilderStatus(), 1);
    }
    BlockBuilder currentBuilder = builder.beginBlockEntry();

    ImmutableList.Builder<Supplier> lstBuilder = ImmutableList.builder();
    lstBuilder.add(structData::getHash);
    lstBuilder.add(structData::getNonce);
    lstBuilder.add(structData::getBlockHash);
    lstBuilder.add(structData::getBlockNumber);
    lstBuilder.add(structData::getTransactionIndex);
    lstBuilder.add(structData::getFrom);
    lstBuilder.add(structData::getTo);
    lstBuilder.add(structData::getValue);
    lstBuilder.add(structData::getGas);
    lstBuilder.add(structData::getGasPrice);
    lstBuilder.add(structData::getInput);
    ImmutableList<Supplier> txColumns = lstBuilder.build();

    for (int i = 0; i < typeParameters.size(); i++) {
        serializeObject(typeParameters.get(i), currentBuilder, txColumns.get(i).get());
    }

    builder.closeEntry();
    if (builderSynthesized) {
        return (Block) type.getObject(builder, 0);
    } else {
        return null;
    }
}
 

private static Block serializeMap(Type type, BlockBuilder builder, Object object) {
    Map<?, ?> map = (Map) object;
    if (map == null) {
        requireNonNull(builder, "parent builder is null").appendNull();
        return null;
    }

    List<Type> typeParameters = type.getTypeParameters();
    checkArgument(typeParameters.size() == 2, "map must have exactly 2 type parameter");
    Type keyType = typeParameters.get(0);
    Type valueType = typeParameters.get(1);
    boolean builderSynthesized = false;

    if (builder == null) {
        builderSynthesized = true;
        builder = type.createBlockBuilder(new BlockBuilderStatus(), 1);
    }
    BlockBuilder currentBuilder = builder.beginBlockEntry();

    for (Map.Entry<?, ?> entry : map.entrySet()) {
        // Hive skips map entries with null keys
        if (entry.getKey() != null) {
            serializeObject(keyType, currentBuilder, entry.getKey());
            serializeObject(valueType, currentBuilder, entry.getValue());
        }
    }

    builder.closeEntry();
    if (builderSynthesized) {
        return (Block) type.getObject(builder, 0);
    } else {
        return null;
    }
}
 

private static Block serializeList(Type type, BlockBuilder builder, Object object) {
    List<?> list = (List) object;
    if (list == null) {
        requireNonNull(builder, "parent builder is null").appendNull();
        return null;
    }

    List<Type> typeParameters = type.getTypeParameters();
    checkArgument(typeParameters.size() == 1, "list must have exactly 1 type parameter");
    Type elementType = typeParameters.get(0);

    BlockBuilder currentBuilder;
    if (builder != null) {
        currentBuilder = builder.beginBlockEntry();
    } else {
        currentBuilder = elementType.createBlockBuilder(new BlockBuilderStatus(), list.size());
    }

    for (Object element : list) {
        serializeObject(elementType, currentBuilder, element);
    }

    if (builder != null) {
        builder.closeEntry();
        return null;
    } else {
        Block resultBlock = currentBuilder.build();
        return resultBlock;
    }
}
 

private static Block serializeObject(Type type, BlockBuilder builder, Object object) {
    if (!isStructuralType(type)) {
        serializePrimitive(type, builder, object);
        return null;
    } else if (isArrayType(type)) {
        return serializeList(type, builder, object);
    } else if (isMapType(type)) {
        return serializeMap(type, builder, object);
    } else if (isRowType(type)) {
        return serializeStruct(type, builder, object);
    }
    throw new RuntimeException("Unknown object type: " + type);
}
 

@Override
public Object getObjectValue(ConnectorSession session, Block block, int position)
{
    // call is null in case position is out of bounds
    checkArgument(block.isNull(position), "Expected NULL value for UnknownType");
    return null;
}
 

@Override
public int compareTo(Block leftBlock, int leftPosition, Block rightBlock, int rightPosition)
{
    // Check that the position is valid
    checkArgument(leftBlock.isNull(leftPosition), "Expected NULL value for UnknownType");
    checkArgument(rightBlock.isNull(rightPosition), "Expected NULL value for UnknownType");
    return 0;
}
 

@Override
public long hash(Block block, int position)
{
    // Check that the position is valid
    checkArgument(block.isNull(position), "Expected NULL value for UnknownType");
    return 0;
}
 

public ParaflowPageSource(
        ParquetReader parquetReader,
        ParquetDataSource dataSource,
        MessageType fileSchema,
        MessageType requestedSchema,
        long totalBytes,
        List<ParaflowColumnHandle> columns,
        TypeManager typeManager)
{
    checkArgument(totalBytes >= 0, "totalBytes is negative");

    this.parquetReader = requireNonNull(parquetReader, "parquetReader is null");
    this.dataSource = requireNonNull(dataSource, "dataSource is null");
    this.fileSchema = requireNonNull(fileSchema, "fileSchema is null");
    this.requestedSchema = requireNonNull(requestedSchema, "requestedSchema is null");
    this.totalBytes = totalBytes;

    this.columnSize = columns.size();
    this.constantBlocks = new Block[columnSize];
    ImmutableList.Builder<String> namesBuilder = ImmutableList.builder();
    ImmutableList.Builder<Type> typesBuilder = ImmutableList.builder();
    for (int columnIndex = 0; columnIndex < columnSize; columnIndex++) {
        ParaflowColumnHandle column = columns.get(columnIndex);
        String name = column.getName();
        Type type = typeManager.getType(column.getType().getTypeSignature());

        namesBuilder.add(name);
        typesBuilder.add(type);

        if (getParquetType(column, fileSchema) == null) {
            constantBlocks[columnIndex] = RunLengthEncodedBlock.create(type, null, MAX_VECTOR_LENGTH);
        }
    }
    columnNames = namesBuilder.build();
    types = typesBuilder.build();
}
 

public Block readStruct(Type type, List<String> path)
        throws IOException
{
    List<TypeSignatureParameter> parameters = type.getTypeSignature().getParameters();
    Block[] blocks = new Block[parameters.size()];
    for (int i = 0; i < parameters.size(); i++) {
        NamedTypeSignature namedTypeSignature = parameters.get(i).getNamedTypeSignature();
        Type fieldType = typeManager.getType(namedTypeSignature.getTypeSignature());
        String name = namedTypeSignature.getName();
        path.add(name);
        Optional<RichColumnDescriptor> columnDescriptor = getDescriptor(fileSchema, requestedSchema, path);
        if (!columnDescriptor.isPresent()) {
            path.remove(name);
            blocks[i] = RunLengthEncodedBlock.create(type, null, batchSize);
            continue;
        }

        if (ROW.equals(fieldType.getTypeSignature().getBase())) {
            blocks[i] = readStruct(fieldType, path);
        }
        else {
            blocks[i] = readPrimitive(columnDescriptor.get(), fieldType);
        }
        path.remove(name);
    }

    int blockSize = blocks[0].getPositionCount();
    int[] offsets = new int[batchSize];
    for (int i = 0; i < batchSize; i++) {
        offsets[i] = (i + 1) * parameters.size();
    }
    return new RowBlock(0, blockSize, new boolean[blockSize], offsets, blocks);
}
 

/**
 * Encodes the given list into a Block.
 *
 * @param elementType Element type of the array
 * @param bytes hbase byte[]
 * @return Presto Block
 */
public static Block getBlockFromArray(Type elementType, byte[] bytes)
        throws IOException
{
    List<?> array = MAPPER.readValue(bytes, List.class);
    BlockBuilder builder = elementType.createBlockBuilder(null, array.size());
    for (Object item : array) {
        writeObject(builder, elementType, item);
    }
    return builder.build();
}
 

/**
 * Given the map type and Presto Block, decodes the Block into a map of values.
 *
 * @param type Map type
 * @param block Map block
 * @return List of values
 */
private static Map<Object, Object> getMapFromBlock(Type type, Block block)
{
    Map<Object, Object> map = new HashMap<>(block.getPositionCount() / 2);
    Type keyType = Types.getKeyType(type);
    Type valueType = Types.getValueType(type);
    for (int i = 0; i < block.getPositionCount(); i += 2) {
        map.put(readObject(keyType, block, i), readObject(valueType, block, i + 1));
    }
    return map;
}
 

/**
 * Given the array element type and Presto Block, decodes the Block into a list of values.
 *
 * @param elementType Array element type
 * @param block Array block
 * @return List of values
 */
private static List<Object> getArrayFromBlock(Type elementType, Block block)
{
    ImmutableList.Builder<Object> arrayBuilder = ImmutableList.builder();
    for (int i = 0; i < block.getPositionCount(); ++i) {
        arrayBuilder.add(readObject(elementType, block, i));
    }
    return arrayBuilder.build();
}
 

@Override
public void appendTo(Block block, int position, BlockBuilder blockBuilder)
{
    blockBuilder.appendNull();
}
 

private void appendTo(Type type, Object value, BlockBuilder output)
{
    if (value == null) {
        output.appendNull();
        return;
    }
    Class<?> javaType = type.getJavaType();
    try {
        if (javaType == boolean.class) {
            type.writeBoolean(output, (Boolean) value);
        }
        else if (javaType == long.class) {
            if (type.equals(BIGINT)) {
                type.writeLong(output, ((Number) value).longValue());
            }
            else if (type.equals(INTEGER)) {
                type.writeLong(output, ((Number) value).intValue());
            }
            else if (type.equals(DATE)) {
                long millis = (Long) value;
                type.writeLong(output, TimeUnit.MILLISECONDS.toDays(millis));
            }
            else if (type.equals(TIME)) {
                type.writeLong(output, (Long) value);
            }
            else if (type.equals(TIMESTAMP)) {
                type.writeLong(output, (Long) value);
            }
            else {
                throw new PrestoException(GENERIC_INTERNAL_ERROR, "Unhandled type for " + javaType.getSimpleName() + ":" + type.getTypeSignature());
            }
        }
        else if (javaType == double.class) {
            type.writeDouble(output, ((Number) value).doubleValue());
        }
        else if (javaType == Slice.class) {
            writeSlice(output, type, value);
        }
        else if (javaType == Block.class) {
            writeBlock(output, type, value);
        }
        else {
            throw new PrestoException(GENERIC_INTERNAL_ERROR, "Unhandled type for " + javaType.getSimpleName() + ":" + type.getTypeSignature());
        }
    }
    catch (ClassCastException ignore) {
        // returns null instead of raising exception
        output.appendNull();
    }
}
 

/**
 * Gets the next page of data.  This method is allowed to return null.
 */
@Override
public Page getNextPage()
{
    try {
        batchId++;
        long start = System.nanoTime();

        int batchSize = parquetReader.nextBatch();

        readTimeNanos += System.nanoTime() - start;

        if (closed || batchSize <= 0) {
            close();
            return null;
        }

        Block[] blocks = new Block[columnSize];
        for (int fieldId = 0; fieldId < blocks.length; fieldId++) {
            if (constantBlocks[fieldId] != null) {
                blocks[fieldId] = constantBlocks[fieldId].getRegion(0, batchSize);
            }
            else {
                Type type = types.get(fieldId);
                int fieldIndex = getFieldIndex(fileSchema, columnNames.get(fieldId));
                if (fieldIndex == -1) {
                    blocks[fieldId] = RunLengthEncodedBlock.create(type, null, batchSize);
                    continue;
                }

                List<String> path = new ArrayList<>();
                path.add(fileSchema.getFields().get(fieldIndex).getName());
                if (StandardTypes.ROW.equals(type.getTypeSignature().getBase())) {
                    blocks[fieldId] = parquetReader.readStruct(type, path);
                }
                else {
                    Optional<RichColumnDescriptor> descriptor = ParquetTypeUtils.getDescriptor(fileSchema, requestedSchema, path);
                    if (descriptor.isPresent()) {
                        blocks[fieldId] = new LazyBlock(batchSize, new ParquetBlockLoader(descriptor.get(), type));
                    }
                    else {
                        blocks[fieldId] = RunLengthEncodedBlock.create(type, null, batchSize);
                    }
                }
            }
        }
        return new Page(batchSize, blocks);
    }
    catch (IOException e) {
        closeWithSupression(e);
        throw new ParaflowCursorException();
    }
}
 

public Block readPrimitive(ColumnDescriptor columnDescriptor, Type type)
        throws IOException
{
    return readPrimitive(columnDescriptor, type, new IntArrayList());
}
 

@Override
public Slice getSlice(Block block, int position)
{
    return block.getSlice(position, 0, block.getSliceLength(position));
}