Java Code Examples for org.apache.calcite.plan.RelOptRuleCall#getMetadataQuery()

@Override public void onMatch(RelOptRuleCall call) {
  final Project project = call.rel(0);
  final RelMetadataQuery mq = call.getMetadataQuery();
  final RelOptPredicateList predicates =
      mq.getPulledUpPredicates(project.getInput());
  final List<RexNode> expList =
      Lists.newArrayList(project.getProjects());
  if (reduceExpressions(project, expList, predicates, false,
      matchNullability)) {
    call.transformTo(
        call.builder()
            .push(project.getInput())
            .project(expList, project.getRowType().getFieldNames())
            .build());

    // New plan is absolutely better than old plan.
    call.getPlanner().setImportance(project, 0.0);
  }
}
 

@Override
public boolean matches(RelOptRuleCall call) {
	final Aggregate aggregate = call.rel(0);
	final RelNode input = call.rel(1);
	if (aggregate.getGroupCount() == 0 || aggregate.indicator ||
			aggregate.getGroupType() != Aggregate.Group.SIMPLE) {
		return false;
	}
	for (AggregateCall aggCall : aggregate.getAggCallList()) {
		SqlKind aggCallKind = aggCall.getAggregation().getKind();
		// TODO supports more AggregateCalls
		boolean isAllowAggCall = aggCallKind == SqlKind.SUM ||
				aggCallKind == SqlKind.MIN ||
				aggCallKind == SqlKind.MAX ||
				aggCall.getAggregation() instanceof SqlAuxiliaryGroupAggFunction;
		if (!isAllowAggCall || aggCall.filterArg >= 0 || aggCall.getArgList().size() != 1) {
			return false;
		}
	}

	final RelMetadataQuery mq = call.getMetadataQuery();
	return SqlFunctions.isTrue(mq.areColumnsUnique(input, aggregate.getGroupSet()));
}
 

public void onMatch(RelOptRuleCall call) {
  final Aggregate aggregate = call.rel(0);
  final RelNode input = call.rel(1);
  if (!aggregate.getAggCallList().isEmpty() || aggregate.indicator) {
    return;
  }
  final RelMetadataQuery mq = call.getMetadataQuery();
  if (!SqlFunctions.isTrue(mq.areColumnsUnique(input, aggregate.getGroupSet()))) {
    return;
  }
  // Distinct is "GROUP BY c1, c2" (where c1, c2 are a set of columns on
  // which the input is unique, i.e. contain a key) and has no aggregate
  // functions. It can be removed.
  final RelNode newInput = convert(input, aggregate.getTraitSet().simplify());

  // If aggregate was projecting a subset of columns, add a project for the
  // same effect.
  final RelBuilder relBuilder = call.builder();
  relBuilder.push(newInput);
  if (newInput.getRowType().getFieldCount()
      > aggregate.getRowType().getFieldCount()) {
    relBuilder.project(relBuilder.fields(aggregate.getGroupSet().asList()));
  }
  call.transformTo(relBuilder.build());
}
 

@Override
public boolean matches(RelOptRuleCall call) {
	final Aggregate aggregate = call.rel(0);
	final RelNode input = call.rel(1);
	if (aggregate.getGroupCount() == 0 || aggregate.indicator ||
			aggregate.getGroupType() != Aggregate.Group.SIMPLE) {
		return false;
	}
	for (AggregateCall aggCall : aggregate.getAggCallList()) {
		SqlKind aggCallKind = aggCall.getAggregation().getKind();
		// TODO supports more AggregateCalls
		boolean isAllowAggCall = aggCallKind == SqlKind.SUM ||
				aggCallKind == SqlKind.MIN ||
				aggCallKind == SqlKind.MAX ||
				aggCall.getAggregation() instanceof SqlAuxiliaryGroupAggFunction;
		if (!isAllowAggCall || aggCall.filterArg >= 0 || aggCall.getArgList().size() != 1) {
			return false;
		}
	}

	final RelMetadataQuery mq = call.getMetadataQuery();
	return SqlFunctions.isTrue(mq.areColumnsUnique(input, aggregate.getGroupSet()));
}
 

