求高性能Python/Numpy/Numba归约算法：按entityId求最值

方案1：Numpy原生分组聚合（无额外依赖，最快之一） #

利用Numpy的排序和reduceat实现高效分组求极值，完全基于C底层实现，避免Python循环开销。

步骤：

• 提取实体ID和对应的值
• 按实体ID排序，让同一实体的记录连续
• 找到每个实体的起始索引，用reduceat批量计算极值

代码示例（求最大值）：

import numpy as np

# 提取实体ID（取第一个元素）和值
entity_ids = iarr['e'][:, 0].astype(np.uint32)
values = iarr['r'].astype(np.float32)

# 按实体ID排序，获取排序索引
sorted_indices = np.argsort(entity_ids)
sorted_ids = entity_ids[sorted_indices]
sorted_vals = values[sorted_indices]

# 找到每个唯一实体的起始位置（第一个出现的索引）
unique_positions = np.where(np.diff(sorted_ids, prepend=-1) != 0)[0]

# 用reduceat计算每个分组的最大值
max_vals = np.reduceat(sorted_vals, unique_positions, op=np.max)

# 构造结果数组
oarr = np.array(list(zip(sorted_ids[unique_positions], max_vals)), 
                dtype=[('e', '<u4'), ('r', '<f4')])

优势：

• 完全依赖Numpy原生函数，无额外编译或线程安全问题
• 时间复杂度主要由排序决定（O(N log N)），百万级数据处理效率极高
• 内存使用可控，无需额外字典结构

方案2：Numba并行无锁数组归约（线程安全，接近C速度） #

放弃线程不安全的numba.typed.Dict，改用连续数组+原子操作或分块归约，利用Numba的并行优化。

子方案2.1：原子操作并行更新 #

利用Numba的原子max操作，直接在数组上并行更新，避免线程冲突：

import numba as nb
import numpy as np

@nb.njit(parallel=True)
def parallel_reduce_max(entity_ids, values, unique_ids, id_to_idx):
    result = np.full(len(unique_ids), -np.inf, dtype=np.float32)
    for i in nb.prange(len(entity_ids)):
        eid = entity_ids[i]
        val = values[i]
        idx = id_to_idx[eid]
        # 原子max操作，线程安全
        nb.atomic.max(result, idx, val)
    return result

# 预处理：获取唯一实体ID并建立映射
entity_ids = iarr['e'][:, 0].astype(np.uint32)
values = iarr['r'].astype(np.float32)
unique_ids = np.unique(entity_ids)
# 建立实体ID到数组索引的映射
id_to_idx = np.searchsorted(unique_ids, entity_ids)

# 执行并行归约
max_vals = parallel_reduce_max(entity_ids, values, unique_ids, id_to_idx)

# 构造结果数组
oarr = np.array(list(zip(unique_ids, max_vals)), dtype=[('e', '<u4'), ('r', '<f4')])

子方案2.2：分块归约再合并 #

先分块计算每个块内的局部极值，再合并全局极值，避免原子操作的开销：

import numba as nb
import numpy as np

@nb.njit(parallel=True)
def block_reduce_max(entity_ids, values):
    num_threads = nb.get_num_threads()
    blocks = np.array_split(np.arange(len(entity_ids)), num_threads)
    
    # 每个线程生成局部结果字典
    local_results = []
    for tid in nb.prange(num_threads):
        local_dict = nb.typed.Dict.empty(key_type=nb.uint32, value_type=nb.float32)
        for i in blocks[tid]:
            eid = entity_ids[i]
            val = values[i]
            if eid not in local_dict or val > local_dict[eid]:
                local_dict[eid] = val
        local_results.append(local_dict)
    
    # 合并所有局部结果到全局字典
    global_dict = nb.typed.Dict.empty(key_type=nb.uint32, value_type=nb.float32)
    for local_dict in local_results:
        for eid, val in local_dict.items():
            if eid not in global_dict or val > global_dict[eid]:
                global_dict[eid] = val
    
    # 转换为数组
    res_ids = np.empty(len(global_dict), dtype=np.uint32)
    res_vals = np.empty(len(global_dict), dtype=np.float32)
    for i, (eid, val) in enumerate(global_dict.items()):
        res_ids[i] = eid
        res_vals[i] = val
    return res_ids, res_vals

