如何通过将 pandas 的聚合应用到 Polars 来转换嵌套组？

Question

我正在从 pandas 过渡，所以请原谅我的非并行大脑。假设我们有以下 pandas 代码：

import numpy as np
import pandas as pd

df = pd.DataFrame({
    val: np.random.randint(1,5,100) for val in ['a','b','c','d','x','y','z']
})    

df.groupby('a').apply(lambda df:
    df.sort_values('c')
      .groupby('d')
      [['x','y','z']]
      .agg(['max','mean','median'])
)

输出（使用平滑的多重索引将其粘贴到此处）：

a	d	0	1	2	3	4	5	6	7	8
		x	x	x	y	y	y	z	z	z
		总和	意思是	中位数	总和	意思是	中位数	总和	意思是	中位数
1	1	15.0	3.75	4.0	12.0	3.0	3.5	12.0	3.0	3.0
1	2	9.0	3.0	3.0	5.0	1.666667	1.0	9.0	3.0	4.0
1	3	33.0	3.0	3.0	30.0	2.727273	3.0	27.0	2.454545	2.0
1	4	23.0	2.8750	3.0	16.0	2.0	2.0	15.0	1.8750	1.0
2	1	18.0	2.571429	2.0	13.0	1.857143	2.0	18.0	2.571429	3.0
2	2	18.0	2.0	1.0	23.0	2.555556	2.0	25.0	2.777778	3.0
2	3	11.0	3.666667	4.0	9.0	3.0	3.0	9.0	3.0	4.0
2	4	3.0	1.50	1.50	6.0	3.0	3.0	4.0	2.0	2.0
3	1	28.0	2.80	3.0	21.0	2.10	2.0	29.0	2.90	3.0
3	2	13.0	2.166667	2.0	19.0	3.166667	3.0	18.0	3.0	3.0
3	3	16.0	1.777778	2.0	22.0	2.444444	3.0	32.0	3.555556	4.0
3	4	20.0	2.222222	2.0	23.0	2.555556	2.0	23.0	2.555556	3.0
4	1	9.0	2.250	2.0	10.0	2.50	2.50	5.0	1.250	1.0
4	2	19.0	3.166667	3.0	8.0	1.333333	1.0	22.0	3.666667	4.0
4	3	10.0	2.0	1.0	14.0	2.80	3.0	15.0	3.0	3.0
4	4	9.0	2.250	2.0	12.0	3.0	3.0	10.0	2.50	2.50

如何用极坐标重写它？

练习的核心思想是，在

apply

中，我可以对整个数据框组做一些事情，例如对它进行排序，然后聚合（我知道这没有意义，但这个想法是自由地做任何事）。如果我希望我的代码可并行化，我是否会失去这种自由，或者有没有办法捕获整个组？我尝试过

pl.all()

但无法找出至少对每个 sub-df 进行排序的技巧

Answer 1

从问题的后半部分开始，我认为使用子数据帧的 Polars 方法是使用

over

来窗口：

dfplx = pl.DataFrame(dfx)
dfplx.select([
        'a', 'b',
        pl.col(["x", "y", "z"]).sort_by('c').over(['a', 'd']),
        pl.col("x").sort_by('c').mean().over(['a', 'd'].alias("x mean")
])

如果您想在 Windows 中使用 apply，这是完全可能的（

.list()

使极坐标期望来自 lambda 的多个结果）：

dfplx.select([
        'a', 'b', 'c', 
        pl.col(['x', 'y', 'z'])
           .sort_by('c')
           .apply(lambda x: [x.mean(), x.median(), x.max()])
           .list()
           .over(['a', 'd'])
    ])

对于完整的翻译，我发现的最简洁的方法是简单地生成所需的列列表：

dfplx.select(['a', 'b'] + 
            [pl.col(a)
                .sort_by('c') # sort for some reason ;)
                .apply(func) # apply the [max, mean, median] function
                .over(['a', 'd']) # window by a then d
                .alias(f"{a} {label}")  # rename the result
            for a in ['x', 'y', 'z']
            for label, func in [("max", pl.max), ("mean", pl.mean), ("median", pl.median)]
        ])