pandas项目中的json_normalize功能改进探讨

背景介绍

在pandas数据处理过程中，经常会遇到需要处理嵌套JSON数据的情况。pandas提供的json_normalize函数是一个非常实用的工具，它能够将嵌套的JSON结构展平为表格形式。然而，在实际应用中，我们发现当对一个包含嵌套JSON列的DataFrame使用json_normalize时，输出结果会丢失原始DataFrame中的其他列信息。

当前问题分析

假设我们有一个包含交易信息的DataFrame，其中"dados"列存储了嵌套的JSON数据：

data = {
    "id de transação": [1, 2, 3, 4, 5],
    "nome": ["Alice", "Bob", "Charlie", "David", "Eve"],
    "dados": [
        {"data": "2024-01-01", "local": "São Paulo", "valor": 100.50},
        {"data": "2024-01-02", "local": "Rio de Janeiro", "valor": 200.75},
        {"data": "2024-01-03", "local": "Belo Horizonte", "valor": 300.00},
        {"data": "2024-01-04", "local": "Curitiba", "valor": 400.25},
        {"data": "2024-01-05", "local": "Porto Alegre", "valor": 500.50}
    ]
}
df = pd.DataFrame(data)

当我们直接对"dados"列使用json_normalize时，输出结果只包含展平后的JSON字段，而丢失了"id de transação"和"nome"列。

现有解决方案

目前，开发者需要手动将展平后的数据与原始DataFrame的其他列合并。常见的方法有两种：

1. 使用concat合并：

pd.concat(
    [
        df[['id de transação', 'nome']],
        pd.json_normalize(data=df['dados'], record_path=None),
    ],
    axis=1,
)

2. 使用join合并：

df[['id de transação', 'nome']].join(
    pd.json_normalize(data=df['dados'], record_path=None)
)

这两种方法都能达到目的，但在处理大规模数据或分布式计算环境(如Dask)时，可能会遇到性能或数据一致性方面的问题。

改进建议

可以考虑在json_normalize函数中增加一个参数，允许用户指定要保留的原始列。这样函数内部可以自动处理列合并，既简化了代码，又提高了性能。

例如，可以设计如下接口：

pd.json_normalize(
    data=df['dados'],
    record_path=None,
    keep_columns=['id de transação', 'nome']
)

分布式计算场景的考量

在Dask等分布式计算框架中，由于数据被分区处理，索引顺序可能无法保证。在这种情况下，保留一个唯一标识列(如交易ID)尤为重要。改进后的json_normalize可以确保即使在分布式环境下，也能正确关联展平后的数据与原始标识信息。

总结

json_normalize是pandas中处理嵌套JSON数据的强大工具，但在实际应用中存在一些不便之处。通过增加保留原始列的功能，可以显著提升用户体验，特别是在处理大规模分布式数据时。这种改进既保持了函数的简洁性，又增强了其在实际应用场景中的实用性。