Spark：如何展开长度可变且含空值的多数组列并满足指定规则？

需求规则 #

需要在Spark中处理数组列展开需求，满足以下规则：

• 相同索引的数组元素映射到同一行；
• 若某列仅含1个元素，需将该元素应用到所有展开后的行；
• 若数组为null，需将null应用到所有行。

输入数据 #

+----+------------+--------------+--------------------+
|col1|        col2|          col3|                col4|
+----+------------+--------------+--------------------+
|   1|[id_1, id_2]|  [tim, steve]|       [apple, pear]|
|   2|[id_3, id_4]|       [jenny]|           [avocado]|
|   3|        null|[tommy, megan]| [apple, strawberry]|
|   4|        null|          null|[banana, strawberry]|
+----+------------+--------------+--------------------+

期望输出 #

+----+----+-----+----------+
|col1|col2|col3 |col4      |
+----+----+-----+----------+
|1   |id_1|tim  |apple     |
|1   |id_2|steve|pear      |
|2   |id_3|jenny|avocado   |
|2   |id_4|jenny|avocado   |
|3   |null|tommy|apple     |
|3   |null|megan|strawberry|
|4   |null|null |banana    |
|4   |null|null |strawberry|
+----+----+-----+----------+

现有实现代码 #

已通过以下代码实现需求，但希望找到更简洁的写法：

df = spark.createDataFrame(
    [
        (1, ["id_1", "id_2"], ["tim", "steve"], ["apple", "pear"]),
        (2, ["id_3", "id_4"], ["jenny"], ["avocado"]),
        (3, None, ["tommy", "megan"], ["apple", "strawberry"]),
        (4, None, None, ["banana", "strawberry"])
    ],
    ["col1", "col2", "col3", "col4"]
)

df.createOrReplaceTempView("my_table")


spark.sql("""
with cte as (
  SELECT
  col1,
  col2,
  col3,
  col4,
  greatest(size(col2), size(col3), size(col4)) as max_array_len
  FROM my_table
), arrays_extended as (
select 
col1,
case 
  when col2 is null then array_repeat(null, max_array_len) 
  else col2 
end as col2, 
case
  when size(col3) = 1 then array_repeat(col3[0], max_array_len)
  when col3 is null then array_repeat(null, max_array_len)
  else col3
end as col3,
case
  when size(col4) = 1 then array_repeat(col4[0], max_array_len)
  when col4 is null then array_repeat(null, max_array_len)
  else col4
end as col4
from cte), 
arrays_zipped as (
select *, explode(arrays_zip(col2, col3, col4)) as zipped
from arrays_extended
)
select 
  col1,
  zipped.col2,
  zipped.col3,
  zipped.col4
from arrays_zipped
""").show(truncate=False)

更简洁的实现方案 #

可以通过**posexplode**结合条件判断简化逻辑，无需提前扩展数组，直接基于最长数组的位置索引匹配其他列元素：

SQL版本 #

df.createOrReplaceTempView("my_table")

spark.sql("""
SELECT 
    col1,
    CASE 
        WHEN col2 IS NULL THEN NULL
        WHEN size(col2) = 1 THEN col2[0]
        ELSE col2[pos]
    END AS col2,
    CASE 
        WHEN col3 IS NULL THEN NULL
        WHEN size(col3) = 1 THEN col3[0]
        ELSE col3[pos]
    END AS col3,
    CASE 
        WHEN col4 IS NULL THEN NULL
        WHEN size(col4) = 1 THEN col4[0]
        ELSE col4[pos]
    END AS col4
FROM my_table
LATERAL VIEW posexplode(
    sequence(0, greatest(
        size(coalesce(col2, array())),
        size(coalesce(col3, array())),
        size(coalesce(col4, array()))
    ) - 1)
) AS pos, val
""").show(truncate=False)

DataFrame API版本 #

from pyspark.sql import functions as F

# 计算每行的最大数组长度
max_len_expr = F.greatest(
    F.size(F.coalesce(F.col("col2"), F.array())),
    F.size(F.coalesce(F.col("col3"), F.array())),
    F.size(F.coalesce(F.col("col4"), F.array()))
)

# 生成位置序列并展开
df = df.withColumn("pos", F.explode(F.sequence(F.lit(0), max_len_expr - 1)))

# 处理各列的元素映射
df = df.withColumn("col2", 
    F.when(F.col("col2").isNull(), None)
     .when(F.size("col2") == 1, F.col("col2")[0])
     .otherwise(F.col("col2")[F.col("pos")])
).withColumn("col3", 
    F.when(F.col("col3").isNull(), None)
     .when(F.size("col3") == 1, F.col("col3")[0])
     .otherwise(F.col("col3")[F.col("pos")])
).withColumn("col4", 
    F.when(F.col("col4").isNull(), None)
     .when(F.size("col4") == 1, F.col("col4")[0])
     .otherwise(F.col("col4")[F.col("pos")])
).drop("pos")

df.show(truncate=False)

方案优势 #

1. 省去多层CTE嵌套，逻辑线性直观；
2. 无需提前扩展数组，减少内存开销；
3. 统一处理null数组、单元素数组和多元素数组的逻辑，复用性更强。

内容的提问来源于stack exchange，提问作者Dumbledore__