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);
    }
  }
}

```

Reply via email to