如何在 Polars 中使用 groupby 和 apply？

Question

我正在绞尽脑汁尝试在Python的Polars库中使用groupby和apply。从Pandas转移而来时，我曾经使用：

def get_score(df):
   return spearmanr(df["prediction"], df["target"]).correlation

correlations = df.groupby("era").apply(get_score)

但是在 polars 中，这种方法行不通。

我尝试了几种方法，主要集中在:

correlations = df.groupby("era").apply(get_score)

但是这会出现以下错误信息：

“无法获取DataFrame属性'_df'。确保返回DataFrame对象：PyErr {type: <class 'AttributeError'>, value: AttributeError（“'numpy.float64' object has no attribute '_df'”）,

有什么想法吗？

- jbssm

---

请问您能提供一个样例吗？ - Corralien

Answer 1

从 polars>=0.10.4 版本开始，您可以使用 pl.spearman_rank_corr 函数。

如果您想使用自定义函数，可以按照以下方式进行操作：

在多列/表达式上使用自定义函数

import polars as pl
from typing import List
from scipy import stats

df = pl.DataFrame({
    "g": [1, 1, 1, 2, 2, 2, 5],
    "a": [2, 4, 5, 190, 1, 4, 1],
    "b": [1, 3, 2, 1, 43, 3, 1]
})

def get_score(args: List[pl.Series]) -> pl.Series:
    return pl.Series([stats.spearmanr(args[0], args[1]).correlation], dtype=pl.Float64)

(df.group_by("g", maintain_order=True)
 .agg(
    pl.apply(
        exprs=["a", "b"], 
        function=get_score).alias("corr")
 ))

Polars提供的函数

(df.group_by("g", maintain_order=True)
 .agg(
     pl.spearman_rank_corr("a", "b").alias("corr")
 ))

两者的输出：

shape: (3, 2)
┌─────┬──────┐
│ g   ┆ corr │
│ --- ┆ ---  │
│ i64 ┆ f64  │
╞═════╪══════╡
│ 1   ┆ 0.5  │
├╌╌╌╌╌┼╌╌╌╌╌╌┤
│ 2   ┆ -1e0 │
├╌╌╌╌╌┼╌╌╌╌╌╌┤
│ 5   ┆ NaN  │
└─────┴──────┘

在单个列/表达式上应用自定义函数

我们还可以通过.apply或.map在单个表达式上应用自定义函数。

下面是一个示例，展示了如何使用自定义函数和普通的polar表达式来对一个列进行平方运算。应该始终优先选择表达式语法，因为它的速度更快。

(df.group_by("g")
 .agg(
     pl.col("a").apply(lambda group: group**2).alias("squared1"),
     (pl.col("a")**2).alias("squared2")
 ))

apply

和map之间有什么区别？

map适用于整个列series。而apply根据上下文适用于单个值或单个组。

select上下文：

• map
  • 输入/输出类型：Series
  • 输入的语义含义：列值
• apply
  • 输入/输出类型：Union[int, float, str, bool]
  • 输入的语义含义：列中的单个值

group_by上下文：

• map
  • 输入/输出类型：Series
  • 输入的语义含义：一个列表列，其中值是组
• apply
  • 输入/输出类型：Series
  • 输入的语义含义：组

- ritchie46

---

感谢 @ritchie46。这个完美解决了。 - jbssm

我正在努力理解 polars，它看起来很棒。你的函数为什么要使用 args[0] 和 args[1] 而不是 x 和 y？两者都可以使用，通常命名参数被认为更符合 Python 风格。你的方式是否符合 polars 的惯用法？ - SamR