Java Code Examples for org.apache.calcite.rel.RelNode#copy()

/** Helper for {@link #replace}. */
private static RelNode replaceRecurse(
    RelNode query, RelNode find, RelNode replace) {
  if (query == find) {
    return replace;
  }
  final List<RelNode> inputs = query.getInputs();
  if (!inputs.isEmpty()) {
    final List<RelNode> newInputs = new ArrayList<>();
    for (RelNode input : inputs) {
      newInputs.add(replaceRecurse(input, find, replace));
    }
    if (!newInputs.equals(inputs)) {
      return query.copy(query.getTraitSet(), newInputs);
    }
  }
  return query;
}
 

/**
 * Visits a particular child of a parent.
 */
protected RelNode visitChild(RelNode parent, int i, RelNode child) {
  appendPath.add(i);
  try {
    RelNode child2 = child.accept(this);
    if (child2 != child) {
      final List<RelNode> newInputs = new ArrayList<>(parent.getInputs());
      newInputs.set(i, child2);
      return parent.copy(parent.getTraitSet(), newInputs);
    }
    return parent;
  } finally {
    // Remove the last element.
    appendPath.remove(appendPath.size() - 1);
  }
}
 

/** Helper for {@link #replace}. */
private static RelNode replaceRecurse(RelNode query, RelNode find, RelNode replace) {
    if (query == find) {
        return replace;
    }
    final List<RelNode> inputs = query.getInputs();
    if (!inputs.isEmpty()) {
        final List<RelNode> newInputs = new ArrayList<>();
        for (RelNode input : inputs) {
            newInputs.add(replaceRecurse(input, find, replace));
        }
        if (!newInputs.equals(inputs)) {
            return query.copy(query.getTraitSet(), newInputs);
        }
    }
    return query;
}
 

/** Fallback if none of the other {@code decorrelateRel} methods match. */
public Frame decorrelateRel(RelNode rel) {
	RelNode newRel = rel.copy(rel.getTraitSet(), rel.getInputs());
	if (rel.getInputs().size() > 0) {
		List<RelNode> oldInputs = rel.getInputs();
		List<RelNode> newInputs = new ArrayList<>();
		for (int i = 0; i < oldInputs.size(); ++i) {
			final Frame frame = getInvoke(oldInputs.get(i));
			if (frame == null || frame.c != null) {
				// if input is not rewritten, or if it produces correlated variables, terminate rewrite
				return null;
			}
			newInputs.add(frame.r);
			newRel.replaceInput(i, frame.r);
		}

		if (!Util.equalShallow(oldInputs, newInputs)) {
			newRel = rel.copy(rel.getTraitSet(), newInputs);
		}
	}
	// the output position should not change since there are no corVars coming from below.
	return new Frame(rel, newRel, null, identityMap(rel.getRowType().getFieldCount()));
}
 

@Override
protected RelNode visitChild(RelNode parent, int i, RelNode child) {
  RelNode child2 = child.accept(this);
  if (child2 != child) {
    final List<RelNode> newInputs = new ArrayList<>(parent.getInputs());
    newInputs.set(i, child2);
    return parent.copy(parent.getTraitSet(), newInputs);
  }
  return parent;
}
 

/**
 * Returns a shallow copy of a relational expression with a particular
 * input replaced.
 */
public static RelNode replaceInput(
    RelNode parent, int ordinal, RelNode newInput) {
  final List<RelNode> inputs = new ArrayList<>(parent.getInputs());
  if (inputs.get(ordinal) == newInput) {
    return parent;
  }
  inputs.set(ordinal, newInput);
  return parent.copy(parent.getTraitSet(), inputs);
}
 

