批量写入与批次大小

🎉摘要:本文讲解Chroma向量数据库的单次批次上限查询方法,说明推荐的舒适写入区间,提供可复用的Python分批写入工具代码,介绍numpy数组索引的避坑技巧,帮助开发者规避内存溢出、超时、请求被服务端拒绝、切片维度丢失等问题,高效完成向量数据写入。

Chroma 对单次调用的记录数有上限,用 client.get_max_batch_size() 能查到:

import chromadb

client = chromadb.EphemeralClient()
print(client.get_max_batch_size()) # 5461

这个数不是给你当目标用的。实际经验是每批 100~1000 条比较舒服,再大内存和超时都会出问题(尤其是 HttpClient 模式,payload 太大会被服务端拒绝,默认上限 40 MiB)。

一个通用的分批工具:

from typing import Iterator, Sequence, TypeVar

T = TypeVar("T")

def chunked(seq: Sequence[T], size: int) -> Iterator[Sequence[T]]:
    """把序列按 size 切成一段一段,最后一段可能不够 size
    例如:list(chunked([1, 2, 3, 4, 5], 2))
         [[1, 2], [3, 4], [5]]
    """
    # 步长取 size,起点 0 / size / 2*size ...
    # 末尾不足 size 时自动只切到结尾
    for i in range(0, len(seq), size):
        yield seq[i:i + size]


def add_in_batches(col, ids, embeddings, documents=None, metadatas=None,
                   batch_size=500, verbose=False):
    """分批写入,返回实际写入条数"""
    written = 0
    # ids / embeddings / documents / metadatas 是按下标一一对应的,
    # 所以先切下标、再按同一批下标去取各自的分片,保证一批里各字段条数一致
    for idx_batch in chunked(range(len(ids)), batch_size):
        sl = slice(idx_batch[0], idx_batch[-1] + 1)
        col.add(
            ids=ids[sl],
            embeddings=embeddings[sl] if embeddings else None,
            documents=documents[sl] if documents else None,
            metadatas=metadatas[sl] if metadatas else None,
        )
        written += len(idx_batch)
        if verbose:
            print(f"已写入 {written}/{len(ids)}")
    return written


if __name__ == "__main__":
    import chromadb

    # 20 条数据按 6 条一批写入
    ids = [f"id{i}" for i in range(20)]
    embeddings = [[float(i), float(i % 3), 1.0] for i in range(20)]
    documents = [f"第 {i} 条文档" for i in range(20)]
    metadatas = [{"idx": i, "even": i % 2 == 0} for i in range(20)]

    # 自己给向量,就不用下载默认的嵌入模型了
    col = chromadb.EphemeralClient().get_or_create_collection(
        "batch_demo", embedding_function=None
    )
    print("写入条数:", add_in_batches(
        col, ids, embeddings, documents, metadatas, batch_size=6, verbose=True
    ))
    print("集合里现在有:", col.count(), "条")

运行脚本,输出如下:

已写入 6/20
已写入 12/20
已写入 18/20
已写入 20/20
写入条数: 20
集合里现在有: 20 条

注意:分批的时候千万别用 embeddings[i:i+n] 这种写法去索引 numpy 的 2D 数组再 tolist(),容易把维度搞丢。建议统一转成 List[List[float]] 再切片。例如:

import numpy as np

# 2维数组:3条向量,每条向量长度4
embeddings = np.array([
    [0.1,0.2,0.3,0.4],
    [0.5,0.6,0.7,0.8],
    [0.9,1.0,1.1,1.2]
])
print(embeddings.shape) # (3,4)

使用 embeddings[i:i+n] 进行 numpy 原生切片:

batch = embeddings[0:1]
print(batch.shape) # (1,4) 二维,没问题

# 如果用单索引 embeddings[0],不是切片
single = embeddings[0]
print(single.shape) # (4,) 变成一维数组

很多人代码写着写着混淆区间切片 i:i+n 和单点索引 [i],为了避开这个坑,建议统一使用分片处理。

  

说说我的看法
全部评论()
没有评论
关于
本网站专注于 Java、数据库(MySQL、Oracle)、Linux、软件架构及大数据等多领域技术知识分享。涵盖丰富的原创与精选技术文章,助力技术传播与交流。无论是技术新手渴望入门,还是资深开发者寻求进阶,这里都能为您提供深度见解与实用经验,让复杂编码变得轻松易懂,携手共赴技术提升新高度。如有侵权,请来信告知:hxstrive@outlook.com
其他应用
公众号