Java Code Examples for org.apache.calcite.plan.hep.HepPlanner#findBestExp()

/** Creates a program that executes a {@link HepProgram}. */
public static Program of(final HepProgram hepProgram, final boolean noDag,
    final RelMetadataProvider metadataProvider) {
  return (planner, rel, requiredOutputTraits, materializations, lattices) -> {
    final HepPlanner hepPlanner = new HepPlanner(hepProgram,
        null, noDag, null, RelOptCostImpl.FACTORY);

    List<RelMetadataProvider> list = new ArrayList<>();
    if (metadataProvider != null) {
      list.add(metadataProvider);
    }
    hepPlanner.registerMetadataProviders(list);
    for (RelOptMaterialization materialization : materializations) {
      hepPlanner.addMaterialization(materialization);
    }
    for (RelOptLattice lattice : lattices) {
      hepPlanner.addLattice(lattice);
    }
    RelMetadataProvider plannerChain =
        ChainedRelMetadataProvider.of(list);
    rel.getCluster().setMetadataProvider(plannerChain);

    hepPlanner.setRoot(rel);
    return hepPlanner.findBestExp();
  };
}
 

protected void apply(RelOptRuleCall call, Filter filter, TableScan scan) {
  final RelOptPlanner planner = call.getPlanner();
  final List<RelOptMaterialization> materializations =
      planner.getMaterializations();
  if (!materializations.isEmpty()) {
    RelNode root = filter.copy(filter.getTraitSet(),
        Collections.singletonList((RelNode) scan));
    List<RelOptMaterialization> applicableMaterializations =
        RelOptMaterializations.getApplicableMaterializations(root, materializations);
    for (RelOptMaterialization materialization : applicableMaterializations) {
      if (RelOptUtil.areRowTypesEqual(scan.getRowType(),
          materialization.queryRel.getRowType(), false)) {
        RelNode target = materialization.queryRel;
        final HepPlanner hepPlanner =
            new HepPlanner(program, planner.getContext());
        hepPlanner.setRoot(target);
        target = hepPlanner.findBestExp();
        List<RelNode> subs = new SubstitutionVisitor(target, root)
            .go(materialization.tableRel);
        for (RelNode s : subs) {
          call.transformTo(s);
        }
      }
    }
  }
}
 

/** Creates a program that executes a {@link HepProgram}. */
public static Program of(final HepProgram hepProgram, final boolean noDag,
    final RelMetadataProvider metadataProvider) {
  return (planner, rel, requiredOutputTraits, materializations, lattices) -> {
    final HepPlanner hepPlanner = new HepPlanner(hepProgram,
        null, noDag, null, RelOptCostImpl.FACTORY);

    List<RelMetadataProvider> list = new ArrayList<>();
    if (metadataProvider != null) {
      list.add(metadataProvider);
    }
    hepPlanner.registerMetadataProviders(list);
    RelMetadataProvider plannerChain =
        ChainedRelMetadataProvider.of(list);
    rel.getCluster().setMetadataProvider(plannerChain);

    hepPlanner.setRoot(rel);
    return hepPlanner.findBestExp();
  };
}
 

@Test void testHintsPropagateWithDifferentKindOfRels() {
  final String sql = "select /*+ AGG_STRATEGY(TWO_PHASE) */\n"
      + "ename, avg(sal)\n"
      + "from emp group by ename";
  final RelNode rel = tester.convertSqlToRel(sql).rel;
  final RelHint hint = RelHint.builder("AGG_STRATEGY")
      .inheritPath(0)
      .hintOption("TWO_PHASE")
      .build();
  // AggregateReduceFunctionsRule does the transformation:
  // AGG -> PROJECT + AGG
  HepProgram program = new HepProgramBuilder()
      .addRuleInstance(AggregateReduceFunctionsRule.INSTANCE)
      .build();
  HepPlanner planner = new HepPlanner(program);
  planner.setRoot(rel);
  RelNode newRel = planner.findBestExp();
  new ValidateHintVisitor(hint, Aggregate.class).go(newRel);
}
 

