如何在PySpark中将变长Array类型列拆分为两个子数组

假设你的DataFrame名为df，包含一个Array类型列（示例中每行是嵌套数组结构，外层数组仅含一个子数组），可以通过PySpark内置函数快速实现拆分需求：

1. 构造示例数据（可选，用于验证） #

from pyspark.sql import SparkSession
from pyspark.sql import functions as F

spark = SparkSession.builder.appName("array_split").getOrCreate()

data = [
    ([[1,2,3,4,6]],),
    ([[0,4,5,4,6,8,9,1]],),
    ([[1,2,3,4,6,2,4,5,6]],)
]
df = spark.createDataFrame(data, ["arr_col"])

2. 执行拆分操作 #

直接利用数组索引和SQL表达式完成切片，逻辑和纯Python的list[0]、list[1:]完全对应：

df_result = df.withColumn("first_element_arr", F.array(F.col("arr_col")[0][0])) \
              .withColumn("remaining_elements_arr", F.expr("arr_col[0][1:]"))

如果你的原列是直接的一维数组（即每行没有外层嵌套，值为[1,2,3,4,6]这类），简化代码即可：

df_result = df.withColumn("first_element_arr", F.array(F.col("arr_col")[0])) \
              .withColumn("remaining_elements_arr", F.expr("arr_col[1:]"))

3. 查看结果 #

执行df_result.show(truncate=False)会输出：

+--------------------------------+----------------+--------------------------------+
|arr_col                         |first_element_arr|remaining_elements_arr          |
+--------------------------------+----------------+--------------------------------+
|[[1, 2, 3, 4, 6]]               |[1]             |[2, 3, 4, 6]                    |
|[[0, 4, 5, 4, 6, 8, 9, 1]]      |[0]             |[4, 5, 4, 6, 8, 9, 1]           |
|[[1, 2, 3, 4, 6, 2, 4, 5, 6]]   |[1]             |[2, 3, 4, 6, 2, 4, 5, 6]        |
+--------------------------------+----------------+--------------------------------+

关键逻辑说明 #

• F.col("arr_col")[0][0]：按索引提取嵌套数组里的第一个元素，再用F.array()包装成数组格式
• F.expr("arr_col[0][1:]")：通过SQL表达式实现数组切片，和Python的list[1:]行为一致，直接获取从第二个元素开始的剩余部分

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