Java Code Examples for io.prestosql.spi.block.BlockBuilder#closeEntry()

@Override
public void setField(int index, StdData value) {
  // TODO: This is not the right way to get this object. The status should be passed in from the invocation of the
  // function and propagated to here. See PRESTO-1359 for more details.
  BlockBuilderStatus blockBuilderStatus = new PageBuilderStatus().createBlockBuilderStatus();
  BlockBuilder mutable = _rowType.createBlockBuilder(blockBuilderStatus, 1);
  BlockBuilder rowBlockBuilder = mutable.beginBlockEntry();
  int i = 0;
  for (RowType.Field field : _rowType.getFields()) {
    if (i == index) {
      ((PrestoData) value).writeToBlock(rowBlockBuilder);
    } else {
      if (_block == null) {
        rowBlockBuilder.appendNull();
      } else {
        field.getType().appendTo(_block, i, rowBlockBuilder);
      }
    }
    i++;
  }
  mutable.closeEntry();
  _block = _rowType.getObject(mutable.build(), 0);
}
 

static Block getBlockFromMap(Type mapType, Map<?, ?> map)
{
    // on functions like COUNT() the Type won't be a MapType
    if (!(mapType instanceof MapType)) {
        return null;
    }
    Type keyType = mapType.getTypeParameters().get(0);
    Type valueType = mapType.getTypeParameters().get(1);

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

/**
 * Recursive helper function used by {@link AccumuloRowSerializer#getBlockFromArray} and
 * {@link AccumuloRowSerializer#getBlockFromMap} to add the given object to the given block
 * builder. Supports nested complex types!
 *
 * @param builder Block builder
 * @param type Presto type
 * @param obj Object to write to the block builder
 */
static void writeObject(BlockBuilder builder, Type type, Object obj)
{
    if (Types.isArrayType(type)) {
        BlockBuilder arrayBldr = builder.beginBlockEntry();
        Type elementType = Types.getElementType(type);
        for (Object item : (List<?>) obj) {
            writeObject(arrayBldr, elementType, item);
        }
        builder.closeEntry();
    }
    else if (Types.isMapType(type)) {
        BlockBuilder mapBlockBuilder = builder.beginBlockEntry();
        for (Entry<?, ?> entry : ((Map<?, ?>) obj).entrySet()) {
            writeObject(mapBlockBuilder, Types.getKeyType(type), entry.getKey());
            writeObject(mapBlockBuilder, Types.getValueType(type), entry.getValue());
        }
        builder.closeEntry();
    }
    else {
        TypeUtils.writeNativeValue(type, builder, obj);
    }
}
 

@Override
public void serialize(T state, BlockBuilder out)
{
    BlockBuilder blockBuilder = out.beginBlockEntry();
    BOOLEAN.writeBoolean(blockBuilder, state.isFirstNull());
    BOOLEAN.writeBoolean(blockBuilder, state.isSecondNull());

    if (state.isFirstNull()) {
        blockBuilder.appendNull();
    }
    else {
        writeFirstField(blockBuilder, state);
    }

    if (state.isSecondNull()) {
        blockBuilder.appendNull();
    }
    else {
        secondType.appendTo(state.getSecondBlock(), state.getSecondPosition(), blockBuilder);
    }
    out.closeEntry();
}
 

private static Block createChannel(int positionCount, int arraySize, ArrayType arrayType)
{
    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 (arrayType.getElementType().getJavaType() == long.class) {
                arrayType.getElementType().writeLong(entryBuilder, ThreadLocalRandom.current().nextLong());
            }
            else {
                throw new UnsupportedOperationException();
            }
        }
        blockBuilder.closeEntry();
    }
    return blockBuilder.build();
}
 

private void createBlockBuilderWithValues(Map<String, Long> map, BlockBuilder mapBlockBuilder)
{
    if (map == null) {
        mapBlockBuilder.appendNull();
    }
    else {
        BlockBuilder elementBlockBuilder = mapBlockBuilder.beginBlockEntry();
        for (Map.Entry<String, Long> entry : map.entrySet()) {
            VARCHAR.writeSlice(elementBlockBuilder, utf8Slice(entry.getKey()));
            if (entry.getValue() == null) {
                elementBlockBuilder.appendNull();
            }
            else {
                BIGINT.writeLong(elementBlockBuilder, entry.getValue());
            }
        }
        mapBlockBuilder.closeEntry();
    }
}
 