public void onMatch(RelOptRuleCall call) {
  final Sort sort = call.rel(0);
  final Union union = call.rel(1);
  List<RelNode> inputs = new ArrayList<>();
  // Thus we use 'ret' as a flag to identify if we have finished pushing the
  // sort past a union.
  boolean ret = true;
  final RelMetadataQuery mq = call.getMetadataQuery();
  for (RelNode input : union.getInputs()) {
    if (!RelMdUtil.checkInputForCollationAndLimit(mq, input,
        sort.getCollation(), sort.offset, sort.fetch)) {
      ret = false;
      Sort branchSort = sort.copy(sort.getTraitSet(), input,
          sort.getCollation(), sort.offset, sort.fetch);
      inputs.add(branchSort);
    } else {
      inputs.add(input);
    }
  }
  // there is nothing to change
  if (ret) {
    return;
  }
  // create new union and sort
  Union unionCopy = (Union) union
      .copy(union.getTraitSet(), inputs, union.all);
  Sort result = sort.copy(sort.getTraitSet(), unionCopy, sort.getCollation(),
      sort.offset, sort.fetch);
  call.transformTo(result);
}
 

public void onMatch(RelOptRuleCall call) {
  final MultiJoin multiJoinRel = call.rel(0);
  final LoptMultiJoin multiJoin = new LoptMultiJoin(multiJoinRel);
  final RelMetadataQuery mq = call.getMetadataQuery();

  findRemovableOuterJoins(mq, multiJoin);

  final RexBuilder rexBuilder = multiJoinRel.getCluster().getRexBuilder();
  final LoptSemiJoinOptimizer semiJoinOpt =
      new LoptSemiJoinOptimizer(call.getMetadataQuery(), multiJoin, rexBuilder);

  // determine all possible semijoins
  semiJoinOpt.makePossibleSemiJoins(multiJoin);

  // select the optimal join filters for semijoin filtering by
  // iteratively calling chooseBestSemiJoin; chooseBestSemiJoin will
  // apply semijoins in sort order, based on the cost of scanning each
  // factor; as it selects semijoins to apply and iterates through the
  // loop, the cost of scanning a factor will decrease in accordance
  // with the semijoins selected
  int iterations = 0;
  do {
    if (!semiJoinOpt.chooseBestSemiJoin(multiJoin)) {
      break;
    }
    if (iterations++ > 10) {
      break;
    }
  } while (true);

  multiJoin.setFactorWeights();

  findRemovableSelfJoins(mq, multiJoin);

  findBestOrderings(mq, call.builder(), multiJoin, semiJoinOpt, call);
}
 

@Override public void onMatch(RelOptRuleCall call) {
  final Sort sort = call.rel(0);
  final RelMetadataQuery mq = call.getMetadataQuery();
  final RelNode input = sort.getInput();
  final RelOptPredicateList predicates = mq.getPulledUpPredicates(input);
  if (predicates == null) {
    return;
  }

  final RexBuilder rexBuilder = sort.getCluster().getRexBuilder();
  final List<RelFieldCollation> collationsList =
      sort.getCollation().getFieldCollations().stream()
          .filter(fc ->
              !predicates.constantMap.containsKey(
                  rexBuilder.makeInputRef(input, fc.getFieldIndex())))
          .collect(Collectors.toList());

  if (collationsList.size() == sort.collation.getFieldCollations().size()) {
    return;
  }

  // No active collations. Remove the sort completely
  if (collationsList.isEmpty() && sort.offset == null && sort.fetch == null) {
    call.transformTo(input);
    call.getPlanner().prune(sort);
    return;
  }

  final Sort result =
      sort.copy(sort.getTraitSet(), input, RelCollations.of(collationsList));
  call.transformTo(result);
  call.getPlanner().prune(sort);
}
 

