Polars 将 array[string] 列聚合到 set[str] 中

我有极地数据框：

df = pl.DataFrame({
    'col1': [["aaa", "aaa"], ["bbb", "ccc"], ["ccc", "ddd", "ddd"], ["ddd", "ddd", "ddd"]],
    'col2': ["a", "a", "a", "a"],
    'col3': ["x", "x", "y", "y"]
})

我想将 col2、col3 分组并将 col1 聚合到 Set[String] 中

(df
.groupby(["col2", "col3"])
.agg([pl.col("col1").flatten().apply(set).alias("result")])
)

当我在 1700 万条记录上运行它时，它的执行速度非常慢。 10分钟后它仍然没有完成。

如何加快速度？

EDIT:

这就是我解决它的方法，而且速度非常快：

df = (
    df
    .with_columns([
        pl.col("col1").arr.join(",")
    ])
    .groupby(["col2", "col3"])
    .agg([
        pl.col("col1").alias("col1")
    ])
    .with_columns([
        pl.col("col1").arr.join(",")
    ])
    .with_columns([
        pl.col("col1").str.split(",").arr.unique().alias("col1")
    ])
)

┌──────┬──────┬───────────────────────┐
│ col2 ┆ col3 ┆ col1                  │
│ ---  ┆ ---  ┆ ---                   │
│ str  ┆ str  ┆ list[str]             │
╞══════╪══════╪═══════════════════════╡
│ a    ┆ x    ┆ ["aaa", "bbb", "ccc"] │
│ a    ┆ y    ┆ ["ccc", "ddd"]        │
└──────┴──────┴───────────────────────┘

这是我提出的最终解决方案：它使用 Polars 本机函数并且速度非常快。

预处理：将 array[str] 列连接到带有分隔符的字符串中 >> 将字符串列聚合到 array[string] >> 后处理：首先重复将 array[str] 列连接到带有分隔符的字符串中，然后在分隔符上拆分字符串并获取唯一值。

df = (
       df
       .with_columns([
           pl.col("col1").arr.join(",")
       ])
       .groupby(["col2", "col3"])
       .agg([
           pl.col("col1").alias("col1")
       ])
       .with_columns([
            pl.col("col1").arr.join(",")
       ])
       .with_columns([
           pl.col("col1").str.split(",").arr.unique().alias("col1")
       ])
   )
    

这是输出：

┌──────┬──────┬───────────────────────┐
│ col2 ┆ col3 ┆ col1                  │
│ ---  ┆ ---  ┆ ---                   │
│ str  ┆ str  ┆ list[str]             │
╞══════╪══════╪═══════════════════════╡
│ a    ┆ x    ┆ ["aaa", "bbb", "ccc"] │
│ a    ┆ y    ┆ ["ccc", "ddd"]        │
└──────┴──────┴───────────────────────┘