How to select the top_k of one column, sorted by another column, within a third column in Polars? - python-polars

Question

I was writing some code and realized this might be a reasonably common operation. I also realized I don't know a clean way to do it.
The questions is: Get the top 5 entries in column1, sorted by column2, within groups given by column3.
If I had to intuit how this would be writen in Polars it'd be:
df.select(pl.col('column1').top_k(n=5, by='column2').over('column3'))
But note that is made up code; it does not work.
Consider this sample data:
import numpy as np
import pandas as pd
import polars as pl
data_size = 10_000_000
np.random.seed = 1
saleValue = np.random.randint(0, 100, data_size)
storeId = np.random.choice([f'Store: {i}' for i in range(200_000)], replace=True, size=data_size)
customerId = np.random.choice([f'Customer: {i}' for i in range(1_000)], replace=True, size=data_size)
df = pd.DataFrame(
dict(storeId=storeId, customerId=customerId, saleValue=saleValue)
).pipe(pl.from_pandas)
It generates a dataframe of the form:
┌───────────────┬───────────────┬───────────┐
│ storeId ┆ customerId ┆ saleValue │
│ --- ┆ --- ┆ --- │
│ str ┆ str ┆ i32 │
╞═══════════════╪═══════════════╪═══════════╡
│ Store: 161472 ┆ Customer: 960 ┆ 29 │
├╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌┤
│ Store: 168620 ┆ Customer: 814 ┆ 21 │
├╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌┤
│ Store: 37904 ┆ Customer: 80 ┆ 61 │
├╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌┤
│ Store: 166077 ┆ Customer: 516 ┆ 23 │
├╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌┤
│ Store: 141748 ┆ Customer: 549 ┆ 58 │
└─///───────────┴─///───────────┴─///───────┘
I'm curious how one would get the top 5 customer per store, sorted by their total spend.
One solution is:
(df
# This part is essential; we need to get the total spend (sales)
.groupby(['storeId','customerId'])
.agg(pl.col('saleValue').sum().alias('totalSales'))
# This is the part I think could be cleaner
.sort('totalSales', reverse=True)
.groupby('storeId')
.agg(pl.col('customerId').head(5).list().alias('customerIds'))
)
┌───────────────┬─────────────────────────────────────┐
│ storeId ┆ customerIds │
│ --- ┆ --- │
│ str ┆ list[str] │
╞═══════════════╪═════════════════════════════════════╡
│ Store: 78152 ┆ ["Customer: 753", "Customer: 170... │
├╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
│ Store: 67676 ┆ ["Customer: 957", "Customer: 896... │
├╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
│ Store: 45152 ┆ ["Customer: 118", "Customer: 127... │
├╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
│ Store: 183339 ┆ ["Customer: 370", "Customer: 227... │
├╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
│ Store: 144688 ┆ ["Customer: 328", "Customer: 294... │
└─///───────────┴─///─────────────────────────────────┘
But I wonder if there is something cleaner using .top_k

Answer 1

It can be done a bit cleaner indeed, although you are close.
Define the first step as df_agg (for the sake of this answer, you can chain):
df_agg = df.groupby(['storeId','customerId']).agg(pl.col('saleValue').sum().alias('totalSales'))
Then we can do:
df_agg.groupby('storeId').agg(pl.col('customerId').sort_by('totalSales', reverse=True).slice(0,5))
Which reads as:
group by store
take column customer id
we sort the values of this column by the column totalSales, from high to low (reverse=True)
we take the first 5 values of customer id
So we do the groupby like you proposed, but the sorting is done inside the aggregation using sort_by, rather than on the full dataframe. Also, I use slice rather than head + list.
On your point on top_k: this function returns the largest elements of itself, not by another. Polars has sufficient ways of achieving what you want, notably sort_by, that I dont think there is a need to complicate the implementation of top_k by adding a by argument.

Answer 2

Would using top_k in a filter work for you? For example:
(
df.groupby(["storeId", "customerId"])
.agg(pl.col("saleValue").sum().alias("totalSales"))
.filter(
pl.col("totalSales")
>= pl.col("totalSales").top_k(k=5).list().over("storeId").arr.last()
)
)
shape: (1048652, 3)
┌───────────────┬───────────────┬────────────┐
│ storeId ┆ customerId ┆ totalSales │
│ --- ┆ --- ┆ --- │
│ str ┆ str ┆ i64 │
╞═══════════════╪═══════════════╪════════════╡
│ Store: 92626 ┆ Customer: 829 ┆ 98 │
├╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┤
│ Store: 56532 ┆ Customer: 840 ┆ 93 │
├╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┤
│ Store: 159073 ┆ Customer: 684 ┆ 88 │
├╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┤
│ Store: 131292 ┆ Customer: 836 ┆ 98 │
├╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┤
│ Store: 73245 ┆ Customer: 545 ┆ 93 │
├╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┤
│ Store: 31163 ┆ Customer: 554 ┆ 91 │
├╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┤
│ Store: 128047 ┆ Customer: 971 ┆ 89 │
├╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┤
│ Store: 41563 ┆ Customer: 85 ┆ 92 │
├╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┤
│ Store: 157951 ┆ Customer: 45 ┆ 97 │
├╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┤
│ Store: 7677 ┆ Customer: 390 ┆ 88 │
├╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┤
│ ... ┆ ... ┆ ... │
Adding a sort to make the result easier to inspect:
(
df.groupby(["storeId", "customerId"])
.agg(pl.col("saleValue").sum().alias("totalSales"))
.filter(
pl.col("totalSales")
>= pl.col("totalSales").top_k(k=5).list().over("storeId").arr.last()
)
.sort(["storeId", "totalSales"], reverse=[False, True])
)
shape: (1048652, 3)
┌──────────────┬───────────────┬────────────┐
│ storeId ┆ customerId ┆ totalSales │
│ --- ┆ --- ┆ --- │
│ str ┆ str ┆ i64 │
╞══════════════╪═══════════════╪════════════╡
│ Store: 0 ┆ Customer: 46 ┆ 151 │
├╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┤
│ Store: 0 ┆ Customer: 267 ┆ 102 │
├╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┤
│ Store: 0 ┆ Customer: 354 ┆ 94 │
├╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┤
│ Store: 0 ┆ Customer: 416 ┆ 93 │
├╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┤
│ Store: 0 ┆ Customer: 729 ┆ 93 │
├╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┤
│ Store: 1 ┆ Customer: 459 ┆ 99 │
├╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┤
│ Store: 1 ┆ Customer: 417 ┆ 90 │
├╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┤
│ Store: 1 ┆ Customer: 982 ┆ 89 │
├╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┤
│ Store: 1 ┆ Customer: 337 ┆ 86 │
├╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┤
│ Store: 1 ┆ Customer: 202 ┆ 84 │
├╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┤
│ ... ┆ ... ┆ ... │
├╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┤
│ Store: 99998 ┆ Customer: 536 ┆ 99 │
├╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┤
│ Store: 99998 ┆ Customer: 295 ┆ 99 │
├╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┤
│ Store: 99998 ┆ Customer: 841 ┆ 98 │
├╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┤
│ Store: 99998 ┆ Customer: 782 ┆ 94 │
├╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┤
│ Store: 99999 ┆ Customer: 29 ┆ 96 │
├╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┤
│ Store: 99999 ┆ Customer: 84 ┆ 96 │
├╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┤
│ Store: 99999 ┆ Customer: 557 ┆ 96 │
├╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┤
│ Store: 99999 ┆ Customer: 885 ┆ 91 │
├╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┤
│ Store: 99999 ┆ Customer: 866 ┆ 89 │
├╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┤
│ Store: 99999 ┆ Customer: 695 ┆ 89 │
└──────────────┴───────────────┴────────────┘
From the above, you can proceed to create a list for each store (or whatever you need).
Note: this approach may produce more than 5 customers per store. (It does not break ties, as shown for Store 99999 at the bottom).