[FEA][AUDIT][SPARK-52921][SQL] Specify outputPartitioning for UnionExec for same output partitoning as children operators

[abellina]

This is a spark 4.1 audit task.

UnionExec supports outputPartitioning, when children have the same output partitioning. I am not 100% sure why we would want this, but filing in case we need it. According to the issue, Union has "unknown" partitioning otherwise. My guess is that if we don't follow suit, spark-rapids could add shuffles when the CPU case wouldn't. That said, I am not 100% sure, and wanted to get some comments from @revans2 and others.

There is a follow up I was looking at for audit (apache/spark@c0acf45023f), so I created the issue for the main spark issue to discuss.

Here's another related follow up: apache/spark@8edc7685b97

[revans2]

We kind of do need this, but it might just be a nice to have. Spark uses the output partitioning and the required partitioning to decide if it needs to insert in a shuffle to make the partitioning what it needs. I think all of that has been planned out before us already, but it would be good to try and match it just in case. Especially because Union is a logical operation only. We don't do anything on the GPU to make the union happen.

[nartal1]

We are seeing output mismatch due to this in Spark-4.1. Some of the integration tests are failing due to this.
Minimum repro case. Below query has a union of 2 identical tables and then a group by. GPU result has double the output rows compared to CPU output rows. Reason is that shuffle is not inserted so we end up showing the union of 2 tables as the result.
I have posted the Spark plan below:
cc: @gerashegalov

spark.conf.set("spark.sql.adaptive.enabled", "false")
val factData = for {
  key <- 0 until 10
  mult <- 1 to 10
} yield (key, key % 5, 3, key * 100 + mult)

// dim: 10 keys (0-9), all with ex_key=3, filter=42
val dimData = (0 until 10).map(key => (key, key % 5, 3, key * 10, 42))
val factDF = factData.toDF("key", "skey", "ex_key", "value")
val dimDF = dimData.toDF("key", "skey", "ex_key", "value", "filter")
factDF.createOrReplaceTempView("fact_table_raw")
dimDF.createOrReplaceTempView("dim_table_raw")

val testQuery = """
  WITH fact_table1 AS (
      SELECT key, sum(value) as value FROM fact_table_raw GROUP BY key
  ),
  fact_table2 AS (
      SELECT key, sum(value) as value FROM fact_table_raw GROUP BY key
  ),
  dim_table AS (
      SELECT key, value, filter FROM dim_table_raw WHERE ex_key = 3
  )
  SELECT key, sum(value) FROM (
    SELECT f.key as key, f.value as value
    FROM fact_table1 f
    JOIN dim_table d ON f.key = d.key
    WHERE d.filter = 42
    UNION ALL
    SELECT f.key as key, f.value as value
    FROM fact_table2 f
    JOIN dim_table d ON f.key = d.key
    WHERE d.filter = 42
  ) GROUP BY key
"""

val result = spark.sql(testQuery)
result.collect()

GPU result and Physical Plan:

Array([1,1055], [6,6055], [3,3055], [5,5055], [9,9055], [4,4055], [8,8055], [7,7055], [2,2055], [0,55], [1,1055], [6,6055], [3,3055], [5,5055], [9,9055], [4,4055], [8,8055], [7,7055], [2,2055], [0,55])

