io.prestosql.spi.block.Block Java Examples

@Setup
public void setup()
{
    Metadata metadata = createTestMetadataManager();
    metadata.addFunctions(extractFunctions(BenchmarkArraySort.class));
    ExpressionCompiler compiler = new ExpressionCompiler(metadata, new PageFunctionCompiler(metadata, 0));
    ImmutableList.Builder<RowExpression> projectionsBuilder = ImmutableList.builder();
    Block[] blocks = new Block[TYPES.size()];
    for (int i = 0; i < TYPES.size(); i++) {
        Type elementType = TYPES.get(i);
        ArrayType arrayType = new ArrayType(elementType);
        projectionsBuilder.add(new CallExpression(
                metadata.resolveFunction(QualifiedName.of(name), fromTypes(arrayType)),
                arrayType,
                ImmutableList.of(field(i, arrayType))));
        blocks[i] = createChannel(POSITIONS, ARRAY_SIZE, arrayType);
    }

    ImmutableList<RowExpression> projections = projectionsBuilder.build();
    pageProcessor = compiler.compilePageProcessor(Optional.empty(), projections).get();
    page = new Page(blocks);
}
 

@Test
public void testMisaligned()
{
    int[] unnestedLengths = {1, 2, 0, 0};
    int[] requiredOutputCounts = {1, 3, 0, 1};

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

    Block[] blocks = testMapUnnester(requiredOutputCounts, unnestedLengths, elements);

    // Check final state. Misalignment has occurred.
    assertEquals(blocks.length, 2);
    // Misaligned, does not have a null entry.
    assertFalse(blocks[0] instanceof DictionaryBlock);
    // Misaligned, but has a null entry elements[1][1][1]
    assertTrue(blocks[1] instanceof DictionaryBlock);
}
 

public static Block createStringArraysBlock(Iterable<? extends Iterable<String>> values)
{
    ArrayType arrayType = new ArrayType(VARCHAR);
    BlockBuilder builder = arrayType.createBlockBuilder(null, 100);

    for (Iterable<String> value : values) {
        if (value == null) {
            builder.appendNull();
        }
        else {
            arrayType.writeObject(builder, createStringsBlock(value));
        }
    }

    return builder.build();
}
 

@Test
public void testPrimitiveByteSerialization()
{
    AccumulatorStateFactory<ByteState> factory = StateCompiler.generateStateFactory(ByteState.class);
    AccumulatorStateSerializer<ByteState> serializer = StateCompiler.generateStateSerializer(ByteState.class);
    ByteState state = factory.createSingleState();
    ByteState deserializedState = factory.createSingleState();

    state.setByte((byte) 3);

    BlockBuilder builder = TINYINT.createBlockBuilder(null, 1);
    serializer.serialize(state, builder);

    Block block = builder.build();
    serializer.deserialize(block, 0, deserializedState);
    assertEquals(deserializedState.getByte(), state.getByte());
}
 

public void compact()
{
    if (eagerCompact) {
        return;
    }
    for (int channel = 0; channel < types.size(); channel++) {
        ObjectArrayList<Block> blocks = channels[channel];
        for (int i = nextBlockToCompact; i < blocks.size(); i++) {
            Block block = blocks.get(i);

            // Copy the block to compact its size
            Block compactedBlock = block.copyRegion(0, block.getPositionCount());
            blocks.set(i, compactedBlock);
            pagesMemorySize -= block.getRetainedSizeInBytes();
            pagesMemorySize += compactedBlock.getRetainedSizeInBytes();
        }
    }
    nextBlockToCompact = channels[0].size();
    estimatedSize = calculateEstimatedSize();
}
 

private static Block serializeList(BlockBuilder parentBlockBuilder, Object value, Type type, String columnName)
{
    if (value == null) {
        checkState(parentBlockBuilder != null, "parentBlockBuilder is null");
        parentBlockBuilder.appendNull();
        return null;
    }
    List<?> list = (List<?>) value;
    List<Type> typeParameters = type.getTypeParameters();
    Type elementType = typeParameters.get(0);

    BlockBuilder blockBuilder = elementType.createBlockBuilder(null, list.size());
    for (Object element : list) {
        serializeObject(blockBuilder, element, elementType, columnName);
    }
    if (parentBlockBuilder != null) {
        type.writeObject(parentBlockBuilder, blockBuilder.build());
        return null;
    }
    return blockBuilder.build();
}
 

