如何在peewee中使用带FORMAT BINARY的COPY FROM导入稀疏向量？

核心问题解答 #

1. Peewee是否支持带FORMAT BINARY的COPY FROM？
   Peewee本身没有提供高层API直接封装FORMAT BINARY的COPY FROM操作，但可以通过获取底层原生的psycopg连接，手动实现该功能——这也是导入大体积稀疏向量最高效的方式。

2. bulk_create()和insert_many()能否使用FORMAT BINARY？
   不能。这两个方法本质是生成多值INSERT语句（如INSERT INTO ... VALUES (...), (...)），属于SQL层面的批量插入，和COPY FROM的二进制协议机制完全不同，无法支持FORMAT BINARY。如果数据量不大可以用这两个方法，但大规模数据导入推荐用COPY FROM BINARY。

具体实现方案 #

1. 用Peewee结合原生psycopg实现COPY FROM BINARY #

首先定义Peewee模型（对应你的稀疏向量表）：

from peewee import PostgresqlDatabase, Model, IntegerField
from playhouse.postgres_ext import SparseVectorField

# 初始化数据库连接
db = PostgresqlDatabase(
    "your_database_name",
    user="your_user",
    password="your_password",
    host="localhost"
)

class TFIDFVector(Model):
    id = IntegerField(primary_key=True)
    vector = SparseVectorField()  # 对应pgvector的sparsevec类型

    class Meta:
        database = db
        table_name = "tfidf_vectors"

然后借助psycopg的二进制COPY能力实现批量导入：

import psycopg2.sql
from sklearn.feature_extraction.text import TfidfVectorizer
from pgvector.psycopg2 import SparseVector  # 需要安装pgvector包：pip install pgvector

# 假设已生成TF-IDF稀疏矩阵和对应ID
tfidf = TfidfVectorizer()
sparse_matrix = tfidf.fit_transform(["text sample 1", "text sample 2", "text sample 3"])
vector_ids = [1, 2, 3]  # 每个向量对应的唯一ID

# 通过Peewee获取原生连接并执行COPY FROM BINARY
with db.connection_context():
    conn = db.connection
    cur = conn.cursor()

    # 构造COPY语句
    copy_sql = psycopg2.sql.SQL("COPY tfidf_vectors (id, vector) FROM STDIN WITH FORMAT BINARY")
    
    # 开启二进制COPY会话
    with cur.copy(copy_sql) as copy:
        for vec_id, row in zip(vector_ids, sparse_matrix):
            # 将scipy稀疏行转换为pgvector兼容的SparseVector对象
            # psycopg会自动处理二进制序列化
            sparse_vec = SparseVector(row)
            copy.write_row((vec_id, sparse_vec))
    
    conn.commit()

2. 用bulk_create()实现批量插入（非二进制，适合小数据量） #

如果数据规模不大，可直接用Peewee的bulk_create方法：

# 构造模型实例列表
instances = []
for vec_id, row in zip(vector_ids, sparse_matrix):
    sparse_vec = SparseVector(row)
    instances.append(TFIDFVector(id=vec_id, vector=sparse_vec))

# 批量插入，可指定batch_size控制每批次插入数量
TFIDFVector.bulk_create(instances, batch_size=1000)

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