Spark中对同一DataFrame并行执行独立随机选词操作的实现方法

没问题，我来帮你搞定这个Spark抽样的需求！咱们先理清楚核心要求：从带概率权重的单词DataFrame里，独立生成X组单词对，每组两个不同的单词，而且抽样要跟着prob列的权重走（比如你示例里green的概率最高，出现次数最多）。下面给你两种可行的方案，分别适合不同的数据规模场景：

方案一：小数据量场景（单词数量不多） #

如果你的单词总数不多（比如几百几千个），直接把单词和权重收集到本地处理会更高效，代码也更简洁：

Python 代码示例 #

from pyspark.sql import functions as F
from pyspark.sql.types import StructType, StructField, StringType
import random

# 假设你的原DataFrame叫df
df = spark.createDataFrame(
    [(0.1,"blue"),(0.2,"yellow"),(0.1,"red"),(0.6,"green")],
    ["prob", "word"]
)

# 收集单词和权重到本地列表
word_prob_pairs = df.select("word", "prob").collect()
words = [row.word for row in word_prob_pairs]
probs = [row.prob for row in word_prob_pairs]

# 广播变量，避免每个Executor重复加载数据
broadcast_data = spark.sparkContext.broadcast((words, probs))

# 定义生成单词对的UDF
def generate_unique_pair():
    words_list, probs_list = broadcast_data.value
    # 第一次按权重抽样
    first_word = random.choices(words_list, weights=probs_list, k=1)[0]
    # 过滤掉已选中的单词，调整权重（归一化）
    filtered_items = [(w, p) for w, p in zip(words_list, probs_list) if w != first_word]
    filtered_words = [item[0] for item in filtered_items]
    filtered_probs = [item[1] for item in filtered_items]
    total_prob = sum(filtered_probs)
    adjusted_probs = [p / total_prob for p in filtered_probs]
    # 第二次抽样
    second_word = random.choices(filtered_words, weights=adjusted_probs, k=1)[0]
    # 如果需要无序对（blue,green和green,blue算同一组）就排序，否则直接返回元组
    return tuple(sorted((first_word, second_word)))

# 注册UDF
pair_udf = F.udf(generate_unique_pair, StructType([
    StructField("word1", StringType()),
    StructField("word2", StringType())
]))

# 生成X组结果
X = 5
result_df = spark.range(X).withColumn("pair", pair_udf()).select("pair.*")
result_df.show()

Scala 代码示例 #

import org.apache.spark.sql.functions._
import org.apache.spark.sql.types._
import scala.util.Random
import scala.util.control.Breaks._

// 原DataFrame
val df = spark.createDataFrame(Seq(
  (0.1, "blue"),
  (0.2, "yellow"),
  (0.1, "red"),
  (0.6, "green")
)).toDF("prob", "word")

// 收集单词和权重
val wordProbList = df.select("word", "prob").collect()
val words = wordProbList.map(_.getString(0))
val probs = wordProbList.map(_.getDouble(1))

// 广播变量
val broadcastData = spark.sparkContext.broadcast((words, probs))

// 定义UDF
val generatePair = udf(() => {
  val (wordsArr, probsArr) = broadcastData.value
  // 第一次加权抽样
  val rand1 = Random.nextDouble()
  var cumProb = 0.0
  var firstWord = ""
  breakable {
    for ((w, p) <- wordsArr.zip(probsArr)) {
      cumProb += p
      if (rand1 <= cumProb) {
        firstWord = w
        break
      }
    }
  }
  // 过滤并调整权重
  val filtered = wordsArr.zip(probsArr).filter(_._1 != firstWord)
  val total = filtered.map(_._2).sum
  val adjustedProbs = filtered.map(_._2 / total)
  // 第二次抽样
  val rand2 = Random.nextDouble()
  cumProb = 0.0
  var secondWord = ""
  breakable {
    for (((w, _), p) <- filtered.zip(adjustedProbs)) {
      cumProb += p
      if (rand2 <= cumProb) {
        secondWord = w
        break
      }
    }
  }
  // 返回无序对（如需有序则去掉.sorted）
  (firstWord, secondWord).sorted
})

// 生成X组结果
val X = 5
val resultDF = spark.range(X).withColumn("pair", generatePair())
  .select($"pair._1".alias("word1"), $"pair._2".alias("word2"))
resultDF.show()

方案二：大数据量分布式场景（单词数量极多） #

如果你的单词数量非常大（比如百万级），不能收集到本地处理，那就用纯分布式的Spark操作来实现：

Python 代码示例 #

from pyspark.sql import functions as F
from pyspark.sql.window import Window

df = spark.createDataFrame(
    [(0.1,"blue"),(0.2,"yellow"),(0.1,"red"),(0.6,"green")],
    ["prob", "word"]
)

# 第一步：计算累积概率，用于加权抽样
cum_window = Window.orderBy("word")
df_with_cum = df.withColumn("cum_prob", F.sum("prob").over(cum_window))

# 第二步：生成X行的基础DataFrame，带两个随机数（分别用于两次抽样）
X = 5
random_df = spark.range(X).withColumn("rand1", F.rand()).withColumn("rand2", F.rand())

# 第三步：第一次抽样，找到每个rand1对应的单词
def sample_single_word(base_df, rand_col, alias_name):
    # 交叉连接后，找到第一个累积概率大于等于随机数的单词
    cross_df = base_df.crossJoin(df_with_cum)
    rank_window = Window.partitionBy("id").orderBy(F.abs(F.col("cum_prob") - F.col(rand_col)))
    return cross_df.withColumn("rank", F.rank().over(rank_window))\
                   .filter(F.col("rank") == 1)\
                   .select("id", F.col("word").alias(alias_name))

sampled_first = sample_single_word(random_df, "rand1", "word1")

# 第四步：第二次抽样，排除第一次抽到的单词
# 关联已抽样结果，过滤掉相同单词
filtered_df = sampled_first.join(df_with_cum, sampled_first["word1"] != df_with_cum["word"], "left")
# 计算过滤后的累积概率
filtered_cum_window = Window.partitionBy("id").orderBy("word")
filtered_df_with_cum = filtered_df.withColumn("filtered_cum", F.sum("prob").over(filtered_cum_window))
# 第二次抽样
cross_filtered = filtered_df_with_cum.join(random_df.select("id", "rand2"), on="id")
rank_window2 = Window.partitionBy("id").orderBy(F.abs(F.col("filtered_cum") - F.col("rand2")))
sampled_second = cross_filtered.withColumn("rank", F.rank().over(rank_window2))\
                               .filter(F.col("rank") == 1)\
                               .select("id", F.col("word").alias("word2"))

# 第五步：合并结果
final_result = sampled_first.join(sampled_second, on="id").orderBy("id")
final_result.show()

注意事项 #

• 如果原DataFrame的prob列总和不是1，需要先做归一化处理（比如除以sum(prob)），否则加权抽样的结果会有偏差。
• 如果不需要加权抽样（均匀随机选），只需要去掉代码中weights相关的参数即可。
• 如果认为blue,green和green,blue是不同的组，去掉代码中sorted的逻辑即可。

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