/** Fallback if none of the other {@code decorrelateRel} methods match. */
public Frame decorrelateRel(RelNode rel) {
  RelNode newRel = rel.copy(rel.getTraitSet(), rel.getInputs());

  if (rel.getInputs().size() > 0) {
    List<RelNode> oldInputs = rel.getInputs();
    List<RelNode> newInputs = new ArrayList<>();
    for (int i = 0; i < oldInputs.size(); ++i) {
      final Frame frame = getInvoke(oldInputs.get(i), rel);
      if (frame == null || !frame.corDefOutputs.isEmpty()) {
        // if input is not rewritten, or if it produces correlated
        // variables, terminate rewrite
        return null;
      }
      newInputs.add(frame.r);
      newRel.replaceInput(i, frame.r);
    }

    if (!Util.equalShallow(oldInputs, newInputs)) {
      newRel = rel.copy(rel.getTraitSet(), newInputs);
    }
  }

  // the output position should not change since there are no corVars
  // coming from below.
  return register(rel, newRel, identityMap(rel.getRowType().getFieldCount()),
      ImmutableSortedMap.of());
}
 

public void onMatch(RelOptRuleCall call) {
  SingleRel singleRel = call.rel(0);
  RelNode emptyValues = call.builder().push(singleRel).empty().build();
  RelTraitSet traits = singleRel.getTraitSet();
  // propagate all traits (except convention) from the original singleRel into the empty values
  if (emptyValues.getConvention() != null) {
    traits = traits.replace(emptyValues.getConvention());
  }
  emptyValues = emptyValues.copy(traits, Collections.emptyList());
  call.transformTo(emptyValues);
}
 

@Override
public RelNode visit(RelNode other) {
  if(!data.add(other)) {
    if (other instanceof LogicalTableScan) {
      // LogicalTableScan does not have implementation of a deep copy. Create a new instance.
      other = LogicalTableScan.create(other.getCluster(), other.getTable());
    } else {
      other = other.copy(other.getTraitSet(), other.getInputs());
    }
  }

  return super.visit(other);
}
 

/** Fallback if none of the other {@code decorrelateRel} methods match. */
public Frame decorrelateRel(RelNode rel) {
  RelNode newRel = rel.copy(rel.getTraitSet(), rel.getInputs());

  if (rel.getInputs().size() > 0) {
    List<RelNode> oldInputs = rel.getInputs();
    List<RelNode> newInputs = new ArrayList<>();
    for (int i = 0; i < oldInputs.size(); ++i) {
      final Frame frame = getInvoke(oldInputs.get(i), rel);
      if (frame == null || !frame.corDefOutputs.isEmpty()) {
        // if input is not rewritten, or if it produces correlated
        // variables, terminate rewrite
        return null;
      }
      newInputs.add(frame.r);
      newRel.replaceInput(i, frame.r);
    }

    if (!Util.equalShallow(oldInputs, newInputs)) {
      newRel = rel.copy(rel.getTraitSet(), newInputs);
    }
  }

  // the output position should not change since there are no corVars
  // coming from below.
  return register(rel, newRel, identityMap(rel.getRowType().getFieldCount()),
      ImmutableSortedMap.of());
}
 

/** Fallback if none of the other {@code decorrelateRel} methods match. */
public Frame decorrelateRel(RelNode rel) {
  RelNode newRel = rel.copy(rel.getTraitSet(), rel.getInputs());

  if (rel.getInputs().size() > 0) {
    List<RelNode> oldInputs = rel.getInputs();
    List<RelNode> newInputs = new ArrayList<>();
    for (int i = 0; i < oldInputs.size(); ++i) {
      final Frame frame = getInvoke(oldInputs.get(i), rel);
      if (frame == null || !frame.corDefOutputs.isEmpty()) {
        // if input is not rewritten, or if it produces correlated
        // variables, terminate rewrite
        return null;
      }
      newInputs.add(frame.r);
      newRel.replaceInput(i, frame.r);
    }

    if (!Util.equalShallow(oldInputs, newInputs)) {
      newRel = rel.copy(rel.getTraitSet(), newInputs);
    }
  }

  // the output position should not change since there are no corVars
  // coming from below.
  return register(rel, newRel, identityMap(rel.getRowType().getFieldCount()),
      ImmutableSortedMap.of());
}
 

/**
 * Visits a particular child of a parent.
 */
