Java Code Examples for org.apache.calcite.rex.RexInputRef#of()

/**
 * Expression lineage from {@link TableScan}.
 *
 * <p>We extract the fields referenced by the expression and we express them
 * using {@link RexTableInputRef}.
 */
public Set<RexNode> getExpressionLineage(TableScan rel,
    RelMetadataQuery mq, RexNode outputExpression) {
  final RexBuilder rexBuilder = rel.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<>();
  for (int idx : inputFieldsUsed) {
    final RexNode inputRef = RexTableInputRef.of(
        RelTableRef.of(rel.getTable(), 0),
        RexInputRef.of(idx, rel.getRowType().getFieldList()));
    final RexInputRef ref = RexInputRef.of(idx, rel.getRowType().getFieldList());
    mapping.put(ref, ImmutableSet.of(inputRef));
  }

  // Return result
  return createAllPossibleExpressions(rexBuilder, outputExpression, mapping);
}
 

@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()));
}
 

/**
 * Expression lineage from Project.
 */
public Set<RexNode> getExpressionLineage(Project rel,
    final RelMetadataQuery mq, RexNode outputExpression) {
  final RelNode input = rel.getInput();
  final RexBuilder rexBuilder = rel.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<>();
  for (int idx : inputFieldsUsed) {
    final RexNode inputExpr = rel.getProjects().get(idx);
    final Set<RexNode> originalExprs = mq.getExpressionLineage(input, inputExpr);
    if (originalExprs == null) {
      // Bail out
      return null;
    }
    final RexInputRef ref = RexInputRef.of(idx, rel.getRowType().getFieldList());
    mapping.put(ref, originalExprs);
  }

  // Return result
  return createAllPossibleExpressions(rexBuilder, outputExpression, mapping);
}
 

/**
 * Expression lineage from {@link TableScan}.
 *
 * <p>We extract the fields referenced by the expression and we express them
 * using {@link RexTableInputRef}.
 */
public Set<RexNode> getExpressionLineage(TableScan rel,
    RelMetadataQuery mq, RexNode outputExpression) {
  final RexBuilder rexBuilder = rel.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<>();
  for (int idx : inputFieldsUsed) {
    final RexNode inputRef = RexTableInputRef.of(
        RelTableRef.of(rel.getTable(), 0),
        RexInputRef.of(idx, rel.getRowType().getFieldList()));
    final RexInputRef ref = RexInputRef.of(idx, rel.getRowType().getFieldList());
    mapping.put(ref, ImmutableSet.of(inputRef));
  }

  // Return result
  return createAllPossibleExpressions(rexBuilder, outputExpression, mapping);
}
 

/**
 * Test {@link RelOptUtil#splitJoinCondition(RelNode, RelNode, RexNode, List, List, List)}
 * where the join condition contains an expanded version of IS NOT DISTINCT using CASE
 */
@Test void testSplitJoinConditionExpandedIsNotDistinctFromUsingCase() {
  int leftJoinIndex = empScan.getRowType().getFieldNames().indexOf("DEPTNO");
  int rightJoinIndex = deptRow.getFieldNames().indexOf("DEPTNO");

  RexInputRef leftKeyInputRef = RexInputRef.of(leftJoinIndex, empDeptJoinRelFields);
  RexInputRef rightKeyInputRef =
      RexInputRef.of(empRow.getFieldCount() + rightJoinIndex, empDeptJoinRelFields);
  RexNode joinCond = RelOptUtil.isDistinctFrom(
      relBuilder.getRexBuilder(),
      leftKeyInputRef,
      rightKeyInputRef,
      true);


  splitJoinConditionHelper(
      joinCond,
      Collections.singletonList(leftJoinIndex),
      Collections.singletonList(rightJoinIndex),
      Collections.singletonList(false),
      relBuilder.literal(true));
}
 

@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));
}
 

