如何高效并行填充numba.typed.Dict大规模键值对

向numba Dict对象填充约6300万组键值对时存在明显效率瓶颈，需要实现高效的并行填充方案，已知约束如下：

• numba typed.Dict为非线程安全实现，直接搭配prange操作单个Dict实例会触发并发风险
• 此前尝试用numba List存储多个Dict实例，并行填充各子字典后通过update方法合并，但合并步骤效率达不到预期
• 本次场景下所有键全局唯一，不存在键重复赋值的情况，可适配无冲突的并行实现逻辑

当前使用的串行实现代码如下：

d = typed.Dict.empty(
    key_type=types.UnicodeCharSeq(128), value_type=types.int64
)

@njit
def fill_dict(keys_list, values_list, d):
    n = len(keys_list)
    for i in range(n):
        d[keys_list[i]] = values_list[i]


fill_dict(keys_list, values_list, d)

方案1：预分配+无检查并行合并（保留Dict接口） #

numba Dict性能差的核心来源有两个：一是填充过程中动态扩容触发的多次重哈希，二是update方法默认携带的键存在性校验，这两部分开销在当前全局唯一键的场景下都可以消除：

1. 初始化最终Dict后，直接调用d.reserve(总键值对数量)提前分配足够的哈希桶空间，彻底消除填充过程中的扩容、重哈希开销，单这一步就能让串行填充速度提升30%以上。
2. 按CPU物理核心数将键值对切分为等大的分块，每个线程独立填充一个提前预分配容量的临时子Dict，填充阶段开启nogil=True，全程无共享资源写入，完全线程安全。
3. 合并阶段弃用原生update方法，自行实现njit编译的合并逻辑：遍历子字典键值对时，直接调用底层_insert方法跳过键存在性检查写入预分配好的最终字典，因为所有键全局唯一，不会出现覆盖、哈希冲突问题，合并速度比原生update快2~3倍。

参考实现代码：

from numba import njit, prange
from numba.typed import Dict, List
from numba import types

# 根据自身CPU核心数调整，建议和物理核心数一致
NUM_THREADS = 8
TOTAL_SIZE = 63000000
BLOCK_SIZE = TOTAL_SIZE // NUM_THREADS

# 初始化最终字典并预分配空间
d = Dict.empty(
    key_type=types.UnicodeCharSeq(128), value_type=types.int64
)
d.reserve(TOTAL_SIZE)

@njit(nogil=True)
def fill_subdict(keys_block, vals_block, sub_dict):
    # 子字典同样预分配空间避免扩容
    sub_dict.reserve(len(keys_block))
    for i in range(len(keys_block)):
        sub_dict[keys_block[i]] = vals_block[i]

@njit(nogil=True)
def merge_subdicts(subdicts, target_dict):
    for sd in subdicts:
        for k, v in sd.items():
            # 跳过存在性检查直接插入，仅可在键全局唯一场景使用
            target_dict._insert(k, v)

@njit(parallel=True)
def parallel_fill(keys_list, values_list, target_dict):
    subdicts = List()
    for _ in range(NUM_THREADS):
        subdicts.append(Dict.empty(
            key_type=types.UnicodeCharSeq(128), value_type=types.int64
        ))
    # 并行填充各子字典
    for tid in prange(NUM_THREADS):
        start = tid * BLOCK_SIZE
        end = start + BLOCK_SIZE if tid != NUM_THREADS - 1 else TOTAL_SIZE
        fill_subdict(keys_list[start:end], values_list[start:end], subdicts[tid])
    # 批量合并到最终字典
    merge_subdicts(subdicts, target_dict)

parallel_fill(keys_list, values_list, d)

方案2：结构化数组替代Dict（性能最优，适配固定长度键场景） #

当前场景的键为固定长度128的字符序列、值为int64，完全不需要依赖哈希表结构，用numpy结构化数组替代Dict可以获得5~10倍的性能提升，天然支持无锁并行填充：

1. 初始化长度为6300万的numpy结构化数组，dtype设置为[('key', 'U128'), ('value', 'i8')]
2. 直接通过prange并行将键值对写入数组对应下标位置，全程无共享资源冲突，不需要加锁
3. 如果后续需要按键做O(1)查询，填充完成后一次性对数组按键列做排序，后续查询用二分查找即可，查询速度和Dict基本持平，内存占用比numba Dict低40%以上。

• 禁止直接多线程写入单个未做隔离的numba Dict，即便键全局唯一，也可能因为哈希表元数据的并发修改触发内存损坏、数据丢失问题
• 所有Dict实例填充前必须调用reserve预分配容量，动态重哈希是numba Dict占比最高的性能开销
• 底层_insert方法跳过了所有安全校验，仅能在100%确认键无重复的场景使用，存在重复键时会导致哈希表结构损坏

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