Есть ли способ записать имя группы для group_by в полярах? ⇐ Python

[Anonymous]

Есть ли способ получить имя текущей группы для любой заданной итерации экземпляра 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 полярников. В попытке доказать это, вот (надеюсь) действительный тест.
Нативный вызов с использованием аргумента по умолчанию.
%%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, который при необходимости генерирует некоторые отчеты контроля качества на диске, и наличие имени группы необходимо для различения отчеты каждой группы.
# 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

))
)



Подробнее здесь: https://stackoverflow.com/questions/791 ... -in-polars