Polars修改符合条件的列值

听说了polars比pandas快,今天使用了polars,然后有个需求是要修改某些符合条件的行数据,找了好多文章都没介绍这块. 后来问了AI发现一个方法,在此记录一下.

AI给的示例:

import polars as pl

df = pl.DataFrame(
    {
        "id": [0, 1, 2, 3],
        "left": [10, 11, 12, 13],
        "right": [14, 15, 16, 17],
    }
)

n = 15

# 创建条件列
df = df.with_column(pl.when(df.right < n).then(df.right + 2).otherwise(df.right))

# 过滤出符合条件的行
df = df[df.right < n]

# 设置符合条件的列的 right 为 right + 2
df = df.with_columns(
    [
        pl.col("id"),
        pl.col("left"),
        pl.col("right").alias("right_updated").when(df.right < n).then(df.right + 2),
    ]
)

print(df)

我的代码:

df = df.with_columns(
        [
            pl
            .when(pl.col('left')>start)
            .then(pl.col('left')+2)
            .otherwise(pl.col('left')),
            pl
            .when(pl.col('right')>=start)
            .then(pl.col('right')+2)
            .otherwise(pl.col('right'))
        ]
    )

这样还支持循环中修改df的值