/*
 * 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.drill.exec.planner.logical;

import org.apache.drill.exec.planner.sql.DrillCalciteSqlSumEmptyIsZeroAggFunctionWrapper;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.Lists;
import com.google.common.collect.Maps;
import org.apache.calcite.plan.Convention;
import org.apache.calcite.plan.RelOptCluster;
import org.apache.calcite.plan.RelOptRule;
import org.apache.calcite.plan.RelOptRuleCall;
import org.apache.calcite.plan.RelOptRuleOperand;
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.core.Window;
import org.apache.calcite.rel.logical.LogicalAggregate;
import org.apache.calcite.rel.type.RelDataType;
import org.apache.calcite.rel.type.RelDataTypeFactory;
import org.apache.calcite.rel.type.RelDataTypeField;
import org.apache.calcite.rex.RexBuilder;
import org.apache.calcite.rex.RexLiteral;
import org.apache.calcite.rex.RexNode;
import org.apache.calcite.sql.SqlAggFunction;
import org.apache.calcite.sql.SqlKind;
import org.apache.calcite.sql.SqlOperator;
import org.apache.calcite.sql.SqlOperatorBinding;
import org.apache.calcite.sql.fun.SqlAvgAggFunction;
import org.apache.calcite.sql.fun.SqlCountAggFunction;
import org.apache.calcite.sql.fun.SqlStdOperatorTable;
import org.apache.calcite.sql.fun.SqlSumAggFunction;
import org.apache.calcite.sql.fun.SqlSumEmptyIsZeroAggFunction;
import org.apache.calcite.sql.type.SqlReturnTypeInference;
import org.apache.calcite.sql.type.SqlTypeName;
import org.apache.calcite.util.CompositeList;
import org.apache.calcite.util.ImmutableIntList;
import org.apache.calcite.util.Util;
import org.apache.drill.exec.planner.physical.PlannerSettings;
import org.apache.drill.exec.planner.sql.DrillCalciteSqlAggFunctionWrapper;
import org.apache.drill.exec.planner.sql.DrillSqlOperator;
import org.apache.drill.exec.planner.sql.TypeInferenceUtils;
import org.apache.drill.exec.planner.sql.parser.DrillCalciteWrapperUtility;

import java.math.BigDecimal;
import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * Rule to reduce aggregates to simpler forms. Currently only AVG(x) to
 * SUM(x)/COUNT(x), but eventually will handle others such as STDDEV.
 */
public class DrillReduceAggregatesRule extends RelOptRule {
  //~ Static fields/initializers ---------------------------------------------

  /**
   * The singleton.
   */
  public static final DrillReduceAggregatesRule INSTANCE =
      new DrillReduceAggregatesRule(operand(LogicalAggregate.class, any()));
  public static final DrillConvertSumToSumZero INSTANCE_SUM =
      new DrillConvertSumToSumZero(operand(DrillAggregateRel.class, any()));

  public static final DrillConvertWindowSumToSumZero INSTANCE_WINDOW_SUM =
          new DrillConvertWindowSumToSumZero(operand(DrillWindowRel.class, any()));

  private static final DrillSqlOperator CastHighOp = new DrillSqlOperator("CastHigh", 1, false,
      new SqlReturnTypeInference() {
        @Override
        public RelDataType inferReturnType(SqlOperatorBinding opBinding) {
          return TypeInferenceUtils.createCalciteTypeWithNullability(
              opBinding.getTypeFactory(),
              SqlTypeName.ANY,
              opBinding.getOperandType(0).isNullable());
      }
  }, false);

  //~ Constructors -----------------------------------------------------------

  protected DrillReduceAggregatesRule(RelOptRuleOperand operand) {
    super(operand, DrillRelFactories.LOGICAL_BUILDER, null);
  }

  //~ Methods ----------------------------------------------------------------

  @Override
  public boolean matches(RelOptRuleCall call) {
    if (!super.matches(call)) {
      return false;
    }
    Aggregate oldAggRel = (Aggregate) call.rels[0];
    return containsAvgStddevVarCall(oldAggRel.getAggCallList());
  }

