org.apache.calcite.tools.Program Java Examples

/**
 * Converts a relational expression to a form where
 * {@link org.apache.calcite.rel.logical.LogicalJoin}s are
 * as close to leaves as possible.
 */
public static RelNode toLeafJoinForm(RelNode rel) {
  final Program program = Programs.hep(
      ImmutableList.of(
          JoinProjectTransposeRule.RIGHT_PROJECT,
          JoinProjectTransposeRule.LEFT_PROJECT,
          FilterJoinRule.FilterIntoJoinRule.FILTER_ON_JOIN,
          ProjectRemoveRule.INSTANCE,
          ProjectMergeRule.INSTANCE),
      false,
      DefaultRelMetadataProvider.INSTANCE);
  if (CalciteSystemProperty.DEBUG.value()) {
    System.out.println(
        RelOptUtil.dumpPlan("before", rel, SqlExplainFormat.TEXT,
            SqlExplainLevel.DIGEST_ATTRIBUTES));
  }
  final RelNode rel2 = program.run(null, rel, null,
      ImmutableList.of(),
      ImmutableList.of());
  if (CalciteSystemProperty.DEBUG.value()) {
    System.out.println(
        RelOptUtil.dumpPlan("after", rel2, SqlExplainFormat.TEXT,
            SqlExplainLevel.DIGEST_ATTRIBUTES));
  }
  return rel2;
}
 

@Test void testUseMergeJoin() {
  final String sql = "select /*+ use_merge_join(emp, dept) */\n"
      + "ename, job, sal, dept.name\n"
      + "from emp join dept on emp.deptno = dept.deptno";
  RelOptPlanner planner = new VolcanoPlanner();
  planner.addRelTraitDef(ConventionTraitDef.INSTANCE);
  planner.addRelTraitDef(RelCollationTraitDef.INSTANCE);
  Tester tester1 = tester.withDecorrelation(true)
      .withClusterFactory(
          relOptCluster -> RelOptCluster.create(planner, relOptCluster.getRexBuilder()));
  final RelNode rel = tester1.convertSqlToRel(sql).rel;
  RuleSet ruleSet = RuleSets.ofList(
      EnumerableRules.ENUMERABLE_MERGE_JOIN_RULE,
      EnumerableRules.ENUMERABLE_JOIN_RULE,
      EnumerableRules.ENUMERABLE_PROJECT_RULE,
      EnumerableRules.ENUMERABLE_TABLE_SCAN_RULE,
      EnumerableRules.ENUMERABLE_SORT_RULE,
      AbstractConverter.ExpandConversionRule.INSTANCE);
  Program program = Programs.of(ruleSet);
  RelTraitSet toTraits = rel
      .getCluster()
      .traitSet()
      .replace(EnumerableConvention.INSTANCE);

  RelNode relAfter = program.run(planner, rel, toTraits,
      Collections.emptyList(), Collections.emptyList());

  String planAfter = NL + RelOptUtil.toString(relAfter);
  getDiffRepos().assertEquals("planAfter", "${planAfter}", planAfter);
}
 

private static Planner getPlanner(List<RelTraitDef> traitDefs,
    Config parserConfig, Program... programs) {
  final SchemaPlus rootSchema = Frameworks.createRootSchema(true);
  rootSchema.add("TMP", new AbstractTable() {
    @Override public RelDataType getRowType(RelDataTypeFactory typeFactory) {
      return typeFactory.createStructType(
          ImmutableList.of(typeFactory.createSqlType(SqlTypeName.VARCHAR),
              typeFactory.createSqlType(SqlTypeName.INTEGER)),
          ImmutableList.of("localtime", "current_timestamp"));
    }
  });
  final FrameworkConfig config = Frameworks.newConfigBuilder()
      .parserConfig(parserConfig)
      .defaultSchema(rootSchema)
      .traitDefs(traitDefs)
      .programs(programs)
      .operatorTable(SqlStdOperatorTable.instance())
      .build();
  return Frameworks.getPlanner(config);
}
 

/**
 * Converts a relational expression to a form where
 * {@link org.apache.calcite.rel.logical.LogicalJoin}s are
 * as close to leaves as possible.
 */
