向Chroma DB导入预向量化数据时遇bad allocation错误的技术求助
2026-7-17
解决Chroma向量数据库导入预向量化数据时的内存分配错误
问题描述
本地磁盘已有一个包含50万行数据的Chroma集合，在Jupyter Notebook中重新加载该集合并添加更多预向量化文本数据时，触发了bad allocation内存分配错误。推测是客户端无法将大文件加载到内存导致的，尝试过每批次导入后和单个10万行文件导入后执行persist()，但问题仍存在。需要优化导入和持久化的代码设计来避免该错误。
现有代码
客户端初始化及获取现有集合
client_wiki_ja = chromadb.Client(Settings( chroma_db_impl="duckdb+parquet", persist_directory= database_directory )) cohere_ef = embedding_functions.CohereEmbeddingFunction( api_key=api_key_embedding_cohere, model_name="multilingual-22-12") collection_wiki_ja = client_wiki_ja.get_collection(name=collection_name,embedding_function=cohere_ef) # 从现有集合获取对象，不存在则抛出异常
数据导入函数
def load_files(batch_size=500,max_files_load=2,startFile_index=4): # startFile_index从0开始（例如train-00003-of-00035-99c97d6c1602f565.parquet对应startFile_index=3） directory = r"C:\Users\liewg\Desktop\p005-butumon-go\wikipedia-22-12-ja-embeddings\data" # 维基数据存放路径 file_list = os.listdir(directory) counter = 0 # 已导入文件计数 for idx, filename in enumerate(file_list): # 遍历多个文件 if idx < startFile_index: continue # 跳过前几个文件 file_path = os.path.join(directory, filename) # 单个parquet文件路径 print(file_path) df = pd.read_parquet(file_path, engine='pyarrow') # df = df.head(3) df['emb'] = df['emb'].apply(lambda x: x.tolist()) print("finished convert to list...") # print("number of files: ",df.count()) rows = df# .head(10000) for i in tqdm(range(0, len(rows), batch_size)): # 单文件内分批次导入 i_end = min(len(rows), i+batch_size) # 计算批次结束位置 batch = rows[i:i_end] embeds_batch = [ emb for emb in batch['emb']] text_batch = [ text for text in batch['text']] metadatas_batch = [ {"source": title} for title in batch['title']] ids_batch = [ str(id) for id in batch['id']] collection_wiki_ja.add( embeddings= embeds_batch, documents= text_batch, metadatas=metadatas_batch, ids = ids_batch ) print("batch:", i) print("persist starts") client_wiki_ja.persist() # 每批次导入后持久化 print("persist ends") # if i == 2: # break print("completed") counter += 1 if counter == max_files_load : break
错误信息
Error Traceback (most recent call last) Cell In[9], line 1 ----> 1 load_files() Cell In[8], line 42, in load_files(batch_size, max_files_load, startFile_index) 40 print("batch:", i) 41 print("persist starts") ---> 42 client_wiki_ja.persist() # for after initial 43 print("persist ends") 45 # if i == 2: 46 # break File ~\AppData\Local\Packages\PythonSoftwareFoundation.Python.3.11_qbz5n2kfra8p0\LocalCache\local-packages\Python311\site-packages\chromadb\api\local.py:544, in LocalAPI.persist(self) 536 @override 537 def persist(self) -> bool: 538 """Persist the database to disk. 539 540 Returns: 541 True if the database was persisted successfully 542 543 """ --> 544 self._db.persist() 545 return True File ~\AppData\Local\Packages\PythonSoftwareFoundation.Python.3.11_qbz5n2kfra8p0\LocalCache\local-packages\Python311\site-packages\chromadb\db\duckdb.py:473, in PersistentDuckDB.persist(self) 470 if self._conn.query("SELECT COUNT() FROM embeddings") == 0: 471 return --> 473 self._conn.execute( 474 f"""``` 475 COPY 476 (SELECT * FROM embeddings) 477 TO '{self._save_folder}/chroma-embeddings.parquet' 478 (FORMAT PARQUET); 479 """ 480 ) 482 self._conn.execute( 483 f""" 484 COPY (...) 488 """ 489 ) Error: Invalid Error: bad allocation
优化后的代码设计方案
核心优化思路
- 降低
persist()调用频率：每批次调用会重复加载全量数据集写入磁盘，改为每导入完1个或多个文件后执行一次，减少内存重复占用。 - 分块读取文件：避免一次性加载整个Parquet文件到内存，用分块读取逐步处理。
- 限制DuckDB内存：通过配置限制DuckDB的内存使用上限，防止其耗尽系统内存。
- 主动清理内存：处理完数据块后主动删除临时变量并触发垃圾回收，释放内存。
优化后的代码
1. 客户端初始化（增加内存配置）
from chromadb.config import Settings import chromadb from chromadb.utils import embedding_functions import pandas as pd import os from tqdm import tqdm import gc # 初始化客户端时配置DuckDB内存上限，根据系统内存调整（例如设为8GB） client_wiki_ja = chromadb.Client(Settings( chroma_db_impl="duckdb+parquet", persist_directory=database_directory, duckdb_settings={"memory_limit": "8GB"} )) cohere_ef = embedding_functions.CohereEmbeddingFunction( api_key=api_key_embedding_cohere, model_name="multilingual-22-12") collection_wiki_ja = client_wiki_ja.get_collection(name=collection_name, embedding_function=cohere_ef)
2. 优化后的导入函数
def load_files(batch_size=1000, max_files_load=2, startFile_index=4, persist_interval=1): directory = r"C:\Users\liewg\Desktop\p005-butumon-go\wikipedia-22-12-ja-embeddings\data" file_list = os.listdir(directory) counter = 0 for idx, filename in enumerate(file_list): if idx < startFile_index: continue file_path = os.path.join(directory, filename) print(f"Processing file: {file_path}") # 分块读取Parquet文件，避免一次性加载全量数据 chunk_iter = pd.read_parquet(file_path, engine='pyarrow', chunksize=batch_size*5) for chunk_df in tqdm(chunk_iter): # 转换emb为列表 chunk_df['emb'] = chunk_df['emb'].apply(lambda x: x.tolist()) # 再分小批次导入 for i in range(0, len(chunk_df), batch_size): i_end = min(len(chunk_df), i + batch_size) batch = chunk_df[i:i_end] embeds_batch = batch['emb'].tolist() text_batch = batch['text'].tolist() metadatas_batch = [{"source": title} for title in batch['title']] ids_batch = batch['id'].astype(str).tolist() # 批量添加数据 collection_wiki_ja.add( embeddings=embeds_batch, documents=text_batch, metadatas=metadatas_batch, ids=ids_batch ) # 处理完一个大chunk后清理内存 del chunk_df gc.collect() print(f"Completed processing file: {filename}") # 每导入persist_interval个文件后执行一次持久化 counter += 1 if counter % persist_interval == 0: print("Starting persist...") client_wiki_ja.persist() print("Persist finished") if counter == max_files_load: break # 最后执行一次持久化，确保所有数据写入磁盘 print("Final persist...") client_wiki_ja.persist() print("All data imported and persisted")
额外建议
- 调整batch_size：根据系统内存灵活调整，内存充足可增大到2000，内存紧张则缩小到300。
- 使用64位Python：32位Python有内存上限限制，必须使用64位版本处理大数据集。
- 监控内存占用：导入过程中用任务管理器监控内存使用，及时调整分块大小和批次参数。
- 关闭Jupyter自动保存：暂时关闭Jupyter的自动保存功能，减少额外内存消耗。
内容的提问来源于stack exchange，提问作者Mike
相关产品推荐
相关产品推荐
