如何使用 group_by 并通过 Polars 应用自定义函数？

Question

我正在绞尽脑汁试图弄清楚如何使用

group_by

并使用 Polars 应用自定义功能。

来自 Pandas，我正在使用：

import pandas as pd
from scipy.stats import spearmanr
 
def get_score(df):
   return spearmanr(df["prediction"], df["target"]).correlation

df = pd.DataFrame({
    "era": [1, 1, 1, 2, 2, 2, 5],
    "prediction": [2, 4, 5, 190, 1, 4, 1],
    "target": [1, 3, 2, 1, 43, 3, 1]
})

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

Polars 有

map_groups()

可以在组上应用自定义函数，我尝试过：

correlations = df.group_by("era").map_groups(get_score)

但是失败并显示错误消息：

'无法获取DataFrame属性'_df'。确保返回一个 DataFrame 对象。： PyErr { type: , value: AttributeError("'float' object has no attribute '_df'"), Traceback: None }

有什么想法吗？

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

对单个表达式应用自定义函数。

下面是我们如何使用自定义函数和普通极坐标表达式对列求平方的示例。表达式语法 应该始终是首选，因为它要快得多。

(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

  • 输入的语义含义：组