如何用SparkSQL或PySpark计算各时间点全位置最近状态总和？

问题 #

现有如下结构的Spark DataFrame：

+-------+---------+-----+-------------------+
|product| location|  num|                 ts|
+-------+---------+-----+-------------------+
|      1|        A|    3|2024-01-31 04:28:27|
|      1|        B|   12|2024-01-31 04:28:27|
|      1|        C|   19|2024-01-31 04:28:27|
|      1|        D|    1|2024-01-31 04:28:27|
|      1|        E|    4|2024-01-31 04:28:27|
|      1|        D|    2|2024-02-02 04:30:08|
|      1|        A|    4|2024-02-02 04:30:08|
|      1|        C|   20|2024-02-02 04:30:08|
|      1|        E|    1|2024-02-02 04:30:08|
|      1|        F|   20|2024-02-02 04:30:08|
|      1|        E|    2|2024-02-03 04:32:39|
|      1|        D|    5|2024-02-03 04:32:39|
|      1|        A|    3|2024-02-03 04:32:39|
+-------+---------+-----+-------------------+

需要计算每个时间点下，所有位置的最近状态的产品总数。目前仅能计算同timestamp的总和，无法包含最近前置时间点的其他位置数据。

期望结果如下：

+-------+---------+-----+-------------------+-------+
|product| location|  num|                 ts|  total|
+-------+---------+-----+-------------------+-------+
|      1|        A|    3|2024-01-31 04:28:27|     39| --> A(3) + B(12) + C(19) + D(1) + E(4)
|      1|        B|   12|2024-01-31 04:28:27|     39| --> A(3) + B(12) + C(19) + D(1) + E(4)
|      1|        C|   19|2024-01-31 04:28:27|     39| --> A(3) + B(12) + C(19) + D(1) + E(4)
|      1|        D|    1|2024-01-31 04:28:27|     39| --> A(3) + B(12) + C(19) + D(1) + E(4)
|      1|        E|    4|2024-01-31 04:28:27|     39| --> A(3) + B(12) + C(19) + D(1) + E(4)
|      1|        D|    2|2024-02-02 04:30:08|     59| --> A(4) + B(12) + C(20) + D(2) + E(1) + F(20)
|      1|        A|    4|2024-02-02 04:30:08|     59| --> A(4) + B(12) + C(20) + D(2) + E(1) + F(20)
|      1|        C|   20|2024-02-02 04:30:08|     59| --> A(4) + B(12) + C(20) + D(2) + E(1) + F(20)
|      1|        E|    1|2024-02-02 04:30:08|     59| --> A(4) + B(12) + C(20) + D(2) + E(1) + F(20)
|      1|        F|   20|2024-02-02 04:30:08|     59| --> A(4) + B(12) + C(20) + D(2) + E(1) + F(20)
|      1|        E|    2|2024-02-03 04:32:39|     62| --> A(3) + B(12) + C(20) + D(5) + E(2) + F(20)
|      1|        D|    5|2024-02-03 04:32:39|     62| --> A(3) + B(12) + C(20) + D(5) + E(2) + F(20)
|      1|        A|    3|2024-02-03 04:32:39|     62| --> A(3) + B(12) + C(20) + D(5) + E(2) + F(20)
+-------+---------+-----+-------------------+-------+

能否用SparkSQL或PySpark实现该需求？

解决方案 #

方法一：SparkSQL实现 #

核心逻辑是先补全每个时间点下的所有位置组合，再找到每个位置在对应时间点的最近记录，最后按时间点求和并关联回原表：

-- 获取所有唯一时间点
WITH unique_ts AS (
    SELECT DISTINCT ts FROM your_table
),
-- 获取所有唯一的产品-位置组合
unique_locations AS (
    SELECT DISTINCT product, location FROM your_table
),
-- 交叉关联得到每个时间点的全量产品-位置组合
all_combinations AS (
    SELECT u.product, u.location, t.ts
    FROM unique_locations u
    CROSS JOIN unique_ts t
),
-- 找到每个组合下不晚于当前时间点的最近记录
latest_records AS (
    SELECT 
        ac.product,
        ac.location,
        ac.ts,
        FIRST_VALUE(y.num) OVER (
            PARTITION BY ac.product, ac.location 
            ORDER BY y.ts DESC 
            ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW
        ) AS latest_num
    FROM all_combinations ac
    LEFT JOIN your_table y 
        ON ac.product = y.product 
        AND ac.location = y.location 
        AND y.ts <= ac.ts
),
-- 按时间点计算总数量
total_per_ts AS (
    SELECT product, ts, SUM(latest_num) AS total
    FROM latest_records
    GROUP BY product, ts
)
-- 关联回原表得到最终结果
SELECT 
    y.product,
    y.location,
    y.num,
    y.ts,
    t.total
FROM your_table y
JOIN total_per_ts t 
    ON y.product = t.product 
    AND y.ts = t.ts
ORDER BY y.ts, y.location;

方法二：PySpark DataFrame API实现 #

用DataFrame API复刻上述逻辑，步骤清晰直观：

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

# 初始化示例DataFrame（实际使用时替换为你的数据源）
df = spark.createDataFrame(
    [
        (1, "A", 3, "2024-01-31 04:28:27"),
        (1, "B", 12, "2024-01-31 04:28:27"),
        (1, "C", 19, "2024-01-31 04:28:27"),
        (1, "D", 1, "2024-01-31 04:28:27"),
        (1, "E", 4, "2024-01-31 04:28:27"),
        (1, "D", 2, "2024-02-02 04:30:08"),
        (1, "A", 4, "2024-02-02 04:30:08"),
        (1, "C", 20, "2024-02-02 04:30:08"),
        (1, "E", 1, "2024-02-02 04:30:08"),
        (1, "F", 20, "2024-02-02 04:30:08"),
        (1, "E", 2, "2024-02-03 04:32:39"),
        (1, "D", 5, "2024-02-03 04:32:39"),
        (1, "A", 3, "2024-02-03 04:32:39"),
    ],
    ["product", "location", "num", "ts"]
)
# 转换时间字符串为时间类型
df = df.withColumn("ts", F.to_timestamp("ts"))

# 1. 获取唯一时间点和产品-位置组合
unique_ts = df.select("ts").distinct()
unique_locations = df.select("product", "location").distinct()

# 2. 交叉关联生成全量组合
all_combinations = unique_locations.crossJoin(unique_ts)

# 3. 关联原始数据，用窗口函数取每个位置的最近记录
window_spec = Window.partitionBy("product", "location").orderBy(F.desc("ts"))
latest_records = all_combinations.join(
    df,
    (all_combinations.product == df.product) &
    (all_combinations.location == df.location) &
    (df.ts <= all_combinations.ts),
    "left"
).withColumn(
    "latest_num",
    F.first("num", ignorenulls=True).over(window_spec)
).select(
    all_combinations.product,
    all_combinations.location,
    all_combinations.ts,
    "latest_num"
)

# 4. 按时间点计算总数量
total_per_ts = latest_records.groupBy("product", "ts").agg(F.sum("latest_num").alias("total"))

# 5. 关联回原表得到最终结果
final_df = df.join(total_per_ts, on=["product", "ts"], how="inner")
final_df.orderBy("ts", "location").show(truncate=False)

内容的提问来源于stack exchange，提问作者Praiwan N.