[SPARK-28782][SQL] Generator support in aggregate expressions

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/Analyzer.scala

@@ -2018,6 +2018,62 @@ class Analyzer(
         throw new AnalysisException("Only one generator allowed per select clause but found " +
           generators.size + ": " + generators.map(toPrettySQL).mkString(", "))
 
+      case Aggregate(_, aggList, _) if aggList.exists(hasNestedGenerator) =>
+        val nestedGenerator = aggList.find(hasNestedGenerator).get
+        throw new AnalysisException("Generators are not supported when it's nested in " +
+          "expressions, but got: " + toPrettySQL(trimAlias(nestedGenerator)))
+
+      case Aggregate(_, aggList, _) if aggList.count(hasGenerator) > 1 =>
+        val generators = aggList.filter(hasGenerator).map(trimAlias)
+        throw new AnalysisException("Only one generator allowed per aggregate clause but found " +
+          generators.size + ": " + generators.map(toPrettySQL).mkString(", "))
+
+      case agg @ Aggregate(groupList, aggList, child) if aggList.forall {
+        case AliasedGenerator(generator, _, _) => generator.childrenResolved
+        case other => other.resolved
+      } =>
+        // Holds the resolved generator, if one exists in the project list.
+        var generatorVisited = false
+
+        val projectExprs = Array.ofDim[NamedExpression](aggList.length)
+        val newAggList = aggList
+          .map(CleanupAliases.trimNonTopLevelAliases(_).asInstanceOf[NamedExpression])
+          .zipWithIndex
+          .flatMap {
+            case (AliasedGenerator(generator, names, outer), idx) =>
+              // It's a sanity check, this should not happen as the previous case will throw
+              // exception earlier.
+              assert(!generatorVisited, "More than one generator found in aggregate.")
+              generatorVisited = true
+
+              val newGenChildren: Seq[Expression] = generator.children.zipWithIndex.map {
+                case (e, idx) => if (e.foldable) e else Alias(e, s"_gen_input_${idx}")()
+              }
+              val newGenerator = {
+                val g = generator.withNewChildren(newGenChildren.map { e =>
+                  if (e.foldable) e else e.asInstanceOf[Alias].toAttribute
+                }).asInstanceOf[Generator]
+                if (outer) GeneratorOuter(g) else g
+              }
+              val newAliasedGenerator = if (names.length == 1) {
+                Alias(newGenerator, names(0))()
+              } else {
+                MultiAlias(newGenerator, names)
+              }
+              projectExprs(idx) = newAliasedGenerator
+              newGenChildren.filter(!_.foldable).asInstanceOf[Seq[NamedExpression]]
+            case (other, idx) =>
+              projectExprs(idx) = other.toAttribute
+              other :: Nil
+          }
+
+        if (generatorVisited) {
+          val newAgg = Aggregate(groupList, newAggList, child)
+          Project(projectExprs.toList, newAgg)
+        } else {
+          agg
+        }
+
       case p @ Project(projectList, child) =>
         // Holds the resolved generator, if one exists in the project list.
         var resolvedGenerator: Generate = null

sql/core/src/test/scala/org/apache/spark/sql/GeneratorFunctionSuite.scala

@@ -308,6 +308,23 @@ class GeneratorFunctionSuite extends QueryTest with SharedSQLContext {
       sql("select * from values 1, 2 lateral view outer empty_gen() a as b"),
       Row(1, null) :: Row(2, null) :: Nil)
   }
+
+  test("generator in aggregate expression") {
+    Seq((1, 1), (1, 2), (2, 3)).toDF("c1", "c2").createOrReplaceTempView("t1")
+    checkAnswer(
+      sql("select explode(array(min(c2), max(c2))) from t1"),
+      Row(1) :: Row(3) :: Nil
+    )
+    checkAnswer(
+      sql("select posexplode(array(min(c2), max(c2))) from t1 group by c1"),
+      Row(0, 1) :: Row(1, 2) :: Row(0, 3) :: Row(1, 3) :: Nil
+    )
+    // test generator "stack" which require foldable argument
+    checkAnswer(
+      sql("select stack(2, min(c1), max(c1), min(c2), max(c2)) from t1"),
+      Row(1, 2) :: Row(1, 3) :: Nil
+    )
+  }
 }
 
 case class EmptyGenerator() extends Generator {