Как отфильтровать фрейм данных Polars на две части: строки, соответствующие условию, и строки, которые не соответствуют условию?
Я могу разделить фрейм данных df на две части следующим образом:
predicate = pl.col("a").gt(pl.col("b"))
dfx = df.filter(predicate)
dfy = df.filter(~predicate)
Это расточительно, потому что мы не хотим/не должны проходить df дважды, чтобы произвести dfx и dfy. Есть ли лучший метод?
@Догберт, думаю, этого будет достаточно. Насколько я понимаю, добавление столбца обходится дешево (в конце концов, DataFrame — это всего лишь набор серий, и мы можем дешево удалять и добавлять серии (хотя создание серии может быть недешевым, но в данном случае это будет дешево) потому что мы вычисляем предикат только один раз, а не два))?
Ответы 2
Вы можете использовать with_columns, чтобы добавить свой предикат, а затем partition_by с include_key, установленным в False, чтобы удалить этот столбец. Кроме того, вам нужно выполнить сортировку по этому столбцу и использовать maintain_order, чтобы надежно получить результаты там, где вы хотите.
Все вместе вы сможете сделать
dfx, dfy = (
df
.with_columns(_z=pl.col("a").gt(2))
.sort('_z')
.partition_by('_z',maintain_order=True, include_key=False)
)
Вы можете group_by любое выражение, а не только существующие столбцы, поэтому вы можете group_by использовать предикат. Затем вы можете просто перебирать объект GroupBy.
import polars as pl
import numpy as np
df = pl.DataFrame(
{"a": np.random.randint(5, size=10), "b": np.random.randint(2, size=10)},
)
for (key, group) in df.group_by(pl.col("a") > pl.col("b")):
print(f"{key=}")
print(group)
print()
key=(True,)
shape: (3, 2)
┌─────┬─────┐
│ a ┆ b │
│ --- ┆ --- │
│ i64 ┆ i64 │
╞═════╪═════╡
│ 3 ┆ 0 │
│ 2 ┆ 1 │
│ 2 ┆ 1 │
└─────┴─────┘
key=(False,)
shape: (7, 2)
┌─────┬─────┐
│ a ┆ b │
│ --- ┆ --- │
│ i64 ┆ i64 │
╞═════╪═════╡
│ 1 ┆ 1 │
│ 0 ┆ 0 │
│ 0 ┆ 1 │
│ 0 ┆ 0 │
│ 0 ┆ 1 │
│ 0 ┆ 1 │
│ 1 ┆ 1 │
└─────┴─────┘
Есть part_by (docs.pola.rs/api/python/stable/reference/dataframe/api/…), но он принимает только существующие столбцы, поэтому вам придется добавить столбец для предиката, затем разбить, затем брось это, оно того не стоит.