Return 生成器而不是来自 df.to_dict() 的列表

Question

我正在处理一个大型 Pandas DataFrame，它需要在被另一个 API 处理之前转换成字典。

需要的词典可以通过调用.to_dict(orient='records')方法生成。如文档中所述，returned 值取决于 orient 选项：

Returns: dict, list or collections.abc.Mapping

Return a collections.abc.Mapping object representing the DataFrame. The resulting transformation depends on the orient parameter.

就我而言，通过 orient='records'，return 编辑了一个字典列表。在处理列表时，存储列表项所需的完整内存是 reserved/allocated。由于我的数据帧可能会变得相当大，这可能会导致内存问题，尤其是当代码可能在较低规格的目标系统上执行时。

我当然可以通过按块处理数据帧并为每个块生成字典列表来绕过这个问题，然后将其传递给 API。此外，调用 iter(df.to_dict(orient='records')) 会 return 所需的生成器，但不会减少所需的内存占用，因为列表是在中间创建的。

有没有办法从 df.to_dict(orient='records') 直接 return 生成器表达式而不是列表以减少内存占用？

Answer 1

无法直接从 to_dict(orient='records') 获取生成器。但是，可以将 to_dict source code 修改为生成器而不是返回列表理解：

from pandas.core.common import standardize_mapping
from pandas.core.dtypes.cast import maybe_box_native


def dataframe_records_gen(df_):
    columns = df_.columns.tolist()
    into_c = standardize_mapping(dict)

    for row in df_.itertuples(index=False, name=None):
        yield into_c(
            (k, maybe_box_native(v)) for k, v in dict(zip(columns, row)).items()
        )

示例代码：

import pandas as pd

df = pd.DataFrame({
    'A': [1, 2],
    'B': [3, 4]
})

# Using Generator
for row in dataframe_records_gen(df):
    print(row)

# For Comparison with to_dict function
print("to_dict", df.to_dict(orient='records'))

输出：

{'A': 1, 'B': 3}
{'A': 2, 'B': 4}
to_dict [{'A': 1, 'B': 3}, {'A': 2, 'B': 4}]

为了更自然的语法，也可以 register a custom accessor:

import pandas as pd
from pandas.core.common import standardize_mapping
from pandas.core.dtypes.cast import maybe_box_native


@pd.api.extensions.register_dataframe_accessor("gen")
class GenAccessor:
    def __init__(self, pandas_obj):
        self._obj = pandas_obj

    def records(self):
        columns = self._obj.columns.tolist()
        into_c = standardize_mapping(dict)

        for row in self._obj.itertuples(index=False, name=None):
            yield into_c(
                (k, maybe_box_native(v))
                for k, v in dict(zip(columns, row)).items()
            )

在这种情况下，这使得可以通过 gen 访问器访问此生成器：

df = pd.DataFrame({
        'A': [1, 2],
        'B': [3, 4]
    })

# Using Generator through registered custom accessor
for row in df.gen.records():
    print(row)

# For Comparison with to_dict function
print("to_dict", df.to_dict(orient='records'))

输出：

{'A': 1, 'B': 3}
{'A': 2, 'B': 4}
to_dict [{'A': 1, 'B': 3}, {'A': 2, 'B': 4}]

Return 生成器而不是来自 df.to_dict() 的列表

Return generator instead of list from df.to_dict()

python

pandas

dataframe

dictionary

generator