如何将DataFrame分组后指定列的数据聚合转换为元组格式

以下提供两种常用数据处理环境下的实现方案：

Pandas 实现 #

直接通过groupby+多列聚合即可实现需求：

import pandas as pd

# 构造示例DataFrame
df = pd.DataFrame({
    'Id': [1,1,1,1,2],
    'fomrid': ['x1','x1','x1','x2','p1'],
    'values': [22.0, 'test', 11, 21, 1],
    'occ': [1,2,3,0,1],
    'comments': ['text1','text2','text3','text4','text5']
})

# 分组聚合
result = df.groupby(['Id', 'fomrid'], as_index=False).agg(
    **{
        'tuple(values)': ('values', tuple),
        'tuple(occ)': ('occ', tuple),
        'comments': ('comments', 'first')
    }
)

print(result)

PySpark 实现 #

Spark原生collect_list返回数组格式，额外加一个简单UDF转换为元组即可：

from pyspark.sql import SparkSession
from pyspark.sql.functions import collect_list, first, udf

# 初始化Spark会话
spark = SparkSession.builder.appName("group_to_tuple").getOrCreate()

# 构造示例DataFrame
data = [
    (1, "x1", 22.0, 1, "text1"),
    (1, "x1", "test", 2, "text2"),
    (1, "x1", 11, 3, "text3"),
    (1, "x2", 21, 0, "text4"),
    (2, "p1", 1, 1, "text5")
]
df = spark.createDataFrame(data, schema=["Id", "fomrid", "values", "occ", "comments"])

# 定义列表转元组的UDF
to_tuple = udf(lambda x: tuple(x))

# 分组聚合
result = df.groupBy("Id", "fomrid").agg(
    to_tuple(collect_list("values")).alias("tuple(values)"),
    to_tuple(collect_list("occ")).alias("tuple(occ)"),
    first("comments").alias("comments")
)

result.show()

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