@Override public void onMatch(RelOptRuleCall call) {
  final Exchange exchange = call.rel(0);
  final RelMetadataQuery mq = call.getMetadataQuery();
  final RelNode input = exchange.getInput();
  final RelOptPredicateList predicates = mq.getPulledUpPredicates(input);
  if (predicates == null) {
    return;
  }

  final Set<Integer> constants = new HashSet<>();
  predicates.constantMap.keySet().forEach(key -> {
    if (key instanceof RexInputRef) {
      constants.add(((RexInputRef) key).getIndex());
    }
  });
  if (constants.isEmpty()) {
    return;
  }

  final List<Integer> distributionKeys = simplifyDistributionKeys(
      exchange.getDistribution(), constants);

  if (distributionKeys.size() != exchange.getDistribution().getKeys()
      .size()) {
    call.transformTo(call.builder()
        .push(exchange.getInput())
        .exchange(distributionKeys.isEmpty()
            ? RelDistributions.SINGLETON
            : RelDistributions.hash(distributionKeys))
        .build());
    call.getPlanner().prune(exchange);
  }
}
 

@Override public void onMatch(RelOptRuleCall call) {
  final Join join = call.rel(0);
  final List<RexNode> expList = Lists.newArrayList(join.getCondition());
  final int fieldCount = join.getLeft().getRowType().getFieldCount();
  final RelMetadataQuery mq = call.getMetadataQuery();
  final RelOptPredicateList leftPredicates =
      mq.getPulledUpPredicates(join.getLeft());
  final RelOptPredicateList rightPredicates =
      mq.getPulledUpPredicates(join.getRight());
  final RexBuilder rexBuilder = join.getCluster().getRexBuilder();
  final RelOptPredicateList predicates =
      leftPredicates.union(rexBuilder,
          rightPredicates.shift(rexBuilder, fieldCount));
  if (!reduceExpressions(join, expList, predicates, true,
      matchNullability)) {
    return;
  }
  call.transformTo(
      join.copy(
          join.getTraitSet(),
          expList.get(0),
          join.getLeft(),
          join.getRight(),
          join.getJoinType(),
          join.isSemiJoinDone()));

  // New plan is absolutely better than old plan.
  call.getPlanner().prune(join);
}
 

@Override public void onMatch(RelOptRuleCall call) {
  final Sort sort = call.rel(0);
  final RelMetadataQuery mq = call.getMetadataQuery();
  final RelNode input = sort.getInput();
  final RelOptPredicateList predicates = mq.getPulledUpPredicates(input);
  if (predicates == null) {
    return;
  }

  final RexBuilder rexBuilder = sort.getCluster().getRexBuilder();
  final List<RelFieldCollation> collationsList =
      sort.getCollation().getFieldCollations().stream()
          .filter(fc ->
              !predicates.constantMap.containsKey(
                  rexBuilder.makeInputRef(input, fc.getFieldIndex())))
          .collect(Collectors.toList());

  if (collationsList.size() == sort.collation.getFieldCollations().size()) {
    return;
  }

  // No active collations. Remove the sort completely
  if (collationsList.isEmpty() && sort.offset == null && sort.fetch == null) {
    call.transformTo(input);
    call.getPlanner().setImportance(sort, 0.0);
    return;
  }

  final Sort result =
      sort.copy(sort.getTraitSet(), input, RelCollations.of(collationsList));
  call.transformTo(result);
  call.getPlanner().setImportance(sort, 0.0);
}
 

public void onMatch(RelOptRuleCall call) {
  final Sort sort = call.rel(0);
  final Union union = call.rel(1);
  List<RelNode> inputs = new ArrayList<>();
  // Thus we use 'ret' as a flag to identify if we have finished pushing the
  // sort past a union.
  boolean ret = true;
  final RelMetadataQuery mq = call.getMetadataQuery();
  for (RelNode input : union.getInputs()) {
    if (!RelMdUtil.checkInputForCollationAndLimit(mq, input,
        sort.getCollation(), sort.offset, sort.fetch)) {
      ret = false;
      Sort branchSort = sort.copy(sort.getTraitSet(), input,
          sort.getCollation(), sort.offset, sort.fetch);
      inputs.add(branchSort);
    } else {
      inputs.add(input);
    }
  }
  // there is nothing to change
  if (ret) {
    return;
  }
  // create new union and sort
  Union unionCopy = (Union) union
      .copy(union.getTraitSet(), inputs, union.all);
  Sort result = sort.copy(sort.getTraitSet(), unionCopy, sort.getCollation(),
      sort.offset, sort.fetch);
  call.transformTo(result);
}
 

