如何在PySpark DataFrame中生成饮品与食物的交叉比例表

数据准备 #

首先创建示例DataFrame：

from pyspark.sql import SparkSession
from pyspark.sql.functions import col, max, sum, when

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

# 原始订单数据
data = [
    ("A", "Coffee"), ("A", "Tea"), ("A", "Burger"), ("A", "Fried Chicken"),
    ("B", "Coffee"), ("B", "Tea"), ("B", "Hot Dog"), ("B", "Fried Chicken"),
    ("C", "Coffee"), ("C", "Hot Dog"), ("C", "Fried Chicken"),
    ("D", "Tea"), ("D", "Burger")
]

df = spark.createDataFrame(data, ["Name", "Order"])

定义食物与饮品列表 #

Food = ['Hot Dog', 'Burger', 'Fried Chicken']
Drink = ['Coffee', 'Tea']

步骤1：生成用户购买标记宽表 #

为每个用户标记是否购买了目标食物和饮品，再按用户聚合得到最终购买状态：

# 标记用户是否购买每种食物
for food in Food:
    col_name = f"has_{food.replace(' ', '_')}"
    df = df.withColumn(col_name, when(col("Order") == food, 1).otherwise(0))

# 标记用户是否购买每种饮品
for drink in Drink:
    col_name = f"has_{drink.replace(' ', '_')}"
    df = df.withColumn(col_name, when(col("Order") == drink, 1).otherwise(0))

# 按用户分组，聚合得到每个用户的最终购买状态（购买过则记为1）
user_purchases = df.groupBy("Name").agg(
    *[max(f"has_{food.replace(' ', '_')}").alias(f"has_{food.replace(' ', '_')}") for food in Food],
    *[max(f"has_{drink.replace(' ', '_')}").alias(f"has_{drink.replace(' ', '_')}") for drink in Drink]
)

步骤2：计算每种食物的总购买用户数 #

统计购买过每种食物的用户总量：

food_total_users = user_purchases.agg(
    *[sum(f"has_{food.replace(' ', '_')}").alias(f"{food}_total") for food in Food]
).collect()[0].asDict()

步骤3：计算同时购买食物与饮品的用户数 #

通过用户购买状态的乘积求和，得到同时购买两者的用户数：

# 创建临时视图用于SQL查询
user_purchases.createOrReplaceTempView("user_purchases")

# 构建查询语句，计算每个食物-饮品组合的共同购买用户数
query_items = []
for drink in Drink:
    drink_col = f"has_{drink.replace(' ', '_')}"
    for food in Food:
        food_col = f"has_{food.replace(' ', '_')}"
        query_items.append(f"sum({food_col} * {drink_col}) as {drink}_{food.replace(' ', '_')}")

joint_counts = spark.sql(f"select {', '.join(query_items)} from user_purchases").collect()[0].asDict()

步骤4：计算比例并整理成目标格式 #

用共同购买用户数除以食物总购买用户数，得到比例并转换为目标表格：

# 计算比例并构建结果字典
ratio_result = {}
for drink in Drink:
    ratio_result[drink] = {}
    for food in Food:
        total = food_total_users[f"{food}_total"]
        joint = joint_counts[f"{drink}_{food.replace(' ', '_')}"]
        ratio = round(joint / total, 2) if total != 0 else 0.0
        ratio_result[drink][food] = ratio

# 转换为DataFrame展示
result_df = spark.createDataFrame(
    [(drink, *[ratio_result[drink][food] for food in Food]) for drink in Drink],
    ["Drink"] + Food
)

result_df.show()

最终输出 #

运行代码后将得到符合要求的交叉比例表：

+-------+--------+------+---------------+
|  Drink|Hot Dog|Burger|Fried Chicken|
+-------+--------+------+---------------+
| Coffee|     1.0|   0.5|            1.0|
|    Tea|     0.5|   1.0|           0.67|
+-------+--------+------+---------------+

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