public static RelNode toLeafJoinForm(RelNode rel) {
  final Program program = Programs.hep(
      ImmutableList.of(
          JoinProjectTransposeRule.RIGHT_PROJECT,
          JoinProjectTransposeRule.LEFT_PROJECT,
          FilterJoinRule.FilterIntoJoinRule.FILTER_ON_JOIN,
          ProjectRemoveRule.INSTANCE,
          ProjectMergeRule.INSTANCE),
      false,
      DefaultRelMetadataProvider.INSTANCE);
  if (CalciteSystemProperty.DEBUG.value()) {
    System.out.println(
        RelOptUtil.dumpPlan("before", rel, SqlExplainFormat.TEXT,
            SqlExplainLevel.DIGEST_ATTRIBUTES));
  }
  final RelNode rel2 = program.run(null, rel, null,
      ImmutableList.of(),
      ImmutableList.of());
  if (CalciteSystemProperty.DEBUG.value()) {
    System.out.println(
        RelOptUtil.dumpPlan("after", rel2, SqlExplainFormat.TEXT,
            SqlExplainLevel.DIGEST_ATTRIBUTES));
  }
  return rel2;
}
 

private List<RelNode> optimizePlans(List<RelNode> originalRels,
    List<RelOptRule> rules) {
  final RelOptPlanner planner = originalRels.get(0).getCluster().getPlanner();
  // Remember old rule set of the planner before resetting it with new rules
  final List<RelOptRule> oldRules = planner.getRules();
  resetPlannerRules(planner, rules);
  final Program program = Programs.of(RuleSets.ofList(planner.getRules()));
  final List<RelNode> optimizedPlans = new ArrayList<>();
  for (RelNode rel : originalRels) {
    final RelCollation collation = rel instanceof Sort
        ? ((Sort) rel).collation
        : RelCollations.EMPTY;
    // Apply the planner to obtain the physical plan
    final RelNode physicalPlan = program.run(planner, rel,
        rel.getTraitSet().replace(EnumerableConvention.INSTANCE)
            .replace(collation).simplify(),
        ImmutableList.of(), ImmutableList.of());

    // Then convert the physical plan back to logical plan
    final RelNode logicalPlan = new ToLogicalConverter(builder).visit(physicalPlan);
    optimizedPlans.add(logicalPlan);
  }
  resetPlannerRules(planner, oldRules);
  return optimizedPlans;
}
 

public static RelNode toPhysical(RelNode rel, Consumer<RelOptPlanner> setting) {
    final RelOptPlanner planner = rel.getCluster().getPlanner();
    planner.clear();
    setting.accept(planner);
    planner.addRule(new RelOptRule(operand(MycatTransientSQLTableScan.class, none()), RelFactories.LOGICAL_BUILDER, "MycatTransientSQLTableScan") {

        @Override
        public void onMatch(RelOptRuleCall call) {
            final MycatTransientSQLTableScan scan = call.rel(0);
            final RelOptTable table = scan.getTable();
            if (Bindables.BindableTableScan.canHandle(table)) {
                call.transformTo(
                        Bindables.BindableTableScan.create(scan.getCluster(), table));
            }
        }
    });
    final Program program = Programs.of(RuleSets.ofList(planner.getRules()));
    return program.run(planner, rel, rel.getTraitSet().replace(EnumerableConvention.INSTANCE),
            ImmutableList.of(), ImmutableList.of());
}
 

@Test
public void testFactoryWillAutomaticallyAddRules() {
	// This test changes the global state of Calcite! It shouldn't cause issues elsewhere since the rules avoid
	// changing unrecognised tables, so will not apply to their own output.
	JournalledJdbcSchema.Factory.INSTANCE.setAutomaticallyAddRules(true);
	JournalledJdbcSchema.Factory.INSTANCE.create(null, "my-parent", options);
	try {
		Program def = Mockito.mock(Program.class);
		Holder<Program> holder = Holder.of(def);
		Hook.PROGRAM.run(holder);
		Assert.assertTrue(holder.get() instanceof SequenceProgram);
		Assert.assertTrue(((SequenceProgram) holder.get()).getPrograms().get(0) instanceof ForcedRulesProgram);
	} finally {
		// ensure no gap where another hook may be added
		JournalledJdbcSchema.Factory.INSTANCE.setAutomaticallyAddRules(false);
		JournalledJdbcRuleManager.removeHook();
	}
}
 

