本文介绍一种针对超大数据集（6亿行 × 200万行）的多不等式条件 asof join（1个等值 + 2个 ≥ 不等式）的高性能解决方案，通过 numba 构建哈希索引显著降低时间复杂度，实测提速达30倍。
本文介绍一种针对超大数据集（6亿行 × 200万行）的多不等式条件 asof join（1个等值 + 2个 ≥ 不等式）的高性能解决方案，通过 numba 构建哈希索引显著降低时间复杂度，实测提速达30倍。
在处理大规模时序或事件对齐任务时，常需基于一个等值键（如用户ID、设备ID）和多个不等式约束（如 a_ts >= b_ts 且 a_val >= b_threshold）进行关联匹配。然而，主流数据引擎（Polars、DuckDB）的 asof_join 仅支持单不等式条件；暴力双重循环（O(n×m)）在 6 亿行数据上不可行；而原始 Numba 实现虽避免 Python 解释开销，却因每次遍历全量 b 表查找匹配项，导致严重内存带宽瓶颈与缓存失效。
核心优化思路是：将 O(m) 的线性扫描降为 O(k)，其中 k 是等值键对应子集的平均大小。具体分两步：
-
预构建哈希索引字典：遍历
b表一次，以b_1值为键，存储所有满足b_1[j] == key的行索引列表； -
定向检索 + 逆序遍历：对每个
a[i]，仅在其对应的b_1子集中（而非全表）按j从大到小检查a_2[i] >= b_2[j] and a_3[i] >= b_3[j]，首次命中即终止——这保证了“最晚满足条件”的语义（类 ASOF 行为），同时最小化内层迭代次数。
以下是完整可运行的优化实现（兼容 Polars 数据输入，支持 int32 高效运算）：
import numba as nb
import numpy as np
# 定义 Numba 兼容的 List 类型
IntList = nb.types.ListType(nb.types.int32)
@nb.njit(nb.int32[:](nb.int32[:], nb.int32[:], nb.int32[:],
nb.int32[:], nb.int32[:], nb.int32[:], nb.int32[:]),
parallel=True)
def join_multi_ineq_fast(a_1, a_2, a_3, b_1, b_2, b_3, b_4):
n = len(a_1)
output = np.zeros(n, dtype=np.int32)
# Step 1: 构建 b_1 → [j indices] 的哈希索引
b1_indices = nb.typed.Dict.empty(key_type=nb.types.int32, value_type=IntList)
for j in range(len(b_1)):
key = b_1[j]
if key in b1_indices:
b1_indices[key].append(j)
else:
lst = nb.typed.List.empty_list(item_type=nb.types.int32)
lst.append(j)
b1_indices[key] = lst
# Step 2: 并行处理每个 a[i]
for i in nb.prange(n):
key = a_1[i]
if key not in b1_indices:
continue # 无匹配键，保持 output[i] = 0（或按需设为 null 等价值）
indices = b1_indices[key]
v2, v3 = a_2[i], a_3[i]
# 逆序遍历：优先取最大 j（隐含“最近”语义），提升命中率
for k in range(len(indices) - 1, -1, -1):
j = indices[k] # 注意：Numba 中 list 索引直接支持 uint32/int32
if v2 >= b_2[j] and v3 >= b_3[j]:
output[i] = b_4[j]
break
return output
# 使用示例（适配 Polars DataFrame）
# df_a = pl.read_parquet("a.parquet")
# df_b = pl.read_parquet("b.parquet")
# result = join_multi_ineq_fast(
# df_a["a_1"].to_numpy().astype(np.int32),
# df_a["a_2"].to_numpy().astype(np.int32),
# df_a["a_3"].to_numpy().astype(np.int32),
# df_b["b_1"].to_numpy().astype(np.int32),
# df_b["b_2"].to_numpy().astype(np.int32),
# df_b["b_3"].to_numpy().astype(np.int32),
# df_b["b_4"].to_numpy().astype(np.int32)
# )✅ 关键优势：
- 时间效率：实测 500 万 × 200 万随机数据下耗时仅 0.83 秒（原版 24.85 秒），提速约 30×；
- 内存友好：索引仅存储整数索引列表，内存占用远低于笛卡尔积中间态；
-
语义保真：逆序遍历确保返回
b中满足条件的最后一个匹配项，符合 ASOF “向后填充”直觉； -
可扩展性强：若需支持更多不等式（如
a_4 ），只需扩展条件判断逻辑，无需改变索引结构。
⚠️ 注意事项：
- 输入数组必须为
np.int32（或显式 cast），避免 Numba 类型推断失败； - 若
b_1值分布极度倾斜（如某 key 占b表 80% 行），最坏情况仍为 O(m)，建议先分析键分布，必要时对高频 key 单独优化（如二级排序 + 二分）； - 输出数组默认用
0填充未匹配项，生产环境建议返回Optional[int32]或配合pl.Series的null处理； - 对于浮点型字段，需替换为
nb.float64[:]并注意精度比较（建议用>=而非==）。
该方案在保持代码简洁性与工程可维护性的前提下，逼近理论最优性能边界，是当前处理“多条件 ASOF JOIN”这一经典难题的实用工业级解法。