public void onMatch(RelOptRuleCall call) {
  final MultiJoin multiJoinRel = call.rel(0);
  final LoptMultiJoin multiJoin = new LoptMultiJoin(multiJoinRel);
  final RelMetadataQuery mq = call.getMetadataQuery();

  findRemovableOuterJoins(mq, multiJoin);

  final RexBuilder rexBuilder = multiJoinRel.getCluster().getRexBuilder();
  final LoptSemiJoinOptimizer semiJoinOpt =
      new LoptSemiJoinOptimizer(call.getMetadataQuery(), multiJoin, rexBuilder);

  // determine all possible semijoins
  semiJoinOpt.makePossibleSemiJoins(multiJoin);

  // select the optimal join filters for semijoin filtering by
  // iteratively calling chooseBestSemiJoin; chooseBestSemiJoin will
  // apply semijoins in sort order, based on the cost of scanning each
  // factor; as it selects semijoins to apply and iterates through the
  // loop, the cost of scanning a factor will decrease in accordance
  // with the semijoins selected
  int iterations = 0;
  do {
    if (!semiJoinOpt.chooseBestSemiJoin(multiJoin)) {
      break;
    }
    if (iterations++ > 10) {
      break;
    }
  } while (true);

  multiJoin.setFactorWeights();

  findRemovableSelfJoins(mq, multiJoin);

  findBestOrderings(mq, call.builder(), multiJoin, semiJoinOpt, call);
}
 

@Override public void onMatch(RelOptRuleCall call) {
  final Sort sort = call.rel(0);
  final Join join = call.rel(1);

  // We create a new sort operator on the corresponding input
  final RelNode newLeftInput;
  final RelNode newRightInput;
  final RelMetadataQuery mq = call.getMetadataQuery();
  if (join.getJoinType() == JoinRelType.LEFT) {
    // If the input is already sorted and we are not reducing the number of tuples,
    // we bail out
    if (RelMdUtil.checkInputForCollationAndLimit(mq, join.getLeft(),
        sort.getCollation(), sort.offset, sort.fetch)) {
      return;
    }
    newLeftInput = sort.copy(sort.getTraitSet(), join.getLeft(), sort.getCollation(),
        sort.offset, sort.fetch);
    newRightInput = join.getRight();
  } else {
    final RelCollation rightCollation =
        RelCollationTraitDef.INSTANCE.canonize(
            RelCollations.shift(sort.getCollation(),
                -join.getLeft().getRowType().getFieldCount()));
    // If the input is already sorted and we are not reducing the number of tuples,
    // we bail out
    if (RelMdUtil.checkInputForCollationAndLimit(mq, join.getRight(),
        rightCollation, sort.offset, sort.fetch)) {
      return;
    }
    newLeftInput = join.getLeft();
    newRightInput = sort.copy(sort.getTraitSet().replace(rightCollation),
        join.getRight(), rightCollation, sort.offset, sort.fetch);
  }
  // We copy the join and the top sort operator
  final RelNode joinCopy = join.copy(join.getTraitSet(), join.getCondition(), newLeftInput,
      newRightInput, join.getJoinType(), join.isSemiJoinDone());
  final RelNode sortCopy = sort.copy(sort.getTraitSet(), joinCopy, sort.getCollation(),
      sort.offset, sort.fetch);

  call.transformTo(sortCopy);
}
 