protected RelRoot getSqlPlanRel(String sql) throws Throwable {
  try (Planner planner = Frameworks.getPlanner(frameworkConfig)) {
    RelRoot root;
    final SqlNode parsedSqlNode = planner.parse(sql);
    final Pair<SqlNode, RelDataType> validatedSqlNodeAndType = planner
        .validateAndGetType(
            parsedSqlNode);
    root = planner.rel(validatedSqlNodeAndType.getKey());
    final Program program = createProgram();
    //getDesiredTraits
    final RelTraitSet desiredTraits = root.rel.getTraitSet()
        .replace(EnumerableConvention.INSTANCE)
        .replace(root.collation)
        .simplify();

    RelNode logicalRelNode = root.rel;
    final RelNode optimizedRelNode = program.run(
        root.rel.getCluster().getPlanner(), logicalRelNode, desiredTraits,
        Collections.emptyList(), Collections.emptyList());
    root = root.withRel(optimizedRelNode);
    return root;
  }

}
 

private static Consumer<Holder<Program>> handler(
    final boolean bushy, final int minJoinCount) {
  return holder -> holder.set(
      Programs.sequence(
          Programs.heuristicJoinOrder(Programs.RULE_SET, bushy,
              minJoinCount),
          Programs.CALC_PROGRAM));
}
 

static void testPrograms(RelNode relNode) {
    final RelOptPlanner planner = relNode.getCluster().getPlanner();
    final Program program = Programs.ofRules(ReduceExpressionsRule.PROJECT_INSTANCE);
    relNode = program.run(planner, relNode, relNode.getTraitSet(), ImmutableList.of(), ImmutableList.of());
    String plan = RelOptUtil.toString(relNode);
    System.out.println(plan);
}
 

@Test void testHintsPropagationInVolcanoPlannerRules() {
  final String sql = "select /*+ use_hash_join(r, s), use_hash_join(emp, dept) */\n"
      + "ename, job, sal, dept.name\n"
      + "from emp join dept on emp.deptno = dept.deptno";
  RelOptPlanner planner = new VolcanoPlanner();
  planner.addRelTraitDef(ConventionTraitDef.INSTANCE);
  Tester tester1 = tester.withDecorrelation(true)
      .withClusterFactory(
        relOptCluster -> RelOptCluster.create(planner, relOptCluster.getRexBuilder()));
  final RelNode rel = tester1.convertSqlToRel(sql).rel;
  final RelHint hint = RelHint.builder("USE_HASH_JOIN")
      .inheritPath(0)
      .hintOption("EMP")
      .hintOption("DEPT")
      .build();
  // Validate Volcano planner.
  RuleSet ruleSet = RuleSets.ofList(
      new MockEnumerableJoinRule(hint), // Rule to validate the hint.
      FilterProjectTransposeRule.INSTANCE,
      FilterMergeRule.INSTANCE,
      ProjectMergeRule.INSTANCE,
      EnumerableRules.ENUMERABLE_JOIN_RULE,
      EnumerableRules.ENUMERABLE_PROJECT_RULE,
      EnumerableRules.ENUMERABLE_FILTER_RULE,
      EnumerableRules.ENUMERABLE_SORT_RULE,
      EnumerableRules.ENUMERABLE_LIMIT_RULE,
      EnumerableRules.ENUMERABLE_TABLE_SCAN_RULE);
  Program program = Programs.of(ruleSet);
  RelTraitSet toTraits = rel
      .getCluster()
      .traitSet()
      .replace(EnumerableConvention.INSTANCE);

  program.run(planner, rel, toTraits,
      Collections.emptyList(), Collections.emptyList());
}
 

protected Program createProgram() {
  final Program program1 =
      (planner, rel, requiredOutputTraits, materializations, lattices) -> {
        //override the default ruleSet of planner here
        planner.setRoot(rel);
        for (RelOptMaterialization materialization : materializations) {
          planner.addMaterialization(materialization);
        }
        for (RelOptLattice lattice : lattices) {
          planner.addLattice(lattice);
        }

        final RelNode rootRel2 =
            rel.getTraitSet().equals(requiredOutputTraits)
                ? rel
                : planner.changeTraits(rel, requiredOutputTraits);
        assert rootRel2 != null;

        planner.setRoot(rootRel2);
        final RelOptPlanner planner2 = planner.chooseDelegate();
        final RelNode rootRel3 = planner2.findBestExp();
        assert rootRel3 != null : "could not implement exp";
        return rootRel3;
      };
  DefaultRelMetadataProvider metadataProvider = DefaultRelMetadataProvider
      .INSTANCE;
  return Programs.sequence(Programs.subQuery(metadataProvider),
      new Programs.DecorrelateProgram(),
      new Programs.TrimFieldsProgram(),
      program1,

      // Second planner pass to do physical "tweaks". This the first time
      // that EnumerableCalcRel is introduced.
      Programs.calc(metadataProvider));
}
 

