use of io.confluent.ksql.execution.expression.tree.LambdaFunctionCall in project ksql by confluentinc.
the class SqlToJavaVisitorTest method shouldGenerateCorrectCodeForLambdaExpression.
@Test
public void shouldGenerateCorrectCodeForLambdaExpression() {
// Given:
final UdfFactory udfFactory = mock(UdfFactory.class);
final KsqlScalarFunction udf = mock(KsqlScalarFunction.class);
givenUdf("ABS", udfFactory, udf, SqlTypes.STRING);
givenUdf("TRANSFORM", udfFactory, udf, SqlTypes.STRING);
when(udf.parameters()).thenReturn(ImmutableList.of(ArrayType.of(ParamTypes.DOUBLE), LambdaType.of(ImmutableList.of(ParamTypes.DOUBLE), ParamTypes.DOUBLE)));
final Expression expression = new FunctionCall(FunctionName.of("TRANSFORM"), ImmutableList.of(ARRAYCOL, new LambdaFunctionCall(ImmutableList.of("x"), (new FunctionCall(FunctionName.of("ABS"), ImmutableList.of(new LambdaVariable("X")))))));
// When:
final String javaExpression = sqlToJavaVisitor.process(expression);
// Then
assertThat(javaExpression, equalTo("((String) TRANSFORM_0.evaluate(COL4, new Function() {\n @Override\n public Object apply(Object arg1) {\n final Double x = (Double) arg1;\n return ((String) ABS_1.evaluate(X));\n }\n}))"));
}
use of io.confluent.ksql.execution.expression.tree.LambdaFunctionCall in project ksql by confluentinc.
the class FunctionArgumentsUtil method getFunctionTypeInfo.
/**
* Compute type information given a function call node. Specifically, computes
* the function return type, the types of all arguments, and the types of all
* lambda parameters for arguments that are lambda expressions.
*
* <p>Given a function call node, we have to do a two pass processing of the
* function arguments in order to properly handle any potential lambda functions.</p>
*
* <p>In the first pass, if there are lambda functions, we create a SqlLambda that only contains
* the number of input arguments for the lambda. We pass this first argument list
* to UdfFactory in order to get the correct function. We can make this assumption
* due to Java's handling of type erasure (Function(T,R) is considered the same as
* Function(U,R)).</p>
*
* <p>In the second pass, we use the LambdaType inputTypes field to construct SqlLambdaResolved
* that has the proper input type list and return type. We also need to construct a list of
* lambda type mapping that should be used when processing each function argument subtree.</p>
*
* @param expressionTypeManager an expression type manager
* @param functionCall the function expression
* @param udfFactory udf factory for the function in the expression
* @param lambdaMapping a type context
*
* @return a wrapper that contains a list of function arguments
* (any lambdas are SqlLambdaResolved), the ksql function,
* type contexts for use in further processing the function
* argument child nodes, and the return type of the ksql function
*/
// CHECKSTYLE_RULES.OFF: CyclomaticComplexity
public static FunctionTypeInfo getFunctionTypeInfo(final ExpressionTypeManager expressionTypeManager, final FunctionCall functionCall, final UdfFactory udfFactory, final Map<String, SqlType> lambdaMapping) {
// CHECKSTYLE_RULES.ON: CyclomaticComplexity
final List<Expression> arguments = functionCall.getArguments();
final List<SqlArgument> functionArgumentTypes = firstPassOverFunctionArguments(arguments, expressionTypeManager, lambdaMapping);
final KsqlScalarFunction function = udfFactory.getFunction(functionArgumentTypes);
final SqlType returnSchema;
final List<ArgumentInfo> argumentInfoForFunction = new ArrayList<>();
if (!functionCall.hasLambdaFunctionCallArguments()) {
returnSchema = function.getReturnType(functionArgumentTypes);
return FunctionTypeInfo.of(functionArgumentTypes.stream().map(argument -> ArgumentInfo.of(argument, new HashMap<>(lambdaMapping))).collect(Collectors.toList()), returnSchema, function);
} else {
final List<ParamType> paramTypes = function.parameters();
final Map<GenericType, SqlType> reservedGenerics = new HashMap<>();
final List<SqlArgument> functionArgumentTypesWithResolvedLambdaType = new ArrayList<>();
// second pass over the function arguments to properly do lambda type checking
for (int i = 0; i < arguments.size(); i++) {
final Expression expression = arguments.get(i);
final ParamType parameter = paramTypes.get(i);
if (expression instanceof LambdaFunctionCall) {
// lambda node at this index in the function node argument list
if (!(parameter instanceof LambdaType)) {
throw new RuntimeException(String.format("Error while processing lambda function." + "Expected lambda parameter but was %s" + "This is most likely an internal error and a " + "Github issue should be filed for debugging. " + "Include the function name, the parameters passed in, the expected " + "signature, and any other relevant information.", parameter.toString()));
}
final ArrayList<SqlType> lambdaSqlTypes = new ArrayList<>();
final Map<String, SqlType> variableTypeMapping = mapLambdaParametersToTypes((LambdaFunctionCall) expression, (LambdaType) parameter, reservedGenerics, lambdaSqlTypes);
final Map<String, SqlType> updateLambdaMapping = LambdaMappingUtil.resolveOldAndNewLambdaMapping(variableTypeMapping, lambdaMapping);
final SqlType resolvedLambdaReturnType = expressionTypeManager.getExpressionSqlType(expression, updateLambdaMapping);
final SqlArgument lambdaArgument = SqlArgument.of(SqlLambdaResolved.of(lambdaSqlTypes, resolvedLambdaReturnType));
functionArgumentTypesWithResolvedLambdaType.add(lambdaArgument);
argumentInfoForFunction.add(ArgumentInfo.of(lambdaArgument, new HashMap<>(updateLambdaMapping)));
} else {
functionArgumentTypesWithResolvedLambdaType.add(functionArgumentTypes.get(i));
argumentInfoForFunction.add(ArgumentInfo.of(functionArgumentTypes.get(i), new HashMap<>(lambdaMapping)));
}
if (GenericsUtil.hasGenerics(parameter)) {
final Pair<Boolean, Optional<KsqlException>> success = GenericsUtil.reserveGenerics(parameter, functionArgumentTypesWithResolvedLambdaType.get(i), reservedGenerics);
if (!success.getLeft() && success.getRight().isPresent()) {
throw success.getRight().get();
}
}
}
returnSchema = function.getReturnType(functionArgumentTypesWithResolvedLambdaType);
return new FunctionTypeInfo(argumentInfoForFunction, returnSchema, function);
}
}
use of io.confluent.ksql.execution.expression.tree.LambdaFunctionCall in project ksql by confluentinc.
the class ExpressionTypeManagerTest method shouldEvaluateAnyNumberOfArgumentLambda.
@Test
public void shouldEvaluateAnyNumberOfArgumentLambda() {
// Given:
givenUdfWithNameAndReturnType("TRANSFORM", SqlTypes.STRING);
when(function.parameters()).thenReturn(ImmutableList.of(ArrayType.of(DoubleType.INSTANCE), StringType.INSTANCE, MapType.of(LongType.INSTANCE, DoubleType.INSTANCE), LambdaType.of(ImmutableList.of(DoubleType.INSTANCE, StringType.INSTANCE, LongType.INSTANCE, DoubleType.INSTANCE), StringType.INSTANCE)));
final Expression expression = new FunctionCall(FunctionName.of("TRANSFORM"), ImmutableList.of(ARRAYCOL, new StringLiteral("Q"), MAPCOL, new LambdaFunctionCall(ImmutableList.of("A", "B", "C", "D"), new ArithmeticBinaryExpression(Operator.ADD, new LambdaVariable("C"), new IntegerLiteral(5)))));
// When:
final SqlType exprType = expressionTypeManager.getExpressionSqlType(expression);
// Then:
assertThat(exprType, is(SqlTypes.STRING));
verify(udfFactory).getFunction(ImmutableList.of(SqlArgument.of(SqlTypes.array(SqlTypes.DOUBLE)), SqlArgument.of(SqlTypes.STRING), SqlArgument.of(SqlTypes.map(SqlTypes.BIGINT, SqlTypes.DOUBLE)), SqlArgument.of(SqlLambda.of(4))));
verify(function).getReturnType(ImmutableList.of(SqlArgument.of(SqlTypes.array(SqlTypes.DOUBLE)), SqlArgument.of(SqlTypes.STRING), SqlArgument.of(SqlTypes.map(SqlTypes.BIGINT, SqlTypes.DOUBLE)), SqlArgument.of(SqlLambdaResolved.of(ImmutableList.of(SqlTypes.DOUBLE, SqlTypes.STRING, SqlTypes.BIGINT, SqlTypes.DOUBLE), SqlTypes.BIGINT))));
}
use of io.confluent.ksql.execution.expression.tree.LambdaFunctionCall in project ksql by confluentinc.
the class ExpressionTypeManagerTest method shouldHandleMultipleLambdasInSameFunctionCallWithDifferentVariableNames.
@Test
public void shouldHandleMultipleLambdasInSameFunctionCallWithDifferentVariableNames() {
// Given:
givenUdfWithNameAndReturnType("TRANSFORM", SqlTypes.INTEGER);
when(function.parameters()).thenReturn(ImmutableList.of(MapType.of(LongType.INSTANCE, DoubleType.INSTANCE), IntegerType.INSTANCE, LambdaType.of(ImmutableList.of(DoubleType.INSTANCE, DoubleType.INSTANCE), StringType.INSTANCE), LambdaType.of(ImmutableList.of(DoubleType.INSTANCE, DoubleType.INSTANCE), StringType.INSTANCE)));
final Expression expression = new ArithmeticBinaryExpression(Operator.ADD, new FunctionCall(FunctionName.of("TRANSFORM"), ImmutableList.of(MAPCOL, new IntegerLiteral(0), new LambdaFunctionCall(ImmutableList.of("A", "B"), new ArithmeticBinaryExpression(Operator.ADD, new LambdaVariable("A"), new LambdaVariable("B"))), new LambdaFunctionCall(ImmutableList.of("K", "V"), new ArithmeticBinaryExpression(Operator.ADD, new LambdaVariable("K"), new LambdaVariable("V"))))), new IntegerLiteral(5));
// When:
final SqlType result = expressionTypeManager.getExpressionSqlType(expression);
assertThat(result, is(SqlTypes.INTEGER));
}
use of io.confluent.ksql.execution.expression.tree.LambdaFunctionCall in project ksql by confluentinc.
the class ExpressionTypeManagerTest method shouldEvaluateLambdaInUDFWithMap.
@Test
public void shouldEvaluateLambdaInUDFWithMap() {
// Given:
givenUdfWithNameAndReturnType("TRANSFORM", SqlTypes.DOUBLE);
when(function.parameters()).thenReturn(ImmutableList.of(MapType.of(DoubleType.INSTANCE, DoubleType.INSTANCE), LambdaType.of(ImmutableList.of(LongType.INSTANCE, DoubleType.INSTANCE), DoubleType.INSTANCE)));
final Expression expression = new FunctionCall(FunctionName.of("TRANSFORM"), ImmutableList.of(MAPCOL, new LambdaFunctionCall(ImmutableList.of("X", "Y"), new ArithmeticBinaryExpression(Operator.ADD, new LambdaVariable("X"), new IntegerLiteral(5)))));
// When:
final SqlType exprType = expressionTypeManager.getExpressionSqlType(expression);
// Then:
assertThat(exprType, is(SqlTypes.DOUBLE));
verify(udfFactory).getFunction(ImmutableList.of(SqlArgument.of(SqlTypes.map(SqlTypes.BIGINT, SqlTypes.DOUBLE)), SqlArgument.of(SqlLambda.of(2))));
verify(function).getReturnType(ImmutableList.of(SqlArgument.of(SqlTypes.map(SqlTypes.BIGINT, SqlTypes.DOUBLE)), SqlArgument.of(SqlLambdaResolved.of(ImmutableList.of(SqlTypes.BIGINT, SqlTypes.DOUBLE), SqlTypes.BIGINT))));
}
Aggregations