将大尺寸数据集高效加载至Python字典的内存优化问题
2026-6-27
问题
我尝试将一个约560MB的大型数据集加载到Python字典中，用于生成3D图形，但遇到内存不足问题，进程被直接"Killed"。之后我添加了分块读取并定期将字典写入JSON文件的逻辑，想避免内存耗尽，但处理到约4.00M/558.0M进度时还是被终止了。
我想知道：
- 为什么仅560MB的文件，剔除无用列并转成字典后会占用数GB内存？
- 有没有更高效的方法，能得到可高效提取坐标集及其对应值的数据对象？
现有代码与示例数据
处理代码
import json import logging import os import pandas as pd from tqdm import tqdm def create_grid_dict(file_path, chunk_size=500000): """ :param file_path: Path to a grid file. :param chunk_size: Number of lines to process before dumping into json :return: Dictionary object containing the gist grid data with as index the voxel number and as values the x, y and z coordinates, and the value """ # Read the data from the file with open(file_path, 'r') as file: # Read the first line header = file.readline().strip() header2 = file.readline().strip() # Log the header logging.info(header) columns = header2.split(' ') # Get the file size file_size = os.path.getsize(file_path) output_file = 'datasets/cache.json' # Check if the output file already exists if os.path.exists(output_file): with open(output_file, 'r') as f: grid_dict = json.load(f) return grid_dict else: # Create an empty dictionary to store the grid data grid_dict = {} logging.info(f"Reading file size {file_size} in chunks of {chunk_size} lines.") # Read the file in chunks with tqdm(total=file_size, unit='B', unit_scale=True, desc="Processing") as pbar: for chunk in pd.read_csv(file_path, delim_whitespace=True, skiprows=2, names=columns, chunksize=chunk_size): # Filter out the columns you need chunk = chunk[['voxel', 'xcoord', 'ycoord', 'zcoord', 'val1', 'val2']] # Iterate through each row in the chunk for index, row in chunk.iterrows(): voxel = row['voxel'] # Store the values in the dictionary grid_dict[voxel] = { 'xcoord': row['xcoord'], 'ycoord': row['ycoord'], 'zcoord': row['zcoord'], 'val': row['val1'] + 2 * row['val2'] } pbar.update(chunk_size) # Write the grid dictionary to the output file after processing each chunk with open(output_file, 'w') as f: json.dump(grid_dict, f) return grid_dict
示例数据集
voxel xcoord ycoord zcoord val1 val2 1 0.1 0.2 0.3 10 5 2 0.2 0.3 0.4 8 4 3 0.3 0.4 0.5 12 6 4 0.4 0.5 0.6 15 7 5 0.5 0.6 0.7 9 3 6 0.6 0.7 0.8 11 5 7 0.7 0.8 0.9 13 6 8 0.8 0.9 1.0 14 7 9 0.9 1.0 1.1 16 8 10 1.0 1.1 1.2 18 9
分析与解决方案
一、内存暴涨的原因
- Python字典的内存开销极高：每个键值对都需要存储哈希表结构、对象引用等额外信息，嵌套小字典时开销会进一步放大。560MB的文本数据转成Python对象后，内存占用通常会膨胀5-20倍。
- 分块逻辑存在误区：你虽然分块读取数据，但
grid_dict始终在内存中累积所有记录，写入JSON只是备份，内存中的字典并没有被清空，随着处理的块数增加，内存占用只会持续上升。 - Pandas的额外消耗：DataFrame本身就有不小的内存开销，加上
iterrows()逐行处理时会将每行转为Series对象，又增加了一层内存消耗。
二、高效解决方案
1. 改用HDF5列式存储（推荐）
HDF5支持增量写入、随机读取，无需加载全量数据到内存，且存储空间远小于JSON：
import pandas as pd import os def create_hdf_cache(file_path, output_file='datasets/cache.h5'): if os.path.exists(output_file): # 返回只读的HDF5数据接口，不会加载全量数据 return pd.read_hdf(output_file, key='voxels', mode='r') # 读取表头 with open(file_path, 'r') as f: f.readline() # 跳过第一行注释表头 columns = f.readline().strip().split() # 分块处理并写入HDF5 chunk_iter = pd.read_csv( file_path, delim_whitespace=True, skiprows=2, names=columns, chunksize=500000 ) for i, chunk in enumerate(chunk_iter): # 计算目标值并保留所需列 chunk['val'] = chunk['val1'] + 2 * chunk['val2'] chunk = chunk[['voxel', 'xcoord', 'ycoord', 'zcoord', 'val']] # 第一次写入用'w'模式，后续用追加模式 mode = 'w' if i == 0 else 'a' chunk.to_hdf(output_file, key='voxels', mode=mode, format='table') return pd.read_hdf(output_file, key='voxels', mode='r') # 使用示例 voxel_data = create_hdf_cache('your_dataset.txt') # 按需提取单个voxel数据 target_row = voxel_data[voxel_data['voxel'] == 100] x, y, z, val = target_row[['xcoord', 'ycoord', 'zcoord', 'val']].values[0] # 批量提取所有坐标和值数组（按需加载，不会占满内存） coords = voxel_data[['xcoord', 'ycoord', 'zcoord']].values vals = voxel_data['val'].values
2. 用SQLite数据库存储
适合需要复杂查询的场景，同样支持增量写入和按需读取：
import sqlite3 import pandas as pd import os def create_sqlite_cache(file_path, output_file='datasets/cache.db'): if os.path.exists(output_file): conn = sqlite3.connect(output_file) return pd.read_sql('SELECT * FROM voxels', conn, index_col='voxel') conn = sqlite3.connect(output_file) # 读取表头 with open(file_path, 'r') as f: f.readline() columns = f.readline().strip().split() chunk_iter = pd.read_csv( file_path, delim_whitespace=True, skiprows=2, names=columns, chunksize=500000 ) for chunk in chunk_iter: chunk['val'] = chunk['val1'] + 2 * chunk['val2'] chunk = chunk[['voxel', 'xcoord', 'ycoord', 'zcoord', 'val']] # 追加写入数据库 chunk.to_sql('voxels', conn, if_exists='append', index=False) # 创建索引加速voxel查询 conn.execute('CREATE INDEX idx_voxel ON voxels(voxel)') conn.close() conn = sqlite3.connect(output_file) return pd.read_sql('SELECT * FROM voxels', conn, index_col='voxel')
3. 优化字典存储逻辑（应急方案）
如果必须用字典，需避免内存累积，采用逐行JSON写入（读取时按需加载）：
import json import os import pandas as pd from tqdm import tqdm def create_grid_dict(file_path, chunk_size=500000): output_file = 'datasets/cache.json' if os.path.exists(output_file): return output_file with open(file_path, 'r') as f: f.readline() columns = f.readline().strip().split() file_size = os.path.getsize(file_path) # 逐行写入JSON对象，避免内存累积 with open(output_file, 'w') as out_f, tqdm(total=file_size, unit='B', unit_scale=True, desc="Processing") as pbar: for chunk in pd.read_csv( file_path, delim_whitespace=True, skiprows=2, names=columns, chunksize=chunk_size ): chunk['val'] = chunk['val1'] + 2 * chunk['val2'] chunk = chunk[['voxel', 'xcoord', 'ycoord', 'zcoord', 'val']] for _, row in chunk.iterrows(): json.dump({int(row['voxel']): { 'xcoord': row['xcoord'], 'ycoord': row['ycoord'], 'zcoord': row['zcoord'], 'val': row['val'] }}, out_f) out_f.write('\n') pbar.update(chunk_size) return output_file # 按需读取指定voxel数据 def get_voxel_data(cache_file, voxel_id): with open(cache_file, 'r') as f: for line in f: data = json.loads(line) if voxel_id in data: return data[voxel_id] return None
内容的提问来源于stack exchange，提问作者Zyzyx
相关产品推荐
相关产品推荐
