В Polars: каков правильный эквивалентный код для row_number() over(partition by) в BigQuery SQL? ⇐ Python

[Anonymous]

Я пытаюсь выполнить рефакторинг (перевести) заданный SQL-запрос в скрипт Python, используя библиотеку Polars.

Я застрял в одной строке запроса, где используется функция ROW_NUMBER(), за которой следует функция OVER(PARTITION BY).
Ниже приведена схема таблицы:
product_id (INTEGER)
variant_id (INTEGER)
client_code (VARCHAR)
transaction_date (DATE)
customer_id (INTEGER)
store_id (INTEGER)
invoice_id (VARCHAR)
invoice_line_id (INTEGER)
quantity (NUMERIC)
net_sales_price (NUMERIC)

Ниже приведен SQL-запрос:
SELECT
product_id,
variant_id,
client_code,
transaction_date,

ROW_NUMBER() OVER(
PARTITION BY
product_id, variant_id, store_id, customer_id, client_code
ORDER BY
transaction_date ASC,
invoice_id ASC,
invoice_line_id ASC,
quantity DESC,
net_sales_price ASC
) AS repeat_purchase_seq

FROM transactions

Я пробовал несколько способов, например:
пример 1: использование pl.first().cum_count().over()
new_df = (
df
.sort(['product_id', 'variant_id', 'store_id', 'customer_id', 'client_code','transaction_date', 'invoice_id', 'invoice_line_id',pl.col('quantity').reverse(), 'net_sales_price'])
.with_columns(repeat_purchase_seq = pl.first().cum_count().over(['product_id', 'variant_id', 'store_id', 'customer_id', 'client_code']).flatten())
)

пример 2: использование pl.rank('ordinal').over()
new_df = (
df
.sort(['transaction_date', 'invoice_id', 'invoice_line_id', 'quantity', 'net_sales_price'], descending=[False, False, False, True, False])
.with_columns(repeat_purchase_seq = pl.struct('transaction_date', 'invoice_id', 'invoice_line_id', 'quantity', 'net_sales_price').rank('ordinal').over(['product_id', 'variant_id', 'store_id', 'customer_id', 'client_code']))
)

Оба примера имеют ту или иную проблему.

Я попытался сравнить таблицу, созданную SQL, с фреймом данных, созданным с помощью Polars, из 17 миллионов строк. , имеется около 250 000 несовпадающих строк.
Итак, есть ли лучший способ справиться с этой ситуацией ROW_NUMBER() OVER(PARTITION BY)?
Изменить. Ниже приведен ответ @roman, который помог в моем случае:
partition_by_keys = ["product_id", "variant_id", "store_id", "customer_id", "client_code"]
order_by_keys = ["transaction_date", "invoice_id", "invoice_line_id", "quantity", "net_sales_price"]
order_by_descending = [False, False, False, True, False]

order_by = [-pl.col(col) if desc else pl.col(col) for col, desc in zip(order_by_keys, order_by_descending)]

df.with_columns(
pl.struct(order_by)
.rank("ordinal")
.over(partition_by_keys)
.alias("rn")
)


Подробнее здесь: https://stackoverflow.com/questions/792 ... artition-b