protected RelNode visitChild(RelNode parent, int i, RelNode child) {
  inheritPaths.forEach(inheritPath -> inheritPath.right.push(i));
  try {
    RelNode child2 = child.accept(this);
    if (child2 != child) {
      final List<RelNode> newInputs = new ArrayList<>(parent.getInputs());
      newInputs.set(i, child2);
      return parent.copy(parent.getTraitSet(), newInputs);
    }
    return parent;
  } finally {
    inheritPaths.forEach(inheritPath -> inheritPath.right.pop());
  }
}
 

public void rewriteGeneric(RelNode rel) {
    RelNode newRel = rel.copy(rel.getTraitSet(), rel.getInputs());
    List<RelNode> oldInputs = rel.getInputs();
    for (int i = 0; i < oldInputs.size(); ++i) {
        newRel.replaceInput(i, getNewForOldRel(oldInputs.get(i)));
    }
    setNewForOldRel(rel, newRel);
}
 

/** Fallback if none of the other {@code decorrelateRel} methods match. */
public Frame decorrelateRel(RelNode rel) {
    RelNode newRel = rel.copy(rel.getTraitSet(), rel.getInputs());

    if (rel.getInputs().size() > 0) {
        List<RelNode> oldInputs = rel.getInputs();
        List<RelNode> newInputs = new ArrayList<>();
        for (int i = 0; i < oldInputs.size(); ++i) {
            final Frame frame = getInvoke(oldInputs.get(i), rel);
            if (frame == null || !frame.corDefOutputs.isEmpty()) {
                // if input is not rewritten, or if it produces correlated
                // variables, terminate rewrite
                return null;
            }
            newInputs.add(frame.r);
            newRel.replaceInput(i, frame.r);
        }

        if (!Util.equalShallow(oldInputs, newInputs)) {
            newRel = rel.copy(rel.getTraitSet(), newInputs);
        }
    }

    // the output position should not change since there are no corVars
    // coming from below.
    return register(rel, newRel, identityMap(rel.getRowType().getFieldCount()), ImmutableSortedMap.of());
}
 

@Override
protected RelNode rewriteQuery(RelBuilder relBuilder, RexBuilder rexBuilder, RexSimplify simplify,
        RelMetadataQuery mq, RexNode compensationColumnsEquiPred, RexNode otherCompensationPred,
        Project topProject, RelNode node, BiMap<RelTableRef, RelTableRef> viewToQueryTableMapping,
        EquivalenceClasses viewEC, EquivalenceClasses queryEC) {
    // Our target node is the node below the root, which should have the maximum
    // number of available expressions in the tree in order to maximize our
    // number of rewritings.
    // We create a project on top. If the program is available, we execute
    // it to maximize rewriting opportunities. For instance, a program might
    // pull up all the expressions that are below the aggregate so we can
    // introduce compensation filters easily. This is important depending on
    // the planner strategy.
    RelNode newNode = node;
    RelNode target = node;
    if (unionRewritingPullProgram != null) {
        final HepPlanner tmpPlanner = new HepPlanner(unionRewritingPullProgram);
        tmpPlanner.setRoot(newNode);
        newNode = tmpPlanner.findBestExp();
        target = newNode.getInput(0);
    }

    // All columns required by compensating predicates must be contained
    // in the query.
    List<RexNode> queryExprs = extractReferences(rexBuilder, target);

    if (!compensationColumnsEquiPred.isAlwaysTrue()) {
        compensationColumnsEquiPred = rewriteExpression(rexBuilder, mq, target, target, queryExprs,
                viewToQueryTableMapping.inverse(), queryEC, false, compensationColumnsEquiPred);
        if (compensationColumnsEquiPred == null) {
            // Skip it
            return null;
        }
    }
    // For the rest, we use the query equivalence classes
    if (!otherCompensationPred.isAlwaysTrue()) {
        otherCompensationPred = rewriteExpression(rexBuilder, mq, target, target, queryExprs,
                viewToQueryTableMapping.inverse(), viewEC, true, otherCompensationPred);
        if (otherCompensationPred == null) {
            // Skip it
            return null;
        }
    }
    final RexNode queryCompensationPred = RexUtil.not(RexUtil.composeConjunction(rexBuilder,
            ImmutableList.of(compensationColumnsEquiPred, otherCompensationPred)));

    // Generate query rewriting.
    RelNode rewrittenPlan = relBuilder.push(target)
            .filter(simplify.simplifyUnknownAsFalse(queryCompensationPred)).build();
    if (unionRewritingPullProgram != null) {
        rewrittenPlan = newNode.copy(newNode.getTraitSet(), ImmutableList.of(rewrittenPlan));
    }
    if (topProject != null) {
        return topProject.copy(topProject.getTraitSet(), ImmutableList.of(rewrittenPlan));
    }
    return rewrittenPlan;
}
 

