pythonindexingslicepython-polars

Python polars: modify every nth row


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'

pandas version of the question


Solution

  • 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   │
    └───────┴─────┘