在 Polars 中：BigQuery SQL 中 row_number() over(partition by) 的正确等效代码是什么？

Question

我正在尝试使用 Polars 库将给定的 SQL 查询重构（翻译）为 Python 脚本。
我陷入一行查询，其中使用

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：

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：

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 创建的数据框进行比较，在 1700 万行中，大约有 250,000 行不匹配。

那么有没有更好的方法来处理这种

ROW_NUMBER() OVER(PARTITION BY)

情况？

Answer 1

你可以使用

df.with_columns(
    rn = pl.int_range(pl.len()).over(…)
)