如何在 Python 中对大量字典进行排序而不加载到内存中

Posted 2023-02-25

【中文标题】如何在 Python 中对大量字典进行排序而不加载到内存中

【英文标题】：How to sort a large list of dictionaries without loading into memory in Python

【发布时间】：2021-10-29 00:09:16

【问题描述】：

我有一个包含大约 5000 万行的 CSV 文件，我正在尝试操作数据并写入一个新的 CSV 文件。下面是代码：

import csv
import itertools

def main():
    with open("input.csv", "r") as csvfile:
        rows = csv.DictReader(csvfile)
        sorted_rows = sorted(rows, key=lambda row: row["name"])
        grouping = groupby(sorted_rows, lambda row: row["name"])

        with open("output.csv", "w") as final_csvfile:
            fieldnames = ["name", "number"]
            writer = csv.DictWriter(final_csvfile, fieldnames=fieldnames)

            for group, items in grouping:
                total = sum(int(item["number"]) for item in items)
                writer.writerow(
                    
                        "name": group,
                        "number": str(total),
                    
                )


if __name__ == "__main__":
    main()

这在不太多的行上运行良好，但是当我运行具有 5000 万行的实际 CSV 时，它变得非常慢并且程序最终会被终止。

现在这一行：sorted_rows = sorted(rows, key=lambda row: row["name"]) 是主要问题，因为它将 5000 万行加载到内存（列表）中，因此可以对其进行排序。我已经明白sorted() 所做的第一件事就是将给它的任何生成器转换为一个列表，那么我该怎么做呢？有什么指点吗？

【参考方案1】：

@python_user 上述方法的问题在于，它会一直附加到字典中，在你知道它之前，字典会变得非常大，并且可能会在内存方面搞砸。

@Bharel 说了一些关于 cmets 中的外部排序，我研究了一下，找到了一种方法。

我发现 UNIX 排序命令可以对非常大的文件执行外部合并排序，所以我编写了一个脚本来对非常大的 CSV 文件进行排序，然后将排序后的 CSV 文件传递​​给问题中的 python 代码。这样，没有什么太大的东西被写入内存。

代码如下：

sort.sh

echo "sorting CSV"
sort input.csv -o input.csv
echo "Done!"

上述脚本运行后，将排序后的CSV传入程序：

import csv
from itertools import groupby


def main():
    with open("input.csv", "r") as csvfile:
        rows = csv.DictReader(csvfile)
        grouping = groupby(rows, lambda row: row["name"])

        with open("output.csv", "w") as final_csvfile:
            fieldnames = ["name", "number"]
            writer = csv.DictWriter(final_csvfile, fieldnames=fieldnames)

            for group, items in grouping:
                total = sum(int(item["number"]) for item in items)
                writer.writerow(
                    
                        "name": group,
                        "number": str(total),
                    
                )


if __name__ == "__main__":
    main()

请注意，问题中使用 sorted() 的行已消失。我认为这是一个更有效的解决方案。

【讨论】：

您是否考虑了多行 CSV 单元格？尽管您的解决方案非常出色，但 sort 很遗憾会破坏它。

欲了解更多信息，请参阅RFC4180 section 2 clause 6。

@Bharel 我将通读这篇文章。谢谢！