private HepRelVertex addRelToGraph(
    RelNode rel) {
  // Check if a transformation already produced a reference
  // to an existing vertex.
  if (graph.vertexSet().contains(rel)) {
    return (HepRelVertex) rel;
  }

  // Recursively add children, replacing this rel's inputs
  // with corresponding child vertices.
  final List<RelNode> inputs = rel.getInputs();
  final List<RelNode> newInputs = new ArrayList<>();
  for (RelNode input1 : inputs) {
    HepRelVertex childVertex = addRelToGraph(input1);
    newInputs.add(childVertex);
  }

  if (!Util.equalShallow(inputs, newInputs)) {
    RelNode oldRel = rel;
    rel = rel.copy(rel.getTraitSet(), newInputs);
    onCopy(oldRel, rel);
  }
  // Compute digest first time we add to DAG,
  // otherwise can't get equivVertex for common sub-expression
  rel.recomputeDigest();

  // try to find equivalent rel only if DAG is allowed
  if (!noDag) {
    // Now, check if an equivalent vertex already exists in graph.
    HepRelVertex equivVertex = mapDigestToVertex.get(rel.getRelDigest());
    if (equivVertex != null) {
      // Use existing vertex.
      return equivVertex;
    }
  }

  // No equivalence:  create a new vertex to represent this rel.
  HepRelVertex newVertex = new HepRelVertex(rel);
  graph.addVertex(newVertex);
  updateVertex(newVertex, rel);

  for (RelNode input : rel.getInputs()) {
    graph.addEdge(newVertex, (HepRelVertex) input);
  }

  nTransformations++;
  return newVertex;
}
 

@Override protected RelNode rewriteQuery(
    RelBuilder relBuilder,
    RexBuilder rexBuilder,
    RexSimplify simplify,
    RelMetadataQuery mq,
    RexNode compensationColumnsEquiPred,
    RexNode otherCompensationPred,
    Project topProject,
    RelNode node,
    BiMap<RelTableRef, RelTableRef> viewToQueryTableMapping,
    EquivalenceClasses viewEC, EquivalenceClasses queryEC) {
  // Our target node is the node below the root, which should have the maximum
  // number of available expressions in the tree in order to maximize our
  // number of rewritings.
  // We create a project on top. If the program is available, we execute
  // it to maximize rewriting opportunities. For instance, a program might
  // pull up all the expressions that are below the aggregate so we can
  // introduce compensation filters easily. This is important depending on
  // the planner strategy.
  RelNode newNode = node;
  RelNode target = node;
  if (unionRewritingPullProgram != null) {
    final HepPlanner tmpPlanner = new HepPlanner(unionRewritingPullProgram);
    tmpPlanner.setRoot(newNode);
    newNode = tmpPlanner.findBestExp();
    target = newNode.getInput(0);
  }

  // All columns required by compensating predicates must be contained
  // in the query.
  List<RexNode> queryExprs = extractReferences(rexBuilder, target);

  if (!compensationColumnsEquiPred.isAlwaysTrue()) {
    compensationColumnsEquiPred = rewriteExpression(rexBuilder, mq,
        target, target, queryExprs, viewToQueryTableMapping.inverse(), queryEC, false,
        compensationColumnsEquiPred);
    if (compensationColumnsEquiPred == null) {
      // Skip it
      return null;
    }
  }
  // For the rest, we use the query equivalence classes
  if (!otherCompensationPred.isAlwaysTrue()) {
    otherCompensationPred = rewriteExpression(rexBuilder, mq,
        target, target, queryExprs, viewToQueryTableMapping.inverse(), viewEC, true,
        otherCompensationPred);
    if (otherCompensationPred == null) {
      // Skip it
      return null;
    }
  }
  final RexNode queryCompensationPred = RexUtil.not(
      RexUtil.composeConjunction(rexBuilder,
          ImmutableList.of(compensationColumnsEquiPred,
              otherCompensationPred)));

  // Generate query rewriting.
  RelNode rewrittenPlan = relBuilder
      .push(target)
      .filter(simplify.simplifyUnknownAsFalse(queryCompensationPred))
      .build();
  if (unionRewritingPullProgram != null) {
    rewrittenPlan = newNode.copy(
        newNode.getTraitSet(), ImmutableList.of(rewrittenPlan));
  }
  if (topProject != null) {
    return topProject.copy(topProject.getTraitSet(), ImmutableList.of(rewrittenPlan));
  }
  return rewrittenPlan;
}
 

