Pandas DataFrame按market分组生成按年排序的JSON字符串

原数据与初始DataFrame #

构造数据并生成DataFrame的代码：

import pandas as pd
# 构造数据列表
data = [['market1', 2023, 100, 200],
        ['market2', 2022, 300, 400],
        ['market1', 2021, 500, 600],
        ['market2', 2020, 700, 800]]
# 创建DataFrame
df = pd.DataFrame(data, columns=['market', 'year', 'val1', 'val2'])

打印后的数据结构：

market  year  val1  val2
0  market1  2023   100   200
1  market2  2022   300   400
2  market1  2021   500   600
3  market2  2020   700   800

需求说明 #

按market列分组，生成year-val1和year-val2的映射列并保存为JSON字符串（year作为键，对应值为val1或val2），且JSON中的year键需按升序排序，期望输出如下：

market  val1_json                   val2_json
market1 {"2021": 500, "2023": 100}  {"2021": 600, "2023": 200}
market2 {"2020": 700, "2022": 300}  {"2020": 800, "2022": 400}

解决方案 #

直接通过groupby结合自定义聚合函数实现，代码如下：

import pandas as pd
import json

# 构造数据
data = [['market1', 2023, 100, 200],
        ['market2', 2022, 300, 400],
        ['market1', 2021, 500, 600],
        ['market2', 2020, 700, 800]]
df = pd.DataFrame(data, columns=['market', 'year', 'val1', 'val2'])

# 自定义聚合函数：按year升序排序后生成JSON字符串
def to_sorted_json(series):
    # 关联原year列并按年份升序排序
    sorted_df = series.reset_index().sort_values('year')
    # 构造year:value的字典，年份转为字符串
    mapping = {str(row['year']): row[series.name] for _, row in sorted_df.iterrows()}
    # 转为JSON字符串
    return json.dumps(mapping)

# 按market分组，对val1、val2应用聚合函数
result = df.groupby('market').agg(
    val1_json=('val1', to_sorted_json),
    val2_json=('val2', to_sorted_json)
).reset_index()

print(result)

执行后输出结果与期望一致：

market                  val1_json                  val2_json
0  market1  {"2021": 500, "2023": 100}  {"2021": 600, "2023": 200}
1  market2  {"2020": 700, "2022": 300}  {"2020": 800, "2022": 400}

代码说明 #

• to_sorted_json函数：接收分组后的val1/val2序列，先关联原year列并按年份升序排序，再构造键为字符串类型年份的字典，最后转为JSON字符串；
• groupby.agg：指定要聚合的列和对应的自定义函数，直接生成目标结果；
• reset_index()：将market从索引转为普通列，匹配期望的输出格式。

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