开发Polars分组描述统计扩展，简化group_by.agg统计量生成

以下是一个演示用的Polars DataFrame：

import polars as pl

df = pl.DataFrame(
    {
        "groups": ["A", "A", "A", "B", "B", "B"],
        "values": [1, 2, 3, 4, 5, 6],
    }
)

当前使用group_by.agg()生成描述统计量的方式较为繁琐：

print(
    df.group_by("groups").agg(
        pl.len().alias("count"),
        pl.col("values").mean().alias("mean"),
        pl.col("values").std().alias("std"),
        pl.col("values").min().alias("min"),
        pl.col("values").quantile(0.25).alias("25%"),
        pl.col("values").quantile(0.5).alias("50%"),
        pl.col("values").quantile(0.75).alias("75%"),
        pl.col("values").max().alias("max"),
        pl.col("values").skew().alias("skew"),
        pl.col("values").kurtosis().alias("kurtosis"),
    )
)

输出结果：

shape: (2, 11)
┌────────┬───────┬──────┬─────┬───┬─────┬─────┬──────┬──────────┐
│ groups ┆ count ┆ mean ┆ std ┆ … ┆ 75% ┆ max ┆ skew ┆ kurtosis │
│ ---    ┆ ---   ┆ ---  ┆ --- ┆   ┆ --- ┆ --- ┆ ---  ┆ ---      │
│ str    ┆ u32   ┆ f64  ┆ f64 ┆   ┆ f64 ┆ i64 ┆ f64  ┆ f64      │
╞════════╪═══════╪══════╪═════╪═══╪═════╪═════╪══════╪══════════╡
│ B      ┆ 3     ┆ 5.0  ┆ 1.0 ┆ … ┆ 6.0 ┆ 6   ┆ 0.0  ┆ -1.5     │
│ A      ┆ 3     ┆ 2.0  ┆ 1.0 ┆ … ┆ 3.0 ┆ 3   ┆ 0.0  ┆ -1.5     │
└────────┴───────┴──────┴─────┴───┴─────┴─────┴──────┴──────────┘

我们可以通过扩展Polars的方法，实现两种更简洁的调用方式来得到相同结果：

• 直接在DataFrame上调用：df.describe(by="groups", percentiles=[0.25, 0.5, 0.75], skew=True, kurt=True)
• 在分组对象上调用：df.group_by("groups").describe(percentiles=[0.25, 0.5, 0.75], skew=True, kurt=True)

自定义扩展实现代码 #

from typing import Optional, List, Union
from polars import Expr

# 生成描述统计表达式的工具函数
def _get_describe_exprs(
    percentiles: List[float] = [0.25, 0.5, 0.75],
    skew: bool = False,
    kurt: bool = False
) -> List[Expr]:
    exprs = [
        pl.len().alias("count"),
        pl.col(pl.NUMERIC_DTYPES).mean().alias("mean"),
        pl.col(pl.NUMERIC_DTYPES).std().alias("std"),
        pl.col(pl.NUMERIC_DTYPES).min().alias("min"),
        pl.col(pl.NUMERIC_DTYPES).max().alias("max"),
    ]
    # 添加自定义分位数
    for p in percentiles:
        exprs.append(pl.col(pl.NUMERIC_DTYPES).quantile(p).alias(f"{int(p*100)}%"))
    # 添加偏度和峰度（按需开启）
    if skew:
        exprs.append(pl.col(pl.NUMERIC_DTYPES).skew().alias("skew"))
    if kurt:
        exprs.append(pl.col(pl.NUMERIC_DTYPES).kurtosis().alias("kurtosis"))
    return exprs

# 给DataFrame/LazyFrame添加describe方法
def df_describe(
    self,
    by: Optional[Union[str, List[str]]] = None,
    percentiles: List[float] = [0.25, 0.5, 0.75],
    skew: bool = False,
    kurt: bool = False
) -> pl.DataFrame:
    exprs = _get_describe_exprs(percentiles, skew, kurt)
    if by is not None:
        return self.group_by(by).agg(exprs)
    # 无分组时返回长格式统计结果
    return self.select(exprs).melt(variable_name="statistic")

# 给GroupBy/LazyGroupBy添加describe方法
def groupby_describe(
    self,
    percentiles: List[float] = [0.25, 0.5, 0.75],
    skew: bool = False,
    kurt: bool = False
) -> pl.DataFrame:
    exprs = _get_describe_exprs(percentiles, skew, kurt)
    return self.agg(exprs)

# 注册方法到Polars类
pl.DataFrame.describe = df_describe
pl.LazyFrame.describe = df_describe
pl.GroupBy.describe = groupby_describe
pl.LazyGroupBy.describe = groupby_describe

验证效果 #

现在就可以用两种方式调用：

方式1：DataFrame直接调用 #

result1 = df.describe(by="groups", percentiles=[0.25, 0.5, 0.75], skew=True, kurt=True)
print(result1)

方式2：分组后调用 #

result2 = df.group_by("groups").describe(percentiles=[0.25, 0.5, 0.75], skew=True, kurt=True)
print(result2)

两种方式都会输出和原始繁琐写法完全一致的结果。

扩展说明 #

• 自动识别所有数值类型列，无需手动指定统计列
• 支持自定义分位数列表，默认值为[0.25, 0.5, 0.75]
• 通过skew和kurt参数控制是否计算偏度和峰度
• 同时支持DataFrame、LazyFrame以及对应的分组对象

内容的提问来源于stack exchange，提问作者Kevin Li