Java Code Examples for org.apache.calcite.rel.metadata.RelMetadataQuery#getExpressionLineage()

@Test void testExpressionLineageTwoColumns() {
  // mgr is column 3 in catalog.sales.emp
  // deptno is column 7 in catalog.sales.emp
  final RelNode rel = convertSql("select mgr, deptno from emp");
  final RelMetadataQuery mq = rel.getCluster().getMetadataQuery();

  final RexNode ref1 = RexInputRef.of(0, rel.getRowType().getFieldList());
  final Set<RexNode> r1 = mq.getExpressionLineage(rel, ref1);
  assertThat(r1.size(), is(1));
  final RexTableInputRef result1 = (RexTableInputRef) r1.iterator().next();
  assertThat(result1.getQualifiedName(), is(EMP_QNAME));
  assertThat(result1.getIndex(), is(3));

  final RexNode ref2 = RexInputRef.of(1, rel.getRowType().getFieldList());
  final Set<RexNode> r2 = mq.getExpressionLineage(rel, ref2);
  assertThat(r2.size(), is(1));
  final RexTableInputRef result2 = (RexTableInputRef) r2.iterator().next();
  assertThat(result2.getQualifiedName(), is(EMP_QNAME));
  assertThat(result2.getIndex(), is(7));

  assertThat(result1.getIdentifier(), is(result2.getIdentifier()));
}
 

@Test void testExpressionLineageTwoColumnsSwapped() {
  // deptno is column 7 in catalog.sales.emp
  // mgr is column 3 in catalog.sales.emp
  final RelNode rel = convertSql("select deptno, mgr from emp");
  final RelMetadataQuery mq = rel.getCluster().getMetadataQuery();

  final RexNode ref1 = RexInputRef.of(0, rel.getRowType().getFieldList());
  final Set<RexNode> r1 = mq.getExpressionLineage(rel, ref1);
  assertThat(r1.size(), is(1));
  final RexTableInputRef result1 = (RexTableInputRef) r1.iterator().next();
  assertThat(result1.getQualifiedName(), is(EMP_QNAME));
  assertThat(result1.getIndex(), is(7));

  final RexNode ref2 = RexInputRef.of(1, rel.getRowType().getFieldList());
  final Set<RexNode> r2 = mq.getExpressionLineage(rel, ref2);
  assertThat(r2.size(), is(1));
  final RexTableInputRef result2 = (RexTableInputRef) r2.iterator().next();
  assertThat(result2.getQualifiedName(), is(EMP_QNAME));
  assertThat(result2.getIndex(), is(3));

  assertThat(result1.getIdentifier(), is(result2.getIdentifier()));
}
 

@Test void testExpressionLineageCombineTwoColumns() {
  // empno is column 0 in catalog.sales.emp
  // deptno is column 7 in catalog.sales.emp
  final RelNode rel = convertSql("select empno + deptno from emp");
  final RelMetadataQuery mq = rel.getCluster().getMetadataQuery();

  final RexNode ref = RexInputRef.of(0, rel.getRowType().getFieldList());
  final Set<RexNode> r = mq.getExpressionLineage(rel, ref);

  assertThat(r.size(), is(1));
  final RexNode result = r.iterator().next();
  assertThat(result.getKind(), is(SqlKind.PLUS));
  final RexCall call = (RexCall) result;
  assertThat(call.getOperands().size(), is(2));
  final RexTableInputRef inputRef1 = (RexTableInputRef) call.getOperands().get(0);
  assertThat(inputRef1.getQualifiedName(), is(EMP_QNAME));
  assertThat(inputRef1.getIndex(), is(0));
  final RexTableInputRef inputRef2 = (RexTableInputRef) call.getOperands().get(1);
  assertThat(inputRef2.getQualifiedName(), is(EMP_QNAME));
  assertThat(inputRef2.getIndex(), is(7));
  assertThat(inputRef1.getIdentifier(), is(inputRef2.getIdentifier()));
}
 

@Test void testExpressionLineageUnion() {
  // sal is column 5 in catalog.sales.emp
  final RelNode rel = convertSql("select sal from (\n"
      + "  select * from emp union all select * from emp) "
      + "where deptno = 10");
  final RelNode tableRel = convertSql("select * from emp");
  final RelMetadataQuery mq = rel.getCluster().getMetadataQuery();

  final RexNode ref = RexInputRef.of(0, rel.getRowType().getFieldList());
  final Set<RexNode> r = mq.getExpressionLineage(rel, ref);
  final String inputRef = RexInputRef.of(5, tableRel.getRowType().getFieldList()).toString();
  assertThat(r.size(), is(2));
  for (RexNode result : r) {
    final String resultString = result.toString();
    assertThat(resultString, startsWith(EMP_QNAME.toString()));
    assertThat(resultString, endsWith(inputRef));
  }

  Iterator<RexNode> it = r.iterator();
  assertThat(((RexTableInputRef) it.next()).getIdentifier(),
      not(((RexTableInputRef) it.next()).getIdentifier()));
}
 

