Examples with PageBuilder io.trino.spi.PageBuilder used on opensource projects

Example 31 with PageBuilder

use of io.trino.spi.PageBuilder in project trino by trinodb.

the class PinotSegmentPageSource method getNextPage.

/**
 * @return constructed page for pinot data.
 */
@Override
public Page getNextPage() {
    if (isFinished()) {
        close();
        return null;
    }
    if (!isPinotDataFetched) {
        fetchPinotData();
    }
    // To reduce memory usage, remove dataTable from dataTableList once it's processed.
    if (currentDataTable != null) {
        estimatedMemoryUsageInBytes -= currentDataTable.getEstimatedSizeInBytes();
    }
    if (dataTableList.size() == 0) {
        close();
        return null;
    }
    currentDataTable = dataTableList.pop();
    PageBuilder pageBuilder = new PageBuilder(columnTypes);
    // Note that declared positions in the Page should be the same with number of rows in each Block
    pageBuilder.declarePositions(currentDataTable.getDataTable().getNumberOfRows());
    for (int columnHandleIdx = 0; columnHandleIdx < columnHandles.size(); columnHandleIdx++) {
        BlockBuilder blockBuilder = pageBuilder.getBlockBuilder(columnHandleIdx);
        Type columnType = columnTypes.get(columnHandleIdx);
        // Write a block for each column in the original order.
        writeBlock(blockBuilder, columnType, columnHandleIdx);
    }
    return pageBuilder.build();
}

Example 32 with PageBuilder

use of io.trino.spi.PageBuilder in project trino by trinodb.

the class ArrayTransformFunction method generateTransform.

private static Class<?> generateTransform(Type inputType, Type outputType) {
    CallSiteBinder binder = new CallSiteBinder();
    Class<?> inputJavaType = Primitives.wrap(inputType.getJavaType());
    Class<?> outputJavaType = Primitives.wrap(outputType.getJavaType());
    ClassDefinition definition = new ClassDefinition(a(PUBLIC, FINAL), makeClassName("ArrayTransform"), type(Object.class));
    definition.declareDefaultConstructor(a(PRIVATE));
    // define createPageBuilder
    MethodDefinition createPageBuilderMethod = definition.declareMethod(a(PUBLIC, STATIC), "createPageBuilder", type(PageBuilder.class));
    createPageBuilderMethod.getBody().append(newInstance(PageBuilder.class, constantType(binder, new ArrayType(outputType)).invoke("getTypeParameters", List.class)).ret());
    // define transform method
    Parameter pageBuilder = arg("pageBuilder", PageBuilder.class);
    Parameter block = arg("block", Block.class);
    Parameter function = arg("function", UnaryFunctionInterface.class);
    MethodDefinition method = definition.declareMethod(a(PUBLIC, STATIC), "transform", type(Block.class), ImmutableList.of(pageBuilder, block, function));
    BytecodeBlock body = method.getBody();
    Scope scope = method.getScope();
    Variable positionCount = scope.declareVariable(int.class, "positionCount");
    Variable position = scope.declareVariable(int.class, "position");
    Variable blockBuilder = scope.declareVariable(BlockBuilder.class, "blockBuilder");
    Variable inputElement = scope.declareVariable(inputJavaType, "inputElement");
    Variable outputElement = scope.declareVariable(outputJavaType, "outputElement");
    // invoke block.getPositionCount()
    body.append(positionCount.set(block.invoke("getPositionCount", int.class)));
    // reset page builder if it is full
    body.append(new IfStatement().condition(pageBuilder.invoke("isFull", boolean.class)).ifTrue(pageBuilder.invoke("reset", void.class)));
    // get block builder
    body.append(blockBuilder.set(pageBuilder.invoke("getBlockBuilder", BlockBuilder.class, constantInt(0))));
    BytecodeNode loadInputElement;
    if (!inputType.equals(UNKNOWN)) {
        loadInputElement = new IfStatement().condition(block.invoke("isNull", boolean.class, position)).ifTrue(inputElement.set(constantNull(inputJavaType))).ifFalse(inputElement.set(constantType(binder, inputType).getValue(block, position).cast(inputJavaType)));
    } else {
        loadInputElement = new BytecodeBlock().append(inputElement.set(constantNull(inputJavaType)));
    }
    BytecodeNode writeOutputElement;
    if (!outputType.equals(UNKNOWN)) {
        writeOutputElement = new IfStatement().condition(equal(outputElement, constantNull(outputJavaType))).ifTrue(blockBuilder.invoke("appendNull", BlockBuilder.class).pop()).ifFalse(constantType(binder, outputType).writeValue(blockBuilder, outputElement.cast(outputType.getJavaType())));
    } else {
        writeOutputElement = new BytecodeBlock().append(blockBuilder.invoke("appendNull", BlockBuilder.class).pop());
    }
    body.append(new ForLoop().initialize(position.set(constantInt(0))).condition(lessThan(position, positionCount)).update(incrementVariable(position, (byte) 1)).body(new BytecodeBlock().append(loadInputElement).append(outputElement.set(function.invoke("apply", Object.class, inputElement.cast(Object.class)).cast(outputJavaType))).append(writeOutputElement)));
    body.append(pageBuilder.invoke("declarePositions", void.class, positionCount));
    body.append(blockBuilder.invoke("getRegion", Block.class, subtract(blockBuilder.invoke("getPositionCount", int.class), positionCount), positionCount).ret());
    return defineClass(definition, Object.class, binder.getBindings(), ArrayTransformFunction.class.getClassLoader());
}

