使用np.unique处理大型二维数组时遭遇内存错误的解决方法

首先注意到你代码中的关键笔误：你要提取唯一行，应该使用axis=0而非axis=1。axis=1是针对列去重，而你的数组有243327132行、仅3列，执行np.unique(arr, axis=1)会尝试将每一列（长度243327132的一维数组）转换为结构化数组的字段，这必然触发内存爆炸——这也是你报错的核心原因。

以下是针对提取唯一行（需保留计数和索引）的低内存解决方案：

方案1：将行转换为哈希值（内存友好的快速方法） #

由于你的每行是3个uint8值，可直接将其编码为一个uint32整数（3×8=24位，完全适配uint32的32位空间，无哈希碰撞风险），将二维数组转换为一维数组后再用np.unique处理，内存占用大幅降低：

import numpy as np

# 加载数组（用memmap避免全量加载到内存）
arr = np.memmap("memarr.memmap", mode='r', dtype=np.uint8, shape=(243327132, 3))

# 将每行的3个uint8编码为一个uint32
hash_vals = arr[:,0].astype(np.uint32) << 16 | arr[:,1].astype(np.uint32) << 8 | arr[:,2].astype(np.uint32)

# 提取唯一哈希值、计数、首次出现索引
unique_hashes, counts, indices = np.unique(hash_vals, return_counts=True, return_index=True)

# 从哈希值还原唯一行
unique_rows = np.column_stack([
    (unique_hashes >> 16).astype(np.uint8),
    (unique_hashes >> 8).astype(np.uint8),
    unique_hashes.astype(np.uint8)
])

此方法仅需将原数组的3列转换为1列uint32（内存占用约970MB），远低于np.unique(axis=0)的结构化数组内存开销，16GB内存完全可以承载。

方案2：分块哈希统计（极端内存紧张时使用） #

如果方案1仍有内存压力，可通过分块处理，用字典逐步统计哈希值的计数和首次索引，全程仅加载小块数据到内存：

import numpy as np

arr = np.memmap("memarr.memmap", mode='r', dtype=np.uint8, shape=(243327132, 3))

hash_counts = {}
hash_first_idx = {}
chunk_size = 1_000_000  # 每次处理100万行，可根据内存调整
n_chunks = (arr.shape[0] + chunk_size - 1) // chunk_size

for i in range(n_chunks):
    start = i * chunk_size
    end = min((i+1)*chunk_size, arr.shape[0])
    chunk = arr[start:end]
    
    # 编码当前块的行哈希值
    chunk_hash = chunk[:,0].astype(np.uint32) << 16 | chunk[:,1].astype(np.uint32) << 8 | chunk[:,2].astype(np.uint32)
    
    # 统计块内唯一哈希值及计数
    unique_chunk, cnt_chunk = np.unique(chunk_hash, return_counts=True)
    
    # 更新全局统计
    for h, cnt in zip(unique_chunk, cnt_chunk):
        if h in hash_counts:
            hash_counts[h] += cnt
        else:
            hash_counts[h] = cnt
            # 记录首次出现的全局索引
            hash_first_idx[h] = start + np.where(chunk_hash == h)[0][0]

# 转换为Numpy数组格式
unique_hashes = np.array(list(hash_counts.keys()), dtype=np.uint32)
counts = np.array(list(hash_counts.values()), dtype=np.int64)
indices = np.array([hash_first_idx[h] for h in unique_hashes], dtype=np.int64)

# 还原唯一行
unique_rows = np.column_stack([
    (unique_hashes >> 16).astype(np.uint8),
    (unique_hashes >> 8).astype(np.uint8),
    unique_hashes.astype(np.uint8)
])

此方法每次仅加载约3MB数据（100万行×3字节），内存占用可忽略不计，适合极端内存受限场景。

关键注意事项 #

• 始终用mode='r'的memmap加载数组，避免将全量数据拷贝到内存；
• 若确实需要对列去重（axis=1），由于你的数组仅3列，直接手动对比即可，无需调用np.unique。

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