io.prestosql.spi.block.PageBuilderStatus Java Examples

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

private PageBuilder(int initialExpectedEntries, int maxPageBytes, List<? extends Type> types, Optional<BlockBuilder[]> templateBlockBuilders)
{
    this.types = List.copyOf(requireNonNull(types, "types is null"));

    pageBuilderStatus = new PageBuilderStatus(maxPageBytes);
    blockBuilders = new BlockBuilder[types.size()];

    if (templateBlockBuilders.isPresent()) {
        BlockBuilder[] templates = templateBlockBuilders.get();
        checkArgument(templates.length == types.size(), "Size of templates and types should match");
        for (int i = 0; i < blockBuilders.length; i++) {
            blockBuilders[i] = templates[i].newBlockBuilderLike(pageBuilderStatus.createBlockBuilderStatus());
        }
    }
    else {
        for (int i = 0; i < blockBuilders.length; i++) {
            blockBuilders[i] = types.get(i).createBlockBuilder(pageBuilderStatus.createBlockBuilderStatus(), initialExpectedEntries);
        }
    }
}
 

@Override
public BlockBuilder createBlockBuilder(BlockBuilderStatus blockBuilderStatus, int expectedEntries, int expectedBytesPerEntry)
{
    int maxBlockSizeInBytes;
    if (blockBuilderStatus == null) {
        maxBlockSizeInBytes = PageBuilderStatus.DEFAULT_MAX_PAGE_SIZE_IN_BYTES;
    }
    else {
        maxBlockSizeInBytes = blockBuilderStatus.getMaxPageSizeInBytes();
    }

    // it is guaranteed Math.min will not overflow; safe to cast
    int expectedBytes = (int) min((long) expectedEntries * expectedBytesPerEntry, maxBlockSizeInBytes);
    return new VariableWidthBlockBuilder(
            blockBuilderStatus,
            expectedBytesPerEntry == 0 ? expectedEntries : Math.min(expectedEntries, maxBlockSizeInBytes / expectedBytesPerEntry),
            expectedBytes);
}
 