  @Override
  public void onMatch(RelOptRuleCall ruleCall) {
    Aggregate oldAggRel = (Aggregate) ruleCall.rels[0];
    reduceAggs(ruleCall, oldAggRel);
  }

  /**
   * Returns whether any of the aggregates are calls to AVG, STDDEV_*, VAR_*.
   *
   * @param aggCallList List of aggregate calls
   */
  private boolean containsAvgStddevVarCall(List<AggregateCall> aggCallList) {
    for (AggregateCall call : aggCallList) {
      SqlAggFunction sqlAggFunction = DrillCalciteWrapperUtility.extractSqlOperatorFromWrapper(call.getAggregation());
      if (sqlAggFunction instanceof SqlAvgAggFunction
          || sqlAggFunction instanceof SqlSumAggFunction) {
        return true;
      }
    }
    return false;
  }

  /*
  private boolean isMatch(AggregateCall call) {
    if (call.getAggregation() instanceof SqlAvgAggFunction) {
      final SqlAvgAggFunction.Subtype subtype =
          ((SqlAvgAggFunction) call.getAggregation()).getSubtype();
      return (subtype == SqlAvgAggFunction.Subtype.AVG);
    }
    return false;
  }
 */

  /**
   * Reduces all calls to AVG, STDDEV_POP, STDDEV_SAMP, VAR_POP, VAR_SAMP in
   * the aggregates list to.
   *
   * <p>It handles newly generated common subexpressions since this was done
   * at the sql2rel stage.
   */
  private void reduceAggs(
      RelOptRuleCall ruleCall,
      Aggregate oldAggRel) {
    RexBuilder rexBuilder = oldAggRel.getCluster().getRexBuilder();

    List<AggregateCall> oldCalls = oldAggRel.getAggCallList();
    final int nGroups = oldAggRel.getGroupCount();

    List<AggregateCall> newCalls = new ArrayList<>();
    Map<AggregateCall, RexNode> aggCallMapping =
        new HashMap<>();

    List<RexNode> projList = new ArrayList<>();

    // pass through group key
    for (int i = 0; i < nGroups; ++i) {
      projList.add(
          rexBuilder.makeInputRef(
              getFieldType(oldAggRel, i),
              i));
    }

    // List of input expressions. If a particular aggregate needs more, it
    // will add an expression to the end, and we will create an extra
    // project.
    RelNode input = oldAggRel.getInput();
    List<RexNode> inputExprs = new ArrayList<>();
    for (RelDataTypeField field : input.getRowType().getFieldList()) {
      inputExprs.add(
          rexBuilder.makeInputRef(
              field.getType(), inputExprs.size()));
    }

    // create new agg function calls and rest of project list together
    for (AggregateCall oldCall : oldCalls) {
      projList.add(
          reduceAgg(
              oldAggRel, oldCall, newCalls, aggCallMapping, inputExprs));
    }

    final int extraArgCount =
        inputExprs.size() - input.getRowType().getFieldCount();
    if (extraArgCount > 0) {
      input =
          relBuilderFactory
              .create(input.getCluster(), null)
              .push(input)
              .projectNamed(
                  inputExprs,
                  CompositeList.of(
                  input.getRowType().getFieldNames(),
                  Collections.nCopies(
                      extraArgCount,
                      null)),
                  true)
              .build();
    }
    Aggregate newAggRel =
        newAggregateRel(
            oldAggRel, input, newCalls);

    RelNode projectRel =
        relBuilderFactory
            .create(newAggRel.getCluster(), null)
            .push(newAggRel)
            .projectNamed(projList, oldAggRel.getRowType().getFieldNames(), true)
            .build();

    ruleCall.transformTo(projectRel);
  }

