Python下如何高效存储10000×10000大型数据集至NetCDF/HDF5？

我有一个10000×10000的数据集需要存储，原本考虑用HDF5，后来转向Xarray，但电脑性能有限，没法全量加载数据完成转换。试过逐行读取的方式，但发现NetCDF和Xarray不支持这种写入方式，目前只能把每行数据存成一个新group，感觉这不是最优方案。

数据集是组间对比数据，对角线值都是1，理论上可以只存一半数据保留全部信息，但查了Xarray资料后暂时放弃了这个想法。

当前使用的代码如下：

label_data = pd.read_excel('./my10000_label.xlsx')
file_local_path = Path("./my10000_data.txt")
hf = h5py.File("./savedata.h5")


local_idx = 0
with file_local_path.open() as my_data:
    while True:
        line = my_data.readline()
        if not line:
            break

        line_array = np.fromstring(line, dtype=np.float64, sep=",")
        label = label_data.loc[local_idx]

        xdata = xarray.DataArray(line_array, dims=["lbl"], coords={"lbl":label_data["label"]})
        hf.create_dataset(label, data=xdata)
        local_idx += 1

hf.close()       

想请教有没有更优的浮点数据压缩数组存储方法？了解过numpy.memmap可以分片加载数据，但好像没法用类似方式写入。

1. HDF5直接创建大数组逐行写入（带压缩） #

不用Xarray中转，直接用h5py预设对应形状的数据集，开启压缩后逐行写入，既能避免大量小group的冗余，又能利用HDF5的压缩特性减少存储空间。推荐用gzip压缩，可通过compression_opts调整压缩级别（1-9，级别越高压缩率越高但速度越慢）。

import h5py
import numpy as np
import pandas as pd
from pathlib import Path

label_data = pd.read_excel('./my10000_label.xlsx')
n_rows = len(label_data)
file_local_path = Path("./my10000_data.txt")

# 创建带压缩的HDF5数据集
with h5py.File("./savedata_opt.h5", "w") as hf:
    dset = hf.create_dataset(
        "comparison_data",
        shape=(n_rows, n_rows),
        dtype=np.float64,
        compression="gzip",
        compression_opts=5
    )
    
    # 逐行读取写入
    with file_local_path.open() as my_data:
        for idx, line in enumerate(my_data):
            if not line.strip():
                continue
            line_array = np.fromstring(line, dtype=np.float64, sep=",")
            dset[idx, :] = line_array

# 后续用Xarray读取该文件
import xarray as xr
ds = xr.open_dataset("./savedata_opt.h5", engine="h5netcdf")

2. 利用对称矩阵特性压缩存储（节省50%空间） #

你的数据集是对称矩阵（对角线为1，a_ij = a_ji），可以只存储下三角（或上三角）部分（含对角线），存储量直接从1亿个浮点降到约5000万个，大幅节省空间。

import h5py
import numpy as np
import pandas as pd
from pathlib import Path

label_data = pd.read_excel('./my10000_label.xlsx')
n_rows = len(label_data)
file_local_path = Path("./my10000_data.txt")

# 计算下三角元素总数：n*(n+1)/2
tri_size = n_rows * (n_rows + 1) // 2

with h5py.File("./savedata_tri.h5", "w") as hf:
    dset = hf.create_dataset(
        "tri_comparison",
        shape=(tri_size,),
        dtype=np.float64,
        compression="gzip",
        compression_opts=5
    )
    
    offset = 0
    with file_local_path.open() as my_data:
        for row_idx, line in enumerate(my_data):
            if not line.strip():
                continue
            line_array = np.fromstring(line, dtype=np.float64, sep=",")
            # 取当前行的前row_idx+1个元素（下三角部分）
            tri_part = line_array[:row_idx+1]
            dset[offset:offset+len(tri_part)] = tri_part
            offset += len(tri_part)

# 读取时恢复对称矩阵
def restore_symmetric_matrix(tri_data, n):
    mat = np.zeros((n, n), dtype=np.float64)
    idx = 0
    for i in range(n):
        mat[i, :i+1] = tri_data[idx:idx+i+1]
        mat[:i+1, i] = tri_data[idx:idx+i+1]
        idx += i+1
    return mat

# 恢复并转为Xarray对象
with h5py.File("./savedata_tri.h5", "r") as hf:
    tri_data = hf["tri_comparison"][:]
    full_mat = restore_symmetric_matrix(tri_data, n_rows)

xr_data = xr.DataArray(
    full_mat,
    dims=["row_lbl", "col_lbl"],
    coords={"row_lbl": label_data["label"], "col_lbl": label_data["label"]}
)

3. NetCDF分块写入（结合Xarray） #

Xarray支持分块（chunked）存储，先创建空的分块数据集，再逐块写入数据，避免全量加载内存。分块大小可根据你的内存情况调整，比如设为100行一块。

import xarray as xr
import numpy as np
import pandas as pd
from pathlib import Path

label_data = pd.read_excel('./my10000_label.xlsx')
labels = label_data["label"].values
n_rows = len(labels)
file_local_path = Path("./my10000_data.txt")

# 创建分块数据集，chunk_size根据内存调整
chunk_size = 100
ds = xr.Dataset(
    {
        "comparison": xr.DataArray(
            data=np.zeros((n_rows, n_rows), dtype=np.float64),
            dims=["row", "col"],
            coords={"row": labels, "col": labels},
            chunks={"row": chunk_size, "col": chunk_size}
        )
    }
)

# 逐块读取写入
with file_local_path.open() as my_data:
    for start_idx in range(0, n_rows, chunk_size):
        end_idx = min(start_idx + chunk_size, n_rows)
        chunk_data = []
        for _ in range(start_idx, end_idx):
            line = my_data.readline()
            if not line:
                break
            line_array = np.fromstring(line, dtype=np.float64, sep=",")
            chunk_data.append(line_array)
        chunk_data = np.array(chunk_data)
        ds["comparison"][start_idx:end_idx, :] = chunk_data

# 保存为带压缩的NetCDF文件
ds.to_netcdf(
    "./savedata_netcdf.nc",
    encoding={"comparison": {"compression": "gzip", "compression_opts": 5}}
)

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