过滤大量（数百）条件

Question

我有一个较大的数据框（550 万行，四列）。第一列（我们称之为 A 列）有 235 个不同的条目。第二列 (B) 有 100 个不同的条目，整数从 0 到 99，所有条目都以不同的比例出现在 A 中的每个条目中。我通过随机 selecting 来自 B 的值对 A 进行分组。类似于

df.groupby("A").agg(
    pl.col("B").unique().apply(np.random.choice)
)

在这样做时，A 中的每个值都被赋予一个随机整数。我的目标是 select，从数据帧中，列 C 和 D 对应于如此生成的 A、B 对。

目前我的方法是

choice = df.groupby("A").agg(
    pl.col("B").unique().apply(np.random.choice)
).to_numpy()
lchoice = ((df["A"] == arr[0]) & (df["B"] == arr[1]) for arr in choice)
mask = functools.reduce(operator.or_, lchoice)
sdf = df[mask].select(["C", "D"])

它完成了工作，但感觉不是很地道。

我的第一次尝试是

sdf = df.filter(functools
                .reduce(operator.or_,
                        [(pl.col("A") == arr[0]) & (pl.col("B") == arr[1])
                         for arr in choice]))

但它一直挂起直到我杀死它（我等了大约 30 分钟，而另一种方法需要 1.6 秒）。

df.filter(
    (pl.col("period") == choice[0, 0]) & (pl.col("exp_id") == choice[0, 1])
)

工作正常，正如预期的那样，并且我在过去成功地使用 functools.reduce 构造作为参数来过滤。显然，我不想全部手写；我可以遍历 choice 的行，一次过滤 df 一个，然后连接数据帧，但这听起来比它应该的要昂贵得多。

关于无需创建临时对象、数组等即可实现我的 sdf“极地方式”的任何提示？正如我所说，我有一个可行的解决方案，但它有点不稳定，我有兴趣学习更好的极地。

编辑：一些模拟数据

df = pl.DataFrame({"A": [1.3, 8.9, 6.7]*3 + [3.6, 4.1]*2,
                   "B": [1]*3 + [2]*3 + [3]*3 + [1]*2 + [2]*2,
                   "C": [21.5, 24.3, 21.8, 20.8, 23.6, 15.6, 23.5,
                         16.1, 15.6, 14.8, 14.7, 23.8, 20.],                   
                   "D": [6.9, 7.6, 6.4, 6.2, 7.6, 6.2,
                         6.3, 7.1, 7.8,7.7, 6.5, 6.6, 7.1]})

对已接受的答案略作改动：

df.sort(by=['A', 'B'], in_place=True)
sdf = (df
       .join(df
             .groupby('A', maintain_order=True)
             .agg(pl.col('B')
                  .unique()
                  .sort()
                  .shuffle(seed)
                  .first()
                  .alias('B')),   
             on=['A', 'B'])
       .select(['C','D']))

我需要多次执行此操作，我想确保随机生成的可重复性，因此排序和 maintain_order=True。

Answer 1

看起来您可以用 join 完成您需要的事情。

让我们看看我是否理解你的问题。让我们从这些数据开始：

import polars as pl

df = pl.DataFrame(
    {
        "A": ["a", "b", "c", "c", "b", "a"] * 2,
        "B": [1, 2, 3, 2] * 3,
        "C": [1, 2, 3, 4, 5, 6] * 2,
        "D": ["x", "y", "y", "x"] * 3,
    }
).sort(["A", "B"])
print(df)

shape: (12, 4)
┌─────┬─────┬─────┬─────┐
│ A   ┆ B   ┆ C   ┆ D   │
│ --- ┆ --- ┆ --- ┆ --- │
│ str ┆ i64 ┆ i64 ┆ str │
╞═════╪═════╪═════╪═════╡
│ a   ┆ 1   ┆ 1   ┆ x   │
├╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌┤
│ a   ┆ 2   ┆ 6   ┆ y   │
├╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌┤
│ a   ┆ 2   ┆ 6   ┆ x   │
├╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌┤
│ a   ┆ 3   ┆ 1   ┆ y   │
├╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌┤
│ b   ┆ 1   ┆ 5   ┆ x   │
├╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌┤
│ b   ┆ 2   ┆ 2   ┆ y   │
├╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌┤
│ b   ┆ 2   ┆ 2   ┆ x   │
├╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌┤
│ b   ┆ 3   ┆ 5   ┆ y   │
├╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌┤
│ c   ┆ 1   ┆ 3   ┆ x   │
├╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌┤
│ c   ┆ 2   ┆ 4   ┆ x   │
├╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌┤
│ c   ┆ 2   ┆ 4   ┆ y   │
├╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌┤
│ c   ┆ 3   ┆ 3   ┆ y   │
└─────┴─────┴─────┴─────┘

接下来，我们随机 select A 的每个值对应一个 B 的值。

我将稍微更改您的代码以消除对 numpy 的使用，而是使用 Polars 的 shuffle 表达式。这样，我们就得到了一个 Polars DataFrame，我们将在即将到来的连接中使用它。（shuffle 表达式的文档是 here。如果没有提供种子，它会使用 numpy。）

choice_df = df.groupby("A").agg(pl.col("B").unique().shuffle().first())
print(choice_df)

shape: (3, 2)
┌─────┬─────┐
│ A   ┆ B   │
│ --- ┆ --- │
│ str ┆ i64 │
╞═════╪═════╡
│ b   ┆ 3   │
├╌╌╌╌╌┼╌╌╌╌╌┤
│ a   ┆ 2   │
├╌╌╌╌╌┼╌╌╌╌╌┤
│ c   ┆ 1   │
└─────┴─────┘

如果我对你的问题的理解正确，我们现在想要获取原始数据集中 C 列和 D 列的值，这些值对应于我们在上一步中 select 编辑的 A 和 B 的三种组合.我们可以用 join 最简单地完成此操作，然后是 select。（select 只是为了从结果中删除列 A 和 B）。

df.join(choice_df, on=['A','B']).select(['C','D'])

shape: (4, 2)
┌─────┬─────┐
│ C   ┆ D   │
│ --- ┆ --- │
│ i64 ┆ str │
╞═════╪═════╡
│ 6   ┆ y   │
├╌╌╌╌╌┼╌╌╌╌╌┤
│ 6   ┆ x   │
├╌╌╌╌╌┼╌╌╌╌╌┤
│ 5   ┆ y   │
├╌╌╌╌╌┼╌╌╌╌╌┤
│ 3   ┆ x   │
└─────┴─────┘

这是否满足您的需求？生成的代码干净、简洁，并且是 Polars API.

的典型用法

过滤大量（数百）条件

Filtering on a large number (hundreds) of conditions

python

python-polars