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());
}
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);
}
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);
}
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);
}
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));
}
Aggregations