Spark 2.3中如何实现arrays_zip等效功能？

Spark 2.4才正式引入arrays_zip函数，2.3版本里确实没有这个内置工具，但我们可以通过两种实用的方式复刻它的核心功能，下面分别介绍：

方法一：利用posexplode和分组聚合（无UDF方案） #

这种方法完全依赖Spark原生SQL函数，避免了UDF可能带来的性能开销，核心思路是把数组按位置炸开后重新聚合：

from pyspark.sql import SparkSession
from pyspark.sql.functions import posexplode, monotonically_increasing_id, collect_list, struct

# 初始化SparkSession
spark = SparkSession.builder.appName("arrays_zip_emulation").getOrCreate()

# 创建和官方示例一致的测试数据
df = spark.createDataFrame([(([1, 2, 3], [2, 3, 4]))], ['vals1', 'vals2'])

# 步骤1：添加唯一行标识，用来区分不同数据行
df_with_id = df.withColumn("row_id", monotonically_increasing_id())

# 步骤2：对每个数组列执行posexplode，提取位置索引和对应元素
df_explode_vals1 = df_with_id.select("row_id", posexplode("vals1").alias("pos", "vals1_val"))
df_explode_vals2 = df_with_id.select("row_id", posexplode("vals2").alias("pos", "vals2_val"))

# 步骤3：按行ID和位置关联，打包成struct后再按行ID聚合回数组
zipped_df = df_explode_vals1.join(df_explode_vals2, on=["row_id", "pos"], how="inner") \
    .groupBy("row_id") \
    .agg(collect_list(struct("vals1_val", "vals2_val")).alias("zipped")) \
    .drop("row_id")

# 查看结果
zipped_df.show(truncate=False)

运行后会得到和Spark 2.4arrays_zip完全一致的输出：

+----------------------------------------+
|zipped                                  |
+----------------------------------------+
|[{1, 2}, {2, 3}, {3, 4}]                |
+----------------------------------------+

方法二：使用自定义UDF（Python层面处理） #

如果更习惯Python的逻辑，也可以写一个UDF来模拟arrays_zip的行为，注意要和官方函数保持一致：当输入数组长度不一致时，取最短数组的长度作为结果数组的长度。

from pyspark.sql import SparkSession
from pyspark.sql.functions import udf
from pyspark.sql.types import ArrayType, StructType, StructField, IntegerType

# 初始化SparkSession
spark = SparkSession.builder.appName("arrays_zip_udf").getOrCreate()

# 定义和输入列匹配的StructType，这里对应两个Integer类型的字段
zip_schema = ArrayType(StructType([
    StructField("vals1", IntegerType()),
    StructField("vals2", IntegerType())
]))

# 实现UDF逻辑：用Python的zip函数打包，再转成字典列表（对应Spark的struct）
def arrays_zip_udf(arr1, arr2):
    # 处理空数组的情况，和官方行为对齐
    if not arr1 or not arr2:
        return []
    # 取最短数组的长度，和Spark 2.4的arrays_zip行为一致
    min_len = min(len(arr1), len(arr2))
    return [{"vals1": arr1[i], "vals2": arr2[i]} for i in range(min_len)]

# 注册UDF
zip_udf = udf(arrays_zip_udf, zip_schema)

# 创建测试数据
df = spark.createDataFrame([(([1, 2, 3], [2, 3, 4]))], ['vals1', 'vals2'])

# 调用UDF
result_df = df.select(zip_udf(df.vals1, df.vals2).alias("zipped"))

# 查看结果
result_df.show(truncate=False)

这个方法的输出和官方函数完全一致，而且可以轻松扩展到更多数组列的情况（只需要修改schema和UDF的参数即可）。

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