private RexNode restructure(RelDataType structuredType) {
    // Access null indicator for entire structure.
    RexInputRef nullIndicator = RexInputRef.of(iRestructureInput++, flattenedRootType.getFieldList());

    // Use NEW to put flattened data back together into a structure.
    List<RexNode> inputExprs = restructureFields(structuredType);
    RexNode newInvocation = rexBuilder.makeNewInvocation(structuredType, inputExprs);

    if (!structuredType.isNullable()) {
        // Optimize away the null test.
        return newInvocation;
    }

    // Construct a CASE expression to handle the structure-level null
    // indicator.
    RexNode[] caseOperands = new RexNode[3];

    // WHEN StructuredType.Indicator IS NULL
    caseOperands[0] = rexBuilder.makeCall(SqlStdOperatorTable.IS_NULL, nullIndicator);

    // THEN CAST(NULL AS StructuredType)
    caseOperands[1] = rexBuilder.makeCast(structuredType, rexBuilder.constantNull());

    // ELSE NEW StructuredType(inputs...) END
    caseOperands[2] = newInvocation;

    return rexBuilder.makeCall(SqlStdOperatorTable.CASE, caseOperands);
}
 

@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 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);
}
 

/**
 * Add the Filter condition to the list obtained from the input.
 */
public RelOptPredicateList getAllPredicates(Filter filter, RelMetadataQuery mq) {
  final RelNode input = filter.getInput();
  final RexBuilder rexBuilder = filter.getCluster().getRexBuilder();
  final RexNode pred = filter.getCondition();

  final RelOptPredicateList predsBelow = mq.getAllPredicates(input);
  if (predsBelow == null) {
    // Safety check
    return null;
  }

  // Extract input fields referenced by Filter condition
  final Set<RelDataTypeField> inputExtraFields = new LinkedHashSet<>();
  final RelOptUtil.InputFinder inputFinder = new RelOptUtil.InputFinder(inputExtraFields);
  pred.accept(inputFinder);
  final ImmutableBitSet inputFieldsUsed = inputFinder.build();

  // Infer column origin expressions for given references
  final Map<RexInputRef, Set<RexNode>> mapping = new LinkedHashMap<>();
  for (int idx : inputFieldsUsed) {
    final RexInputRef ref = RexInputRef.of(idx, filter.getRowType().getFieldList());
    final Set<RexNode> originalExprs = mq.getExpressionLineage(filter, ref);
    if (originalExprs == null) {
      // Bail out
      return null;
    }
    mapping.put(ref, originalExprs);
  }

  // Replace with new expressions and return union of predicates
  final Set<RexNode> allExprs =
      RelMdExpressionLineage.createAllPossibleExpressions(rexBuilder, pred, mapping);
  if (allExprs == null) {
    return null;
  }
  return predsBelow.union(rexBuilder, RelOptPredicateList.of(rexBuilder, allExprs));
}
 

@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 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 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 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));
}
 

@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 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));
}
 

/**
 * Add the Join condition to the list obtained from the input.
 */
