获取第一行时对 groupby 表达式进行排序的最佳方法

Question

我想按“foo”对数据帧进行分组，并获取 sorted 组的第一个值。

我有一个解决方案，但它涉及对整个数据框进行排序，我相信在组内排序会快得多。排序的实例会更多，但每个组的大小

n

要小得多，并且排序执行为

n*log(n)

。

df = pl.DataFrame({"foo": [1, 1, 1, 2, 2, 2, 3], "bar": [5, 7, 6, 4, 2, 3, 1]})

df_desired = pl.DataFrame({"foo": [1, 2, 3], "bar": [5, 2, 1]})

df_solution = df.sort("bar").groupby("foo", maintain_order=True).first().sort(by="foo")

assert df_desired.frame_equal(df_solution)

我的建议是一种对每个组进行排序的方法。有这种事吗？

df_suggestion = df.groupby("foo").<sort_groupby(by="bar")>.first()

Answer 1

尝试：

df.group_by("foo").agg(pl.col("bar").sort().first()).sort(by="foo")