Я писал код и понял, что это может быть довольно распространенная операция. Я также понял, что не знаю чистого способа сделать это.
Вопросы: получить 5 лучших записей в столбце 1, отсортированных по столбцу 2, в группах, заданных столбцом 3.
Если бы мне пришлось интуитивно догадываться, как это будет написано в полярах, это было бы так:
df.select(pl.col('column1').top_k(n=5, by='column2').over('column3'))
Но обратите внимание, что это вымышленный код; это не работает.
Рассмотрим этот образец данных:
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)
Он генерирует фрейм данных вида:
┌───────────────┬───────────────┬───────────┐
│ 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 │
└─///───────────┴─///───────────┴─///───────┘
Мне любопытно, как можно получить 5 лучших клиентов в магазине, отсортированных по их общим расходам.
Одно решение:
(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... │
└─///───────────┴─///─────────────────────────────────┘
Но мне интересно, есть ли что-то более чистое, используя .top_k
Это действительно можно сделать немного чище, хотя вы близки.
Определите первый шаг как df_agg
(ради этого ответа вы можете связать):
df_agg = df.groupby(['storeId','customerId']).agg(pl.col('saleValue').sum().alias('totalSales'))
Тогда мы можем сделать:
df_agg.groupby('storeId').agg(pl.col('customerId').sort_by('totalSales', reverse=True).slice(0,5))
Что читается как:
Итак, мы делаем groupby, как вы предложили, но сортировка выполняется внутри агрегации с использованием sort_by
, а не на полном фрейме данных. Кроме того, я использую slice
, а не head
+ list
.
К вашему сведению top_k
: эта функция возвращает самые большие элементы самой себя, а не другой. У Polars есть достаточно способов добиться того, чего вы хотите, особенно sort_by
, поэтому я не думаю, что есть необходимость усложнять реализацию top_k
, добавляя аргумент by
.
Подойдет ли вам использование top_k
в фильтре? Например:
(
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 │
├╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┤
│ ... ┆ ... ┆ ... │
Добавление сортировки для облегчения проверки результата:
(
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 │
└──────────────┴───────────────┴────────────┘
Исходя из вышеизложенного, вы можете приступить к созданию списка для каждого магазина (или того, что вам нужно).
Примечание: этот подход может привлечь более 5 покупателей в магазин. (Это не разрывает связи, как показано для Магазина 99999 внизу).