Groupby 将两列聚合到 Polars 中的字典中

Question

鉴于以下数据，我正在寻找 groupby 并将两列合并为一列，并持有一本字典。一列提供键，而值来自另一列，该列首先聚合到列表中。

import polars as pl

data = pl.DataFrame(
    {
        "names": ["foo", "ham", "spam", "cheese", "egg", "foo"],
        "dates": ["1", "1", "2", "3", "3", "4"],
        "groups": ["A", "A", "B", "B", "B", "C"],
    }
)

>>> print(data)
    names dates groups
0     foo     1      A
1     ham     1      A
2    spam     2      B
3  cheese     3      B
4     egg     3      B
5     foo     4      C


# This is what i'm trying to do:
  groups                                 combined
0      A                    {'1': ['foo', 'ham']}
1      B  {'2': ['spam'], '3': ['cheese', 'egg']}
2      C                           {'4': ['foo']}

在 pandas 中，我可以使用两个 groupby 语句来做到这一点，在 pyspark 中使用围绕“map_from_entries”的一组操作，但尽管进行了各种尝试，我还没有找到在 polars 中的方法。到目前为止，我使用 agg_list()，转换为 pandas 并使用 lambda。虽然这有效，但肯定感觉不对。

data = data.groupby(["groups", "dates"])["names"].agg_list()

data = (
    data.to_pandas()
    .groupby(["groups"])
    .apply(lambda x: dict(zip(x["dates"], x["names_agg_list"])))
    .reset_index(name="combined")
    )

或者，受此启发，我尝试了多种类似于以下内容的变体，包括将字典转换为 json 字符串等。

data = data.groupby(["groups"]).agg(
    pl.apply(exprs=["dates", "names_agg_list"], f=build_dict).alias("combined")
    )

Answer 1

随着 polars>=0.12.10 的发布，您可以这样做：

print(data
    .groupby(["groups", "dates"]).agg(pl.col("names").list().keep_name())
    .groupby("groups")
    .agg([
        pl.apply([pl.col("dates"), pl.col("names")], lambda s: dict(zip(s[0], s[1].to_list())))
    ])
)

shape: (3, 2)
┌────────┬─────────────────────────────────────┐
│ groups ┆ dates                               │
│ ---    ┆ ---                                 │
│ str    ┆ object                              │
╞════════╪═════════════════════════════════════╡
│ A      ┆ {'1': ['foo', 'ham']}               │
├╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
│ C      ┆ {'4': ['foo']}                      │
├╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
│ B      ┆ {'3': ['cheese', 'egg'], '2': ['... │
└────────┴─────────────────────────────────────┘

虽然这不是您应该如何使用 DataFrames。可能有一种解决方案可以让您处理更扁平的数据帧，并且不需要您将缓慢的 python 对象放入数据帧中。

Groupby 将两列聚合到 Polars 中的字典中

Groupby aggregate two columns into a dictionary in Polars

python-polars