Polars: How to find value in previous values by predicate?

Question

How to find "price" of previous 'top'?

df = DataFrame({
    'price': [1, 2, 4, 2, 1, 2, 3, 2],
    'spikes': [None, None, 'top', None, 'bottom', None, 'top', None]
})

In result I want to find value of previous top/bottom price. Expected result would be:

df = DataFrame({
    'price': [1, 2, 4, 2, 1, 2, 3, 2],
    'spikes': [None, None, 'top', None, 'bottom', None, 'top', None],
    'prev_spikes_prices: [None, None, None, None, None, None, 4, None]
})

A `fill_null` with an expression is probably the way to go. What is the expected output? I don't understand "third row of previous price". — Wayoshi, Apr 25 '23 at 23:55

score 2 · Accepted Answer · answered Apr 26 '23 at 09:36

So this can be solved by grouping by spikes and shifting the price column by one and since you only want the prices where spikes is not null we add this as pl.when condition. See the following code:

(df
 .with_columns(
    pl.when(pl.col("spikes").is_not_null())
    .then(pl.col("price").shift(1).over("spikes"))
    .otherwise(None).alias("prev_spikes_prices"))
)



shape: (8, 3)
┌───────┬────────┬────────────────────┐
│ price ┆ spikes ┆ prev_spikes_prices │
│ ---   ┆ ---    ┆ ---                │
│ i64   ┆ str    ┆ i64                │
╞═══════╪════════╪════════════════════╡
│ 1     ┆ null   ┆ null               │
│ 2     ┆ null   ┆ null               │
│ 4     ┆ top    ┆ null               │
│ 2     ┆ null   ┆ null               │
│ 1     ┆ bottom ┆ null               │
│ 2     ┆ null   ┆ null               │
│ 3     ┆ top    ┆ 4                  │
│ 2     ┆ null   ┆ null               │
└───────┴────────┴────────────────────┘

Polars: How to find value in previous values by predicate?

1 Answers1