@Override protected RelNode createUnion(RelBuilder relBuilder, RexBuilder rexBuilder,
    RelNode topProject, RelNode unionInputQuery, RelNode unionInputView) {
  // Union
  relBuilder.push(unionInputQuery);
  relBuilder.push(unionInputView);
  relBuilder.union(true);
  List<RexNode> exprList = new ArrayList<>(relBuilder.peek().getRowType().getFieldCount());
  List<String> nameList = new ArrayList<>(relBuilder.peek().getRowType().getFieldCount());
  for (int i = 0; i < relBuilder.peek().getRowType().getFieldCount(); i++) {
    // We can take unionInputQuery as it is query based.
    RelDataTypeField field = unionInputQuery.getRowType().getFieldList().get(i);
    exprList.add(
        rexBuilder.ensureType(
            field.getType(),
            rexBuilder.makeInputRef(relBuilder.peek(), i),
            true));
    nameList.add(field.getName());
  }
  relBuilder.project(exprList, nameList);
  // Rollup aggregate
  Aggregate aggregate = (Aggregate) unionInputQuery;
  final ImmutableBitSet groupSet = ImmutableBitSet.range(aggregate.getGroupCount());
  final List<AggCall> aggregateCalls = new ArrayList<>();
  for (int i = 0; i < aggregate.getAggCallList().size(); i++) {
    AggregateCall aggCall = aggregate.getAggCallList().get(i);
    if (aggCall.isDistinct()) {
      // Cannot ROLLUP distinct
      return null;
    }
    SqlAggFunction rollupAgg =
        getRollup(aggCall.getAggregation());
    if (rollupAgg == null) {
      // Cannot rollup this aggregate, bail out
      return null;
    }
    final RexInputRef operand =
        rexBuilder.makeInputRef(relBuilder.peek(),
            aggregate.getGroupCount() + i);
    aggregateCalls.add(
        relBuilder.aggregateCall(rollupAgg, operand)
            .distinct(aggCall.isDistinct())
            .approximate(aggCall.isApproximate())
            .as(aggCall.name));
  }
  RelNode prevNode = relBuilder.peek();
  RelNode result = relBuilder
      .aggregate(relBuilder.groupKey(groupSet), aggregateCalls)
      .build();
  if (prevNode == result && groupSet.cardinality() != result.getRowType().getFieldCount()) {
    // Aggregate was not inserted but we need to prune columns
    result = relBuilder
        .push(result)
        .project(relBuilder.fields(groupSet))
        .build();
  }
  if (topProject != null) {
    // Top project
    return topProject.copy(topProject.getTraitSet(), ImmutableList.of(result));
  }
  // Result
  return result;
}
 