public RelOptPredicateList getAllPredicates(Join join, RelMetadataQuery mq) {
  if (join.getJoinType() != JoinRelType.INNER) {
    // We cannot map origin of this expression.
    return null;
  }

  final RexBuilder rexBuilder = join.getCluster().getRexBuilder();
  final RexNode pred = join.getCondition();

  final Multimap<List<String>, RelTableRef> qualifiedNamesToRefs = HashMultimap.create();
  RelOptPredicateList newPreds = RelOptPredicateList.EMPTY;
  for (RelNode input : join.getInputs()) {
    final RelOptPredicateList inputPreds = mq.getAllPredicates(input);
    if (inputPreds == null) {
      // Bail out
      return null;
    }
    // Gather table references
    final Set<RelTableRef> tableRefs = mq.getTableReferences(input);
    if (input == join.getLeft()) {
      // Left input references remain unchanged
      for (RelTableRef leftRef : tableRefs) {
        qualifiedNamesToRefs.put(leftRef.getQualifiedName(), leftRef);
      }
      newPreds = newPreds.union(rexBuilder, inputPreds);
    } else {
      // Right input references might need to be updated if there are table name
      // clashes with left input
      final Map<RelTableRef, RelTableRef> currentTablesMapping = new HashMap<>();
      for (RelTableRef rightRef : tableRefs) {
        int shift = 0;
        Collection<RelTableRef> lRefs = qualifiedNamesToRefs.get(
            rightRef.getQualifiedName());
        if (lRefs != null) {
          shift = lRefs.size();
        }
        currentTablesMapping.put(rightRef,
            RelTableRef.of(rightRef.getTable(), shift + rightRef.getEntityNumber()));
      }
      final List<RexNode> updatedPreds = Lists.newArrayList(
          Iterables.transform(inputPreds.pulledUpPredicates,
              e -> RexUtil.swapTableReferences(rexBuilder, e,
                  currentTablesMapping)));
      newPreds = newPreds.union(rexBuilder,
          RelOptPredicateList.of(rexBuilder, updatedPreds));
    }
  }

  // Extract input fields referenced by Join condition
  final Set<RelDataTypeField> inputExtraFields = new LinkedHashSet<>();
  final RelOptUtil.InputFinder inputFinder = new RelOptUtil.InputFinder(inputExtraFields);
  pred.accept(inputFinder);
  final ImmutableBitSet inputFieldsUsed = inputFinder.inputBitSet.build();

  // Infer column origin expressions for given references
  final Map<RexInputRef, Set<RexNode>> mapping = new LinkedHashMap<>();
  for (int idx : inputFieldsUsed) {
    final RexInputRef inputRef = RexInputRef.of(idx, join.getRowType().getFieldList());
    final Set<RexNode> originalExprs = mq.getExpressionLineage(join, inputRef);
    if (originalExprs == null) {
      // Bail out
      return null;
    }
    final RexInputRef ref = RexInputRef.of(idx, join.getRowType().getFieldList());
    mapping.put(ref, originalExprs);
  }

  // Replace with new expressions and return union of predicates
  final Set<RexNode> allExprs =
      RelMdExpressionLineage.createAllPossibleExpressions(rexBuilder, pred, mapping);
  if (allExprs == null) {
    return null;
  }
  return newPreds.union(rexBuilder, RelOptPredicateList.of(rexBuilder, allExprs));
}
 

/**
 * Given an {@link org.apache.calcite.rel.core.Aggregate}
 * and the ordinals of the arguments to a
 * particular call to an aggregate function, creates a 'select distinct'
 * relational expression which projects the group columns and those
 * arguments but nothing else.
 *
 * <p>For example, given
 *
 * <blockquote>
 * <pre>select f0, count(distinct f1), count(distinct f2)
 * from t group by f0</pre>
 * </blockquote>
 *
 * <p>and the argument list
 *
 * <blockquote>{2}</blockquote>
 *
 * <p>returns
 *
 * <blockquote>
 * <pre>select distinct f0, f2 from t</pre>
 * </blockquote>
 *
 * <p>The <code>sourceOf</code> map is populated with the source of each
 * column; in this case sourceOf.get(0) = 0, and sourceOf.get(1) = 2.
 *
 * @param relBuilder Relational expression builder
 * @param aggregate Aggregate relational expression
 * @param argList   Ordinals of columns to make distinct
 * @param filterArg Ordinal of column to filter on, or -1
 * @param sourceOf  Out parameter, is populated with a map of where each
 *                  output field came from
 * @return Aggregate relational expression which projects the required
 * columns
 */
