如何高效合并多个 NumPy .npz 文件

import numpy as np
import os

def merge_npz_files(file_list, output_filename='merged_data.npz'):
    """
    合并多个 .npz 文件到一个新的 .npz 文件中。

    参数:
    file_list (list): 包含所有待合并 .npz 文件路径的列表。
    output_filename (str): 合并后 .npz 文件的输出路径和名称。
    """
    if not file_list:
        print("文件列表为空，无法合并。")
        return

    # 1. 加载所有 .npz 文件
    # np.load 返回一个 NpzFile 对象，可以像字典一样访问其内部数组
    data_all = [np.load(fname) for fname in file_list]

    merged_data = {}

    # 2. 遍历第一个文件的所有键（假设所有文件有相同的键结构）
    # 对于每个键，收集所有文件中的对应数组并进行拼接
    for key in data_all[0].keys():
        # 收集所有文件中对应当前键的数组
        # 使用 d[key] 访问 NpzFile 对象中的数组
        arrays_to_concatenate = [d[key] for d in data_all]

        # 3. 使用 np.concatenate 拼接数组
        # 默认沿轴0拼接，即在第一个维度上堆叠
        try:
            merged_data[key] = np.concatenate(arrays_to_concatenate, axis=0)
            print(f"键 '{key}' 已成功合并。新形状: {merged_data[key].shape}")
        except ValueError as e:
            print(f"警告: 键 '{key}' 的数组无法合并。错误: {e}")
            print(f"跳过键 '{key}' 的合并。")
            # 如果无法合并，可以选择跳过或采取其他处理
            # 例如，如果数组形状不兼容，可能需要进行 reshape 或 padding

    # 4. 保存合并后的数据到新的 .npz 文件
    if merged_data:
        np.savez_compressed(output_filename, **merged_data)
        print(f"所有数据已成功合并并保存到 '{output_filename}'。")
    else:
        print("没有可合并的数据。")

# --- 示例用法 ---
if __name__ == "__main__":
    # 假设我们有以下文件：
    # 创建一些示例 .npz 文件
    if not os.path.exists("temp_npz_files"):
        os.makedirs("temp_npz_files")

    for i in range(3):
        arr_features = np.random.rand(10 + i, 5) # 模拟不同数量的样本
        arr_labels = np.random.randint(0, 2, size=(10 + i,))
        data_dict = {'features': arr_features, 'labels': arr_labels}
        np.savez_compressed(f'temp_npz_files/data_{i}.npz', **data_dict)
        print(f"创建文件: temp_npz_files/data_{i}.npz (features shape: {arr_features.shape}, labels shape: {arr_labels.shape})")

    # 获取所有 .npz 文件名
    file_names = [os.path.join("temp_npz_files", f) for f in os.listdir("temp_npz_files") if f.endswith('.npz')]
    file_names.sort() # 确保顺序一致

    # 执行合并
    merge_npz_files(file_names, 'merged_output.npz')

    # 验证合并结果
    print("\n验证合并结果:")
    merged_file = np.load('merged_output.npz')
    for key in merged_file.keys():
        print(f"合并文件中键 '{key}' 的形状: {merged_file[key].shape}")

    # 清理临时文件
    import shutil
    shutil.rmtree("temp_npz_files")
    os.remove("merged_output.npz")
    print("\n清理完成。")