[SPARK-42168][3.2][SQL][PYTHON] Fix required child distribution of FlatMapCoGroupsInPandas (as in CoGroup)

[EnricoMi]

What changes were proposed in this pull request?

Make FlatMapCoGroupsInPandas (used by PySpark) report its required child distribution as HashClusteredDistribution, rather than ClusteredDistribution. That is the same distribution as reported by CoGroup (used by Scala).

Why are the changes needed?

This allows the EnsureRequirements rule to correctly recognizes that FlatMapCoGroupsInPandas requiring HashClusteredDistribution(id, day) is not compatible with HashPartitioning(day, id), while ClusteredDistribution(id, day) is compatible with HashPartitioning(day, id).

The following example returns an incorrect result in Spark 3.0, 3.1, and 3.2.

import org.apache.spark.sql.expressions.Window
import org.apache.spark.sql.functions.{col, lit, sum}

val ids = 1000
val days = 1000
val parts = 10

val id_df = spark.range(ids)
val day_df = spark.range(days).withColumnRenamed("id", "day")
val id_day_df = id_df.join(day_df)
// these redundant aliases are needed to workaround bug SPARK-42132
val left_df = id_day_df.select($"id".as("id"), $"day".as("day"), lit("left").as("side")).repartition(parts).cache()
val right_df = id_day_df.select($"id".as("id"), $"day".as("day"), lit("right").as("side")).repartition(parts).cache()  //.withColumnRenamed("id", "id2")

// note the column order is different to the groupBy("id", "day") column order below
val window = Window.partitionBy("day", "id")

case class Key(id: BigInt, day: BigInt)
case class Value(id: BigInt, day: BigInt, side: String)
case class Sum(id: BigInt, day: BigInt, side: String, day_sum: BigInt)

val left_grouped_df = left_df.groupBy("id", "day").as[Key, Value]
val right_grouped_df = right_df.withColumn("day_sum", sum(col("day")).over(window)).groupBy("id", "day").as[Key, Sum]

val df = left_grouped_df.cogroup(right_grouped_df)((key: Key, left: Iterator[Value], right: Iterator[Sum]) => left)

df.explain()
df.show(5)

Output was

== Physical Plan ==
AdaptiveSparkPlan isFinalPlan=false
+- FlatMapCoGroupsInPandas [id#8L, day#9L], [id#29L, day#30L], cogroup(id#8L, day#9L, side#10, id#29L, day#30L, side#31, day_sum#54L), [id#64L, day#65L, lefts#66, rights#67]
   :- Sort [id#8L ASC NULLS FIRST, day#9L ASC NULLS FIRST], false, 0
   :  +- Exchange hashpartitioning(id#8L, day#9L, 200), ENSURE_REQUIREMENTS, [plan_id=117]
   :     +- ...
   +- Sort [id#29L ASC NULLS FIRST, day#30L ASC NULLS FIRST], false, 0
      +- Project [id#29L, day#30L, id#29L, day#30L, side#31, day_sum#54L]
         +- Window [sum(day#30L) windowspecdefinition(day#30L, id#29L, specifiedwindowframe(RowFrame, unboundedpreceding$(), unboundedfollowing$())) AS day_sum#54L], [day#30L, id#29L]
            +- Sort [day#30L ASC NULLS FIRST, id#29L ASC NULLS FIRST], false, 0
               +- Exchange hashpartitioning(day#30L, id#29L, 200), ENSURE_REQUIREMENTS, [plan_id=112]
                  +- ...

+---+---+-----+------+
| id|day|lefts|rights|
+---+---+-----+------+
|  0|  3|    0|     1|
|  0|  4|    0|     1|
|  0| 13|    1|     0|
|  0| 27|    0|     1|
|  0| 31|    0|     1|
+---+---+-----+------+
only showing top 5 rows

Output now is

== Physical Plan ==
AdaptiveSparkPlan isFinalPlan=false
+- FlatMapCoGroupsInPandas [id#8L, day#9L], [id#29L, day#30L], cogroup(id#8L, day#9L, side#10, id#29L, day#30L, side#31, day_sum#54L), [id#64L, day#65L, lefts#66, rights#67]
   :- Sort [id#8L ASC NULLS FIRST, day#9L ASC NULLS FIRST], false, 0
   :  +- Exchange hashpartitioning(id#8L, day#9L, 200), ENSURE_REQUIREMENTS, [plan_id=117]
   :     +- ...
   +- Sort [id#29L ASC NULLS FIRST, day#30L ASC NULLS FIRST], false, 0
      +- Exchange hashpartitioning(id#29L, day#30L, 200), ENSURE_REQUIREMENTS, [plan_id=118]
         +- Project [id#29L, day#30L, id#29L, day#30L, side#31, day_sum#54L]
            +- Window [sum(day#30L) windowspecdefinition(day#30L, id#29L, specifiedwindowframe(RowFrame, unboundedpreceding$(), unboundedfollowing$())) AS day_sum#54L], [day#30L, id#29L]
               +- Sort [day#30L ASC NULLS FIRST, id#29L ASC NULLS FIRST], false, 0
                  +- Exchange hashpartitioning(day#30L, id#29L, 200), ENSURE_REQUIREMENTS, [plan_id=112]
                     +- ...

+---+---+-----+------+
| id|day|lefts|rights|
+---+---+-----+------+
|  0| 13|    1|     1|
|  0| 63|    1|     1|
|  0| 89|    1|     1|
|  0| 95|    1|     1|
|  0| 96|    1|     1|
+---+---+-----+------+
only showing top 5 rows

Spark 3.3 reworked HashClusteredDistribution, and is not sensitive to using ClusteredDistribution: #32875

Does this PR introduce any user-facing change?

This fixes correctness.

How was this patch tested?

A unit test in EnsureRequirementsSuite.

[HyukjinKwon]

cc @sunchao

[sunchao]

Looks OK to me, just a few questions:

• We probably should change the title to use [3.2] instead of [3.3]?
• Is it easy to add a E2E test for this? perhaps in test_pandas_udf_cogrouped_map.py?

[EnricoMi]

@sunchao good catch!

I have renamed the PR and added the Python example as a unit test.

[HyukjinKwon]

I am good w/ this. will defer to @sunchao

[EnricoMi]

Tests are green: https://github.com/G-Research/spark/actions/runs/4003829481

[sunchao]

LGTM

[HyukjinKwon]

Merged to branch-3.2.

[dongjoon-hyun]

cc @kazuyukitanimura since this landed to branch-3.2.