@Override
public void setField(String name, StdData value) {
  BlockBuilder mutable = _rowType.createBlockBuilder(new PageBuilderStatus().createBlockBuilderStatus(), 1);
  BlockBuilder rowBlockBuilder = mutable.beginBlockEntry();
  int i = 0;
  for (RowType.Field field : _rowType.getFields()) {
    if (field.getName().isPresent() && name.equals(field.getName().get())) {
      ((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);
}
 

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

@Override
public Object getObject(int field) {
    Object fieldValue;
    Type type;
    PageBuilderStatus pageBuilderStatus = new PageBuilderStatus(MAX_BLOCK_SIZE);
    type = this.getType(field);
    fieldValue = this.getFieldValue(field);
    if (fieldValue == null || "".equals(fieldValue)) {
        return this.getType(field).createBlockBuilder(pageBuilderStatus.createBlockBuilderStatus(), 0).build();
    }

    if (type.getTypeSignature().getBase().equals(StandardTypes.ARRAY)) {
        String[] values = ((String) fieldValue).split(Constant.ARRAY_STRING_SPLITTER);
        Type elementType = type.getTypeParameters().get(0);
        BlockBuilder builder = elementType.createBlockBuilder(null, values.length);
        for (String value : values) {
            if (value != null && !"".equals(value)) {
                TypeUtils.writeNativeValue(elementType, builder, value);
            }
        }
        return builder.build();
    } else {
        throw new UnsupportedOperationException("OOPS！UNSUPPORTED TYPE：" + type.getDisplayName());
    }
}
 

@Test
public void testCapacityIncrease()
{
    Slice[] values = new Slice[100];
    for (int i = 0; i < values.length; i++) {
        values[i] = utf8Slice("a");
    }

    UnnestBlockBuilder unnestBlockBuilder = new UnnestBlockBuilder(VARCHAR);
    Block valuesBlock = createSimpleBlock(values);
    unnestBlockBuilder.resetInputBlock(valuesBlock);

    unnestBlockBuilder.startNewOutput(new PageBuilderStatus(), 20);
    unnestBlockBuilder.appendRange(0, values.length);
    assertBlock(unnestBlockBuilder.buildOutputAndFlush(), values);

    unnestBlockBuilder.clearCurrentOutput();
}
 

@Override
public final BlockBuilder createBlockBuilder(BlockBuilderStatus blockBuilderStatus, int expectedEntries, int expectedBytesPerEntry)
{
    int maxBlockSizeInBytes;
    if (blockBuilderStatus == null) {
        maxBlockSizeInBytes = PageBuilderStatus.DEFAULT_MAX_PAGE_SIZE_IN_BYTES;
    }
    else {
        maxBlockSizeInBytes = blockBuilderStatus.getMaxPageSizeInBytes();
    }
    return new LongArrayBlockBuilder(
            blockBuilderStatus,
            Math.min(expectedEntries, maxBlockSizeInBytes / Double.BYTES));
}
 

public void reset()
{
    if (isEmpty()) {
        return;
    }
    pageBuilderStatus = new PageBuilderStatus(pageBuilderStatus.getMaxPageSizeInBytes());

    declaredPositions = 0;

    for (int i = 0; i < types.size(); i++) {
        blockBuilders[i] = blockBuilders[i].newBlockBuilderLike(pageBuilderStatus.createBlockBuilderStatus());
    }
}
 

@Override
public BlockBuilder createBlockBuilder(BlockBuilderStatus blockBuilderStatus, int expectedEntries, int expectedBytesPerEntry)
{
    int maxBlockSizeInBytes;
    if (blockBuilderStatus == null) {
        maxBlockSizeInBytes = PageBuilderStatus.DEFAULT_MAX_PAGE_SIZE_IN_BYTES;
    }
    else {
        maxBlockSizeInBytes = blockBuilderStatus.getMaxPageSizeInBytes();
    }
    return new ByteArrayBlockBuilder(
            blockBuilderStatus,
            Math.min(expectedEntries, maxBlockSizeInBytes / Byte.BYTES));
}
 

@Override
public final BlockBuilder createBlockBuilder(BlockBuilderStatus blockBuilderStatus, int expectedEntries, int expectedBytesPerEntry)
{
    int maxBlockSizeInBytes;
    if (blockBuilderStatus == null) {
        maxBlockSizeInBytes = PageBuilderStatus.DEFAULT_MAX_PAGE_SIZE_IN_BYTES;
    }
    else {
        maxBlockSizeInBytes = blockBuilderStatus.getMaxPageSizeInBytes();
    }
    return new LongArrayBlockBuilder(
            blockBuilderStatus,
            Math.min(expectedEntries, maxBlockSizeInBytes / Long.BYTES));
}
 

@Override
public BlockBuilder createBlockBuilder(BlockBuilderStatus blockBuilderStatus, int expectedEntries, int expectedBytesPerEntry)
{
    int maxBlockSizeInBytes;
    if (blockBuilderStatus == null) {
        maxBlockSizeInBytes = PageBuilderStatus.DEFAULT_MAX_PAGE_SIZE_IN_BYTES;
    }
    else {
        maxBlockSizeInBytes = blockBuilderStatus.getMaxPageSizeInBytes();
    }
    return new Int96ArrayBlockBuilder(
            blockBuilderStatus,
            Math.min(expectedEntries, maxBlockSizeInBytes / getFixedSize()));
}
 

@Override
public BlockBuilder createBlockBuilder(BlockBuilderStatus blockBuilderStatus, int expectedEntries, int expectedBytesPerEntry)
{
    int maxBlockSizeInBytes;
    if (blockBuilderStatus == null) {
        maxBlockSizeInBytes = PageBuilderStatus.DEFAULT_MAX_PAGE_SIZE_IN_BYTES;
    }
    else {
        maxBlockSizeInBytes = blockBuilderStatus.getMaxPageSizeInBytes();
    }
    return new ShortArrayBlockBuilder(
            blockBuilderStatus,
            Math.min(expectedEntries, maxBlockSizeInBytes / Short.BYTES));
}
 

@Override
public BlockBuilder createBlockBuilder(BlockBuilderStatus blockBuilderStatus, int expectedEntries, int expectedBytesPerEntry)
{
    int maxBlockSizeInBytes;
    if (blockBuilderStatus == null) {
        maxBlockSizeInBytes = PageBuilderStatus.DEFAULT_MAX_PAGE_SIZE_IN_BYTES;
    }
    else {
        maxBlockSizeInBytes = blockBuilderStatus.getMaxPageSizeInBytes();
    }
    return new ByteArrayBlockBuilder(
            blockBuilderStatus,
            Math.min(expectedEntries, maxBlockSizeInBytes / Byte.BYTES));
}
 

@Override
public final BlockBuilder createBlockBuilder(BlockBuilderStatus blockBuilderStatus, int expectedEntries, int expectedBytesPerEntry)
{
    int maxBlockSizeInBytes;
    if (blockBuilderStatus == null) {
        maxBlockSizeInBytes = PageBuilderStatus.DEFAULT_MAX_PAGE_SIZE_IN_BYTES;
    }
    else {
        maxBlockSizeInBytes = blockBuilderStatus.getMaxPageSizeInBytes();
    }
    return new IntArrayBlockBuilder(
            blockBuilderStatus,
            Math.min(expectedEntries, maxBlockSizeInBytes / Integer.BYTES));
}
 

@Override
public BlockBuilder createBlockBuilder(BlockBuilderStatus blockBuilderStatus, int expectedEntries, int expectedBytesPerEntry)
{
    int maxBlockSizeInBytes;
    if (blockBuilderStatus == null) {
        maxBlockSizeInBytes = PageBuilderStatus.DEFAULT_MAX_PAGE_SIZE_IN_BYTES;
    }
    else {
        maxBlockSizeInBytes = blockBuilderStatus.getMaxPageSizeInBytes();
    }
    return new LongArrayBlockBuilder(
            blockBuilderStatus,
            Math.min(expectedEntries, maxBlockSizeInBytes / getFixedSize()));
}
 

@Override
public BlockBuilder createBlockBuilder(BlockBuilderStatus blockBuilderStatus, int expectedEntries, int expectedBytesPerEntry)
{
    int maxBlockSizeInBytes;
    if (blockBuilderStatus == null) {
        maxBlockSizeInBytes = PageBuilderStatus.DEFAULT_MAX_PAGE_SIZE_IN_BYTES;
    }
    else {
        maxBlockSizeInBytes = blockBuilderStatus.getMaxPageSizeInBytes();
    }
    return new Int128ArrayBlockBuilder(
            blockBuilderStatus,
            Math.min(expectedEntries, maxBlockSizeInBytes / getFixedSize()));
}
 

@Override
public final BlockBuilder createBlockBuilder(BlockBuilderStatus blockBuilderStatus, int expectedEntries, int expectedBytesPerEntry)
{
    int maxBlockSizeInBytes;
    if (blockBuilderStatus == null) {
        maxBlockSizeInBytes = PageBuilderStatus.DEFAULT_MAX_PAGE_SIZE_IN_BYTES;
    }
    else {
        maxBlockSizeInBytes = blockBuilderStatus.getMaxPageSizeInBytes();
    }
    return new LongArrayBlockBuilder(
            blockBuilderStatus,
            Math.min(expectedEntries, maxBlockSizeInBytes / Long.BYTES));
}
 

@Override
public BlockBuilder createBlockBuilder(BlockBuilderStatus blockBuilderStatus, int expectedEntries, int expectedBytesPerEntry)
{
    int maxBlockSizeInBytes;
    if (blockBuilderStatus == null) {
        maxBlockSizeInBytes = PageBuilderStatus.DEFAULT_MAX_PAGE_SIZE_IN_BYTES;
    }
    else {
        maxBlockSizeInBytes = blockBuilderStatus.getMaxPageSizeInBytes();
    }
    return new Int96ArrayBlockBuilder(
            blockBuilderStatus,
            Math.min(expectedEntries, maxBlockSizeInBytes / getFixedSize()));
}
 

@Override
public Page getNextPage()
{
    if (finished) {
        return null;
    }
    if (resultIterator == null) {
        long start = System.nanoTime();
        resultIterator = pinotClient.createResultIterator(session, query, columnHandles);
        readTimeNanos = System.nanoTime() - start;
    }

    if (!resultIterator.hasNext()) {
        finished = true;
        return null;
    }
    long size = 0;
    while (size < PageBuilderStatus.DEFAULT_MAX_PAGE_SIZE_IN_BYTES && resultIterator.hasNext()) {
        BrokerResultRow row = resultIterator.next();
        for (int i = 0; i < decoders.size(); i++) {
            int fieldIndex = i;
            decoders.get(i).decode(() -> row.getField(fieldIndex), columnBuilders[i]);
        }
        size = Arrays.stream(columnBuilders)
                .mapToLong(BlockBuilder::getSizeInBytes)
                .sum();
    }
    completedBytes += size;
    Block[] blocks = new Block[columnBuilders.length];
    for (int i = 0; i < columnBuilders.length; i++) {
        blocks[i] = columnBuilders[i].build();
        columnBuilders[i] = columnBuilders[i].newBlockBuilderLike(null);
    }
    return new Page(blocks);
}
 

public PrestoArray(ArrayType arrayType, int expectedEntries, StdFactory stdFactory) {
  _block = null;
  _elementType = arrayType.getElementType();
  _mutable = _elementType.createBlockBuilder(new PageBuilderStatus().createBlockBuilderStatus(), expectedEntries);
  _stdFactory = stdFactory;
  _arrayType = arrayType;
}
 

@Override
public void add(StdData e) {
  if (_mutable == null) {
    _mutable = _elementType.createBlockBuilder(new PageBuilderStatus().createBlockBuilderStatus(), 1);
  }
  ((PrestoData) e).writeToBlock(_mutable);
}
 

public PrestoMap(Type mapType, StdFactory stdFactory) {
  BlockBuilder mutable = mapType.createBlockBuilder(new PageBuilderStatus().createBlockBuilderStatus(), 1);
  mutable.beginBlockEntry();
  mutable.closeEntry();
  _block = ((MapType) mapType).getObject(mutable.build(), 0);

  _keyType = ((MapType) mapType).getKeyType();
  _valueType = ((MapType) mapType).getValueType();
  _mapType = mapType;

  _stdFactory = stdFactory;
  _keyEqualsMethod = ((PrestoFactory) stdFactory).getScalarFunctionImplementation(
          ((PrestoFactory) stdFactory).resolveOperator(OperatorType.EQUAL, ImmutableList.of(_keyType, _keyType)))
      .getMethodHandle();
}
 

@Test
public void testTrimmedBlocks()
{
    int[] unnestedLengths = {1, 2, 1};

    Slice[][][] elements = column(
            array(toSlices("0.0.0", "0.0.1")),
            array(toSlices("1.0.0", "1.0.1"), toSlices("1.1.0", "1.1.1")),
            array(toSlices("2.0.0", "2.0.1")));

    Block mapBlock = createBlockBuilderWithValues(elements).build();

    // Remove the first element from the arrayBlock for testing
    int startElement = 1;

    Slice[][][] truncatedSlices = Arrays.copyOfRange(elements, startElement, elements.length - startElement + 1);
    int[] truncatedUnnestedLengths = Arrays.copyOfRange(unnestedLengths, startElement, elements.length - startElement + 1);
    Block truncatedBlock = mapBlock.getRegion(startElement, elements.length - startElement);
    assertBlock(truncatedBlock, truncatedSlices);

    Unnester mapUnnester = new MapUnnester(VARCHAR, VARCHAR);
    mapUnnester.resetInput(truncatedBlock);

    mapUnnester.startNewOutput(new PageBuilderStatus(), 20);

    // Process all input entries in the truncated block
    for (int i = 0; i < truncatedBlock.getPositionCount(); i++) {
        mapUnnester.processCurrentAndAdvance(truncatedUnnestedLengths[i]);
    }

    Block[] output = mapUnnester.buildOutputBlocksAndFlush();
    assertEquals(Arrays.asList(truncatedSlices).stream().mapToInt(slice -> slice.length).sum(), output[0].getPositionCount());

    Slice[] expectedOutput0 = computeExpectedUnnestedOutput(getFieldElements(truncatedSlices, 0), truncatedUnnestedLengths, 0, truncatedSlices.length);
    assertBlock(output[0], expectedOutput0);

    Slice[] expectedOutput1 = computeExpectedUnnestedOutput(getFieldElements(truncatedSlices, 1), truncatedUnnestedLengths, 0, truncatedSlices.length);
    assertBlock(output[1], expectedOutput1);
}
 

/**
 * Prepares for a new output block.
 * The caller has to ensure that the source is not null when this method is invoked.
 */
public void startNewOutput(PageBuilderStatus pageBuilderStatus, int expectedEntries)
{
    checkState(source != null, "source is null");
    this.pageBuilderStatus = pageBuilderStatus;

    // Prepare for a new dictionary block
    this.ids = new int[expectedEntries];
    this.positionCount = 0;
    this.usingCopiedBlock = false;
}
 

@Override
public Page getOutput()
{
    if (currentPage == null) {
        return null;
    }

    PageBuilderStatus pageBuilderStatus = new PageBuilderStatus(MAX_BYTES_PER_PAGE);
    prepareForNewOutput(pageBuilderStatus);

    int outputRowCount = 0;

    while (currentPosition < currentPage.getPositionCount()) {
        outputRowCount += processCurrentPosition();
        currentPosition++;

        if (outputRowCount >= MAX_ROWS_PER_BLOCK || pageBuilderStatus.isFull()) {
            break;
        }
    }

    Block[] outputBlocks = buildOutputBlocks();

    if (currentPosition == currentPage.getPositionCount()) {
        currentPage = null;
        currentPosition = 0;
    }

    return new Page(outputBlocks);
}
 

private void prepareForNewOutput(PageBuilderStatus pageBuilderStatus)
{
    unnesters.forEach(unnester -> unnester.startNewOutput(pageBuilderStatus, estimatedMaxRowsPerBlock));
    replicatedBlockBuilders.forEach(replicatedBlockBuilder -> replicatedBlockBuilder.startNewOutput(estimatedMaxRowsPerBlock));

    if (withOrdinality) {
        ordinalityBlockBuilder = BIGINT.createBlockBuilder(pageBuilderStatus.createBlockBuilderStatus(), estimatedMaxRowsPerBlock);
    }
}
 

@TypeParameter("T")
@SqlType("array(array(T))")
public static Block combinations(
        @TypeParameter("T") Type elementType,
        @SqlType("array(T)") Block array,
        @SqlType(INTEGER) long n)
{
    int arrayLength = array.getPositionCount();
    int combinationLength = toIntExact(n);
    checkCondition(combinationLength >= 0, INVALID_FUNCTION_ARGUMENT, "combination size must not be negative: %s", combinationLength);
    checkCondition(combinationLength <= MAX_COMBINATION_LENGTH, INVALID_FUNCTION_ARGUMENT, "combination size must not exceed %s: %s", MAX_COMBINATION_LENGTH, combinationLength);

    ArrayType arrayType = new ArrayType(elementType);
    if (combinationLength > arrayLength) {
        return arrayType.createBlockBuilder(new PageBuilderStatus().createBlockBuilderStatus(), 0).build();
    }

    int combinationCount = combinationCount(arrayLength, combinationLength);
    checkCondition(combinationCount * (long) combinationLength <= MAX_RESULT_ELEMENTS, INVALID_FUNCTION_ARGUMENT, "combinations exceed max size");

    int[] ids = new int[combinationCount * combinationLength];
    int idsPosition = 0;

    int[] combination = firstCombination(arrayLength, combinationLength);
    do {
        arraycopy(combination, 0, ids, idsPosition, combinationLength);
        idsPosition += combinationLength;
    }
    while (nextCombination(combination, combinationLength));
    verify(idsPosition == ids.length, "idsPosition != ids.length, %s and %s respectively", idsPosition, ids.length);

    int[] offsets = new int[combinationCount + 1];
    setAll(offsets, i -> i * combinationLength);

    return ArrayBlock.fromElementBlock(combinationCount, Optional.empty(), offsets, new DictionaryBlock(array, ids));
}
 

@Override
public BlockBuilder createBlockBuilder(BlockBuilderStatus blockBuilderStatus, int expectedEntries, int expectedBytesPerEntry)
{
    int maxBlockSizeInBytes;
    if (blockBuilderStatus == null) {
        maxBlockSizeInBytes = PageBuilderStatus.DEFAULT_MAX_PAGE_SIZE_IN_BYTES;
    }
    else {
        maxBlockSizeInBytes = blockBuilderStatus.getMaxPageSizeInBytes();
    }
    return new Int128ArrayBlockBuilder(
            blockBuilderStatus,
            Math.min(expectedEntries, maxBlockSizeInBytes / getFixedSize()));
}
 

@Override
public BlockBuilder createBlockBuilder(BlockBuilderStatus blockBuilderStatus, int expectedEntries, int expectedBytesPerEntry)
{
    int maxBlockSizeInBytes;
    if (blockBuilderStatus == null) {
        maxBlockSizeInBytes = PageBuilderStatus.DEFAULT_MAX_PAGE_SIZE_IN_BYTES;
    }
    else {
        maxBlockSizeInBytes = blockBuilderStatus.getMaxPageSizeInBytes();
    }
    return new Int128ArrayBlockBuilder(
            blockBuilderStatus,
            Math.min(expectedEntries, maxBlockSizeInBytes / getFixedSize()));
}