PySpark中能否提取数组任意元素？无需UDF实现按索引列取值？

可以不用编写UDF，直接通过Spark内置函数组合实现需求，以下是两种常用方案：

方案一：使用transform + array_get（推荐，高效无Shuffle） #

利用transform遍历indices数组中的每个索引，再通过array_get从text数组中提取对应位置的元素，这是最简洁高效的方式：

from pyspark.sql import functions as F

# 修正原示例代码的语法错误（补充逗号）
df = spark.createDataFrame(
   [
      {
         'text': ['0', '1', '2', '3', '4', '5'],
         'indices': [0, 2, 4],
      },
   ]
)

# 生成提取后的列
result_df = df.withColumn('extracted', F.transform('indices', lambda idx: F.array_get('text', idx)))
result_df.show(truncate=False)

执行后输出：

+----------------------+---------+-----------+
|text                  |indices  |extracted  |
+----------------------+---------+-----------+
|[0, 1, 2, 3, 4, 5]|[0, 2, 4]|[0, 2, 4]|
+----------------------+---------+-----------+

方案二：使用posexplode + 过滤 + 聚合 #

如果需要对中间元素做额外处理，可以先将text拆分为索引和元素，过滤出索引在indices中的元素后再聚合：

result_df = df.select(
    '*',
    F.posexplode('text').alias('pos', 'val')
).filter(F.array_contains('indices', F.col('pos')))
.groupBy('text', 'indices')
.agg(F.collect_list('val').alias('extracted'))

result_df.show(truncate=False)

这个方案会产生数据炸开和聚合的Shuffle操作，适合复杂场景，简单提取需求优先用方案一。

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