private HepRelVertex addRelToGraph(
    RelNode rel) {
  // Check if a transformation already produced a reference
  // to an existing vertex.
  if (graph.vertexSet().contains(rel)) {
    return (HepRelVertex) rel;
  }

  // Recursively add children, replacing this rel's inputs
  // with corresponding child vertices.
  final List<RelNode> inputs = rel.getInputs();
  final List<RelNode> newInputs = new ArrayList<>();
  for (RelNode input1 : inputs) {
    HepRelVertex childVertex = addRelToGraph(input1);
    newInputs.add(childVertex);
  }

  if (!Util.equalShallow(inputs, newInputs)) {
    RelNode oldRel = rel;
    rel = rel.copy(rel.getTraitSet(), newInputs);
    onCopy(oldRel, rel);
  }
  // Compute digest first time we add to DAG,
  // otherwise can't get equivVertex for common sub-expression
  rel.recomputeDigest();

  // try to find equivalent rel only if DAG is allowed
  if (!noDag) {
    // Now, check if an equivalent vertex already exists in graph.
    String digest = rel.getDigest();
    HepRelVertex equivVertex = mapDigestToVertex.get(digest);
    if (equivVertex != null) {
      // Use existing vertex.
      return equivVertex;
    }
  }

  // No equivalence:  create a new vertex to represent this rel.
  HepRelVertex newVertex = new HepRelVertex(rel);
  graph.addVertex(newVertex);
  updateVertex(newVertex, rel);

  for (RelNode input : rel.getInputs()) {
    graph.addEdge(newVertex, (HepRelVertex) input);
  }

  nTransformations++;
  return newVertex;
}
 

@Override
protected RelNode createUnion(RelBuilder relBuilder, RexBuilder rexBuilder, RelNode topProject,
        RelNode unionInputQuery, RelNode unionInputView) {
    // Union
    relBuilder.push(unionInputQuery);
    relBuilder.push(unionInputView);
    relBuilder.union(true);
    List<RexNode> exprList = new ArrayList<>(relBuilder.peek().getRowType().getFieldCount());
    List<String> nameList = new ArrayList<>(relBuilder.peek().getRowType().getFieldCount());
    for (int i = 0; i < relBuilder.peek().getRowType().getFieldCount(); i++) {
        // We can take unionInputQuery as it is query based.
        RelDataTypeField field = unionInputQuery.getRowType().getFieldList().get(i);
        exprList.add(
                rexBuilder.ensureType(field.getType(), rexBuilder.makeInputRef(relBuilder.peek(), i), true));
        nameList.add(field.getName());
    }
    relBuilder.project(exprList, nameList);
    // Rollup aggregate
    Aggregate aggregate = (Aggregate) unionInputQuery;
    final ImmutableBitSet groupSet = ImmutableBitSet.range(aggregate.getGroupCount());
    final List<AggCall> aggregateCalls = new ArrayList<>();
    for (int i = 0; i < aggregate.getAggCallList().size(); i++) {
        AggregateCall aggCall = aggregate.getAggCallList().get(i);
        if (aggCall.isDistinct()) {
            // Cannot ROLLUP distinct
            return null;
        }
        SqlAggFunction rollupAgg = getRollup(aggCall.getAggregation());
        if (rollupAgg == null) {
            // Cannot rollup this aggregate, bail out
            return null;
        }
        final RexInputRef operand = rexBuilder.makeInputRef(relBuilder.peek(), aggregate.getGroupCount() + i);
        aggregateCalls.add(
                // TODO: handle aggregate ordering
                relBuilder.aggregateCall(rollupAgg, operand).distinct(aggCall.isDistinct())
                        .approximate(aggCall.isApproximate()).as(aggCall.name));
    }
    RelNode prevNode = relBuilder.peek();
    RelNode result = relBuilder.aggregate(relBuilder.groupKey(groupSet), aggregateCalls).build();
    if (prevNode == result && groupSet.cardinality() != result.getRowType().getFieldCount()) {
        // Aggregate was not inserted but we need to prune columns
        result = relBuilder.push(result).project(relBuilder.fields(groupSet.asList())).build();
    }
    if (topProject != null) {
        // Top project
        return topProject.copy(topProject.getTraitSet(), ImmutableList.of(result));
    }
    // Result
    return result;
}