Given a polars DataFrame in Python, how can I modify every nth element in a series?
# have
df = pl.DataFrame(pl.Series("a", [1, -1, 1, -1, 1]))
# want
# [1, 1, 1, 1, 1]
# selecting works fine:
df["a", 1::2]
shape: (2,)
Series: 'a' [i64]
[
-1
-1
]
# but modification fails:
df["a", 1::2] *= -1
Traceback (most recent call last):
File "/tmp/ipykernel_103522/957012809.py", line 1, in <cell line: 1>
df["a", 1::2] *= -1
File "/home/.../.pyenv/versions/3.10.9/lib/python3.10/site-packages/polars/internals/dataframe/frame.py", line 1439, in __setitem__
raise ValueError(f"column selection not understood: {col_selection}")
ValueError: column selection not understood: slice(1, None, 2)
pl.__version__
'0.15.14'
You could add in the row index and use modulo arithmetic:
(df.with_row_index()
.with_columns(
pl.when((pl.col("index") + 1) % 2 == 0)
.then(pl.col("a") * -1)
.otherwise(pl.col("a"))
)
)
shape: (5, 2)
┌───────┬─────┐
│ index ┆ a │
│ --- ┆ --- │
│ u32 ┆ i64 │
╞═══════╪═════╡
│ 0 ┆ 1 │
│ 1 ┆ 1 │
│ 2 ┆ 1 │
│ 3 ┆ 1 │
│ 4 ┆ 1 │
└───────┴─────┘