private static <T> void verifyBlock(Block block, T[] expectedValues)
{
    assertBlock(block, expectedValues);

    assertColumnarArray(block, expectedValues);
    assertDictionaryBlock(block, expectedValues);
    assertRunLengthEncodedBlock(block, expectedValues);

    int offset = 1;
    int length = expectedValues.length - 2;
    Block blockRegion = block.getRegion(offset, length);
    T[] expectedValuesRegion = Arrays.copyOfRange(expectedValues, offset, offset + length);

    assertBlock(blockRegion, expectedValuesRegion);

    assertColumnarArray(blockRegion, expectedValuesRegion);
    assertDictionaryBlock(blockRegion, expectedValuesRegion);
    assertRunLengthEncodedBlock(blockRegion, expectedValuesRegion);
}
 

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 Block createTestBlock()
{
    BlockBuilder blockBuilder = TINYINT.createBlockBuilder(null, 15);
    TINYINT.writeLong(blockBuilder, 111);
    TINYINT.writeLong(blockBuilder, 111);
    TINYINT.writeLong(blockBuilder, 111);
    TINYINT.writeLong(blockBuilder, 22);
    TINYINT.writeLong(blockBuilder, 22);
    TINYINT.writeLong(blockBuilder, 22);
    TINYINT.writeLong(blockBuilder, 22);
    TINYINT.writeLong(blockBuilder, 22);
    TINYINT.writeLong(blockBuilder, 33);
    TINYINT.writeLong(blockBuilder, 33);
    TINYINT.writeLong(blockBuilder, 44);
    return blockBuilder.build();
}
 

@Override
public ScalarFunctionImplementation specialize(BoundVariables boundVariables, int arity, Metadata metadata)
{
    if (arity < 2) {
        throw new PrestoException(INVALID_FUNCTION_ARGUMENT, "There must be two or more arguments to " + FUNCTION_NAME);
    }

    Type elementType = boundVariables.getTypeVariable("E");

    VarArgsToArrayAdapterGenerator.MethodHandleAndConstructor methodHandleAndConstructor = generateVarArgsToArrayAdapter(
            Block.class,
            Block.class,
            arity,
            METHOD_HANDLE.bindTo(elementType),
            USER_STATE_FACTORY.bindTo(elementType));

    return new ScalarFunctionImplementation(
            false,
            nCopies(arity, valueTypeArgumentProperty(RETURN_NULL_ON_NULL)),
            methodHandleAndConstructor.getMethodHandle(),
            Optional.of(methodHandleAndConstructor.getConstructor()));
}
 

public static void input(BlockComparator comparator, Type valueType, Type keyType, MinMaxByNState state, Block value, Block key, int blockIndex, long n)
{
    TypedKeyValueHeap heap = state.getTypedKeyValueHeap();
    if (heap == null) {
        if (n <= 0) {
            throw new PrestoException(INVALID_FUNCTION_ARGUMENT, "third argument of max_by/min_by must be a positive integer");
        }
        checkCondition(n <= MAX_NUMBER_OF_VALUES, INVALID_FUNCTION_ARGUMENT, "third argument of max_by/min_by must be less than or equal to %s; found %s", MAX_NUMBER_OF_VALUES, n);
        heap = new TypedKeyValueHeap(comparator, keyType, valueType, toIntExact(n));
        state.setTypedKeyValueHeap(heap);
    }

    long startSize = heap.getEstimatedSize();
    if (!key.isNull(blockIndex)) {
        heap.add(key, value, blockIndex);
    }
    state.addMemoryUsage(heap.getEstimatedSize() - startSize);
}
 

