极性将列表[str]列聚合到集合[str]中

Question

我有极坐标数据框：

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
.group_by("col2", "col3")
.agg(pl.col("col1").flatten().map_elements(set).alias("result"))
)

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

如何加快速度？

编辑：

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

df = (
   df
   .with_columns(
       pl.col("col1").list.join(",")
   )
   .group_by("col2", "col3")
   .agg(
       pl.col("col1").alias("col1")
   )
   .with_columns(
       pl.col("col1").list.join(",")
   )
   .with_columns(
       pl.col("col1").str.split(",").list.unique().alias("col1")
   )
)

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

Answer 1

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

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

df = (
   df
   .with_columns(
       pl.col("col1").list.join(",")
   )
   .group_by("col2", "col3")
   .agg(
       pl.col("col1").alias("col1")
   )
   .with_columns(
       pl.col("col1").list.join(",")
   )
   .with_columns(
       pl.col("col1").str.split(",").list.unique().alias("col1")
   )
)

这是输出：

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

Answer 2

如果您不需要集合而只需要元素的统一，那么您可以尝试使用

.unique()

而不是

.apply(set)

:

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

样品结果：

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

我对更大的数据帧（最多 100,000 行）进行了

timeit

基准测试，这个版本始终快了约 10 倍。