/*
 * 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.spark.sql.catalyst.analysis.resolver

import org.apache.spark.sql.catalyst.analysis.GroupingAnalyticsTransformer
import org.apache.spark.sql.catalyst.expressions.{
  BaseGroupingSets,
  Expression,
  GroupingAnalyticsExtractor,
  SortOrder
}
import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, Expand, Sort}
import org.apache.spark.sql.errors.QueryCompilationErrors

/**
 * Resolves [[Aggregate]] node with grouping analytics (i.e., ROLLUP, CUBE, GROUPING SETS).
 */
class GroupingAnalyticsResolver(resolver: Resolver, expressionResolver: ExpressionResolver)
    extends TreeNodeResolver[Aggregate, Aggregate] {
  private val scopes: NameScopeStack = resolver.getNameScopes
  private val autoGeneratedAliasProvider = expressionResolver.getAutoGeneratedAliasProvider
  private val operatorResolutionContextStack = resolver.getOperatorResolutionContextStack

  /**
   * Resolution of the [[Aggregate]] is done in following steps:
   *  1. Extract grouping analytics from grouping expressions using [[GroupingAnalyticsExtractor]].
   *  2. Transform the [[Aggregate]] node with grouping analytics using
   *     [[GroupingAnalyticsTransformer]] (see its scala doc for more details).
   *  3. If the child of the transformed [[Aggregate]] is an [[Expand]] node, update its output
   *     attributes to have correct expression IDs.
   *
   * For example for this query:
   *
   * {{{
   *   SELECT col1 FROM values(1) GROUP BY grouping sets ((col1), ());
   * }}}
   *
   * The parsed plan would be:
   *
   * {{{
   *   'Aggregate [groupingsets(Vector(0), Vector(), 'col1)], ['col1]
   *   +- LocalRelation [col1#0]
   * }}}
   *
   * Whereas the resolved plan would be:
   *
   * {{{
   *   Aggregate [col1#3, spark_grouping_id#2L], [col1#3]
   *   +- Expand [[col1#0, col1#1, 0], [col1#0, null, 1]], [col1#0, col1#3, spark_grouping_id#2L]
   *      +- Project [col1#0, col1#0 AS col1#1]
   *         +- LocalRelation [col1#0]
   * }}}
   *
   * As it can be seen [[Expand]] node propagates two `col1` attributes which have different
   * properties. The first one (`col1#0`) comes from the child of the original [[Aggregate]] and it
   * is passed unchanged. The second one (`col1#1`) is generated by the
   * [[GroupingAnalyticsTransformer]] to represent `col1` in the grouping set and it has updated
   * nullability. It's nullability is set to `true` since in the grouping set `col1` can be absent
   * (i.e., null) as in the second grouping set `()`. The second one is used later to resolve `col1`
   * in the [[Aggregate]] node.
   */
  override def resolve(aggregate: Aggregate): Aggregate = {
    val groupingAnalytics = tryExtractGroupingAnalytics(aggregate)

    groupingAnalytics match {
      case Some((selectedGroupByExpressions, groupByExpressions)) =>
        operatorResolutionContextStack.current.hasGroupingAnalytics = true

        val transformedAggregate = GroupingAnalyticsTransformer(
          newAlias = (child, name, qualifier) =>
            autoGeneratedAliasProvider.newAlias(child, name, qualifier = qualifier),
          childOutput = scopes.current.output,
          groupByExpressions = groupByExpressions,
          selectedGroupByExpressions = selectedGroupByExpressions,
          child = aggregate.child,
          aggregationExpressions = aggregate.aggregateExpressions
        )

        val newAggregateChild = transformedAggregate.child match {
          case expand: Expand =>
            handleExpandBelowAggregate(expand)
          case other => other
        }

        val resolvedAggregate = transformedAggregate.copy(child = newAggregateChild)
        // copy() does not carry tree-node tags, so preserve the grand-total marker set by
        // GroupingAnalyticsTransformer (used to fold grouping_id() to 0 in HAVING/ORDER BY).
        resolvedAggregate.copyTagsFrom(transformedAggregate)
        resolvedAggregate
      case None =>
        aggregate
    }
  }

  /**
   * Handles [[SortOrder]] expressions that contain grouping analytics expressions. This is done by:
   *  1. Validating that there is a base [[Aggregate]] in the current scope:
   *      - If there is, collect grouping expressions from it using
   *        [[GroupingAnalyticsTransformer.collectGroupingExpressions]].
   *      - If there isn't, throw `groupingMustWithGroupingSetsOrCubeOrRollupError` exception.
   *  2. Compute the grouping id via [[GroupingAnalyticsTransformer.groupingIdExpression]] (the
   *     constant 0 for a grand total, otherwise the `spark_grouping_id` attribute) and resolve it
   *     using [[ExpressionResolver]].
   *  3. Replace [[SortOrder]] expressions using
   *     [[GroupingAnalyticsTransformer.replaceGroupingFunction]] (see its scala doc for more
   *     details).
   */
  def handleSortOrderExpressionsWithGroupingAnalytics(
      orderExpressions: Seq[SortOrder],
      sort: Sort): Seq[SortOrder] = {
    if (!scopes.current.baseAggregate.isDefined) {
      throw QueryCompilationErrors.groupingMustWithGroupingSetsOrCubeOrRollupError()
    }

    val groupingExpressions =
      GroupingAnalyticsTransformer.collectGroupingExpressions(
        scopes.current.baseAggregate.get
      )

    val groupingId = GroupingAnalyticsTransformer.groupingIdExpression(groupingExpressions)
    val resolvedGroupingId =
      expressionResolver.resolveExpressionTreeInOperator(groupingId, sort)

    val orderExpressionsWithGroupingAnalytics = orderExpressions.map { orderExpression =>
      GroupingAnalyticsTransformer
        .replaceGroupingFunction(
          expression = orderExpression,
          groupByExpressions = groupingExpressions,
          gid = resolvedGroupingId,
          newAlias = (child, name, qualifier) =>
            autoGeneratedAliasProvider.newAlias(child, name, qualifier = qualifier)
        )
        .asInstanceOf[SortOrder]
    }

    orderExpressionsWithGroupingAnalytics
  }

  /**
   * Extract expressions that contain grouping analytics operations from an
   * [[Aggregate.groupingExpressions]] using [[GroupingAnalyticsExtractor]]. See its scala doc for
   * more details.
   */
  private def tryExtractGroupingAnalytics(
      aggregate: Aggregate): Option[(Seq[Seq[Expression]], Seq[Expression])] = {
    aggregate.groupingExpressions match {
      case exprs if exprs.exists(_.isInstanceOf[BaseGroupingSets]) =>
        GroupingAnalyticsExtractor(exprs)
      case _ =>
        None
    }
  }

  /**
   * Handles [[Expand]] node below an [[Aggregate]]. This includes:
   *  - Updating output attributes to have correct expression IDs.
   *  - Overwriting the current scope output and extending hidden output with the new output
   *    since operators above the [[Aggregate]] might resolve attributes using hidden output
   *    produces by the [[Expand]].
   *  - Copying tags from the original [[Expand]].
   */
  private def handleExpandBelowAggregate(expand: Expand): Expand = {
    val mappedOutput = expand.output.map { attribute =>
      expressionResolver.getExpressionIdAssigner
        .mapExpression(attribute, allowUpdatesForAttributeReferences = true)
    }

    val newExpand = expand.copy(output = mappedOutput)

    scopes.overwriteOutputAndExtendHiddenOutput(output = newExpand.output)
    newExpand.copyTagsFrom(expand)

    newExpand
  }
}

/**
 * Helper object for [[GroupingAnalyticsResolver]].
 */
object GroupingAnalyticsResolver {

  /**
   * In case there are grouping analytics outside of an [[Aggregate]] node (e.g. [[Filter]]), throw
   * an [[ExplicitlyUnsupportedResolverFeature]] exception.
   */
  def restrictGroupingAnalyticsBelowSortAndFilter(
      operatorResolutionContext: OperatorResolutionContext): Unit = {
    if (operatorResolutionContext.hasGroupingAnalytics) {
      throw new ExplicitlyUnsupportedResolverFeature("grouping analytics outside of Aggregate")
    }
  }
}
