Polars用特定值替换前n行

Question

在Pandas中，我们可以用特定的值替换前面n行的数据：

import pandas as pd
import numpy as np

df = pd.DataFrame({"test": np.arange(0, 10)})
df.iloc[0:5] = np.nan
df

out:
   test
0   NaN
1   NaN
2   NaN
3   NaN
4   NaN
5   5.0
6   6.0
7   7.0
8   8.0
9   9.0

那么在Python Polars中，怎么做呢？

Answer 1

为了完整性，这里再提供几个可能的解决方案。

这个方法只适用于 null，除非你之后再填充这些空值。我们只需把这一列向上移动，然后再向下移动，这样顶部的 n 行就会留空：

df.with_columns(pl.col('test').shift(-5).shift(5))

┌──────┐
│ test │
│ ---  │
│ f64  │
╞══════╡
│ Null │
│ Null │
│ Null │
│ Null │
│ Null │
│ 6    │
│ 7    │
│ 8    │
│ 9    │
│ 10   │
└──────┘

另一个方法是创建一个新列，这个新列的值是给定的值重复 n 次，然后我们使用 append() 把它和现有列的 [tail()] 连接起来：

df.with_columns(
    pl.repeat(None, 5).append(pl.col('test').tail(df.count() - 5))
)

┌──────┐
│ test │
│ ---  │
│ f64  │
╞══════╡
│ Null │
│ Null │
│ Null │
│ Null │
│ Null │
│ 6    │
│ 7    │
│ 8    │
│ 9    │
│ 10   │
└──────┘

我用 timeit 做了一个快速的性能测试，结果显示使用 with_row_index 的方法是所有方法中最快的：

%timeit df.with_row_index().with_columns(pl.when(pl.col("index") < 5).then(None).otherwise(pl.col("test")).alias("test")).drop("index")
65.3 µs ± 3.4 µs per loop (mean ± std. dev. of 7 runs, 10,000 loops each)

%timeit df.with_columns(pl.repeat(None, 5).append(pl.col('test').tail(df.count() - 5)))
286 µs ± 14.5 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)

%timeit df.with_columns(pl.when(pl.int_range(0, pl.len()) >= 5).then("test"))
140 µs ± 20.5 µs per loop (mean ± std. dev. of 7 runs, 10,000 loops each)

%timeit df.with_columns(pl.col('test').shift(-5).shift(5))
142 µs ± 13.5 µs per loop (mean ± std. dev. of 7 runs, 10,000 loops each)

也可以把使用 with_row_index 的解决方案稍微调整一下，让它变得更简洁（再次强调，这种方法只适合你想用 null 作为替代值的情况）：

(
    df.with_row_index()
    .with_columns(
        pl.when(pl.col("index") >= 5).then(pl.col("test"))
    ).drop("index")
)

Answer 2

试试这个一行代码：

df.with_columns(pl.when(pl.int_range(0, pl.len()) >= 5).then("test"))

Answer 3

你可以使用 DataFrame.with_row_index() 这个方法：

import polars as pl

df = pl.DataFrame({"test": np.arange(1, 11)})

print(
    df.with_row_index()
    .with_columns(
        pl.when(pl.col("index") < 5)
        .then(None)
        .otherwise(pl.col("test"))
        .alias("test")
    )
    .drop("index")
)

输出结果是：

shape: (10, 1)
┌──────┐
│ test │
│ ---  │
│ f64  │
╞══════╡
│ Null │
│ Null │
│ Null │
│ Null │
│ Null │
│ 6    │
│ 7    │
│ 8    │
│ 9    │
│ 10   │
└──────┘

或者你也可以用 pl.int_range()，这个方法在处理分组时会更灵活。

print(
(df
 .with_columns(
    pl.int_range(0, pl.len(), dtype=pl.UInt32).over(True).alias("index")
    )
 .with_columns(
     pl.when(pl.col("index") < 5)
     .then(None)
     .otherwise(pl.col("test"))
     .alias("test")
     )
 .drop("index")
 )
)