  private RexNode reduceAgg(
      Aggregate oldAggRel,
      AggregateCall oldCall,
      List<AggregateCall> newCalls,
      Map<AggregateCall, RexNode> aggCallMapping,
      List<RexNode> inputExprs) {
    final SqlAggFunction sqlAggFunction = DrillCalciteWrapperUtility.extractSqlOperatorFromWrapper(oldCall.getAggregation());
    if (sqlAggFunction instanceof SqlSumAggFunction) {
      // replace original SUM(x) with
      // case COUNT(x) when 0 then null else SUM0(x) end
      return reduceSum(oldAggRel, oldCall, newCalls, aggCallMapping);
    }
    if (sqlAggFunction instanceof SqlAvgAggFunction) {
      // for DECIMAL data types does not produce rewriting of complex calls,
      // since SUM returns value with 38 precision and further handling of the value
      // causes the loss of the scale
      if (oldCall.getType().getSqlTypeName() == SqlTypeName.DECIMAL) {
        return oldAggRel.getCluster().getRexBuilder().addAggCall(
            oldCall,
            oldAggRel.getGroupCount(),
            newCalls,
            aggCallMapping,
            ImmutableList.of(getFieldType(
                oldAggRel.getInput(),
                oldCall.getArgList().get(0))));
      }
      final SqlKind subtype = sqlAggFunction.getKind();
      switch (subtype) {
      case AVG:
        // replace original AVG(x) with SUM(x) / COUNT(x)
        return reduceAvg(
            oldAggRel, oldCall, newCalls, aggCallMapping);
      case STDDEV_POP:
        // replace original STDDEV_POP(x) with
        //   SQRT(
        //     (SUM(x * x) - SUM(x) * SUM(x) / COUNT(x))
        //     / COUNT(x))
        return reduceStddev(
            oldAggRel, oldCall, true, true, newCalls, aggCallMapping,
            inputExprs);
      case STDDEV_SAMP:
        // replace original STDDEV_SAMP(x) with
        //   SQRT(
        //     (SUM(x * x) - SUM(x) * SUM(x) / COUNT(x))
        //     / CASE COUNT(x) WHEN 1 THEN NULL ELSE COUNT(x) - 1 END)
        return reduceStddev(
            oldAggRel, oldCall, false, true, newCalls, aggCallMapping,
            inputExprs);
      case VAR_POP:
        // replace original VAR_POP(x) with
        //     (SUM(x * x) - SUM(x) * SUM(x) / COUNT(x))
        //     / COUNT(x)
        return reduceStddev(
            oldAggRel, oldCall, true, false, newCalls, aggCallMapping,
            inputExprs);
      case VAR_SAMP:
        // replace original VAR_SAMP(x) with
        //     (SUM(x * x) - SUM(x) * SUM(x) / COUNT(x))
        //     / CASE COUNT(x) WHEN 1 THEN NULL ELSE COUNT(x) - 1 END
        return reduceStddev(
            oldAggRel, oldCall, false, false, newCalls, aggCallMapping,
            inputExprs);
      default:
        throw Util.unexpected(subtype);
      }
    } else {
      // anything else:  preserve original call
      RexBuilder rexBuilder = oldAggRel.getCluster().getRexBuilder();
      final int nGroups = oldAggRel.getGroupCount();

      List<RelDataType> oldArgTypes = new ArrayList<>();
      List<Integer> ordinals = oldCall.getArgList();

      assert ordinals.size() <= inputExprs.size();

      for (int ordinal : ordinals) {
        oldArgTypes.add(inputExprs.get(ordinal).getType());
      }
      //to solve AggregateCall returns true with equals method but has
      // different RelDataTypes
      if (aggCallMapping.containsKey(oldCall) && !aggCallMapping.get(oldCall)
        .getType().equals(oldCall.getType())) {
        int index = newCalls.size() + nGroups;
        newCalls.add(oldCall);
        return rexBuilder.makeInputRef(oldCall.getType(), index);
      }
      return rexBuilder.addAggCall(
          oldCall,
          nGroups,
          newCalls,
          aggCallMapping,
          oldArgTypes);
    }
  }

