时间戳

首页 > 解决方案 > 最大限度地减少大型 CSV 文件中 python 的搜索时间

python - 最大限度地减少大型 CSV 文件中 python 的搜索时间

问题描述

我有一个包含大约 700 行左右和 8 列的 CSV 文件，但是最后一列有一个非常大的文本块（每个里面有足够的多个长段落）。

我想通过 python 实现一个文本搜索功能，它可以让我返回所有文本与第 8 列数据内部匹配的行（这意味着它需要遍历整个内容）。

解决这个问题并最小化搜索时间的最快方法可能是什么？

标签： pythonpython-3.xsqlitecsvsearch

解决方案

您可以将 csv 文件转储到sqlite数据库中，并使用 sqlite 的全文搜索功能为您进行搜索。

此示例代码显示了它是如何完成的。有几点需要注意：

  • 它假定 csv 文件有一个标题行。如果不是这种情况，您需要提供列名（或者只使用通用名称，如“col1”、“col2”等）。
  • 它搜索csv 中的所有列；如果不希望这样做，请在创建 SQL 语句之前过滤掉其他列（和标题值）。
  • 如果您希望能够将结果与 csv 文件中的行进行匹配，则需要创建一个包含行号的列。
import csv                                                                                                                                      
import sqlite3                                                                                                                                  
import sys                                                                                                                                      


def create_table(conn, cols, name='mytable'):                                                                                                   
    stmt = f"""CREATE VIRTUAL TABLE "{name}" USING fts5({cols})"""                                                                              
    with conn:                                                                                                                                  
        conn.execute(stmt)                                                                                                                      
    return                                                                                                                                      


def populate_table(conn, reader, cols, ncols, name='mytable'):                                                                                  
    placeholders = ', '.join(['?'] * ncols)                                                                                                     
    stmt = f"""INSERT INTO "{name}" ({cols})                                                                                                    
    VALUES ({placeholders})
    """
    # Filter out any blank rows in the csv
    reader = filter(None, reader)
    with conn:
        conn.executemany(stmt, reader)
    return


def search(conn, term, cols, name='mytable'):
    stmt = f"""SELECT {cols}
    FROM "{name}"
    WHERE "{name}" MATCH ?
    """
    with conn:
        cursor = conn.cursor()
        cursor.execute(stmt, (term,))
        result = cursor.fetchall()
    return result


def main(path, term):
    result = 'NO RESULT SET'
    try:
        conn = sqlite3.connect(':memory:')
        with open(path, 'r') as f:
            reader = csv.reader(f)
            # Assume headers are in the first row
            headers = next(reader)
            ncols = len(headers)
            cols = ', '.join([f'"{x.strip()}"' for x in headers])
            create_table(conn, cols)
            populate_table(conn, reader, cols, ncols)
        result = search(conn, term, cols)
    finally:
        conn.close()
    return result


if __name__ == '__main__':
    print(main(*sys.argv[1:]))

推荐阅读