@Override public void onMatch(RelOptRuleCall call) {
  final Union union = call.rel(0);

  final RexBuilder rexBuilder = union.getCluster().getRexBuilder();
  final RelMetadataQuery mq = call.getMetadataQuery();
  final RelOptPredicateList predicates = mq.getPulledUpPredicates(union);
  if (predicates == null) {
    return;
  }

  final Map<Integer, RexNode> constants = new HashMap<>();
  for (Map.Entry<RexNode, RexNode> e : predicates.constantMap.entrySet()) {
    if (e.getKey() instanceof RexInputRef) {
      constants.put(((RexInputRef) e.getKey()).getIndex(), e.getValue());
    }
  }

  // None of the expressions are constant. Nothing to do.
  if (constants.isEmpty()) {
    return;
  }

  // Create expressions for Project operators before and after the Union
  List<RelDataTypeField> fields = union.getRowType().getFieldList();
  List<RexNode> topChildExprs = new ArrayList<>();
  List<String> topChildExprsFields = new ArrayList<>();
  List<RexNode> refs = new ArrayList<>();
  ImmutableBitSet.Builder refsIndexBuilder = ImmutableBitSet.builder();
  for (RelDataTypeField field : fields) {
    final RexNode constant = constants.get(field.getIndex());
    if (constant != null) {
      topChildExprs.add(constant);
      topChildExprsFields.add(field.getName());
    } else {
      final RexNode expr = rexBuilder.makeInputRef(union, field.getIndex());
      topChildExprs.add(expr);
      topChildExprsFields.add(field.getName());
      refs.add(expr);
      refsIndexBuilder.set(field.getIndex());
    }
  }
  ImmutableBitSet refsIndex = refsIndexBuilder.build();

  // Update top Project positions
  final Mappings.TargetMapping mapping =
      RelOptUtil.permutation(refs, union.getInput(0).getRowType()).inverse();
  topChildExprs = RexUtil.apply(mapping, topChildExprs);

  // Create new Project-Union-Project sequences
  final RelBuilder relBuilder = call.builder();
  for (RelNode input : union.getInputs()) {
    List<Pair<RexNode, String>> newChildExprs = new ArrayList<>();
    for (int j : refsIndex) {
      newChildExprs.add(
          Pair.of(rexBuilder.makeInputRef(input, j),
              input.getRowType().getFieldList().get(j).getName()));
    }
    if (newChildExprs.isEmpty()) {
      // At least a single item in project is required.
      newChildExprs.add(
          Pair.of(topChildExprs.get(0), topChildExprsFields.get(0)));
    }
    // Add the input with project on top
    relBuilder.push(input);
    relBuilder.project(Pair.left(newChildExprs), Pair.right(newChildExprs));
  }
  relBuilder.union(union.all, union.getInputs().size());
  // Create top Project fixing nullability of fields
  relBuilder.project(topChildExprs, topChildExprsFields);
  relBuilder.convert(union.getRowType(), false);

  call.transformTo(relBuilder.build());
}
 

public void onMatch(RelOptRuleCall call) {
  final Aggregate aggregate = call.rel(0);
  final RelNode input = aggregate.getInput();
  final RelMetadataQuery mq = call.getMetadataQuery();
  if (!SqlFunctions.isTrue(mq.areColumnsUnique(input, aggregate.getGroupSet()))) {
    return;
  }

  final RelBuilder relBuilder = call.builder();
  final RexBuilder rexBuilder = relBuilder.getRexBuilder();
  final List<RexNode> projects = new ArrayList<>();
  for (AggregateCall aggCall : aggregate.getAggCallList()) {
    final SqlAggFunction aggregation = aggCall.getAggregation();
    if (aggregation.getKind() == SqlKind.SUM0) {
      // Bail out for SUM0 to avoid potential infinite rule matching,
      // because it may be generated by transforming SUM aggregate
      // function to SUM0 and COUNT.
      return;
    }
    final SqlSplittableAggFunction splitter =
        Objects.requireNonNull(
            aggregation.unwrap(SqlSplittableAggFunction.class));
    final RexNode singleton = splitter.singleton(
        rexBuilder, input.getRowType(), aggCall);
    projects.add(singleton);
  }

  final RelNode newInput = convert(input, aggregate.getTraitSet().simplify());
  relBuilder.push(newInput);
  if (!projects.isEmpty()) {
    projects.addAll(0, relBuilder.fields(aggregate.getGroupSet()));
    relBuilder.project(projects);
  } else if (newInput.getRowType().getFieldCount()
      > aggregate.getRowType().getFieldCount()) {
    // If aggregate was projecting a subset of columns, and there were no
    // aggregate functions, add a project for the same effect.
    relBuilder.project(relBuilder.fields(aggregate.getGroupSet()));
  }
  call.getPlanner().prune(aggregate);
  call.transformTo(relBuilder.build());
}
 

