百万级数据：列表与DataFrame搜索效率对比及df.loc优化问询

我原本将数据存储在列表中，使用如下代码进行搜索：

next((sub for sub in curr_data_arr if (sub == query_string)), None)

但处理200万条数据时耗时过长。因此计划将数据存入DataFrame，设置索引后查询目标值是否存在，具体代码如下：

curr_data_df = pd.DataFrame(list(audit_collection.find({"source_name":default_attributes['source_name']}, project_keys)))
curr_data_df = curr_data_df.set_index(queryString)
audit_data = {}
try:
    print(tuple(query_string.values))
    stark_audit_data = curr_data_df.loc(tuple(query_string.values))
except Exception as e:
    pass

想了解使用DataFrame搭配df.loc搜索是否会大幅提速？

补充说明
单条数据示例如下：

{
  "asin_product_title": "title 1",
  "contextHelperUrls": "url1",
  "d_category": "d1",
  "landing_page_url": "l1",
  "query": "how is it?",
  "leftProductTitle": "l1",
  "rightProductTitle": "r1",
  "moderator": "kodaga",
  "source_name": "SPRSPQA",
  "type": "Audit",
  "status": "Unworked",
  "uploadTime": "2023-01-24T02:42:27.825918+0000",
  "batchId": "638eafff21657a1f04c1faf0_fvj61wj51uxggv9t67efp8p8q7gdiokj_24012023024227"
}

是的，用DataFrame搭配loc基于索引查询会大幅提速，核心原因和优化点如下：

• 时间复杂度差异：列表遍历搜索的时间复杂度是O(n)，200万条数据需要逐条比对，最坏情况要扫描全部数据；而pandas的索引底层采用哈希表结构，查询的时间复杂度接近O(1)，定位数据的效率是数量级的提升。
• 代码修正：注意你的代码里curr_data_df.loc(tuple(query_string.values))写法有误，loc是属性访问器，应该用方括号调用：curr_data_df.loc[tuple(query_string.values)]。
• 额外优化建议：
  • 如果queryString是多列组合索引，确保这些列的数据类型统一，避免隐性类型转换带来的额外耗时；
  • 从MongoDB读取数据时，可直接用pd.DataFrame.from_records(audit_collection.find(...))生成DataFrame，减少中间列表转换的开销；
  • 若仅需判断目标值是否存在，无需返回整条数据，建议用tuple(query_string.values) in curr_data_df.index，比捕获异常的方式更高效、可读性更强。

内容的提问来源于stack exchange，提问作者Abcd Efgh