Examples with Block io.trino.spi.block.Block used on opensource projects

Example 11 with Block

use of io.trino.spi.block.Block in project trino by trinodb.

the class MapAggregationFunction method combine.

public static void combine(KeyValuePairsState state, KeyValuePairsState otherState) {
    if (state.get() != null && otherState.get() != null) {
        Block keys = otherState.get().getKeys();
        Block values = otherState.get().getValues();
        KeyValuePairs pairs = state.get();
        long startSize = pairs.estimatedInMemorySize();
        for (int i = 0; i < keys.getPositionCount(); i++) {
            pairs.add(keys, values, i, i);
        }
        state.addMemoryUsage(pairs.estimatedInMemorySize() - startSize);
    } else if (state.get() == null) {
        state.set(otherState.get());
    }
}

Example 12 with Block

use of io.trino.spi.block.Block in project trino by trinodb.

the class ArrayAggregationStateSerializer method deserialize.

@Override
public void deserialize(Block block, int index, ArrayAggregationState state) {
    state.reset();
    Block stateBlock = (Block) arrayType.getObject(block, index);
    for (int i = 0; i < stateBlock.getPositionCount(); i++) {
        state.add(stateBlock, i);
    }
}

Example 13 with Block

use of io.trino.spi.block.Block in project trino by trinodb.

the class SplitToMapFunction method splitToMap.

@SqlType("map(varchar,varchar)")
public Block splitToMap(@TypeParameter("map(varchar,varchar)") Type mapType, @SqlType(StandardTypes.VARCHAR) Slice string, @SqlType(StandardTypes.VARCHAR) Slice entryDelimiter, @SqlType(StandardTypes.VARCHAR) Slice keyValueDelimiter) {
    checkCondition(entryDelimiter.length() > 0, INVALID_FUNCTION_ARGUMENT, "entryDelimiter is empty");
    checkCondition(keyValueDelimiter.length() > 0, INVALID_FUNCTION_ARGUMENT, "keyValueDelimiter is empty");
    checkCondition(!entryDelimiter.equals(keyValueDelimiter), INVALID_FUNCTION_ARGUMENT, "entryDelimiter and keyValueDelimiter must not be the same");
    Map<Slice, Slice> map = new HashMap<>();
    int entryStart = 0;
    while (entryStart < string.length()) {
        // Extract key-value pair based on current index
        // then add the pair if it can be split by keyValueDelimiter
        Slice keyValuePair;
        int entryEnd = string.indexOf(entryDelimiter, entryStart);
        if (entryEnd >= 0) {
            keyValuePair = string.slice(entryStart, entryEnd - entryStart);
        } else {
            // The rest of the string is the last possible pair.
            keyValuePair = string.slice(entryStart, string.length() - entryStart);
        }
        int keyEnd = keyValuePair.indexOf(keyValueDelimiter);
        if (keyEnd >= 0) {
            int valueStart = keyEnd + keyValueDelimiter.length();
            Slice key = keyValuePair.slice(0, keyEnd);
            Slice value = keyValuePair.slice(valueStart, keyValuePair.length() - valueStart);
            if (value.indexOf(keyValueDelimiter) >= 0) {
                throw new TrinoException(INVALID_FUNCTION_ARGUMENT, "Key-value delimiter must appear exactly once in each entry. Bad input: '" + keyValuePair.toStringUtf8() + "'");
            }
            if (map.containsKey(key)) {
                throw new TrinoException(INVALID_FUNCTION_ARGUMENT, format("Duplicate keys (%s) are not allowed", key.toStringUtf8()));
            }
            map.put(key, value);
        } else {
            throw new TrinoException(INVALID_FUNCTION_ARGUMENT, "Key-value delimiter must appear exactly once in each entry. Bad input: '" + keyValuePair.toStringUtf8() + "'");
        }
        if (entryEnd < 0) {
            // No more pairs to add
            break;
        }
        // Next possible pair is placed next to the current entryDelimiter
        entryStart = entryEnd + entryDelimiter.length();
    }
    if (pageBuilder.isFull()) {
        pageBuilder.reset();
    }
    BlockBuilder blockBuilder = pageBuilder.getBlockBuilder(0);
    BlockBuilder singleMapBlockBuilder = blockBuilder.beginBlockEntry();
    for (Map.Entry<Slice, Slice> entry : map.entrySet()) {
        VARCHAR.writeSlice(singleMapBlockBuilder, entry.getKey());
        VARCHAR.writeSlice(singleMapBlockBuilder, entry.getValue());
    }
    blockBuilder.closeEntry();
    pageBuilder.declarePosition();
    return (Block) mapType.getObject(blockBuilder, blockBuilder.getPositionCount() - 1);
}

