PySpark：如何将列名列表传入函数并合并聚合结果为单个DataFrame

你现在遇到的问题是用map生成了多个独立的DataFrame，而想要的是把这些结果纵向合并成一个包含列名和对应最大值的单DataFrame。flatMap在这里并不适用——它主要用于RDD中展开元素，而我们需要的是DataFrame的合并或者宽表转长表的操作。下面提供两种可行的方案：

方案1：逐个生成小DataFrame后合并 #

这个思路是为每个列生成一个包含Column_name和Max的小DataFrame，然后通过union把所有小DataFrame合并起来：

from pyspark import SparkContext
from pyspark.sql import SQLContext
from pyspark.sql.functions import max as sparkMax, lit

sc = SparkContext("local[2]", "Count App")
sqlContext = SQLContext(sc)
df = sqlContext.createDataFrame(
    [(1, 100, 200), (100, 200, 100), (100, 200, 100), (-100, 50, 200)],
    ("col1", "col2", "col3"))

colnames = ['col1','col2','col3']

# 为每个列生成带列名和最大值的DataFrame
df_list = []
for col in colnames:
    col_max_df = df.agg(
        lit(f'max_of_{col}').alias('Column_name'),
        sparkMax(df[col]).alias('Max')
    )
    df_list.append(col_max_df)

# 合并所有DataFrame
final_df = df_list[0]
for df_item in df_list[1:]:
    final_df = final_df.union(df_item)

final_df.show()

运行后输出：

+-----------+---+
|Column_name|Max|
+-----------+---+
|max_of_col1|100|
|max_of_col2|200|
|max_of_col3|200|
+-----------+---+

方案2：一次聚合后转长表（更高效） #

如果你的Spark版本在3.0及以上，可以用stack函数直接把宽表转成你需要的长表，这种方式只需要一次聚合操作，性能更好：

from pyspark import SparkContext
from pyspark.sql import SQLContext
from pyspark.sql.functions import max as sparkMax

sc = SparkContext("local[2]", "Count App")
sqlContext = SQLContext(sc)
df = sqlContext.createDataFrame(
    [(1, 100, 200), (100, 200, 100), (100, 200, 100), (-100, 50, 200)],
    ("col1", "col2", "col3"))

colnames = ['col1','col2','col3']

# 一次计算所有列的最大值，得到宽表
wide_max_df = df.agg(*[sparkMax(c).alias(f'max_of_{c}') for c in colnames])

# 动态生成stack表达式，适配任意数量的列
stack_expr = f"""
stack({len(colnames)}, {', '.join([f"'max_of_{c}', max_of_{c}" for c in colnames])}) 
as (Column_name, Max)
"""

# 转成长表
final_df = wide_max_df.selectExpr(stack_expr)

final_df.show()

这个方案的输出和方案1完全一致，但只需要执行一次聚合，在数据量大的时候优势明显。

为什么原来的map不行？ #

你之前的map会为每个列生成一个只有最大值列的DataFrame，这些DataFrame的结构不同（列名不一样），无法直接合并。而我们需要的是每个结果都包含相同的两列（Column_name和Max），这样才能通过union合并成一个统一的DataFrame。

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