Baes on Index (or column value), update values in a multiple column in Polars - python-polars

Question

Based on the index, I want to update 3 columns in a data frame. In pandas, I would do the following:
'''
index = 5
df.loc[index, ['a', 'b', 'c']] = var1, var2, var3
'''
What is the polras equivalent notation?

Answer 1

One possible approach is to use when/then/otherwise combined with .map()
df = pl.DataFrame({
"a": [1, 2, 3],
"b": [4, 5, 6],
"c": [7, 8, 9],
"d": [10, 11, 12]
})
index = 1
df.with_row_count().with_columns(
pl.when(pl.col("row_nr") == index)
.then(pl.all().map(lambda column:
{
"a": 100,
"b": 200,
"c": 300
}.get(column.name, column)))
.otherwise(pl.all())
)
shape: (3, 5)
┌────────┬─────┬─────┬─────┬─────┐
│ row_nr | a | b | c | d │
│ --- | --- | --- | --- | --- │
│ u32 | i64 | i64 | i64 | i64 │
╞════════╪═════╪═════╪═════╪═════╡
│ 0 | 1 | 4 | 7 | 10 │
├────────┼─────┼─────┼─────┼─────┤
│ 1 | 100 | 200 | 300 | 11 │
├────────┼─────┼─────┼─────┼─────┤
│ 2 | 3 | 6 | 9 | 12 │
└────────┴─────┴─────┴─────┴─────┘
Another is .set_at_idx()
df.with_columns([
df["a"].set_at_idx(index, 100),
df["b"].set_at_idx(index, 200),
df["c"].set_at_idx(index, 300)
])
shape: (3, 4)
┌─────┬─────┬─────┬─────┐
│ a | b | c | d │
│ --- | --- | --- | --- │
│ i64 | i64 | i64 | i64 │
╞═════╪═════╪═════╪═════╡
│ 1 | 4 | 7 | 10 │
├─────┼─────┼─────┼─────┤
│ 100 | 200 | 300 | 11 │
├─────┼─────┼─────┼─────┤
│ 3 | 6 | 9 | 12 │
└─────┴─────┴─────┴─────┘