usdn's questions

我有两个数据框：a (~600M 行)和b (~2M 行) 。当在相应列上使用 1 个相等条件和2 个不等条件时，将 b 连接到 a 的最佳方法是什么？

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

我目前探索了以下路径：

• 極色：
  • join_asof()：仅允许 1 个不等式条件
  • join_where() 与 filter()：即使容差窗口较小，标准 Polars 安装在连接期间也会用尽行数（4.3B 行限制），并且 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")

我有一个长度约为 100M 行的数据框，其中包含不同组中的 ID。其中一些是错误的（以 99 表示）。我正尝试使用滚动模式窗口来纠正它们，类似于下面的代码示例。有没有更好的方法来做到这一点，因为 rolling_map() 非常慢？

import polars as pl
from scipy import stats

def dummy(input):
    return stats.mode(input)[0]

df = pl.DataFrame({'group': [10, 10, 10, 10, 10, 10, 10, 20, 20, 20, 20],
                   'id': [1, 1, 99, 1, 1, 2, 2, 3, 3, 99, 3]})

df.with_columns(pl.col('id')
                 .rolling_map(function=dummy,
                              window_size=3,
                              min_periods=1,
                              center=True)
                 .over('group')
                 .alias('id_mode'))

shape: (11, 3)
╭───────┬─────┬─────────╮
│ group ┆  id ┆ id_mode │
│   i64 ┆ i64 ┆     i64 │
╞═══════╪═════╪═════════╡
│    10 ┆   1 ┆       1 │
│    10 ┆   1 ┆       1 │
│    10 ┆  99 ┆       1 │
│    10 ┆   1 ┆       1 │
│    10 ┆   1 ┆       1 │
│    10 ┆   2 ┆       2 │
│    10 ┆   2 ┆       2 │
│    20 ┆   3 ┆       3 │
│    20 ┆   3 ┆       3 │
│    20 ┆  99 ┆       3 │
│    20 ┆   3 ┆       3 │
╰───────┴─────┴─────────╯