如何使用 Python 中列表中的值有条件地更新 Polars DataFrame 中的列？

Question

我正在尝试更新 polars DataFrame 中的特定行，其中两列（“

Season"

”和

"Wk"

）满足某些条件，使用列表或系列中应与过滤行对齐的值。在 pandas ，我会使用

.loc[]

来执行此操作，但我还没有找到使用 Polars 实现相同结果的方法。

import polars as pl

# Sample DataFrame
df = pl.DataFrame({
    "Season": [2024, 2024, 2024, 2024],
    "Wk": [28, 28, 29, 29],
    "position": [1, 2, 3, 4]
})

# List of new values for the filtered rows
new_positions = [10, 20]  # Example values aligned with the filtered rows

# Condition to update only where Season is 2024 and Wk is 29
df = df.with_columns(
    # Expected update logic here
)

我尝试过使用

when().then()

和

zip_with()

的方法，但它们要么不直接接受列表/系列，要么不将新值与过滤后的行正确对齐。是否有推荐的方法使用列表或系列中的值以与过滤器相同的顺序有条件地更新 polars 中的列？

Answer 1

您可以使用

when/then

和

.cum_sum()

生成索引

df.with_columns(
   idx = pl.when(Season=2024, Wk=29).then(1).cum_sum() - 1
)

shape: (4, 4)
┌────────┬─────┬──────────┬──────┐
│ Season ┆ Wk  ┆ position ┆ idx  │
│ ---    ┆ --- ┆ ---      ┆ ---  │
│ i64    ┆ i64 ┆ i64      ┆ i32  │
╞════════╪═════╪══════════╪══════╡
│ 2024   ┆ 28  ┆ 1        ┆ null │
│ 2024   ┆ 28  ┆ 2        ┆ null │
│ 2024   ┆ 29  ┆ 3        ┆ 0    │
│ 2024   ┆ 29  ┆ 4        ┆ 1    │
└────────┴─────┴──────────┴──────┘

您可以

.get()

然后

.fill_null()

剩余的原始值。

df.with_columns(
   pl.lit(pl.Series(new_positions))
     .get(pl.when(Season=2024, Wk=29).then(1).cum_sum() - 1)
     .fill_null(pl.col.position)
     .alias("position")
)

shape: (4, 3)
┌────────┬─────┬──────────┐
│ Season ┆ Wk  ┆ position │
│ ---    ┆ --- ┆ ---      │
│ i64    ┆ i64 ┆ i64      │
╞════════╪═════╪══════════╡
│ 2024   ┆ 28  ┆ 1        │
│ 2024   ┆ 28  ┆ 2        │
│ 2024   ┆ 29  ┆ 10       │
│ 2024   ┆ 29  ┆ 20       │
└────────┴─────┴──────────┘