# 执行分块归约
entity_ids = iarr['e'][:, 0].astype(np.uint32)
values = iarr['r'].astype(np.float32)
res_ids, res_vals = block_reduce_max(entity_ids, values)

oarr = np.array(list(zip(res_ids, res_vals)), dtype=[('e', '<u4'), ('r', '<f4')])

优势：

• 充分利用Numba的并行优化，性能接近原生C
• 分块归约避免全局字典的线程冲突，原子操作方案更简洁
• 内存开销可控，适配百万级数据规模

方案3：Pandas分组聚合（代码简洁，适合快速实现） #

如果项目已引入Pandas，利用其高效的groupby聚合功能，代码极简：

import pandas as pd
import numpy as np

# 转换为DataFrame
df = pd.DataFrame({
    'e': iarr['e'][:, 0].astype(np.uint32),
    'r': iarr['r'].astype(np.float32)
})

# 分组求最大值
result_df = df.groupby('e')['r'].max().reset_index()

# 转换为目标结构化数组
oarr = result_df.to_records(index=False)

优势：

• 代码可读性强，实现成本极低
• Pandas的groupby内部优化充分，性能接近Numpy方案
• 支持灵活扩展（如同时计算最大最小值）

方案4：Cython原生实现（极致性能，适合超大规模数据） #

若上述方案仍无法满足性能需求，可使用Cython编写底层循环，完全规避Python GIL，实现纯C级性能。

示例Cython代码（reduce_ext.pyx）：

import numpy as np
cimport numpy as np
from cython.parallel import prange
from libc.stdint cimport uint32_t, float_t
from libcpp.unordered_map cimport unordered_map

def cython_reduce_max(np.ndarray[uint32_t, ndim=1] entity_ids, 
                      np.ndarray[float_t, ndim=1] values):
    cdef:
        uint32_t *eids = entity_ids.data
        float_t *vals = values.data
        int n = entity_ids.shape[0]
        unordered_map[uint32_t, float_t] result_map
        int i
        uint32_t eid
        float_t val
    
    # 并行遍历，每个线程维护局部映射
    cdef unordered_map[uint32_t, float_t] local_maps[prange.num_threads()]
    for i in prange(n, nogil=True):
        eid = eids[i]
        val = vals[i]
        if local_maps[prange.thread_num()].find(eid) == local_maps[prange.thread_num()].end():
            local_maps[prange.thread_num()][eid] = val
        else:
            if val > local_maps[prange.thread_num()][eid]:
                local_maps[prange.thread_num()][eid] = val
    
    # 合并局部映射到全局结果
    for i in range(prange.num_threads()):
        for auto &item : local_maps[i]:
            if result_map.find(item.first) == result_map.end():
                result_map[item.first] = item.second
            else:
                if item.second > result_map[item.first]:
                    result_map[item.first] = item.second
    
    # 转换为numpy结构化数组
    cdef np.ndarray[uint32_t, ndim=1] res_eids = np.empty(len(result_map), dtype=np.uint32)
    cdef np.ndarray[float_t, ndim=1] res_vals = np.empty(len(result_map), dtype=np.float32)
    cdef int idx = 0
    for auto &item : result_map:
        res_eids[idx] = item.first
        res_vals[idx] = item.second
        idx += 1
    
    return np.rec.fromarrays([res_eids, res_vals], dtype=[('e', '<u4'), ('r', '<f4')])

编译后调用即可，性能接近纯C实现，适配超大规模数据处理场景。

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