Asof-join с множественными условиями неравенства ⇐ Python

[Anonymous]

У меня есть два кадра данных: a (~600 млн строк) и b (~2 млн строк). Как лучше всего соединить b с a при использовании 1 условия равенства и 2 условий неравенства в соответствующих столбцах?

• a_1 = b_1
• a_2 >= b_2
• a_3 >= b_3

На данный момент я исследовал следующие пути:

• Polars:
  
  join_asof(): допускает только 1 условие неравенства
• join_where() с фильтром(): даже при небольшом окне допуска в стандартной установке Polars не хватает строк (ограничение строк 4,3 Б) во время объединения, а при установке Polars-u64-idx не хватает памяти (512 ГБ)

[*]DuckDB: ASOF LEFT JOIN: также допускает только одно условие неравенства
[*]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 == b_1[j]:

if a_2 >= b_2[j]:

if a_3 >= b_3[j]:
output = 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")


Подробнее здесь: https://stackoverflow.com/questions/793 ... conditions