@Test void testCalcColumnOriginsTable() {
  final String sql = "select name,deptno from dept where deptno > 10";
  final RelNode relNode = convertSql(sql);
  final HepProgram program = new HepProgramBuilder().
      addRuleInstance(ProjectToCalcRule.INSTANCE).build();
  final HepPlanner planner = new HepPlanner(program);
  planner.setRoot(relNode);
  final RelNode calc = planner.findBestExp();
  final RelMetadataQuery mq = calc.getCluster().getMetadataQuery();
  final RelColumnOrigin nameColumn = mq.getColumnOrigin(calc, 0);
  assertThat(nameColumn.getOriginColumnOrdinal(), is(1));
  final RelColumnOrigin deptnoColumn = mq.getColumnOrigin(calc, 1);
  assertThat(deptnoColumn.getOriginColumnOrdinal(), is(0));
}
 

/** Adds a query.
 *
 * <p>It may fit within an existing lattice (or lattices). Or it may need a
 * new lattice, or an extension to an existing lattice.
 *
 * @param r Relational expression for a query
 *
 * @return A list of join graphs: usually 1; more if the query contains a
 * cartesian product; zero if the query graph is cyclic
 */
public List<Lattice> addQuery(RelNode r) {
  // Push filters into joins and towards leaves
  final HepPlanner planner =
      new HepPlanner(PROGRAM, null, true, null, RelOptCostImpl.FACTORY);
  planner.setRoot(r);
  final RelNode r2 = planner.findBestExp();

  final Query q = new Query(space);
  final List<Frame> frameList = new ArrayList<>();
  frames(frameList, q, r2);
  final List<Lattice> lattices = new ArrayList<>();
  frameList.forEach(frame -> addFrame(q, frame, lattices));
  return ImmutableList.copyOf(lattices);
}
 

@Test void testInterpretSemiJoin() throws Exception {
  final String sql = "select x, y from (values (1, 'a'), (2, 'b'), (3, 'c')) as t(x, y)\n"
      + "where x in\n"
      + "(select x from (values (1, 'd'), (3, 'g')) as t2(x, y))";
  SqlNode validate = planner.validate(planner.parse(sql));
  RelNode convert = planner.rel(validate).rel;
  final HepProgram program = new HepProgramBuilder()
      .addRuleInstance(SemiJoinRule.PROJECT)
      .build();
  final HepPlanner hepPlanner = new HepPlanner(program);
  hepPlanner.setRoot(convert);
  final RelNode relNode = hepPlanner.findBestExp();
  final Interpreter interpreter = new Interpreter(dataContext, relNode);
  assertRows(interpreter, true, "[1, a]", "[3, c]");
}
 

@Test void testHintsForCalc() {
  final String sql = "select /*+ resource(mem='1024MB')*/ ename, sal, deptno from emp";
  final RelNode rel = tester.convertSqlToRel(sql).rel;
  final RelHint hint = RelHint.builder("RESOURCE")
      .hintOption("MEM", "1024MB")
      .build();
  // planner rule to convert Project to Calc.
  HepProgram program = new HepProgramBuilder()
      .addRuleInstance(ProjectToCalcRule.INSTANCE)
      .build();
  HepPlanner planner = new HepPlanner(program);
  planner.setRoot(rel);
  RelNode newRel = planner.findBestExp();
  new ValidateHintVisitor(hint, Calc.class).go(newRel);
}
 

public RelNode removeCorrelationViaRule(RelNode root) {
  final RelBuilderFactory f = relBuilderFactory();
  HepProgram program = HepProgram.builder()
      .addRuleInstance(new RemoveSingleAggregateRule(f))
      .addRuleInstance(new RemoveCorrelationForScalarProjectRule(f))
      .addRuleInstance(new RemoveCorrelationForScalarAggregateRule(f))
      .build();

  HepPlanner planner = createPlanner(program);

  planner.setRoot(root);
  return planner.findBestExp();
}
 

@Test void testCalcFieldTrimmer1() {
  final RelBuilder builder = RelBuilder.create(config().build());
  final RelNode root =
      builder.scan("EMP")
          .project(builder.field("EMPNO"), builder.field("ENAME"), builder.field("DEPTNO"))
          .exchange(RelDistributions.SINGLETON)
          .filter(
              builder.call(SqlStdOperatorTable.GREATER_THAN,
                  builder.field("EMPNO"), builder.literal(100)))
          .build();

  final HepProgram hepProgram = new HepProgramBuilder()
      .addRuleInstance(ProjectToCalcRule.INSTANCE)
      .addRuleInstance(FilterToCalcRule.INSTANCE)
      .build();

  final HepPlanner hepPlanner = new HepPlanner(hepProgram);
  hepPlanner.setRoot(root);
  final RelNode relNode = hepPlanner.findBestExp();
  final RelFieldTrimmer fieldTrimmer = new RelFieldTrimmer(null, builder);
  final RelNode trimmed = fieldTrimmer.trim(relNode);

  final String expected = ""
      + "LogicalCalc(expr#0..2=[{inputs}], expr#3=[100], expr#4=[>($t0, $t3)], proj#0."
      + ".2=[{exprs}], $condition=[$t4])\n"
      + "  LogicalExchange(distribution=[single])\n"
      + "    LogicalCalc(expr#0..2=[{inputs}], proj#0..2=[{exprs}])\n"
      + "      LogicalProject(EMPNO=[$0], ENAME=[$1], DEPTNO=[$7])\n"
      + "        LogicalTableScan(table=[[scott, EMP]])\n";
  assertThat(trimmed, hasTree(expected));
}
 

