如何在PySpark中按元素合并多个等长数组？

要实现按元素逐个合并多个等长数组列，核心思路是拆分数组到对应位置的元素，对每个位置的元素执行coalesce，再重新组合为数组，具体步骤如下：

1. 确定数组长度：由于所有输入列的数组等长，取任意一列的长度即可。
2. 生成逐元素coalesce表达式：遍历数组的每个索引位置，对所有输入列的该位置元素执行coalesce，得到该位置的结果值。
3. 组合结果为数组：将所有位置的coalesce结果用array函数组合成新的数组列。

代码实现 #

修改测试用例中的关键部分，替换原来无效的coalesce调用：

def test_coalesce_elements():
    """
    Test array coalescing on a per-element basis
    """
    from pyspark.sql import SparkSession
    import pyspark.sql.types as t
    import pyspark.sql.functions as f

    spark = SparkSession.builder.getOrCreate()

    data = [
        {
            "a": [None, 1, None, None],
            "b": [2, 3, None, None],
            "c": [5, 6, 7, None],
        }
    ]

    schema = t.StructType([
        t.StructField('a', t.ArrayType(t.IntegerType())),
        t.StructField('b', t.ArrayType(t.IntegerType())),
        t.StructField('c', t.ArrayType(t.IntegerType())),
    ])
    df = spark.createDataFrame(data, schema)

    input_cols = ['a', 'b', 'c']
    
    # 获取数组长度（所有列等长，取第一列即可）
    array_length = df.select(f.size(f.col(input_cols[0]))).collect()[0][0]
    
    # 生成每个位置的coalesce表达式
    element_coalesces = [
        f.coalesce(*[f.col(col_name)[i] for col_name in input_cols])
        for i in range(array_length)
    ]
    
    # 组合为新数组列
    df = df.withColumn('d', f.array(*element_coalesces))

    # 验证结果
    assert df.collect()[0]['d'] == [2, 1, 7, None]

说明 #

• 动态适配任意列数：input_cols可以是任意数量的数组列，代码会自动遍历所有列的对应位置元素。
• 适配任意数组长度：通过size函数获取数组长度，无需硬编码索引范围。
• 性能优化：优先使用内置函数组合方案，PySpark优化器会自动处理这类逻辑，性能优于自定义UDF；如果遇到复杂业务场景，再考虑UDF方案。

备选UDF方案（适合复杂场景） #

如果内置函数组合无法满足需求，可以使用自定义UDF处理：

# 定义UDF：接收多个数组，返回逐元素coalesce后的数组
from pyspark.sql.functions import udf
from pyspark.sql.types import ArrayType, IntegerType

def element_wise_coalesce(*arrays):
    # 按位置zip所有数组，对每个位置的元素取第一个非空值
    return [next((x for x in elements if x is not None), None) for elements in zip(*arrays)]

element_coalesce_udf = udf(element_wise_coalesce, ArrayType(IntegerType()))

# 使用UDF生成新列
df = df.withColumn('d', element_coalesce_udf(*input_cols))

注意：UDF需要确保输入的所有数组长度一致，否则zip会截断到最短数组的长度，可根据业务场景添加长度校验逻辑。

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