public void onMatch(RelOptRuleCall call) {
  Aggregate aggRel = call.rel(0);
  Union union = call.rel(1);

  if (!union.all) {
    // This transformation is only valid for UNION ALL.
    // Consider t1(i) with rows (5), (5) and t2(i) with
    // rows (5), (10), and the query
    // select sum(i) from (select i from t1) union (select i from t2).
    // The correct answer is 15.  If we apply the transformation,
    // we get
    // select sum(i) from
    // (select sum(i) as i from t1) union (select sum(i) as i from t2)
    // which yields 25 (incorrect).
    return;
  }

  int groupCount = aggRel.getGroupSet().cardinality();

  List<AggregateCall> transformedAggCalls =
      transformAggCalls(
          aggRel.copy(aggRel.getTraitSet(), aggRel.getInput(), false,
              aggRel.getGroupSet(), null, aggRel.getAggCallList()),
          groupCount, aggRel.getAggCallList());
  if (transformedAggCalls == null) {
    // we've detected the presence of something like AVG,
    // which we can't handle
    return;
  }

  // create corresponding aggregates on top of each union child
  final RelBuilder relBuilder = call.builder();
  int transformCount = 0;
  final RelMetadataQuery mq = call.getMetadataQuery();
  for (RelNode input : union.getInputs()) {
    boolean alreadyUnique =
        RelMdUtil.areColumnsDefinitelyUnique(mq, input,
            aggRel.getGroupSet());

    relBuilder.push(input);
    if (!alreadyUnique) {
      ++transformCount;
      relBuilder.aggregate(relBuilder.groupKey(aggRel.getGroupSet()),
          aggRel.getAggCallList());
    }
  }

  if (transformCount == 0) {
    // none of the children could benefit from the push-down,
    // so bail out (preventing the infinite loop to which most
    // planners would succumb)
    return;
  }

  // create a new union whose children are the aggregates created above
  relBuilder.union(true, union.getInputs().size());
  relBuilder.aggregate(
      relBuilder.groupKey(aggRel.getGroupSet(), aggRel.getGroupSets()),
      transformedAggCalls);
  call.transformTo(relBuilder.build());
}
 

@Override public void onMatch(RelOptRuleCall call) {
  final Sort sort = call.rel(0);
  final Join join = call.rel(1);

  // We create a new sort operator on the corresponding input
  final RelNode newLeftInput;
  final RelNode newRightInput;
  final RelMetadataQuery mq = call.getMetadataQuery();
  if (join.getJoinType() == JoinRelType.LEFT) {
    // If the input is already sorted and we are not reducing the number of tuples,
    // we bail out
    if (RelMdUtil.checkInputForCollationAndLimit(mq, join.getLeft(),
        sort.getCollation(), sort.offset, sort.fetch)) {
      return;
    }
    newLeftInput = sort.copy(sort.getTraitSet(), join.getLeft(), sort.getCollation(),
        sort.offset, sort.fetch);
    newRightInput = join.getRight();
  } else {
    final RelCollation rightCollation =
        RelCollationTraitDef.INSTANCE.canonize(
            RelCollations.shift(sort.getCollation(),
                -join.getLeft().getRowType().getFieldCount()));
    // If the input is already sorted and we are not reducing the number of tuples,
    // we bail out
    if (RelMdUtil.checkInputForCollationAndLimit(mq, join.getRight(),
        rightCollation, sort.offset, sort.fetch)) {
      return;
    }
    newLeftInput = join.getLeft();
    newRightInput = sort.copy(sort.getTraitSet().replace(rightCollation),
        join.getRight(), rightCollation, sort.offset, sort.fetch);
  }
  // We copy the join and the top sort operator
  final RelNode joinCopy = join.copy(join.getTraitSet(), join.getCondition(), newLeftInput,
      newRightInput, join.getJoinType(), join.isSemiJoinDone());
  final RelNode sortCopy = sort.copy(sort.getTraitSet(), joinCopy, sort.getCollation(),
      sort.offset, sort.fetch);

  call.transformTo(sortCopy);
}
 

