Search in sources :

Example 56 with ArrayType

use of io.trino.spi.type.ArrayType in project trino by trinodb.

the class AbstractMinMaxNAggregationFunction method specialize.

@Override
public AggregationMetadata specialize(BoundSignature boundSignature, FunctionDependencies functionDependencies) {
    Type type = boundSignature.getArgumentTypes().get(0);
    MethodHandle compare = getMinMaxCompare(functionDependencies, type, simpleConvention(FAIL_ON_NULL, BLOCK_POSITION, BLOCK_POSITION), min);
    MinMaxNStateSerializer stateSerializer = new MinMaxNStateSerializer(compare, type);
    ArrayType outputType = new ArrayType(type);
    MethodHandle inputFunction = INPUT_FUNCTION.bindTo(compare).bindTo(type);
    inputFunction = normalizeInputMethod(inputFunction, boundSignature, STATE, BLOCK_INPUT_CHANNEL, INPUT_CHANNEL, BLOCK_INDEX);
    return new AggregationMetadata(inputFunction, Optional.empty(), Optional.of(COMBINE_FUNCTION), OUTPUT_FUNCTION.bindTo(outputType), ImmutableList.of(new AccumulatorStateDescriptor<>(MinMaxNState.class, stateSerializer, new MinMaxNStateFactory())));
}
Also used : ArrayType(io.trino.spi.type.ArrayType) MinMaxNStateSerializer(io.trino.operator.aggregation.state.MinMaxNStateSerializer) Type(io.trino.spi.type.Type) ArrayType(io.trino.spi.type.ArrayType) AccumulatorStateDescriptor(io.trino.operator.aggregation.AggregationMetadata.AccumulatorStateDescriptor) MinMaxNStateFactory(io.trino.operator.aggregation.state.MinMaxNStateFactory) MethodHandle(java.lang.invoke.MethodHandle)

Example 57 with ArrayType

use of io.trino.spi.type.ArrayType in project trino by trinodb.

the class AbstractMinMaxNAggregationFunction method output.

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();
}
Also used : Type(io.trino.spi.type.Type) ArrayType(io.trino.spi.type.ArrayType) BlockBuilder(io.trino.spi.block.BlockBuilder)

Example 58 with ArrayType

use of io.trino.spi.type.ArrayType in project trino by trinodb.

the class MapEntriesFunction method mapFromEntries.

@TypeParameter("K")
@TypeParameter("V")
@SqlType("array(row(K,V))")
public Block mapFromEntries(@TypeParameter("row(K,V)") RowType rowType, @SqlType("map(K,V)") Block block) {
    verify(rowType.getTypeParameters().size() == 2);
    verify(block.getPositionCount() % 2 == 0);
    Type keyType = rowType.getTypeParameters().get(0);
    Type valueType = rowType.getTypeParameters().get(1);
    ArrayType arrayType = new ArrayType(rowType);
    if (pageBuilder.isFull()) {
        pageBuilder.reset();
    }
    int entryCount = block.getPositionCount() / 2;
    BlockBuilder blockBuilder = pageBuilder.getBlockBuilder(0);
    BlockBuilder entryBuilder = blockBuilder.beginBlockEntry();
    for (int i = 0; i < entryCount; i++) {
        BlockBuilder rowBuilder = entryBuilder.beginBlockEntry();
        keyType.appendTo(block, 2 * i, rowBuilder);
        valueType.appendTo(block, 2 * i + 1, rowBuilder);
        entryBuilder.closeEntry();
    }
    blockBuilder.closeEntry();
    pageBuilder.declarePosition();
    return arrayType.getObject(blockBuilder, blockBuilder.getPositionCount() - 1);
}
Also used : ArrayType(io.trino.spi.type.ArrayType) RowType(io.trino.spi.type.RowType) SqlType(io.trino.spi.function.SqlType) Type(io.trino.spi.type.Type) ArrayType(io.trino.spi.type.ArrayType) BlockBuilder(io.trino.spi.block.BlockBuilder) TypeParameter(io.trino.spi.function.TypeParameter) SqlType(io.trino.spi.function.SqlType)

Example 59 with ArrayType

use of io.trino.spi.type.ArrayType in project trino by trinodb.

the class MultimapFromEntriesFunction method multimapFromEntries.

