[Python] Rolling Kurtosis with Group_by in Polars

As far as I know there's no built-in rolling_kurtosis function in Polars. (There has been one ticket requesting for this feature back in 2022 but has since been stale.) I tried to implement rolling kurtosis with functionalities in the existing Polars library but would like people who have more experience to see if the solutions make sense, and whether there could be improvements.

For testing purposes I created the following dataframe:

import polars as pl
import polars_xdt as xdt

holidays = [datetime.date(y, 7, 4) for y in range(2000, 2025)]
dates = pl.select(
xdt.date_range(datetime.date(2000, 7, 1), datetime.date(2024, 7, 31), interval='1bd', holidays=holidays)
.alias('date')
)

values = np.random.rand(len(dates) * 500)
values[523], values[1040] = np.nan, np.nan # just throwing in some nans

df = pl.LazyFrame(
{
't': itertools.chain.from_iterable([
[d] * 500
for d in xdt.date_range(
datetime.date(2000, 7, 1), datetime.date(2024, 7, 31),
interval='1bd',
holidays=holidays,
eager=True
)]),
'id': itertools.chain.from_iterable(range(1, 501) for _ in range(len(dates))),
'value': values
}
).with_columns(
pl.col('value').fill_nan(None)
).collect()


As a baseline, I used pandas and checked its performance:

df_pd = df.to_pandas().set_index(['t', 'id'])

# 437 ms ± 10 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
df_pd.groupby('id').apply(
lambda s: s.rolling(window=5, min_periods=4).kurt()
).reset_index(level=0, drop=True).sort_index()


Now here are what I tried in polars:

1. First method: Expr.rolling_map + Expr.over

This took almost 12 minutes to finish, which is a bit intolerable...

# 11min 42s ± 1min 46s per loop (mean ± std. dev. of 7 runs, 1 loop each)

q_rolling_map = (
df.lazy().with_columns(
pl.col('value').rolling_map(
functools.partial(st.kurtosis, bias=False, nan_policy='omit'),
window_size=5,
min_periods=4
).over('id')
)
)

out = q_rolling_map.collect()

1. Second Method: DataFrame.rolling (with groupby).agg + Expr.kurt

This is less ideal though because ideally I'd like to have a polars expression returned, such that I can easily use it in a Select context.

It's also a bit complicated to handle business days, as there's no built-in period='5bd' in polars, and there's also no easy way to do rolling simply based on the number of rows (like what pandas does). I also had to use conditional expressions to implement the min_periods.

# 2.47 s ± 231 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
q_rolling = (
df.lazy().join(
df.lazy().select(pl.col('t').unique()).sort('t').with_row_index('t_idx'), on='t'
)
.rolling(index_column='t_idx', period='5i', group_by='id').agg(
pl.col('t').last(),
pl.when(pl.col('value').count() >= 4)
.then(pl.col('value').kurtosis(bias=False))
.otherwise(None)
.alias('kurt')
)
.select('t', 'id', 'kurt')
.sort(['t', 'id'])
)

q_rolling.collect()


An alternative way is to skip the join in the very beginning, but it's not faster probably because I had to do the sorting first.

q_rolling_1 = (
df.lazy().sort(['id', 't'])
.rolling(index_column=pl.int_range(pl.len()).alias('t_idx'), period='5i', group_by='id').agg(
pl.col('t').last(),
pl.when(pl.col('value').count() >= 4)
.then(pl.col('value').kurtosis(bias=False))
.otherwise(None)
.alias('kurt')
)
.select('t', 'id', 'kurt')
.sort(['t', 'id'])
)

q_rolling_1.collect()

1. Compare with rolling_std and rolling_skew

There's a built-in rolling_std function in polars (I understand maybe it takes slightly longer to compute kurtosis), which is much faster than all the implementations above:

# 33.2 ms ± 449 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

q_rolling_std = (
df.lazy().with_columns(
pl.col('value').rolling_std(window_size=5, min_periods=4).over('id').alias('std')
)
)

q_rolling_std.collect()


@jcurious pointed out that rolling_skew in polars is slower than pandas. I tested that and indeed that was the case. Also there's no min_periods argument in rolling_skew, which is inconsistent with the other rolling methods in polars.

# 1.63 s ± 11.3 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
q_rolling_skew = (
df.lazy().with_columns(
pl.col('value').rolling_skew(window_size=5, bias=False).over('id').alias('skew') # why no min_periods arg?
)
)

q_rolling_skew.collect()


# 422 ms ± 12.8 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
df_pd.groupby('id').apply(
lambda s: s.rolling(window=5, min_periods=4).skew()
).reset_index(level=0, drop=True).sort_index()


My questions are:

1. 
   Do my implementations make sense? Have I I made some stupid mistakes that led to slower than expected performance (in either the First or Second method)? Any best practices to follow in groupby rolling in general?
   
   
2. 
   Is there actually any way to do df.rolling with window size as rolling window, but without having to manually create an index column? (Or maybe it's a feature to be requested?)
   
   
3. 
   What's the best way to escalate the need for rolling_kurtosis, now that there's already a stale feature request?
   
   
4. 
   Is there a general need for the functionality where chained polars expressions can be used in rolling + over? e.g. I could implement std as
   

((pl.col('value') - pl.col('value').mean()) ** 2).mean()


but I need to apply it in a rolling+over setting.

Thanks for your inputs in advance!

Continue reading...