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.lang.sqlpp.rewrites.visitor;
  
  import java.util.ArrayList;
  import java.util.Collection;
  import java.util.HashMap;
  import java.util.HashSet;
  import java.util.List;
  import java.util.Map;
  import java.util.Set;
  
  import org.apache.asterix.common.exceptions.CompilationException;
  import org.apache.asterix.common.exceptions.ErrorCode;
  import org.apache.asterix.lang.common.base.AbstractClause;
  import org.apache.asterix.lang.common.base.Expression;
  import org.apache.asterix.lang.common.base.ILangExpression;
  import org.apache.asterix.lang.common.clause.GroupbyClause;
  import org.apache.asterix.lang.common.clause.LetClause;
  import org.apache.asterix.lang.common.clause.LimitClause;
  import org.apache.asterix.lang.common.clause.OrderbyClause;
  import org.apache.asterix.lang.common.context.Scope;
  import org.apache.asterix.lang.common.expression.GbyVariableExpressionPair;
  import org.apache.asterix.lang.common.expression.VariableExpr;
  import org.apache.asterix.lang.common.rewrites.LangRewritingContext;
  import org.apache.asterix.lang.common.struct.Identifier;
  import org.apache.asterix.lang.common.struct.VarIdentifier;
  import org.apache.asterix.lang.sqlpp.clause.FromClause;
  import org.apache.asterix.lang.sqlpp.clause.SelectBlock;
  import org.apache.asterix.lang.sqlpp.clause.SelectClause;
  import org.apache.asterix.lang.sqlpp.expression.SelectExpression;
  import org.apache.asterix.lang.sqlpp.util.SqlppRewriteUtil;
  import org.apache.asterix.lang.sqlpp.util.SqlppVariableUtil;
  import org.apache.asterix.lang.sqlpp.visitor.base.AbstractSqlppExpressionScopingVisitor;
  import org.apache.hyracks.algebricks.common.utils.Pair;
  
  /**
   * An AST pre-processor to rewrite group-by sugar queries, which does the following transformations:
   * 1. Rewrite the argument expression of an aggregation function into a subquery
   * 2. Turn a SQL-92 aggregate function into a SQL++ core aggregate function when performing 1.
   * <p>
   * <p>
   * For example, this visitor turns the following query:
   * <pre>
   * FROM Employee e
   * JOIN Incentive i ON e.job_category = i.job_category
   * JOIN SuperStars s ON e.id = s.id
   * GROUP BY e.department_id AS deptId
   * GROUP AS eis(e AS e, i AS i, s AS s)
   * SELECT deptId as deptId, SUM(e.salary + i.bonus) AS star_cost;
   * </pre>
   * into the following core-version query:
   * <pre>
   * FROM Employee e
   * JOIN Incentive i ON e.job_category = i.job_category
   * JOIN SuperStars s ON e.id = s.id
   * GROUP BY e.department_id AS deptId
   * GROUP AS eis(e AS e, i AS i, s AS s)
   * SELECT ELEMENT {
   *          'deptId': deptId,
   *          'star_cost': array_sum( (FROM eis AS p SELECT ELEMENT p.e.salary + p.i.bonus) )
   * };
   * </pre>
   * where <code>SUM(e.salary + i.bonus)</code>
   * is turned into <code>array_sum( (FROM eis AS p SELECT ELEMENT p.e.salary + p.i.bonus) )</code>
   */
  public class SqlppGroupByAggregationSugarVisitor extends AbstractSqlppExpressionScopingVisitor {
  
      private final Collection<VarIdentifier> externalVars;
  
      public SqlppGroupByAggregationSugarVisitor(LangRewritingContext context, Collection<VarIdentifier> externalVars) {
          super(context);
          this.externalVars = externalVars;
      }
  
      @Override
      public Expression visit(SelectBlock selectBlock, ILangExpression arg) throws CompilationException {
  
          Set<VariableExpr> outerVars = scopeChecker.getCurrentScope().getLiveVariables().keySet();
  
          // Traverses the select block in the order of "from", "let/where"s, "group by", "let/having"s and "select".
          FromClause fromClause = selectBlock.getFromClause();
          if (selectBlock.hasFromClause()) {
              fromClause.accept(this, arg);
          }
          if (selectBlock.hasLetWhereClauses()) {
              for (AbstractClause letWhereClause : selectBlock.getLetWhereList()) {
                  letWhereClause.accept(this, arg);
              }
          }
          if (selectBlock.hasGroupbyClause()) {
              Map<VariableExpr, Set<? extends Scope.SymbolAnnotation>> preGroupAnnotatedVars =
                      scopeChecker.getCurrentScope().getLiveVariables();
              Set<VariableExpr> preGroupVars = preGroupAnnotatedVars.keySet();
  
              GroupbyClause groupbyClause = selectBlock.getGroupbyClause();
              groupbyClause.accept(this, arg);
              Collection<VariableExpr> groupByBindingVars = SqlppVariableUtil.getBindingVariables(groupbyClause);
  
              VariableExpr groupVar = groupbyClause.getGroupVar();
              if (!groupbyClause.hasGroupFieldList()) {
                  throw new CompilationException(ErrorCode.COMPILATION_ILLEGAL_STATE, groupbyClause.getSourceLocation());
              }
              Map<VariableExpr, Identifier> groupVarFieldMap = createGroupVarFieldMap(groupbyClause.getGroupFieldList());
              Set<VariableExpr> preGroupMappedVars = groupVarFieldMap.keySet();
              Set<VariableExpr> preGroupContextVars = Scope.findVariablesAnnotatedBy(preGroupMappedVars,
                      SqlppVariableAnnotation.CONTEXT_VARIABLE, preGroupAnnotatedVars, groupbyClause.getSourceLocation());
              Set<VariableExpr> preGroupUnmappedVars = getUnmappedVariables(preGroupVars, preGroupMappedVars, outerVars);
  
              Collection<VariableExpr> freeVariables = new HashSet<>();
              Collection<VariableExpr> freeVariablesInGbyLets = new HashSet<>();
              if (selectBlock.hasLetHavingClausesAfterGroupby()) {
                  for (AbstractClause letHavingClause : selectBlock.getLetHavingListAfterGroupby()) {
                      letHavingClause.accept(this, arg);
                      // Rewrites each let/having clause after the group-by.
                      rewriteExpressionUsingGroupVariable(letHavingClause, groupVar, groupVarFieldMap,
                              preGroupContextVars, preGroupUnmappedVars, outerVars);
                      switch (letHavingClause.getClauseType()) {
                          case LET_CLAUSE:
                              LetClause letClause = (LetClause) letHavingClause;
                              Collection<VariableExpr> freeVariablesInClause =
                                      SqlppVariableUtil.getFreeVariables(letClause.getBindingExpr());
                              freeVariablesInClause.removeAll(groupByBindingVars);
                              freeVariablesInGbyLets.addAll(freeVariablesInClause);
                              groupByBindingVars.add(letClause.getVarExpr());
                              break;
                          case HAVING_CLAUSE:
                              freeVariables.addAll(SqlppVariableUtil.getFreeVariables(letHavingClause));
                              break;
                          default:
                              throw new IllegalStateException(String.valueOf(letHavingClause.getClauseType()));
                      }
                  }
              }
  
              SelectExpression parentSelectExpression = (SelectExpression) arg;
              // We cannot rewrite ORDER BY and LIMIT if it's a SET operation query.
              if (!parentSelectExpression.getSelectSetOperation().hasRightInputs()) {
                  if (parentSelectExpression.hasOrderby()) {
                      // Rewrites the ORDER BY clause.
                      OrderbyClause orderbyClause = parentSelectExpression.getOrderbyClause();
                      orderbyClause.accept(this, arg);
                      rewriteExpressionUsingGroupVariable(orderbyClause, groupVar, groupVarFieldMap, preGroupContextVars,
                              preGroupUnmappedVars, outerVars);
                      freeVariables.addAll(SqlppVariableUtil.getFreeVariables(orderbyClause));
                  }
                  if (parentSelectExpression.hasLimit()) {
                      // Rewrites the LIMIT clause.
                      LimitClause limitClause = parentSelectExpression.getLimitClause();
                      limitClause.accept(this, arg);
                      rewriteExpressionUsingGroupVariable(limitClause, groupVar, groupVarFieldMap, preGroupContextVars,
                              preGroupUnmappedVars, outerVars);
                      freeVariables.addAll(SqlppVariableUtil.getFreeVariables(limitClause));
                  }
              }
  
              // Visits the select clause.
              SelectClause selectClause = selectBlock.getSelectClause();
              selectClause.accept(this, arg);
              // Rewrites the select clause.
              rewriteExpressionUsingGroupVariable(selectClause, groupVar, groupVarFieldMap, preGroupContextVars,
                      preGroupUnmappedVars, outerVars);
              freeVariables.addAll(SqlppVariableUtil.getFreeVariables(selectClause));
              freeVariables.removeAll(groupByBindingVars);
  
              // Gets the final free variables.
              freeVariables.addAll(freeVariablesInGbyLets);
              removeExternalVariables(freeVariables);
  
              if (!groupbyClause.isGroupAll()) {
                  // Gets outer scope variables.
                  Collection<VariableExpr> decorVars = scopeChecker.getCurrentScope().getLiveVariables().keySet();
                  decorVars.removeAll(groupByBindingVars);
                  // Only retains used free variables.
                  if (!decorVars.containsAll(freeVariables)) {
                      throw new CompilationException(ErrorCode.COMPILATION_ILLEGAL_STATE,
                              groupbyClause.getSourceLocation(), decorVars + ":" + freeVariables);
                  }
                  decorVars.retainAll(freeVariables);
  
                  if (!decorVars.isEmpty()) {
                      // Adds necessary decoration variables for the GROUP BY.
                      // NOTE: we need to include outer binding variables so as they can be evaluated before
                      // the GROUP BY instead of being inlined as part of nested pipepline. The current optimzier
                      // is not able to optimize the latter case. The following query is such an example:
                      // asterixdb/asterix-app/src/test/resources/runtimets/queries_sqlpp/dapd/q2-11
                      List<GbyVariableExpressionPair> decorList = new ArrayList<>();
                      if (groupbyClause.hasDecorList()) {
                          decorList.addAll(groupbyClause.getDecorPairList());
                      }
                      for (VariableExpr var : decorVars) {
                          decorList.add(new GbyVariableExpressionPair((VariableExpr) SqlppRewriteUtil.deepCopy(var),
                                  (Expression) SqlppRewriteUtil.deepCopy(var)));
                      }
                      groupbyClause.setDecorPairList(decorList);
                  }
              }
          } else {
              selectBlock.getSelectClause().accept(this, arg);
          }
          return null;
      }
  
      private void removeExternalVariables(Collection<VariableExpr> freeVariables) {
          if (!externalVars.isEmpty()) {
              freeVariables.removeIf(ve -> externalVars.contains(ve.getVar()));
          }
      }
  
      static Map<VariableExpr, Identifier> createGroupVarFieldMap(List<Pair<Expression, Identifier>> fieldList) {
          Map<VariableExpr, Identifier> fieldVars = new HashMap<>();
          for (Pair<Expression, Identifier> p : fieldList) {
              if (p.first.getKind() == Expression.Kind.VARIABLE_EXPRESSION) {
                  fieldVars.put((VariableExpr) p.first, p.second);
              }
          }
          return fieldVars;
      }
  
      /**
       * Returns variables of the current SELECT block that were defined before GROUP BY clause but were not mapped by
       * GROUP AS sub-clause. These variables cannot be used by SQL aggregate functions after the GROUP BY
       */
      private Set<VariableExpr> getUnmappedVariables(Set<VariableExpr> preGroupByVars,
              Set<VariableExpr> preGroupByMappedVars, Set<VariableExpr> outerVars) {
          Set<VariableExpr> result = new HashSet<>(preGroupByVars);
          result.removeAll(preGroupByMappedVars);
          result.removeAll(outerVars);
          return result;
      }
  
      // Applying sugar rewriting for group-by.
      private void rewriteExpressionUsingGroupVariable(ILangExpression expr, VariableExpr groupVar,
              Map<VariableExpr, Identifier> groupVarFieldMap, Set<VariableExpr> preGroupContextVars,
              Set<VariableExpr> preGroupUnmappedVars, Set<VariableExpr> outerVars) throws CompilationException {
          Sql92AggregateFunctionVisitor visitor = new Sql92AggregateFunctionVisitor(context, groupVar, groupVarFieldMap,
                  preGroupContextVars, preGroupUnmappedVars, outerVars);
          expr.accept(visitor, null);
      }
  }
  