[WIP] Add support for incremental rewrite of nested queries

coral-incremental/src/main/java/com/linkedin/coral/incremental/IncrementalTransformerResults.java

@@ -0,0 +1,76 @@
+/**
+ * Copyright 2023 LinkedIn Corporation. All rights reserved.
+ * Licensed under the BSD-2 Clause license.
+ * See LICENSE in the project root for license information.
+ */
+package com.linkedin.coral.incremental;
+
+import java.util.HashMap;
+import java.util.Map;
+
+import org.apache.calcite.rel.RelNode;
+
+
+public class IncrementalTransformerResults {
+
+  private RelNode incrementalRelNode;
+  private RelNode refreshRelNode;
+  private Map<String, RelNode> intermediateQueryRelNodes;
+
+  public IncrementalTransformerResults() {
+    incrementalRelNode = null;
+    refreshRelNode = null;
+    intermediateQueryRelNodes = new HashMap<>();
+  }
+
+  public boolean existsIncrementalRelNode() {
+    return incrementalRelNode != null;
+  }
+
+  public RelNode getIncrementalRelNode() {
+    return incrementalRelNode;
+  }
+
+  public boolean existsRefreshRelNode() {
+    return refreshRelNode != null;
+  }
+
+  public RelNode getRefreshRelNode() {
+    return refreshRelNode;
+  }
+
+  public Map<String, RelNode> getIntermediateQueryRelNodes() {
+    return intermediateQueryRelNodes;
+  }
+
+  public boolean containsIntermediateQueryRelNodeKey(String name) {
+    return intermediateQueryRelNodes.containsKey(name);
+  }
+
+  public RelNode getIntermediateQueryRelNodeCorrespondingToKey(String name) {
+    return intermediateQueryRelNodes.get(name);
+  }
+
+  public void setIncrementalRelNode(RelNode incrementalRelNode) {
+    this.incrementalRelNode = incrementalRelNode;
+  }
+
+  public void setRefreshRelNode(RelNode refreshRelNode) {
+    this.refreshRelNode = refreshRelNode;
+  }
+
+  public void setIntermediateQueryRelNodes(Map<String, RelNode> intermediateQueryRelNodes) {
+    this.intermediateQueryRelNodes = intermediateQueryRelNodes;
+  }
+
+  public void addIntermediateQueryRelNode(String name, RelNode intermediateRelNode) {
+    this.intermediateQueryRelNodes.put(name, intermediateRelNode);
+  }
+
+  public void addMultipleIntermediateQueryRelNodes(Map<String, RelNode> intermediateQueryRelNodes) {
+    if (intermediateQueryRelNodes != null) {
+      this.intermediateQueryRelNodes.putAll(intermediateQueryRelNodes);
+    }
+  }
+
+}

coral-incremental/src/main/java/com/linkedin/coral/incremental/RelNodeIncrementalTransformer.java

@@ -7,10 +7,12 @@
 
 import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.Collections;
 import java.util.List;
 import java.util.stream.Collectors;
 import java.util.stream.IntStream;
 
+import org.apache.calcite.plan.RelOptSchema;
 import org.apache.calcite.plan.RelOptTable;
 import org.apache.calcite.prepare.RelOptTableImpl;
 import org.apache.calcite.rel.RelNode;