public int putIfAbsent(Block block, int position)
{
    requireNonNull(block, "block must not be null");

    if (block.isNull(position)) {
        containsNullElement = true;
        return NULL_POSITION;
    }

    int blockPosition;
    long hashPosition = getHashPositionOfElement(block, position);
    if (blockPositionByHash.get(hashPosition) != EMPTY_SLOT) {
        blockPosition = blockPositionByHash.get(hashPosition);
    }
    else {
        blockPosition = addNewElement(hashPosition, block, position);
    }
    verify(blockPosition != NULL_POSITION);
    return blockPosition;
}
 

@UsedByGeneratedCode
public static Slice sliceArrayMinMax(MethodHandle compareMethodHandle, Type elementType, Block block)
{
    try {
        if (block.getPositionCount() == 0) {
            return null;
        }

        Slice selectedValue = elementType.getSlice(block, 0);
        for (int i = 0; i < block.getPositionCount(); i++) {
            if (block.isNull(i)) {
                return null;
            }
            Slice value = elementType.getSlice(block, i);
            if ((boolean) compareMethodHandle.invokeExact(value, selectedValue)) {
                selectedValue = value;
            }
        }

        return selectedValue;
    }
    catch (Throwable t) {
        throw internalError(t);
    }
}
 

@Override
public void encodeColumn(Block block, SliceOutput output, EncodeOutput encodeOutput)
{
    for (int position = 0; position < block.getPositionCount(); position++) {
        if (block.isNull(position)) {
            output.writeBytes(nullSequence);
        }
        else {
            long value = type.getLong(block, position);
            buffer.setLength(0);
            buffer.append(value);
            for (int index = 0; index < buffer.length(); index++) {
                output.writeByte(buffer.charAt(index));
            }
        }
        encodeOutput.closeEntry();
    }
}
 

public Page build()
{
    if (blockBuilders.length == 0) {
        return new Page(declaredPositions);
    }

    Block[] blocks = new Block[blockBuilders.length];
    for (int i = 0; i < blocks.length; i++) {
        blocks[i] = blockBuilders[i].build();
        if (blocks[i].getPositionCount() != declaredPositions) {
            throw new IllegalStateException(format("Declared positions (%s) does not match block %s's number of entries (%s)", declaredPositions, i, blocks[i].getPositionCount()));
        }
    }

    return Page.wrapBlocksWithoutCopy(declaredPositions, blocks);
}
 

private void evaluateAndFlushGroup(Page page, int position)
{
    pageBuilder.declarePosition();
    for (int i = 0; i < groupByTypes.size(); i++) {
        Block block = page.getBlock(groupByChannels[i]);
        Type type = groupByTypes.get(i);
        type.appendTo(block, position, pageBuilder.getBlockBuilder(i));
    }
    int offset = groupByTypes.size();
    for (int i = 0; i < aggregates.size(); i++) {
        aggregates.get(i).evaluate(pageBuilder.getBlockBuilder(offset + i));
    }

    if (pageBuilder.isFull()) {
        outputPages.add(pageBuilder.build());
        pageBuilder.reset();
    }

    aggregates = setupAggregates(step, accumulatorFactories);
}
 

public static Function<Block, Block> createDoubleToDecimalCoercer(DecimalType toType)
{
    if (toType.isShort()) {
        return new DoubleToShortDecimalCoercer(toType);
    }
    else {
        return new DoubleToLongDecimalCoercer(toType);
    }
}
 

@Override
public void write(Block block)
{
    for (int i = 0; i < block.getPositionCount(); i++) {
        if (!block.isNull(i)) {
            Slice slice = type.getSlice(block, i);
            Binary binary = Binary.fromConstantByteBuffer(slice.toByteBuffer());
            valuesWriter.writeBytes(binary);
            getStatistics().updateStats(binary);
        }
    }
}
 

@Override
public void encodeColumn(Block block, SliceOutput output, EncodeOutput encodeOutput)
{
    for (int position = 0; position < block.getPositionCount(); position++) {
        if (block.isNull(position)) {
            output.writeBytes(nullSequence);
        }
        else {
            encodeValue(block, position, output);
        }
        encodeOutput.closeEntry();
    }
}
 

