Examples with FAIL_ON_NULL io.trino.spi.function.InvocationConvention.InvocationReturnConvention.FAIL_ON_NULL used on opensource projects

Example 1 with FAIL_ON_NULL

use of io.trino.spi.function.InvocationConvention.InvocationReturnConvention.FAIL_ON_NULL in project trino by trinodb.

the class ScalarFunctionAdapter method adaptParameter.

private MethodHandle adaptParameter(MethodHandle methodHandle, int parameterIndex, Type argumentType, InvocationArgumentConvention actualArgumentConvention, InvocationArgumentConvention expectedArgumentConvention, InvocationReturnConvention returnConvention) {
    if (actualArgumentConvention == expectedArgumentConvention) {
        return methodHandle;
    }
    if (actualArgumentConvention == BLOCK_POSITION) {
        throw new IllegalArgumentException("Block and position argument can not be adapted");
    }
    if (actualArgumentConvention == FUNCTION) {
        throw new IllegalArgumentException("Function argument can not be adapted");
    }
    // caller will never pass null
    if (expectedArgumentConvention == NEVER_NULL) {
        if (actualArgumentConvention == BOXED_NULLABLE) {
            // if actual argument is boxed primitive, change method handle to accept a primitive and then box to actual method
            if (isWrapperType(methodHandle.type().parameterType(parameterIndex))) {
                MethodType targetType = methodHandle.type().changeParameterType(parameterIndex, unwrap(methodHandle.type().parameterType(parameterIndex)));
                methodHandle = explicitCastArguments(methodHandle, targetType);
            }
            return methodHandle;
        }
        if (actualArgumentConvention == NULL_FLAG) {
            // actual method takes value and null flag, so change method handle to not have the flag and always pass false to the actual method
            return insertArguments(methodHandle, parameterIndex + 1, false);
        }
        throw new IllegalArgumentException("Unsupported actual argument convention: " + actualArgumentConvention);
    }
    // caller will pass Java null for SQL null
    if (expectedArgumentConvention == BOXED_NULLABLE) {
        if (actualArgumentConvention == NEVER_NULL) {
            if (nullAdaptationPolicy == UNSUPPORTED) {
                throw new IllegalArgumentException("Not null argument can not be adapted to nullable");
            }
            // box argument
            Class<?> boxedType = wrap(methodHandle.type().parameterType(parameterIndex));
            MethodType targetType = methodHandle.type().changeParameterType(parameterIndex, boxedType);
            methodHandle = explicitCastArguments(methodHandle, targetType);
            if (nullAdaptationPolicy == UNDEFINED_VALUE_FOR_NULL) {
                // currently, we just perform unboxing, which converts nulls to Java primitive default value
                return methodHandle;
            }
            if (nullAdaptationPolicy == RETURN_NULL_ON_NULL) {
                if (returnConvention == FAIL_ON_NULL) {
                    throw new IllegalArgumentException("RETURN_NULL_ON_NULL adaptation can not be used with FAIL_ON_NULL return convention");
                }
                return guardWithTest(isNullArgument(methodHandle.type(), parameterIndex), returnNull(methodHandle.type()), methodHandle);
            }
            if (nullAdaptationPolicy == THROW_ON_NULL) {
                MethodType adapterType = methodType(boxedType, boxedType);
                MethodHandle adapter = guardWithTest(isNullArgument(adapterType, 0), throwTrinoNullArgumentException(adapterType), identity(boxedType));
                return collectArguments(methodHandle, parameterIndex, adapter);
            }
        }
        if (actualArgumentConvention == NULL_FLAG) {
            // The conversion is described below in reverse order as this is how method handle adaptation works.  The provided example
            // signature is based on a boxed Long argument.
            // 3. unbox the value (if null the java default is sent)
            // long, boolean => Long, boolean
            Class<?> parameterType = methodHandle.type().parameterType(parameterIndex);
            methodHandle = explicitCastArguments(methodHandle, methodHandle.type().changeParameterType(parameterIndex, wrap(parameterType)));
            // 2. replace second argument with the result of isNull
            // long, boolean => Long, Long
            methodHandle = filterArguments(methodHandle, parameterIndex + 1, explicitCastArguments(IS_NULL_METHOD, methodType(boolean.class, wrap(parameterType))));
            // 1. Duplicate the argument, so we have two copies of the value
            // Long, Long => Long
            int[] reorder = IntStream.range(0, methodHandle.type().parameterCount()).map(i -> i <= parameterIndex ? i : i - 1).toArray();
            MethodType newType = methodHandle.type().dropParameterTypes(parameterIndex + 1, parameterIndex + 2);
            methodHandle = permuteArguments(methodHandle, newType, reorder);
            return methodHandle;
        }
        throw new IllegalArgumentException("Unsupported actual argument convention: " + actualArgumentConvention);
    }
    // caller will pass boolean true in the next argument for SQL null
    if (expectedArgumentConvention == NULL_FLAG) {
        if (actualArgumentConvention == NEVER_NULL) {
            if (nullAdaptationPolicy == UNSUPPORTED) {
                throw new IllegalArgumentException("Not null argument can not be adapted to nullable");
            }
            if (nullAdaptationPolicy == UNDEFINED_VALUE_FOR_NULL) {
                // add null flag to call
                methodHandle = dropArguments(methodHandle, parameterIndex + 1, boolean.class);
                return methodHandle;
            }
            // if caller sets null flag, return null, otherwise invoke target
            if (nullAdaptationPolicy == RETURN_NULL_ON_NULL) {
                if (returnConvention == FAIL_ON_NULL) {
                    throw new IllegalArgumentException("RETURN_NULL_ON_NULL adaptation can not be used with FAIL_ON_NULL return convention");
                }
                // add null flag to call
                methodHandle = dropArguments(methodHandle, parameterIndex + 1, boolean.class);
                return guardWithTest(isTrueNullFlag(methodHandle.type(), parameterIndex), returnNull(methodHandle.type()), methodHandle);
            }
            if (nullAdaptationPolicy == THROW_ON_NULL) {
                MethodHandle adapter = identity(methodHandle.type().parameterType(parameterIndex));
                adapter = dropArguments(adapter, 1, boolean.class);
                adapter = guardWithTest(isTrueNullFlag(adapter.type(), 0), throwTrinoNullArgumentException(adapter.type()), adapter);
                return collectArguments(methodHandle, parameterIndex, adapter);
            }
        }
        if (actualArgumentConvention == BOXED_NULLABLE) {
            return collectArguments(methodHandle, parameterIndex, boxedToNullFlagFilter(methodHandle.type().parameterType(parameterIndex)));
        }
        throw new IllegalArgumentException("Unsupported actual argument convention: " + actualArgumentConvention);
    }
    // caller will pass boolean true in the next argument for SQL null
    if (expectedArgumentConvention == BLOCK_POSITION) {
        MethodHandle getBlockValue = getBlockValue(argumentType, methodHandle.type().parameterType(parameterIndex));
        if (actualArgumentConvention == NEVER_NULL) {
            if (nullAdaptationPolicy == UNDEFINED_VALUE_FOR_NULL) {
                // Current, null is not checked, so whatever type returned is passed through
                methodHandle = collectArguments(methodHandle, parameterIndex, getBlockValue);
                return methodHandle;
            }
            if (nullAdaptationPolicy == RETURN_NULL_ON_NULL && returnConvention != FAIL_ON_NULL) {
                // if caller sets null flag, return null, otherwise invoke target
                methodHandle = collectArguments(methodHandle, parameterIndex, getBlockValue);
                return guardWithTest(isBlockPositionNull(methodHandle.type(), parameterIndex), returnNull(methodHandle.type()), methodHandle);
            }
            if (nullAdaptationPolicy == THROW_ON_NULL || nullAdaptationPolicy == UNSUPPORTED || nullAdaptationPolicy == RETURN_NULL_ON_NULL) {
                MethodHandle adapter = guardWithTest(isBlockPositionNull(getBlockValue.type(), 0), throwTrinoNullArgumentException(getBlockValue.type()), getBlockValue);
                return collectArguments(methodHandle, parameterIndex, adapter);
            }
        }
        if (actualArgumentConvention == BOXED_NULLABLE) {
            getBlockValue = explicitCastArguments(getBlockValue, getBlockValue.type().changeReturnType(wrap(getBlockValue.type().returnType())));
            getBlockValue = guardWithTest(isBlockPositionNull(getBlockValue.type(), 0), returnNull(getBlockValue.type()), getBlockValue);
            methodHandle = collectArguments(methodHandle, parameterIndex, getBlockValue);
            return methodHandle;
        }
        if (actualArgumentConvention == NULL_FLAG) {
            // long, boolean => long, Block, int
            MethodHandle isNull = isBlockPositionNull(getBlockValue.type(), 0);
            methodHandle = collectArguments(methodHandle, parameterIndex + 1, isNull);
            // long, Block, int => Block, int, Block, int
            getBlockValue = guardWithTest(isBlockPositionNull(getBlockValue.type(), 0), returnNull(getBlockValue.type()), getBlockValue);
            methodHandle = collectArguments(methodHandle, parameterIndex, getBlockValue);
            int[] reorder = IntStream.range(0, methodHandle.type().parameterCount()).map(i -> i <= parameterIndex + 1 ? i : i - 2).toArray();
            MethodType newType = methodHandle.type().dropParameterTypes(parameterIndex + 2, parameterIndex + 4);
            methodHandle = permuteArguments(methodHandle, newType, reorder);
            return methodHandle;
        }
        throw new IllegalArgumentException("Unsupported actual argument convention: " + actualArgumentConvention);
    }
    throw new IllegalArgumentException("Unsupported expected argument convention: " + expectedArgumentConvention);
}

