在 polars 中,我如何使用 rank() 来获得每个用户最受欢迎的类别

in polars, how could i use rank() to get most popular category per user

假设我有一个 csv

transaction_id,user,book
1,bob,bookA
2,bob,bookA
3,bob,bookB
4,tim,bookA
5,lucy,bookA
6,lucy,bookC
7,lucy,bookC
8,lucy,bookC

每个用户,我想找到他们最喜欢的书。例如,输出应该是;

shape: (3, 2)
┌──────┬──────────┐
│ user ┆ fav_book │
│ ---  ┆ ---      │
│ str  ┆ str      │
╞══════╪══════════╡
│ bob  ┆ bookA    │
├╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌┤
│ tim  ┆ bookA    │
├╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌┤
│ lucy ┆ bookC    │
└──────┴──────────┘

现在我已经弄清楚了如何做到这一点

import polars as pl

df = pl.read_csv("book_aggs.csv")

print(df)

df2 = df.groupby(["user", "book"]).agg([
  pl.col("book").count(),
  pl.col("transaction_id") # just so we can double check where it all came from - TODO: how to output this to csv?
  ])

print(df2)

df3 = df2.sort(["user", "book_count"], reverse=True).groupby("user").agg([
  pl.col("book").first().alias("fav_book")
])

print(df3)

但真正正常的 sql 方法是 dense_rank 按书籍数量降序排列,其中 rank = 1。我已经尝试了几个小时来让它工作,但我做不到在文档中找到相关示例。

问题是在文档中,none 的聚合示例引用了另一个聚合的输出 - 在这种情况下,它需要引用每个用户每本书的计数,然后将这些计数降序排列然后根据该排序顺序进行排名。

请提供一个示例来说明如何使用排名来执行此任务,以及如何有效地嵌套聚合。

方法一

我们可以先按 user 和 'book' 分组以获得所有 user -> book 组合并计算出现次数最多的组合。

这会给这个中间体 DataFrame:

shape: (5, 3)
┌──────┬───────┬────────────┐
│ user ┆ book  ┆ book_count │
│ ---  ┆ ---   ┆ ---        │
│ str  ┆ str   ┆ u32        │
╞══════╪═══════╪════════════╡
│ lucy ┆ bookC ┆ 3          │
├╌╌╌╌╌╌┼╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┤
│ lucy ┆ bookA ┆ 1          │
├╌╌╌╌╌╌┼╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┤
│ bob  ┆ bookB ┆ 1          │
├╌╌╌╌╌╌┼╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┤
│ tim  ┆ bookA ┆ 1          │
├╌╌╌╌╌╌┼╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┤
│ bob  ┆ bookA ┆ 2          │
└──────┴───────┴────────────┘

然后我们可以做另一个 groupby user 计算 maximum book_count 的索引并将该索引用于 take 正确的 book .

整个查询如下所示:

df = pl.DataFrame({'book': ['bookA',
          'bookA',
          'bookB',
          'bookA',
          'bookA',
          'bookC',
          'bookC',
          'bookC'],
 'transaction_id': [1, 2, 3, 4, 5, 6, 7, 8],
 'user': ['bob', 'bob', 'bob', 'tim', 'lucy', 'lucy', 'lucy', 'lucy']
})

(df.groupby(["user", "book"])
 .agg([
     pl.col("book").count()
 ])
 .groupby("user")
 .agg([
     pl.col("book").take(pl.col("book_count").arg_max()).alias("fav_book")
 ])
)

并创建此输出:

shape: (3, 2)
┌──────┬──────────┐
│ user ┆ fav_book │
│ ---  ┆ ---      │
│ str  ┆ str      │
╞══════╪══════════╡
│ tim  ┆ bookA    │
├╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌┤
│ bob  ┆ bookA    │
├╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌┤
│ lucy ┆ bookC    │
└──────┴──────────┘

方法二

另一种方法是创建一个带有 window_expressionbook_count 列,然后使用最大值的索引 take 聚合中的正确书籍:

(df
 .with_column(pl.count("book").over(["user", "book"]).alias("book_count"))
 .groupby("user")
 .agg([
     pl.col("book").take(pl.col("book_count").arg_max())
 ])
)