PySpark DataFrame 多列多函数聚合结果的行式展示与优化实践

+--------+-------+-------+-------+-------+
|agg_type| col_1 | col_2 | col_3 | col_4 |
+--------+-------+-------+-------+-------+
|     min|  (min)|  (min)|  (min)|  (min)|
|     max|  (max)|  (max)|  (max)|  (max)|
+--------+-------+-------+-------+-------+

import operator
from pyspark.sql import SparkSession
from pyspark.sql import functions as F

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

# 示例数据
_data = [
    (4, 123, 18, 29),
    (8, 5, 26, 187),
    (2, 97, 18, 29),
]
_schema = ['col_1', 'col2', 'col3', 'col_4']
df = spark.createDataFrame(_data, _schema)

print("原始 DataFrame:")
df.show()

# 为每列创建最小值和最大值的聚合表达式
min_vals_exprs = [F.min(c).alias(f'min_{c}') for c in df.columns]
max_vals_exprs = [F.max(c).alias(f'max_{c}') for c in df.columns]

# 将所有聚合表达式合并，并通过 select 执行
# 此时 df2 将是一个单行 DataFrame，包含所有 min_ 和 max_ 列
df2 = df.select(min_vals_exprs + max_vals_exprs)

print("中间聚合结果 (单行):")
df2.show()
# +-----+----+----+-----+------+-----+------+-------+
# |min_col_1|min_col2|min_col3|min_col_4|max_col_1|max_col2|max_col3|max_col_4|
# +-----+----+----+-----+------+-----+------+-------+
# |        2|     5|      18|       29|        8|   123|      26|      187|
# +-----+----+----+-----+------+-----+------+-------+

# 缓存中间结果以提高性能
df2.cache()

# 准备最小值行的数据：添加 'agg_type' 列，并重命名聚合列
min_cols_selection = operator.add(
    [F.lit('min').alias('agg_type')],  # 添加聚合类型标识
    [F.col(f'min_{c}').alias(c) for c in df.columns] # 重命名 min_col 为原列名
)
min_df = df2.select(min_cols_selection)

print("最小值 DataFrame:")
min_df.show()
# +--------+-----+----+----+-----+
# |agg_type|col_1|col2|col3|col_4|
# +--------+-----+----+----+-----+
# |     min|    2|   5|  18|   29|
# +--------+-----+----+----+-----+

# 准备最大值行的数据：同样添加 'agg_type' 列并重命名
max_cols_selection = operator.add(
    [F.lit('max').alias('agg_type')],  # 添加聚合类型标识
    [F.col(f'max_{c}').alias(c) for c in df.columns] # 重命名 max_col 为原列名
)
max_df = df2.select(max_cols_selection)

print("最大值 DataFrame:")
max_df.show()
# +--------+-----+----+----+-----+
# |agg_type|col_1|col2|col3|col_4|
# +--------+-----+----+----+-----+
# |     max|    8| 123|  26|  187|
# +--------+-----+----+----+-----+

# 使用 unionByName 合并最小值和最大值 DataFrame
result = min_df.unionByName(max_df)

print("最终行式聚合结果:")
result.show()
# +--------+-----+----+----+-----+
# |agg_type|col_1|col2|col3|col_4|
# +--------+-----+----+----+-----+
# |     min|    2|   5|  18|   29|
# |     max|    8| 123|  26|  187|
# +--------+-----+----+----+-----+

# 停止 SparkSession
spark.stop()

import operator
from pyspark.sql import SparkSession
from pyspark.sql import functions as F

def aggregate_multiple_functions_row_wise(dataframe, functions_map):
    """
    对 PySpark DataFrame 的所有列应用多个聚合函数，并以行式结构展示结果。

    Args:
        dataframe (DataFrame): 输入的 PySpark DataFrame。
        functions_map (dict): 字典，键为聚合类型字符串（如 'min', 'max'），
                              值为对应的 PySpark 聚合函数（如 F.min, F.max）。

    Returns:
        DataFrame: 包含行式聚合结果的 DataFrame。
    """
    all_aggs_exprs = []
    for agg_type, func in functions_map.items():
        all_aggs_exprs.extend([func(c).alias(f'{agg_type}_{c}') for c in dataframe.columns])

    # 1. 执行所有聚合操作，生成单行中间结果
    intermediate_df = dataframe.select(all_aggs_exprs)
    intermediate_df.cache() # 缓存中间结果以提高性能

    result_dfs = []
    for agg_type in functions_map.keys():
        # 2. 为每种聚合类型重塑数据
        selection_cols = operator.add(
            [F.lit(agg_type).alias('agg_type')],
            [F.col(f'{agg_type}_{c}').alias(c) for c in dataframe.columns]
        )
        agg_df = intermediate_df.select(selection_cols)
        result_dfs.append(agg_df)

    # 3. 合并所有重塑后的结果
    if not result_dfs:
        return spark.createDataFrame([], schema=['agg_type'] + dataframe.columns)

    final_result = result_dfs[0]
    for i in range(1, len(result_dfs)):
        final_result = final_result.unionByName(result_dfs[i])

    intermediate_df.unpersist() # 释放缓存
    return final_result

if __name__ == "__main__":
    spark = SparkSession.builder.appName("MultiFunctionAggregationTutorial").getOrCreate()

    _data = [
        (4, 123, 18, 29),
        (8, 5, 26, 187),
        (2, 97, 18, 29),
    ]
    _schema = ['col_1', 'col2', 'col3', 'col_4']
    df = spark.createDataFrame(_data, _schema)

    print("原始 DataFrame:")
    df.show()

    # 定义要应用的聚合函数
    functions_to_apply = {
        'min': F.min,
        'max': F.max,
        'avg': F.avg,
        'sum': F.sum
    }

    # 调用函数获取行式聚合结果
    final_agg_df = aggregate_multiple_functions_row_wise(df, functions_to_apply)
    print("最终行式聚合结果:")
    final_agg_df.show()

    spark.stop()