  private RexNode reduceAvg(
      Aggregate oldAggRel,
      AggregateCall oldCall,
      List<AggregateCall> newCalls,
      Map<AggregateCall, RexNode> aggCallMapping) {
    final PlannerSettings plannerSettings = (PlannerSettings) oldAggRel.getCluster().getPlanner().getContext();
    final boolean isInferenceEnabled = plannerSettings.isTypeInferenceEnabled();
    final int nGroups = oldAggRel.getGroupCount();
    RelDataTypeFactory typeFactory =
        oldAggRel.getCluster().getTypeFactory();
    RexBuilder rexBuilder = oldAggRel.getCluster().getRexBuilder();
    int iAvgInput = oldCall.getArgList().get(0);
    RelDataType avgInputType =
        getFieldType(
            oldAggRel.getInput(),
            iAvgInput);
    RelDataType sumType =
        TypeInferenceUtils.getDrillSqlReturnTypeInference(SqlKind.SUM.name(),
            ImmutableList.of())
          .inferReturnType(oldCall.createBinding(oldAggRel));
    sumType =
        typeFactory.createTypeWithNullability(
            sumType,
            sumType.isNullable() || nGroups == 0);
    SqlAggFunction sumAgg =
        new DrillCalciteSqlSumEmptyIsZeroAggFunctionWrapper(
          new SqlSumEmptyIsZeroAggFunction(), sumType);
    AggregateCall sumCall = getAggCall(oldCall, sumAgg, sumType);
    final SqlCountAggFunction countAgg = (SqlCountAggFunction) SqlStdOperatorTable.COUNT;
    final RelDataType countType = countAgg.getReturnType(typeFactory);
    AggregateCall countCall = getAggCall(oldCall, countAgg, countType);

    RexNode tmpsumRef =
        rexBuilder.addAggCall(
            sumCall,
            nGroups,
            newCalls,
            aggCallMapping,
            ImmutableList.of(avgInputType));

    RexNode tmpcountRef =
        rexBuilder.addAggCall(
            countCall,
            nGroups,
            newCalls,
            aggCallMapping,
            ImmutableList.of(avgInputType));

    RexNode n = rexBuilder.makeCall(SqlStdOperatorTable.CASE,
        rexBuilder.makeCall(SqlStdOperatorTable.EQUALS,
            tmpcountRef, rexBuilder.makeExactLiteral(BigDecimal.ZERO)),
            rexBuilder.constantNull(),
            tmpsumRef);

    // NOTE:  these references are with respect to the output
    // of newAggRel
    /*
    RexNode numeratorRef =
        rexBuilder.makeCall(CastHighOp,
          rexBuilder.addAggCall(
              sumCall,
              nGroups,
              newCalls,
              aggCallMapping,
              ImmutableList.of(avgInputType))
        );
    */
    RexNode numeratorRef = rexBuilder.makeCall(CastHighOp,  n);

    RexNode denominatorRef =
        rexBuilder.addAggCall(
            countCall,
            nGroups,
            newCalls,
            aggCallMapping,
            ImmutableList.of(avgInputType));
    if (isInferenceEnabled) {
      return rexBuilder.makeCall(
          new DrillSqlOperator(
              "divide",
              2,
              true,
              oldCall.getType(), false),
          numeratorRef,
          denominatorRef);
    } else {
      final RexNode divideRef =
          rexBuilder.makeCall(
              SqlStdOperatorTable.DIVIDE,
              numeratorRef,
              denominatorRef);
      return rexBuilder.makeCast(
          typeFactory.createSqlType(SqlTypeName.ANY), divideRef);
    }
  }

  private static AggregateCall getAggCall(AggregateCall oldCall,
      SqlAggFunction aggFunction,
      RelDataType sumType) {
    return AggregateCall.create(aggFunction,
        oldCall.isDistinct(),
        oldCall.isApproximate(),
        oldCall.ignoreNulls(),
        oldCall.getArgList(),
        oldCall.filterArg,
        oldCall.distinctKeys,
        oldCall.getCollation(),
        sumType,
        null);
  }

