从 df.to_dict() 返回生成器而不是列表

Question

我正在开发一个大型Pandas DataFrame，在由另一个 API 处理之前需要将其转换为字典。

调用该.to_dict(orient='records')方法即可生成所需的词典。如文档中所述，返回值取决于选项orient：

返回：dict、list 或 collections.abc.Mapping

返回表示 DataFrame 的 collections.abc.Mapping 对象。生成的变换取决于方向参数。

对于我的情况，传递orient='records'，返回字典列表。处理列表时，存储列表项所需的完整内存将被保留/分配。由于我的数据帧可能变得相当大，这可能会导致内存问题，特别是当代码可能在较低规格的目标系统上执行时。

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

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

Answer 1

没有办法直接从 获取生成器to_dict(orient='records')。但是，可以将to_dict 源代码修改为生成器而不是返回列表理解：

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}]

为了更自然的语法，还可以注册自定义访问器：

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}]