-
Notifications
You must be signed in to change notification settings - Fork 95
feat(isthmus): observe scalar function types #1015
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Changes from all commits
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -28,8 +28,10 @@ | |
| import java.math.BigDecimal; | ||
| import java.util.Collections; | ||
| import java.util.List; | ||
| import java.util.Objects; | ||
| import java.util.Set; | ||
| import java.util.concurrent.TimeUnit; | ||
| import java.util.function.Supplier; | ||
| import java.util.stream.Collectors; | ||
| import java.util.stream.IntStream; | ||
| import java.util.stream.Stream; | ||
|
|
@@ -97,6 +99,9 @@ public class ExpressionRexConverter | |
| /** Converter for Substrait window function invocations to Calcite {@link SqlOperator}s. */ | ||
| protected final WindowFunctionConverter windowFunctionConverter; | ||
|
|
||
| /** Observer for supplied and inferred expression types. */ | ||
| protected final TypeObserver typeObserver; | ||
|
|
||
| /** Converter for Substrait relational nodes to Calcite {@link RelNode}s, used for subqueries. */ | ||
| protected SubstraitRelNodeConverter relNodeConverter; | ||
|
|
||
|
|
@@ -115,11 +120,35 @@ public ExpressionRexConverter( | |
| ScalarFunctionConverter scalarFunctionConverter, | ||
| WindowFunctionConverter windowFunctionConverter, | ||
| TypeConverter typeConverter) { | ||
| this( | ||
| typeFactory, | ||
| scalarFunctionConverter, | ||
| windowFunctionConverter, | ||
| typeConverter, | ||
| TypeObserver.NOOP); | ||
| } | ||
|
|
||
| /** | ||
| * Creates an {@code ExpressionRexConverter} with type observation enabled. | ||
| * | ||
| * @param typeFactory Calcite type factory for type creation | ||
| * @param scalarFunctionConverter converter for scalar function invocations | ||
| * @param windowFunctionConverter converter for window function invocations | ||
| * @param typeConverter converter for Substrait and Calcite type mappings | ||
| * @param typeObserver observer for supplied and independently inferred expression types | ||
| */ | ||
| public ExpressionRexConverter( | ||
| RelDataTypeFactory typeFactory, | ||
| ScalarFunctionConverter scalarFunctionConverter, | ||
| WindowFunctionConverter windowFunctionConverter, | ||
| TypeConverter typeConverter, | ||
| TypeObserver typeObserver) { | ||
| this.typeFactory = typeFactory; | ||
| this.typeConverter = typeConverter; | ||
| this.rexBuilder = new RexBuilder(typeFactory); | ||
| this.scalarFunctionConverter = scalarFunctionConverter; | ||
| this.windowFunctionConverter = windowFunctionConverter; | ||
| this.typeObserver = Objects.requireNonNull(typeObserver, "typeObserver"); | ||
| } | ||
|
|
||
| /** | ||
|
|
@@ -522,13 +551,49 @@ public RexNode visit(Expression.ScalarFunctionInvocation expr, Context context) | |
|
|
||
| RelDataType returnType = typeConverter.toCalcite(typeFactory, expr.outputType()); | ||
| if (operator == SqlStdOperatorTable.CONCAT && args.size() > 2) { | ||
| return args.stream() | ||
| .skip(1) | ||
| .reduce( | ||
| args.get(0), | ||
| (left, right) -> rexBuilder.makeCall(returnType, operator, List.of(left, right))); | ||
| RexNode suppliedCall = | ||
| args.stream() | ||
| .skip(1) | ||
| .reduce( | ||
| args.get(0), | ||
| (left, right) -> rexBuilder.makeCall(returnType, operator, List.of(left, right))); | ||
| if (typeObserver == TypeObserver.NOOP) { | ||
| return suppliedCall; | ||
| } | ||
| observeScalarType( | ||
| expr, | ||
| () -> | ||
| args.stream() | ||
| .skip(1) | ||
| .reduce( | ||
| args.get(0), | ||
| (left, right) -> rexBuilder.makeCall(operator, List.of(left, right)))); | ||
| return suppliedCall; | ||
| } | ||
| RexNode suppliedCall = rexBuilder.makeCall(returnType, operator, args); | ||
|
Member
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. minor: I would just suggest calling this something like |
||
| if (typeObserver == TypeObserver.NOOP) { | ||
|
Member
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. A comment to the effect of |
||
| return suppliedCall; | ||
| } | ||
| observeScalarType(expr, () -> rexBuilder.makeCall(operator, args)); | ||
| return suppliedCall; | ||
| } | ||
|
|
||
| private void observeScalarType( | ||
| Expression.ScalarFunctionInvocation expression, Supplier<RexNode> inferredCallSupplier) { | ||
| TypeObservation observation; | ||
| RexNode inferredCall; | ||
| try { | ||
| inferredCall = inferredCallSupplier.get(); | ||
| } catch (RuntimeException exception) { | ||
| observation = | ||
| TypeObservation.failure(TypeObservation.Source.SCALAR_FUNCTION, expression, exception); | ||
|
Member
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. minor: I might suggest I don't feel super strongly about it. |
||
| typeObserver.observe(observation); | ||
| return; | ||
| } | ||
| return rexBuilder.makeCall(returnType, operator, args); | ||
| observation = | ||
| TypeObservation.success( | ||
| TypeObservation.Source.SCALAR_FUNCTION, expression, inferredCall.getType()); | ||
| typeObserver.observe(observation); | ||
| } | ||
|
|
||
| private String callConversionFailureMessage( | ||
|
|
||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,109 @@ | ||
| package io.substrait.isthmus.expression; | ||
|
|
||
| import io.substrait.expression.Expression; | ||
| import io.substrait.type.Type; | ||
| import java.util.Objects; | ||
| import java.util.Optional; | ||
| import org.apache.calcite.rel.type.RelDataType; | ||
|
|
||
| /** | ||
| * The result of attempting to independently infer a Calcite type during expression conversion. | ||
|
Member
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. I would specify Substrait to Calcite expression conversion... though could we ever do TypeObservations in the opposite direction (not something to implement in this PR). |
||
| * Exactly one of {@link #inferredType()} and {@link #inferenceFailure()} is present. | ||
| */ | ||
| public final class TypeObservation { | ||
| /** The expression category that produced an observation. */ | ||
| public enum Source { | ||
| /** A scalar function invocation. */ | ||
| SCALAR_FUNCTION | ||
| } | ||
|
|
||
| private final Source source; | ||
| private final Expression expression; | ||
| private final RelDataType inferredType; | ||
| private final RuntimeException inferenceFailure; | ||
|
|
||
| /** | ||
| * Creates a successful type observation. | ||
| * | ||
| * @param source expression category that produced the observation | ||
| * @param expression expression that produced the observation | ||
| * @param inferredType type independently inferred by Calcite | ||
| * @return a successful type observation | ||
| */ | ||
| static TypeObservation success(Source source, Expression expression, RelDataType inferredType) { | ||
| return new TypeObservation(source, expression, inferredType, null); | ||
| } | ||
|
|
||
| /** | ||
| * Creates a failed type observation. | ||
| * | ||
| * @param source expression category that produced the observation | ||
| * @param expression expression that produced the observation | ||
| * @param inferenceFailure failure to independently infer a Calcite type | ||
| * @return a failed type observation | ||
| */ | ||
| static TypeObservation failure( | ||
| Source source, Expression expression, RuntimeException inferenceFailure) { | ||
| return new TypeObservation(source, expression, null, inferenceFailure); | ||
| } | ||
|
|
||
| private TypeObservation( | ||
| Source source, | ||
| Expression expression, | ||
| RelDataType inferredType, | ||
| RuntimeException inferenceFailure) { | ||
| this.source = Objects.requireNonNull(source, "source"); | ||
| this.expression = Objects.requireNonNull(expression, "expression"); | ||
| if ((inferredType == null) == (inferenceFailure == null)) { | ||
| throw new IllegalArgumentException( | ||
| "Exactly one of inferredType and inferenceFailure must be present"); | ||
| } | ||
| this.inferredType = inferredType; | ||
| this.inferenceFailure = inferenceFailure; | ||
| } | ||
|
|
||
| /** | ||
| * Returns the expression category that produced this observation. | ||
| * | ||
| * @return the expression category | ||
| */ | ||
| public Source source() { | ||
| return source; | ||
| } | ||
|
|
||
| /** | ||
| * Returns the expression that produced this observation. | ||
| * | ||
| * @return the observed expression | ||
| */ | ||
| public Expression expression() { | ||
| return expression; | ||
| } | ||
|
|
||
| /** | ||
| * Returns the type supplied by Substrait. | ||
| * | ||
| * @return the supplied type | ||
| */ | ||
| public Type suppliedType() { | ||
| return expression.getType(); | ||
| } | ||
|
|
||
| /** | ||
| * Returns the type independently inferred by Calcite. | ||
| * | ||
| * @return the inferred type, or empty if inference failed | ||
| */ | ||
| public Optional<RelDataType> inferredType() { | ||
| return Optional.ofNullable(inferredType); | ||
| } | ||
|
|
||
| /** | ||
| * Returns the Calcite inference failure. | ||
| * | ||
| * @return the inference failure, or empty if inference succeeded | ||
| */ | ||
| public Optional<RuntimeException> inferenceFailure() { | ||
| return Optional.ofNullable(inferenceFailure); | ||
| } | ||
| } | ||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,17 @@ | ||
| package io.substrait.isthmus.expression; | ||
|
|
||
| /** Receives type observations while converting Substrait expressions to Calcite. */ | ||
| @FunctionalInterface | ||
| public interface TypeObserver { | ||
| /** Observer that disables type inference and discards all observations. */ | ||
| TypeObserver NOOP = observation -> {}; | ||
|
|
||
| /** | ||
| * Receives the result of attempting to observe an expression's inferred type. | ||
| * | ||
| * <p>Exceptions thrown by an observer are propagated to the conversion caller. | ||
| * | ||
| * @param observation supplied type and either an inferred type or inference failure | ||
| */ | ||
| void observe(TypeObservation observation); | ||
| } |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I got a little nerdsniped seeing that you had to work around this special case CONCAT handling, and am moving to just fully remove the handling which is both cleaner overall for dealing with CONCAT, and should simplify your work here.
#1025