[SPARK-31670][SQL] Trim unnecessary Struct field alias in Aggregate/GroupingSets

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

@@ -510,10 +510,12 @@ class Analyzer(
       // collect all the found AggregateExpression, so we can check an expression is part of
       // any AggregateExpression or not.
       val aggsBuffer = ArrayBuffer[Expression]()
+
       // Returns whether the expression belongs to any expressions in `aggsBuffer` or not.
       def isPartOfAggregation(e: Expression): Boolean = {
         aggsBuffer.exists(a => a.find(_ eq e).isDefined)
       }
+
       replaceGroupingFunc(_, groupByExprs, gid).transformDown {
         // AggregateExpression should be computed on the unmodified value of its argument
         // expressions, so we should not replace any references to grouping expression
@@ -1259,6 +1261,11 @@ class Analyzer(
       attr.withExprId(exprId)
     }
 
+    private def dedupStructField(attr: Alias, structFieldMap: Map[String, Attribute]) = {
+      val exprId = structFieldMap.getOrElse(attr.child.sql, attr).exprId
+      Alias(attr.child, attr.name)(exprId, attr.qualifier, attr.explicitMetadata)
+    }
+
     /**
      * The outer plan may have been de-duplicated and the function below updates the
      * outer references to refer to the de-duplicated attributes.
@@ -1479,11 +1486,70 @@ class Analyzer(
       // Skip the having clause here, this will be handled in ResolveAggregateFunctions.
       case h: UnresolvedHaving => h
 
+      case p: LogicalPlan if needResolveStructField(p) =>
+        logTrace(s"Attempting to resolve ${p.simpleString(SQLConf.get.maxToStringFields)}")
+        val resolved = p.mapExpressions(resolveExpressionTopDown(_, p))
+        val structFieldMap = new mutable.HashMap[String, Alias]
+        resolved.transformExpressions {
+          case a @ Alias(struct: GetStructField, _) =>
+            if (structFieldMap.contains(struct.sql)) {
+              val exprId = structFieldMap.getOrElse(struct.sql, a).exprId
+              Alias(a.child, a.name)(exprId, a.qualifier, a.explicitMetadata)
+            } else {
+              structFieldMap.put(struct.sql, a)
+              a
+            }
+          case e => e
+        }
+
       case q: LogicalPlan =>
         logTrace(s"Attempting to resolve ${q.simpleString(SQLConf.get.maxToStringFields)}")
         q.mapExpressions(resolveExpressionTopDown(_, q))
     }
 
+    def needResolveStructField(plan: LogicalPlan): Boolean = {
+      plan match {
+        case UnresolvedHaving(havingCondition, a: Aggregate)
+          if containSameStructFields(a.groupingExpressions.flatMap(_.references),
+            a.aggregateExpressions.flatMap(_.references),
+            Some(havingCondition.references.toSeq)) => true
+        case Aggregate(groupingExpressions, aggregateExpressions, _)
+          if containSameStructFields(groupingExpressions.flatMap(_.references),
+            aggregateExpressions.flatMap(_.references)) => true
+        case GroupingSets(selectedGroupByExprs, groupByExprs, _, aggregations)
+          if containSameStructFields(groupByExprs.flatMap(_.references),
+            aggregations.flatMap(_.references),
+            Some(selectedGroupByExprs.flatMap(_.flatMap(_.references)))) => true
+        case _ => false
+      }
+    }
+
+    def containSameStructFields(
+        grpExprs: Seq[Attribute],
+        aggExprs: Seq[Attribute],
+        extra: Option[Seq[Attribute]] = None): Boolean = {
+
+      def isStructField(attr: Attribute): Boolean = {
+        attr.isInstanceOf[UnresolvedAttribute] &&
+          attr.asInstanceOf[UnresolvedAttribute].nameParts.size == 2
+      }
+
+      val grpAttrs = grpExprs.filter(isStructField)
+        .map(_.asInstanceOf[UnresolvedAttribute].name)
+      val aggAttrs = aggExprs.filter(isStructField)
+        .map(_.asInstanceOf[UnresolvedAttribute].name)
+      val havingAttrs = extra.getOrElse(Seq.empty[Attribute]).filter(isStructField)
+        .map(_.asInstanceOf[UnresolvedAttribute].name)
+
+      if (extra.isDefined) {
+        grpAttrs.exists(aggAttrs.contains)
+      } else {
+        grpAttrs.exists(aggAttrs.contains) ||
+          grpAttrs.exists(havingAttrs.contains) ||
+          aggAttrs.exists(havingAttrs.contains)
+      }
+    }
+
     def resolveAssignments(
         assignments: Seq[Assignment],
         mergeInto: MergeIntoTable,

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

@@ -3496,6 +3496,88 @@ class SQLQuerySuite extends QueryTest with SharedSparkSession with AdaptiveSpark
     checkIfSeedExistsInExplain(df2)
   }
 
+  test("SPARK-31670: Struct Field in groupByExpr with CUBE") {
+    withTable("t") {
+      sql(
+        """CREATE TABLE t(
+          |a STRING,
+          |b INT,
+          |c ARRAY<STRUCT<row_id:INT,json_string:STRING>>,
+          |d ARRAY<ARRAY<STRING>>,
+          |e ARRAY<MAP<STRING, INT>>)
+          |USING ORC""".stripMargin)
+
+      checkAnswer(
+        sql(
+          """
+            |SELECT a, each.json_string, SUM(b)
+            |FROM t
+            |LATERAL VIEW EXPLODE(c) x AS each
+            |GROUP BY a, each.json_string
+            |WITH CUBE
+            |""".stripMargin), Nil)
+
+      checkAnswer(
+        sql(
+          """
+            |SELECT a, get_json_object(each.json_string, '$.i'), SUM(b)
+            |FROM t
+            |LATERAL VIEW EXPLODE(c) X AS each
+            |GROUP BY a, get_json_object(each.json_string, '$.i')
+            |WITH CUBE
+            |""".stripMargin), Nil)
+
+      checkAnswer(
+        sql(
+          """
+            |SELECT a, each.json_string AS json_string, SUM(b)
+            |FROM t
+            |LATERAL VIEW EXPLODE(c) x AS each
+            |GROUP BY a, each.json_string
+            |WITH CUBE
+            |""".stripMargin), Nil)
+
+      checkAnswer(
+        sql(
+          """
+            |SELECT a, each.json_string as js, SUM(b)
+            |FROM t
+            |LATERAL VIEW EXPLODE(c) X AS each
+            |GROUP BY a, each.json_string
+            |WITH CUBE
+            |""".stripMargin), Nil)
+
+      checkAnswer(
+        sql(
+          """
+            |SELECT a, each.json_string as js, SUM(b)
+            |FROM t
+            |LATERAL VIEW EXPLODE(c) X AS each
+            |GROUP BY a, each.json_string
+            |WITH ROLLUP
+            |""".stripMargin), Nil)
+
+      sql(
+        """
+          |SELECT a, each.json_string, SUM(b)
+          |FROM t
+          |LATERAL VIEW EXPLODE(c) X AS each
+          |GROUP BY a, each.json_string
+          |GROUPING sets((a),(a, each.json_string))
+          |""".stripMargin).explain(true)
+
+      checkAnswer(
+        sql(
+          """
+            |SELECT a, each.json_string, SUM(b)
+            |FROM t
+            |LATERAL VIEW EXPLODE(c) X AS each
+            |GROUP BY a, each.json_string
+            |GROUPING sets((a),(a, each.json_string))
+            |""".stripMargin), Nil)
+    }
+  }
+
   test("SPARK-31761: test byte, short, integer overflow for (Divide) integral type") {
     checkAnswer(sql("Select -2147483648 DIV -1"), Seq(Row(Integer.MIN_VALUE.toLong * -1)))
     checkAnswer(sql("select CAST(-128 as Byte) DIV CAST (-1 as Byte)"),