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

Question

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

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

返回：字典、列表或 collections.abc.Mapping

返回表示 DataFrame 的 collections.abc.Mapping object。生成的转换取决于 orient 参数。

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

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

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

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'))

Output：

{'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'))

Output：

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