[Python] How to rewrite a DataFrame-based grouped rolling aggregation function to use Polars...

I need to apply a custom grouped rolling aggregation function to every nth row for specific columns of a DataFrame. My current implementation works with a DataFrame as an argument and returns the modified DataFrame.

I have a few issues with the current approach:

• The function requires specifying many column names.
• Irrelevant columns (if present) are processed, even though they aren't needed.

I would like to rewrite this function so that it works with expressions as arguments and also returns the result as an expression. I'm not sure if this approach is more efficient, but it seems like it would make the handling cleaner.

Unfortunately, I’m having trouble figuring out how to implement this with Polars expressions.

Can anyone guide me on how to convert this DataFrame-based approach into an expression-based one?

Here’s my current function that works with DataFrames:

from typing import Callable, Sequence
import numpy as np
import polars as pl
from numba import guvectorize

@guvectorize(['(float64[:], int64, float64[:])'], '(n),()->(n)')
def rolling_func(input_array, window_size, out):
"""Example for a custom rolling function with a specified window size."""
n = len(input_array)
for i in range(n):
start = max(i - window_size + 1, 0)
out = np.mean(input_array[start:i+1])


def apply_rolling_gathered_agg(
df,
func: Callable,
window_size: int,
*func_args,
group_col: str | list[str] | None = None,
value_col: str | None = None,
result_col: str = 'result',
every_nth: int = 1,
window_buffer: int = 0,
return_dtype: pl.DataType = pl.Float64) -> pl.DataFrame:
"""
Apply a custom rolling aggregation function to a DataFrame, with grouping and every nth value selection.

This function performs a rolling aggregation on a specified value column in a Polars DataFrame. It allows
grouping by one or more columns, gathering every nth value, and applying a custom aggregation function
(e.g., `rolling_func`) with a specified window size and optional buffer.

Args:
df (pl.DataFrame): The DataFrame to operate on.
func (Callable): The aggregation function to apply to each rolling window.
window_size (int): The size of the window over which to apply the aggregation function.
*func_args: Additional arguments to pass to the custom function.
group_col (str | list[str] | None, optional): The column(s) to group by. If `None`, the first column is used.
value_col (str | None, optional): The column to apply the rolling function to. If `None`, the last column is used.
result_col (str, optional): The name of the result column in the output DataFrame. Default is 'result'.
every_nth (int, optional): The step size for gathering values within each group. Default is 1.
window_buffer (int, optional): A buffer to add around the rolling window, extending the window on both ends. Default is 0.
return_dtype (pl.DataType, optional): The desired data type for the result column. Default is `pl.Float64`.

Returns:
pl.DataFrame: A DataFrame containing the results of the rolling aggregation, with one row per group.


Example:
# Create a sample DataFrame with two groups 'A' and 'B', and values from 0 to 99
df = pl.DataFrame({
'group': np.repeat(['A', 'B'], 100), # Repeat 'A' and 'B' for each group
'value': np.tile(np.arange(100), 2) # Tile the values 0 to 99 for each group
})
func_args = []
res = apply_rolling_gathered_agg(
df,
func=rolling_func,
window_size=3,
*func_args,
group_col='group',
value_col='value',
every_nth=10,
window_buffer=0,
return_dtype=pl.Float64,
)
print(res)
res_pd = res.to_pandas()
"""
# Handle cases where group_col or value_col might not be passed
cols = df.columns
group_col = group_col or cols[0]
value_col = value_col or cols[-1]

# If group_col is a list, ensure it is processed correctly
if isinstance(group_col, list):
group_by = group_col
else:
group_by = [group_col]

# Temporary index column for rolling aggregation
index_col = '_index'

# Calculate the total window size
total_window = every_nth * (window_size + window_buffer)
period = f'{total_window}i'

# Apply rolling aggregation
result = (
df
.with_row_index(name=index_col)
.rolling(index_column=index_col, period=period, group_by=group_by)
.agg(
pl.all().last(), # pass the last element of all present columns
pl.col(value_col)
.gather_every(every_nth, offset=every_nth-1)
.map_batches(lambda batch: func(batch, window_size, *func_args), return_dtype=return_dtype)
.last().alias(result_col)) # This is the desired expression
.drop(index_col)
)
return result


I want to convert this function into an expression-based function, similar to:

def expr_apply_rolling_gathered_agg(
group_expr: pl.Expr | Sequence[pl.Expr], # Single or list of group column expressions
value_expr: pl.Expr, # Expression for the value column (series/column)
func: Callable, # The rolling aggregation function
window_size: int, # Size of the rolling window
*func_args, # Additional arguments for the rolling function
every_nth: int = 1, # Step size for gathering values
window_buffer: int = 0, # Buffer size around the window
return_dtype: pl.DataType = pl.Float64 # Output data type
) -> pl.Expr:
pass

Continue reading...