@Override
public boolean equalTo(Block leftBlock, int leftPosition, Block rightBlock, int rightPosition)
{
    float leftValue = intBitsToFloat(leftBlock.getInt(leftPosition, 0));
    float rightValue = intBitsToFloat(rightBlock.getInt(rightPosition, 0));

    // direct equality is correct here
    // noinspection FloatingPointEquality
    return leftValue == rightValue;
}
 

@TypeParameter("T")
@TypeParameterSpecialization(name = "T", nativeContainerType = long.class)
@SqlType(StandardTypes.BOOLEAN)
@SqlNullable
public static Boolean anyMatchLong(
        @TypeParameter("T") Type elementType,
        @SqlType("array(T)") Block arrayBlock,
        @SqlType("function(T, boolean)") LongToBooleanFunction function)
{
    boolean hasNullResult = false;
    int positionCount = arrayBlock.getPositionCount();
    for (int i = 0; i < positionCount; i++) {
        Long element = null;
        if (!arrayBlock.isNull(i)) {
            element = elementType.getLong(arrayBlock, i);
        }
        Boolean match = function.apply(element);
        if (TRUE.equals(match)) {
            return true;
        }
        if (match == null) {
            hasNullResult = true;
        }
    }
    if (hasNullResult) {
        return null;
    }
    return false;
}
 

@Override
public BytecodeNode generateExpression(ResolvedFunction resolvedFunction, BytecodeGeneratorContext context, Type rowType, List<RowExpression> arguments)
{
    BytecodeBlock block = new BytecodeBlock().setDescription("Constructor for " + rowType.toString());
    CallSiteBinder binder = context.getCallSiteBinder();
    Scope scope = context.getScope();
    List<Type> types = rowType.getTypeParameters();

    block.comment("Create new RowBlockBuilder; beginBlockEntry;");
    Variable blockBuilder = scope.createTempVariable(BlockBuilder.class);
    Variable singleRowBlockWriter = scope.createTempVariable(BlockBuilder.class);
    block.append(blockBuilder.set(
            constantType(binder, rowType).invoke(
                    "createBlockBuilder",
                    BlockBuilder.class,
                    constantNull(BlockBuilderStatus.class),
                    constantInt(1))));
    block.append(singleRowBlockWriter.set(blockBuilder.invoke("beginBlockEntry", BlockBuilder.class)));

    for (int i = 0; i < arguments.size(); ++i) {
        Type fieldType = types.get(i);
        Variable field = scope.createTempVariable(fieldType.getJavaType());
        block.comment("Clean wasNull and Generate + " + i + "-th field of row");
        block.append(context.wasNull().set(constantFalse()));
        block.append(context.generate(arguments.get(i)));
        block.putVariable(field);
        block.append(new IfStatement()
                .condition(context.wasNull())
                .ifTrue(singleRowBlockWriter.invoke("appendNull", BlockBuilder.class).pop())
                .ifFalse(constantType(binder, fieldType).writeValue(singleRowBlockWriter, field).pop()));
    }
    block.comment("closeEntry; slice the SingleRowBlock; wasNull = false;");
    block.append(blockBuilder.invoke("closeEntry", BlockBuilder.class).pop());
    block.append(constantType(binder, rowType).invoke("getObject", Object.class, blockBuilder.cast(Block.class), constantInt(0))
            .cast(Block.class));
    block.append(context.wasNull().set(constantFalse()));
    return block;
}
 

@Override
public void encodeColumn(Block block, SliceOutput output, EncodeOutput encodeOutput)
{
    for (int position = 0; position < block.getPositionCount(); position++) {
        if (!block.isNull(position)) {
            encodeValueInto(block, position, output);
        }
        encodeOutput.closeEntry();
    }
}
 

@Benchmark
@OperationsPerInvocation(POSITIONS)
public BlockBuilder benchmarkMapCopy(BenchmarkData data)
{
    Block block = data.getDataBlock();
    BlockBuilder blockBuilder = data.getBlockBuilder();
    MapType mapType = mapType(VARCHAR, BIGINT);

    for (int i = 0; i < POSITIONS; i++) {
        mapType.appendTo(block, i, blockBuilder);
    }

    return blockBuilder;
}
 

