PySpark如何从数组列中找到匹配另一列子串的首个元素？

问题分析 #

你原代码的错误主要有几点：

• df.getItem()用法错误：getItem是列对象的方法（如df.col_a.getItem(0)），不能直接调用DataFrame的getItem
• 数组过滤函数误用：PySpark中过滤数组元素需要用array_filter，而非原生filter
• 变量引用错误：lambda里的变量大小写不一致（i写成I），且'col_b'是字符串常量，没有引用列对象，应该用F.col("col_b")
• 数组索引语法错误：PySpark中不能直接用[0]访问数组元素，需用element_at或getItem方法

正确实现方法 #

步骤1：构建示例数据（可跳过，直接用你的DataFrame） #

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

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

data = [
    (["AB1 2Z", "CD3 4Y", "AB5 6X"], "AB"),
    (["GH7 8W", "EF9 0V", "EF1 2U"], "EF"),
    (["IJ3 4T", "KL5 6S"], "KL")
]

df = spark.createDataFrame(data, ["col_a", "col_b"])

步骤2：提取匹配的首个元素 #

使用array_filter过滤出符合条件的数组元素，再用element_at取第一个元素（PySpark中数组索引从1开始）：

df = df.withColumn(
    "col_c",
    F.element_at(
        F.array_filter("col_a", lambda x: F.startsWith(x, F.col("col_b"))),
        1
    )
)

# 查看结果
df.select("col_c").show(truncate=False)

执行后输出：

+-------+
|col_c  |
+-------+
|AB1 2Z |
|EF9 0V |
|KL5 6S |
+-------+

备选方案：SQL表达式实现 #

如果你更习惯SQL语法，可以通过临时视图实现：

df.createOrReplaceTempView("temp_table")
df = spark.sql("""
    SELECT 
        element_at(array_filter(col_a, x -> startsWith(x, col_b)), 1) AS col_c
    FROM temp_table
""")

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