PySpark中提取数组元素：获取address_line1数组指定位置元素

首先，你的代码报错的原因很明确：df.selectExpr("address_line1[1]")返回的是一个DataFrame对象，但withColumn方法的第二个参数需要的是一个Column表达式，两者类型不匹配，所以Spark会抛出"cannot resolve"的错误。

接下来针对你的需求——保留原数组列的同时，按address_line1[1]、address_line1[2]、address_line1[3]的方式提取元素（不存在的位置返回null），给你几种可行的实现方法：

方法一：使用Spark SQL表达式（最贴合你的需求写法） #

Spark SQL中数组是1-based索引，直接用expr函数编写SQL风格的索引表达式即可，不存在的索引会自动返回null，完美匹配你的要求：

from pyspark.sql.functions import expr

# 新增列提取数组的第1、2、3个元素，原数组列保留
df = df.withColumn("address_line1_1", expr("address_line1[1]")) \
       .withColumn("address_line1_2", expr("address_line1[2]")) \
       .withColumn("address_line1_3", expr("address_line1[3]"))

# 查看结果
df.show()

方法二：使用PySpark原生API的getItem方法 #

如果你更习惯用PySpark的面向对象写法，getItem方法是0-based索引（和Python列表一致），要对应SQL的address_line1[1]需要用getItem(0)，以此类推：

from pyspark.sql.functions import col

df = df.withColumn("address_line1_1", col("address_line1").getItem(0)) \
       .withColumn("address_line1_2", col("address_line1").getItem(1)) \
       .withColumn("address_line1_3", col("address_line1").getItem(2))

方法三：封装成通用函数（复用性更强） #

如果需要多次提取不同数组列的元素，可以封装成一个通用函数，支持指定列名和要提取的索引：

from pyspark.sql.functions import expr

def extract_array_elements(df, col_name, indices):
    # 遍历索引，逐个新增提取列
    for idx in indices:
        df = df.withColumn(f"{col_name}_{idx}", expr(f"{col_name}[{idx}]"))
    return df

# 使用示例：提取address_line1的第1、2、3个元素
df = extract_array_elements(df, "address_line1", [1, 2, 3])

结果验证 #

针对你的示例数据，运行上述代码后会得到如下结果（原数组列保留，新增列对应提取的元素，不存在的位置返回null）：

+--------------------+-------------------+-------------------+-------------------+
|       address_line1|address_line1_1    |address_line1_2    |address_line1_3    |
+--------------------+-------------------+-------------------+-------------------+
| [atmosphere e 20]  |atmosphere e 20    |null               |null               |
| [tennesse row 3]   |tennesse row 3     |null               |null               |
|                null|null               |null               |null               |
+--------------------+-------------------+-------------------+-------------------+

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