  private RexNode reduceSum(
      Aggregate oldAggRel,
      AggregateCall oldCall,
      List<AggregateCall> newCalls,
      Map<AggregateCall, RexNode> aggCallMapping) {
    final PlannerSettings plannerSettings = (PlannerSettings) oldAggRel.getCluster().getPlanner().getContext();
    final boolean isInferenceEnabled = plannerSettings.isTypeInferenceEnabled();
    final int nGroups = oldAggRel.getGroupCount();
    RelDataTypeFactory typeFactory =
        oldAggRel.getCluster().getTypeFactory();
    RexBuilder rexBuilder = oldAggRel.getCluster().getRexBuilder();
    int arg = oldCall.getArgList().get(0);
    RelDataType argType =
        getFieldType(
            oldAggRel.getInput(),
            arg);
    final RelDataType sumType;
    final SqlAggFunction sumZeroAgg;
    if (isInferenceEnabled) {
      sumType = oldCall.getType();
    } else {
      sumType =
          typeFactory.createTypeWithNullability(
              oldCall.getType(), argType.isNullable());
    }
    sumZeroAgg = new DrillCalciteSqlSumEmptyIsZeroAggFunctionWrapper(
        new SqlSumEmptyIsZeroAggFunction(), sumType);
    AggregateCall sumZeroCall = getAggCall(oldCall, sumZeroAgg, sumType);
    final SqlCountAggFunction countAgg = (SqlCountAggFunction) SqlStdOperatorTable.COUNT;
    final RelDataType countType = countAgg.getReturnType(typeFactory);
    AggregateCall countCall = getAggCall(oldCall, countAgg, countType);
    // NOTE:  these references are with respect to the output
    // of newAggRel
    RexNode sumZeroRef =
        rexBuilder.addAggCall(
            sumZeroCall,
            nGroups,
            newCalls,
            aggCallMapping,
            ImmutableList.of(argType));
    if (!oldCall.getType().isNullable()) {
      // If SUM(x) is not nullable, the validator must have determined that
      // nulls are impossible (because the group is never empty and x is never
      // null). Therefore we translate to SUM0(x).
      return sumZeroRef;
    }
    RexNode countRef =
        rexBuilder.addAggCall(
            countCall,
            nGroups,
            newCalls,
            aggCallMapping,
            ImmutableList.of(argType));
    return rexBuilder.makeCall(SqlStdOperatorTable.CASE,
        rexBuilder.makeCall(SqlStdOperatorTable.EQUALS,
            countRef, rexBuilder.makeExactLiteral(BigDecimal.ZERO)),
            rexBuilder.constantNull(),
            sumZeroRef);
  }

