Pivot for polars LazyFrame

lazyframe_pivot.py

import polars as pl

def lazy_pivot(lf, on, index, values, aggregate_function=None, maintain_order=True):
    on = pl.col(on)
    index = pl.col(index)
    values = pl.col(values)
    unq_values = lf.select(on).unique(maintain_order=maintain_order).collect(engine='streaming').to_series().to_list()
    if aggregate_function is None:
        # FIXME: this does not replicate eager pivot default value which raises when there are several values in a cell
        aggregate_function = 'first'

    aggs = [
        getattr(values.filter(on == value), aggregate_function)().alias(str(value))
        for value in unq_values
    ]
    return lf.group_by(index, maintain_order=maintain_order).agg(aggs)