如何高效将TSV文件中的数组列分别存入单个NPZ文件？

问题背景 #

有如下格式的TSV数据文件，列依次为：MD5哈希值、目标浮点输出值、需转为np.array的浮点数组、另一需转为np.array的浮点数组：

58f0965a62d62099f5c0771d35dbc218        0.868632614612579       [0.028979932889342308, 0.004080114420503378, 0.03757167607545853]       [-0.006008833646774292, -0.010409083217382431, 0.01565541699528694]
36f7859ce47417470bc28384694f0ac4        0.835115909576416       [0.026130573824048042, -0.00358427781611681, 0.06635218113660812]       [-0.06970945745706558, 0.03816794604063034, 0.03491008281707764]
59f7d617bb662155b0d49ce3f27093ed        0.907200276851654       [0.009903069585561752, -0.009721670299768448, 0.0151780480518937]       [-0.03264783322811127, 0.0035394825972616673, -0.05089104175567627]

原代码通过eval解析数组、存入DataFrame再转numpy数组的方式，在小数据集可行，但面对2亿行+大维度数组时，会触发内存爆炸、速度极慢的问题：

import numpy as np
from tqdm import tqdm
import pandas as pd

lol = []
with open('data.tsv') as fin:
    for line in tqdm(fin):
        md5hash, score, vector1, vector2 = line.strip().split('\t')
        row = {'md5_hash': md5hash, 'score':float(score), 
               'vector1': np.array(eval(vector1)), 
               'vector2': np.array(eval(vector2))
              }
        lol.append(row)
        
df = pd.DataFrame(lol)

training_vector1 = np.array(list(df['vector1']))
np.save('vector1.npz', training_vector1)

training_vector2 = np.array(list(df['vector2']))
np.save('vector1.npz', training_vector2)  # 原代码存在文件名错误，应为vector2.npz

高效解决方案 #

核心思路：避免全量加载到内存，逐行解析并预分配内存，直接写入二进制格式

方法1：预分配内存+逐行解析（无Pandas） #

适用于内存可容纳全量数组的场景：

1. 先统计总行数和固定数组维度
2. 预分配numpy数组内存
3. 用高效字符串解析替代eval
4. 直接填充数组后压缩保存

import numpy as np
from tqdm import tqdm

# 第一步：获取总行数和数组维度（假设所有数组维度一致）
total_lines = 0
vec_dim = None
with open('data.tsv', 'r') as fin:
    for line in fin:
        total_lines +=1
        if not vec_dim:
            _, _, vec1_str, _ = line.strip().split('\t')
            vec_dim = len(vec1_str.strip('[]').split(','))
            print(f"数组维度：{vec_dim}，总行数：{total_lines}")

# 第二步：预分配内存（用float32节省空间，精度要求高则用float64）
vector1_arr = np.zeros((total_lines, vec_dim), dtype=np.float32)
vector2_arr = np.zeros((total_lines, vec_dim), dtype=np.float32)

# 第三步：逐行解析填充数组
with open('data.tsv', 'r') as fin:
    for idx, line in enumerate(tqdm(fin, total=total_lines)):
        parts = line.strip().split('\t')
        # 解析vector1：去掉[]后直接转numpy数组
        vec1_str = parts[2].strip('[]')
        vector1_arr[idx] = np.fromstring(vec1_str, sep=',', dtype=np.float32)
        
        # 解析vector2
        vec2_str = parts[3].strip('[]')
        vector2_arr[idx] = np.fromstring(vec2_str, sep=',', dtype=np.float32)

# 第四步：压缩保存为NPZ
np.savez_compressed('vectors.npz', vector1=vector1_arr, vector2=vector2_arr)
# 如需分开保存：
# np.savez_compressed('vector1.npz', data=vector1_arr)
# np.savez_compressed('vector2.npz', data=vector2_arr)

方法2：分块处理（内存不足时） #

当内存无法容纳全量数组时，拆分数据为小块处理，最后可选合并：

import numpy as np
from tqdm import tqdm

chunk_size = 1_000_000  # 每块100万行，可根据内存调整
vec_dim = None
chunk_idx = 0

with open('data.tsv', 'r') as fin:
    while True:
        chunk_lines = []
        # 读取单块数据
        for _ in range(chunk_size):
            line = fin.readline()
            if not line:
                break
            chunk_lines.append(line)
        if not chunk_lines:
            break
        
        # 确定数组维度
        if not vec_dim:
            _, _, vec1_str, _ = chunk_lines[0].strip().split('\t')
            vec_dim = len(vec1_str.strip('[]').split(','))
        
        # 初始化块数组
        vec1_chunk = np.zeros((len(chunk_lines), vec_dim), dtype=np.float32)
        vec2_chunk = np.zeros((len(chunk_lines), vec_dim), dtype=np.float32)
        
        # 填充块数组
        for idx, line in enumerate(chunk_lines):
            parts = line.strip().split('\t')
            vec1_str = parts[2].strip('[]')
            vec1_chunk[idx] = np.fromstring(vec1_str, sep=',', dtype=np.float32)
            
            vec2_str = parts[3].strip('[]')
            vec2_chunk[idx] = np.fromstring(vec2_str, sep=',', dtype=np.float32)
        
        # 保存当前块
        np.savez_compressed(f'vector1_chunk_{chunk_idx}.npz', data=vec1_chunk)
        np.savez_compressed(f'vector2_chunk_{chunk_idx}.npz', data=vec2_chunk)
        chunk_idx +=1

# 可选：合并所有块为全量数组
vector1_arr = np.concatenate([np.load(f'vector1_chunk_{i}.npz')['data'] for i in range(chunk_idx)])
vector2_arr = np.concatenate([np.load(f'vector2_chunk_{i}.npz')['data'] for i in range(chunk_idx)])
np.savez_compressed('vectors_merged.npz', vector1=vector1_arr, vector2=vector2_arr)

关键优化点 #

• 替换eval：用np.fromstring直接解析字符串，速度更快且避免安全风险
• 预分配内存：避免动态列表扩容带来的内存开销和碎片
• 低精度 dtype：精度允许时用float32，内存占用减半
• 压缩存储：用savez_compressed替代savez，大幅节省磁盘空间
• 分块拆分：内存不足时拆分数据，避免OOM（内存溢出）

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