private RelBuilder createSelectDistinct(RelBuilder relBuilder,
		Aggregate aggregate, List<Integer> argList, int filterArg,
		Map<Integer, Integer> sourceOf) {
	relBuilder.push(aggregate.getInput());
	final List<Pair<RexNode, String>> projects = new ArrayList<>();
	final List<RelDataTypeField> childFields =
			relBuilder.peek().getRowType().getFieldList();
	for (int i : aggregate.getGroupSet()) {
		sourceOf.put(i, projects.size());
		projects.add(RexInputRef.of2(i, childFields));
	}
	if (filterArg >= 0) {
		sourceOf.put(filterArg, projects.size());
		projects.add(RexInputRef.of2(filterArg, childFields));
	}
	for (Integer arg : argList) {
		if (filterArg >= 0) {
			// Implement
			//   agg(DISTINCT arg) FILTER $f
			// by generating
			//   SELECT DISTINCT ... CASE WHEN $f THEN arg ELSE NULL END AS arg
			// and then applying
			//   agg(arg)
			// as usual.
			//
			// It works except for (rare) agg functions that need to see null
			// values.
			final RexBuilder rexBuilder = aggregate.getCluster().getRexBuilder();
			final RexInputRef filterRef = RexInputRef.of(filterArg, childFields);
			final Pair<RexNode, String> argRef = RexInputRef.of2(arg, childFields);
			RexNode condition =
					rexBuilder.makeCall(SqlStdOperatorTable.CASE, filterRef,
							argRef.left,
							rexBuilder.makeNullLiteral(argRef.left.getType()));
			sourceOf.put(arg, projects.size());
			projects.add(Pair.of(condition, "i$" + argRef.right));
			continue;
		}
		if (sourceOf.get(arg) != null) {
			continue;
		}
		sourceOf.put(arg, projects.size());
		projects.add(RexInputRef.of2(arg, childFields));
	}
	relBuilder.project(Pair.left(projects), Pair.right(projects));

	// Get the distinct values of the GROUP BY fields and the arguments
	// to the agg functions.
	relBuilder.push(
			aggregate.copy(aggregate.getTraitSet(), relBuilder.build(),
					ImmutableBitSet.range(projects.size()),
					null, com.google.common.collect.ImmutableList.<AggregateCall>of()));
	return relBuilder;
}
 

/**
 * Computes the project group set.
 *
 * @param projectNode Project under the Aggregates if any
 * @param groupSet Ids of grouping keys as they are listed in {@code projects} list
 * @param inputRowType Input row type under the project
 * @param druidQuery Druid Query
 *
 * @return A list of {@link DimensionSpec} containing the group by dimensions,
 * and a list of {@link VirtualColumn} containing Druid virtual column
 * projections; or null, if translation is not possible.
 * Note that the size of lists can be different.
 */
@Nullable
protected static Pair<List<DimensionSpec>, List<VirtualColumn>> computeProjectGroupSet(
    @Nullable Project projectNode, ImmutableBitSet groupSet,
    RelDataType inputRowType, DruidQuery druidQuery) {
  final List<DimensionSpec> dimensionSpecList = new ArrayList<>();
  final List<VirtualColumn> virtualColumnList = new ArrayList<>();
  final Set<String> usedFieldNames = new HashSet<>();
  for (int groupKey : groupSet) {
    final DimensionSpec dimensionSpec;
    final RexNode project;
    if (projectNode == null) {
      project =  RexInputRef.of(groupKey, inputRowType);
    } else {
      project = projectNode.getProjects().get(groupKey);
    }

    Pair<String, ExtractionFunction> druidColumn =
        toDruidColumn(project, inputRowType, druidQuery);
    if (druidColumn.left != null && druidColumn.right == null) {
      // SIMPLE INPUT REF
      dimensionSpec = new DefaultDimensionSpec(druidColumn.left, druidColumn.left,
          DruidExpressions.EXPRESSION_TYPES.get(project.getType().getSqlTypeName()));
      usedFieldNames.add(druidColumn.left);
    } else if (druidColumn.left != null && druidColumn.right != null) {
     // CASE it is an extraction Dimension
      final String columnPrefix;
      //@TODO Remove it! if else statement is not really needed it is here to make tests pass.
      if (project.getKind() == SqlKind.EXTRACT) {
        columnPrefix =
            EXTRACT_COLUMN_NAME_PREFIX + "_" + Objects
                .requireNonNull(DruidDateTimeUtils
                    .extractGranularity(project, druidQuery.getConnectionConfig().timeZone())
                    .getType().lowerName);
      } else if (project.getKind() == SqlKind.FLOOR) {
        columnPrefix =
            FLOOR_COLUMN_NAME_PREFIX + "_" + Objects
                .requireNonNull(DruidDateTimeUtils
                    .extractGranularity(project, druidQuery.getConnectionConfig().timeZone())
                    .getType().lowerName);
      } else {
        columnPrefix = "extract";
      }
      final String uniqueExtractColumnName = SqlValidatorUtil
          .uniquify(columnPrefix, usedFieldNames,
              SqlValidatorUtil.EXPR_SUGGESTER);
      dimensionSpec = new ExtractionDimensionSpec(druidColumn.left,
          druidColumn.right, uniqueExtractColumnName);
      usedFieldNames.add(uniqueExtractColumnName);
    } else {
      // CASE it is Expression
      final String expression = DruidExpressions
          .toDruidExpression(project, inputRowType, druidQuery);
      if (Strings.isNullOrEmpty(expression)) {
        return null;
      }
      final String name = SqlValidatorUtil
          .uniquify("vc", usedFieldNames,
              SqlValidatorUtil.EXPR_SUGGESTER);
      VirtualColumn vc = new VirtualColumn(name, expression,
          DruidExpressions.EXPRESSION_TYPES.get(project.getType().getSqlTypeName()));
      virtualColumnList.add(vc);
      dimensionSpec = new DefaultDimensionSpec(name, name,
          DruidExpressions.EXPRESSION_TYPES.get(project.getType().getSqlTypeName()));
      usedFieldNames.add(name);

    }

    dimensionSpecList.add(dimensionSpec);
  }
  return Pair.of(dimensionSpecList, virtualColumnList);
}
 

