[SPARK-31915][SQL][PYTHON] Remove projection that adds grouping keys in grouped and cogrouped pandas UDFs

python/pyspark/sql/tests/test_pandas_cogrouped_map.py

@@ -19,7 +19,7 @@
 import sys
 
 from pyspark.sql.functions import array, explode, col, lit, udf, sum, pandas_udf, PandasUDFType
-from pyspark.sql.types import DoubleType, StructType, StructField
+from pyspark.sql.types import DoubleType, StructType, StructField, Row
 from pyspark.testing.sqlutils import ReusedSQLTestCase, have_pandas, have_pyarrow, \
     pandas_requirement_message, pyarrow_requirement_message
 from pyspark.testing.utils import QuietTest
@@ -193,6 +193,22 @@ def test_wrong_args(self):
             left.groupby('id').cogroup(right.groupby('id')) \
                 .applyInPandas(lambda: 1, StructType([StructField("d", DoubleType())]))
 
+    def test_case_insensitive_grouping_column(self):
+        # SPARK-31915: case-insensitive grouping column should work.
+        df1 = self.spark.createDataFrame([(1, 1)], ("column", "value"))
+
+        row = df1.groupby("ColUmn").cogroup(
+            df1.groupby("COLUMN")
+        ).applyInPandas(lambda r, l: r + l, "column long, value long").first()
+        self.assertEquals(row.asDict(), Row(column=2, value=2).asDict())
+
+        df2 = self.spark.createDataFrame([(1, 1)], ("column", "value"))
+
+        row = df1.groupby("ColUmn").cogroup(
+            df2.groupby("COLUMN")
+        ).applyInPandas(lambda r, l: r + l, "column long, value long").first()
+        self.assertEquals(row.asDict(), Row(column=2, value=2).asDict())
+
     @staticmethod
     def _test_with_key(left, right, isLeft):
 

python/pyspark/sql/tests/test_pandas_grouped_map.py

@@ -587,6 +587,16 @@ def f(key, pdf):
         # Check that all group and window_range values from udf matched expected
         self.assertTrue(all([r[0] for r in result]))
 
+    def test_case_insensitive_grouping_column(self):
+        # SPARK-31915: case-insensitive grouping column should work.
+        @pandas_udf("column integer, Score float", PandasUDFType.GROUPED_MAP)
+        def my_pandas_udf(pdf):
+            return pdf.assign(Score=0.5)
+
+        df = self.spark.createDataFrame([[1, 1]], ["column", "Score"])
+        row = df.groupby('COLUMN').apply(my_pandas_udf).first()
+        self.assertEquals(row.asDict(), Row(column=1, Score=0.5).asDict())
+
 
 if __name__ == "__main__":
     from pyspark.sql.tests.test_pandas_grouped_map import *

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