private static void writePartitionSummaries(BlockBuilder blockBuilder, List<PartitionFieldSummary> summaries, PartitionSpec partitionSpec)
{
    for (int i = 0; i < summaries.size(); i++) {
        PartitionFieldSummary summary = summaries.get(i);
        PartitionField field = partitionSpec.fields().get(i);
        Type nestedType = partitionSpec.partitionType().fields().get(i).type();

        BlockBuilder rowBuilder = blockBuilder.beginBlockEntry();
        BOOLEAN.writeBoolean(rowBuilder, summary.containsNull());
        VARCHAR.writeString(rowBuilder, field.transform().toHumanString(
                Conversions.fromByteBuffer(nestedType, summary.lowerBound())));
        VARCHAR.writeString(rowBuilder, field.transform().toHumanString(
                Conversions.fromByteBuffer(nestedType, summary.upperBound())));
        blockBuilder.closeEntry();
    }
}
 

@Override
public void writeSlice(BlockBuilder blockBuilder, Slice value, int offset, int length)
{
    if (length != INT128_BYTES) {
        throw new IllegalStateException("Expected entry size to be exactly " + INT128_BYTES + " but was " + length);
    }
    blockBuilder.writeLong(value.getLong(offset));
    blockBuilder.writeLong(value.getLong(offset + SIZE_OF_LONG));
    blockBuilder.closeEntry();
}
 

private static Block createChannel(int positionCount, int arraySize)
{
    ArrayType arrayType = new ArrayType(BIGINT);

    BlockBuilder blockBuilder = arrayType.createBlockBuilder(null, positionCount);
    for (int position = 0; position < positionCount; position++) {
        BlockBuilder entryBuilder = blockBuilder.beginBlockEntry();
        for (int i = 0; i < arraySize; i++) {
            arrayType.getElementType().writeLong(entryBuilder, ThreadLocalRandom.current().nextLong());
        }
        blockBuilder.closeEntry();
    }
    return blockBuilder.build();
}
 

private static Block serializeMap(BlockBuilder parentBlockBuilder, Object value, Type type, String columnName)
{
    if (value == null) {
        checkState(parentBlockBuilder != null, "parentBlockBuilder is null");
        parentBlockBuilder.appendNull();
        return null;
    }

    Map<?, ?> map = (Map<?, ?>) value;
    List<Type> typeParameters = type.getTypeParameters();
    Type keyType = typeParameters.get(0);
    Type valueType = typeParameters.get(1);

    BlockBuilder blockBuilder;
    if (parentBlockBuilder != null) {
        blockBuilder = parentBlockBuilder;
    }
    else {
        blockBuilder = type.createBlockBuilder(null, 1);
    }

    BlockBuilder entryBuilder = blockBuilder.beginBlockEntry();
    for (Map.Entry<?, ?> entry : map.entrySet()) {
        if (entry.getKey() != null) {
            keyType.writeSlice(entryBuilder, truncateToLength(utf8Slice(entry.getKey().toString()), keyType));
            serializeObject(entryBuilder, entry.getValue(), valueType, columnName);
        }
    }
    blockBuilder.closeEntry();

    if (parentBlockBuilder == null) {
        return blockBuilder.getObject(0, Block.class);
    }
    return null;
}
 

private static Block getColumnMetricBlock(RowType columnMetricType, Object min, Object max, Long nullCount)
{
    BlockBuilder rowBlockBuilder = columnMetricType.createBlockBuilder(null, 1);
    BlockBuilder builder = rowBlockBuilder.beginBlockEntry();
    List<RowType.Field> fields = columnMetricType.getFields();
    TypeUtils.writeNativeValue(fields.get(0).getType(), builder, min);
    TypeUtils.writeNativeValue(fields.get(1).getType(), builder, max);
    TypeUtils.writeNativeValue(fields.get(2).getType(), builder, nullCount);

    rowBlockBuilder.closeEntry();
    return columnMetricType.getObject(rowBlockBuilder, 0);
}
 

public void appendBigintArray(Iterable<Long> values)
{
    BlockBuilder column = nextColumn();
    BlockBuilder array = column.beginBlockEntry();
    for (Long value : values) {
        BIGINT.writeLong(array, value);
    }
    column.closeEntry();
}
 