Example 2 with FAIL_ON_NULL

use of io.trino.spi.function.InvocationConvention.InvocationReturnConvention.FAIL_ON_NULL in project trino by trinodb.

the class FormatFunction method valueConverter.

private static BiFunction<ConnectorSession, Block, Object> valueConverter(FunctionDependencies functionDependencies, Type type, int position) {
    if (type.equals(UNKNOWN)) {
        return (session, block) -> null;
    }
    if (type.equals(BOOLEAN)) {
        return (session, block) -> type.getBoolean(block, position);
    }
    if (type.equals(TINYINT) || type.equals(SMALLINT) || type.equals(INTEGER) || type.equals(BIGINT)) {
        return (session, block) -> type.getLong(block, position);
    }
    if (type.equals(REAL)) {
        return (session, block) -> intBitsToFloat(toIntExact(type.getLong(block, position)));
    }
    if (type.equals(DOUBLE)) {
        return (session, block) -> type.getDouble(block, position);
    }
    if (type.equals(DATE)) {
        return (session, block) -> LocalDate.ofEpochDay(type.getLong(block, position));
    }
    if (type instanceof TimestampWithTimeZoneType) {
        return (session, block) -> toZonedDateTime(((TimestampWithTimeZoneType) type), block, position);
    }
    if (type instanceof TimestampType) {
        return (session, block) -> toLocalDateTime(((TimestampType) type), block, position);
    }
    if (type instanceof TimeType) {
        return (session, block) -> toLocalTime(type.getLong(block, position));
    }
    // TODO: support TIME WITH TIME ZONE by https://github.com/trinodb/trino/issues/191 + mapping to java.time.OffsetTime
    if (type.equals(JSON)) {
        MethodHandle handle = functionDependencies.getFunctionInvoker(QualifiedName.of("json_format"), ImmutableList.of(JSON), simpleConvention(FAIL_ON_NULL, NEVER_NULL)).getMethodHandle();
        return (session, block) -> convertToString(handle, type.getSlice(block, position));
    }
    if (isShortDecimal(type)) {
        int scale = ((DecimalType) type).getScale();
        return (session, block) -> BigDecimal.valueOf(type.getLong(block, position), scale);
    }
    if (isLongDecimal(type)) {
        int scale = ((DecimalType) type).getScale();
        return (session, block) -> new BigDecimal(((Int128) type.getObject(block, position)).toBigInteger(), scale);
    }
    if (type instanceof VarcharType) {
        return (session, block) -> type.getSlice(block, position).toStringUtf8();
    }
    if (type instanceof CharType) {
        CharType charType = (CharType) type;
        return (session, block) -> padSpaces(type.getSlice(block, position), charType).toStringUtf8();
    }
    BiFunction<ConnectorSession, Block, Object> function;
    if (type.getJavaType() == long.class) {
        function = (session, block) -> type.getLong(block, position);
    } else if (type.getJavaType() == double.class) {
        function = (session, block) -> type.getDouble(block, position);
    } else if (type.getJavaType() == boolean.class) {
        function = (session, block) -> type.getBoolean(block, position);
    } else if (type.getJavaType() == Slice.class) {
        function = (session, block) -> type.getSlice(block, position);
    } else {
        function = (session, block) -> type.getObject(block, position);
    }
    MethodHandle handle = functionDependencies.getCastInvoker(type, VARCHAR, simpleConvention(FAIL_ON_NULL, NEVER_NULL)).getMethodHandle();
    return (session, block) -> convertToString(handle, function.apply(session, block));
}

