diff --git a/isthmus/src/main/java/io/substrait/isthmus/SubstraitRelNodeConverter.java b/isthmus/src/main/java/io/substrait/isthmus/SubstraitRelNodeConverter.java index 3f5b14519..98b2bfdf8 100644 --- a/isthmus/src/main/java/io/substrait/isthmus/SubstraitRelNodeConverter.java +++ b/isthmus/src/main/java/io/substrait/isthmus/SubstraitRelNodeConverter.java @@ -320,6 +320,7 @@ public RelNode visit(Aggregate aggregate, Context context) throws RuntimeExcepti } RelNode child = aggregate.getInput().accept(this, context); + context.enterScope(AnchoredInput.of(aggregate.getInput().getRelAnchor(), child.getRowType())); List> groupExprLists = aggregate.getGroupings().stream() .map( @@ -383,6 +384,8 @@ public RelNode visit(Aggregate aggregate, Context context) throws RuntimeExcepti } } + exitUncorrelatedScope(context, "aggregate"); + // RelBuilder deduplicates equal aggregate calls, and AggregateCall equality ignores the stored // type: two measures of the same function that differ only by their declared output type would // collapse into one column. Opt out of deduplication for exactly those aggregates — narrowly, @@ -605,10 +608,12 @@ private AggregateCall fromMeasure( @Override public RelNode visit(Sort sort, Context context) throws RuntimeException { RelNode child = sort.getInput().accept(this, context); + context.enterScope(AnchoredInput.of(sort.getInput().getRelAnchor(), child.getRowType())); List sortExpressions = sort.getSortFields().stream() .map(sortField -> directedRexNode(sortField, context)) .collect(Collectors.toList()); + exitUncorrelatedScope(context, "sort"); RelNode node = relBuilder.push(child).sort(sortExpressions).build(); return applyRemap(node, sort.getRemap()); } @@ -641,6 +646,7 @@ private RexNode directedRexNode(Expression.SortField sortField, Context context) @Override public RelNode visit(Fetch fetch, Context context) throws RuntimeException { RelNode child = fetch.getInput().accept(this, context); + context.enterScope(AnchoredInput.of(fetch.getInput().getRelAnchor(), child.getRowType())); // Offset/count are expressions; pass them through to Calcite as RexNodes so non-literal (e.g. // dynamic-parameter) offset/count are preserved. An unset offset means 0 and an unset count // means LIMIT ALL. @@ -648,6 +654,7 @@ public RelNode visit(Fetch fetch, Context context) throws RuntimeException { fetch.getOffset().map(e -> e.accept(expressionRexConverter, context)).orElse(null); RexNode count = fetch.getCount().map(e -> e.accept(expressionRexConverter, context)).orElse(null); + exitUncorrelatedScope(context, "fetch"); RelNode node = relBuilder.push(child).sortLimit(offset, count, ImmutableList.of()).build(); return applyRemap(node, fetch.getRemap()); } @@ -688,9 +695,8 @@ private RelFieldCollation toRelFieldCollation(Expression.SortField sortField, Co @Override public RelNode visit(NamedUpdate update, Context context) { relBuilder.scan(update.getNames()); + context.enterScope(AnchoredInput.of(Optional.empty(), relBuilder.peek().getRowType())); RexNode condition = update.getCondition().accept(expressionRexConverter, context); - relBuilder.filter(condition); - RelNode inputForModify = relBuilder.build(); NamedStruct tableSchema = update.getTableSchema(); List fieldNames = tableSchema.names(); @@ -705,6 +711,10 @@ public RelNode visit(NamedUpdate update, Context context) { transform.getTransformation().accept(expressionRexConverter, context)); } + java.util.Set correlationIds = context.exitScope(); + relBuilder.filter(correlationIds, condition); + RelNode inputForModify = relBuilder.build(); + assert relBuilder.getRelOptSchema() != null; final RelOptTable table = relBuilder.getRelOptSchema().getTableForMember(update.getNames()); @@ -929,6 +939,23 @@ protected RelNode applyRemap(RelNode relNode, Optional remap) { return relNode; } + /** + * Exits a scope entered only to resolve input field types, on an operator that cannot carry + * correlation variables. Calcite's Aggregate, Sort and Fetch have no {@code variablesSet}, so a + * correlation resolved against such an operator's input has nowhere to be declared and is + * rejected rather than silently dropped. + * + * @param context the conversion context whose innermost scope is exited + * @param relName the relation kind, used in the failure message + */ + private static void exitUncorrelatedScope(Context context, String relName) { + java.util.Set correlationIds = context.exitScope(); + if (!correlationIds.isEmpty()) { + throw new UnsupportedOperationException( + "Outer references bound to the " + relName + " input are not supported"); + } + } + private RelNode applyRemap(RelNode relNode, Rel.Remap remap) { RelDataType rowType = relNode.getRowType(); List fieldNames = rowType.getFieldNames(); @@ -987,6 +1014,9 @@ public static class Context implements VisitationContext { /** Maps a {@code rel_anchor} to the single {@link CorrelationId} minted for it. */ private final Map correlationIdByAnchor = new HashMap<>(); + /** Lambda parameter types by nesting level, innermost on top. */ + private final Deque> lambdaParameterTypes = new ArrayDeque<>(); + /** * Every {@code rel_anchor} that has entered a scope. Resolution keys {@link #scopeByAnchor} and * {@link #correlationIdByAnchor} purely by anchor value, which is only sound if anchors are @@ -998,6 +1028,7 @@ public static class Context implements VisitationContext { /** One correlation scope per enclosing relational operator. */ private static final class Scope { + final List inputRowTypes = new ArrayList<>(); final Map rowTypeByAnchor = new HashMap<>(); final java.util.Set correlationIds = new HashSet<>(); } @@ -1013,13 +1044,14 @@ public static Context newContext() { /** * Enters a correlation scope for a relational operator, recording the {@code rel_anchor} (if - * any) carried by each of its inputs. + * any) and row type carried by each of its inputs. * * @param inputs the operator's inputs paired with their anchors */ public void enterScope(final AnchoredInput... inputs) { final Scope scope = new Scope(); for (final AnchoredInput input : inputs) { + scope.inputRowTypes.add(input.rowType); if (input.anchor.isPresent()) { final int anchor = input.anchor.get(); if (!seenAnchors.add(anchor)) { @@ -1036,6 +1068,63 @@ public void enterScope(final AnchoredInput... inputs) { scopes.push(scope); } + /** + * Returns the type of a field in the current operator's flattened input row. + * + * @param fieldIndex zero-based field index across all current inputs + * @return the Calcite field type from the input relation + * @throws IllegalStateException if expression conversion has no current relational input + * @throws IndexOutOfBoundsException if the field index is outside the flattened input row + */ + public RelDataType getInputFieldType(final int fieldIndex) { + if (scopes.isEmpty()) { + throw new IllegalStateException("No input row type is available for field reference"); + } + + int remainingIndex = fieldIndex; + int fieldCount = 0; + for (final RelDataType inputRowType : scopes.peek().inputRowTypes) { + final int inputFieldCount = inputRowType.getFieldCount(); + fieldCount += inputFieldCount; + if (remainingIndex < inputFieldCount) { + return inputRowType.getFieldList().get(remainingIndex).getType(); + } + remainingIndex -= inputFieldCount; + } + throw new IndexOutOfBoundsException( + "Field index " + fieldIndex + " is outside input row with " + fieldCount + " fields"); + } + + /** + * Enters a lambda scope with its Calcite parameter types. + * + * @param parameterTypes parameter types in declaration order + */ + public void enterLambdaScope(final List parameterTypes) { + lambdaParameterTypes.push(List.copyOf(parameterTypes)); + } + + /** Exits the innermost lambda scope. */ + public void exitLambdaScope() { + lambdaParameterTypes.pop(); + } + + /** + * Returns a parameter type from the innermost lambda scope. + * + * @param parameterIndex zero-based parameter index + * @return the Calcite parameter type + * @throws IllegalStateException if expression conversion has no current lambda + * @throws IndexOutOfBoundsException if the index is outside the lambda's parameter list + */ + public RelDataType getLambdaParameterType(final int parameterIndex) { + if (lambdaParameterTypes.isEmpty()) { + throw new IllegalStateException( + "No lambda parameter type is available for field reference"); + } + return lambdaParameterTypes.peek().get(parameterIndex); + } + /** * Exits the innermost correlation scope, returning the correlation ids to attach to the * operator that owns it. 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 f64ff9184..089425e34 100644 --- a/isthmus/src/main/java/io/substrait/isthmus/expression/ExpressionRexConverter.java +++ b/isthmus/src/main/java/io/substrait/isthmus/expression/ExpressionRexConverter.java @@ -498,17 +498,22 @@ public RexNode visit(Expression.IfThen expr, Context context) throws RuntimeExce @Override public RexNode visit(Expression.Lambda expr, Context context) throws RuntimeException { + List parameterTypes = + expr.parameters().fields().stream() + .map(parameter -> typeConverter.toCalcite(typeFactory, parameter)) + .collect(Collectors.toList()); List parameters = - IntStream.range(0, expr.parameters().fields().size()) - .mapToObj( - i -> - new RexLambdaRef( - i, - "p" + i, - typeConverter.toCalcite(typeFactory, expr.parameters().fields().get(i)))) + IntStream.range(0, parameterTypes.size()) + .mapToObj(i -> new RexLambdaRef(i, "p" + i, parameterTypes.get(i))) .collect(Collectors.toList()); - RexNode body = expr.body().accept(this, context); + context.enterLambdaScope(parameterTypes); + RexNode body; + try { + body = expr.body().accept(this, context); + } finally { + context.exitLambdaScope(); + } return rexBuilder.makeLambdaCall(body, parameters); } @@ -816,6 +821,10 @@ public RexNode visit(FieldReference expr, Context context) throws RuntimeExcepti final FieldReference.StructField field = (FieldReference.StructField) segment; rexInputRef = new RexInputRef(field.offset(), typeConverter.toCalcite(typeFactory, expr.getType())); + observeType( + expr, + TypeObservation.Source.FIELD_REFERENCE, + () -> context.getInputFieldType(field.offset())); } else { throw new IllegalArgumentException("Unhandled type: " + segment); } @@ -834,8 +843,11 @@ public RexNode visit(FieldReference expr, Context context) throws RuntimeExcepti final CorrelationId correlationId = context.correlationIdForAnchor( anchor, () -> relNodeConverter.getRelBuilder().getCluster().createCorrel()); - return rexBuilder.makeFieldAccess( - rexBuilder.makeCorrel(rowType, correlationId), field.offset()); + RexNode fieldAccess = + rexBuilder.makeFieldAccess( + rexBuilder.makeCorrel(rowType, correlationId), field.offset()); + observeType(expr, TypeObservation.Source.FIELD_REFERENCE, fieldAccess::getType); + return fieldAccess; } else { throw new IllegalArgumentException("Unhandled type: " + segment); } @@ -852,7 +864,13 @@ public RexNode visit(FieldReference expr, Context context) throws RuntimeExcepti if (segment instanceof FieldReference.StructField) { final FieldReference.StructField field = (FieldReference.StructField) segment; RelDataType calciteType = typeConverter.toCalcite(typeFactory, expr.getType()); - return new RexLambdaRef(field.offset(), "p" + field.offset(), calciteType); + RexLambdaRef lambdaRef = + new RexLambdaRef(field.offset(), "p" + field.offset(), calciteType); + observeType( + expr, + TypeObservation.Source.FIELD_REFERENCE, + () -> context.getLambdaParameterType(field.offset())); + return lambdaRef; } else { throw new IllegalArgumentException("Unhandled type: " + segment); } diff --git a/isthmus/src/main/java/io/substrait/isthmus/expression/TypeObservation.java b/isthmus/src/main/java/io/substrait/isthmus/expression/TypeObservation.java index 764c99032..cc875855a 100644 --- a/isthmus/src/main/java/io/substrait/isthmus/expression/TypeObservation.java +++ b/isthmus/src/main/java/io/substrait/isthmus/expression/TypeObservation.java @@ -18,7 +18,10 @@ public enum Source { SCALAR_FUNCTION, /** A window function invocation. */ - WINDOW_FUNCTION + WINDOW_FUNCTION, + + /** A field reference. */ + FIELD_REFERENCE } private final Source source; diff --git a/isthmus/src/test/java/io/substrait/isthmus/SubstraitExpressionConverterTest.java b/isthmus/src/test/java/io/substrait/isthmus/SubstraitExpressionConverterTest.java index 0a6a36d45..84ae622e6 100644 --- a/isthmus/src/test/java/io/substrait/isthmus/SubstraitExpressionConverterTest.java +++ b/isthmus/src/test/java/io/substrait/isthmus/SubstraitExpressionConverterTest.java @@ -10,8 +10,11 @@ import io.substrait.expression.Expression; import io.substrait.expression.Expression.Switch; +import io.substrait.expression.FieldReference; +import io.substrait.expression.LambdaBuilder; import io.substrait.expression.WindowBound; import io.substrait.extension.DefaultExtensionCatalog; +import io.substrait.isthmus.SubstraitRelNodeConverter.AnchoredInput; import io.substrait.isthmus.SubstraitRelNodeConverter.Context; import io.substrait.isthmus.expression.ExpressionRexConverter; import io.substrait.isthmus.expression.ScalarFunctionConverter; @@ -168,6 +171,142 @@ Rel createSubQueryRel() { sb.filter(input -> sb.equal(sb.fieldReference(input, 2), sb.str("EUROPE")), commonTable)); } + @Test + void observeSuppliedAndInputRootFieldReferenceTypes() { + AtomicReference observed = new AtomicReference<>(); + ExpressionRexConverter observingConverter = observingConverter(observed::set); + Context context = Context.newContext(); + RelDataType rowType = + typeFactory.builder().add("actual", SqlTypeName.INTEGER).nullable(false).build(); + context.enterScope(AnchoredInput.of(Optional.empty(), rowType)); + FieldReference expr = FieldReference.newRootStructReference(0, R.FP32); + + RexNode calciteExpr = expr.accept(observingConverter, context); + + TypeObservation observation = observed.get(); + assertEquals(TypeObservation.Source.FIELD_REFERENCE, observation.source()); + assertSame(expr, observation.expression()); + assertEquals(R.FP32, observation.suppliedType()); + assertTrue(observation.inferenceFailure().isEmpty()); + assertEquals( + TypeConverter.DEFAULT.toCalcite(typeFactory, R.I32), + observation.inferredType().orElseThrow()); + assertEquals(TypeConverter.DEFAULT.toCalcite(typeFactory, R.FP32), calciteExpr.getType()); + context.exitScope(); + } + + @Test + void observeMatchingRootFieldReferenceType() { + AtomicReference observed = new AtomicReference<>(); + ConverterProvider observingProvider = + ConverterProvider.builder().typeObserver(observed::set).build(); + FieldReference expr = FieldReference.newRootStructReference(0, R.I32); + Project query = sb.project(input -> List.of(expr), commonTable); + + new SubstraitToCalcite(observingProvider).convert(query); + + assertSame(expr, observed.get().expression()); + assertEquals(R.I32, observed.get().suppliedType()); + assertEquals( + TypeConverter.DEFAULT.toCalcite(observingProvider.getTypeFactory(), R.I32), + observed.get().inferredType().orElseThrow()); + } + + @Test + void observeAggregateFieldReferenceTypesFromInput() { + List observations = new ArrayList<>(); + ConverterProvider observingProvider = + ConverterProvider.builder().typeObserver(observations::add).build(); + Rel query = + sb.aggregate( + input -> sb.grouping(input, 0), input -> List.of(sb.count(input, 0)), commonTable); + + new SubstraitToCalcite(observingProvider).convert(query); + + assertEquals(2, observations.size()); + assertTrue( + observations.stream() + .allMatch( + observation -> + observation.source() == TypeObservation.Source.FIELD_REFERENCE + && observation.inferredType().isPresent())); + } + + @Test + void observeSortFieldReferenceTypeFromInput() { + List observations = new ArrayList<>(); + ConverterProvider observingProvider = + ConverterProvider.builder().typeObserver(observations::add).build(); + Rel query = sb.sort(input -> sb.sortFields(input, 0), commonTable); + + new SubstraitToCalcite(observingProvider).convert(query); + + assertEquals(1, observations.size()); + assertEquals(TypeObservation.Source.FIELD_REFERENCE, observations.get(0).source()); + assertTrue(observations.get(0).inferredType().isPresent()); + } + + @Test + void reportMissingInputTypeWithoutFailingFieldReferenceConversion() { + AtomicReference observed = new AtomicReference<>(); + ExpressionRexConverter observingConverter = observingConverter(observed::set); + FieldReference expr = FieldReference.newRootStructReference(0, R.FP32); + + RexNode calciteExpr = expr.accept(observingConverter, Context.newContext()); + + assertEquals(TypeConverter.DEFAULT.toCalcite(typeFactory, R.FP32), calciteExpr.getType()); + assertSame(expr, observed.get().expression()); + assertTrue(observed.get().inferredType().isEmpty()); + IllegalStateException failure = + assertInstanceOf( + IllegalStateException.class, observed.get().inferenceFailure().orElseThrow()); + assertEquals("No input row type is available for field reference", failure.getMessage()); + } + + @Test + void observeOuterFieldReferenceType() { + AtomicReference observed = new AtomicReference<>(); + ConverterProvider observingProvider = + ConverterProvider.builder().typeObserver(observed::set).build(); + SubstraitRelNodeConverter observingRelNodeConverter = + new SubstraitRelNodeConverter(builder, observingProvider); + ExpressionRexConverter observingConverter = + observingProvider.getExpressionRexConverter(observingRelNodeConverter); + Context context = Context.newContext(); + RelDataType rowType = + typeFactory.builder().add("actual", SqlTypeName.INTEGER).nullable(false).build(); + context.enterScope(AnchoredInput.of(Optional.of(17), rowType)); + FieldReference expr = FieldReference.newRootStructOuterReferenceByRelReference(0, R.FP32, 17); + + RexNode calciteExpr = expr.accept(observingConverter, context); + + assertEquals(TypeObservation.Source.FIELD_REFERENCE, observed.get().source()); + assertSame(expr, observed.get().expression()); + assertEquals( + rowType.getFieldList().get(0).getType(), observed.get().inferredType().orElseThrow()); + assertEquals(rowType.getFieldList().get(0).getType(), calciteExpr.getType()); + context.exitScope(); + } + + @Test + void observeLambdaParameterReferenceType() { + AtomicReference observed = new AtomicReference<>(); + ExpressionRexConverter observingConverter = observingConverter(observed::set); + LambdaBuilder lambdaBuilder = new LambdaBuilder(); + Expression.Lambda lambda = + lambdaBuilder.lambda( + List.of(R.I32), + parameters -> FieldReference.builder().from(parameters.ref(0)).type(R.FP32).build()); + + lambda.accept(observingConverter, Context.newContext()); + + assertEquals(TypeObservation.Source.FIELD_REFERENCE, observed.get().source()); + assertEquals(R.FP32, observed.get().suppliedType()); + assertEquals( + TypeConverter.DEFAULT.toCalcite(typeFactory, R.I32), + observed.get().inferredType().orElseThrow()); + } + @Test void useSubstraitReturnTypeDuringScalarFunctionConversion() { Expression.ScalarFunctionInvocation expr = diff --git a/isthmus/src/test/java/io/substrait/isthmus/expression/SubqueryConversionTest.java b/isthmus/src/test/java/io/substrait/isthmus/expression/SubqueryConversionTest.java index 88475a4df..470029fe8 100644 --- a/isthmus/src/test/java/io/substrait/isthmus/expression/SubqueryConversionTest.java +++ b/isthmus/src/test/java/io/substrait/isthmus/expression/SubqueryConversionTest.java @@ -3,6 +3,7 @@ import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertThrows; +import io.substrait.expression.Expression; import io.substrait.expression.FieldReference; import io.substrait.isthmus.PlanTestBase; import io.substrait.isthmus.sql.SubstraitSqlDialect; @@ -161,6 +162,40 @@ void duplicateRelAnchorIsRejected() { assertThrows(UnsupportedOperationException.class, () -> substraitToCalcite.convert(root)); } + @Test + void correlationBoundToSortInputIsRejected() { + /* + * A sort key holding a correlated scalar subquery that binds to the sort's own input. Calcite's + * Sort carries no variablesSet, so there is nowhere to declare the correlation: conversion must + * reject the plan rather than emit a correlation variable no operator owns. + */ + final Rel root = + sb.sort( + input -> + List.of( + sb.sortField( + sb.scalarSubquery( + sb.project( + input2 -> List.of(sb.fieldReference(input2, 1)), + Remap.of(List.of(1)), + sb.filter( + input2 -> + sb.equal( + sb.fieldReference(input2, 0), + FieldReference + .newRootStructOuterReferenceByRelReference( + 1, TypeCreator.REQUIRED.I64, 1)), + customerTableScan)), + TypeCreator.NULLABLE.I64), + Expression.SortDirection.ASC_NULLS_LAST)), + orderTableScan.withRelAnchor(1)); + + final UnsupportedOperationException failure = + assertThrows(UnsupportedOperationException.class, () -> substraitToCalcite.convert(root)); + assertEquals( + "Outer references bound to the sort input are not supported", failure.getMessage()); + } + @Test void testOuterFieldReferenceTwoSteps() { /*