使用另一个列表作为极坐标中的布尔掩码的过滤器列表

我有一个包含两列的极坐标数据框，其中两列都是列表。

df = pl.DataFrame({
    'a': [[True, False], [False, True]],
    'b': [['name1', 'name2'], ['name3', 'name4']]
})
df
shape: (2, 2)
┌───────────────┬────────────────────┐
│ a             ┆ b                  │
│ ---           ┆ ---                │
│ list[bool]    ┆ list[str]          │
╞═══════════════╪════════════════════╡
│ [true, false] ┆ ["name1", "name2"] │
├╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
│ [false, true] ┆ ["name3", "name4"] │
└───────────────┴────────────────────┘

我想过滤列b使用列a作为布尔掩码。列中每个列表的长度a始终与列中每个列表的长度相同b.

我可以考虑使用explode，然后过滤、聚合并执行join，但在某些情况下，连接列不可用，为了简单起见，我宁愿避免使用此方法。

是否有其他方法使用另一个列表作为布尔掩码来过滤列表？我尝试过使用.arr.eval，但它似乎不接受涉及其他列的操作。

任何帮助，将不胜感激！

这不是最理想的解决方案，因为我们整理数据，为每个分解为元素的列表创建一个组。然后我们再次按该组进行分组并应用过滤器。

df = pl.DataFrame({
    'a': [[True, False], [False, True]],
    'b': [['name1', 'name2'], ['name3', 'name4']]
})

(df.with_row_count()
   .explode(["a", "b"])
   .groupby("row_nr")
   .agg([
       pl.col("b").filter(pl.col("a"))
   ])
)

shape: (2, 2)
┌────────┬───────────┐
│ row_nr ┆ b         │
│ ---    ┆ ---       │
│ u32    ┆ list[str] │
╞════════╪═══════════╡
│ 1      ┆ ["name4"] │
├╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌┤
│ 0      ┆ ["name1"] │
└────────┴───────────┘

也许我们可以在极地想出更好的办法。如果arr.eval可以访问其他列。待定！

编辑 02-06-2022

In polars-0.13.41这不会像您想象的那么贵。北极星知道row_count已排序并在整个查询中保持排序。列表列的爆炸也是免费的。

当 Polars 知道 groupby 键已排序时，groupby 操作将快约 15 倍。

在上面的查询中，您只需支付：

• 行数爆炸式增长
• 对排序后的键进行分组（速度非常快）
• 遍历列表（无论如何，这是我们需要支付的费用）。

为了确保它运行速度快，您可以使用以下命令运行查询POLARS_VERBOSE=1。这会将以下文本写入 stderr：

could fast explode column a
could fast explode column b
keys/aggregates are not partitionable: running default HASH AGGREGATION
groupby keys are sorted; running sorted key fast path