== Physical Plan ==
GpuColumnarToRow false, [loreId=34]
+- GpuHashAggregate (keys=[key#1028], functions=[gpubasicsum(value#1029L, LongType, true)], output=[key#1028, sum(value)#1048L]) [loreId=33]
   +- GpuUnion [loreId=32]
      :- GpuProject [key#356 AS key#1028, value#1032L AS value#1029L], true, [loreId=22]
      :  +- GpuBroadcastHashJoin [key#356], [key#381], Inner, GpuBuildRight, false, [loreId=21]
      :     :- GpuHashAggregate (keys=[key#356], functions=[gpubasicsum(value#359, LongType, true)], output=[key#356, value#1032L]) [loreId=18]
      :     :  +- GpuShuffleCoalesce 1073741824, [loreId=17]
      :     :     +- GpuColumnarExchange gpuhashpartitioning(key#356, 200, Murmur3Mode), ENSURE_REQUIREMENTS, [plan_id=1428], [loreId=16]
      :     :        +- GpuHashAggregate (keys=[key#356], functions=[partial_gpubasicsum(value#359, LongType, true)], output=[key#356, sum#1072L]) [loreId=15]
      :     :           +- GpuRowToColumnar targetsize(1073741824), [loreId=14]
      :     :              +- LocalTableScan [key#356, value#359]
      :     +- GpuBroadcastExchange HashedRelationBroadcastMode(List(cast(input[0, int, false] as bigint)),false), [plan_id=1432], [loreId=20]
      :        +- GpuRowToColumnar targetsize(1073741824), [loreId=19]
      :           +- LocalTableScan [key#381]
      +- GpuProject [key#1038, value#1033L], true, [loreId=31]
         +- GpuBroadcastHashJoin [key#1038], [key#1044], Inner, GpuBuildRight, false, [loreId=30]
            :- GpuHashAggregate (keys=[key#1038], functions=[gpubasicsum(value#1041, LongType, true)], output=[key#1038, value#1033L]) [loreId=27]
            :  +- GpuShuffleCoalesce 1073741824, [loreId=26]
            :     +- ReusedExchange [key#1038, sum#1074L], GpuColumnarExchange gpuhashpartitioning(key#356, 200, Murmur3Mode), ENSURE_REQUIREMENTS, [plan_id=1428], [loreId=16]
            +- ReusedExchange [key#1044], GpuBroadcastExchange HashedRelationBroadcastMode(List(cast(input[0, int, false] as bigint)),false), [plan_id=1432], [loreId=20]

CPU result and physical plan

val res17: Array[org.apache.spark.sql.Row] = Array([1,2110], [6,12110], [3,6110], [5,10110], [9,18110], [4,8110], [8,16110], [7,14110], [2,4110], [0,110])

== Physical Plan ==
*(5) HashAggregate(keys=[key#1292], functions=[sum(value#1293L)], output=[key#1292, sum(value)#1312L])
+- *(5) HashAggregate(keys=[key#1292], functions=[partial_sum(value#1293L)], output=[key#1292, sum#1334L])
   +- Union
      :- *(2) Project [key#356 AS key#1292, value#1296L AS value#1293L]
      :  +- *(2) BroadcastHashJoin [key#356], [key#381], Inner, BuildRight, false
      :     :- *(2) HashAggregate(keys=[key#356], functions=[sum(value#359)], output=[key#356, value#1296L])
      :     :  +- Exchange hashpartitioning(key#356, 200), ENSURE_REQUIREMENTS, [plan_id=1584]
      :     :     +- *(1) HashAggregate(keys=[key#356], functions=[partial_sum(value#359)], output=[key#356, sum#1336L])
      :     :        +- *(1) LocalTableScan [key#356, value#359]
      :     +- BroadcastExchange HashedRelationBroadcastMode(List(cast(input[0, int, false] as bigint)),false), [plan_id=1533]
      :        +- LocalTableScan [key#381]
      +- *(4) Project [key#1302, value#1297L]
         +- *(4) BroadcastHashJoin [key#1302], [key#1308], Inner, BuildRight, false
            :- *(4) HashAggregate(keys=[key#1302], functions=[sum(value#1305)], output=[key#1302, value#1297L])
            :  +- ReusedExchange [key#1302, sum#1338L], Exchange hashpartitioning(key#356, 200), ENSURE_REQUIREMENTS, [plan_id=1584]
            +- ReusedExchange [key#1308], BroadcastExchange HashedRelationBroadcastMode(List(cast(input[0, int, false] as bigint)),false), [plan_id=1533]

@revans2 - I think we should add Gpu version of PartitionerAwareUnionRDD /SQLPartitioningAwareUnionRDD for Spark-4.1. Please let me know if that makes sense. I can work on it. Other option is to add a Rule in spark-rapids to insert a shuffle whenever there is an Union(and no shuffle in hashPartitioing).

[nartal1]

Created the PR to include GpuPartitionerAwareUnionRDD - #14164