@Override public void onMatch(RelOptRuleCall call) {
  final Project project = call.rel(0);
  final Join join = call.rel(1);
  final boolean isLeftJoin = join.getJoinType() == JoinRelType.LEFT;
  int lower = isLeftJoin
      ? join.getLeft().getRowType().getFieldCount() - 1 : 0;
  int upper = isLeftJoin
      ? join.getRowType().getFieldCount()
      : join.getLeft().getRowType().getFieldCount();

  // Check whether the project uses columns whose index is between
  // lower(included) and upper(excluded).
  for (RexNode expr: project.getProjects()) {
    if (RelOptUtil.InputFinder.bits(expr).asList().stream().anyMatch(
        i -> i >= lower && i < upper)) {
      return;
    }
  }

  final List<Integer> leftKeys = new ArrayList<>();
  final List<Integer> rightKeys = new ArrayList<>();
  RelOptUtil.splitJoinCondition(join.getLeft(), join.getRight(),
      join.getCondition(), leftKeys, rightKeys,
      new ArrayList<>());

  final List<Integer> joinKeys = isLeftJoin ? rightKeys : leftKeys;
  final ImmutableBitSet.Builder columns = ImmutableBitSet.builder();
  joinKeys.forEach(key -> columns.set(key));

  final RelMetadataQuery mq = call.getMetadataQuery();
  if (!mq.areColumnsUnique(isLeftJoin ? join.getRight() : join.getLeft(),
      columns.build())) {
    return;
  }

  RelNode node;
  if (isLeftJoin) {
    node = project
        .copy(project.getTraitSet(), join.getLeft(), project.getProjects(),
            project.getRowType());
  } else {
    final int offset = join.getLeft().getRowType().getFieldCount();
    final List<RexNode> newExprs = project.getProjects().stream()
        .map(expr -> RexUtil.shift(expr, -offset))
        .collect(Collectors.toList());
    node = project.copy(project.getTraitSet(), join.getRight(), newExprs,
        project.getRowType());
  }
  call.transformTo(node);
}
 

@Override public void onMatch(RelOptRuleCall call) {
  final Project project = call.rel(0);
  final Join topJoin = call.rel(1);
  final Join bottomJoin = call.rel(2);
  int leftBottomChildSize = bottomJoin.getLeft().getRowType().getFieldCount();

  // Check whether the project uses columns in the right input of bottom join.
  for (RexNode expr: project.getProjects()) {
    if (RelOptUtil.InputFinder.bits(expr).asList().stream().anyMatch(
        i -> i >= leftBottomChildSize
            && i < bottomJoin.getRowType().getFieldCount())) {
      return;
    }
  }

  // Check whether the top join uses columns in the right input of bottom join.
  final List<Integer> leftKeys = new ArrayList<>();
  RelOptUtil.splitJoinCondition(topJoin.getLeft(), topJoin.getRight(),
      topJoin.getCondition(), leftKeys, new ArrayList<>(),
      new ArrayList<>());
  if (leftKeys.stream().anyMatch(s -> s >= leftBottomChildSize)) {
    return;
  }

  // Check whether left join keys in top join and bottom join are equal.
  final List<Integer> leftChildKeys = new ArrayList<>();
  final List<Integer> rightChildKeys = new ArrayList<>();
  RelOptUtil.splitJoinCondition(bottomJoin.getLeft(), bottomJoin.getRight(),
      bottomJoin.getCondition(), leftChildKeys, rightChildKeys,
      new ArrayList<>());
  if (!leftKeys.equals(leftChildKeys)) {
    return;
  }

  // Make sure that right keys of bottom join are unique.
  final ImmutableBitSet.Builder columns = ImmutableBitSet.builder();
  rightChildKeys.forEach(key -> columns.set(key));
  final RelMetadataQuery mq = call.getMetadataQuery();
  if (!mq.areColumnsUnique(bottomJoin.getRight(), columns.build())) {
    return;
  }

  int offset = bottomJoin.getRight().getRowType().getFieldCount();
  final RelBuilder relBuilder = call.builder();

  final RexNode condition = RexUtil.shift(topJoin.getCondition(),
      leftBottomChildSize, -offset);
  final RelNode join = relBuilder.push(bottomJoin.getLeft())
      .push(topJoin.getRight())
      .join(topJoin.getJoinType(), condition)
      .build();

  final List<RexNode> newExprs = project.getProjects().stream()
      .map(expr -> RexUtil.shift(expr, leftBottomChildSize, -offset))
      .collect(Collectors.toList());
  relBuilder.push(join).project(newExprs);
  call.transformTo(relBuilder.build());
}
 