private static List<RelNode> substitute(
    RelNode root, RelOptMaterialization materialization) {
  // First, if the materialization is in terms of a star table, rewrite
  // the query in terms of the star table.
  if (materialization.starTable != null) {
    RelNode newRoot = RelOptMaterialization.tryUseStar(root,
        materialization.starRelOptTable);
    if (newRoot != null) {
      root = newRoot;
    }
  }

  // Push filters to the bottom, and combine projects on top.
  RelNode target = materialization.queryRel;
  HepProgram program =
      new HepProgramBuilder()
          .addRuleInstance(FilterProjectTransposeRule.INSTANCE)
          .addRuleInstance(ProjectMergeRule.INSTANCE)
          .addRuleInstance(ProjectRemoveRule.INSTANCE)
          .build();

  final HepPlanner hepPlanner = new HepPlanner(program);
  hepPlanner.setRoot(target);
  target = hepPlanner.findBestExp();

  hepPlanner.setRoot(root);
  root = hepPlanner.findBestExp();

  return new MaterializedViewSubstitutionVisitor(target, root)
          .go(materialization.tableRel);
}
 

public RelNode removeCorrelationViaRule(RelNode root) {
    final RelBuilderFactory f = relBuilderFactory();
    HepProgram program = HepProgram.builder().addRuleInstance(new RemoveSingleAggregateRule(f))
            .addRuleInstance(new RemoveCorrelationForScalarProjectRule(f))
            .addRuleInstance(new RemoveCorrelationForScalarAggregateRule(f)).build();

    HepPlanner planner = createPlanner(program);

    planner.setRoot(root);
    return planner.findBestExp();
}
 

/**
 * Ensures {@link org.apache.calcite.rel.AbstractRelNode} digest does not include
 * full digest tree.
 */
@Test void relDigestLength() {
  HepProgramBuilder programBuilder = HepProgram.builder();
  HepPlanner planner =
      new HepPlanner(
          programBuilder.build());
  StringBuilder sb = new StringBuilder();
  final int n = 10;
  sb.append("select * from (");
  sb.append("select name from sales.dept");
  for (int i = 0; i < n; i++) {
    sb.append(" union all select name from sales.dept");
  }
  sb.append(")");
  RelRoot root = tester.convertSqlToRel(sb.toString());
  planner.setRoot(root.rel);
  RelNode best = planner.findBestExp();

  // Good digest should look like rel#66:LogicalProject(input=rel#64:LogicalUnion)
  // Bad digest includes full tree like rel#66:LogicalProject(input=rel#64:LogicalUnion(...))
  // So the assertion is to ensure digest includes LogicalUnion exactly once

  assertIncludesExactlyOnce("best.getDescription()",
      best.toString(), "LogicalUnion");
  assertIncludesExactlyOnce("best.getDigest()",
      best.getDigest(), "LogicalUnion");
}
 

@Override
public Pair<RelNode, RelNode> pushFilterToOriginalViewPlan(RelBuilder builder, RelNode topViewProject,
        RelNode viewNode, RexNode cond) {
    // We add (and push) the filter to the view plan before triggering the rewriting.
    // This is useful in case some of the columns can be folded to same value after
    // filter is added.
    HepProgramBuilder pushFiltersProgram = new HepProgramBuilder();
    if (topViewProject != null) {
        pushFiltersProgram.addRuleInstance(filterProjectTransposeRule);
    }
    pushFiltersProgram.addRuleInstance(this.filterAggregateTransposeRule)
            .addRuleInstance(this.aggregateProjectPullUpConstantsRule).addRuleInstance(this.projectMergeRule);
    final HepPlanner tmpPlanner = new HepPlanner(pushFiltersProgram.build());
    // Now that the planner is created, push the node
    RelNode topNode = builder.push(topViewProject != null ? topViewProject : viewNode).filter(cond).build();
    tmpPlanner.setRoot(topNode);
    topNode = tmpPlanner.findBestExp();
    RelNode resultTopViewProject = null;
    RelNode resultViewNode = null;
    while (topNode != null) {
        if (topNode instanceof Project) {
            if (resultTopViewProject != null) {
                // Both projects could not be merged, we will bail out
                return Pair.of(topViewProject, viewNode);
            }
            resultTopViewProject = topNode;
            topNode = topNode.getInput(0);
        } else if (topNode instanceof Aggregate) {
            resultViewNode = topNode;
            topNode = null;
        } else {
            // We move to the child
            topNode = topNode.getInput(0);
        }
    }
    return Pair.of(resultTopViewProject, resultViewNode);
}
 

