Request for rolling_top_k

[leoknuth]

Description

Request for a rolling_top_k expression.

Currently, I can do this with very slow code using rolling_map and lambda.

data = data.with_columns(
    roll_top_k_mean = pl.col("a").rolling_map(lambda x: x.top_k(5).mean(), window_size=10)
)

The method I hope to use is like:

data = data.with_columns(
    roll_top_k_mean = pl.col("a").rolling_top_k(k=5, w=10).mean()
)

You can use 2 heaps and 1 queue to implement this.

[cmdlineluser]

As an aside, I think the slow rolling_map can be replaced with Expr.rolling - right?

(data
  .with_row_index()
  .with_columns(
      pl.col("a").top_k(5).mean().rolling(index_column="index", period="5i")
  )
)