Make Your Naive Code Fast with Polars

The Naive pandas Approach

This is our first take: the naive approach. This is how someone who is unfamiliar with pandas might approach this problem. We want to calculate a weighted average for each group? Well Pandas has a .groupby(...).apply method that lets us pass each group into a UDF, so we'll go ahead and write a function that operates on 2 Series–one of weights and another of values–and returns an average after multiplying them together. Easy, right?

def weighted_mean(weights, values):
    return (weights * values).sum() / weights.sum()


with timer.time('pandas naive udf') as t:
    df.groupby('groups').apply(
        lambda g: weighted_mean(g['weights'], g['values']),
        include_groups=False
    )
    
print(f'{t.name} {t.elapsed:.6f}s')

pandas naive udf 0.147143s

Making the most of vectorized operations

Considering the above was a little slow, we might come to realize that our multiplication step is fairly redundant. Why multiply each of these Series per group when we can simply create an intermediate column of the weights * values?

If we sacrifice a little bit of extra memory and then pass a UDF that calculates the mean of these intermediates, we get a HUGE boost in speed!

with timer.time('pandas intermediate UDF') as t:
    (
        df.assign(
            intermediate=lambda d: d['weights'] * d['values']
        )
        .groupby('groups').apply(
            lambda group: group['intermediate'].sum() / group['weights'].sum(),
            include_groups=False
        )
    )

print(f'{t.name} {t.elapsed:.6f}s')

pandas intermediate UDF 0.095506s

Removing All Fastpath Barriers

When it comes to restricted computation domains like pandas, we pay a heavy price for crossing the barrier between Python functions and functions that exist in within that domain. We typically want to push as much of our computation into that domain to best leverage any optimizations or fastpaths that are available.

In the above example, we passed a user-defined function that simply called .mean on each group of the intermediate column we created. pandas cannot "see" into this UDF, and thus cannot leverage any fast paths when performing this operation. If instead we use .groupby(...).mean() or .groupby(...).agg('mean') we can make much better use of pandas internals than we did before to get even better performance.

with timer.time('pandas intermediate fastpath') as t:
    (
        df.assign(
            intermediate=lambda d: d['weights'] * d['values']
        )
        .groupby('groups').agg({'intermediate': 'sum', 'weights': 'sum'})
        .eval('intermediate / weights')
    )

print(f'{t.name} {t.elapsed:.6f}s')

pandas intermediate fastpath 0.010572s

Polars Naive Expressions

The thing I enjoy about Polars expression syntax is that they help ensure that the naive approach (like the one we took in the first pandas example) is fast. From a logical understanding, "I want to calculate a weighted mean for each group", it makes sense to use that approach. The expression API helps this method of thinking translate into fast operations.

from polars import from_pandas, col

pl_df = from_pandas(df)

with timer.time('polars naive expression') as t:
    (
        pl_df
        .group_by('groups', maintain_order=True)
        .agg((col('values') * col('weights')).sum() / col('weights').sum())
    )

print(f'{t.name} {t.elapsed:.6f}s')

polars naive expression 0.008849s

If you have been paying close attention- you'll notice that I just recreated my weighted_mean function just using Polars expressions. As it turns out, I did not need to perform this rewrite and instead can use my previous UDF directly provided that it works with expressions. We can verify that it does since calling weighted_mean with column expressions as the input returns a new valid expression:

wmean_expr = weighted_mean(weights=col('weights'), values=col('values'))
wmean_expr

Which means we can use this expression directly in our Polars syntax!

(
    pl_df.group_by('groups', maintain_order=True).agg(wmean_expr)
).head(5)

We can still create and use UDFs with no performance impact, so long as they take in expressions and output expressions.

Polars is not better because it was written in Rust

Both of these tools have similar performance hits when they encounter opaque barriers. If we use the same .groupby(...).apply(UDF) approach from pandas and have it evaluated by Polars, we see that we lose nearly ALL benefit of the gains that Polars provided to us. Let this be a note that when you're operating within a restricted computation domain, and want your code to run fast, you must play by the rules of that restricted computation domain. In this case, Polars wants you to use its expression API as much as possible, so use it.

from polars import from_pandas, col

with timer.time('polars naive UDF') as t:
    (
        pl_df
        .group_by('groups', maintain_order=True)
        .map_groups(lambda d: weighted_mean(d.select('weights'), d.select('values')))
    )

print(f'{t.name} {t.elapsed:.6f}s')

polars naive UDF 0.297142s