Example 3 with FAIL_ON_NULL

use of io.trino.spi.function.InvocationConvention.InvocationReturnConvention.FAIL_ON_NULL in project trino by trinodb.

the class TestPolymorphicScalarFunction method testSelectsMultipleChoiceWithBlockPosition.

@Test
public void testSelectsMultipleChoiceWithBlockPosition() throws Throwable {
    Signature signature = Signature.builder().operatorType(IS_DISTINCT_FROM).argumentTypes(DECIMAL_SIGNATURE, DECIMAL_SIGNATURE).returnType(BOOLEAN.getTypeSignature()).build();
    SqlScalarFunction function = new PolymorphicScalarFunctionBuilder(TestMethods.class).signature(signature).argumentNullability(true, true).deterministic(true).choice(choice -> choice.argumentProperties(NULL_FLAG, NULL_FLAG).implementation(methodsGroup -> methodsGroup.methods("shortShort", "longLong"))).choice(choice -> choice.argumentProperties(BLOCK_POSITION, BLOCK_POSITION).implementation(methodsGroup -> methodsGroup.methodWithExplicitJavaTypes("blockPositionLongLong", asList(Optional.of(Int128.class), Optional.of(Int128.class))).methodWithExplicitJavaTypes("blockPositionShortShort", asList(Optional.of(long.class), Optional.of(long.class))))).build();
    BoundSignature shortDecimalBoundSignature = new BoundSignature(signature.getName(), BOOLEAN, ImmutableList.of(SHORT_DECIMAL_BOUND_TYPE, SHORT_DECIMAL_BOUND_TYPE));
    ChoicesScalarFunctionImplementation functionImplementation = (ChoicesScalarFunctionImplementation) function.specialize(shortDecimalBoundSignature, new FunctionDependencies(FUNCTION_MANAGER::getScalarFunctionInvoker, ImmutableMap.of(), ImmutableSet.of()));
    assertEquals(functionImplementation.getChoices().size(), 2);
    assertEquals(functionImplementation.getChoices().get(0).getInvocationConvention(), new InvocationConvention(ImmutableList.of(NULL_FLAG, NULL_FLAG), FAIL_ON_NULL, false, false));
    assertEquals(functionImplementation.getChoices().get(1).getInvocationConvention(), new InvocationConvention(ImmutableList.of(BLOCK_POSITION, BLOCK_POSITION), FAIL_ON_NULL, false, false));
    Block block1 = new LongArrayBlock(0, Optional.empty(), new long[0]);
    Block block2 = new LongArrayBlock(0, Optional.empty(), new long[0]);
    assertFalse((boolean) functionImplementation.getChoices().get(1).getMethodHandle().invoke(block1, 0, block2, 0));
    BoundSignature longDecimalBoundSignature = new BoundSignature(signature.getName(), BOOLEAN, ImmutableList.of(LONG_DECIMAL_BOUND_TYPE, LONG_DECIMAL_BOUND_TYPE));
    functionImplementation = (ChoicesScalarFunctionImplementation) function.specialize(longDecimalBoundSignature, new FunctionDependencies(FUNCTION_MANAGER::getScalarFunctionInvoker, ImmutableMap.of(), ImmutableSet.of()));
    assertTrue((boolean) functionImplementation.getChoices().get(1).getMethodHandle().invoke(block1, 0, block2, 0));
}