  private RexNode reduceStddev(
      Aggregate oldAggRel,
      AggregateCall oldCall,
      boolean biased,
      boolean sqrt,
      List<AggregateCall> newCalls,
      Map<AggregateCall, RexNode> aggCallMapping,
      List<RexNode> inputExprs) {
    // stddev_pop(x) ==>
    //   power(
    //     (sum(x * x) - sum(x) * sum(x) / count(x))
    //     / count(x),
    //     .5)
    //
    // stddev_samp(x) ==>
    //   power(
    //     (sum(x * x) - sum(x) * sum(x) / count(x))
    //     / nullif(count(x) - 1, 0),
    //     .5)
    final PlannerSettings plannerSettings = (PlannerSettings) oldAggRel.getCluster().getPlanner().getContext();
    final boolean isInferenceEnabled = plannerSettings.isTypeInferenceEnabled();
    final int nGroups = oldAggRel.getGroupCount();
    RelDataTypeFactory typeFactory =
        oldAggRel.getCluster().getTypeFactory();
    final RexBuilder rexBuilder = oldAggRel.getCluster().getRexBuilder();

    assert oldCall.getArgList().size() == 1 : oldCall.getArgList();
    final int argOrdinal = oldCall.getArgList().get(0);
    final RelDataType argType =
        getFieldType(
            oldAggRel.getInput(),
            argOrdinal);

    // final RexNode argRef = inputExprs.get(argOrdinal);
    RexNode argRef = rexBuilder.makeCall(CastHighOp, inputExprs.get(argOrdinal));
    inputExprs.set(argOrdinal, argRef);

    final RexNode argSquared =
        rexBuilder.makeCall(
            SqlStdOperatorTable.MULTIPLY, argRef, argRef);
    final int argSquaredOrdinal = lookupOrAdd(inputExprs, argSquared);

    RelDataType sumType =
        TypeInferenceUtils.getDrillSqlReturnTypeInference(SqlKind.SUM.name(),
            ImmutableList.of())
          .inferReturnType(oldCall.createBinding(oldAggRel));
    sumType = typeFactory.createTypeWithNullability(sumType, true);
    final AggregateCall sumArgSquaredAggCall =
        AggregateCall.create(
            new DrillCalciteSqlAggFunctionWrapper(
                new SqlSumAggFunction(sumType), sumType),
            oldCall.isDistinct(),
            oldCall.isApproximate(),
            oldCall.ignoreNulls(),
            ImmutableIntList.of(argSquaredOrdinal),
            oldCall.filterArg,
            oldCall.distinctKeys,
            oldCall.getCollation(),
            sumType,
            null);
    final RexNode sumArgSquared =
        rexBuilder.addAggCall(
            sumArgSquaredAggCall,
            nGroups,
            newCalls,
            aggCallMapping,
            ImmutableList.of(argType));

    final AggregateCall sumArgAggCall =
        AggregateCall.create(
            new DrillCalciteSqlAggFunctionWrapper(
                new SqlSumAggFunction(sumType), sumType),
            oldCall.isDistinct(),
            oldCall.isApproximate(),
            oldCall.ignoreNulls(),
            ImmutableIntList.of(argOrdinal),
            oldCall.filterArg,
            oldCall.distinctKeys,
            oldCall.getCollation(),
            sumType,
            null);
    final RexNode sumArg =
          rexBuilder.addAggCall(
              sumArgAggCall,
              nGroups,
              newCalls,
              aggCallMapping,
              ImmutableList.of(argType));

    final RexNode sumSquaredArg =
          rexBuilder.makeCall(
              SqlStdOperatorTable.MULTIPLY, sumArg, sumArg);

    final SqlCountAggFunction countAgg = (SqlCountAggFunction) SqlStdOperatorTable.COUNT;
    final RelDataType countType = countAgg.getReturnType(typeFactory);
    final AggregateCall countArgAggCall = getAggCall(oldCall, countAgg, countType);
    final RexNode countArg =
        rexBuilder.addAggCall(
            countArgAggCall,
            nGroups,
            newCalls,
            aggCallMapping,
            ImmutableList.of(argType));

    final RexNode avgSumSquaredArg =
        rexBuilder.makeCall(
            SqlStdOperatorTable.DIVIDE,
            sumSquaredArg, countArg);

    final RexNode diff =
        rexBuilder.makeCall(
            SqlStdOperatorTable.MINUS,
            sumArgSquared, avgSumSquaredArg);

    final RexNode denominator;
    if (biased) {
      denominator = countArg;
    } else {
      final RexLiteral one =
          rexBuilder.makeExactLiteral(BigDecimal.ONE);
      final RexNode nul =
          rexBuilder.makeNullLiteral(countArg.getType());
      final RexNode countMinusOne =
          rexBuilder.makeCall(
              SqlStdOperatorTable.MINUS, countArg, one);
      final RexNode countEqOne =
          rexBuilder.makeCall(
              SqlStdOperatorTable.EQUALS, countArg, one);
      denominator =
          rexBuilder.makeCall(
              SqlStdOperatorTable.CASE,
              countEqOne, nul, countMinusOne);
    }

    final SqlOperator divide;
    if (isInferenceEnabled) {
      divide = new DrillSqlOperator(
          "divide",
          2,
          true,
          oldCall.getType(), false);
    } else {
      divide = SqlStdOperatorTable.DIVIDE;
    }

    final RexNode div =
        rexBuilder.makeCall(
            divide, diff, denominator);

    RexNode result = div;
    if (sqrt) {
      final RexNode half =
          rexBuilder.makeExactLiteral(new BigDecimal("0.5"));
      result =
          rexBuilder.makeCall(
              SqlStdOperatorTable.POWER, div, half);
    }

    if (isInferenceEnabled) {
      return result;
    } else {
     /*
      * Currently calcite's strategy to infer the return type of aggregate functions
      * is wrong because it uses the first known argument to determine output type. For
      * instance if we are performing stddev on an integer column then it interprets the
      * output type to be integer which is incorrect as it should be double. So based on
      * this if we add cast after rewriting the aggregate we add an additional cast which
      * would cause wrong results. So we simply add a cast to ANY.
      */
      return rexBuilder.makeCast(
          typeFactory.createSqlType(SqlTypeName.ANY), result);
    }
  }

