1

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

FObersteiner
  • 22,500
  • 8
  • 42
  • 72

1 Answers1

2

You could add in the row count and use the modulo operator:

df.with_row_count().select(
   pl.when((pl.col("row_nr") + 1) % 2 == 0)
     .then(pl.col("a") * -1)
     .otherwise(pl.col("a"))
)

shape: (5, 1)
┌─────┐
│ a   │
│ --- │
│ i64 │
╞═════╡
│ 1   │
├─────┤
│ 1   │
├─────┤
│ 1   │
├─────┤
│ 1   │
├─────┤
│ 1   │
└─────┘
jqurious
  • 9,953
  • 1
  • 4
  • 14
  • 1
    There is `.set_at_idx()` e.g. `idx = range(1, df.height, 2); df["a"].set_at_idx(idx, df["a", idx] * -1)` but I think that type of indexing/modification is discouraged. – jqurious Jan 16 '23 at 16:08