Search in sources :

Example 16 with TypeParameter

use of io.trino.spi.function.TypeParameter in project trino by trinodb.

the class ArrayRemoveFunction method remove.

@TypeParameter("E")
@SqlType("array(E)")
public Block remove(@OperatorDependency(operator = EQUAL, argumentTypes = { "E", "E" }, convention = @Convention(arguments = { NEVER_NULL, NEVER_NULL }, result = NULLABLE_RETURN)) MethodHandle equalsFunction, @TypeParameter("E") Type type, @SqlType("array(E)") Block array, @SqlType("E") Object value) {
    List<Integer> positions = new ArrayList<>();
    for (int i = 0; i < array.getPositionCount(); i++) {
        Object element = readNativeValue(type, array, i);
        try {
            if (element == null) {
                positions.add(i);
                continue;
            }
            Boolean result = (Boolean) equalsFunction.invoke(element, value);
            if (result == null) {
                throw new TrinoException(NOT_SUPPORTED, "array_remove does not support arrays with elements that are null or contain null");
            }
            if (!result) {
                positions.add(i);
            }
        } catch (Throwable t) {
            throw internalError(t);
        }
    }
    if (array.getPositionCount() == positions.size()) {
        return array;
    }
    if (pageBuilder.isFull()) {
        pageBuilder.reset();
    }
    BlockBuilder blockBuilder = pageBuilder.getBlockBuilder(0);
    for (int position : positions) {
        type.appendTo(array, position, blockBuilder);
    }
    pageBuilder.declarePositions(positions.size());
    return blockBuilder.getRegion(blockBuilder.getPositionCount() - positions.size(), positions.size());
}
Also used : ArrayList(java.util.ArrayList) TrinoException(io.trino.spi.TrinoException) BlockBuilder(io.trino.spi.block.BlockBuilder) TypeParameter(io.trino.spi.function.TypeParameter) SqlType(io.trino.spi.function.SqlType)

Example 17 with TypeParameter

use of io.trino.spi.function.TypeParameter in project trino by trinodb.

the class ArrayReverseFunction method reverse.

@TypeParameter("E")
@SqlType("array(E)")
public Block reverse(@TypeParameter("E") Type type, @SqlType("array(E)") Block block) {
    int arrayLength = block.getPositionCount();
    if (arrayLength < 2) {
        return block;
    }
    if (pageBuilder.isFull()) {
        pageBuilder.reset();
    }
    BlockBuilder blockBuilder = pageBuilder.getBlockBuilder(0);
    for (int i = arrayLength - 1; i >= 0; i--) {
        type.appendTo(block, i, blockBuilder);
    }
    pageBuilder.declarePositions(arrayLength);
    return blockBuilder.getRegion(blockBuilder.getPositionCount() - arrayLength, arrayLength);
}
Also used : BlockBuilder(io.trino.spi.block.BlockBuilder) TypeParameter(io.trino.spi.function.TypeParameter) SqlType(io.trino.spi.function.SqlType)

Example 18 with TypeParameter

use of io.trino.spi.function.TypeParameter in project trino by trinodb.

the class ArrayShuffleFunction method shuffle.

@TypeParameter("E")
@SqlType("array(E)")
public Block shuffle(@TypeParameter("E") Type type, @SqlType("array(E)") Block block) {
    int length = block.getPositionCount();
    if (positions.length < length) {
        positions = new int[length];
    }
    for (int i = 0; i < length; i++) {
        positions[i] = i;
    }
    // Randomly swap a pair of positions
    for (int i = length - 1; i > 0; i--) {
        int index = ThreadLocalRandom.current().nextInt(i + 1);
        int swap = positions[i];
        positions[i] = positions[index];
        positions[index] = swap;
    }
    if (pageBuilder.isFull()) {
        pageBuilder.reset();
    }
    BlockBuilder blockBuilder = pageBuilder.getBlockBuilder(0);
    for (int i = 0; i < length; i++) {
        type.appendTo(block, positions[i], blockBuilder);
    }
    pageBuilder.declarePositions(length);
    return blockBuilder.getRegion(blockBuilder.getPositionCount() - length, length);
}
Also used : BlockBuilder(io.trino.spi.block.BlockBuilder) TypeParameter(io.trino.spi.function.TypeParameter) SqlType(io.trino.spi.function.SqlType)

