如何用现有值组合与指定日期列表补全PySpark DataFrame缺失行

原始PySpark DataFrame：

Product    Location    Date        Stock
Apple      L1          2023-01-01  100
Apple      L2          2023-01-15  30
Banana     L1          2023-01-01  10
Banana     L1          2023-01-08  20
Orange     L2          2023-01-01  50

给定日期列表：

dates = ["2023-01-01", "2023-01-08", "2023-01-15"]

期望输出DataFrame：

Product    Location    Date        Stock
Apple      L1          2023-01-01  100
Apple      L1          2023-01-08  0
Apple      L1          2023-01-15  0
Apple      L2          2023-01-15  30
Apple      L2          2023-01-01  0
Apple      L2          2023-01-08  0
Banana     L1          2023-01-01  10
Banana     L1          2023-01-08  20
Banana     L1          2023-01-15  0
Orange     L2          2023-01-01  50
Orange     L2          2023-01-08  0
Orange     L2          2023-01-15  0

注意事项：

• 不添加原本不存在的Product和Location组合（如(Banana, L2)或(Orange, L1)）
• 直接使用给定的dates列表，无需考虑日期属性

不需要窗口函数，通过生成全量组合+左连接的方式即可实现，步骤如下：

1. 提取原始DataFrame中所有存在的(Product, Location)唯一组合
2. 将这些组合与给定的dates列表做笛卡尔积，得到每个组合对应的所有日期行
3. 将全量组合表与原始DataFrame左连接，缺失的Stock值填充为0

具体代码实现：

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

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

# 构造原始DataFrame
data = [
    ("Apple", "L1", "2023-01-01", 100),
    ("Apple", "L2", "2023-01-15", 30),
    ("Banana", "L1", "2023-01-01", 10),
    ("Banana", "L1", "2023-01-08", 20),
    ("Orange", "L2", "2023-01-01", 50)
]
df = spark.createDataFrame(data, ["Product", "Location", "Date", "Stock"])

# 构造日期列表对应的DataFrame
dates = ["2023-01-01", "2023-01-08", "2023-01-15"]
dates_df = spark.createDataFrame([(d,) for d in dates], ["Date"])

# 提取唯一的(Product, Location)组合
product_loc_unique = df.select("Product", "Location").distinct()

# 生成全量组合（笛卡尔积）
full_combinations = product_loc_unique.crossJoin(dates_df)

# 左连接原始数据，填充缺失的Stock为0
result_df = full_combinations.join(df, on=["Product", "Location", "Date"], how="left") \
    .withColumn("Stock", col("Stock").cast("int")) \
    .fillna(0, subset=["Stock"])

# 按指定顺序展示结果
result_df.orderBy("Product", "Location", "Date").show()

代码说明： #

• distinct()：确保只保留原始数据中存在的(Product, Location)组合，不会生成不存在的无效组合
• crossJoin()：实现组合与日期的全量匹配，得到每个组合对应的所有日期行
• left join：保留全量组合的所有行，匹配原始数据中存在的Stock值，缺失值用fillna(0)填充为0
• orderBy()：让结果按Product、Location、Date排序，与示例输出格式一致

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