protected RelNode decorrelate(RelNode root) {
  // first adjust count() expression if any
  final RelBuilderFactory f = relBuilderFactory();
  HepProgram program = HepProgram.builder()
      .addRuleInstance(new AdjustProjectForCountAggregateRule(false, f))
      .addRuleInstance(new AdjustProjectForCountAggregateRule(true, f))
      .addRuleInstance(
          new FilterJoinRule.FilterIntoJoinRule(true, f,
              FilterJoinRule.TRUE_PREDICATE))
      .addRuleInstance(
          new FilterProjectTransposeRule(Filter.class, Project.class, true,
              true, f))
      .addRuleInstance(new FilterCorrelateRule(f))
      .build();

  HepPlanner planner = createPlanner(program);

  planner.setRoot(root);
  root = planner.findBestExp();

  // Perform decorrelation.
  map.clear();

  final Frame frame = getInvoke(root, null);
  if (frame != null) {
    // has been rewritten; apply rules post-decorrelation
    final HepProgram program2 = HepProgram.builder()
        .addRuleInstance(
            new FilterJoinRule.FilterIntoJoinRule(
                true, f,
                FilterJoinRule.TRUE_PREDICATE))
        .addRuleInstance(
            new FilterJoinRule.JoinConditionPushRule(
                f,
                FilterJoinRule.TRUE_PREDICATE))
        .build();

    final HepPlanner planner2 = createPlanner(program2);
    final RelNode newRoot = frame.r;
    planner2.setRoot(newRoot);
    return planner2.findBestExp();
  }

  return root;
}
 