@TypeParameter("K")
@TypeParameter("V")
@SqlType("map(K,array(V))")
@SqlNullable
public Block multimapFromEntries(@TypeParameter("map(K,array(V))") MapType mapType, @OperatorDependency(operator = EQUAL, argumentTypes = { "K", "K" }, convention = @Convention(arguments = { BLOCK_POSITION, BLOCK_POSITION }, result = NULLABLE_RETURN)) BlockPositionEqual keyEqual, @OperatorDependency(operator = HASH_CODE, argumentTypes = "K", convention = @Convention(arguments = BLOCK_POSITION, result = FAIL_ON_NULL)) BlockPositionHashCode keyHashCode, @SqlType("array(row(K,V))") Block mapEntries) {
    Type keyType = mapType.getKeyType();
    Type valueType = ((ArrayType) mapType.getValueType()).getElementType();
    RowType mapEntryType = RowType.anonymous(ImmutableList.of(keyType, valueType));
    if (pageBuilder.isFull()) {
        pageBuilder.reset();
    }
    int entryCount = mapEntries.getPositionCount();
    if (entryCount > entryIndicesList.length) {
        initializeEntryIndicesList(entryCount);
    }
    TypedSet keySet = createEqualityTypedSet(keyType, keyEqual, keyHashCode, entryCount, NAME);
    for (int i = 0; i < entryCount; i++) {
        if (mapEntries.isNull(i)) {
            clearEntryIndices(keySet.size());
            throw new TrinoException(INVALID_FUNCTION_ARGUMENT, "map entry cannot be null");
        }
        Block mapEntryBlock = mapEntryType.getObject(mapEntries, i);
        if (mapEntryBlock.isNull(0)) {
            clearEntryIndices(keySet.size());
            throw new TrinoException(INVALID_FUNCTION_ARGUMENT, "map key cannot be null");
        }
        if (keySet.add(mapEntryBlock, 0)) {
            entryIndicesList[keySet.size() - 1].add(i);
        } else {
            entryIndicesList[keySet.positionOf(mapEntryBlock, 0)].add(i);
        }
    }
    BlockBuilder multimapBlockBuilder = pageBuilder.getBlockBuilder(0);
    BlockBuilder mapWriter = multimapBlockBuilder.beginBlockEntry();
    for (int i = 0; i < keySet.size(); i++) {
        keyType.appendTo(mapEntryType.getObject(mapEntries, entryIndicesList[i].getInt(0)), 0, mapWriter);
        BlockBuilder valuesArray = mapWriter.beginBlockEntry();
        for (int entryIndex : entryIndicesList[i]) {
            valueType.appendTo(mapEntryType.getObject(mapEntries, entryIndex), 1, valuesArray);
        }
        mapWriter.closeEntry();
    }
    multimapBlockBuilder.closeEntry();
    pageBuilder.declarePosition();
    clearEntryIndices(keySet.size());
    return mapType.getObject(multimapBlockBuilder, multimapBlockBuilder.getPositionCount() - 1);
}
Also used : ArrayType(io.trino.spi.type.ArrayType) Type(io.trino.spi.type.Type) SqlType(io.trino.spi.function.SqlType) RowType(io.trino.spi.type.RowType) MapType(io.trino.spi.type.MapType) ArrayType(io.trino.spi.type.ArrayType) RowType(io.trino.spi.type.RowType) TypedSet(io.trino.operator.aggregation.TypedSet) TypedSet.createEqualityTypedSet(io.trino.operator.aggregation.TypedSet.createEqualityTypedSet) TrinoException(io.trino.spi.TrinoException) Block(io.trino.spi.block.Block) BlockBuilder(io.trino.spi.block.BlockBuilder) SqlNullable(io.trino.spi.function.SqlNullable) TypeParameter(io.trino.spi.function.TypeParameter) SqlType(io.trino.spi.function.SqlType)

Example 60 with ArrayType

use of io.trino.spi.type.ArrayType in project trino by trinodb.

the class JsonToArrayCast method specialize.

@Override
protected ScalarFunctionImplementation specialize(BoundSignature boundSignature) {
    checkArgument(boundSignature.getArity() == 1, "Expected arity to be 1");
    ArrayType arrayType = (ArrayType) boundSignature.getReturnType();
    checkCondition(canCastFromJson(arrayType), INVALID_CAST_ARGUMENT, "Cannot cast JSON to %s", arrayType);
    BlockBuilderAppender arrayAppender = BlockBuilderAppender.createBlockBuilderAppender(arrayType);
    MethodHandle methodHandle = METHOD_HANDLE.bindTo(arrayType).bindTo(arrayAppender);
    return new ChoicesScalarFunctionImplementation(boundSignature, NULLABLE_RETURN, ImmutableList.of(NEVER_NULL), methodHandle);
}
Also used : ArrayType(io.trino.spi.type.ArrayType) BlockBuilderAppender(io.trino.util.JsonUtil.BlockBuilderAppender) MethodHandle(java.lang.invoke.MethodHandle)

Aggregations

ArrayType (io.trino.spi.type.ArrayType)289 Test (org.testng.annotations.Test)205 Type (io.trino.spi.type.Type)92 RowType (io.trino.spi.type.RowType)86 ImmutableList (com.google.common.collect.ImmutableList)66 List (java.util.List)62 ArrayList (java.util.ArrayList)59 MapType (io.trino.spi.type.MapType)43 Arrays.asList (java.util.Arrays.asList)36 Collections.singletonList (java.util.Collections.singletonList)34 VarcharType (io.trino.spi.type.VarcharType)32 VarcharType.createUnboundedVarcharType (io.trino.spi.type.VarcharType.createUnboundedVarcharType)32 BlockBuilder (io.trino.spi.block.BlockBuilder)31 MessageType (org.apache.parquet.schema.MessageType)31 DecimalType.createDecimalType (io.trino.spi.type.DecimalType.createDecimalType)30 MessageTypeParser.parseMessageType (org.apache.parquet.schema.MessageTypeParser.parseMessageType)27 DecimalType (io.trino.spi.type.DecimalType)26 StructuralTestUtil.mapType (io.trino.testing.StructuralTestUtil.mapType)24 Block (io.trino.spi.block.Block)23 Map (java.util.Map)23