/**
 * It swaps the column references and then the table references of the input
 * expressions using the equivalence classes and the table mapping.
 */
private static NodeLineage generateSwapColumnTableReferencesLineage(RexBuilder rexBuilder, RelMetadataQuery mq,
        RelNode node, BiMap<RelTableRef, RelTableRef> tableMapping, EquivalenceClasses ec,
        List<RexNode> nodeExprs) {
    final Map<RexNode, Integer> exprsLineage = new HashMap<>();
    final Map<RexNode, Integer> exprsLineageLosslessCasts = new HashMap<>();
    for (int i = 0; i < nodeExprs.size(); i++) {
        final Set<RexNode> s = mq.getExpressionLineage(node, nodeExprs.get(i));
        if (s == null) {
            // Next expression
            continue;
        }
        // We only support project - filter - join, thus it should map to
        // a single expression
        final RexNode node2 = Iterables.getOnlyElement(s);
        // Rewrite expr. First we take first element from the corresponding equivalence class,
        // then we swap the table references following the table mapping
        final RexNode e = RexUtil.swapColumnTableReferences(rexBuilder, node2, ec.getEquivalenceClassesMap(),
                tableMapping);
        exprsLineage.put(e, i);
        if (RexUtil.isLosslessCast(e)) {
            exprsLineageLosslessCasts.put(((RexCall) e).getOperands().get(0), i);
        }
    }
    return new NodeLineage(exprsLineage, exprsLineageLosslessCasts);
}
 

@Test void testExpressionLineageRightJoinRight() {
  // ename is column 0 in catalog.sales.bonus
  final RelNode rel = convertSql("select bonus.ename from emp right join bonus using (ename)");
  final RelMetadataQuery mq = rel.getCluster().getMetadataQuery();

  final RexNode ref = RexInputRef.of(0, rel.getRowType().getFieldList());
  final Set<RexNode> r = mq.getExpressionLineage(rel, ref);
  assertThat(r.size(), is(1));
  final RexTableInputRef result = (RexTableInputRef) r.iterator().next();
  assertThat(result.getQualifiedName(), equalTo(ImmutableList.of("CATALOG", "SALES", "BONUS")));
  assertThat(result.getIndex(), is(0));
}
 

@Test void testExpressionLineageValues() {
  // lineage cannot be determined
  final RelNode rel = convertSql("select * from (values (1), (2)) as t(c)");
  final RelMetadataQuery mq = rel.getCluster().getMetadataQuery();

  final RexNode ref = RexInputRef.of(0, rel.getRowType().getFieldList());
  final Set<RexNode> r = mq.getExpressionLineage(rel, ref);
  assertNull(r);
}
 

@Test void testExpressionLineageMultiUnion() {
  // empno is column 0 in catalog.sales.emp
  // sal is column 5 in catalog.sales.emp
  final RelNode rel = convertSql("select a.empno + b.sal from\n"
      + " (select empno, ename from emp,dept) a join "
      + " (select * from emp union all select * from emp) b\n"
      + " on a.empno = b.empno\n"
      + " where b.deptno = 10");
  final RelMetadataQuery mq = rel.getCluster().getMetadataQuery();

  final RexNode ref = RexInputRef.of(0, rel.getRowType().getFieldList());
  final Set<RexNode> r = mq.getExpressionLineage(rel, ref);

  // With the union, we should get two origins
  // The first one should be the same one: join
  // The second should come from each union input
  final Set<List<String>> set = new HashSet<>();
  assertThat(r.size(), is(2));
  for (RexNode result : r) {
    assertThat(result.getKind(), is(SqlKind.PLUS));
    final RexCall call = (RexCall) result;
    assertThat(call.getOperands().size(), is(2));
    final RexTableInputRef inputRef1 = (RexTableInputRef) call.getOperands().get(0);
    assertThat(inputRef1.getQualifiedName(), is(EMP_QNAME));
    // Add join alpha to set
    set.add(inputRef1.getQualifiedName());
    assertThat(inputRef1.getIndex(), is(0));
    final RexTableInputRef inputRef2 = (RexTableInputRef) call.getOperands().get(1);
    assertThat(inputRef2.getQualifiedName(), is(EMP_QNAME));
    assertThat(inputRef2.getIndex(), is(5));
    assertThat(inputRef1.getIdentifier(), not(inputRef2.getIdentifier()));
  }
  assertThat(set.size(), is(1));
}
 

