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.asterix.optimizer.rules;
  
  import java.util.ArrayList;
  import java.util.List;
  
  import org.apache.asterix.lang.sqlpp.util.SqlppVariableUtil;
  import org.apache.asterix.om.functions.BuiltinFunctions;
  import org.apache.commons.lang3.mutable.Mutable;
  import org.apache.commons.lang3.mutable.MutableObject;
  import org.apache.hyracks.algebricks.common.exceptions.AlgebricksException;
  import org.apache.hyracks.algebricks.core.algebra.base.ILogicalExpression;
  import org.apache.hyracks.algebricks.core.algebra.base.ILogicalOperator;
  import org.apache.hyracks.algebricks.core.algebra.base.IOptimizationContext;
  import org.apache.hyracks.algebricks.core.algebra.base.LogicalExpressionTag;
  import org.apache.hyracks.algebricks.core.algebra.base.LogicalOperatorTag;
  import org.apache.hyracks.algebricks.core.algebra.base.LogicalVariable;
  import org.apache.hyracks.algebricks.core.algebra.expressions.AbstractFunctionCallExpression;
  import org.apache.hyracks.algebricks.core.algebra.expressions.AggregateFunctionCallExpression;
  import org.apache.hyracks.algebricks.core.algebra.expressions.ConstantExpression;
  import org.apache.hyracks.algebricks.core.algebra.expressions.VariableReferenceExpression;
  import org.apache.hyracks.algebricks.core.algebra.functions.FunctionIdentifier;
  import org.apache.hyracks.algebricks.core.algebra.operators.logical.AbstractBinaryJoinOperator;
  import org.apache.hyracks.algebricks.core.algebra.operators.logical.AbstractLogicalOperator;
  import org.apache.hyracks.algebricks.core.algebra.operators.logical.AggregateOperator;
  import org.apache.hyracks.algebricks.core.algebra.operators.logical.AssignOperator;
  import org.apache.hyracks.algebricks.core.algebra.operators.logical.GroupByOperator;
  import org.apache.hyracks.algebricks.core.algebra.operators.logical.visitors.VariableUtilities;
  import org.apache.hyracks.algebricks.core.algebra.util.OperatorManipulationUtil;
  import org.apache.hyracks.algebricks.core.rewriter.base.IAlgebraicRewriteRule;
  
  /**
   * Pushes aggregate functions into a stand alone aggregate operator (no group by).
   */
  public class PushAggFuncIntoStandaloneAggregateRule implements IAlgebraicRewriteRule {
  
      @Override
      public boolean rewritePre(Mutable<ILogicalOperator> opRef, IOptimizationContext context)
              throws AlgebricksException {
          return false;
      }
  
      @Override
      public boolean rewritePost(Mutable<ILogicalOperator> opRef, IOptimizationContext context)
              throws AlgebricksException {
          // Pattern to match: assign <-- aggregate <-- !(group-by)
          AbstractLogicalOperator op = (AbstractLogicalOperator) opRef.getValue();
          if (op.getOperatorTag() != LogicalOperatorTag.ASSIGN) {
              return false;
          }
          AssignOperator assignOp = (AssignOperator) op;
  
          Mutable<ILogicalOperator> opRef2 = op.getInputs().get(0);
          AbstractLogicalOperator op2 = (AbstractLogicalOperator) opRef2.getValue();
          if (op2.getOperatorTag() == LogicalOperatorTag.AGGREGATE) {
              AggregateOperator aggOp = (AggregateOperator) op2;
              // Make sure the agg expr is a listify.
              return pushAggregateFunction(assignOp, aggOp, context);
          } else if (op2.getOperatorTag() == LogicalOperatorTag.INNERJOIN
                  || op2.getOperatorTag() == LogicalOperatorTag.LEFTOUTERJOIN) {
              AbstractBinaryJoinOperator join = (AbstractBinaryJoinOperator) op2;
              // Tries to push aggregates through the join.
              if (containsAggregate(assignOp.getExpressions()) && pushableThroughJoin(join)) {
                  return pushAggregateFunctionThroughJoin(assignOp, join, context);
              }
          }
          return false;
      }
  
      /**
       * Recursively check whether the list of expressions contains an aggregate function.
       *
       * @return true if the list contains an aggregate function and false otherwise.
       */
      private boolean containsAggregate(List<Mutable<ILogicalExpression>> exprRefs) {
          for (Mutable<ILogicalExpression> exprRef : exprRefs) {
              ILogicalExpression expr = exprRef.getValue();
              if (expr.getExpressionTag() != LogicalExpressionTag.FUNCTION_CALL) {
                  continue;
              }
              AbstractFunctionCallExpression funcExpr = (AbstractFunctionCallExpression) expr;
              FunctionIdentifier funcIdent = BuiltinFunctions.getAggregateFunction(funcExpr.getFunctionIdentifier());
              if (funcIdent == null) {
                  // Recursively look in func args.
                  if (containsAggregate(funcExpr.getArguments())) {
                      return true;
                  }
              } else {
                  // This is an aggregation function.
                  return true;
              }
          }
          return false;
      }
  
      /**
       * Check whether the join is aggregate-pushable, that is,
       * 1) the join condition is true;
       * 2) each join branch produces only one tuple.
       *
       * @return true if pushable
       */
      private boolean pushableThroughJoin(AbstractBinaryJoinOperator join) {
          ILogicalExpression condition = join.getCondition().getValue();
          if (condition.equals(ConstantExpression.TRUE)) {
              // Checks if the aggregation functions are pushable through the join
              boolean pushable = true;
              for (Mutable<ILogicalOperator> branchRef : join.getInputs()) {
                  AbstractLogicalOperator branch = (AbstractLogicalOperator) branchRef.getValue();
                  if (branch.getOperatorTag() == LogicalOperatorTag.AGGREGATE) {
                      pushable &= true;
                  } else if (branch.getOperatorTag() == LogicalOperatorTag.INNERJOIN
                          || branch.getOperatorTag() == LogicalOperatorTag.LEFTOUTERJOIN) {
                      AbstractBinaryJoinOperator childJoin = (AbstractBinaryJoinOperator) branch;
                      pushable &= pushableThroughJoin(childJoin);
                  } else {
                      pushable &= false;
                  }
              }
              return pushable;
          }
          return false;
      }
  
      /**
       * Does the actual push of aggregates for qualified joins.
       *
       * @param assignOp
       *            that contains aggregate function calls.
       * @param context
       * @throws AlgebricksException
       */
      private boolean pushAggregateFunctionThroughJoin(AssignOperator assignOp, AbstractBinaryJoinOperator join,
              IOptimizationContext context) throws AlgebricksException {
          boolean applied = false;
          for (Mutable<ILogicalOperator> branchRef : join.getInputs()) {
              AbstractLogicalOperator branch = (AbstractLogicalOperator) branchRef.getValue();
              if (branch.getOperatorTag() == LogicalOperatorTag.AGGREGATE) {
                  AggregateOperator aggOp = (AggregateOperator) branch;
                  applied |= pushAggregateFunction(assignOp, aggOp, context);
              } else if (branch.getOperatorTag() == LogicalOperatorTag.INNERJOIN
                      || branch.getOperatorTag() == LogicalOperatorTag.LEFTOUTERJOIN) {
                  AbstractBinaryJoinOperator childJoin = (AbstractBinaryJoinOperator) branch;
                  applied |= pushAggregateFunctionThroughJoin(assignOp, childJoin, context);
              }
          }
          return applied;
      }
  
      private boolean pushAggregateFunction(AssignOperator assignOp, AggregateOperator aggOp,
              IOptimizationContext context) throws AlgebricksException {
          Mutable<ILogicalOperator> aggChilldOpRef = aggOp.getInputs().get(0);
          AbstractLogicalOperator aggChildOp = (AbstractLogicalOperator) aggChilldOpRef.getValue();
          // If there's a group by below the agg, then we want to have the agg pushed into the group by
          if (aggChildOp.getOperatorTag() == LogicalOperatorTag.GROUP
                  && !((GroupByOperator) aggChildOp).getNestedPlans().isEmpty()) {
              return false;
          }
  
          List<LogicalVariable> assignUsedVars = new ArrayList<>();
          VariableUtilities.getUsedVariables(assignOp, assignUsedVars);
  
          List<Mutable<ILogicalExpression>> assignScalarAggExprRefs = new ArrayList<>();
          List<LogicalVariable> aggAddVars = null;
          List<Mutable<ILogicalExpression>> aggAddExprs = null;
  
          for (int i = 0, n = aggOp.getVariables().size(); i < n; i++) {
              LogicalVariable aggVar = aggOp.getVariables().get(i);
              Mutable<ILogicalExpression> aggExprRef = aggOp.getExpressions().get(i);
              ILogicalExpression aggExpr = aggExprRef.getValue();
              if (aggExpr.getExpressionTag() != LogicalExpressionTag.FUNCTION_CALL) {
                  continue;
              }
              AbstractFunctionCallExpression listifyCandidateExpr = (AbstractFunctionCallExpression) aggExpr;
              if (listifyCandidateExpr.getFunctionIdentifier() != BuiltinFunctions.LISTIFY) {
                  continue;
              }
              if (!assignUsedVars.contains(aggVar)) {
                  continue;
              }
              assignScalarAggExprRefs.clear();
              findScalarAggFuncExprRef(assignOp.getExpressions(), aggVar, assignScalarAggExprRefs);
              if (assignScalarAggExprRefs.isEmpty()) {
                  continue;
              }
              // perform rewrite
              if (aggAddVars == null) {
                  aggAddVars = new ArrayList<>();
                  aggAddExprs = new ArrayList<>();
              }
              for (Mutable<ILogicalExpression> assignScalarAggExprRef : assignScalarAggExprRefs) {
                  AbstractFunctionCallExpression assignScalarAggExpr =
                          (AbstractFunctionCallExpression) assignScalarAggExprRef.getValue();
                  FunctionIdentifier aggFuncIdent =
                          BuiltinFunctions.getAggregateFunction(assignScalarAggExpr.getFunctionIdentifier());
  
                  // Push the scalar aggregate function into the aggregate op.
                  int nArgs = assignScalarAggExpr.getArguments().size();
                  List<Mutable<ILogicalExpression>> aggArgs = new ArrayList<>(nArgs);
                  aggArgs.add(
                          new MutableObject<>(listifyCandidateExpr.getArguments().get(0).getValue().cloneExpression()));
                  aggArgs.addAll(OperatorManipulationUtil
                          .cloneExpressions(assignScalarAggExpr.getArguments().subList(1, nArgs)));
                  AggregateFunctionCallExpression aggFuncExpr =
                          BuiltinFunctions.makeAggregateFunctionExpression(aggFuncIdent, aggArgs);
                  aggFuncExpr.setSourceLocation(assignScalarAggExpr.getSourceLocation());
  
                  LogicalVariable newVar = context.newVar();
                  aggAddVars.add(newVar);
                  aggAddExprs.add(new MutableObject<>(aggFuncExpr));
                  // The assign now just "renames" the variable to make sure the upstream plan still works.
                  VariableReferenceExpression newVarRef = new VariableReferenceExpression(newVar);
                  newVarRef.setSourceLocation(assignScalarAggExpr.getSourceLocation());
                  assignScalarAggExprRef.setValue(newVarRef);
              }
          }
  
          if (aggAddVars == null) {
              return false;
          }
  
          // add new variables and expressions to the aggregate operator.
          aggOp.getVariables().addAll(aggAddVars);
          aggOp.getExpressions().addAll(aggAddExprs);
  
          // Note: we retain the original listify() call in the aggregate operator because
          // the variable it is assigned to might be used upstream by other operators.
          // If the variable is not used upstream then it'll later be removed
          // by {@code RemoveUnusedAssignAndAggregateRule}
  
          context.computeAndSetTypeEnvironmentForOperator(aggOp);
          context.computeAndSetTypeEnvironmentForOperator(assignOp);
          return true;
      }
  
      private void findScalarAggFuncExprRef(List<Mutable<ILogicalExpression>> exprRefs, LogicalVariable aggVar,
              List<Mutable<ILogicalExpression>> outScalarAggExprRefs) {
          for (Mutable<ILogicalExpression> exprRef : exprRefs) {
              ILogicalExpression expr = exprRef.getValue();
              if (expr.getExpressionTag() == LogicalExpressionTag.FUNCTION_CALL) {
                  AbstractFunctionCallExpression funcExpr = (AbstractFunctionCallExpression) expr;
                  FunctionIdentifier funcIdent = BuiltinFunctions.getAggregateFunction(funcExpr.getFunctionIdentifier());
                  if (funcIdent != null
                          && aggVar.equals(SqlppVariableUtil.getVariable(funcExpr.getArguments().get(0).getValue()))) {
                      outScalarAggExprRefs.add(exprRef);
                  } else {
                      // Recursively look in func args.
                      findScalarAggFuncExprRef(funcExpr.getArguments(), aggVar, outScalarAggExprRefs);
                  }
              }
          }
      }
  }
  