PySpark如何按用户分组并在正负采样率下采样

from pyspark.sql import SparkSession
from pyspark.sql.functions import col

spark = SparkSession.builder.appName("UserSampling").getOrCreate()

data = spark.read.csv("data.csv", header=True, inferSchema=True)

user_counts = data.groupBy("user_id").count()

positive_rate = 0.2  # 正样本采样率
negative_rate = 0.1  # 负样本采样率

sample_rates = user_counts.withColumn("sample_rate", 
                                      col("count") * (positive_rate + negative_rate))

sampled_data = data.join(sample_rates, on="user_id", how="inner") \
                   .where(col("sample_rate") >= 1.0 or 
                          (col("sample_rate") < 1.0 and col("sample_rate") >= rand()))

sampled_data.show()