@@ -1374,12 +1374,10 @@ class Analyzer(
       // To resolve duplicate expression IDs for Join and Intersect
       case j @ Join(left, right, _, _, _) if !j.duplicateResolved =>
         j.copy(right = dedupRight(left, right))
-      case f @ FlatMapCoGroupsInPandas(leftAttributes, rightAttributes, _, _, left, right) =>
-        val leftRes = leftAttributes
-          .map(x => resolveExpressionBottomUp(x, left).asInstanceOf[Attribute])
-        val rightRes = rightAttributes
-          .map(x => resolveExpressionBottomUp(x, right).asInstanceOf[Attribute])
-        f.copy(leftAttributes = leftRes, rightAttributes = rightRes)
+      case f @ FlatMapCoGroupsInPandas(leftExprs, rightExprs, _, _, left, right) =>
+        f.copy(
+          leftExprs = leftExprs.map(x => resolveExpressionBottomUp(x, left)),
+          rightExprs = rightExprs.map(x => resolveExpressionBottomUp(x, right)))
       // intersect/except will be rewritten to join at the begininng of optimizer. Here we need to
       // deduplicate the right side plan, so that we won't produce an invalid self-join later.
       case i @ Intersect(left, right, _) if !i.duplicateResolved =>

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/namedExpressions.scala

@@ -23,14 +23,18 @@ import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.catalyst.analysis.UnresolvedAttribute
 import org.apache.spark.sql.catalyst.expressions.codegen._
 import org.apache.spark.sql.catalyst.plans.logical.EventTimeWatermark
-import org.apache.spark.sql.catalyst.util.quoteIdentifier
+import org.apache.spark.sql.catalyst.util.{quoteIdentifier, toPrettySQL}
 import org.apache.spark.sql.types._
 
 object NamedExpression {
   private val curId = new java.util.concurrent.atomic.AtomicLong()
   private[expressions] val jvmId = UUID.randomUUID()
   def newExprId: ExprId = ExprId(curId.getAndIncrement(), jvmId)
   def unapply(expr: NamedExpression): Option[(String, DataType)] = Some((expr.name, expr.dataType))
+  def fromExpression(expr: Expression): NamedExpression = expr match {
+    case ne: NamedExpression => ne
+    case _: Expression => Alias(expr, toPrettySQL(expr))()
+  }
 }
 
 /**

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/pythonLogicalOperators.scala

@@ -24,7 +24,7 @@ import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeSet, Expre
  * This is used by DataFrame.groupby().apply().
  */
 case class FlatMapGroupsInPandas(
-    groupingAttributes: Seq[Attribute],
+    groupingExprs: Seq[Expression],
     functionExpr: Expression,
     output: Seq[Attribute],
     child: LogicalPlan) extends UnaryNode {
@@ -56,8 +56,8 @@ case class MapInPandas(
  * This is used by DataFrame.groupby().cogroup().apply().
  */
 case class FlatMapCoGroupsInPandas(
-    leftAttributes: Seq[Attribute],
-    rightAttributes: Seq[Attribute],
+    leftExprs: Seq[Expression],
+    rightExprs: Seq[Expression],
     functionExpr: Expression,
     output: Seq[Attribute],
     left: LogicalPlan,

sql/core/src/main/scala/org/apache/spark/sql/RelationalGroupedDataset.scala

@@ -541,17 +541,8 @@ class RelationalGroupedDataset protected[sql](
       "Must pass a grouped map udf")
     require(expr.dataType.isInstanceOf[StructType],
       s"The returnType of the udf must be a ${StructType.simpleString}")
-
-    val groupingNamedExpressions = groupingExprs.map {
-      case ne: NamedExpression => ne
-      case other => Alias(other, other.toString)()
-    }
-    val groupingAttributes = groupingNamedExpressions.map(_.toAttribute)
-    val child = df.logicalPlan
-    val project = Project(groupingNamedExpressions ++ child.output, child)
     val output = expr.dataType.asInstanceOf[StructType].toAttributes
-    val plan = FlatMapGroupsInPandas(groupingAttributes, expr, output, project)
-
+    val plan = FlatMapGroupsInPandas(groupingExprs, expr, output, df.logicalPlan)
     Dataset.ofRows(df.sparkSession, plan)
   }
 
@@ -572,28 +563,9 @@ class RelationalGroupedDataset protected[sql](
       "Must pass a cogrouped map udf")
     require(expr.dataType.isInstanceOf[StructType],
       s"The returnType of the udf must be a ${StructType.simpleString}")
-
-    val leftGroupingNamedExpressions = groupingExprs.map {
-      case ne: NamedExpression => ne
-      case other => Alias(other, other.toString)()
-    }
-
-    val rightGroupingNamedExpressions = r.groupingExprs.map {
-      case ne: NamedExpression => ne
-      case other => Alias(other, other.toString)()
-    }
-
-    val leftAttributes = leftGroupingNamedExpressions.map(_.toAttribute)
-    val rightAttributes = rightGroupingNamedExpressions.map(_.toAttribute)
-
-    val leftChild = df.logicalPlan
-    val rightChild = r.df.logicalPlan
-
-    val left = Project(leftGroupingNamedExpressions ++ leftChild.output, leftChild)
-    val right = Project(rightGroupingNamedExpressions ++ rightChild.output, rightChild)
-
     val output = expr.dataType.asInstanceOf[StructType].toAttributes
-    val plan = FlatMapCoGroupsInPandas(leftAttributes, rightAttributes, expr, output, left, right)
+    val plan = FlatMapCoGroupsInPandas(
+      groupingExprs, r.groupingExprs, expr, output, df.logicalPlan, r.df.logicalPlan)
     Dataset.ofRows(df.sparkSession, plan)
   }
 

sql/core/src/main/scala/org/apache/spark/sql/execution/SparkStrategies.scala

@@ -608,10 +608,14 @@ abstract class SparkStrategies extends QueryPlanner[SparkPlan] {
         execution.MapPartitionsInRWithArrowExec(
           f, p, b, is, ot, planLater(child)) :: Nil
       case logical.FlatMapGroupsInPandas(grouping, func, output, child) =>
-        execution.python.FlatMapGroupsInPandasExec(grouping, func, output, planLater(child)) :: Nil
-      case logical.FlatMapCoGroupsInPandas(leftGroup, rightGroup, func, output, left, right) =>
+        val groupingExprs = grouping.map(NamedExpression.fromExpression)
+        execution.python.FlatMapGroupsInPandasExec(
+          groupingExprs, func, output, planLater(child)) :: Nil
+      case logical.FlatMapCoGroupsInPandas(leftExprs, rightExprs, func, output, left, right) =>
+        val leftAttrs = leftExprs.map(NamedExpression.fromExpression)
+        val rightAttrs = rightExprs.map(NamedExpression.fromExpression)
         execution.python.FlatMapCoGroupsInPandasExec(
-          leftGroup, rightGroup, func, output, planLater(left), planLater(right)) :: Nil
+          leftAttrs, rightAttrs, func, output, planLater(left), planLater(right)) :: Nil
       case logical.MapInPandas(func, output, child) =>
         execution.python.MapInPandasExec(func, output, planLater(child)) :: Nil
       case logical.MapElements(f, _, _, objAttr, child) =>

sql/core/src/main/scala/org/apache/spark/sql/execution/python/FlatMapCoGroupsInPandasExec.scala

@@ -48,8 +48,8 @@ import org.apache.spark.sql.util.ArrowUtils
  * is left as future work.
  */
 case class FlatMapCoGroupsInPandasExec(
-    leftGroup: Seq[Attribute],
-    rightGroup: Seq[Attribute],
+    leftGroupingExprs: Seq[NamedExpression],
+    rightGroupingExprs: Seq[NamedExpression],
     func: Expression,
     output: Seq[Attribute],
     left: SparkPlan,
@@ -60,42 +60,51 @@ case class FlatMapCoGroupsInPandasExec(
   private val pythonRunnerConf = ArrowUtils.getPythonRunnerConfMap(conf)
   private val pandasFunction = func.asInstanceOf[PythonUDF].func
   private val chainedFunc = Seq(ChainedPythonFunctions(Seq(pandasFunction)))
+  private val inputExprs =
+    func.asInstanceOf[PythonUDF].children.map(_.asInstanceOf[NamedExpression])
+  private val leftExprs =
+    left.output.filter(e => inputExprs.exists(_.semanticEquals(e)))
+  private val rightExprs =
+    right.output.filter(e => inputExprs.exists(_.semanticEquals(e)))
 
   override def producedAttributes: AttributeSet = AttributeSet(output)
 
   override def outputPartitioning: Partitioning = left.outputPartitioning
 
   override def requiredChildDistribution: Seq[Distribution] = {
-    val leftDist = if (leftGroup.isEmpty) AllTuples else ClusteredDistribution(leftGroup)
-    val rightDist = if (rightGroup.isEmpty) AllTuples else ClusteredDistribution(rightGroup)
+    val leftDist =
+      if (leftGroupingExprs.isEmpty) AllTuples else ClusteredDistribution(leftGroupingExprs)
+    val rightDist =
+      if (rightGroupingExprs.isEmpty) AllTuples else ClusteredDistribution(rightGroupingExprs)
     leftDist :: rightDist :: Nil
   }
 
   override def requiredChildOrdering: Seq[Seq[SortOrder]] = {
-    leftGroup
-      .map(SortOrder(_, Ascending)) :: rightGroup.map(SortOrder(_, Ascending)) :: Nil
+    leftGroupingExprs
+      .map(SortOrder(_, Ascending)) :: rightGroupingExprs.map(SortOrder(_, Ascending)) :: Nil
   }
 
   override protected def doExecute(): RDD[InternalRow] = {
 
-    val (leftDedup, leftArgOffsets) = resolveArgOffsets(left, leftGroup)
-    val (rightDedup, rightArgOffsets) = resolveArgOffsets(right, rightGroup)
+    val (leftDedup, leftArgOffsets) = resolveArgOffsets(leftExprs, leftGroupingExprs)
+    val (rightDedup, rightArgOffsets) = resolveArgOffsets(rightExprs, rightGroupingExprs)
 
     // Map cogrouped rows to ArrowPythonRunner results, Only execute if partition is not empty
     left.execute().zipPartitions(right.execute())  { (leftData, rightData) =>
       if (leftData.isEmpty && rightData.isEmpty) Iterator.empty else {
 
-        val leftGrouped = groupAndProject(leftData, leftGroup, left.output, leftDedup)
-        val rightGrouped = groupAndProject(rightData, rightGroup, right.output, rightDedup)
-        val data = new CoGroupedIterator(leftGrouped, rightGrouped, leftGroup)
+        val leftGrouped = groupAndProject(leftData, leftGroupingExprs, left.output, leftDedup)
+        val rightGrouped = groupAndProject(rightData, rightGroupingExprs, right.output, rightDedup)
+        val data = new CoGroupedIterator(
+          leftGrouped, rightGrouped, leftGroupingExprs.map(_.toAttribute))
           .map { case (_, l, r) => (l, r) }
 
         val runner = new CoGroupedArrowPythonRunner(
           chainedFunc,
           PythonEvalType.SQL_COGROUPED_MAP_PANDAS_UDF,
           Array(leftArgOffsets ++ rightArgOffsets),
-          StructType.fromAttributes(leftDedup),
-          StructType.fromAttributes(rightDedup),
+          StructType.fromAttributes(leftDedup.map(_.toAttribute)),
+          StructType.fromAttributes(rightDedup.map(_.toAttribute)),
           sessionLocalTimeZone,
           pythonRunnerConf)
 

sql/core/src/main/scala/org/apache/spark/sql/execution/python/FlatMapGroupsInPandasExec.scala

@@ -46,7 +46,7 @@ import org.apache.spark.sql.util.ArrowUtils
  * is left as future work.
  */
 case class FlatMapGroupsInPandasExec(
-    groupingAttributes: Seq[Attribute],
+    groupingExprs: Seq[NamedExpression],
     func: Expression,
     output: Seq[Attribute],
     child: SparkPlan)
@@ -56,38 +56,39 @@ case class FlatMapGroupsInPandasExec(
   private val pythonRunnerConf = ArrowUtils.getPythonRunnerConfMap(conf)
   private val pandasFunction = func.asInstanceOf[PythonUDF].func
   private val chainedFunc = Seq(ChainedPythonFunctions(Seq(pandasFunction)))
+  private val inputExprs =
+    func.asInstanceOf[PythonUDF].children.map(_.asInstanceOf[NamedExpression])
 
   override def producedAttributes: AttributeSet = AttributeSet(output)
 
   override def outputPartitioning: Partitioning = child.outputPartitioning
 
   override def requiredChildDistribution: Seq[Distribution] = {
-    if (groupingAttributes.isEmpty) {
+    if (groupingExprs.isEmpty) {
       AllTuples :: Nil
     } else {
-      ClusteredDistribution(groupingAttributes) :: Nil
+      ClusteredDistribution(groupingExprs) :: Nil
     }
   }
 
   override def requiredChildOrdering: Seq[Seq[SortOrder]] =
-    Seq(groupingAttributes.map(SortOrder(_, Ascending)))
+    Seq(groupingExprs.map(SortOrder(_, Ascending)))
 
   override protected def doExecute(): RDD[InternalRow] = {
     val inputRDD = child.execute()
 
-    val (dedupAttributes, argOffsets) = resolveArgOffsets(child, groupingAttributes)
-
+    val (dedupExprs, argOffsets) = resolveArgOffsets(inputExprs, groupingExprs)
     // Map grouped rows to ArrowPythonRunner results, Only execute if partition is not empty
     inputRDD.mapPartitionsInternal { iter => if (iter.isEmpty) iter else {
 
-      val data = groupAndProject(iter, groupingAttributes, child.output, dedupAttributes)
+      val data = groupAndProject(iter, groupingExprs, child.output, dedupExprs)
         .map { case (_, x) => x }
 
       val runner = new ArrowPythonRunner(
         chainedFunc,
         PythonEvalType.SQL_GROUPED_MAP_PANDAS_UDF,
         Array(argOffsets),
-        StructType.fromAttributes(dedupAttributes),
+        StructType.fromAttributes(dedupExprs.map(_.toAttribute)),
         sessionLocalTimeZone,
         pythonRunnerConf)