asof-join com múltiplas condições de desigualdade

Você pode usar distinct ona cláusula DuckDB (Postgresql):

import duckdb

df_res = duckdb.sql("""
    select distinct on (a.a1, a.a2, a.a3)
        a.a1,
        a.a2,
        a.a3,
        b.b4
    from df_a as a
        inner join df_b as b on
            a.a1 = b.b1 and
            a.a2 >= b.b2 and
            a.a3 >= b.b3
    order by
        b.b2 desc,
        b.b3 desc
""").pl()

Você também pode tentar usar pl.DataFrame.join_where()mas no modo lazy. Estou assumindo que seu dataframe 'a' tem uma chave única, neste caso de exemplo - a1,a2,a3.

• pl.DataFrame.lazy()para tratar DataFrame como LazyFrame.
• pl.LazyFrame.join_where()para unir quadros preguiçosos.
• pl.LazyFrame.sort()para classificar os resultados.
• pl.LazyFrame.drop()para remover b2,b3colunas.
• pl.LazyFrame.unique()para deixar apenas uma linha por a1,a2,a3.
• pl.LazyFrame.collect().

df_res = (
    df_a.lazy()
    .join_where(
        df_b.lazy(),
        pl.col.a1 == pl.col.b1,
        pl.col.a2 >= pl.col.b2,
        pl.col.a3 >= pl.col.b3
    )
    .sort(["a1","a2","a3","b2","b3"])
    .drop(["b2","b3"])
    .unique(["a1","a2","a3"], keep="first")
).collect()

Se nada disso funcionar, você pode tentar dividir um dos quadros em N pedaços com pl.DataFrame.partition_by(), processar os pedaços separadamente e então usar pl.concat()para concatená-los novamente.

N = 20

df_a_list = (
    df_a
    .with_columns(r = pl.int_range(pl.len()) * N // pl.len())
    .partition_by("r", include_key=False)
)

df_res = pl.concat([
    df_a_t.join_where(
        df_b,
        pl.col.a1 == pl.col.b1,
        pl.col.a2 >= pl.col.b2,
        pl.col.a3 >= pl.col.b3
    )
    .sort(["a1","a2","a3","b2","b3"])
    .drop(["b2","b3"])
    .unique(["a1","a2","a3"], keep="first")
    for df_a_t in df_a_list
])

Usar Numba aqui é uma boa ideia, já que a operação é particularmente cara. Dito isso, a complexidade do algoritmo é, O(n²)embora não seja fácil fazer muito melhor (sem tornar o código muito mais complexo). Além disso, o array b_1, que pode não caber no cache L3, é lido completamente 5_000_000 vezes, tornando o código bastante limitado à memória.

Podemos acelerar bastante o código construindo um índice para não percorrer todo o array b_1, mas apenas os valores where a_1[i] == b_1[j]. Isso não é suficiente para melhorar a complexidade, pois muitos jvalores atendem a essa condição. Podemos melhorar a complexidade (média) construindo um tipo de árvore para todos os nós do índice, mas, na prática, isso torna o código muito mais complexo e o tempo para construir a árvore seria tão grande que, na verdade, não vale a pena fazer isso na prática. De fato, um índice básico é suficiente para reduzir bastante o tempo de execução no conjunto de dados aleatório fornecido (com números uniformemente distribuídos). Aqui está o código resultante:

import numba as nb
import numpy as np
import time

length_a = 5_000_000
length_b = 2_000_000

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)

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):
    output = np.zeros(len(a_1), dtype=np.int32)
    b1_indices = nb.typed.Dict.empty(key_type=nb.types.int32, value_type=IntList)
    for j in range(len(b_1)):
        val = b_1[j]
        if val in b1_indices:
            b1_indices[val].append(j)
        else:
            lst = nb.typed.List.empty_list(item_type=np.int32)
            lst.append(j)
            b1_indices[val] = lst
    kmean = 0
    for i in nb.prange(len(a_1)):
        if a_1[i] in b1_indices:
            indices = b1_indices[a_1[i]]
            v2 = a_2[i]
            v3 = a_3[i]
            for k in range(len(indices) - 1, -1, -1):
                j = indices[np.uint32(k)]
                #assert a_1[i] == b_1[j]
                if v2 >= b_2[j] and v3 >= b_3[j]:
                    output[i] = b_4[j]
                    break
    return output

%time join_multi_ineq_fast(a_1, a_2, a_3, b_1, b_2, b_3, b_4)

Note que, em média, apenas 32 kvalores são testados (o que é razoável o suficiente para não construir uma estrutura de dados mais eficiente/complicada). Observe também que o resultado é estritamente idêntico ao fornecido pela implementação ingênua.

Referência

Aqui estão os resultados na minha CPU i5-9600KF (6 núcleos):

Roman's code:        >120.00 sec     (require a HUGE amount of RAM: >16 GiB)
Naive Numba code:      24.85 sec
This implementation:    0.83 sec     <-----

Portanto, essa implementação é cerca de 30 vezes mais rápida que o código inicial.