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

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

public static void output(ArrayType outputType, MinMaxNState state, BlockBuilder out)
{
    TypedHeap heap = state.getTypedHeap();
    if (heap == null || heap.isEmpty()) {
        out.appendNull();
        return;
    }

    Type elementType = outputType.getElementType();

    BlockBuilder reversedBlockBuilder = elementType.createBlockBuilder(null, heap.getCapacity());
    long startSize = heap.getEstimatedSize();
    heap.popAll(reversedBlockBuilder);
    state.addMemoryUsage(heap.getEstimatedSize() - startSize);

    BlockBuilder arrayBlockBuilder = out.beginBlockEntry();
    for (int i = reversedBlockBuilder.getPositionCount() - 1; i >= 0; i--) {
        elementType.appendTo(reversedBlockBuilder, i, arrayBlockBuilder);
    }
    out.closeEntry();
}
 

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

public static void output(ArrayType outputType, MinMaxByNState state, BlockBuilder out)
{
    TypedKeyValueHeap heap = state.getTypedKeyValueHeap();
    if (heap == null || heap.isEmpty()) {
        out.appendNull();
        return;
    }

    Type elementType = outputType.getElementType();

    BlockBuilder arrayBlockBuilder = out.beginBlockEntry();
    BlockBuilder reversedBlockBuilder = elementType.createBlockBuilder(null, heap.getCapacity());
    long startSize = heap.getEstimatedSize();
    heap.popAll(reversedBlockBuilder);
    state.addMemoryUsage(heap.getEstimatedSize() - startSize);

    for (int i = reversedBlockBuilder.getPositionCount() - 1; i >= 0; i--) {
        elementType.appendTo(reversedBlockBuilder, i, arrayBlockBuilder);
    }
    out.closeEntry();
}
 

@Override
public void serialize(MinMaxNState state, BlockBuilder out)
{
    TypedHeap heap = state.getTypedHeap();
    if (heap == null) {
        out.appendNull();
        return;
    }

    BlockBuilder blockBuilder = out.beginBlockEntry();
    BIGINT.writeLong(blockBuilder, heap.getCapacity());
    BlockBuilder elements = blockBuilder.beginBlockEntry();
    heap.writeAll(elements);
    blockBuilder.closeEntry();

    out.closeEntry();
}
 

@OutputFunction("array(double)")
public static void output(@AggregationState DigestAndPercentileArrayState state, BlockBuilder out)
{
    QuantileDigest digest = state.getDigest();
    List<Double> percentiles = state.getPercentiles();

    if (percentiles == null || digest == null) {
        out.appendNull();
        return;
    }

    BlockBuilder blockBuilder = out.beginBlockEntry();

    for (int i = 0; i < percentiles.size(); i++) {
        Double percentile = percentiles.get(i);
        DOUBLE.writeDouble(blockBuilder, sortableLongToDouble(digest.getQuantile(percentile)));
    }

    out.closeEntry();
}
 

public static BlockBuilder createBlockBuilderWithValues(Slice[][] expectedValues)
{
    BlockBuilder blockBuilder = createBlockBuilder(null, 100, 100);
    for (Slice[] expectedValue : expectedValues) {
        if (expectedValue == null) {
            blockBuilder.appendNull();
        }
        else {
            BlockBuilder entryBuilder = blockBuilder.beginBlockEntry();
            for (Slice v : expectedValue) {
                if (v == null) {
                    entryBuilder.appendNull();
                }
                else {
                    VARCHAR.writeSlice(entryBuilder, v);
                }
            }
            blockBuilder.closeEntry();
        }
    }
    return blockBuilder;
}
 

@Override
public void decode(SearchHit hit, Supplier<Object> getter, BlockBuilder output)
{
    Object data = getter.get();

    if (data == null) {
        output.appendNull();
    }
    else if (data instanceof Map) {
        BlockBuilder row = output.beginBlockEntry();
        for (int i = 0; i < decoders.size(); i++) {
            String field = fieldNames.get(i);
            decoders.get(i).decode(hit, () -> getField((Map<String, Object>) data, field), row);
        }
        output.closeEntry();
    }
    else {
        throw new PrestoException(TYPE_MISMATCH, format("Expected object for field '%s' of type ROW: %s [%s]", path, data, data.getClass().getSimpleName()));
    }
}
 

@OutputFunction("array(bigint)")
public static void output(@AggregationState DigestAndPercentileArrayState state, BlockBuilder out)
{
    QuantileDigest digest = state.getDigest();
    List<Double> percentiles = state.getPercentiles();

    if (percentiles == null || digest == null) {
        out.appendNull();
        return;
    }

    BlockBuilder blockBuilder = out.beginBlockEntry();

    for (int i = 0; i < percentiles.size(); i++) {
        Double percentile = percentiles.get(i);
        BIGINT.writeLong(blockBuilder, digest.getQuantile(percentile));
    }

    out.closeEntry();
}
 

