有没有办法在极坐标中捕获“group_by”的组名称？

Question

有没有一种方法可以获取

group_by

实例的任何给定迭代的当前组的名称，而不必依赖于

GroupBy.__iter__()

？

pl.Config().set_fmt_str_lengths(100)
df = pl.DataFrame({
    'fruit':['apple', 'pear', 'peach', 'pear'],
    'origin':['mountain', 'valley', 'beach', 'beach'],
    'color':['red', 'green', 'yellow', 'green']
    
})

def udf(df_group, something="arg def"):
    something = something + " gets modified in the udf"
    return df_group.with_columns(pl.lit(something).alias('from_udf'))


print(df.group_by('origin').map_groups(udf))

shape: (4, 4)
┌───────┬──────────┬────────┬──────────────────────────────────┐
│ fruit ┆ origin   ┆ color  ┆ from_udf                         │
│ ---   ┆ ---      ┆ ---    ┆ ---                              │
│ str   ┆ str      ┆ str    ┆ str                              │
╞═══════╪══════════╪════════╪══════════════════════════════════╡
│ pear  ┆ valley   ┆ green  ┆ arg def gets modified in the udf │
│ apple ┆ mountain ┆ red    ┆ arg def gets modified in the udf │
│ peach ┆ beach    ┆ yellow ┆ arg def gets modified in the udf │
│ pear  ┆ beach    ┆ green  ┆ arg def gets modified in the udf │
└───────┴──────────┴────────┴──────────────────────────────────┘


# Desired functionality
print(df.group_by('origin').map_groups(udf, something=__name_of_group__))

shape: (4, 4)
┌───────┬──────────┬────────┬──────────────────────────────────┐
│ fruit ┆ origin   ┆ color  ┆ from_udf                         │
│ ---   ┆ ---      ┆ ---    ┆ ---                              │
│ str   ┆ str      ┆ str    ┆ str                              │
╞═══════╪══════════╪════════╪══════════════════════════════════╡
│ pear  ┆ valley   ┆ green  ┆ valley gets modified in the udf  │
│ apple ┆ mountain ┆ red    ┆ mountain gets modified in the udf│
│ peach ┆ beach    ┆ yellow ┆ beach gets modified in the udf   │
│ pear  ┆ beach    ┆ green  ┆ beach gets modified in the udf   │
└───────┴──────────┴────────┴──────────────────────────────────┘

该功能可以通过手动迭代

.group_by()

对象来实现，但我的直觉告诉我，这对性能来说是不行的，因为迭代被委托给 python 而不是 Polars。试图在这里证明它是一个（希望）有效的基准。

使用默认参数进行本机调用。

%%timeit
df.group_by('origin').map_groups(udf)
308 µs ± 19.5 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)

使用默认参数进行手动迭代。

%%timeit 
pl.concat([ udf(data) for name, data in df.group_by(['origin'])])
473 µs ± 43.5 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)

使用传递的参数进行手动迭代。

%%timeit 
pl.concat([ udf(data, something=name[0])for name, data in df.group_by(['origin'])])
486 µs ± 33.8 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)

这个问题的起源是一个更复杂的问题，其中我编写了一个 Polars 插件，可以选择在磁盘上生成一些 QC 报告，并且需要组的名称来区分每个组的报告。

# real application example (out of context)
(
   df
   .group_by('umi').agg(
              rogtk.assemble_contig( # my plugin function
                expr=pl.col("r2_seq"),
                k=20,
                min_coverage=20,
                export_graphs=True,
                prefix=??? # here is where I would like to pass an argument based on the group's definition

  ))
)

Answer 1

因此，基本上您想使用自定义函数单独处理每个组。您的 udf 还需要知道组名称。但是您可以在应用 udf 之前动态地将组名称注入到每个组中。以下是 oyu 如何做到这一点：

import polars as pl

df = pl.DataFrame({
    'fruit': ['apple', 'pear', 'peach', 'pear'],
    'origin': ['mountain', 'valley', 'beach', 'beach'],
    'color': ['red', 'green', 'yellow', 'green']
})

def udf(df_group, group_name):
    return df_group.with_columns(pl.lit(group_name + " gets modified in the udf").alias('from_udf'))

result = (
    df
    .group_by('origin')
    .map_groups(lambda group: udf(group, group['origin'][0]))
)

print(result)

这给出了

shape: (4, 4)
┌───────┬──────────┬────────┬─────────────────────────────────┐
│ fruit ┆ origin   ┆ color  ┆ from_udf                        │
│ ---   ┆ ---      ┆ ---    ┆ ---                             │
│ str   ┆ str      ┆ str    ┆ str                             │
╞═══════╪══════════╪════════╪═════════════════════════════════╡
│ peach ┆ beach    ┆ yellow ┆ beach gets modified in the udf  │
│ pear  ┆ beach    ┆ green  ┆ beach gets modified in the udf  │
│ apple ┆ mountain ┆ red    ┆ mountain gets modified in the … │
│ pear  ┆ valley   ┆ green  ┆ valley gets modified in the ud… │
└───────┴──────────┴────────┴─────────────────────────────────┘

但是，还有一个但是：它不是 Polars 的原生功能。我的解决方案是一个 Pythonic 解决方法，使用 Polars 现有的 API 来实现目标。