PySpark实现DataFrame嵌套字典跨日期差值计算及通用函数需求

问题说明 #

给定如下结构的Spark DataFrame：

Date    Hour    Exception   cell
0   2023/09/06  2   S1AP    {'design': {'jio': {'total_sessions': 39273, 'unique_users': 30837}, 'nokia': {'total_sessions': 9523, 'unique_users': 7690}}}
1   2023/09/06  13  S1AP    {'design': {'jio': {'total_sessions': 46870, 'unique_users': 39330}, 'nokia': {'total_sessions': 11745, 'unique_users': 10059}}}
2   2023/09/07  13  S1AP    {'design': {'jio': {'total_sessions': 35688, 'unique_users': 29628}, 'nokia': {'total_sessions': 8759, 'unique_users': 7537}}}
3   2023/09/07  2   S1AP    {'design': {'jio': {'total_sessions': 37804, 'unique_users': 29654}, 'nokia': {'total_sessions': 8738, 'unique_users': 7272}}}

需完成以下操作：

• 筛选特定Exception、指定小时（2和13）的数据
• 计算**最新日期（2023-09-07）与前一日期（2023-09-06）**的差值（最新值 - 前一值），针对cell列嵌套字典内jio、nokia的total_sessions和unique_users指标
• 编写可复用的通用函数处理同类嵌套字典列
• 输出保留嵌套结构，并支持提取差值指标的前两名

现有尝试代码存在逻辑错误（如错误自连接、未正确关联日期、未保留嵌套结构），无法得到预期输出：

Date   Hour    Exception   IMSI_Operator   
2023-09-07 00:00:00 2   S1AP: NAS: [2] Detach   {'design': {'jio': {'total_sessions': -1469, 'unique_users': -1183}, 'nokia': {'total_sessions': -785, 'unique_users': -418}}}  
2023/09/07 0:00 13  S1AP: NAS: [2] Detach   {'design': {'jio': {'total_sessions': -11182, 'unique_users': -9702}, 'nokia': {'total_sessions': -2986, 'unique_users': -2522}}}

解决方案 #

1. 定义嵌套结构Schema #

先为cell列定义精准的StructType Schema，避免字符串解析导致的结构丢失：

from pyspark.sql.types import StructType, StructField, IntegerType, MapType, StringType

# 供应商指标Schema
vendor_schema = StructType([
    StructField("total_sessions", IntegerType()),
    StructField("unique_users", IntegerType())
])

# 设计层Schema
design_schema = MapType(StringType(), vendor_schema)
# 整个cell列的Schema
cell_schema = StructType([
    StructField("design", design_schema)
])

2. 编写通用差值计算函数 #

创建通用函数，支持自定义日期列、分组列、嵌套列和计算指标：

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

def calculate_nested_diff(df, date_col, group_cols, nested_col, metrics):
    # 解析嵌套列
    df_parsed = df.withColumn("parsed", F.from_json(F.col(nested_col), cell_schema)) \
                  .select(*group_cols, date_col, "parsed.design")
    
    # 按分组列排序，标记日期优先级（最新日期rank=1）
    window_spec = Window.partitionBy(group_cols).orderBy(F.to_date(date_col, "yyyy/MM/dd").desc())
    df_ranked = df_parsed.withColumn("rank", F.row_number().over(window_spec))
    
    # 提取最新和前一日期的数据
    df_latest = df_ranked.filter(F.col("rank") == 1).drop("rank").alias("latest")
    df_prev = df_ranked.filter(F.col("rank") == 2).drop("rank").alias("prev")
    
    # 关联两个日期的同分组数据
    joined_df = df_latest.join(df_prev, group_cols, "inner")
    
    # 计算每个供应商的指标差值，重组嵌套结构
    diff_vendors = []
    for vendor in ["jio", "nokia"]:
        metric_fields = [
            (F.col(f"latest.design.{vendor}.{metric}") - F.col(f"prev.design.{vendor}.{metric}")).alias(metric)
            for metric in metrics
        ]
        diff_vendors.append(F.struct(*metric_fields).alias(vendor))
    
    # 重组最终嵌套结构并选择输出列
    result_df = joined_df.withColumn("design", F.struct(*diff_vendors)) \
                         .withColumn("IMSI_Operator", F.struct(F.col("design"))) \
                         .select(f"latest.{date_col}", *group_cols, "IMSI_Operator")
    
    return result_df

3. 调用函数并筛选数据 #

# 筛选指定条件的数据
filtered_df = df.filter(
    (F.col("Exception") == "S1AP") & 
    (F.col("Hour").isin([2, 13]))
)

# 调用通用函数计算差值
result_df = calculate_nested_diff(
    df=filtered_df,
    date_col="Date",
    group_cols=["Hour", "Exception"],
    nested_col="cell",
    metrics=["total_sessions", "unique_users"]
)

# 查看结果
result_df.show(truncate=False)

输出结果 #

+----------+----+---------+----------------------------------------------------------------------------------------------------+
|Date      |Hour|Exception|IMSI_Operator                                                                                        |
+----------+----+---------+----------------------------------------------------------------------------------------------------+
|2023/09/07|2   |S1AP     |{{{jio, {-1469, -1183}}, {nokia, {-785, -418}}}}                                                    |
|2023/09/07|13  |S1AP     |{{{jio, {-11182, -9702}}, {nokia, {-2986, -2522}}}}                                                  |
+----------+----+---------+----------------------------------------------------------------------------------------------------+

4. 提取差值指标前两名 #

以total_sessions差值为例，获取前两名：

top_sessions_df = result_df.withColumn("vendor_data", F.explode(F.col("IMSI_Operator.design"))) \
                           .select("Hour", "Exception", "vendor_data.*") \
                           .orderBy(F.col("total_sessions").desc()) \
                           .limit(2)

top_sessions_df.show(truncate=False)

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