Spark DataFrame列转行并生成以gw为键的映射列问题求助

原代码问题分析 #

1. 键值错误：使用collect_list(lit("gw"))会将字符串"gw"作为映射的键，而非gw列实际的IP地址值，正确写法应为collect_list(col("gw"))。
2. 函数缺失：未导入array函数，导致array(pivot_df['rrc'], pivot_df['re_est'])执行报错。
3. 语法错误：最后处理combined_map时，直接调用selectExpr的方式不符合Spark DataFrame操作逻辑，需先将映射列加入DataFrame再拆解。

修正方案1：使用map_from_arrays + explode实现 #

from pyspark.sql import SparkSession
from pyspark.sql.functions import collect_list, map_from_arrays, array, explode, col, lit

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

# 构建原始DataFrame
data = [("210.142.27.137", 1400.0, 26.0),
        ("210.142.27.202", 2300, 12)]
columns = ["gw", "rrc", "re_est"]
df = spark.createDataFrame(data, columns)

# 收集gw列表与各指标值列表
agg_df = df.groupBy().agg(
    collect_list(col("gw")).alias("gw_list"),
    collect_list(col("rrc")).alias("rrc_list"),
    collect_list(col("re_est")).alias("re_est_list")
)

# 生成映射并拆分为目标行结构
result_df = agg_df.select(
    explode(
        map_from_arrays(
            array(lit("rrc"), lit("re_est")),
            array(
                map_from_arrays(col("gw_list"), col("rrc_list")),
                map_from_arrays(col("gw_list"), col("re_est_list"))
            )
        )
    ).alias("index", "gw_mapping")
)

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

修正方案2：使用stack函数简化实现 #

如果需要更简洁的写法，可利用stack直接将多列转成键值对行：

from pyspark.sql import SparkSession
from pyspark.sql.functions import collect_list, map_from_arrays, col

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

# 构建原始DataFrame
data = [("210.142.27.137", 1400.0, 26.0),
        ("210.142.27.202", 2300, 12)]
columns = ["gw", "rrc", "re_est"]
df = spark.createDataFrame(data, columns)

# 直接生成各指标的映射列
agg_df = df.groupBy().agg(
    map_from_arrays(collect_list(col("gw")), collect_list(col("rrc"))).alias("rrc"),
    map_from_arrays(collect_list(col("gw")), collect_list(col("re_est"))).alias("re_est")
)

# 用stack将多列转成行
result_df = agg_df.selectExpr(
    "stack(2, 'rrc', rrc, 're_est', re_est) as (index, gw_mapping)"
)

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

两种方案最终输出均符合预期：

+-------+------------------------------------------------------------------+
|index  |gw_mapping                                                         |
+-------+------------------------------------------------------------------+
|rrc    |{210.142.27.137:1400.0, 210.142.27.202:2300}                      |
|re_est |{210.142.27.137:26.0, 210.142.27.202:12}                          |
+-------+------------------------------------------------------------------+

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