Example 4 with FAIL_ON_NULL

use of io.trino.spi.function.InvocationConvention.InvocationReturnConvention.FAIL_ON_NULL in project trino by trinodb.

the class RowType method getDistinctFromOperatorInvokers.

private static List<OperatorMethodHandle> getDistinctFromOperatorInvokers(TypeOperators typeOperators, List<Field> fields) {
    boolean comparable = fields.stream().allMatch(field -> field.getType().isComparable());
    if (!comparable) {
        return emptyList();
    }
    // for large rows, use a generic loop with a megamorphic call site
    if (fields.size() > MEGAMORPHIC_FIELD_COUNT) {
        List<MethodHandle> distinctFromOperators = fields.stream().map(field -> typeOperators.getDistinctFromOperator(field.getType(), simpleConvention(FAIL_ON_NULL, BLOCK_POSITION, BLOCK_POSITION))).collect(toUnmodifiableList());
        return singletonList(new OperatorMethodHandle(DISTINCT_FROM_CONVENTION, DISTINCT_FROM.bindTo(distinctFromOperators)));
    }
    // (Block, Block):boolean
    MethodHandle distinctFrom = dropArguments(constant(boolean.class, false), 0, Block.class, Block.class);
    for (int fieldId = 0; fieldId < fields.size(); fieldId++) {
        Field field = fields.get(fieldId);
        // (Block, Block, int, MethodHandle, Block, Block):boolean
        distinctFrom = collectArguments(CHAIN_DISTINCT_FROM, 0, distinctFrom);
        // field distinctFrom
        MethodHandle fieldDistinctFromOperator = typeOperators.getDistinctFromOperator(field.getType(), simpleConvention(FAIL_ON_NULL, BLOCK_POSITION, BLOCK_POSITION));
        // (Block, Block, Block, Block):boolean
        distinctFrom = insertArguments(distinctFrom, 2, fieldId, fieldDistinctFromOperator);
        // (Block, Block):boolean
        distinctFrom = permuteArguments(distinctFrom, methodType(boolean.class, Block.class, Block.class), 0, 1, 0, 1);
    }
    distinctFrom = CHAIN_DISTINCT_FROM_START.bindTo(distinctFrom);
    return singletonList(new OperatorMethodHandle(DISTINCT_FROM_CONVENTION, distinctFrom));
}

