在PySpark DataFrame中用withColumn按分区统计唯一user_id数量

步骤说明 #

根据您提供的期望输出，un_count列实际为每个user_id在对应(plan, name)分组内的出现次数。我们可以使用PySpark的窗口函数实现需求：按plan、name和user_id分区，统计每个用户在分组内的出现次数，将结果作为新列添加到原DataFrame中。

如果您实际需要的是每个(plan, name)分组内唯一user_id的总数（即分组内不同用户的数量），可以参考文末的补充说明。

完整代码 #

from pyspark.sql import SparkSession
from pyspark.sql import Window
from pyspark.sql.functions import count

# 初始化SparkSession
spark = SparkSession.builder.appName("user_count_example").getOrCreate()

# 构造示例输入DataFrame
data = [
    ("FWE", "QRS", 123, 10),
    ("FWE", "QRS", 123, 20),
    ("FWE", "QRS", 223, 30),
    ("RWE", "QRS", 223, 40)
]
df = spark.createDataFrame(data, ["plan", "name", "user_id", "dollars"])

# 定义窗口：按plan、name、user_id分区，统计当前用户在分组内的出现次数
window_spec = Window.partitionBy("plan", "name", "user_id")

# 添加un_count列
result_df = df.withColumn("un_count", count("user_id").over(window_spec))

# 查看结果
result_df.show()

输出结果 #

执行代码后会得到与您期望一致的DataFrame：

补充：如果需要分组内唯一user_id的总数 #

如果您的真实需求是统计每个(plan, name)分组内的唯一用户数量（即分组内不同user_id的个数），只需调整窗口和统计函数：

from pyspark.sql.functions import countDistinct

# 窗口改为仅按plan、name分区
window_spec_total = Window.partitionBy("plan", "name")
result_total_df = df.withColumn("un_count", countDistinct("user_id").over(window_spec_total))
result_total_df.show()

此时输出的un_count列会是每个分组的唯一用户总数，比如FWE-QRS分组的所有行un_count都为2。

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