Python Polars: количество строк до следующего значения в группе

IIUC, вам нужно использовать back_fill и инвертировать вычитание:

(data
   .with_row_index()
   .with_columns(distance = 
      pl.when("login").then("index").backward_fill().over("user_id") - pl.col.index
   )
)

Выход:

┌───────┬─────────┬───────┬──────────┐
│ index ┆ user_id ┆ login ┆ distance │
│ ---   ┆ ---     ┆ ---   ┆ ---      │
│ u32   ┆ i64     ┆ bool  ┆ u32      │
╞═══════╪═════════╪═══════╪══════════╡
│ 0     ┆ 1       ┆ false ┆ 1        │
│ 1     ┆ 1       ┆ true  ┆ 0        │
│ 2     ┆ 1       ┆ false ┆ null     │
│ 3     ┆ 2       ┆ false ┆ 2        │
│ 4     ┆ 2       ┆ false ┆ 1        │
│ 5     ┆ 2       ┆ true  ┆ 0        │
└───────┴─────────┴───────┴──────────┘

Вы можете создать обратный индекс, используя параметр step.

i_expr = pl.int_range(pl.len(), 0, step=-1)

(
    df.with_columns(
        (i_expr - pl.when("login").then(i_expr).backward_fill())
        .over('user_id')
        .alias('distance')
    )
)

┌─────────┬───────┬──────────┐
│ user_id ┆ login ┆ distance │
│ ---     ┆ ---   ┆ ---      │
│ i64     ┆ bool  ┆ i64      │
╞═════════╪═══════╪══════════╡
│ 1       ┆ false ┆ 1        │
│ 1       ┆ true  ┆ 0        │
│ 1       ┆ false ┆ null     │
│ 2       ┆ false ┆ 2        │
│ 2       ┆ false ┆ 1        │
│ 2       ┆ true  ┆ 0        │
└─────────┴───────┴──────────┘

Или просто выполните обратное вычитание, как в ответе @mozway:

i_expr = pl.int_range(pl.len())

(
    df.with_columns(
        (pl.when("login").then(i_expr).backward_fill() - i_expr)
        .over('user_id')
        .alias('distance')
    )
)

┌─────────┬───────┬──────────┐
│ user_id ┆ login ┆ distance │
│ ---     ┆ ---   ┆ ---      │
│ i64     ┆ bool  ┆ i64      │
╞═════════╪═══════╪══════════╡
│ 1       ┆ false ┆ 1        │
│ 1       ┆ true  ┆ 0        │
│ 1       ┆ false ┆ null     │
│ 2       ┆ false ┆ 2        │
│ 2       ┆ false ┆ 1        │
│ 2       ┆ true  ┆ 0        │
└─────────┴───────┴──────────┘

Обратите внимание: я также перенес вычисление индекса на отдельную i_expr переменную и сместил операцию over() дальше, так что вам придется использовать ее только один раз, что упрощает корректировку решения.