From fc0c7968561fd43303d79289ecabaeb4fb9e8b53 Mon Sep 17 00:00:00 2001 From: Aleksandr Efimov Date: Tue, 4 Aug 2026 06:54:21 +0300 Subject: [PATCH] feat(isthmus): observe field reference types Report Calcite input-schema types for root, outer and lambda field references through TypeObserver while keeping RexNode construction unchanged. The conversion context now tracks input row types and lambda parameter types; a field reference converted with no input in scope is reported as an inference failure rather than failing the conversion. Aggregate, sort and fetch open a conversion scope so their expressions can resolve input field types. Calcite's Aggregate and Sort carry no variablesSet, so a correlation bound to one of those inputs is rejected instead of being silently dropped. Named update now applies its filter with the correlation ids resolved while converting the condition and the transformations. --- .../isthmus/SubstraitRelNodeConverter.java | 95 +++++++++++- .../expression/ExpressionRexConverter.java | 40 +++-- .../isthmus/expression/TypeObservation.java | 5 +- .../SubstraitExpressionConverterTest.java | 139 ++++++++++++++++++ .../expression/SubqueryConversionTest.java | 35 +++++ 5 files changed, 299 insertions(+), 15 deletions(-) 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() { /*