Question

我想将大量条目（~600k）上传到PostgreSQL DB中的一个简单表中，每个条目有一个外键，一个时间戳和3个浮点数。但是，每个条目执行描述here的核心批量插入需要60毫秒，因此整个执行需要10小时。我发现，这是executemany()方法的性能问题，但它已使用psycopg2 2.7中的execute_values()方法解决。

我运行的代码如下：

#build a huge list of dicts, one dict for each entry
engine.execute(SimpleTable.__table__.insert(),
               values) # around 600k dicts in a list

我发现这是一个常见的问题，但我还没有找到sqlalchemy本身的解决方案。在某些情况下，有没有办法告诉sqlalchemy打电话给execute_values()？有没有其他方法可以实现大型插入而无需自己构建SQL语句？

感谢您的帮助！

Answer 1

同时，use_batch_mode函数上的create_engine()标志成为可能（来自SqlAlchemy 1.2.0）。请参阅docs。它使用execute_batch()中的psycopg.extras函数。

Answer 2

这不是你正在寻找的答案，因为这并不意味着试图指示SQLAlchemy使用psycopg附加组件，并且需要 - 排序 - 手动SQL，但是：你可以从引擎访问底层的psycopg连接raw_connection()，允许使用COPY FROM：

import io
import csv
from psycopg2 import sql

def bulk_copy(engine, table, values):
    csv_file = io.StringIO()
    headers = list(values[0].keys())
    writer = csv.DictWriter(csv_file, headers)
    writer.writerows(values)

    csv_file.seek(0)

    # NOTE: `format()` here is *not* `str.format()`, but
    # `SQL.format()`. Never use plain string formatting.
    copy_stmt = sql.SQL("COPY {} (" +
                        ",".join(["{}"] * len(headers)) +
                        ") FROM STDIN CSV").\
        format(sql.Identifier(str(table.name)),
               *(sql.Identifier(col) for col in headers))

    # Fetch a raw psycopg connection from the SQLAlchemy engine
    conn = engine.raw_connection()
    try:
        with conn.cursor() as cur:
            cur.copy_expert(copy_stmt, csv_file)

    except Exception as e:
        conn.rollback()
        raise e

    else:
        conn.commit()

    finally:
        conn.close()

然后

bulk_copy(engine, SimpleTable.__table__, values)

与执行INSERT语句相比，这应该足够快。在这台机器上移动600,000条记录大约需要8秒，大约13μs/条记录。您还可以将原始连接和光标与extras包一起使用。

我怎样才能在sqlalchemy中使用psycopg2.extras？

2 个答案: