diff --git a/isthmus/src/main/java/io/substrait/isthmus/ConverterProvider.java b/isthmus/src/main/java/io/substrait/isthmus/ConverterProvider.java index 8cd3731b3..30c74397a 100644 --- a/isthmus/src/main/java/io/substrait/isthmus/ConverterProvider.java +++ b/isthmus/src/main/java/io/substrait/isthmus/ConverterProvider.java @@ -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; @@ -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. + * + *

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. diff --git a/isthmus/src/main/java/io/substrait/isthmus/expression/ExpressionRexConverter.java b/isthmus/src/main/java/io/substrait/isthmus/expression/ExpressionRexConverter.java index b5dd28551..308d34d91 100644 --- a/isthmus/src/main/java/io/substrait/isthmus/expression/ExpressionRexConverter.java +++ b/isthmus/src/main/java/io/substrait/isthmus/expression/ExpressionRexConverter.java @@ -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); + if (typeObserver == TypeObserver.NOOP) { + return suppliedCall; + } + observeScalarType(expr, () -> rexBuilder.makeCall(operator, args)); + return suppliedCall; + } + + private void observeScalarType( + Expression.ScalarFunctionInvocation expression, Supplier 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( diff --git a/isthmus/src/main/java/io/substrait/isthmus/expression/TypeObservation.java b/isthmus/src/main/java/io/substrait/isthmus/expression/TypeObservation.java new file mode 100644 index 000000000..b977c1f6d --- /dev/null +++ b/isthmus/src/main/java/io/substrait/isthmus/expression/TypeObservation.java @@ -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 inferredType() { + return Optional.ofNullable(inferredType); + } + + /** + * Returns the Calcite inference failure. + * + * @return the inference failure, or empty if inference succeeded + */ + public Optional inferenceFailure() { + return Optional.ofNullable(inferenceFailure); + } +} diff --git a/isthmus/src/main/java/io/substrait/isthmus/expression/TypeObserver.java b/isthmus/src/main/java/io/substrait/isthmus/expression/TypeObserver.java new file mode 100644 index 000000000..566ea3940 --- /dev/null +++ b/isthmus/src/main/java/io/substrait/isthmus/expression/TypeObserver.java @@ -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. + * + *

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); +} diff --git a/isthmus/src/test/java/io/substrait/isthmus/SubstraitExpressionConverterTest.java b/isthmus/src/test/java/io/substrait/isthmus/SubstraitExpressionConverterTest.java index fe7e24873..0809c577d 100644 --- a/isthmus/src/test/java/io/substrait/isthmus/SubstraitExpressionConverterTest.java +++ b/isthmus/src/test/java/io/substrait/isthmus/SubstraitExpressionConverterTest.java @@ -2,7 +2,10 @@ import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertInstanceOf; +import static org.junit.jupiter.api.Assertions.assertNotEquals; +import static org.junit.jupiter.api.Assertions.assertSame; import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; import io.substrait.expression.Expression; import io.substrait.expression.Expression.Switch; @@ -11,17 +14,26 @@ import io.substrait.isthmus.SubstraitRelNodeConverter.Context; import io.substrait.isthmus.expression.ExpressionRexConverter; import io.substrait.isthmus.expression.ScalarFunctionConverter; +import io.substrait.isthmus.expression.TypeObservation; +import io.substrait.isthmus.expression.TypeObserver; import io.substrait.isthmus.expression.WindowFunctionConverter; import io.substrait.relation.Project; import io.substrait.relation.Rel; import io.substrait.relation.Rel.Remap; import io.substrait.type.Type; +import java.util.ArrayList; import java.util.List; +import java.util.Optional; +import java.util.concurrent.atomic.AtomicReference; import org.apache.calcite.rel.RelNode; import org.apache.calcite.rel.logical.LogicalProject; import org.apache.calcite.rel.type.RelDataType; import org.apache.calcite.rex.RexNode; +import org.apache.calcite.sql.SqlFunction; +import org.apache.calcite.sql.SqlFunctionCategory; import org.apache.calcite.sql.SqlKind; +import org.apache.calcite.sql.SqlOperator; +import org.apache.calcite.sql.type.SqlTypeName; import org.junit.jupiter.api.Test; class SubstraitExpressionConverterTest extends PlanTestBase { @@ -169,6 +181,162 @@ void useSubstraitReturnTypeDuringScalarFunctionConversion() { assertEquals(TypeConverter.DEFAULT.toCalcite(typeFactory, R.FP32), calciteExpr.getType()); } + @Test + void observeSuppliedAndInferredScalarFunctionTypes() { + AtomicReference observed = new AtomicReference<>(); + ExpressionRexConverter observingConverter = observingConverter(observed::set); + Expression.ScalarFunctionInvocation expr = add(R.FP32); + + RexNode calciteExpr = expr.accept(observingConverter, Context.newContext()); + + TypeObservation observation = observed.get(); + assertEquals(TypeObservation.Source.SCALAR_FUNCTION, observation.source()); + assertSame(expr, observation.expression()); + assertEquals(R.FP32, observation.suppliedType()); + assertTrue(observation.inferenceFailure().isEmpty()); + assertNotEquals(calciteExpr.getType(), observation.inferredType().orElseThrow()); + assertEquals( + TypeConverter.DEFAULT.toCalcite(typeFactory, R.I32), + observation.inferredType().orElseThrow()); + } + + @Test + void observeMatchingScalarFunctionTypes() { + AtomicReference observed = new AtomicReference<>(); + ExpressionRexConverter observingConverter = observingConverter(observed::set); + Expression.ScalarFunctionInvocation expr = add(R.I32); + + expr.accept(observingConverter, Context.newContext()); + + assertSame(expr, observed.get().expression()); + assertEquals(R.I32, observed.get().suppliedType()); + assertEquals( + TypeConverter.DEFAULT.toCalcite(typeFactory, R.I32), + observed.get().inferredType().orElseThrow()); + } + + @Test + void reportScalarInferenceFailureWithoutFailingConversion() { + AtomicReference observed = new AtomicReference<>(); + ExpressionRexConverter observingConverter = + observingConverter(failingInferenceScalarFunctionConverter(), observed::set); + Expression.ScalarFunctionInvocation expr = add(R.I32); + + RexNode calciteExpr = expr.accept(observingConverter, Context.newContext()); + + assertEquals(TypeConverter.DEFAULT.toCalcite(typeFactory, R.I32), calciteExpr.getType()); + assertSame(expr, observed.get().expression()); + assertEquals(R.I32, observed.get().suppliedType()); + assertTrue(observed.get().inferredType().isEmpty()); + IllegalStateException failure = + assertInstanceOf( + IllegalStateException.class, observed.get().inferenceFailure().orElseThrow()); + assertEquals("controlled inference failure", failure.getMessage()); + } + + @Test + void observeVariadicConcatOnce() { + List observations = new ArrayList<>(); + ExpressionRexConverter observingConverter = observingConverter(observations::add); + Expression.ScalarFunctionInvocation expr = + sb.scalarFn( + DefaultExtensionCatalog.FUNCTIONS_STRING, + "concat:str", + R.I32, + sb.str("a"), + sb.str("b"), + sb.str("c")); + + RexNode calciteExpr = expr.accept(observingConverter, Context.newContext()); + + assertEquals(1, observations.size()); + assertSame(expr, observations.get(0).expression()); + assertEquals(R.I32, observations.get(0).suppliedType()); + assertEquals(TypeConverter.DEFAULT.toCalcite(typeFactory, R.I32), calciteExpr.getType()); + assertEquals( + SqlTypeName.VARCHAR, observations.get(0).inferredType().orElseThrow().getSqlTypeName()); + } + + @Test + void injectTypeObserverThroughConverterProvider() { + AtomicReference observed = new AtomicReference<>(); + ConverterProvider observingProvider = + new ConverterProvider() { + @Override + public TypeObserver getTypeObserver() { + return observed::set; + } + }; + Expression.ScalarFunctionInvocation expr = + sb.scalarFn( + DefaultExtensionCatalog.FUNCTIONS_ARITHMETIC, + "add:i32_i32", + R.I32, + sb.i32(7), + sb.i32(42)); + Project query = sb.project(input -> List.of(expr), sb.emptyVirtualTableScan()); + + new SubstraitToCalcite(observingProvider).convert(query); + + assertEquals(R.I32, observed.get().suppliedType()); + assertEquals( + TypeConverter.DEFAULT.toCalcite(typeFactory, R.I32), + observed.get().inferredType().orElseThrow()); + } + + @Test + void observeEachNestedScalarFunction() { + List observations = new ArrayList<>(); + ExpressionRexConverter observingConverter = observingConverter(observations::add); + Expression.ScalarFunctionInvocation inner = add(R.I32); + Expression.ScalarFunctionInvocation outer = + sb.scalarFn( + DefaultExtensionCatalog.FUNCTIONS_ARITHMETIC, "add:i32_i32", R.I32, inner, sb.i32(3)); + + outer.accept(observingConverter, Context.newContext()); + + assertEquals(2, observations.size()); + assertTrue(observations.stream().anyMatch(observation -> observation.expression() == inner)); + assertTrue(observations.stream().anyMatch(observation -> observation.expression() == outer)); + assertTrue( + observations.stream().allMatch(observation -> observation.inferredType().isPresent())); + } + + @Test + void propagateObserverException() { + ExpressionRexConverter observingConverter = + observingConverter( + observation -> { + throw new IllegalStateException("observer failure"); + }); + Expression.ScalarFunctionInvocation expr = add(R.I32); + + IllegalStateException failure = + assertThrows( + IllegalStateException.class, + () -> expr.accept(observingConverter, Context.newContext())); + + assertEquals("observer failure", failure.getMessage()); + } + + @Test + void propagateObserverExceptionAfterInferenceFailure() { + ExpressionRexConverter observingConverter = + observingConverter( + failingInferenceScalarFunctionConverter(), + observation -> { + throw new IllegalStateException("failure observer failure"); + }); + Expression.ScalarFunctionInvocation expr = add(R.I32); + + IllegalStateException failure = + assertThrows( + IllegalStateException.class, + () -> expr.accept(observingConverter, Context.newContext())); + + assertEquals("failure observer failure", failure.getMessage()); + } + @Test void useSubstraitReturnTypeDuringWindowFunctionConversion() { Expression.WindowFunctionInvocation expr = @@ -193,4 +361,47 @@ void assertTypeMatch(RelDataType actual, Type expected) { Type type = TypeConverter.DEFAULT.toSubstrait(actual); assertEquals(expected, type); } + + private Expression.ScalarFunctionInvocation add(Type outputType) { + return sb.scalarFn( + DefaultExtensionCatalog.FUNCTIONS_ARITHMETIC, + "add:i32_i32", + outputType, + sb.i32(7), + sb.i32(42)); + } + + private ExpressionRexConverter observingConverter(TypeObserver observer) { + return observingConverter( + new ScalarFunctionConverter(extensions.scalarFunctions(), typeFactory), observer); + } + + private ExpressionRexConverter observingConverter( + ScalarFunctionConverter scalarFunctionConverter, TypeObserver observer) { + return new ExpressionRexConverter( + typeFactory, + scalarFunctionConverter, + new WindowFunctionConverter(extensions.windowFunctions(), typeFactory), + TypeConverter.DEFAULT, + observer); + } + + private ScalarFunctionConverter failingInferenceScalarFunctionConverter() { + SqlFunction failingOperator = + new SqlFunction( + "controlled_inference_failure", + SqlKind.OTHER_FUNCTION, + binding -> { + throw new IllegalStateException("controlled inference failure"); + }, + null, + null, + SqlFunctionCategory.USER_DEFINED_FUNCTION); + return new ScalarFunctionConverter(extensions.scalarFunctions(), typeFactory) { + @Override + public Optional getSqlOperatorFromSubstraitFunc(String key, Type outputType) { + return Optional.of(failingOperator); + } + }; + } }