fix(isthmus): handle subqueries with outer field references

fixes #382

Signed-off-by: Niels Pardon <[email protected]>

Author: nielspardon

core/src/main/java/io/substrait/expression/FieldReference.java

@@ -36,7 +36,13 @@ public <R, C extends VisitationContext, E extends Throwable> R accept(
   }
 
   public boolean isSimpleRootReference() {
-    return segments().size() == 1 && !inputExpression().isPresent();
+    return segments().size() == 1
+        && !inputExpression().isPresent()
+        && !outerReferenceStepsOut().isPresent();
+  }
+
+  public boolean isOuterReference() {
+    return outerReferenceStepsOut().orElse(0) > 0;
   }
 
   public FieldReference dereferenceStruct(int index) {

isthmus/src/main/java/io/substrait/isthmus/SubstraitRelNodeConverter.java

@@ -23,12 +23,14 @@
 import io.substrait.relation.Rel;
 import io.substrait.relation.Set;
 import io.substrait.relation.Sort;
-import io.substrait.util.EmptyVisitationContext;
+import io.substrait.util.VisitationContext;
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.Collections;
+import java.util.HashSet;
 import java.util.List;
 import java.util.Optional;
+import java.util.Stack;
 import java.util.stream.Collectors;
 import java.util.stream.IntStream;
 import java.util.stream.Stream;
@@ -40,6 +42,7 @@
 import org.apache.calcite.rel.RelFieldCollation;
 import org.apache.calcite.rel.RelNode;
 import org.apache.calcite.rel.core.AggregateCall;
+import org.apache.calcite.rel.core.CorrelationId;
 import org.apache.calcite.rel.core.JoinRelType;
 import org.apache.calcite.rel.logical.LogicalValues;
 import org.apache.calcite.rel.type.RelDataType;
@@ -59,7 +62,7 @@
  * visitFallback and throw UnsupportedOperationException.
  */
 public class SubstraitRelNodeConverter
-    extends AbstractRelVisitor<RelNode, EmptyVisitationContext, RuntimeException> {
+    extends AbstractRelVisitor<RelNode, SubstraitRelNodeConverter.Context, RuntimeException> {
 
   protected final RelDataTypeFactory typeFactory;
 
@@ -137,42 +140,44 @@ public static RelNode convert(
     return relRoot.accept(
         new SubstraitRelNodeConverter(
             EXTENSION_COLLECTION, relOptCluster.getTypeFactory(), relBuilder),
-        null);
+        Context.newContext());
   }
 
   @Override
-  public RelNode visit(Filter filter, EmptyVisitationContext context) throws RuntimeException {
+  public RelNode visit(Filter filter, Context context) throws RuntimeException {
     RelNode input = filter.getInput().accept(this, context);
+    context.pushParentRelNodes(input);
     RexNode filterCondition = filter.getCondition().accept(expressionRexConverter, context);
-    RelNode node = relBuilder.push(input).filter(filterCondition).build();
+    RelNode node =
+        relBuilder.push(input).filter(context.popCorrelationIds(), filterCondition).build();
+    context.popParentRelNodes();
     return applyRemap(node, filter.getRemap());
   }
 
   @Override
-  public RelNode visit(NamedScan namedScan, EmptyVisitationContext context)
-      throws RuntimeException {
+  public RelNode visit(NamedScan namedScan, Context context) throws RuntimeException {
     RelNode node = relBuilder.scan(namedScan.getNames()).build();
     return applyRemap(node, namedScan.getRemap());
   }
 
   @Override
-  public RelNode visit(LocalFiles localFiles, EmptyVisitationContext context)
-      throws RuntimeException {
+  public RelNode visit(LocalFiles localFiles, Context context) throws RuntimeException {
     return visitFallback(localFiles, context);
   }
 
   @Override
-  public RelNode visit(EmptyScan emptyScan, EmptyVisitationContext context)
-      throws RuntimeException {
+  public RelNode visit(EmptyScan emptyScan, Context context) throws RuntimeException {
     RelDataType rowType =
         typeConverter.toCalcite(relBuilder.getTypeFactory(), emptyScan.getInitialSchema().struct());
     RelNode node = LogicalValues.create(relBuilder.getCluster(), rowType, ImmutableList.of());
     return applyRemap(node, emptyScan.getRemap());
   }
 
   @Override
-  public RelNode visit(Project project, EmptyVisitationContext context) throws RuntimeException {
+  public RelNode visit(Project project, Context context) throws RuntimeException {
     RelNode child = project.getInput().accept(this, context);
+    context.pushParentRelNodes(child);
+
     Stream<RexNode> directOutputs =
         IntStream.range(0, child.getRowType().getFieldCount())
             .mapToObj(fieldIndex -> rexBuilder.makeInputRef(child, fieldIndex));
@@ -183,12 +188,17 @@ public RelNode visit(Project project, EmptyVisitationContext context) throws Run
     List<RexNode> rexExprs =
         Stream.concat(directOutputs, exprs).collect(java.util.stream.Collectors.toList());
 
-    RelNode node = relBuilder.push(child).project(rexExprs).build();
+    RelNode node =
+        relBuilder
+            .push(child)
+            .project(rexExprs, List.of(), false, context.popCorrelationIds())
+            .build();
+    context.popParentRelNodes();
     return applyRemap(node, project.getRemap());
   }
 
   @Override
-  public RelNode visit(Cross cross, EmptyVisitationContext context) throws RuntimeException {
+  public RelNode visit(Cross cross, Context context) throws RuntimeException {
     RelNode left = cross.getLeft().accept(this, context);
     RelNode right = cross.getRight().accept(this, context);
     // Calcite represents CROSS JOIN as the equivalent INNER JOIN with true condition
@@ -198,14 +208,15 @@ public RelNode visit(Cross cross, EmptyVisitationContext context) throws Runtime
   }
 
   @Override
-  public RelNode visit(Join join, EmptyVisitationContext context) throws RuntimeException {
+  public RelNode visit(Join join, Context context) throws RuntimeException {
     RelNode left = join.getLeft().accept(this, context);
     RelNode right = join.getRight().accept(this, context);
+    context.pushParentRelNodes(left, right);
     RexNode condition =
         join.getCondition()
             .map(c -> c.accept(expressionRexConverter, context))
             .orElse(relBuilder.literal(true));
-    var joinType =
+    JoinRelType joinType =
         switch (join.getJoinType()) {
           case INNER -> JoinRelType.INNER;
           case LEFT -> JoinRelType.LEFT;
@@ -220,12 +231,18 @@ public RelNode visit(Join join, EmptyVisitationContext context) throws RuntimeEx
           default -> throw new UnsupportedOperationException(
               "Unsupported join type: " + join.getJoinType().name());
         };
-    RelNode node = relBuilder.push(left).push(right).join(joinType, condition).build();
+    RelNode node =
+        relBuilder
+            .push(left)
+            .push(right)
+            .join(joinType, condition, context.popCorrelationIds())
+            .build();
+    context.popParentRelNodes();
     return applyRemap(node, join.getRemap());
   }
 
   @Override
-  public RelNode visit(Set set, EmptyVisitationContext context) throws RuntimeException {
+  public RelNode visit(Set set, Context context) throws RuntimeException {
     int numInputs = set.getInputs().size();
     set.getInputs()
         .forEach(
@@ -253,8 +270,7 @@ public RelNode visit(Set set, EmptyVisitationContext context) throws RuntimeExce
   }
 
   @Override
-  public RelNode visit(Aggregate aggregate, EmptyVisitationContext context)
-      throws RuntimeException {
+  public RelNode visit(Aggregate aggregate, Context context) throws RuntimeException {
     if (!PreCalciteAggregateValidator.isValidCalciteAggregate(aggregate)) {
       aggregate =
           PreCalciteAggregateValidator.PreCalciteAggregateTransformer
@@ -282,7 +298,7 @@ public RelNode visit(Aggregate aggregate, EmptyVisitationContext context)
     return applyRemap(node, aggregate.getRemap());
   }
 
-  private AggregateCall fromMeasure(Aggregate.Measure measure, EmptyVisitationContext context) {
+  private AggregateCall fromMeasure(Aggregate.Measure measure, Context context) {
     var eArgs = measure.getFunction().arguments();
     var arguments =
         IntStream.range(0, measure.getFunction().arguments().size())
@@ -357,7 +373,7 @@ private AggregateCall fromMeasure(Aggregate.Measure measure, EmptyVisitationCont
   }
 
   @Override
-  public RelNode visit(Sort sort, EmptyVisitationContext context) throws RuntimeException {
+  public RelNode visit(Sort sort, Context context) throws RuntimeException {
     RelNode child = sort.getInput().accept(this, context);
     List<RexNode> sortExpressions =
         sort.getSortFields().stream()
@@ -367,7 +383,7 @@ public RelNode visit(Sort sort, EmptyVisitationContext context) throws RuntimeEx
     return applyRemap(node, sort.getRemap());
   }
 
-  private RexNode directedRexNode(Expression.SortField sortField, EmptyVisitationContext context) {
+  private RexNode directedRexNode(Expression.SortField sortField, Context context) {
     var expression = sortField.expr();
     var rexNode = expression.accept(expressionRexConverter, context);
     var sortDirection = sortField.direction();
@@ -382,7 +398,7 @@ private RexNode directedRexNode(Expression.SortField sortField, EmptyVisitationC
   }
 
   @Override
-  public RelNode visit(Fetch fetch, EmptyVisitationContext context) throws RuntimeException {
+  public RelNode visit(Fetch fetch, Context context) throws RuntimeException {
     RelNode child = fetch.getInput().accept(this, context);
     var optCount = fetch.getCount();
     long count = optCount.orElse(-1L);
@@ -397,8 +413,7 @@ public RelNode visit(Fetch fetch, EmptyVisitationContext context) throws Runtime
     return applyRemap(node, fetch.getRemap());
   }
 
-  private RelFieldCollation toRelFieldCollation(
-      Expression.SortField sortField, EmptyVisitationContext context) {
+  private RelFieldCollation toRelFieldCollation(Expression.SortField sortField, Context context) {
     var expression = sortField.expr();
     var rex = expression.accept(expressionRexConverter, context);
     var sortDirection = sortField.direction();
@@ -426,7 +441,7 @@ private RelFieldCollation toRelFieldCollation(
   }
 
   @Override
-  public RelNode visitFallback(Rel rel, EmptyVisitationContext context) throws RuntimeException {
+  public RelNode visitFallback(Rel rel, Context context) throws RuntimeException {
     throw new UnsupportedOperationException(
         String.format(
             "Rel %s of type %s not handled by visitor type %s.",
@@ -454,12 +469,49 @@ private RelNode applyRemap(RelNode relNode, Rel.Remap remap) {
     return relBuilder.push(relNode).project(rexList).build();
   }
 
-  private void checkRexInputRefOnly(RexNode rexNode, String context, String aggName) {
-    if (!(rexNode instanceof RexInputRef)) {
-      throw new UnsupportedOperationException(
-          String.format(
-              "Compound expression %s in %s of agg function %s is not implemented yet.",
-              rexNode, context, aggName));
+  public static class Context implements VisitationContext {
+    protected final Stack<RelNode[]> parentRelations = new Stack<>();
+
+    protected final Stack<java.util.Set<CorrelationId>> correlationIds = new Stack<>();
+
+    private int subqueryDepth;
+
+    public static Context newContext() {
+      return new Context();
+    }
+
+    public void pushParentRelNodes(final RelNode... inputs) {
+      parentRelations.push(inputs);
+      this.correlationIds.push(new HashSet<>());
+    }
+
+    public void popParentRelNodes() {
+      parentRelations.pop();
     }
+
+    public RelNode[] getParentRelation(final Integer stepsOut) {
+      return this.parentRelations.get(stepsOut - subqueryDepth);
+    }
+
+    public java.util.Set<CorrelationId> popCorrelationIds() {
+      return correlationIds.pop();
+    }
+
+    public void addCorrelationId(final int stepsOut, final CorrelationId correlationId) {
+      final int index = stepsOut - subqueryDepth;
+      this.correlationIds.get(index).add(correlationId);
+    }
+
+    public void incrementSubqueryDepth() {
+      this.subqueryDepth++;
+    }
+
+    public void decrementSubqueryDepth() {
+      this.subqueryDepth--;
+    }
+  }
+
+  public RelBuilder getRelBuilder() {
+    return relBuilder;
   }
 }

isthmus/src/main/java/io/substrait/isthmus/SubstraitToCalcite.java

@@ -1,6 +1,7 @@
 package io.substrait.isthmus;
 
 import io.substrait.extension.SimpleExtension;
+import io.substrait.isthmus.SubstraitRelNodeConverter.Context;
 import io.substrait.plan.Plan;
 import io.substrait.relation.NamedScan;
 import io.substrait.relation.Rel;
@@ -91,7 +92,7 @@ public RelNode convert(Rel rel) {
     CalciteSchema rootSchema = toSchema(rel);
     RelBuilder relBuilder = createRelBuilder(rootSchema);
     SubstraitRelNodeConverter converter = createSubstraitRelNodeConverter(relBuilder);
-    return rel.accept(converter, null);
+    return rel.accept(converter, Context.newContext());
   }
 
   /**

isthmus/src/main/java/io/substrait/isthmus/expression/CallConverters.java

@@ -5,7 +5,6 @@
 import io.substrait.expression.ExpressionCreator;
 import io.substrait.isthmus.*;
 import io.substrait.type.Type;
-import io.substrait.util.EmptyVisitationContext;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Optional;
@@ -51,7 +50,7 @@ public class CallConverters {
    * org.apache.calcite.rex.RexLiteral} and then re-interpreted to have the correct type.
    *
    * <p>See {@link ExpressionRexConverter#visit(Expression.UserDefinedLiteral,
-   * EmptyVisitationContext)} for this conversion.
+   * SubstraitRelNodeConverter.Context)} for this conversion.
    *
    * <p>When converting from Calcite to Substrait, this call converter extracts the {@link
    * Expression.UserDefinedLiteral} that was stored.