@Override
public void put(StdData key, StdData value) {
  BlockBuilder mutable = _mapType.createBlockBuilder(new PageBuilderStatus().createBlockBuilderStatus(), 1);
  BlockBuilder entryBuilder = mutable.beginBlockEntry();
  Object prestoKey = ((PlatformData) key).getUnderlyingData();
  int valuePosition = seekKey(prestoKey);
  for (int i = 0; i < _block.getPositionCount(); i += 2) {
    // Write the current key to the map
    _keyType.appendTo(_block, i, entryBuilder);
    // Find out if we need to change the corresponding value
    if (i == valuePosition - 1) {
      // Use the user-supplied value
      ((PrestoData) value).writeToBlock(entryBuilder);
    } else {
      // Use the existing value in original _block
      _valueType.appendTo(_block, i + 1, entryBuilder);
    }
  }
  if (valuePosition == -1) {
    ((PrestoData) key).writeToBlock(entryBuilder);
    ((PrestoData) value).writeToBlock(entryBuilder);
  }

  mutable.closeEntry();
  _block = ((MapType) _mapType).getObject(mutable.build(), 0);
}
 

@UsedByGeneratedCode
public static Block zip(List<Type> types, Block... arrays)
{
    int biggestCardinality = 0;
    for (Block array : arrays) {
        biggestCardinality = Math.max(biggestCardinality, array.getPositionCount());
    }
    RowType rowType = RowType.anonymous(types);
    BlockBuilder outputBuilder = rowType.createBlockBuilder(null, biggestCardinality);
    for (int outputPosition = 0; outputPosition < biggestCardinality; outputPosition++) {
        BlockBuilder rowBuilder = outputBuilder.beginBlockEntry();
        for (int fieldIndex = 0; fieldIndex < arrays.length; fieldIndex++) {
            if (arrays[fieldIndex].getPositionCount() <= outputPosition) {
                rowBuilder.appendNull();
            }
            else {
                types.get(fieldIndex).appendTo(arrays[fieldIndex], outputPosition, rowBuilder);
            }
        }
        outputBuilder.closeEntry();
    }
    return outputBuilder.build();
}
 

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

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 if (arrayType.getElementType().getJavaType() == double.class) {
                arrayType.getElementType().writeDouble(entryBuilder, ThreadLocalRandom.current().nextDouble());
            }
            else if (arrayType.getElementType().getJavaType() == boolean.class) {
                arrayType.getElementType().writeBoolean(entryBuilder, ThreadLocalRandom.current().nextBoolean());
            }
            else if (arrayType.getElementType().equals(VARCHAR)) {
                arrayType.getElementType().writeSlice(entryBuilder, Slices.utf8Slice("test_string"));
            }
            else {
                throw new UnsupportedOperationException();
            }
        }
        blockBuilder.closeEntry();
    }
    return blockBuilder.build();
}
 

/**
 * Encodes the given map into a Block.
 *
 * @param mapType Presto type of the map
 * @param map Map of key/value pairs to encode
 * @return Presto Block
 */
static Block getBlockFromMap(Type mapType, Map<?, ?> map)
{
    Type keyType = mapType.getTypeParameters().get(0);
    Type valueType = mapType.getTypeParameters().get(1);

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

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

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

public void appendVarcharArray(Iterable<String> values)
{
    BlockBuilder column = nextColumn();
    BlockBuilder array = column.beginBlockEntry();
    for (String value : values) {
        VARCHAR.writeString(array, value);
    }
    column.closeEntry();
}
 

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

    BlockBuilder mapBuilder = builder.beginBlockEntry();
    if (length > 0) {
        int elementOffset = offset;
        int keyValueSeparatorPosition = -1;
        while (offset < end) {
            byte currentByte = slice.getByte(offset);
            if (currentByte == elementSeparator) {
                decodeEntryInto(depth, mapBuilder, slice, elementOffset, offset - elementOffset, keyValueSeparatorPosition);
                elementOffset = offset + 1;
                keyValueSeparatorPosition = -1;
            }
            else if (currentByte == keyValueSeparator && keyValueSeparatorPosition == -1) {
                keyValueSeparatorPosition = offset;
            }
            else if (isEscapeByte(currentByte) && offset + 1 < length) {
                // ignore the char after escape_char
                offset++;
            }
            offset++;
        }
        decodeEntryInto(depth, mapBuilder, slice, elementOffset, offset - elementOffset, keyValueSeparatorPosition);
    }
    builder.closeEntry();
}
 

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 void writeRow(List<Object> testRow, BlockBuilder rowBlockBuilder)
{
    BlockBuilder singleRowBlockWriter = rowBlockBuilder.beginBlockEntry();
    for (Object fieldValue : testRow) {
        if (fieldValue instanceof String) {
            VARCHAR.writeSlice(singleRowBlockWriter, utf8Slice((String) fieldValue));
        }
        else {
            throw new UnsupportedOperationException();
        }
    }
    rowBlockBuilder.closeEntry();
}
 

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

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 (Map.Entry<?, ?> entry : ((Map<?, ?>) obj).entrySet()) {
            writeObject(mapBlockBuilder, Types.getKeyType(type), entry.getKey());
            writeObject(mapBlockBuilder, Types.getValueType(type), entry.getValue());
        }
        builder.closeEntry();
    }
    else {
        if (BOOLEAN.equals(type)
                || TINYINT.equals(type)
                || SMALLINT.equals(type)
                || INTEGER.equals(type)
                || BIGINT.equals(type)
                || DOUBLE.equals(type)
                || type instanceof VarcharType) {
            TypeUtils.writeNativeValue(type, builder, obj);
        }
    }
}