如何将Pandas DataFrame分组后将指定列转换为JSON字符串？

原始数据与需求 #

现有如下Pandas DataFrame，需按market字段分组，将每组的year、properties、home_values列转换为JSON格式的字符串，最终生成指定结构的结果表。

原始数据代码：

import pandas as pd

# 创建数据列表
data = [
    ['San Francisco', '10', 2023, 1000000],
    ['San Francisco', '8', 2022, 900000],
    ['Los Angeles', '12', 2023, 800000],
    ['Los Angeles', '9', 2022, 700000],
]

# 创建DataFrame
df = pd.DataFrame(data, columns=['market', 'properties', 'year', 'home_values'])

print(df)

预期结果 #

market  year                                       home values
0  San Francisco  2023  {"year": [2023, 2022], "properties": [10, 8], "home_values": [1000000, 900000]}
1    Los Angeles  2023  {"year": [2023, 2022], "properties": [12, 9], "home_values": [800000, 700000]}

解决方案 #

步骤如下：

1. 先将properties列转换为整数类型（匹配预期结果中的数值格式）
2. 定义分组转换函数，将每组的目标列转为列表后再序列化为JSON字符串
3. 按market分组聚合生成JSON列，同时提取每组的最新年份（最大year值）
4. 调整列顺序匹配预期结果

完整代码：

import pandas as pd
import json

# 创建数据列表
data = [
    ['San Francisco', '10', 2023, 1000000],
    ['San Francisco', '8', 2022, 900000],
    ['Los Angeles', '12', 2023, 800000],
    ['Los Angeles', '9', 2022, 700000],
]

# 创建DataFrame
df = pd.DataFrame(data, columns=['market', 'properties', 'year', 'home_values'])

# 将properties列转为整数类型
df['properties'] = df['properties'].astype(int)

# 定义分组转JSON的函数
def group_to_json(group):
    return json.dumps({
        'year': group['year'].tolist(),
        'properties': group['properties'].tolist(),
        'home_values': group['home_values'].tolist()
    })

# 分组生成JSON列
result_df = df.groupby('market').apply(group_to_json).reset_index(name='home values')

# 添加每组的最新年份（最大year值）
result_df['year'] = df.groupby('market')['year'].max().values

# 调整列顺序
result_df = result_df[['market', 'year', 'home values']]

print(result_df)

运行后即可得到与预期一致的结果。

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