@@ -29,14 +31,29 @@
 
 public class RelNodeIncrementalTransformer {
 
+  private static RelOptSchema relOptSchema;
+
   private RelNodeIncrementalTransformer() {
   }
 
-  public static RelNode convertRelIncremental(RelNode originalNode) {
+  public static IncrementalTransformerResults performIncrementalTransformation(RelNode originalNode) {
+    IncrementalTransformerResults incrementalTransformerResults = convertRelIncremental(originalNode);
+    return incrementalTransformerResults;
+  }
+
+  private static IncrementalTransformerResults convertRelIncremental(RelNode originalNode) {
+    IncrementalTransformerResults incrementalTransformerResults = new IncrementalTransformerResults();
     RelShuttle converter = new RelShuttleImpl() {
       @Override
       public RelNode visit(TableScan scan) {
         RelOptTable originalTable = scan.getTable();
+
+        // Set relOptSchema
+        if (relOptSchema == null) {
+          relOptSchema = originalTable.getRelOptSchema();
+        }
+
+        // Create delta scan
         List<String> incrementalNames = new ArrayList<>(originalTable.getQualifiedName());
         String deltaTableName = incrementalNames.remove(incrementalNames.size() - 1) + "_delta";
         incrementalNames.add(deltaTableName);
@@ -49,11 +66,35 @@ public RelNode visit(TableScan scan) {
       public RelNode visit(LogicalJoin join) {
         RelNode left = join.getLeft();
         RelNode right = join.getRight();
-        RelNode incrementalLeft = convertRelIncremental(left);
-        RelNode incrementalRight = convertRelIncremental(right);
+        IncrementalTransformerResults incrementalTransformerResultsLeft = convertRelIncremental(left);
+        IncrementalTransformerResults incrementalTransformerResultsRight = convertRelIncremental(right);
+        RelNode incrementalLeft = incrementalTransformerResultsLeft.getIncrementalRelNode();
+        RelNode incrementalRight = incrementalTransformerResultsRight.getIncrementalRelNode();
+        incrementalTransformerResults
+            .addMultipleIntermediateQueryRelNodes(incrementalTransformerResultsLeft.getIntermediateQueryRelNodes());
+        incrementalTransformerResults
+            .addMultipleIntermediateQueryRelNodes(incrementalTransformerResultsRight.getIntermediateQueryRelNodes());
 
         RexBuilder rexBuilder = join.getCluster().getRexBuilder();
 
+        // Check if we can replace the left and right nodes with a scan of a materialized table
+        if (incrementalTransformerResults.containsIntermediateQueryRelNodeKey(left.getDescription())) {
+          String description = left.getDescription();
+          LogicalProject leftLastProject = createReplacementProjectNodeForGivenRelNode(description, left, rexBuilder);
+          left = leftLastProject;
+          LogicalProject leftDeltaProject =
+              createReplacementProjectNodeForGivenRelNode(description + "_delta", incrementalLeft, rexBuilder);
+          incrementalLeft = leftDeltaProject;
+        }
+        if (incrementalTransformerResults.containsIntermediateQueryRelNodeKey(right.getDescription())) {
+          LogicalProject rightLastProject =
+              createReplacementProjectNodeForGivenRelNode(right.getDescription(), right, rexBuilder);
+          right = rightLastProject;
+          LogicalProject rightDeltaProject = createReplacementProjectNodeForGivenRelNode(
+              right.getDescription() + "_delta", incrementalRight, rexBuilder);
+          incrementalRight = rightDeltaProject;
+        }
+
         LogicalProject p1 = createProjectOverJoin(join, left, incrementalRight, rexBuilder);
         LogicalProject p2 = createProjectOverJoin(join, incrementalLeft, right, rexBuilder);
         LogicalProject p3 = createProjectOverJoin(join, incrementalLeft, incrementalRight, rexBuilder);
@@ -65,45 +106,78 @@ public RelNode visit(LogicalJoin join) {
 
       @Override
       public RelNode visit(LogicalFilter filter) {
-        RelNode transformedChild = convertRelIncremental(filter.getInput());
+        IncrementalTransformerResults incrementalTransformerResultsChild = convertRelIncremental(filter.getInput());
+        RelNode transformedChild = incrementalTransformerResultsChild.getIncrementalRelNode();
+        incrementalTransformerResults
+            .addMultipleIntermediateQueryRelNodes(incrementalTransformerResultsChild.getIntermediateQueryRelNodes());
         return LogicalFilter.create(transformedChild, filter.getCondition());
       }
 
       @Override
       public RelNode visit(LogicalProject project) {
-        RelNode transformedChild = convertRelIncremental(project.getInput());
-        return LogicalProject.create(transformedChild, project.getProjects(), project.getRowType());
+        IncrementalTransformerResults incrementalTransformerResultsChild = convertRelIncremental(project.getInput());
+        RelNode transformedChild = incrementalTransformerResultsChild.getIncrementalRelNode();
+        incrementalTransformerResults
+            .addMultipleIntermediateQueryRelNodes(incrementalTransformerResultsChild.getIntermediateQueryRelNodes());
+        incrementalTransformerResults.addIntermediateQueryRelNode(project.getDescription(), project);
+        LogicalProject transformedProject =
+            LogicalProject.create(transformedChild, project.getProjects(), project.getRowType());
+        incrementalTransformerResults.addIntermediateQueryRelNode(project.getDescription() + "_delta",
+            transformedProject);
+        return transformedProject;
       }
 
       @Override
       public RelNode visit(LogicalUnion union) {
         List<RelNode> children = union.getInputs();
-        List<RelNode> transformedChildren =
+        List<IncrementalTransformerResults> incrementalTransformerResultsChildren =
             children.stream().map(child -> convertRelIncremental(child)).collect(Collectors.toList());
+        List<RelNode> transformedChildren = new ArrayList<>();
+        for (IncrementalTransformerResults incrementalTransformerResultsChild : incrementalTransformerResultsChildren) {
+          transformedChildren.add(incrementalTransformerResultsChild.getIncrementalRelNode());
+          incrementalTransformerResults
+              .addMultipleIntermediateQueryRelNodes(incrementalTransformerResultsChild.getIntermediateQueryRelNodes());
+        }
         return LogicalUnion.create(transformedChildren, union.all);
       }
 
       @Override
       public RelNode visit(LogicalAggregate aggregate) {
-        RelNode transformedChild = convertRelIncremental(aggregate.getInput());
+        IncrementalTransformerResults incrementalTransformerResultsChild = convertRelIncremental(aggregate.getInput());
+        RelNode transformedChild = incrementalTransformerResultsChild.getIncrementalRelNode();
+        incrementalTransformerResults
+            .addMultipleIntermediateQueryRelNodes(incrementalTransformerResultsChild.getIntermediateQueryRelNodes());
         return LogicalAggregate.create(transformedChild, aggregate.getGroupSet(), aggregate.getGroupSets(),
             aggregate.getAggCallList());
       }
     };
-    return originalNode.accept(converter);
+    incrementalTransformerResults.setIncrementalRelNode(originalNode.accept(converter));
+    return incrementalTransformerResults;
   }
 
-  private static LogicalProject createProjectOverJoin(LogicalJoin join, RelNode left, RelNode right,
+  private static LogicalProject createReplacementProjectNodeForGivenRelNode(String relOptTableName, RelNode relNode,
       RexBuilder rexBuilder) {
-    LogicalJoin incrementalJoin =
-        LogicalJoin.create(left, right, join.getCondition(), join.getVariablesSet(), join.getJoinType());
+    RelOptTable table =
+        RelOptTableImpl.create(relOptSchema, relNode.getRowType(), Collections.singletonList(relOptTableName), null);
+    TableScan scan = LogicalTableScan.create(relNode.getCluster(), table);
+    return createProjectOverNode(scan, rexBuilder);
+  }
+
+  private static LogicalProject createProjectOverNode(RelNode relNode, RexBuilder rexBuilder) {
     ArrayList<RexNode> projects = new ArrayList<>();
     ArrayList<String> names = new ArrayList<>();
-    IntStream.range(0, incrementalJoin.getRowType().getFieldList().size()).forEach(i -> {
-      projects.add(rexBuilder.makeInputRef(incrementalJoin, i));
-      names.add(incrementalJoin.getRowType().getFieldNames().get(i));
+    IntStream.range(0, relNode.getRowType().getFieldList().size()).forEach(i -> {
+      projects.add(rexBuilder.makeInputRef(relNode, i));
+      names.add(relNode.getRowType().getFieldNames().get(i));
     });
-    return LogicalProject.create(incrementalJoin, projects, names);
+    return LogicalProject.create(relNode, projects, names);
+  }
+
+  private static LogicalProject createProjectOverJoin(LogicalJoin join, RelNode left, RelNode right,
+      RexBuilder rexBuilder) {
+    LogicalJoin incrementalJoin =
+        LogicalJoin.create(left, right, join.getCondition(), join.getVariablesSet(), join.getJoinType());
+    return createProjectOverNode(incrementalJoin, rexBuilder);
   }
 
 }

coral-incremental/src/test/java/com/linkedin/coral/incremental/RelToIncrementalSqlConverterTest.java

@@ -41,7 +41,9 @@ public void afterClass() throws IOException {
   }
 
   public String convert(RelNode relNode) {
-    RelNode incrementalRelNode = RelNodeIncrementalTransformer.convertRelIncremental(relNode);
+    IncrementalTransformerResults incrementalTransformerResults =
+        RelNodeIncrementalTransformer.performIncrementalTransformation(relNode);
+    RelNode incrementalRelNode = incrementalTransformerResults.getIncrementalRelNode();
     CoralRelToSqlNodeConverter converter = new CoralRelToSqlNodeConverter();
     SqlNode sqlNode = converter.convert(incrementalRelNode);
     return sqlNode.toSqlString(converter.INSTANCE).getSql();
@@ -81,41 +83,6 @@ public void testJoinWithFilter() {
     assertEquals(getIncrementalModification(sql), expected);
   }
 
-  @Test
-  public void testJoinWithNestedFilter() {
-    String sql =
-        "WITH tmp AS (SELECT * from test.bar1 WHERE test.bar1.x > 10), tmp2 AS (SELECT * from test.bar2) SELECT * FROM tmp JOIN tmp2 ON tmp.x = tmp2.x";
-    String expected = "SELECT *\n" + "FROM (SELECT *\n" + "FROM (SELECT *\n" + "FROM test.bar1 AS bar1\n"
-        + "WHERE bar1.x > 10) AS t\n" + "INNER JOIN test.bar2_delta AS bar2_delta ON t.x = bar2_delta.x\n"
-        + "UNION ALL\n" + "SELECT *\n" + "FROM (SELECT *\n" + "FROM test.bar1_delta AS bar1_delta\n"
-        + "WHERE bar1_delta.x > 10) AS t0\n" + "INNER JOIN test.bar2 AS bar2 ON t0.x = bar2.x) AS t1\n" + "UNION ALL\n"
-        + "SELECT *\n" + "FROM (SELECT *\n" + "FROM test.bar1_delta AS bar1_delta0\n"
-        + "WHERE bar1_delta0.x > 10) AS t2\n" + "INNER JOIN test.bar2_delta AS bar2_delta0 ON t2.x = bar2_delta0.x";
-    assertEquals(getIncrementalModification(sql), expected);
-  }
-
-  @Test
-  public void testNestedJoin() {
-    String sql =
-        "WITH tmp AS (SELECT * FROM test.bar1 INNER JOIN test.bar2 ON test.bar1.x = test.bar2.x) SELECT * FROM tmp INNER JOIN test.bar3 ON tmp.x = test.bar3.x";
-    String expected = "SELECT *\n" + "FROM (SELECT *\n" + "FROM test.bar1 AS bar1\n"
-        + "INNER JOIN test.bar2 AS bar2 ON bar1.x = bar2.x\n"
-        + "INNER JOIN test.bar3_delta AS bar3_delta ON bar1.x = bar3_delta.x\n" + "UNION ALL\n" + "SELECT *\n"
-        + "FROM (SELECT *\n" + "FROM (SELECT *\n" + "FROM test.bar1 AS bar10\n"
-        + "INNER JOIN test.bar2_delta AS bar2_delta ON bar10.x = bar2_delta.x\n" + "UNION ALL\n" + "SELECT *\n"
-        + "FROM test.bar1_delta AS bar1_delta\n" + "INNER JOIN test.bar2 AS bar20 ON bar1_delta.x = bar20.x) AS t\n"
-        + "UNION ALL\n" + "SELECT *\n" + "FROM test.bar1_delta AS bar1_delta0\n"
-        + "INNER JOIN test.bar2_delta AS bar2_delta0 ON bar1_delta0.x = bar2_delta0.x) AS t0\n"
-        + "INNER JOIN test.bar3 AS bar3 ON t0.x = bar3.x) AS t1\n" + "UNION ALL\n" + "SELECT *\n" + "FROM (SELECT *\n"
-        + "FROM (SELECT *\n" + "FROM test.bar1 AS bar11\n"
-        + "INNER JOIN test.bar2_delta AS bar2_delta1 ON bar11.x = bar2_delta1.x\n" + "UNION ALL\n" + "SELECT *\n"
-        + "FROM test.bar1_delta AS bar1_delta1\n" + "INNER JOIN test.bar2 AS bar21 ON bar1_delta1.x = bar21.x) AS t2\n"
-        + "UNION ALL\n" + "SELECT *\n" + "FROM test.bar1_delta AS bar1_delta2\n"
-        + "INNER JOIN test.bar2_delta AS bar2_delta2 ON bar1_delta2.x = bar2_delta2.x) AS t3\n"
-        + "INNER JOIN test.bar3_delta AS bar3_delta0 ON t3.x = bar3_delta0.x";
-    assertEquals(getIncrementalModification(sql), expected);
-  }
-
   @Test
   public void testUnion() {
     String sql = "SELECT * FROM test.bar1 UNION SELECT * FROM test.bar2 UNION SELECT * FROM test.bar3";
@@ -143,4 +110,12 @@ public void testSelectSpecificJoin() {
         + "INNER JOIN test.bar2_delta AS bar2_delta0 ON bar1_delta0.x = bar2_delta0.x) AS t0";
     assertEquals(getIncrementalModification(sql), expected);
   }
+
+  @Test
+  public void testJoinOverJoin() {
+    // Debugging only
+    String nestedJoin = "SELECT a1, a2 FROM test.alpha JOIN test.beta ON test.alpha.a1 = test.beta.b1";
+    String sql = "SELECT a2, g1 FROM (" + nestedJoin + ") AS nj JOIN test.gamma ON nj.a2 = test.gamma.g2";
+    getIncrementalModification(sql);
+  }
 }

coral-incremental/src/test/java/com/linkedin/coral/incremental/TestUtils.java

@@ -60,6 +60,10 @@ public static void initializeViews(HiveConf conf) throws HiveException, MetaExce
     run(driver, "CREATE TABLE IF NOT EXISTS test.bar1(x int, y double)");
     run(driver, "CREATE TABLE IF NOT EXISTS test.bar2(x int, y double)");
     run(driver, "CREATE TABLE IF NOT EXISTS test.bar3(x int, y double)");
+
+    run(driver, "CREATE TABLE IF NOT EXISTS test.alpha(a1 int, a2 double)");
+    run(driver, "CREATE TABLE IF NOT EXISTS test.beta(b1 int, b2 double)");
+    run(driver, "CREATE TABLE IF NOT EXISTS test.gamma(g1 int, g2 double)");
   }
 
   public static HiveConf loadResourceHiveConf() {

coral-service/src/main/java/com/linkedin/coral/coralservice/utils/IncrementalUtils.java

@@ -8,6 +8,7 @@
 import org.apache.calcite.rel.RelNode;
 
 import com.linkedin.coral.hive.hive2rel.HiveToRelConverter;
+import com.linkedin.coral.incremental.IncrementalTransformerResults;
 import com.linkedin.coral.incremental.RelNodeIncrementalTransformer;
 import com.linkedin.coral.spark.CoralSpark;
 
@@ -18,7 +19,9 @@ public class IncrementalUtils {
 
   public static String getSparkIncrementalQueryFromUserSql(String query) {
     RelNode originalNode = new HiveToRelConverter(hiveMetastoreClient).convertSql(query);
-    RelNode incrementalRelNode = RelNodeIncrementalTransformer.convertRelIncremental(originalNode);
+    IncrementalTransformerResults incrementalTransformerResults =
+        RelNodeIncrementalTransformer.performIncrementalTransformation(originalNode);
+    RelNode incrementalRelNode = incrementalTransformerResults.getIncrementalRelNode();
     CoralSpark coralSpark = CoralSpark.create(incrementalRelNode);
     return coralSpark.getSparkSql();
   }