Example 19 with TypeParameter

use of io.trino.spi.function.TypeParameter in project trino by trinodb.

the class ArraySortFunction method sort.

@TypeParameter("E")
@SqlType("array(E)")
public Block sort(@OperatorDependency(operator = COMPARISON_UNORDERED_LAST, argumentTypes = { "E", "E" }, convention = @Convention(arguments = { BLOCK_POSITION, BLOCK_POSITION }, result = FAIL_ON_NULL)) BlockPositionComparison comparisonOperator, @TypeParameter("E") Type type, @SqlType("array(E)") Block block) {
    int arrayLength = block.getPositionCount();
    positions.clear();
    for (int i = 0; i < arrayLength; i++) {
        positions.add(i);
    }
    positions.subList(0, arrayLength).sort((left, right) -> {
        boolean nullLeft = block.isNull(left);
        boolean nullRight = block.isNull(right);
        if (nullLeft && nullRight) {
            return 0;
        }
        if (nullLeft) {
            return 1;
        }
        if (nullRight) {
            return -1;
        }
        return (int) comparisonOperator.compare(block, left, block, right);
    });
    if (pageBuilder.isFull()) {
        pageBuilder.reset();
    }
    BlockBuilder blockBuilder = pageBuilder.getBlockBuilder(0);
    for (int i = 0; i < arrayLength; i++) {
        type.appendTo(block, positions.getInt(i), blockBuilder);
    }
    pageBuilder.declarePositions(arrayLength);
    return blockBuilder.getRegion(blockBuilder.getPositionCount() - arrayLength, arrayLength);
}
Also used : BlockBuilder(io.trino.spi.block.BlockBuilder) TypeParameter(io.trino.spi.function.TypeParameter) SqlType(io.trino.spi.function.SqlType)

Example 20 with TypeParameter

use of io.trino.spi.function.TypeParameter in project trino by trinodb.

the class ArrayCombinationsFunction method combinations.

@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));
}
Also used : ArrayType(io.trino.spi.type.ArrayType) PageBuilderStatus(io.trino.spi.block.PageBuilderStatus) DictionaryBlock(io.trino.spi.block.DictionaryBlock) TypeParameter(io.trino.spi.function.TypeParameter) SqlType(io.trino.spi.function.SqlType)

Aggregations

TypeParameter (io.trino.spi.function.TypeParameter)28 SqlType (io.trino.spi.function.SqlType)22 BlockBuilder (io.trino.spi.block.BlockBuilder)20 TypeParameterSpecialization (io.trino.spi.function.TypeParameterSpecialization)9 HyperLogLog (io.airlift.stats.cardinality.HyperLogLog)6 TypedSet (io.trino.operator.aggregation.TypedSet)6 InputFunction (io.trino.spi.function.InputFunction)6 TypedSet.createEqualityTypedSet (io.trino.operator.aggregation.TypedSet.createEqualityTypedSet)4 TrinoException (io.trino.spi.TrinoException)3 Block (io.trino.spi.block.Block)3 SqlNullable (io.trino.spi.function.SqlNullable)3 ArrayType (io.trino.spi.type.ArrayType)3 RowType (io.trino.spi.type.RowType)3 Type (io.trino.spi.type.Type)3 TypedSet.createDistinctTypedSet (io.trino.operator.aggregation.TypedSet.createDistinctTypedSet)2 MapType (io.trino.spi.type.MapType)2 Preconditions.checkArgument (com.google.common.base.Preconditions.checkArgument)1 HashMultimap (com.google.common.collect.HashMultimap)1 ImmutableList (com.google.common.collect.ImmutableList)1 ImmutableList.toImmutableList (com.google.common.collect.ImmutableList.toImmutableList)1