  /**
   * Finds the ordinal of an element in a list, or adds it.
   *
   * @param list    List
   * @param element Element to lookup or add
   * @param <T>     Element type
   * @return Ordinal of element in list
   */
  private static <T> int lookupOrAdd(List<T> list, T element) {
    int ordinal = list.indexOf(element);
    if (ordinal == -1) {
      ordinal = list.size();
      list.add(element);
    }
    return ordinal;
  }

  /**
   * Do a shallow clone of oldAggRel and update aggCalls. Could be refactored
   * into Aggregate and subclasses - but it's only needed for some
   * subclasses.
   *
   * @param oldAggRel AggregateRel to clone.
   * @param inputRel  Input relational expression
   * @param newCalls  New list of AggregateCalls
   * @return shallow clone with new list of AggregateCalls.
   */
  protected Aggregate newAggregateRel(
      Aggregate oldAggRel,
      RelNode inputRel,
      List<AggregateCall> newCalls) {
    RelOptCluster cluster = inputRel.getCluster();
    return new LogicalAggregate(cluster, cluster.traitSetOf(Convention.NONE), Collections.emptyList(),
        inputRel, oldAggRel.getGroupSet(), oldAggRel.getGroupSets(), newCalls);
  }

  private RelDataType getFieldType(RelNode relNode, int i) {
    final RelDataTypeField inputField =
        relNode.getRowType().getFieldList().get(i);
    return inputField.getType();
  }

  private static class DrillConvertSumToSumZero extends RelOptRule {

    public DrillConvertSumToSumZero(RelOptRuleOperand operand) {
      super(operand, DrillRelFactories.LOGICAL_BUILDER, null);
    }

    @Override
    public boolean matches(RelOptRuleCall call) {
      DrillAggregateRel oldAggRel = (DrillAggregateRel) call.rels[0];
      for (AggregateCall aggregateCall : oldAggRel.getAggCallList()) {
        if (isConversionToSumZeroNeeded(aggregateCall.getAggregation(), aggregateCall.getType())) {
          return true;
        }
      }
      return false;
    }

    @Override
    public void onMatch(RelOptRuleCall call) {
      final DrillAggregateRel oldAggRel = (DrillAggregateRel) call.rels[0];

      final Map<AggregateCall, RexNode> aggCallMapping = Maps.newHashMap();
      final List<AggregateCall> newAggregateCalls = Lists.newArrayList();
      for (AggregateCall oldAggregateCall : oldAggRel.getAggCallList()) {
        if (isConversionToSumZeroNeeded(oldAggregateCall.getAggregation(), oldAggregateCall.getType())) {
          final RelDataType argType = oldAggregateCall.getType();
          final RelDataType sumType = oldAggRel.getCluster().getTypeFactory()
              .createTypeWithNullability(argType, argType.isNullable());
          final SqlAggFunction sumZeroAgg = new DrillCalciteSqlSumEmptyIsZeroAggFunctionWrapper(
              new SqlSumEmptyIsZeroAggFunction(), sumType);
          AggregateCall sumZeroCall =
              AggregateCall.create(
                  sumZeroAgg,
                  oldAggregateCall.isDistinct(),
                  oldAggregateCall.isApproximate(),
                  oldAggregateCall.ignoreNulls(),
                  oldAggregateCall.getArgList(),
                  oldAggregateCall.filterArg,
                  oldAggregateCall.distinctKeys,
                  oldAggregateCall.getCollation(),
                  sumType,
                  oldAggregateCall.getName());
          oldAggRel.getCluster().getRexBuilder()
              .addAggCall(sumZeroCall,
                  oldAggRel.getGroupCount(),
                  newAggregateCalls,
                  aggCallMapping,
                  ImmutableList.of(argType));
        } else {
          newAggregateCalls.add(oldAggregateCall);
        }
      }

      call.transformTo(new DrillAggregateRel(
          oldAggRel.getCluster(),
          oldAggRel.getTraitSet(),
          oldAggRel.getInput(),
          oldAggRel.getGroupSet(),
          oldAggRel.getGroupSets(),
          newAggregateCalls));
    }
  }