@Test
public void testJsonSerialization()
        throws Exception
{
    TestingTypeManager typeManager = new TestingTypeManager();
    TestingBlockEncodingSerde blockEncodingSerde = new TestingBlockEncodingSerde();

    ObjectMapper mapper = new ObjectMapperProvider().get()
            .registerModule(new SimpleModule()
                    .addDeserializer(Type.class, new TestingTypeDeserializer(typeManager))
                    .addSerializer(Block.class, new TestingBlockJsonSerde.Serializer(blockEncodingSerde))
                    .addDeserializer(Block.class, new TestingBlockJsonSerde.Deserializer(blockEncodingSerde)));

    EquatableValueSet set = EquatableValueSet.all(ID);
    assertEquals(set, mapper.readValue(mapper.writeValueAsString(set), EquatableValueSet.class));

    set = EquatableValueSet.none(ID);
    assertEquals(set, mapper.readValue(mapper.writeValueAsString(set), EquatableValueSet.class));

    set = EquatableValueSet.of(ID, 1L);
    assertEquals(set, mapper.readValue(mapper.writeValueAsString(set), EquatableValueSet.class));

    set = EquatableValueSet.of(ID, 1L, 2L);
    assertEquals(set, mapper.readValue(mapper.writeValueAsString(set), EquatableValueSet.class));

    set = EquatableValueSet.of(ID, 1L, 2L).complement();
    assertEquals(set, mapper.readValue(mapper.writeValueAsString(set), EquatableValueSet.class));
}
 

@Setup
public void setup()
{
    Metadata metadata = createTestMetadataManager();
    metadata.addFunctions(new FunctionListBuilder().function(EXACT_ARRAY_FILTER_FUNCTION).getFunctions());
    ExpressionCompiler compiler = new ExpressionCompiler(metadata, new PageFunctionCompiler(metadata, 0));
    ImmutableList.Builder<RowExpression> projectionsBuilder = ImmutableList.builder();
    Block[] blocks = new Block[TYPES.size()];
    for (int i = 0; i < TYPES.size(); i++) {
        Type elementType = TYPES.get(i);
        ArrayType arrayType = new ArrayType(elementType);
        ResolvedFunction resolvedFunction = metadata.resolveFunction(
                QualifiedName.of(name),
                fromTypes(arrayType, new FunctionType(ImmutableList.of(BIGINT), BOOLEAN)));
        ResolvedFunction greaterThan = metadata.resolveOperator(GREATER_THAN, ImmutableList.of(BIGINT, BIGINT));
        projectionsBuilder.add(new CallExpression(resolvedFunction, arrayType, ImmutableList.of(
                field(0, arrayType),
                new LambdaDefinitionExpression(
                        ImmutableList.of(BIGINT),
                        ImmutableList.of("x"),
                        new CallExpression(greaterThan, BOOLEAN, ImmutableList.of(new VariableReferenceExpression("x", BIGINT), constant(0L, BIGINT)))))));
        blocks[i] = createChannel(POSITIONS, ARRAY_SIZE, arrayType);
    }

    ImmutableList<RowExpression> projections = projectionsBuilder.build();
    pageProcessor = compiler.compilePageProcessor(Optional.empty(), projections).get();
    page = new Page(blocks);
}
 

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

@Test
public void testBoolean()
{
    InternalAggregationFunction booleanAgg = metadata.getAggregateFunctionImplementation(metadata.resolveFunction(QualifiedName.of("checksum"), fromTypes(BOOLEAN)));
    Block block = createBooleansBlock(null, null, true, false, false);
    assertAggregation(booleanAgg, expectedChecksum(BOOLEAN, block), block);
}
 

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

    return block.getSlice(position, 0, block.getSliceLength(position)).toStringUtf8();
}
 

@UsedByGeneratedCode
public static Slice arrayJoin(
        MethodHandle castFunction,
        Object state,
        ConnectorSession session,
        Block arrayBlock,
        Slice delimiter)
{
    return arrayJoin(castFunction, state, session, arrayBlock, delimiter, null);
}