@TypeParameter("K")
@TypeParameter("V")
@SqlType("array(row(K,V))")
public Block mapFromEntries(
        @TypeParameter("row(K,V)") RowType rowType,
        @SqlType("map(K,V)") Block block)
{
    verify(rowType.getTypeParameters().size() == 2);
    verify(block.getPositionCount() % 2 == 0);

    Type keyType = rowType.getTypeParameters().get(0);
    Type valueType = rowType.getTypeParameters().get(1);
    ArrayType arrayType = new ArrayType(rowType);

    if (pageBuilder.isFull()) {
        pageBuilder.reset();
    }

    int entryCount = block.getPositionCount() / 2;
    BlockBuilder blockBuilder = pageBuilder.getBlockBuilder(0);
    BlockBuilder entryBuilder = blockBuilder.beginBlockEntry();
    for (int i = 0; i < entryCount; i++) {
        BlockBuilder rowBuilder = entryBuilder.beginBlockEntry();
        keyType.appendTo(block, 2 * i, rowBuilder);
        valueType.appendTo(block, 2 * i + 1, rowBuilder);
        entryBuilder.closeEntry();
    }

    blockBuilder.closeEntry();
    pageBuilder.declarePosition();
    return arrayType.getObject(blockBuilder, blockBuilder.getPositionCount() - 1);
}
 

public static Block mapBlockOf(Type keyType, Type valueType, Object[] keys, Object[] values)
{
    checkArgument(keys.length == values.length, "keys/values must have the same length");
    MapType mapType = mapType(keyType, valueType);
    BlockBuilder blockBuilder = mapType.createBlockBuilder(null, 10);
    BlockBuilder singleMapBlockWriter = blockBuilder.beginBlockEntry();
    for (int i = 0; i < keys.length; i++) {
        Object key = keys[i];
        Object value = values[i];
        appendToBlockBuilder(keyType, key, singleMapBlockWriter);
        appendToBlockBuilder(valueType, value, singleMapBlockWriter);
    }
    blockBuilder.closeEntry();
    return mapType.getObject(blockBuilder, 0);
}
 

public static Block mapBlockOf(Type keyType, Type valueType, Object key, Object value)
{
    MapType mapType = mapType(keyType, valueType);
    BlockBuilder blockBuilder = mapType.createBlockBuilder(null, 10);
    BlockBuilder singleMapBlockWriter = blockBuilder.beginBlockEntry();
    appendToBlockBuilder(keyType, key, singleMapBlockWriter);
    appendToBlockBuilder(valueType, value, singleMapBlockWriter);
    blockBuilder.closeEntry();
    return mapType.getObject(blockBuilder, 0);
}
 

@Override
public void decodeValueInto(int depth, BlockBuilder builder, Slice slice, int offset, int length)
        throws RcFileCorruptionException
{
    byte separator = getSeparator(depth);
    int end = offset + length;

    BlockBuilder structBuilder = builder.beginBlockEntry();
    int elementOffset = offset;
    int fieldIndex = 0;
    while (offset < end && (!lastColumnTakesRest || fieldIndex < structFields.size())) {
        byte currentByte = slice.getByte(offset);
        if (currentByte == separator) {
            decodeElementValueInto(depth, fieldIndex, structBuilder, slice, elementOffset, offset - elementOffset);
            elementOffset = offset + 1;
            fieldIndex++;
        }
        else if (isEscapeByte(currentByte) && offset + 1 < length) {
            // ignore the char after escape_char
            offset++;
        }
        offset++;
    }
    decodeElementValueInto(depth, fieldIndex, structBuilder, slice, elementOffset, end - elementOffset);
    fieldIndex++;

    // missing fields are null
    while (fieldIndex < structFields.size()) {
        structBuilder.appendNull();
        fieldIndex++;
    }

    builder.closeEntry();
}
 

public void appendIntegerBigintMap(Map<Integer, Long> values)
{
    BlockBuilder column = nextColumn();
    BlockBuilder map = column.beginBlockEntry();
    values.forEach((key, value) -> {
        INTEGER.writeLong(map, key);
        BIGINT.writeLong(map, value);
    });
    column.closeEntry();
}
 