@Test void testExpressionLineageAggregateAggColumn() {
  // lineage cannot be determined
  final RelNode rel = convertSql("select deptno, count(*) from emp where deptno > 10 "
      + "group by deptno having count(*) = 0");
  final RelMetadataQuery mq = rel.getCluster().getMetadataQuery();

  final RexNode ref = RexInputRef.of(1, rel.getRowType().getFieldList());
  final Set<RexNode> r = mq.getExpressionLineage(rel, ref);
  assertNull(r);
}
 

@Test void testExpressionLineageAggregateGroupColumn() {
  // deptno is column 7 in catalog.sales.emp
  final RelNode rel = convertSql("select deptno, count(*) from emp where deptno > 10 "
      + "group by deptno having count(*) = 0");
  final RelNode tableRel = convertSql("select * from emp");
  final RelMetadataQuery mq = rel.getCluster().getMetadataQuery();

  final RexNode ref = RexInputRef.of(0, rel.getRowType().getFieldList());
  final Set<RexNode> r = mq.getExpressionLineage(rel, ref);
  final String inputRef = RexInputRef.of(7, tableRel.getRowType().getFieldList()).toString();
  assertThat(r.size(), is(1));
  final String resultString = r.iterator().next().toString();
  assertThat(resultString, startsWith(EMP_QNAME.toString()));
  assertThat(resultString, endsWith(inputRef));
}
 

@Test void testExpressionLineageFilter() {
  // ename is column 1 in catalog.sales.emp
  final RelNode rel = convertSql("select ename from emp where deptno = 10");
  final RelNode tableRel = convertSql("select * from emp");
  final RelMetadataQuery mq = rel.getCluster().getMetadataQuery();

  final RexNode ref = RexInputRef.of(0, rel.getRowType().getFieldList());
  final Set<RexNode> r = mq.getExpressionLineage(rel, ref);
  final String inputRef = RexInputRef.of(1, tableRel.getRowType().getFieldList()).toString();
  assertThat(r.size(), is(1));
  final String resultString = r.iterator().next().toString();
  assertThat(resultString, startsWith(EMP_QNAME.toString()));
  assertThat(resultString, endsWith(inputRef));
}
 

@Test void testExpressionLineageOuterJoin() {
  // lineage cannot be determined
  final RelNode rel = convertSql("select name as dname from emp left outer join dept"
      + " on emp.deptno = dept.deptno");
  final RelMetadataQuery mq = rel.getCluster().getMetadataQuery();

  final RexNode ref = RexInputRef.of(0, rel.getRowType().getFieldList());
  final Set<RexNode> r = mq.getExpressionLineage(rel, ref);
  assertNull(r);
}
 

@Test void testExpressionLineageSelfJoin() {
  // deptno is column 7 in catalog.sales.emp
  // sal is column 5 in catalog.sales.emp
  final RelNode rel = convertSql("select a.deptno, b.sal from (select * from emp limit 7) as a\n"
      + "inner join (select * from emp limit 2) as b\n"
      + "on a.deptno = b.deptno");
  final RelNode tableRel = convertSql("select * from emp");
  final RelMetadataQuery mq = tableRel.getCluster().getMetadataQuery();

  final RexNode ref1 = RexInputRef.of(0, rel.getRowType().getFieldList());
  final Set<RexNode> r1 = mq.getExpressionLineage(rel, ref1);
  final String inputRef1 = RexInputRef.of(7, tableRel.getRowType().getFieldList()).toString();
  assertThat(r1.size(), is(1));
  final String resultString1 = r1.iterator().next().toString();
  assertThat(resultString1, startsWith(EMP_QNAME.toString()));
  assertThat(resultString1, endsWith(inputRef1));

  final RexNode ref2 = RexInputRef.of(1, rel.getRowType().getFieldList());
  final Set<RexNode> r2 = mq.getExpressionLineage(rel, ref2);
  final String inputRef2 = RexInputRef.of(5, tableRel.getRowType().getFieldList()).toString();
  assertThat(r2.size(), is(1));
  final String resultString2 = r2.iterator().next().toString();
  assertThat(resultString2, startsWith(EMP_QNAME.toString()));
  assertThat(resultString2, endsWith(inputRef2));

  assertThat(((RexTableInputRef) r1.iterator().next()).getIdentifier(),
      not(((RexTableInputRef) r2.iterator().next()).getIdentifier()));
}
 

