本文介绍如何在 PySpark 中根据 n_relevant 字段动态对固定长度数组（如300元素）执行三种策略：≤100时全量截取、100–299时按模2采样、≥300时统一按模3采样，确保输出数组长度始终≤100。
本文介绍如何在 pyspark 中根据 `n_relevant` 字段动态对固定长度数组（如300元素）执行三种策略：≤100时全量截取、100–299时按模2采样、≥300时统一按模3采样，确保输出数组长度始终≤100。
在 PySpark 数据处理中，常需对数组列进行条件化子集提取——尤其当目标系统（如下游数据库或 API）对数组长度有硬性限制（例如 ≤100）时。本例中，原始数组恒为 300 元素，而 n_relevant 字段指示逻辑上“有效”的元素数量，需据此智能裁剪：
- 若
n_relevant ≤ 100：直接取前n_relevant个元素； - 若
100 ：仅保留索引满足 <code>i % 2 == 0的元素（即偶数索引，等效于每2个取1个，结果约100个）； - 若
n_relevant ≥ 300：统一按i % 3 == 0采样（即每3个取1个，得100个元素）。
关键在于 避免 UDF（性能损耗大、无法下推优化），而应使用原生高阶函数组合：slice() 提前截断范围，再用 filter() 结合索引（index 参数）实现模运算筛选。注意：PySpark 的 filter(array, lambda elem, index: ...) 中 index 从 0 开始，但 slice(array, start, length) 的 start 是 1-based，因此 slice("array", 1, n) 可安全获取前 n 个元素，后续 filter 中的 index 即对应原始 0-based 索引。
以下是完整可运行示例（适配 Azure Databricks 环境）：
from pyspark.sql import functions as F
from pyspark.sql.types import StructType, StructField, IntegerType, ArrayType
# 构造测试数据
df = spark.createDataFrame(
[
(list(range(300)), 4),
(list(range(300)), 200),
(list(range(300)), 300),
(list(range(300)), 800)
],
schema=StructType([
StructField("array", ArrayType(IntegerType())),
StructField("n_relevant", IntegerType())
])
)
# 核心逻辑：分段应用 slice + filter
df_result = df.withColumn(
"result",
F.when(
F.col("n_relevant") <= 100,
F.slice("array", 1, F.col("n_relevant")) # 直接截取前n个
).when(
(F.col("n_relevant") > 100) & (F.col("n_relevant") < 300),
F.filter(
F.slice("array", 1, F.col("n_relevant")),
lambda _, idx: idx % 2 == 0 # 索引0,2,4... → 偶数位（含第0个）
)
).otherwise(
F.filter(
F.slice("array", 1, 300), # n_relevant≥300时，只处理原始300个
lambda _, idx: idx % 3 == 0 # 索引0,3,6... → 每3个取1个
)
)
)
display(df_result.select("array", "n_relevant", "result"))✅ 注意事项：
-
filter的 lambda 函数中，第一个参数是元素值（此处用_忽略），第二个参数idx是 0-based 索引，与 Python 列表行为一致； -
slice(..., 1, n)中n可为列引用（如F.col("n_relevant")），支持动态长度； - 条件判断使用
&（而非and）连接布尔表达式，因 Spark SQL 中and不支持列运算； - 当
n_relevant ≥ 300时，slice("array", 1, 300)显式限定范围，避免无效扩展（虽原数组恒为300，但显式声明更健壮）； - 此方案完全基于 Catalyst 优化器，无序列化开销，性能远优于 Pandas UDF 或
mapPartitions。
该模式可轻松扩展：只需调整 when 分支及对应模数，即可适配任意分段采样规则，是生产环境中处理数组约束的推荐实践。
