如何在Polars中实现BigQuery的ROW_NUMBER() OVER(PARTITION BY)功能？

我正尝试用Polars库将指定SQL查询重构为Python脚本，但在处理包含ROW_NUMBER()及OVER(PARTITION BY)的语句时遇到问题：对比SQL生成的结果表与Polars生成的DataFrame，1700万行数据中有约25万行结果不匹配。

表结构 #

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']))
)

可行解决方案（来自@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")
)

内容的提问来源于stack exchange，提问作者Solomon Papathoti Leo