/**
 * Expression lineage from {@link Union}.
 *
 * <p>For Union operator, we might be able to extract multiple origins for the
 * references in the given expression.
 */
public Set<RexNode> getExpressionLineage(Union rel, RelMetadataQuery mq,
    RexNode outputExpression) {
  final RexBuilder rexBuilder = rel.getCluster().getRexBuilder();

  // Extract input fields referenced by expression
  final ImmutableBitSet inputFieldsUsed = extractInputRefs(outputExpression);

  // Infer column origin expressions for given references
  final Multimap<List<String>, RelTableRef> qualifiedNamesToRefs = HashMultimap.create();
  final Map<RexInputRef, Set<RexNode>> mapping = new LinkedHashMap<>();
  for (RelNode input : rel.getInputs()) {
    // Gather table references
    final Map<RelTableRef, RelTableRef> currentTablesMapping = new HashMap<>();
    final Set<RelTableRef> tableRefs = mq.getTableReferences(input);
    if (tableRefs == null) {
      // Bail out
      return null;
    }
    for (RelTableRef tableRef : tableRefs) {
      int shift = 0;
      Collection<RelTableRef> lRefs = qualifiedNamesToRefs.get(
          tableRef.getQualifiedName());
      if (lRefs != null) {
        shift = lRefs.size();
      }
      currentTablesMapping.put(tableRef,
          RelTableRef.of(tableRef.getTable(), shift + tableRef.getEntityNumber()));
    }
    // Map references
    for (int idx : inputFieldsUsed) {
      final RexInputRef inputRef = RexInputRef.of(idx, input.getRowType().getFieldList());
      final Set<RexNode> originalExprs = mq.getExpressionLineage(input, inputRef);
      if (originalExprs == null) {
        // Bail out
        return null;
      }
      // References might need to be updated
      final RexInputRef ref = RexInputRef.of(idx, rel.getRowType().getFieldList());
      final Set<RexNode> updatedExprs =
          originalExprs.stream()
              .map(e ->
                  RexUtil.swapTableReferences(rexBuilder, e,
                      currentTablesMapping))
              .collect(Collectors.toSet());
      final Set<RexNode> set = mapping.get(ref);
      if (set != null) {
        set.addAll(updatedExprs);
      } else {
        mapping.put(ref, updatedExprs);
      }
    }
    // Add to existing qualified names
    for (RelTableRef newRef : currentTablesMapping.values()) {
      qualifiedNamesToRefs.put(newRef.getQualifiedName(), newRef);
    }
  }

  // Return result
  return createAllPossibleExpressions(rexBuilder, outputExpression, mapping);
}