您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.groupby("g", maintain_order=True)
.agg(
pl.apply(
exprs=["a", "b"],
f=get_score).alias("corr")
))
Polars 提供的功能
(df.groupby("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.groupby("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]
- 输入的语义含义:列中的单个值
groupby
语境:
map
- 输入/输出类型:
Series
- 输入的语义含义:一个列表列,其中值是组
apply