  private static class DrillConvertWindowSumToSumZero extends RelOptRule {
    public DrillConvertWindowSumToSumZero(RelOptRuleOperand operand) {
      super(operand, DrillRelFactories.LOGICAL_BUILDER, null);
    }

    @Override
    public boolean matches(RelOptRuleCall call) {
      final DrillWindowRel oldWinRel = (DrillWindowRel) call.rels[0];
      for (Window.Group group : oldWinRel.groups) {
        for (Window.RexWinAggCall rexWinAggCall : group.aggCalls) {
          if (isConversionToSumZeroNeeded(rexWinAggCall.getOperator(), rexWinAggCall.getType())) {
            return true;
          }
        }
      }
      return false;
    }

    @Override
    public void onMatch(RelOptRuleCall call) {
      final DrillWindowRel oldWinRel = (DrillWindowRel) call.rels[0];
      final ImmutableList.Builder<Window.Group> builder = ImmutableList.builder();

      for (Window.Group group : oldWinRel.groups) {
        final List<Window.RexWinAggCall> aggCalls = Lists.newArrayList();
        for (Window.RexWinAggCall rexWinAggCall : group.aggCalls) {
          if (isConversionToSumZeroNeeded(rexWinAggCall.getOperator(), rexWinAggCall.getType())) {
            final RelDataType argType = rexWinAggCall.getType();
            final RelDataType sumType = oldWinRel.getCluster().getTypeFactory()
                .createTypeWithNullability(argType, argType.isNullable());
            final SqlAggFunction sumZeroAgg = new DrillCalciteSqlSumEmptyIsZeroAggFunctionWrapper(
                new SqlSumEmptyIsZeroAggFunction(), sumType);
            final Window.RexWinAggCall sumZeroCall =
                new Window.RexWinAggCall(
                    sumZeroAgg,
                    sumType,
                    rexWinAggCall.operands,
                    rexWinAggCall.ordinal,
                    rexWinAggCall.distinct);
            aggCalls.add(sumZeroCall);
          } else {
            aggCalls.add(rexWinAggCall);
          }
        }

        final Window.Group newGroup = new Window.Group(
            group.keys,
            group.isRows,
            group.lowerBound,
            group.upperBound,
            group.orderKeys,
            aggCalls);
        builder.add(newGroup);
      }

      call.transformTo(new DrillWindowRel(
          oldWinRel.getCluster(),
          oldWinRel.getTraitSet(),
          oldWinRel.getInput(),
          oldWinRel.constants,
          oldWinRel.getRowType(),
          builder.build()));
    }
  }

  private static boolean isConversionToSumZeroNeeded(SqlOperator sqlOperator, RelDataType type) {
    sqlOperator = DrillCalciteWrapperUtility.extractSqlOperatorFromWrapper(sqlOperator);
    if (sqlOperator instanceof SqlSumAggFunction
        && !type.isNullable()) {
      // If SUM(x) is not nullable, the validator must have determined that
      // nulls are impossible (because the group is never empty and x is never
      // null). Therefore we translate to SUM0(x).
      return true;
    }
    return false;
  }
}