Hello! We've been attempting to implement a simple optimization rule in which we duplicate a limit through a UNION ALL to reduce the amount of data we need to fetch for a query.
Starting from something like ``` LogicalSort[fetch = 10] LogicalUnion[all=true] <Input 1> <Input 2> ``` We're trying to turn it into ``` LogicalSort[fetch = 10] LogicalUnion[all=true] LogicalSort[fetch = 10] <Input 1> LogicalSort[fetch = 10] <Input 2> ``` This, somewhat expectedly, causes issues with the VolcanoPlanner because the newly generated relation is also a candidate for our rule so we end up with an infinite planning loop. We tried to take inspiration from the JoinCommuteRule, which uses the ensureRegistered method to prevent this (or at least that's what we think it's doing). Unfortunately, in our case this appears to be insufficient. I would appreciate any pointers and or suggestions around this. I've included the code for the raw rule below. ``` @Value.Enclosing public class LimitThroughUnionRule extends RelRule<LimitThroughUnionRule.Config> implements SubstitutionRule { public static final LimitThroughUnionRule INSTANCE = LimitThroughUnionRule.Config.DEFAULT.toRule(); protected LimitThroughUnionRule(Config config) { super(config); } private RelNode pushLimitThrough(RelBuilder relBuilder, Sort sort, Union union) { for (RelNode unionInput : union.getInputs()) { relBuilder.push(unionInput).sortLimit(sort.offset, sort.fetch, Collections.emptyList()); } relBuilder.union(true); relBuilder.sortLimit(sort.offset, sort.fetch, sort.getSortExps()); return relBuilder.build(); } @Override public void onMatch(RelOptRuleCall call) { Sort sort = call.rel(0); Union union = call.rel(1); RelNode newNode = pushLimitThrough(call.builder(), sort, union); RelNode nextNode = pushLimitThrough(call.builder(), (Sort) newNode, (Union) newNode.getInput(0)); call.transformTo(newNode); call.getPlanner().ensureRegistered(nextNode, newNode); } @Value.Immutable public interface Config extends RelRule.Config { LimitThroughUnionRule.Config DEFAULT = ImmutableLimitThroughUnionRule.Config.builder() .operandSupplier( b0 -> b0.operand(Sort.class) .predicate(sort -> !(RelOptUtil.isOrder(sort) || RelOptUtil.isOffset(sort))) .oneInput(u -> u.operand(Union.class).anyInputs())) .build(); @Override default LimitThroughUnionRule toRule() { return new LimitThroughUnionRule(this); } } } ```