Scope aggregate columns in SparkConnectPlanner

G-Research · Jul 30, 2024 · c4fce40 · c4fce40
1 parent 2443f8a
commit c4fce40
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 1 deletion.
diff --git a/connect/server/src/main/scala/org/apache/spark/sql/connect/planner/SparkConnectPlanner.scala b/connect/server/src/main/scala/org/apache/spark/sql/connect/planner/SparkConnectPlanner.scala
@@ -2450,6 +2450,9 @@ class SparkConnectPlanner(
     val keyColumn = TypedAggUtils.aggKeyColumn(ds.kEncoder, ds.groupingAttributes)
     val namedColumns = rel.getAggregateExpressionsList.asScala.toSeq
       .map(expr => transformExpressionWithTypedReduceExpression(expr, input))
+      // SPARK-42199: resolve these aggregate expressions only against dataAttributes
+      // this is to hide key column from expression resolution
+      .map(ScopedExpression(_, ds.dataAttributes))
       .map(toNamedExpression)
     logical.Aggregate(ds.groupingAttributes, keyColumn +: namedColumns, ds.analyzed)
   }

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/KeyValueGroupedDataset.scala b/sql/core/src/main/scala/org/apache/spark/sql/KeyValueGroupedDataset.scala
@@ -971,7 +971,8 @@ class KeyValueGroupedDataset[K, V] private[sql](
     val encoders = columns.map(_.encoder)
     val namedColumns =
       columns
-        // SPARK-42199: resolve these sort expressions only against dataAttributes
+        // SPARK-42199: resolve these aggregate expressions only against dataAttributes
+        // this is to hide key column from expression resolution
         .map(scopeTypedColumn(dataAttributes))
         .map(_.withInputType(vExprEnc, dataAttributes).named)
     val keyColumn = TypedAggUtils.aggKeyColumn(kExprEnc, groupingAttributes)