如何加速Python中360万条记录的嵌套循环？

Question

我有一个包含 360 万条记录的 Companies.json 文件（每条记录都包含一个 id 和增值税号）和一个包含 76.000 条记录（+- 20 个属性）的 event.json 文件。我编写了一个脚本，执行以下步骤：

1. 打开两个 JSON 文件
2. 循环遍历 76.000 条事件记录（类型为类 dict）
3. 检查事件的状态是否是新的
4. 如果状态是新的，检查事件是否有companyID
5. 如果事件有companyID，则循环遍历360万条记录以查找匹配的公司ID。
6. 检查匹配的公司记录是否有增值税号
7. 将 companyID 替换为增值税号并添加 companyIDIsVat 布尔值。
8. 完成所有循环后，将事件写入新的 JSON 文件。

该脚本运行良好，但需要 6-7 小时才能完成。有办法加快速度吗？

当前脚本

import json

counter = 0;

with open('companies.json', 'r') as companiesFile:
    with open('events.json', 'r') as eventsFile:
        events = json.load(eventsFile)
        companies = json.load(companiesFile)

        for index, event in enumerate(events):
            print('Counter: ' + str(index))
            if 'status' in event:
                if(event['status'] == 'new'):
                    if 'companyID' in event:
                        for company in companies:
                            if(event['companyID'] == company['_id']):
                                if 'vat' in company:
                                    event['companyID'] = company['vat']
                                    event['companyIDIsVat'] = 1
                                    counter = counter + 1
                                    print('Found matches: ' + str(counter))
        
        with open('new_events.json', 'w', encoding='utf-8') as f:
            json.dump(events, f, ensure_ascii=False, indent=4)

Answer 1

因此，问题在于您要反复搜索整个公司列表。但列表对于搜索来说效率很低，因为在这里，你必须进行线性搜索，即 O(N)。但是如果您使用字典，您可以进行恒定时间搜索。假设你是company['_id']独一无二的。基本上，您想要对您的 ID 建立索引。对于恒定时间查找，请使用字典，即映射（CPython 中的哈希映射，可能还有每个 Python 实现）：

import json

counter = 0

with open('companies.json', 'r') as companiesFile:
    with open('events.json', 'r') as eventsFile:
        events = json.load(eventsFile)
        companies = {
            c["_id"]: c for c in json.load(companiesFile)
        }

        for index, event in enumerate(events):
            print('Counter: ' + str(index))
            if 'status' in event:
                if (
                    event['status'] == 'new' 
                    and 'companyID' in event 
                    and event['companyID'] in companies
                ):
                    company = companies[event['companyID']]
                    if 'vat' in company:
                        event['companyID'] = company['vat']
                        event['companyIDIsVat'] = 1
                        counter = counter + 1
                        print('Found matches: ' + str(counter))
        
        with open('new_events.json', 'w', encoding='utf-8') as f:
            json.dump(events, f, ensure_ascii=False, indent=4)

这是对脚本的最小修改。

您可能应该只保存companies.json在适当的结构中。

同样，它假设公司的 ID 是唯一的。如果没有，那么你可以使用列表字典，只要没有太多重复的 ID，它应该仍然会更快