[calcite] branch master updated: [CALCITE-4273] Support get expression lineage for Calc

[yanlin]

This is an automated email from the ASF dual-hosted git repository.

yanlin pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/calcite.git


The following commit(s) were added to refs/heads/master by this push:
     new 2e9edae  [CALCITE-4273] Support get expression lineage for Calc
2e9edae is described below

commit 2e9edae7fc57ab9c9c7c097008724ac99a1791a3
Author: yanlin-Lynn <1989yanlinw...@163.com>
AuthorDate: Tue Sep 22 18:02:11 2020 +0800

    [CALCITE-4273] Support get expression lineage for Calc
---
 .../rel/metadata/RelMdExpressionLineage.java       | 31 ++++++++++++++++++++++
 .../org/apache/calcite/test/RelMetadataTest.java   | 21 +++++++++++++++
 2 files changed, 52 insertions(+)

diff --git 
a/core/src/main/java/org/apache/calcite/rel/metadata/RelMdExpressionLineage.java
 
b/core/src/main/java/org/apache/calcite/rel/metadata/RelMdExpressionLineage.java
index 41189a0..f1c97cb 100644
--- 
a/core/src/main/java/org/apache/calcite/rel/metadata/RelMdExpressionLineage.java
+++ 
b/core/src/main/java/org/apache/calcite/rel/metadata/RelMdExpressionLineage.java
@@ -21,6 +21,7 @@ import org.apache.calcite.plan.hep.HepRelVertex;
 import org.apache.calcite.plan.volcano.RelSubset;
 import org.apache.calcite.rel.RelNode;
 import org.apache.calcite.rel.core.Aggregate;
+import org.apache.calcite.rel.core.Calc;
 import org.apache.calcite.rel.core.Exchange;
 import org.apache.calcite.rel.core.Filter;
 import org.apache.calcite.rel.core.Join;
@@ -42,6 +43,7 @@ import org.apache.calcite.rex.RexUtil;
 import org.apache.calcite.sql.validate.SqlValidatorUtil;
 import org.apache.calcite.util.BuiltInMethod;
 import org.apache.calcite.util.ImmutableBitSet;
+import org.apache.calcite.util.Pair;
 import org.apache.calcite.util.Util;
 
 import com.google.common.collect.HashMultimap;
@@ -407,6 +409,35 @@ public class RelMdExpressionLineage
   }
 
   /**
+   * Expression lineage from Calc.
+   */
+  public Set<RexNode> getExpressionLineage(Calc calc,
+      RelMetadataQuery mq, RexNode outputExpression) {
+    final RelNode input = calc.getInput();
+    final RexBuilder rexBuilder = calc.getCluster().getRexBuilder();
+    // Extract input fields referenced by expression
+    final ImmutableBitSet inputFieldsUsed = extractInputRefs(outputExpression);
+
+    // Infer column origin expressions for given references
+    final Map<RexInputRef, Set<RexNode>> mapping = new LinkedHashMap<>();
+    Pair<ImmutableList<RexNode>, ImmutableList<RexNode>> calcProjectsAndFilter 
=
+        calc.getProgram().split();
+    for (int idx : inputFieldsUsed) {
+      final RexNode inputExpr = calcProjectsAndFilter.getKey().get(idx);
+      final Set<RexNode> originalExprs = mq.getExpressionLineage(input, 
inputExpr);
+      if (originalExprs == null) {
+        // Bail out
+        return null;
+      }
+      final RexInputRef ref = RexInputRef.of(idx, 
calc.getRowType().getFieldList());
+      mapping.put(ref, originalExprs);
+    }
+
+    // Return result
+    return createAllPossibleExpressions(rexBuilder, outputExpression, mapping);
+  }
+
+  /**
    * Given an expression, it will create all equivalent expressions resulting
    * from replacing all possible combinations of references in the mapping by
    * the corresponding expressions.
diff --git a/core/src/test/java/org/apache/calcite/test/RelMetadataTest.java 
b/core/src/test/java/org/apache/calcite/test/RelMetadataTest.java
index bfda204..e181a0b 100644
--- a/core/src/test/java/org/apache/calcite/test/RelMetadataTest.java
+++ b/core/src/test/java/org/apache/calcite/test/RelMetadataTest.java
@@ -2507,6 +2507,27 @@ public class RelMetadataTest extends SqlToRelTestBase {
     assertNull(r);
   }
 
+  @Test void testExpressionLineageCalc() {
+    final RelNode rel = convertSql("select sal from (\n"
+        + " select deptno, empno, sal + 1 as sal, job from emp) "
+        + "where deptno = 10");
+    final HepProgramBuilder programBuilder = HepProgram.builder();
+    programBuilder.addRuleInstance(CoreRules.PROJECT_TO_CALC);
+    programBuilder.addRuleInstance(CoreRules.FILTER_TO_CALC);
+    programBuilder.addRuleInstance(CoreRules.CALC_MERGE);
+    final HepPlanner planner = new HepPlanner(programBuilder.build());
+    planner.setRoot(rel);
+    final RelNode optimizedRel = planner.findBestExp();
+    final RelMetadataQuery mq = optimizedRel.getCluster().getMetadataQuery();
+
+    final RexNode ref = RexInputRef.of(0, 
optimizedRel.getRowType().getFieldList());
+    final Set<RexNode> r = mq.getExpressionLineage(optimizedRel, ref);
+
+    assertThat(r.size(), is(1));
+    final String resultString = r.iterator().next().toString();
+    assertThat(resultString, is("+([CATALOG, SALES, EMP].#0.$5, 1)"));
+  }
+
   @Test void testAllPredicates() {
     final Project rel = (Project) convertSql("select * from emp, dept");
     final Join join = (Join) rel.getInput();