In Pandas we can replace the first n rows with certain values:
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
What would be the equivalent operations in Python Polars?
You can use 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")
)
Prints:
shape: (10, 1)
┌──────┐
│ test │
│ --- │
│ f64 │
╞══════╡
│ Null │
│ Null │
│ Null │
│ Null │
│ Null │
│ 6 │
│ 7 │
│ 8 │
│ 9 │
│ 10 │
└──────┘
or by pl.int_range() which gives more flexibility in cases of groups.
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")
)
)