Example 33 with PageBuilder

use of io.trino.spi.PageBuilder in project trino by trinodb.

the class MapConstructor method createMap.

@UsedByGeneratedCode
public static Block createMap(MapType mapType, MethodHandle keyIndeterminate, State state, ConnectorSession session, Block keyBlock, Block valueBlock) {
    checkCondition(keyBlock.getPositionCount() == valueBlock.getPositionCount(), INVALID_FUNCTION_ARGUMENT, "Key and value arrays must be the same length");
    PageBuilder pageBuilder = state.getPageBuilder();
    if (pageBuilder.isFull()) {
        pageBuilder.reset();
    }
    MapBlockBuilder mapBlockBuilder = (MapBlockBuilder) pageBuilder.getBlockBuilder(0);
    mapBlockBuilder.strict();
    BlockBuilder blockBuilder = mapBlockBuilder.beginBlockEntry();
    for (int i = 0; i < keyBlock.getPositionCount(); i++) {
        if (keyBlock.isNull(i)) {
            // close block builder before throwing as we may be in a TRY() call
            // so that subsequent calls do not find it in an inconsistent state
            mapBlockBuilder.closeEntry();
            throw new TrinoException(INVALID_FUNCTION_ARGUMENT, "map key cannot be null");
        }
        Object keyObject = readNativeValue(mapType.getKeyType(), keyBlock, i);
        try {
            if ((boolean) keyIndeterminate.invoke(keyObject)) {
                throw new TrinoException(INVALID_FUNCTION_ARGUMENT, "map key cannot be indeterminate: " + mapType.getKeyType().getObjectValue(session, keyBlock, i));
            }
        } catch (Throwable t) {
            mapBlockBuilder.closeEntry();
            throw internalError(t);
        }
        mapType.getKeyType().appendTo(keyBlock, i, blockBuilder);
        mapType.getValueType().appendTo(valueBlock, i, blockBuilder);
    }
    try {
        mapBlockBuilder.closeEntry();
    } catch (DuplicateMapKeyException e) {
        throw e.withDetailedMessage(mapType.getKeyType(), session);
    } finally {
        pageBuilder.declarePosition();
    }
    return mapType.getObject(mapBlockBuilder, mapBlockBuilder.getPositionCount() - 1);
}

Example 34 with PageBuilder

use of io.trino.spi.PageBuilder in project trino by trinodb.

the class GeoFunctions method geometryNearestPoints.

@SqlNullable
@Description("Return the closest points on the two geometries")
@ScalarFunction("geometry_nearest_points")
@SqlType("row(" + GEOMETRY_TYPE_NAME + "," + GEOMETRY_TYPE_NAME + ")")
public static Block geometryNearestPoints(@SqlType(GEOMETRY_TYPE_NAME) Slice left, @SqlType(GEOMETRY_TYPE_NAME) Slice right) {
    Geometry leftGeometry = JtsGeometrySerde.deserialize(left);
    Geometry rightGeometry = JtsGeometrySerde.deserialize(right);
    if (leftGeometry.isEmpty() || rightGeometry.isEmpty()) {
        return null;
    }
    RowType rowType = RowType.anonymous(ImmutableList.of(GEOMETRY, GEOMETRY));
    PageBuilder pageBuilder = new PageBuilder(ImmutableList.of(rowType));
    GeometryFactory geometryFactory = leftGeometry.getFactory();
    Coordinate[] nearestCoordinates = DistanceOp.nearestPoints(leftGeometry, rightGeometry);
    BlockBuilder blockBuilder = pageBuilder.getBlockBuilder(0);
    BlockBuilder entryBlockBuilder = blockBuilder.beginBlockEntry();
    GEOMETRY.writeSlice(entryBlockBuilder, JtsGeometrySerde.serialize(geometryFactory.createPoint(nearestCoordinates[0])));
    GEOMETRY.writeSlice(entryBlockBuilder, JtsGeometrySerde.serialize(geometryFactory.createPoint(nearestCoordinates[1])));
    blockBuilder.closeEntry();
    pageBuilder.declarePosition();
    return rowType.getObject(blockBuilder, blockBuilder.getPositionCount() - 1);
}

Example 35 with PageBuilder

use of io.trino.spi.PageBuilder in project trino by trinodb.

the class MultiChannelGroupByHash method startNewPage.

private void startNewPage() {
    if (currentPageBuilder != null) {
        completedPagesMemorySize += currentPageBuilder.getRetainedSizeInBytes();
        currentPageBuilder = currentPageBuilder.newPageBuilderLike();
    } else {
        currentPageBuilder = new PageBuilder(types);
    }
    for (int i = 0; i < types.size(); i++) {
        channelBuilders.get(i).add(currentPageBuilder.getBlockBuilder(i));
    }
}