在Databricks中不使用UDF聚合DataFrame列信息

原始DataFrame结构如下：

需要为每个客户生成commodity_aggr列，规则如下：

• 条件1：客户仅有1行数据且商品唯一值数量为1，标记为MONO，开始日期取该行开始日期，结束日期取该行结束日期
• 条件2：客户仅有1行数据且商品唯一值数量>1，标记为DUAL，开始日期取该行开始日期，结束日期取该行结束日期
• 条件3：客户有2行数据且商品唯一值数量>1，且第一行开始日期≤第二行开始日期：第一行标记为MONO（开始日期取第一行开始日期，结束日期取第二行开始日期）；第二行标记为DUAL（开始日期取第二行开始日期，结束日期取两行最大结束日期）
• 条件4：客户有2行数据且商品唯一值数量>1，且第一行结束日期≤第二行结束日期：第一行标记为DUAL（开始日期取两行最小开始日期，结束日期取第一行结束日期）；第二行标记为MONO（开始日期取第一行结束日期，结束日期取第二行结束日期）

期望输出：

环境说明：Unity Catalog下的Databricks集群，Runtime版本13.1，不支持UDF，可使用PySpark或SQL实现。

以下是基于PySpark的实现代码，全程使用内置函数完成，无需自定义UDF：

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

# 创建原始DataFrame（Databricks环境可直接读取现有表）
data = [
    ("A", "GAS", "01/03/2023", "19/02/2024"),
    ("A", "POWER", "01/03/2023", "19/02/2024"),
    ("B", "GAS", "01/02/2023", "19/02/2024"),
    ("B", "POWER", "01/03/2023", "19/02/2024"),
    ("C", "GAS", "01/03/2023", "19/12/2023"),
    ("C", "POWER", "01/03/2023", "19/02/2024"),
    ("D", "POWER", "01/03/2023", "19/02/2024")
]
df = spark.createDataFrame(data, ["customer", "commodity", "start_date", "end_date"])

# 转换日期字符串为日期类型，确保比较准确性
df = df.withColumn("start_date", F.to_date("start_date", "dd/MM/yyyy")) \
       .withColumn("end_date", F.to_date("end_date", "dd/MM/yyyy"))

# 定义客户分组窗口
customer_window = Window.partitionBy("customer")

# 计算每个客户的聚合指标
df = df.withColumn("row_count", F.count("*").over(customer_window)) \
       .withColumn("commodity_distinct_count", F.countDistinct("commodity").over(customer_window)) \
       .withColumn("other_starts", F.collect_list("start_date").over(customer_window)) \
       .withColumn("other_ends", F.collect_list("end_date").over(customer_window)) \
       .withColumn("max_end", F.max("end_date").over(customer_window)) \
       .withColumn("min_start", F.min("start_date").over(customer_window))

# 处理单条记录的情况
single_row_df = df.filter(F.col("row_count") == 1).withColumn(
    "commodity_aggr",
    F.when(F.col("commodity_distinct_count") == 1, "MONO").otherwise("DUAL")
).select("customer", "commodity_aggr", "start_date", "end_date")

# 处理两条记录的情况
two_rows_df = df.filter(F.col("row_count") == 2)

# 拆分同行的日期到单独列
two_rows_df = two_rows_df.withColumn("start1", F.element_at(F.col("other_starts"), 1)) \
                         .withColumn("start2", F.element_at(F.col("other_starts"), 2)) \
                         .withColumn("end1", F.element_at(F.col("other_ends"), 1)) \
                         .withColumn("end2", F.element_at(F.col("other_ends"), 2))

# 应用条件生成标记和新日期
two_rows_df = two_rows_df.withColumn(
    "commodity_aggr",
    F.when(
        ((F.col("start_date") == F.col("start1")) & (F.col("start1") <= F.col("start2"))) |
        ((F.col("start_date") == F.col("start2")) & (F.col("start1") <= F.col("start2"))),
        F.when(F.col("start_date") == F.col("start1"), "MONO").otherwise("DUAL")
    ).when(
        ((F.col("end_date") == F.col("end1")) & (F.col("end1") <= F.col("end2"))) |
        ((F.col("end_date") == F.col("end2")) & (F.col("end1") <= F.col("end2"))),
        F.when(F.col("end_date") == F.col("end1"), "DUAL").otherwise("MONO")
    )
).withColumn(
    "new_start",
    F.when(
        (F.col("commodity_aggr") == "MONO") & (F.col("start1") <= F.col("start2")),
        F.col("start_date")
    ).when(
        (F.col("commodity_aggr") == "DUAL") & (F.col("start1") <= F.col("start2")),
        F.col("start2")
    ).when(
        (F.col("commodity_aggr") == "DUAL") & (F.col("end1") <= F.col("end2")),
        F.col("min_start")
    ).when(
        (F.col("commodity_aggr") == "MONO") & (F.col("end1") <= F.col("end2")),
        F.col("end1")
    )
).withColumn(
    "new_end",
    F.when(
        (F.col("commodity_aggr") == "MONO") & (F.col("start1") <= F.col("start2")),
        F.col("start2")
    ).when(
        (F.col("commodity_aggr") == "DUAL") & (F.col("start1") <= F.col("start2")),
        F.col("max_end")
    ).when(
        (F.col("commodity_aggr") == "DUAL") & (F.col("end1") <= F.col("end2")),
        F.col("end_date")
    ).when(
        (F.col("commodity_aggr") == "MONO") & (F.col("end1") <= F.col("end2")),
        F.col("end_date")
    )
).select("customer", "commodity_aggr", F.col("new_start").alias("start_date"), F.col("new_end").alias("end_date"))

