使用 Polars 正确分组和过滤

Question

我的工作有 df 3 个主要列：cid1, cid2, cid3，以及更多 7 个列cid4, cid5, etc。

cid1是cid2，int另一列是float。

cid1和的每个组合cid2都是一个包含某些行的工作集，其中所有其他列的值都不同。我想过滤 df 并接收我的 df ，其中仅包含和cid3的每个组合的列中的最大值。下一列必须保持不变。cid1cid2cid4

这段代码帮助我完成了部分任务：

df = (df
    .groupby(["cid1", "cid2"])
    .agg([pl.max("cid3").alias("max_cid3")])
)

它只接收 3 列：cid1、cid2，并在不是最大值max_cid3时过滤所有行。cid3但我无法找到如何cid4, etc在不进行更改的情况下接收该行的所有其他列 ( )。

df = (df
    .groupby(["cid1", "cid2"])
    .agg([pl.max("cid3").alias("max_cid3"), pl.col("cid4")])
)

我尝试添加pl.col("cid4")到 aggs 列表，但在列中我看到某些值的不同列表cid4。

我怎样才能正确地做到这一点？也许 Polars 有另一种方法可以实现 groupby 功能？

在 Pandas 中我可以做到：

import pandas as pd
import numpy as np

df["max_cid3"] = df.groupby(['cid1', 'cid2'])['cid3'].transform(np.max)

然后在任何地方过滤 dfcid3==max_cid3 但我找不到在 Polars 中制作它的方法。

谢谢你！

Answer 1

在极坐标中，您可以使用窗口函数。

\n

df = pl.DataFrame(\n    {"cid1": [1, 2, 2, 1, 3],\n     "cid2": [1, 2, 2, 1, 3],\n     "cid3": [1, 2, 9, 1, 1],\n     "cid4": [4, 5, 6, 7, 8],\n     "cid5": [4, 5, 4, 9, 3]}\n)\n

\n

df.with_columns(max_cid3 = \n   pl.col("cid3").max().over("cid1", "cid2")\n)\n

\n

shape: (5, 6)\n\xe2\x94\x8c\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\xac\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\xac\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\xac\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\xac\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\xac\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x90\n\xe2\x94\x82 cid1 \xe2\x94\x86 cid2 \xe2\x94\x86 cid3 \xe2\x94\x86 cid4 \xe2\x94\x86 cid5 \xe2\x94\x86 max_cid3 \xe2\x94\x82\n\xe2\x94\x82 ---  \xe2\x94\x86 ---  \xe2\x94\x86 ---  \xe2\x94\x86 ---  \xe2\x94\x86 ---  \xe2\x94\x86 ---      \xe2\x94\x82\n\xe2\x94\x82 i64  \xe2\x94\x86 i64  \xe2\x94\x86 i64  \xe2\x94\x86 i64  \xe2\x94\x86 i64  \xe2\x94\x86 i64      \xe2\x94\x82\n\xe2\x95\x9e\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\xaa\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\xaa\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\xaa\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\xaa\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\xaa\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\xa1\n\xe2\x94\x82 1    \xe2\x94\x86 1    \xe2\x94\x86 1    \xe2\x94\x86 4    \xe2\x94\x86 4    \xe2\x94\x86 1        \xe2\x94\x82\n\xe2\x94\x82 2    \xe2\x94\x86 2    \xe2\x94\x86 2    \xe2\x94\x86 5    \xe2\x94\x86 5    \xe2\x94\x86 9        \xe2\x94\x82\n\xe2\x94\x82 2    \xe2\x94\x86 2    \xe2\x94\x86 9    \xe2\x94\x86 6    \xe2\x94\x86 4    \xe2\x94\x86 9        \xe2\x94\x82\n\xe2\x94\x82 1    \xe2\x94\x86 1    \xe2\x94\x86 1    \xe2\x94\x86 7    \xe2\x94\x86 9    \xe2\x94\x86 1        \xe2\x94\x82\n\xe2\x94\x82 3    \xe2\x94\x86 3    \xe2\x94\x86 1    \xe2\x94\x86 8    \xe2\x94\x86 3    \xe2\x94\x86 1        \xe2\x94\x82\n\xe2\x94\x94\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\xb4\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\xb4\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\xb4\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\xb4\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\xb4\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x98\n

\n

可以直接在里面使用.filter()

\n

df.filter(cid3 = pl.col("cid3").max().over("cid1", "cid2"))\n

\n

shape: (4, 5)\n\xe2\x94\x8c\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\xac\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\xac\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\xac\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\xac\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x90\n\xe2\x94\x82 cid1 \xe2\x94\x86 cid2 \xe2\x94\x86 cid3 \xe2\x94\x86 cid4 \xe2\x94\x86 cid5 \xe2\x94\x82\n\xe2\x94\x82 ---  \xe2\x94\x86 ---  \xe2\x94\x86 ---  \xe2\x94\x86 ---  \xe2\x94\x86 ---  \xe2\x94\x82\n\xe2\x94\x82 i64  \xe2\x94\x86 i64  \xe2\x94\x86 i64  \xe2\x94\x86 i64  \xe2\x94\x86 i64  \xe2\x94\x82\n\xe2\x95\x9e\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\xaa\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\xaa\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\xaa\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\xaa\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\xa1\n\xe2\x94\x82 1    \xe2\x94\x86 1    \xe2\x94\x86 1    \xe2\x94\x86 4    \xe2\x94\x86 4    \xe2\x94\x82\n\xe2\x94\x82 2    \xe2\x94\x86 2    \xe2\x94\x86 9    \xe2\x94\x86 6    \xe2\x94\x86 4    \xe2\x94\x82\n\xe2\x94\x82 1    \xe2\x94\x86 1    \xe2\x94\x86 1    \xe2\x94\x86 7    \xe2\x94\x86 9    \xe2\x94\x82\n\xe2\x94\x82 3    \xe2\x94\x86 3    \xe2\x94\x86 1    \xe2\x94\x86 8    \xe2\x94\x86 3    \xe2\x94\x82\n\xe2\x94\x94\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\xb4\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\xb4\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\xb4\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\xb4\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x98\n

\n

熊猫进行比较

\n

>>> df.to_pandas().groupby(["cid1", "cid2"])["cid3"].transform("max")\n0    1\n1    9\n2    9\n3    1\n4    1\nName: cid3, dtype: int64\n

\n