超大规模DataFrame的1等值+2不等式条件高效连接实现方案咨询

我现在遇到一个棘手的DataFrame连接问题，想请教大家有没有更高效的解决方案：

我有两个DataFrame：

• a（约6亿行）
• b（约200万行）

需要将b连接到a上，连接条件是：

• 等值条件：a_1 = b_1
• 两个不等式条件：a_2 >= b_2 且 a_3 >= b_3

我已经尝试了以下几种方法，但都遇到了瓶颈：

尝试过的方案及问题 #

Polars #

• join_asof()：只支持1个不等式条件，满足不了我的需求
• join_where() + filter()：哪怕设置了很小的筛选窗口，标准Polars安装会触发43亿行的行数限制，换成polars-u64-idx版本后又直接耗尽了512GB内存

DuckDB #

• ASOF LEFT JOIN：同样只支持1个不等式条件，无法适配两个不等式的场景

Numba #

因为上面的方法都不行，我自己写了一个类似join_asof()的自定义函数（代码如下）。小数据量下运行没问题，但随着a的行数增加，速度慢到难以接受。我试过调整for/while循环和过滤逻辑的各种组合，结果都差不多。

现在我有点黔驴技穷了，有没有更高效的实现方式呢？谢谢大家！

import numba as nb
import numpy as np
import polars as pl
import time


@nb.njit(nb.int32[:](nb.int32[:], nb.int32[:], nb.int32[:], nb.int32[:], nb.int32[:], nb.int32[:], nb.int32[:]), parallel=True)
def join_multi_ineq(a_1, a_2, a_3, b_1, b_2, b_3, b_4):
    output = np.zeros(len(a_1), dtype=np.int32)

    for i in nb.prange(len(a_1)):

        for j in range(len(b_1) - 1, -1, -1):

            if a_1[i] == b_1[j]:

                if a_2[i] >= b_2[j]:

                    if a_3[i] >= b_3[j]:
                        output[i] = b_4[j]
                        break

    return output


length_a = 5_000_000
length_b = 2_000_000

start_time = time.time()
output = join_multi_ineq(a_1=np.random.randint(1, 1_000, length_a, dtype=np.int32),
                         a_2=np.random.randint(1, 1_000, length_a, dtype=np.int32),
                         a_3=np.random.randint(1, 1_000, length_a, dtype=np.int32),
                         b_1=np.random.randint(1, 1_000, length_b, dtype=np.int32),
                         b_2=np.random.randint(1, 1_000, length_b, dtype=np.int32),
                         b_3=np.random.randint(1, 1_000, length_b, dtype=np.int32),
                         b_4=np.random.randint(1, 1_000, length_b, dtype=np.int32))
print(f"Duration: {(time.time() - start_time):.2f} seconds")

备注：内容来源于stack exchange，提问作者usdn