Как выбрать 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

Стоит ли изучать PHP в 2023-2024 годах?
Стоит ли изучать PHP в 2023-2024 годах?
Привет всем, сегодня я хочу высказать свои соображения по поводу вопроса, который я уже много раз получал в своем сообществе: "Стоит ли изучать PHP в...
Поведение ключевого слова "this" в стрелочной функции в сравнении с нормальной функцией
Поведение ключевого слова "this" в стрелочной функции в сравнении с нормальной функцией
В JavaScript одним из самых запутанных понятий является поведение ключевого слова "this" в стрелочной и обычной функциях.
Приемы CSS-макетирования - floats и Flexbox
Приемы CSS-макетирования - floats и Flexbox
Здравствуйте, друзья-студенты! Готовы совершенствовать свои навыки веб-дизайна? Сегодня в нашем путешествии мы рассмотрим приемы CSS-верстки - в...
Тестирование функциональных ngrx-эффектов в Angular 16 с помощью Jest
В системе управления состояниями ngrx, совместимой с Angular 16, появились функциональные эффекты. Это здорово и делает код определенно легче для...
Концепция локализации и ее применение в приложениях React ⚡️
Концепция локализации и ее применение в приложениях React ⚡️
Локализация - это процесс адаптации приложения к различным языкам и культурным требованиям. Это позволяет пользователям получить опыт, соответствующий...
Пользовательский скаляр GraphQL
Пользовательский скаляр GraphQL
Листовые узлы системы типов GraphQL называются скалярами. Достигнув скалярного типа, невозможно спуститься дальше по иерархии типов. Скалярный тип...
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 внизу).

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