@Test void testExpressionLineageLeftJoinLeft() {
  // ename is column 1 in catalog.sales.emp
  final RelNode rel = convertSql("select ename from emp left join dept using (deptno)");
  final RelMetadataQuery mq = rel.getCluster().getMetadataQuery();

  final RexNode ref = RexInputRef.of(0, rel.getRowType().getFieldList());
  final Set<RexNode> r = mq.getExpressionLineage(rel, ref);
  assertThat(r.size(), is(1));
  final RexTableInputRef result = (RexTableInputRef) r.iterator().next();
  assertThat(result.getQualifiedName(), is(EMP_QNAME));
  assertThat(result.getIndex(), is(1));
}
 

@Test void testExpressionLineageInnerJoinRight() {
  // ename is column 0 in catalog.sales.bonus
  final RelNode rel = convertSql("select bonus.ename from emp join bonus using (ename)");
  final RelMetadataQuery mq = rel.getCluster().getMetadataQuery();

  final RexNode ref = RexInputRef.of(0, rel.getRowType().getFieldList());
  final Set<RexNode> r = mq.getExpressionLineage(rel, ref);
  assertThat(r.size(), is(1));
  final RexTableInputRef result = (RexTableInputRef) r.iterator().next();
  assertThat(result.getQualifiedName(), equalTo(ImmutableList.of("CATALOG", "SALES", "BONUS")));
  assertThat(result.getIndex(), is(0));
}
 

@Test void testExpressionLineageInnerJoinLeft() {
  // ename is column 1 in catalog.sales.emp
  final RelNode rel = convertSql("select ename from emp,dept");
  final RelMetadataQuery mq = rel.getCluster().getMetadataQuery();

  final RexNode ref = RexInputRef.of(0, rel.getRowType().getFieldList());
  final Set<RexNode> r = mq.getExpressionLineage(rel, ref);
  assertThat(r.size(), is(1));
  final RexTableInputRef result = (RexTableInputRef) r.iterator().next();
  assertThat(result.getQualifiedName(), is(EMP_QNAME));
  assertThat(result.getIndex(), is(1));
}
 

@Test void testExpressionLineageStar() {
  // All columns in output
  final RelNode tableRel = convertSql("select * from emp");
  final RelMetadataQuery mq = tableRel.getCluster().getMetadataQuery();

  final RexNode ref = RexInputRef.of(4, tableRel.getRowType().getFieldList());
  final Set<RexNode> r = mq.getExpressionLineage(tableRel, ref);
  final String inputRef = RexInputRef.of(4, tableRel.getRowType().getFieldList()).toString();
  assertThat(r.size(), is(1));
  final String resultString = r.iterator().next().toString();
  assertThat(resultString, startsWith(EMP_QNAME.toString()));
  assertThat(resultString, endsWith(inputRef));
}
 

/**
 * It swaps the column references and then the table references of the input
 * expressions using the equivalence classes and the table mapping.
 */
protected NodeLineage generateSwapColumnTableReferencesLineage(
    RexBuilder rexBuilder,
    RelMetadataQuery mq,
    RelNode node,
    BiMap<RelTableRef, RelTableRef> tableMapping,
    EquivalenceClasses ec,
    List<RexNode> nodeExprs) {
  final Map<RexNode, Integer> exprsLineage = new HashMap<>();
  final Map<RexNode, Integer> exprsLineageLosslessCasts = new HashMap<>();
  for (int i = 0; i < nodeExprs.size(); i++) {
    final Set<RexNode> s = mq.getExpressionLineage(node, nodeExprs.get(i));
    if (s == null) {
      // Next expression
      continue;
    }
    // We only support project - filter - join, thus it should map to
    // a single expression
    final RexNode node2 = Iterables.getOnlyElement(s);
    // Rewrite expr. First we take first element from the corresponding equivalence class,
    // then we swap the table references following the table mapping
    final RexNode e = RexUtil.swapColumnTableReferences(rexBuilder, node2,
        ec.getEquivalenceClassesMap(), tableMapping);
    exprsLineage.put(e, i);
    if (RexUtil.isLosslessCast(e)) {
      exprsLineageLosslessCasts.put(((RexCall) e).getOperands().get(0), i);
    }
  }
  return new NodeLineage(exprsLineage, exprsLineageLosslessCasts);
}