content: /*
   * Licensed to the Apache Software Foundation (ASF) under one or more
   * contributor license agreements.  See the NOTICE file distributed with
   * this work for additional information regarding copyright ownership.
   * The ASF licenses this file to you under the Apache License, Version 2.0
   * (the "License"); you may not use this file except in compliance with
   * the License.  You may obtain a copy of the License at
   *
   * http://www.apache.org/licenses/LICENSE-2.0
   *
   * Unless required by applicable law or agreed to in writing, software
   * distributed under the License is distributed on an "AS IS" BASIS,
   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   * See the License for the specific language governing permissions and
   * limitations under the License.
   */
  package org.apache.calcite.rel.rules;
  
  import org.apache.calcite.plan.RelOptPredicateList;
  import org.apache.calcite.plan.RelOptRuleCall;
  import org.apache.calcite.plan.RelRule;
  import org.apache.calcite.rel.RelNode;
  import org.apache.calcite.rel.core.Aggregate;
  import org.apache.calcite.rel.core.AggregateCall;
  import org.apache.calcite.rel.logical.LogicalAggregate;
  import org.apache.calcite.rel.logical.LogicalProject;
  import org.apache.calcite.rel.metadata.RelMetadataQuery;
  import org.apache.calcite.rel.type.RelDataType;
  import org.apache.calcite.rel.type.RelDataTypeField;
  import org.apache.calcite.rex.RexBuilder;
  import org.apache.calcite.rex.RexInputRef;
  import org.apache.calcite.rex.RexNode;
  import org.apache.calcite.tools.RelBuilder;
  import org.apache.calcite.tools.RelBuilderFactory;
  import org.apache.calcite.util.ImmutableBitSet;
  import org.apache.calcite.util.Pair;
  
  import org.immutables.value.Value;
  
  import java.util.ArrayList;
  import java.util.List;
  import java.util.NavigableMap;
  import java.util.TreeMap;
  
  /**
   * Planner rule that removes constant keys from an
   * {@link org.apache.calcite.rel.core.Aggregate}.
   *
   * <p>Constant fields are deduced using
   * {@link RelMetadataQuery#getPulledUpPredicates(RelNode)}; the input does not
   * need to be a {@link org.apache.calcite.rel.core.Project}.
   *
   * <p>This rule never removes the last column, because {@code Aggregate([])}
   * returns 1 row even if its input is empty.
   *
   * <p>Since the transformed relational expression has to match the original
   * relational expression, the constants are placed in a projection above the
   * reduced aggregate. If those constants are not used, another rule will remove
   * them from the project.
   */
  @Value.Enclosing
  public class AggregateProjectPullUpConstantsRule
      extends RelRule<AggregateProjectPullUpConstantsRule.Config>
      implements SubstitutionRule {
  
    /** Creates an AggregateProjectPullUpConstantsRule. */
    protected AggregateProjectPullUpConstantsRule(Config config) {
      super(config);
    }
  
    @Deprecated // to be removed before 2.0
    public AggregateProjectPullUpConstantsRule(
        Class<? extends Aggregate> aggregateClass,
        Class<? extends RelNode> inputClass,
        RelBuilderFactory relBuilderFactory, String description) {
      this(Config.DEFAULT
          .withRelBuilderFactory(relBuilderFactory)
          .withDescription(description)
          .as(Config.class)
          .withOperandFor(aggregateClass, inputClass));
    }
  
    //~ Methods ----------------------------------------------------------------
  
    @Override public void onMatch(RelOptRuleCall call) {
      final Aggregate aggregate = call.rel(0);
      final RelNode input = call.rel(1);
  
      final int groupCount = aggregate.getGroupCount();
      if (groupCount == 1) {
        // No room for optimization since we cannot convert from non-empty
        // GROUP BY list to the empty one.
        return;
      }
  
      final RexBuilder rexBuilder = aggregate.getCluster().getRexBuilder();
      final RelMetadataQuery mq = call.getMetadataQuery();
      final RelOptPredicateList predicates =
          mq.getPulledUpPredicates(aggregate.getInput());
      if (RelOptPredicateList.isEmpty(predicates)) {
        return;
      }
      final NavigableMap<Integer, RexNode> map = new TreeMap<>();
      for (int key : aggregate.getGroupSet()) {
        final RexInputRef ref =
            rexBuilder.makeInputRef(aggregate.getInput(), key);
        if (predicates.constantMap.containsKey(ref)) {
          map.put(key, predicates.constantMap.get(ref));
        }
      }
  
      // None of the group expressions are constant. Nothing to do.
      if (map.isEmpty()) {
        return;
      }
  
      if (groupCount == map.size()) {
        // At least a single item in group by is required.
        // Otherwise "GROUP BY 1, 2" might be altered to "GROUP BY ()".
        // Removing of the first element is not optimal here,
        // however it will allow us to use fast path below (just trim
        // groupCount).
        map.remove(map.navigableKeySet().first());
      }
  
      ImmutableBitSet newGroupSet = aggregate.getGroupSet();
      for (int key : map.keySet()) {
        newGroupSet = newGroupSet.clear(key);
      }
  
      // If the constants are on the trailing edge of the group list, we just
      // reduce the group count.
      final RelBuilder relBuilder = call.builder();
      relBuilder.push(input);
  
      // Clone aggregate calls.
      final List<AggregateCall> newAggCalls = new ArrayList<>();
      for (AggregateCall aggCall : aggregate.getAggCallList()) {
        newAggCalls.add(
            aggCall.adaptTo(input, aggCall.getArgList(), aggCall.filterArg,
                aggregate.hasEmptyGroup(), newGroupSet.isEmpty()));
      }
      relBuilder.aggregate(relBuilder.groupKey(newGroupSet), newAggCalls);
  
      // Create a projection back again.
      List<Pair<RexNode, String>> projects = new ArrayList<>();
      int source = 0;
      for (RelDataTypeField field : aggregate.getRowType().getFieldList()) {
        RexNode expr;
        final int i = field.getIndex();
        if (i >= groupCount) {
          // Aggregate expressions' names and positions are unchanged.
          expr = relBuilder.field(i - map.size());
        } else {
          int pos = aggregate.getGroupSet().nth(i);
          RexNode rexNode = map.get(pos);
          if (rexNode != null) {
            // Re-generate the constant expression in the project.
            RelDataType originalType =
                aggregate.getRowType().getFieldList().get(projects.size()).getType();
            if (!originalType.equals(rexNode.getType())) {
              expr = rexBuilder.makeCast(originalType, rexNode, true, false);
            } else {
              expr = rexNode;
            }
          } else {
            // Project the aggregation expression, in its original
            // position.
            expr = relBuilder.field(source);
            ++source;
          }
        }
        projects.add(Pair.of(expr, field.getName()));
      }
      relBuilder.project(Pair.left(projects), Pair.right(projects)); // inverse
      call.transformTo(relBuilder.build());
      call.getPlanner().prune(aggregate);
    }
  
  
    /** Rule configuration. */
    @Value.Immutable
    public interface Config extends RelRule.Config {
      Config DEFAULT = ImmutableAggregateProjectPullUpConstantsRule.Config.of();
      Config ANY = DEFAULT.withOperandFor(LogicalAggregate.class, RelNode.class);
  
      @Override default AggregateProjectPullUpConstantsRule toRule() {
        return new AggregateProjectPullUpConstantsRule(this);
      }
  
      @Override @Value.Default default OperandTransform operandSupplier() {
        return b0 ->
            b0.operand(LogicalAggregate.class)
                .predicate(Aggregate::isSimple)
                .oneInput(b1 ->
                    b1.operand(LogicalProject.class).anyInputs());
      }
  
      /** Defines an operand tree for the given classes.
       *
       * @param aggregateClass Aggregate class
       * @param inputClass Input class, such as {@link LogicalProject}
       */
      default Config withOperandFor(Class<? extends Aggregate> aggregateClass,
          Class<? extends RelNode> inputClass) {
        return withOperandSupplier(b0 ->
            b0.operand(aggregateClass)
                .predicate(Aggregate::isSimple)
                .oneInput(b1 ->
                    b1.operand(inputClass).anyInputs()))
            .as(Config.class);
      }
    }
  }
  