private static Planner getPlanner(List<RelTraitDef> traitDefs,
    SqlParser.Config parserConfig, Program... programs) {
  final SchemaPlus rootSchema = Frameworks.createRootSchema(true);
  final FrameworkConfig config = Frameworks.newConfigBuilder()
      .parserConfig(parserConfig)
      .defaultSchema(CalciteAssert.addSchema(rootSchema, CalciteAssert.SchemaSpec.HR))
      .traitDefs(traitDefs)
      .programs(programs)
      .build();
  return Frameworks.getPlanner(config);
}
 

private static RelNode toPhysical(RelNode rel) {
  final RelOptPlanner planner = rel.getCluster().getPlanner();
  planner.clear();
  for (RelOptRule rule : RULE_SET) {
    planner.addRule(rule);
  }

  final Program program = Programs.of(RuleSets.ofList(planner.getRules()));
  return program.run(planner, rel, rel.getTraitSet().replace(EnumerableConvention.INSTANCE),
      ImmutableList.of(), ImmutableList.of());
}
 

protected Program getProgram() {
  // Allow a test to override the default program.
  final Holder<Program> holder = Holder.of(null);
  Hook.PROGRAM.run(holder);
  if (holder.get() != null) {
    return holder.get();
  }

  return Programs.standard();
}
 

public RelNode transform(int ruleSetIndex, RelTraitSet requiredOutputTraits,
    RelNode rel) {
  ensure(State.STATE_5_CONVERTED);
  rel.getCluster().setMetadataProvider(
      new CachingRelMetadataProvider(
          rel.getCluster().getMetadataProvider(),
          rel.getCluster().getPlanner()));
  Program program = programs.get(ruleSetIndex);
  return program.run(planner, rel, requiredOutputTraits, ImmutableList.of(),
      ImmutableList.of());
}
 

private List<Program> getPrograms() {
  ImmutableList.Builder<Program> builder
      = ImmutableList.builder();
  for (RuleSet ruleSet: getRules()) {
    builder.add(Programs.sequence(
        new EnumerableProgram(ruleSet, this.context, this.plannerHolder),
        Programs.CALC_PROGRAM));
  }
  return builder.build();
}
 

@SuppressWarnings("Guava") // Must conform to Calcite's API
public static Function<Holder<Program>, Void> prepend(Program program) {
	return (holder) -> {
		if (holder == null) {
			throw new IllegalStateException("No program holder");
		}
		Program chain = holder.get();
		if (chain == null) {
			chain = Programs.standard();
		}
		holder.set(new SequenceProgram(program, chain));
		return null;
	};
}
 

@Test
public void testPrepend_usesTheDefaultProgramIfGivenNull() {
	Holder<Program> holder = Holder.of(null);

	SequenceProgram.prepend(subProgram).apply(holder);

	Program newProgram = holder.get();
	Assert.assertTrue(newProgram instanceof SequenceProgram);
	ImmutableList<Program> subPrograms = ((SequenceProgram) newProgram).getPrograms();
	// Not much we can check about the program; it is built inline and has no particular features
	Assert.assertNotNull(subPrograms.get(1));
}
 

@Test
public void testPrepend_returnsAHookForAddingTheProgram() {
	Program originalProgram = Mockito.mock(Program.class);
	Holder<Program> holder = Holder.of(originalProgram);

	SequenceProgram.prepend(subProgram).apply(holder);

	Program newProgram = holder.get();
	Assert.assertTrue(newProgram instanceof SequenceProgram);
	ImmutableList<Program> subPrograms = ((SequenceProgram) newProgram).getPrograms();
	Assert.assertEquals(2, subPrograms.size());
	Assert.assertSame(subProgram, subPrograms.get(0));
	Assert.assertSame(originalProgram, subPrograms.get(1));
}
 

