本文介绍如何在 PySpark 中基于另一列（n_relevant）的值，对固定长度数组执行条件化子集操作：≤100时全量截取，100–299时按动态步长模采样，≥300时统一按模3采样，全程纯 SQL 函数实现，无需 UDF。
本文介绍如何在 pyspark 中基于另一列（`n_relevant`）的值，对固定长度数组列执行条件化子集操作：≤100时全量截取，100–299时按动态步长模采样，≥300时统一按模3采样，全程纯 sql 函数实现，无需 udf。
在大规模数据处理中，常需对数组列施加长度约束（如目标系统要求数组 ≤100 元素），而实际业务逻辑又依赖于动态参数（如 n_relevant）决定“有效范围”。直接使用 transform 或 array_remove 难以实现带索引条件的动态采样——因为 PySpark 的 lambda 表达式不支持传入外部变量，也无法在运行时计算模数。正确解法是组合 slice + filter + when/otherwise，利用 filter 的双参数签名（lambda element, index: ...）在索引维度做条件判断。
以下为完整可运行示例（适配 Databricks 环境）：
from pyspark.sql import functions as F
from pyspark.sql.types import StructType, StructField, IntegerType, ArrayType
# 构造测试数据：array 固定为 0~299，n_relevant 取四种典型值
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())
])
)
# 核心逻辑：分段策略 + 索引过滤
df_result = df.withColumn(
"result",
# 情况1：n_relevant <= 100 → 直接 slice 前 n_relevant 个元素
F.when(F.col("n_relevant") <= 100,
F.slice("array", 1, F.col("n_relevant")))
# 情况2：100 < n_relevant <= 299 → 截取前 n_relevant 个，再按 index % 2 == 0 采样（即每2个取1个，得约100个）
.when((F.col("n_relevant") > 100) & (F.col("n_relevant") <= 299),
F.filter(
F.slice("array", 1, F.col("n_relevant")),
lambda _, idx: idx % 2 == 0 # 注意：PySpark 中索引从 0 开始，但 slice 起始为1 → filter 内 idx 仍从0计
))
# 情况3：n_relevant >= 300 → 截取前 min(n_relevant, 300) 个（实际恒为300），再按 index % 3 == 0 采样（得100个：0,3,6,...297）
.otherwise(
F.filter(
F.slice("array", 1, F.lit(300)), # 显式限定上限，避免无效扩展
lambda _, idx: idx % 3 == 0
)
)
)
display(df_result.select("array", "n_relevant", "result"))✅ 关键要点说明：
-
F.slice(array, start, length)中start=1（1-indexed），length可为列引用（如F.col("n_relevant")），实现动态截断； -
F.filter(array, lambda elem, idx: condition)的idx是 0-based 索引，与 Python 列表一致，可直接用于模运算； - 条件分支需覆盖全部区间：
、<code>(100,299]、>=300，避免NULL输出； - 对
n_relevant ≥ 300场景，slice长度设为F.lit(300)更安全（因原始数组仅300元素），防止越界； - 性能提示：全程使用内置函数，无 UDF 开销，可被 Catalyst 优化器充分下推，适合 TB 级数据。
最终输出中，result 列将严格满足：长度 ≤100，且采样逻辑与需求完全一致（如 n_relevant=200 时返回索引为偶数的前200个元素，共100个）。此方案兼顾表达力、性能与可维护性，是 PySpark 数组条件处理的推荐实践。
