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.adapter.enumerable;
  
  import org.apache.calcite.adapter.enumerable.impl.AggResultContextImpl;
  import org.apache.calcite.adapter.java.JavaTypeFactory;
  import org.apache.calcite.linq4j.Ord;
  import org.apache.calcite.linq4j.function.Function0;
  import org.apache.calcite.linq4j.function.Function2;
  import org.apache.calcite.linq4j.tree.BlockBuilder;
  import org.apache.calcite.linq4j.tree.Expression;
  import org.apache.calcite.linq4j.tree.Expressions;
  import org.apache.calcite.linq4j.tree.ParameterExpression;
  import org.apache.calcite.plan.RelOptCluster;
  import org.apache.calcite.plan.RelTraitSet;
  import org.apache.calcite.rel.RelCollation;
  import org.apache.calcite.rel.RelCollations;
  import org.apache.calcite.rel.RelFieldCollation;
  import org.apache.calcite.rel.RelNode;
  import org.apache.calcite.rel.core.Aggregate;
  import org.apache.calcite.rel.core.AggregateCall;
  import org.apache.calcite.rex.RexUtil;
  import org.apache.calcite.util.BuiltInMethod;
  import org.apache.calcite.util.ImmutableBitSet;
  import org.apache.calcite.util.Pair;
  import org.apache.calcite.util.Util;
  import org.apache.calcite.util.mapping.Mappings;
  
  import com.google.common.collect.ImmutableList;
  
  import org.checkerframework.checker.nullness.qual.Nullable;
  
  import java.lang.reflect.Type;
  import java.util.ArrayList;
  import java.util.List;
  
  import static java.util.Objects.requireNonNull;
  
  /** Sort based physical implementation of {@link Aggregate} in
   * {@link EnumerableConvention enumerable calling convention}. */
  public class EnumerableSortedAggregate extends EnumerableAggregateBase implements EnumerableRel {
    public EnumerableSortedAggregate(
        RelOptCluster cluster,
        RelTraitSet traitSet,
        RelNode input,
        ImmutableBitSet groupSet,
        @Nullable List<ImmutableBitSet> groupSets,
        List<AggregateCall> aggCalls) {
      super(cluster, traitSet, ImmutableList.of(), input, groupSet, groupSets, aggCalls);
      assert getConvention() instanceof EnumerableConvention;
    }
  
    @Override public EnumerableSortedAggregate copy(RelTraitSet traitSet, RelNode input,
        ImmutableBitSet groupSet,
        @Nullable List<ImmutableBitSet> groupSets, List<AggregateCall> aggCalls) {
      return new EnumerableSortedAggregate(getCluster(), traitSet, input,
          groupSet, groupSets, aggCalls);
    }
  
    @Override public @Nullable Pair<RelTraitSet, List<RelTraitSet>> passThroughTraits(
        final RelTraitSet required) {
      if (!isSimple(this)) {
        return null;
      }
  
      RelTraitSet inputTraits = getInput().getTraitSet();
      RelCollation collation =
          requireNonNull(required.getCollation(),
              () -> "collation trait is null, required traits are " + required);
      ImmutableBitSet requiredKeys = ImmutableBitSet.of(RelCollations.ordinals(collation));
      ImmutableBitSet groupKeys = ImmutableBitSet.range(groupSet.cardinality());
  
      Mappings.TargetMapping mapping =
          Mappings.source(groupSet.toList(), input.getRowType().getFieldCount());
  
      if (requiredKeys.equals(groupKeys)) {
        RelCollation inputCollation = RexUtil.apply(mapping, collation);
        return Pair.of(required, ImmutableList.of(inputTraits.replace(inputCollation)));
      } else if (groupKeys.contains(requiredKeys)) {
        // group by a,b,c order by c,b
        List<RelFieldCollation> list = new ArrayList<>(collation.getFieldCollations());
        groupKeys.except(requiredKeys).forEachInt(k ->
            list.add(new RelFieldCollation(k)));
        RelCollation aggCollation = RelCollations.of(list);
        RelCollation inputCollation = RexUtil.apply(mapping, aggCollation);
        return Pair.of(traitSet.replace(aggCollation),
            ImmutableList.of(inputTraits.replace(inputCollation)));
      }
  
      // Group keys doesn't contain all the required keys, e.g.
      // group by a,b order by a,b,c
      // nothing we can do to propagate traits to child nodes.
      return null;
    }
  
    @Override public Result implement(EnumerableRelImplementor implementor, Prefer pref) {
      if (!Aggregate.isSimple(this)) {
        throw Util.needToImplement("EnumerableSortedAggregate");
      }
  
      final JavaTypeFactory typeFactory = implementor.getTypeFactory();
      final BlockBuilder builder = new BlockBuilder();
      final EnumerableRel child = (EnumerableRel) getInput();
      final Result result = implementor.visitChild(this, 0, child, pref);
      Expression childExp =
          builder.append(
              "child",
              result.block);
  
      final PhysType physType =
          PhysTypeImpl.of(
              typeFactory, getRowType(), pref.preferCustom());
  
      final PhysType inputPhysType = result.physType;
  
      ParameterExpression parameter =
          Expressions.parameter(inputPhysType.getJavaRowType(), "a0");
  
      final PhysType keyPhysType =
          inputPhysType.project(groupSet.asList(), getGroupType() != Group.SIMPLE,
              JavaRowFormat.LIST);
      final int groupCount = getGroupCount();
  
      final List<AggImpState> aggs = new ArrayList<>(aggCalls.size());
      for (Ord<AggregateCall> call : Ord.zip(aggCalls)) {
        aggs.add(new AggImpState(call.i, call.e, false));
      }
  
      // Function0<Object[]> accumulatorInitializer =
      //     new Function0<Object[]>() {
      //         public Object[] apply() {
      //             return new Object[] {0, 0};
      //         }
      //     };
      final List<Expression> initExpressions = new ArrayList<>();
      final BlockBuilder initBlock = new BlockBuilder();
  
      final List<Type> aggStateTypes =
          createAggStateTypes(initExpressions, initBlock, aggs, typeFactory);
  
      final PhysType accPhysType =
          PhysTypeImpl.of(typeFactory,
              typeFactory.createSyntheticType(aggStateTypes));
  
      declareParentAccumulator(initExpressions, initBlock, accPhysType);
  
      final Expression accumulatorInitializer =
          builder.append("accumulatorInitializer",
              Expressions.lambda(
                  Function0.class,
                  initBlock.toBlock()));
  
      // Function2<Object[], Employee, Object[]> accumulatorAdder =
      //     new Function2<Object[], Employee, Object[]>() {
      //         public Object[] apply(Object[] acc, Employee in) {
      //              acc[0] = ((Integer) acc[0]) + 1;
      //              acc[1] = ((Integer) acc[1]) + in.salary;
      //             return acc;
      //         }
      //     };
      final ParameterExpression inParameter =
          Expressions.parameter(inputPhysType.getJavaRowType(), "in");
      final ParameterExpression acc_ =
          Expressions.parameter(accPhysType.getJavaRowType(), "acc");
  
      createAccumulatorAdders(
          inParameter, aggs, accPhysType, acc_, inputPhysType, builder, implementor, typeFactory);
  
      final ParameterExpression lambdaFactory =
          Expressions.parameter(AggregateLambdaFactory.class,
              builder.newName("lambdaFactory"));
  
      implementLambdaFactory(builder, inputPhysType, aggs, accumulatorInitializer,
          false, lambdaFactory);
  
      final BlockBuilder resultBlock = new BlockBuilder();
      final List<Expression> results = Expressions.list();
      final ParameterExpression key_;
      final Type keyType = keyPhysType.getJavaRowType();
      key_ = Expressions.parameter(keyType, "key");
      for (int j = 0; j < groupCount; j++) {
        final Expression ref = keyPhysType.fieldReference(key_, j);
        results.add(ref);
      }
  
      for (final AggImpState agg : aggs) {
        results.add(
            agg.implementor.implementResult(
                requireNonNull(agg.context, () -> "agg.context is null for " + agg),
                new AggResultContextImpl(resultBlock, agg.call,
                    requireNonNull(agg.state, () -> "agg.state is null for " + agg),
                    key_,
                    keyPhysType)));
      }
      resultBlock.add(physType.record(results));
  
      final Expression keySelector_ =
          builder.append("keySelector",
              inputPhysType.generateSelector(parameter, groupSet.asList(),
                  keyPhysType.getFormat()));
      // Generate the appropriate key Comparator. In the case of NULL values
      // in group keys, the comparator must be able to support NULL values by giving a
      // consistent sort ordering.
      final Expression comparator =
          keyPhysType.generateComparator(
              requireNonNull(getTraitSet().getCollation(),
                  () -> "getTraitSet().getCollation() is null; traits are "
                      + getTraitSet()));
  
      final Expression resultSelector_ =
          builder.append("resultSelector",
              Expressions.lambda(Function2.class, resultBlock.toBlock(), key_,
                  acc_));
  
      builder.add(
          Expressions.return_(null,
              Expressions.call(childExp,
                  BuiltInMethod.SORTED_GROUP_BY.method,
                  Expressions.list(keySelector_,
                      Expressions.call(lambdaFactory,
                          BuiltInMethod.AGG_LAMBDA_FACTORY_ACC_INITIALIZER.method),
                      Expressions.call(lambdaFactory,
                          BuiltInMethod.AGG_LAMBDA_FACTORY_ACC_ADDER.method),
                      Expressions.call(lambdaFactory,
                          BuiltInMethod.AGG_LAMBDA_FACTORY_ACC_RESULT_SELECTOR.method,
                          resultSelector_), comparator))));
  
      return implementor.result(physType, builder.toBlock());
    }
  }
  