Example 5 with FAIL_ON_NULL

use of io.trino.spi.function.InvocationConvention.InvocationReturnConvention.FAIL_ON_NULL in project trino by trinodb.

the class RowType method getComparisonOperatorInvokers.

private static List<OperatorMethodHandle> getComparisonOperatorInvokers(BiFunction<Type, InvocationConvention, MethodHandle> comparisonOperatorFactory, List<Field> fields) {
    boolean orderable = fields.stream().allMatch(field -> field.getType().isOrderable());
    if (!orderable) {
        return emptyList();
    }
    // for large rows, use a generic loop with a megamorphic call site
    if (fields.size() > MEGAMORPHIC_FIELD_COUNT) {
        List<MethodHandle> comparisonOperators = fields.stream().map(field -> comparisonOperatorFactory.apply(field.getType(), simpleConvention(FAIL_ON_NULL, BLOCK_POSITION, BLOCK_POSITION))).collect(toUnmodifiableList());
        return singletonList(new OperatorMethodHandle(COMPARISON_CONVENTION, COMPARISON.bindTo(comparisonOperators)));
    }
    // (Block, Block):Boolean
    MethodHandle comparison = dropArguments(constant(long.class, 0), 0, Block.class, Block.class);
    for (int fieldId = 0; fieldId < fields.size(); fieldId++) {
        Field field = fields.get(fieldId);
        // (Block, Block, int, MethodHandle, Block, Block):Boolean
        comparison = collectArguments(CHAIN_COMPARISON, 0, comparison);
        // field comparison
        MethodHandle fieldComparisonOperator = comparisonOperatorFactory.apply(field.getType(), simpleConvention(FAIL_ON_NULL, BLOCK_POSITION, BLOCK_POSITION));
        // (Block, Block, Block, Block):Boolean
        comparison = insertArguments(comparison, 2, fieldId, fieldComparisonOperator);
        // (Block, Block):Boolean
        comparison = permuteArguments(comparison, methodType(long.class, Block.class, Block.class), 0, 1, 0, 1);
    }
    return singletonList(new OperatorMethodHandle(COMPARISON_CONVENTION, comparison));
}