PySpark：基于cc值1/2数量，从两列表列按同索引随机选记录

核心思路 #

要实现同索引随机选取n条记录，关键是先把两个列表的对应索引元素绑定，再整体随机打乱后截取指定数量，最后拆分回原列结构。n值由cc列值的1/2取整得到（如cc=4则n=2，cc=9则n=4），同时处理列表长度不足n的边界情况。

代码实现 #

假设你的PySpark DataFrame名为df，执行以下代码：

from pyspark.sql import functions as f

# 1. 计算每个行需要选取的记录数n，避免n超过列表实际长度
df_with_n = df.withColumn(
    "n",
    f.least(
        f.floor(f.col("cc") / 2).cast("int"),  # 取cc/2的整数部分
        f.size(f.col("doc_1"))  # 以列表长度为上限
    )
)

# 2. 绑定索引与对应元素、打乱、截取、拆分
result_df = df_with_n.withColumn(
    # 将索引、doc_1、doc_2的对应元素打包成结构体数组
    "paired_elements",
    f.arrays_zip(
        f.sequence(f.lit(0), f.size(f.col("doc_1")) - 1).alias("idx"),
        f.col("doc_1").alias("d1"),
        f.col("doc_2").alias("d2")
    )
).withColumn(
    # 打乱数组后截取前n个元素
    "sampled_pairs",
    f.slice(f.shuffle(f.col("paired_elements")), 1, f.col("n"))
).select(
    "cc",
    # 从抽样结果中提取doc_1的元素组成新列表
    f.transform(f.col("sampled_pairs"), lambda x: x["d1"]).alias("doc_1"),
    # 从抽样结果中提取doc_2的元素组成新列表
    f.transform(f.col("sampled_pairs"), lambda x: x["d2"]).alias("doc_2")
)

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

关键函数说明 #

• arrays_zip：将索引数组、doc_1、doc_2的对应位置元素绑定，确保后续抽样时索引完全匹配。
• shuffle：对绑定后的数组整体随机打乱，实现无偏随机抽样。
• slice：从打乱后的数组中截取前n个元素，精准控制选取数量。
• transform：从抽样后的结构体数组中拆分出doc_1和doc_2的元素，还原目标列结构。

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