Scope aggregate columns in SparkConnectPlanner

G-Research · Jan 18, 2024 · e396fb0 · e396fb0
1 parent 682c7a3
commit e396fb0
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 1 deletion.
diff --git a/...nect/server/src/main/scala/org/apache/spark/sql/connect/planner/SparkConnectPlanner.scala b/...nect/server/src/main/scala/org/apache/spark/sql/connect/planner/SparkConnectPlanner.scala
@@ -2422,6 +2422,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
@@ -675,7 +675,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)