如何借助第二列在PySpark中获取元素索引？含array_position用法示例

需求说明 #

给定包含数组的col1列和目标值的col2列，需要实现：

• 找到col2的值在col1数组中的0-based索引
• 当col2为空值时，索引列返回空值

输入示例：

col1                                col2             
['apple','sib','mive','banana']     'sib'     
['jasmin','df','nagh']               'nagh'
['jk']                                na

期望输出：

col1                                col2             index 
['apple','sib','mive','banana']     'sib'             1
['jasmin','df','nagh']               'nagh'           2
['jk']                                na              na

实现方案 #

PySpark的array_position函数可直接满足需求，该函数返回目标值在数组中的1-based位置，若目标值为空或不存在于数组中则返回null。结合0-based索引需求，只需对结果减1即可。

步骤1：创建测试DataFrame

from pyspark.sql import SparkSession
from pyspark.sql.functions import col, array_position

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

# 构造测试数据
data = [
    (["apple", "sib", "mive", "banana"], "sib"),
    (["jasmin", "df", "nagh"], "nagh"),
    (["jk"], None)
]

df = spark.createDataFrame(data, schema=["col1", "col2"])

步骤2：添加索引列

# 计算0-based索引，自动处理空值和元素不存在的情况
result_df = df.withColumn(
    "index",
    array_position(col("col1"), col("col2")) - 1
)

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

输出结果

+--------------------------------+-----+-----+
|col1                            |col2 |index|
+--------------------------------+-----+-----+
|[apple, sib, mive, banana]      |sib  |1    |
|[jasmin, df, nagh]              |nagh |2    |
|[jk]                            |null |null |
+--------------------------------+-----+-----+

关键说明 #

1. array_position函数特性：
   • 接受数组列和目标值列两个参数
   • 返回目标值在数组中的1-based位置，目标值为空或不存在时返回null
2. 索引转换：
   • 若需要1-based索引，直接使用array_position(col("col1"), col("col2"))
   • 若需要0-based索引，对结果减1即可，空值减1后仍为null，无需额外判断

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