在PySpark的实际数据处理工作中，经常需要处理数组类型的列，比如用户的历史行为标签、商品的多维度属性等。有时我们需要根据另一个数组列中存储的索引值，从目标数组中提取对应位置的元素，这种需求在特征筛选、数据子集提取场景中非常常见。
需求场景说明
假设我们有如下结构的数据：目标数组列target_array存储了用户的多个标签，索引数组列index_array存储了需要提取的标签对应的位置索引，我们需要生成一个新的数组列extracted_elements，包含目标数组中对应索引的元素。示例如下：
| id | target_array | index_array |
|---|---|---|
| 1 | [a, b, c, d, e] | [0, 2, 4] |
| 2 | [x, y, z] | [1, 2] |
期望的结果是id为1的记录提取出[a, c, e]，id为2的记录提取出[y, z]。
使用内置函数实现
PySpark的pyspark.sql.functions模块提供了transform函数，可以对数组中的每个元素执行指定的表达式，我们可以结合该函数和get方法实现索引提取逻辑。
代码实现示例
首先创建测试数据：
from pyspark.sql import SparkSession
from pyspark.sql.functions import col, transform, get
# 初始化SparkSession
spark = SparkSession.builder
.appName("array_extract_by_index")
.getOrCreate()
# 创建测试数据
data = [
(1, ["a", "b", "c", "d", "e"], [0, 2, 4]),
(2, ["x", "y", "z"], [1, 2]),
(3, ["m", "n"], [0, 3]) # 包含越界索引的情况
]
df = spark.createDataFrame(data, ["id", "target_array", "index_array"])
df.show(truncate=False)
接下来使用transform函数遍历索引数组，通过get方法从目标数组中获取对应元素：
# 基于索引数组提取目标数组元素
result_df = df.withColumn(
"extracted_elements",
transform(col("index_array"), lambda idx: get(col("target_array"), idx))
)
result_df.show(truncate=False)
执行上述代码后，输出结果如下：
+---+-------------------+-----------+-------------------+ |id |target_array |index_array|extracted_elements | +---+-------------------+-----------+-------------------+ |1 |[a, b, c, d, e] |[0, 2, 4] |[a, c, e] | |2 |[x, y, z] |[1, 2] |[y, z] | |3 |[m, n] |[0, 3] |[m, null] | +---+-------------------+-----------+-------------------+
可以看到，当索引越界时，get方法会返回null，如果需要过滤掉越界产生的null值，可以再结合filter函数处理。
使用自定义UDF实现
如果内置函数无法满足复杂的提取逻辑，也可以自定义UDF来实现。自定义UDF的灵活性更高，可以处理各种边界情况。
代码实现示例
from pyspark.sql.functions import udf
from pyspark.sql.types import ArrayType, StringType
# 定义自定义UDF函数
def extract_by_index(target_arr, index_arr):
if not target_arr or not index_arr:
return []
result = []
for idx in index_arr:
if 0 <= idx < len(target_arr):
result.append(target_arr[idx])
return result
# 注册UDF，指定返回类型为字符串数组
extract_udf = udf(extract_by_index, ArrayType(StringType()))
# 使用UDF提取元素
udf_result_df = df.withColumn(
"extracted_elements_udf",
extract_udf(col("target_array"), col("index_array"))
)
udf_result_df.show(truncate=False)
这种实现方式会自动过滤掉越界的索引，执行结果如下：
+---+-------------------+-----------+---------------------+ |id |target_array |index_array|extracted_elements_udf| +---+-------------------+-----------+---------------------+ |1 |[a, b, c, d, e] |[0, 2, 4] |[a, c, e] | |2 |[x, y, z] |[1, 2] |[y, z] | |3 |[m, n] |[0, 3] |[m] | +---+-------------------+-----------+---------------------+
两种方案对比
| 方案 | 优势 | 劣势 |
|---|---|---|
| 内置函数方案 | 无需自定义函数，执行效率更高，符合Spark的优化逻辑 | 越界索引会返回null，需要额外处理 |
| 自定义UDF方案 | 逻辑灵活，可自定义边界处理规则 | 执行效率低于内置函数，需要手动指定返回类型 |
实际使用中，如果不需要特殊的边界处理，优先选择内置函数方案；如果有复杂的过滤、转换逻辑，再考虑使用自定义UDF。