Как выбрать top_k одного столбца, отсортированного по другому столбцу, в третьем столбце в Polars?

Я писал код и понял, что это может быть довольно распространенная операция. Я также понял, что не знаю чистого способа сделать это.

Вопросы: получить 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

Как настроить Tailwind CSS с React.js и Next.js?
Как настроить Tailwind CSS с React.js и Next.js?
Tailwind CSS - единственный фреймворк, который, как я убедился, масштабируется в больших командах. Он легко настраивается, адаптируется к любому...
LeetCode запись решения 2536. Увеличение подматриц на единицу
LeetCode запись решения 2536. Увеличение подматриц на единицу
Увеличение подматриц на единицу - LeetCode
Переключение светлых/темных тем
Переключение светлых/темных тем
В Microsoft Training - Guided Project - Build a simple website with web pages, CSS files and JavaScript files, мы объясняем, как CSS можно...
Отношения "многие ко многим" в Laravel с методами присоединения и отсоединения
Отношения "многие ко многим" в Laravel с методами присоединения и отсоединения
Отношения "многие ко многим" в Laravel могут быть немного сложными, но с помощью Eloquent ORM и его моделей мы можем сделать это с легкостью. В этой...
В PHP
В PHP
В большой кодовой базе с множеством различных компонентов классы, функции и константы могут иметь одинаковые имена. Это может привести к путанице и...
Карта дорог Беладжар PHP Laravel
Карта дорог Беладжар PHP Laravel
Laravel - это PHP-фреймворк, разработанный для облегчения разработки веб-приложений. Laravel предоставляет различные функции, упрощающие разработку...
1
0
141
2
Перейти к ответу Данный вопрос помечен как решенный

Ответы 2

Ответ принят как подходящий

Это действительно можно сделать немного чище, хотя вы близки.

Определите первый шаг как 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))

Что читается как:

  • группировать по магазинам
  • взять идентификатор клиента столбца
  • мы сортируем значения этого столбца по столбцу totalSales, от большего к меньшему (reverse=True)
  • мы берем первые 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 внизу).

Другие вопросы по теме