针对含50列的大型数据集，现有Pandas代码是否可优化？

现有一段将电影数据转换为指定结构化JSON的Pandas代码，计划将其应用到包含约50列的大型数据集，询问该代码是否存在优化空间。

输入数据表格 #

期望输出JSON #

[
  {
    "id": "ab123",
    "name": "Interstellar",
    "classifications": {
      "classification": [
        {
          "genre": "Sci-fi",
          "year": 2020
        },
        {
          "genre": "Space",
          "year": 2020
        }
      ]
    },
    "language": "English",
    "type": "Hollywood"
  },
  {
    "id": "ab125",
    "name": "Hero",
    "classifications": {
      "classification": [
        {
          "genre": "Action",
          "year": 2020
        },
        {
          "genre": "Fighthing",
          "year": 2020
        }
      ]
    },
    "language": "English",
    "type": "Hollywood"
  },
  {
    "id": "ab567",
    "name": "Love",
    "classifications": {
      "classification": [
        {
          "genre": "Dram",
          "year": 2022
        }
      ]
    },
    "language": "English",
    "type": "Hollywood"
  }
]

当前实现代码 #

import pandas
import json

df = pandas.read_excel('C:/Users/pika/Desktop/sample.xlsx', sheet_name='Movie')

# Do a group by and name the newly created column
final = df.groupby(['id', 'name', 'language', 'type']).apply(lambda x: {"classification":x[['genre','year']].to_dict('records')}).reset_index(name="classifications").to_dict(orient="records")

print(json.dumps(final, indent=2))

优化建议 #

针对大型数据集（50列），可以从以下几个方向优化代码的性能和可维护性：

1. 自动生成分组键，避免手动列举 #

当分组列多达几十列时，手动写所有分组键容易出错且繁琐，可通过df.columns.difference自动获取除genre、year外的所有列作为分组键：

group_cols = df.columns.difference(['genre', 'year']).tolist()
final = df.groupby(group_cols).apply(lambda x: {"classification": x[['genre','year']].to_dict('records')}).reset_index(name="classifications").to_dict(orient="records")

2. 替换apply为更高效的聚合方式 #

groupby.apply本质是逐组遍历，性能较差。改用agg进行向量式聚合，再构造目标结构，能大幅提升大数据集下的处理速度：

group_cols = df.columns.difference(['genre', 'year']).tolist()

# 先聚合genre和year为列表
agg_result = df.groupby(group_cols).agg(
    genre_list=('genre', list),
    year_list=('year', list)
).reset_index()

# 构造classification嵌套结构
agg_result['classifications'] = agg_result.apply(
    lambda row: {"classification": [{"genre": g, "year": y} for g, y in zip(row['genre_list'], row['year_list'])]},
    axis=1
)

# 移除临时聚合列，转换为目标格式
agg_result = agg_result.drop(['genre_list', 'year_list'], axis=1)
final = agg_result.to_dict(orient="records")

3. 读取数据时做内存优化 #

大型数据集读取时，通过指定数据类型、只加载必要列减少内存占用，提升后续处理速度：

# 定义需要加载的列（根据实际业务调整）
needed_cols = df.columns.tolist()  # 或手动指定必要列
# 指定数据类型，将枚举类列设为category，数值列设为更紧凑的类型
dtype_spec = {
    'id': 'string',
    'name': 'string',
    'genre': 'category',
    'language': 'category',
    'type': 'category',
    'year': 'int32'
    # 其他50列同理，根据数据类型设置
}

df = pandas.read_excel(
    'C:/Users/pika/Desktop/sample.xlsx',
    sheet_name='Movie',
    usecols=needed_cols,
    dtype=dtype_spec
)

4. 减少中间数据转换开销 #

如果最终只需要JSON输出，可以直接在分组后构造字典，跳过部分DataFrame转换步骤：

group_cols = df.columns.difference(['genre', 'year']).tolist()
final = []

for group_key, group_data in df.groupby(group_cols):
    # 将分组键转为字典
    base_dict = dict(zip(group_cols, group_key))
    # 构造classifications结构
    base_dict['classifications'] = {
        "classification": group_data[['genre', 'year']].to_dict('records')
    }
    final.append(base_dict)

print(json.dumps(final, indent=2))

这种方式避免了reset_index和多次to_dict的额外开销，在数据量较大时更高效。

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