@Override protected RelNode rewriteQuery(
    RelBuilder relBuilder,
    RexBuilder rexBuilder,
    RexSimplify simplify,
    RelMetadataQuery mq,
    RexNode compensationColumnsEquiPred,
    RexNode otherCompensationPred,
    Project topProject,
    RelNode node,
    BiMap<RelTableRef, RelTableRef> queryToViewTableMapping,
    EquivalenceClasses viewEC, EquivalenceClasses queryEC) {
  Aggregate aggregate = (Aggregate) node;

  // 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.
  // 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 newAggregateInput = aggregate.getInput(0);
  RelNode target = aggregate.getInput(0);
  if (unionRewritingPullProgram != null) {
    final HepPlanner tmpPlanner = new HepPlanner(unionRewritingPullProgram);
    tmpPlanner.setRoot(newAggregateInput);
    newAggregateInput = tmpPlanner.findBestExp();
    target = newAggregateInput.getInput(0);
  }

  // We need to check that all columns required by compensating predicates
  // are contained in the query.
  List<RexNode> queryExprs = extractReferences(rexBuilder, target);
  if (!compensationColumnsEquiPred.isAlwaysTrue()) {
    compensationColumnsEquiPred = rewriteExpression(rexBuilder, mq,
        target, target, queryExprs, queryToViewTableMapping, 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, queryToViewTableMapping, 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) {
    return aggregate.copy(aggregate.getTraitSet(),
        ImmutableList.of(
            newAggregateInput.copy(newAggregateInput.getTraitSet(),
                ImmutableList.of(rewrittenPlan))));
  }
  return aggregate.copy(aggregate.getTraitSet(), ImmutableList.of(rewrittenPlan));
}
 

@SneakyThrows
    public static RowBaseIterator run(String sql, MycatCalciteDataContext calciteDataContext, RelNode relNode) {
        SqlRecorder recorder = SqlRecorderRuntime.INSTANCE.getCurrentRecorder();
        Map<String, List<SingeTargetSQLTable>> map = new HashMap<>();
        relNode.accept(new RelShuttleImpl() {
            @Override
            public RelNode visit(TableScan scan) {
                SingeTargetSQLTable unwrap = scan.getTable().unwrap(SingeTargetSQLTable.class);
                if (unwrap != null && !unwrap.existsEnumerable()) {
                    List<SingeTargetSQLTable> tables = map.computeIfAbsent(unwrap.getTargetName(), s -> new ArrayList<>(2));
                    tables.add(unwrap);
                }
                return super.visit(scan);
            }
        });

        HepProgramBuilder hepProgramBuilder = new HepProgramBuilder();
        hepProgramBuilder.addMatchLimit(64);

        hepProgramBuilder.addRuleInstance(StreamUnionRule.INSTANCE);
        final HepPlanner planner2 = new HepPlanner(hepProgramBuilder.build());
        planner2.setRoot(relNode);
        relNode = planner2.findBestExp();

        //check
        relNode.accept(new RelShuttleImpl() {
            @Override
            public RelNode visit(LogicalUnion union) {
                if (union.getInputs().size() > 2) {
                    throw new AssertionError("union input more 2");
                }
                return super.visit(union);
            }
        });
        long startGetConnectionTime = TimeProvider.INSTANCE.now();
        fork(sql, calciteDataContext, map);
        long cbo = TimeProvider.INSTANCE.now();
        recorder.addRecord(SqlRecorderType.GET_CONNECTION, sql, cbo - startGetConnectionTime);
        ArrayBindable bindable1 = Interpreters.bindable(relNode);
        long execution_start = TimeProvider.INSTANCE.now();
        recorder.addRecord(SqlRecorderType.CBO, sql, execution_start - cbo);
//        EnumerableInterpretable.toBindable()
        Enumerable<Object[]> bind = bindable1.bind(calciteDataContext);
        Enumerator<Object[]> enumerator = bind.enumerator();

        return new EnumeratorRowIterator(CalciteConvertors.getMycatRowMetaData(relNode.getRowType()), enumerator,
                () -> {
                    recorder.addRecord(SqlRecorderType.EXECUTION_TIME, sql, TimeProvider.INSTANCE.now()-execution_start);
                    recorder.addRecord(SqlRecorderType.AT_END, sql, TimeProvider.INSTANCE.now());
                });
    }
 

private static List<RelNode> substitute(
    RelNode root, RelOptMaterialization materialization) {
  // First, if the materialization is in terms of a star table, rewrite
  // the query in terms of the star table.
  if (materialization.starTable != null) {
    RelNode newRoot = RelOptMaterialization.tryUseStar(root,
        materialization.starRelOptTable);
    if (newRoot != null) {
      root = newRoot;
    }
  }

  // Push filters to the bottom, and combine projects on top.
  RelNode target = materialization.queryRel;
  // try to trim unused field in relational expressions.
  root = trimUnusedfields(root);
  target = trimUnusedfields(target);
  HepProgram program =
      new HepProgramBuilder()
          .addRuleInstance(FilterProjectTransposeRule.INSTANCE)
          .addRuleInstance(FilterMergeRule.INSTANCE)
          .addRuleInstance(FilterJoinRule.FILTER_ON_JOIN)
          .addRuleInstance(FilterJoinRule.JOIN)
          .addRuleInstance(FilterAggregateTransposeRule.INSTANCE)
          .addRuleInstance(ProjectMergeRule.INSTANCE)
          .addRuleInstance(ProjectRemoveRule.INSTANCE)
          .addRuleInstance(ProjectJoinTransposeRule.INSTANCE)
          .addRuleInstance(ProjectSetOpTransposeRule.INSTANCE)
          .addRuleInstance(FilterToCalcRule.INSTANCE)
          .addRuleInstance(ProjectToCalcRule.INSTANCE)
          .addRuleInstance(FilterCalcMergeRule.INSTANCE)
          .addRuleInstance(ProjectCalcMergeRule.INSTANCE)
          .addRuleInstance(CalcMergeRule.INSTANCE)
          .build();

  // We must use the same HEP planner for the two optimizations below.
  // Thus different nodes with the same digest will share the same vertex in
  // the plan graph. This is important for the matching process.
  final HepPlanner hepPlanner = new HepPlanner(program);
  hepPlanner.setRoot(target);
  target = hepPlanner.findBestExp();

  hepPlanner.setRoot(root);
  root = hepPlanner.findBestExp();

  return new SubstitutionVisitor(target, root).go(materialization.tableRel);
}
 

/**
 * Checks the plan for a SQL statement before/after executing a given rule,
 * with a pre-program to prepare the tree.
 *
 * @param tester     Tester
 * @param preProgram Program to execute before comparing before state
 * @param planner    Planner
 * @param sql        SQL query
 * @param unchanged  Whether the rule is to have no effect
 */
private void checkPlanning(Tester tester, HepProgram preProgram,
    RelOptPlanner planner, String sql, boolean unchanged) {
  final DiffRepository diffRepos = getDiffRepos();
  String sql2 = diffRepos.expand("sql", sql);
  final RelRoot root = tester.convertSqlToRel(sql2);
  final RelNode relInitial = root.rel;

  assertNotNull(relInitial);

  List<RelMetadataProvider> list = new ArrayList<>();
  list.add(DefaultRelMetadataProvider.INSTANCE);
  planner.registerMetadataProviders(list);
  RelMetadataProvider plannerChain =
      ChainedRelMetadataProvider.of(list);
  final RelOptCluster cluster = relInitial.getCluster();
  cluster.setMetadataProvider(plannerChain);

  RelNode relBefore;
  if (preProgram == null) {
    relBefore = relInitial;
  } else {
    HepPlanner prePlanner = new HepPlanner(preProgram);
    prePlanner.setRoot(relInitial);
    relBefore = prePlanner.findBestExp();
  }

  assertThat(relBefore, notNullValue());

  final String planBefore = NL + RelOptUtil.toString(relBefore);
  diffRepos.assertEquals("planBefore", "${planBefore}", planBefore);
  SqlToRelTestBase.assertValid(relBefore);

  if (planner instanceof VolcanoPlanner) {
    relBefore = planner.changeTraits(relBefore,
        relBefore.getTraitSet().replace(EnumerableConvention.INSTANCE));
  }
  planner.setRoot(relBefore);
  RelNode r = planner.findBestExp();
  if (tester.isLateDecorrelate()) {
    final String planMid = NL + RelOptUtil.toString(r);
    diffRepos.assertEquals("planMid", "${planMid}", planMid);
    SqlToRelTestBase.assertValid(r);
    final RelBuilder relBuilder =
        RelFactories.LOGICAL_BUILDER.create(cluster, null);
    r = RelDecorrelator.decorrelateQuery(r, relBuilder);
  }
  final String planAfter = NL + RelOptUtil.toString(r);
  if (unchanged) {
    assertThat(planAfter, is(planBefore));
  } else {
    diffRepos.assertEquals("planAfter", "${planAfter}", planAfter);
    if (planBefore.equals(planAfter)) {
      throw new AssertionError("Expected plan before and after is the same.\n"
          + "You must use unchanged=true or call checkUnchanged");
    }
  }
  SqlToRelTestBase.assertValid(r);
}
 

@Override
protected RelNode rewriteQuery(RelBuilder relBuilder, RexBuilder rexBuilder, RexSimplify simplify,
        RelMetadataQuery mq, RexNode compensationColumnsEquiPred, RexNode otherCompensationPred,
        Project topProject, RelNode node, BiMap<RelTableRef, RelTableRef> queryToViewTableMapping,
        EquivalenceClasses viewEC, EquivalenceClasses queryEC) {
    Aggregate aggregate = (Aggregate) node;

    // 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.
    // 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 newAggregateInput = aggregate.getInput(0);
    RelNode target = aggregate.getInput(0);
    if (unionRewritingPullProgram != null) {
        final HepPlanner tmpPlanner = new HepPlanner(unionRewritingPullProgram);
        tmpPlanner.setRoot(newAggregateInput);
        newAggregateInput = tmpPlanner.findBestExp();
        target = newAggregateInput.getInput(0);
    }

    // We need to check that all columns required by compensating predicates
    // are contained in the query.
    List<RexNode> queryExprs = extractReferences(rexBuilder, target);
    if (!compensationColumnsEquiPred.isAlwaysTrue()) {
        compensationColumnsEquiPred = rewriteExpression(rexBuilder, mq, target, target, queryExprs,
                queryToViewTableMapping, 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,
                queryToViewTableMapping, 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) {
        return aggregate.copy(aggregate.getTraitSet(), ImmutableList
                .of(newAggregateInput.copy(newAggregateInput.getTraitSet(), ImmutableList.of(rewrittenPlan))));
    }
    return aggregate.copy(aggregate.getTraitSet(), ImmutableList.of(rewrittenPlan));
}