Example 14 with Block

use of io.trino.spi.block.Block in project trino by trinodb.

the class RowToRowCast method generateRowCast.

private static Class<?> generateRowCast(Type fromType, Type toType, FunctionDependencies functionDependencies) {
    List<Type> toTypes = toType.getTypeParameters();
    List<Type> fromTypes = fromType.getTypeParameters();
    CallSiteBinder binder = new CallSiteBinder();
    // Embed the MD5 hash code of input and output types into the generated class name instead of the raw type names,
    // which could prevent the class name from hitting the length limitation and invalid characters.
    byte[] md5Suffix = Hashing.md5().hashBytes((fromType + "$" + toType).getBytes(UTF_8)).asBytes();
    ClassDefinition definition = new ClassDefinition(a(PUBLIC, FINAL), makeClassName(Joiner.on("$").join("RowCast", BaseEncoding.base16().encode(md5Suffix))), type(Object.class));
    Parameter session = arg("session", ConnectorSession.class);
    Parameter row = arg("row", Block.class);
    MethodDefinition method = definition.declareMethod(a(PUBLIC, STATIC), "castRow", type(Block.class), session, row);
    Scope scope = method.getScope();
    BytecodeBlock body = method.getBody();
    Variable wasNull = scope.declareVariable(boolean.class, "wasNull");
    Variable blockBuilder = scope.createTempVariable(BlockBuilder.class);
    Variable singleRowBlockWriter = scope.createTempVariable(BlockBuilder.class);
    body.append(wasNull.set(constantBoolean(false)));
    CachedInstanceBinder cachedInstanceBinder = new CachedInstanceBinder(definition, binder);
    // create the row block builder
    body.append(blockBuilder.set(constantType(binder, toType).invoke("createBlockBuilder", BlockBuilder.class, constantNull(BlockBuilderStatus.class), constantInt(1))));
    body.append(singleRowBlockWriter.set(blockBuilder.invoke("beginBlockEntry", BlockBuilder.class)));
    // loop through to append member blocks
    for (int i = 0; i < toTypes.size(); i++) {
        Type fromElementType = fromTypes.get(i);
        Type toElementType = toTypes.get(i);
        Type currentFromType = fromElementType;
        if (currentFromType.equals(UNKNOWN)) {
            body.append(singleRowBlockWriter.invoke("appendNull", BlockBuilder.class).pop());
            continue;
        }
        MethodHandle castMethod = getNullSafeCast(functionDependencies, fromElementType, toElementType);
        MethodHandle writeMethod = getNullSafeWrite(toElementType);
        MethodHandle castAndWrite = collectArguments(writeMethod, 1, castMethod);
        body.append(invokeDynamic(BOOTSTRAP_METHOD, ImmutableList.of(binder.bind(castAndWrite).getBindingId()), "castAndWriteField", castAndWrite.type(), singleRowBlockWriter, scope.getVariable("session"), row, constantInt(i)));
    }
    // call blockBuilder.closeEntry() and return the single row block
    body.append(blockBuilder.invoke("closeEntry", BlockBuilder.class).pop());
    body.append(constantType(binder, toType).invoke("getObject", Object.class, blockBuilder.cast(Block.class), constantInt(0)).cast(Block.class).ret());
    // create constructor
    MethodDefinition constructorDefinition = definition.declareConstructor(a(PUBLIC));
    BytecodeBlock constructorBody = constructorDefinition.getBody();
    Variable thisVariable = constructorDefinition.getThis();
    constructorBody.comment("super();").append(thisVariable).invokeConstructor(Object.class);
    cachedInstanceBinder.generateInitializations(thisVariable, constructorBody);
    constructorBody.ret();
    return defineClass(definition, Object.class, binder.getBindings(), RowToRowCast.class.getClassLoader());
}

Example 15 with Block

use of io.trino.spi.block.Block in project trino by trinodb.

the class ZipFunction method specialize.

@Override
protected ScalarFunctionImplementation specialize(BoundSignature boundSignature) {
    List<Type> types = boundSignature.getArgumentTypes().stream().map(ArrayType.class::cast).map(ArrayType::getElementType).collect(toImmutableList());
    List<Class<?>> javaArgumentTypes = nCopies(types.size(), Block.class);
    MethodHandle methodHandle = METHOD_HANDLE.bindTo(types).asVarargsCollector(Block[].class).asType(methodType(Block.class, javaArgumentTypes));
    return new ChoicesScalarFunctionImplementation(boundSignature, FAIL_ON_NULL, nCopies(types.size(), NEVER_NULL), methodHandle);
}