@Test
public void testRun_chainsSubPrograms() {
	Program subProgram2 = Mockito.mock(Program.class);
	RelNode node2 = Mockito.mock(RelNode.class);
	RelNode node3 = Mockito.mock(RelNode.class);
	program = new SequenceProgram(subProgram, subProgram2);
	Mockito.doReturn(node2).when(subProgram).run(Mockito.any(), Mockito.same(inNode),
			Mockito.any(), Mockito.anyList(), Mockito.anyList());
	Mockito.doReturn(node3).when(subProgram2).run(Mockito.any(), Mockito.same(node2),
			Mockito.any(), Mockito.anyList(), Mockito.anyList());

	RelNode result = program.run(planner, inNode, relTraitSet, relOptMaterializationList, relOptLatticeList);

	Assert.assertSame(node3, result);
}
 

@Before
public void setupMocks() {
	planner = Mockito.mock(RelOptPlanner.class);
	relTraitSet = RelTraitSet.createEmpty();
	subProgram = Mockito.mock(Program.class);
	program = new SequenceProgram(subProgram);
	inNode = Mockito.mock(RelNode.class);
	relOptMaterializationList = Arrays.asList();
	relOptLatticeList = Arrays.asList();
}
 

@Override
public RelNode run(RelOptPlanner relOptPlanner,
		RelNode relNode,
		RelTraitSet relTraitSet,
		List<RelOptMaterialization> relOptMaterializationList,
		List<RelOptLattice> relOptLatticeList) {
	for (Program program : programs) {
		relNode = program.run(relOptPlanner, relNode, relTraitSet, relOptMaterializationList, relOptLatticeList);
		logger.debug("After running " + program + ":\n" + RelOptUtil.toString(relNode));
	}
	return relNode;
}
 

public ImmutableList<Program> getPrograms() {
	return programs;
}
 

@SuppressWarnings("Guava") // Must conform to Calcite's API
public static Function<Holder<Program>, Void> program() {
	return SequenceProgram.prepend(PROGRAM);
}
 

public SequenceProgram(Program... programs) {
	this.programs = ImmutableList.copyOf(programs);
}
 

@Override
public ImmutableList<Program> getPrograms() {
    return MycatCalciteSupport.INSTANCE.config.getPrograms();
}
 

/** Test case for
 * <a href="https://issues.apache.org/jira/browse/CALCITE-2941">[CALCITE-2941]
 * EnumerableLimitRule on Sort with no collation creates EnumerableLimit with
 * wrong traitSet and cluster</a>.
 */
@Test void enumerableLimitOnEmptySort() throws Exception {
  RuleSet prepareRules =
      RuleSets.ofList(
          EnumerableRules.ENUMERABLE_FILTER_RULE,
          EnumerableRules.ENUMERABLE_SORT_RULE,
          EnumerableRules.ENUMERABLE_LIMIT_RULE,
          EnumerableRules.ENUMERABLE_TABLE_SCAN_RULE);
  SchemaPlus rootSchema = Frameworks.createRootSchema(true);
  SchemaPlus defSchema = rootSchema.add("hr", new HrClusteredSchema());
  FrameworkConfig config = Frameworks.newConfigBuilder()
      .parserConfig(SqlParser.Config.DEFAULT)
      .defaultSchema(defSchema)
      .traitDefs(ConventionTraitDef.INSTANCE, RelCollationTraitDef.INSTANCE)
      .programs(Programs.of(prepareRules))
      .build();

  RelBuilder builder = RelBuilder.create(config);
  RelNode planBefore = builder
      .scan("hr", "emps")
      .sort(builder.field(0)) // will produce collation [0] in the plan
      .filter(
          builder.notEquals(
              builder.field(0),
              builder.literal(100)))
      .limit(1, 5) // force a limit inside an "empty" Sort (with no collation)
      .build();

  RelTraitSet desiredTraits = planBefore.getTraitSet()
      .replace(EnumerableConvention.INSTANCE);
  Program program = Programs.of(prepareRules);
  RelNode planAfter = program.run(planBefore.getCluster().getPlanner(), planBefore,
      desiredTraits, ImmutableList.of(), ImmutableList.of());

  // verify that the collation [0] is not lost in the final plan
  final RelCollation collation =
      planAfter.getTraitSet().getTrait(RelCollationTraitDef.INSTANCE);
  assertThat(collation, notNullValue());
  final List<RelFieldCollation> fieldCollationList =
      collation.getFieldCollations();
  assertThat(fieldCollationList, notNullValue());
  assertThat(fieldCollationList.size(), is(1));
  assertThat(fieldCollationList.get(0).getFieldIndex(), is(0));
}