Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@
import io.substrait.isthmus.expression.ScalarFunctionConverter;
import io.substrait.isthmus.expression.SqlArrayValueConstructorCallConverter;
import io.substrait.isthmus.expression.SqlMapValueConstructorCallConverter;
import io.substrait.isthmus.expression.TypeObserver;
import io.substrait.isthmus.expression.WindowFunctionConverter;
import io.substrait.plan.ImmutableExecutionBehavior;
import io.substrait.plan.Plan;
Expand Down Expand Up @@ -321,11 +322,23 @@ public ExpressionRexConverter getExpressionRexConverter(
getTypeFactory(),
getScalarFunctionConverter(),
getWindowFunctionConverter(),
getTypeConverter());
getTypeConverter(),
getTypeObserver());
erc.setRelNodeConverter(relNodeConverter);
return erc;
}

/**
* Returns the observer for supplied and independently inferred expression types.
*
* <p>Override to collect type observations during Substrait-to-Calcite conversion.
*
* @return a no-op observer by default
*/
public TypeObserver getTypeObserver() {
return TypeObserver.NOOP;
}

/**
* A {@link RelBuilder} is a Calcite class used for creating {@link
* org.apache.calcite.rel.RelNode}s.
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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;

Expand All @@ -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");
}

/**
Expand Down Expand Up @@ -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);
if (typeObserver == TypeObserver.NOOP) {
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);
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(
Expand Down
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.
* 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);
}
Loading
Loading