# 合并结果并恢复日期格式
final_df = single_row_df.union(two_rows_df) \
                        .withColumn("start_date", F.date_format("start_date", "dd/MM/yyyy")) \
                        .withColumn("end_date", F.date_format("end_date", "dd/MM/yyyy"))

# 展示结果
final_df.orderBy("customer").show()

代码说明 #

1. 日期转换：先将字符串日期转为日期类型确保比较准确，最后再转回原始字符串格式。
2. 窗口函数计算：通过窗口函数获取每个客户的行数、商品唯一值数量，以及同行的日期集合、最大结束日期、最小开始日期。
3. 分场景处理：将单条记录和两条记录的逻辑分开，避免混乱：
   • 单条记录直接根据商品唯一值数量标记MONO/DUAL。
   • 两条记录拆分出另一行的日期，分别应用条件3和条件4生成对应标记与起止日期。
4. 结果合并：合并两种场景的结果，恢复日期格式后输出。

如果偏好使用SQL，可在Databricks Notebook中执行以下代码：

-- 创建临时视图（可替换为现有表名）
CREATE OR REPLACE TEMP VIEW raw_data AS
SELECT * FROM VALUES
    ("A", "GAS", "01/03/2023", "19/02/2024"),
    ("A", "POWER", "01/03/2023", "19/02/2024"),
    ("B", "GAS", "01/02/2023", "19/02/2024"),
    ("B", "POWER", "01/03/2023", "19/02/2024"),
    ("C", "GAS", "01/03/2023", "19/12/2023"),
    ("C", "POWER", "01/03/2023", "19/02/2024"),
    ("D", "POWER", "01/03/2023", "19/02/2024")
AS t(customer, commodity, start_date, end_date);

-- 分步处理数据
WITH transformed_data AS (
    SELECT
        customer,
        commodity,
        TO_DATE(start_date, 'dd/MM/yyyy') AS start_date,
        TO_DATE(end_date, 'dd/MM/yyyy') AS end_date,
        COUNT(*) OVER (PARTITION BY customer) AS row_count,
        COUNT(DISTINCT commodity) OVER (PARTITION BY customer) AS commodity_distinct_count,
        COLLECT_LIST(start_date) OVER (PARTITION BY customer) AS other_starts,
        COLLECT_LIST(end_date) OVER (PARTITION BY customer) AS other_ends,
        MAX(TO_DATE(end_date, 'dd/MM/yyyy')) OVER (PARTITION BY customer) AS max_end,
        MIN(TO_DATE(start_date, 'dd/MM/yyyy')) OVER (PARTITION BY customer) AS min_start
    FROM raw_data
),
single_row AS (
    SELECT
        customer,
        CASE WHEN commodity_distinct_count = 1 THEN 'MONO' ELSE 'DUAL' END AS commodity_aggr,
        DATE_FORMAT(start_date, 'dd/MM/yyyy') AS start_date,
        DATE_FORMAT(end_date, 'dd/MM/yyyy') AS end_date
    FROM transformed_data
    WHERE row_count = 1
),
two_rows AS (
    SELECT
        customer,
        CASE
            WHEN (start_date = element_at(other_starts, 1) AND element_at(other_starts, 1) <= element_at(other_starts, 2)) OR
                 (start_date = element_at(other_starts, 2) AND element_at(other_starts, 1) <= element_at(other_starts, 2)) THEN
                CASE WHEN start_date = element_at(other_starts, 1) THEN 'MONO' ELSE 'DUAL' END
            WHEN (end_date = element_at(other_ends, 1) AND element_at(other_ends, 1) <= element_at(other_ends, 2)) OR
                 (end_date = element_at(other_ends, 2) AND element_at(other_ends, 1) <= element_at(other_ends, 2)) THEN
                CASE WHEN end_date = element_at(other_ends, 1) THEN 'DUAL' ELSE 'MONO' END
        END AS commodity_aggr,
        CASE
            WHEN (commodity_aggr = 'MONO' AND element_at(other_starts, 1) <= element_at(other_starts, 2)) THEN start_date
            WHEN (commodity_aggr = 'DUAL' AND element_at(other_starts, 1) <= element_at(other_starts, 2)) THEN element_at(other_starts, 2)
            WHEN (commodity_aggr = 'DUAL' AND element_at(other_ends, 1) <= element_at(other_ends, 2)) THEN min_start
            WHEN (commodity_aggr = 'MONO' AND element_at(other_ends, 1) <= element_at(other_ends, 2)) THEN element_at(other_ends, 1)
        END AS new_start,
        CASE
            WHEN (commodity_aggr = 'MONO' AND element_at(other_starts, 1) <= element_at(other_starts, 2)) THEN element_at(other_starts, 2)
            WHEN (commodity_aggr = 'DUAL' AND element_at(other_starts, 1) <= element_at(other_starts, 2)) THEN max_end
            WHEN (commodity_aggr = 'DUAL' AND element_at(other_ends, 1) <= element_at(other_ends, 2)) THEN end_date
            WHEN (commodity_aggr = 'MONO' AND element_at(other_ends, 1) <= element_at(other_ends, 2)) THEN end_date
        END AS new_end
    FROM transformed_data
    WHERE row_count = 2
),
two_rows_formatted AS (
    SELECT
        customer,
        commodity_aggr,
        DATE_FORMAT(new_start, 'dd/MM/yyyy') AS start_date,
        DATE_FORMAT(new_end, 'dd/MM/yyyy') AS end_date
    FROM two_rows
)
SELECT * FROM single_row
UNION ALL
SELECT * FROM two_rows_formatted
ORDER BY customer;

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