大规模数据集下Embeddings存储与分析的最佳实践咨询

1. 分离存储Embedding，关联索引按需加载 #

把embeddings单独存入高效二进制格式（如npy），与原CSV数据通过索引（行号或业务唯一ID）一一对应，避免全量加载到内存：

• 用numpy.save()将embeddings保存为单个大文件，或按固定行数分块存储
• 遍历分析时，采用内存映射方式读取embedding文件，配合Pandas的分块读取功能，每次仅加载部分数据
• 示例代码：

import numpy as np
import pandas as pd

# 保存embeddings
embeddings = np.array(your_embedding_list)
np.save("embeddings_full.npy", embeddings)

# 分块遍历分析
chunk_size = 100000
# 内存映射模式，仅在需要时加载对应块数据
embedding_array = np.load("embeddings_full.npy", mmap_mode="r")

for chunk_idx, df_chunk in enumerate(pd.read_csv("your_data.csv", chunksize=chunk_size)):
    start = chunk_idx * chunk_size
    end = start + chunk_size
    emb_chunk = embedding_array[start:end]
    # 执行当前块的联合分析逻辑
    process_chunk(df_chunk, emb_chunk)

2. 用Parquet格式整合存储 #

Parquet是列式存储格式，对numpy数组（embedding）支持友好，压缩率高且内存占用低：

• 将原CSV数据与embeddings合并后写入Parquet，格式会自动优化embedding列的存储
• 读取时通过分批次迭代，避免全量加载
• 示例代码：

import pandas as pd
import pyarrow.parquet as pq

# 将embedding转为列表存入DataFrame
df["embedding"] = list(embeddings)
# 写入Parquet（snappy压缩平衡速度与体积）
df.to_parquet("data_with_embeddings.parquet", compression="snappy")

# 分批次读取分析
parquet_file = pq.ParquetFile("data_with_embeddings.parquet")
for batch in parquet_file.iter_batches(batch_size=100000):
    df_batch = batch.to_pandas()
    emb_batch = np.array(df_batch["embedding"].tolist())
    # 处理当前批次
    process_batch(df_batch, emb_batch)

3. 用Dask替代Pandas处理超大规模数据 #

Dask支持并行处理超出内存的数据集，自动管理分块与任务调度，适合全量遍历分析：

• 把原CSV转为Dask DataFrame，embedding转为Dask Array，通过索引关联
• 利用map_partitions实现分块并行处理
• 示例代码：

import dask.dataframe as dd
import dask.array as da

# 读取原CSV为Dask DataFrame
ddf = dd.read_csv("your_data.csv")
# 读取embedding为Dask Array，指定分块大小
emb_da = da.from_array(np.load("embeddings_full.npy", mmap_mode="r"), chunks=(100000,))
# 关联embedding列
ddf["embedding"] = emb_da.to_dask_dataframe(columns=["embedding"])

# 定义分块处理逻辑
def process_partition(df_part):
    emb_part = np.array(df_part["embedding"].tolist())
    # 执行分析逻辑，返回结果
    return processed_result

# 触发计算并获取全量结果
final_result = ddf.map_partitions(process_partition).compute()

4. 优化Pandas内的Embedding存储（应急方案） #

如果必须用Pandas，可通过以下方式降低内存占用：

• 将embedding的数值类型从float64转为float32，直接减少一半内存消耗
• 若embedding为稀疏类型，使用pandas.SparseArray存储，仅保留非零值

# 转换为float32
embeddings = embeddings.astype(np.float32)
df["embedding"] = list(embeddings)

# 稀疏embedding存储
from pandas.core.arrays.sparse import SparseArray
sparse_emb = SparseArray(embeddings.ravel(), fill_value=0).reshape(embeddings.shape)
df["embedding"] = list(sparse_emb)

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