@UsedByGeneratedCode
public static Block mapConcat(MapType mapType, Object state, Block[] maps)
{
    int entries = 0;
    int lastMapIndex = maps.length - 1;
    int firstMapIndex = lastMapIndex;
    for (int i = 0; i < maps.length; i++) {
        entries += maps[i].getPositionCount();
        if (maps[i].getPositionCount() > 0) {
            lastMapIndex = i;
            firstMapIndex = min(firstMapIndex, i);
        }
    }
    if (lastMapIndex == firstMapIndex) {
        return maps[lastMapIndex];
    }

    PageBuilder pageBuilder = (PageBuilder) state;
    if (pageBuilder.isFull()) {
        pageBuilder.reset();
    }

    // TODO: we should move TypedSet into user state as well
    Type keyType = mapType.getKeyType();
    Type valueType = mapType.getValueType();
    TypedSet typedSet = new TypedSet(keyType, entries / 2, FUNCTION_NAME);
    BlockBuilder mapBlockBuilder = pageBuilder.getBlockBuilder(0);
    BlockBuilder blockBuilder = mapBlockBuilder.beginBlockEntry();

    // the last map
    Block map = maps[lastMapIndex];
    for (int i = 0; i < map.getPositionCount(); i += 2) {
        typedSet.add(map, i);
        keyType.appendTo(map, i, blockBuilder);
        valueType.appendTo(map, i + 1, blockBuilder);
    }
    // the map between the last and the first
    for (int idx = lastMapIndex - 1; idx > firstMapIndex; idx--) {
        map = maps[idx];
        for (int i = 0; i < map.getPositionCount(); i += 2) {
            if (!typedSet.contains(map, i)) {
                typedSet.add(map, i);
                keyType.appendTo(map, i, blockBuilder);
                valueType.appendTo(map, i + 1, blockBuilder);
            }
        }
    }
    // the first map
    map = maps[firstMapIndex];
    for (int i = 0; i < map.getPositionCount(); i += 2) {
        if (!typedSet.contains(map, i)) {
            keyType.appendTo(map, i, blockBuilder);
            valueType.appendTo(map, i + 1, blockBuilder);
        }
    }

    mapBlockBuilder.closeEntry();
    pageBuilder.declarePosition();
    return mapType.getObject(mapBlockBuilder, mapBlockBuilder.getPositionCount() - 1);
}
 

public static void writeLongTimestamp(BlockBuilder blockBuilder, long epochMicros, int fraction)
{
    blockBuilder.writeLong(epochMicros);
    blockBuilder.writeInt(fraction);
    blockBuilder.closeEntry();
}
 

@TypeParameter("K")
@TypeParameter("V")
@SqlType("map(K,array(V))")
@SqlNullable
public Block multimapFromEntries(
        @TypeParameter("map(K,array(V))") MapType mapType,
        @SqlType("array(row(K,V))") Block block)
{
    Type keyType = mapType.getKeyType();
    Type valueType = ((ArrayType) mapType.getValueType()).getElementType();
    RowType rowType = RowType.anonymous(ImmutableList.of(keyType, valueType));

    if (pageBuilder.isFull()) {
        pageBuilder.reset();
    }

    int entryCount = block.getPositionCount();
    if (entryCount > entryIndicesList.length) {
        initializeEntryIndicesList(entryCount);
    }
    TypedSet keySet = new TypedSet(keyType, entryCount, NAME);

    for (int i = 0; i < entryCount; i++) {
        if (block.isNull(i)) {
            clearEntryIndices(keySet.size());
            throw new PrestoException(INVALID_FUNCTION_ARGUMENT, "map entry cannot be null");
        }
        Block rowBlock = rowType.getObject(block, i);

        if (rowBlock.isNull(0)) {
            clearEntryIndices(keySet.size());
            throw new PrestoException(INVALID_FUNCTION_ARGUMENT, "map key cannot be null");
        }

        if (keySet.contains(rowBlock, 0)) {
            entryIndicesList[keySet.positionOf(rowBlock, 0)].add(i);
        }
        else {
            keySet.add(rowBlock, 0);
            entryIndicesList[keySet.size() - 1].add(i);
        }
    }

    BlockBuilder multimapBlockBuilder = pageBuilder.getBlockBuilder(0);
    BlockBuilder singleMapWriter = multimapBlockBuilder.beginBlockEntry();
    for (int i = 0; i < keySet.size(); i++) {
        keyType.appendTo(rowType.getObject(block, entryIndicesList[i].getInt(0)), 0, singleMapWriter);
        BlockBuilder singleArrayWriter = singleMapWriter.beginBlockEntry();
        for (int entryIndex : entryIndicesList[i]) {
            valueType.appendTo(rowType.getObject(block, entryIndex), 1, singleArrayWriter);
        }
        singleMapWriter.closeEntry();
    }

    multimapBlockBuilder.closeEntry();
    pageBuilder.declarePosition();
    clearEntryIndices(keySet.size());
    return mapType.getObject(multimapBlockBuilder, multimapBlockBuilder.getPositionCount() - 1);
}