@Override public void onMatch(RelOptRuleCall call) {
  final Sort sort = call.rel(0);
  final Join join = call.rel(1);
  final RelMetadataQuery metadataQuery = call.getMetadataQuery();

  final RelNode newLeftInput;
  final RelNode newRightInput;

  final List<RelFieldCollation> leftFieldCollation = new ArrayList<>();
  final List<RelFieldCollation> rightFieldCollation = new ArrayList<>();

  // Decompose sort collations into left and right collations
  for (RelFieldCollation relFieldCollation : sort.getCollation().getFieldCollations()) {
    if (relFieldCollation.getFieldIndex() >= join.getLeft().getRowType().getFieldCount()) {
      rightFieldCollation.add(relFieldCollation);
    } else {
      leftFieldCollation.add(relFieldCollation);
    }
  }

  // Add sort to new left node only if sort collations
  // contained fields from left table
  if (leftFieldCollation.isEmpty()) {
    newLeftInput = join.getLeft();
  } else {
    final RelCollation leftCollation = RelCollationTraitDef.INSTANCE.canonize(
        RelCollations.of(leftFieldCollation));
    // If left table already sorted don't add a sort
    if (RelMdUtil.checkInputForCollationAndLimit(
        metadataQuery,
        join.getLeft(),
        leftCollation,
        null,
        null)) {
      newLeftInput = join.getLeft();
    } else {
      newLeftInput = sort.copy(
          sort.getTraitSet().replaceIf(
              RelCollationTraitDef.INSTANCE,
              () -> leftCollation),
          join.getLeft(),
          leftCollation,
          null,
          null);
    }
  }
  // Add sort to new right node only if sort collations
  // contained fields from right table
  if (rightFieldCollation.isEmpty()) {
    newRightInput = join.getRight();
  } else {
    final RelCollation rightCollation = RelCollationTraitDef.INSTANCE.canonize(
        RelCollations.shift(
            RelCollations.of(rightFieldCollation),
            -join.getLeft().getRowType().getFieldCount()));
    // If right table already sorted don't add a sort
    if (RelMdUtil.checkInputForCollationAndLimit(
        metadataQuery,
        join.getRight(),
        rightCollation,
        null,
        null)) {
      newRightInput = join.getRight();
    } else {
      newRightInput = sort.copy(
          sort.getTraitSet().replaceIf(
              RelCollationTraitDef.INSTANCE,
              () -> rightCollation),
          join.getRight(),
          rightCollation,
          null,
          null);
    }
  }
  // If no change was made no need to apply the rule
  if (newLeftInput == join.getLeft() && newRightInput == join.getRight()) {
    return;
  }

  final RelNode joinCopy = join.copy(
      join.getTraitSet(),
      join.getCondition(),
      newLeftInput,
      newRightInput,
      join.getJoinType(),
      join.isSemiJoinDone());
  final RelNode sortCopy = sort.copy(
      sort.getTraitSet(),
      joinCopy,
      sort.getCollation(),
      sort.offset,
      sort.fetch);

  call.transformTo(sortCopy);
}