Добавьте новый столбец на основе фильтра + group_by в кадр данных Polars. ⇐ Python

[Anonymous]

Я пытаюсь портировать скрипт pandas на Polars. У меня есть набор данных, который выглядит так
import polars as pl

df = pl.from_repr("""
┌────────────┬─────────┬───────────┬───────┬───────┬─────────────┐
│ sid ┆ roi ┆ endpoint ┆ value ┆ std ┆ voxel_count │
│ --- ┆ --- ┆ --- ┆ --- ┆ --- ┆ --- │
│ str ┆ str ┆ str ┆ f64 ┆ f64 ┆ i64 │
╞════════════╪═════════╪═══════════╪═══════╪═══════╪═════════════╡
│ 4213-a3_bl ┆ AF_L ┆ afd_along ┆ 0.4 ┆ 0.21 ┆ 57334 │
│ 4213-a3_bl ┆ AF_L ┆ radfODF ┆ 0.08 ┆ 0.045 ┆ 57334 │
│ 4213-a3_bl ┆ AF_R ┆ afd_along ┆ 0.42 ┆ 0.22 ┆ 53916 │
│ 4213-a3_bl ┆ AF_R ┆ radfODF ┆ 0.08 ┆ 0.04 ┆ 53916 │
│ 4213-a3_bl ┆ CC_1 ┆ afd_along ┆ null ┆ null ┆ null │
│ 4213-a3_bl ┆ CC_2a ┆ afd_along ┆ 0.54 ┆ 0.3 ┆ 3264 │
│ 4213-a3_bl ┆ wm_mask ┆ null ┆ null ┆ null ┆ 602620 │
│ 4225-a3_bl ┆ CC_2a ┆ radfODF ┆ 0.06 ┆ 0.04 ┆ 3264 │
│ 4225-a3_bl ┆ CC_2b ┆ afd_along ┆ 0.47 ┆ 0.24 ┆ 18833 │
│ 4225-a3_bl ┆ wm_mask ┆ null ┆ null ┆ null ┆ 718758 │
└────────────┴─────────┴───────────┴───────┴───────┴─────────────┘
""")

Я хочу добавить столбец на основе фильтра, за которым следует group_by
df.filter(roi = "wm_mask").group_by("sid").first()

shape: (2, 6)
┌────────────┬─────────┬──────────┬───────┬──────┬─────────────┐
│ sid ┆ roi ┆ endpoint ┆ value ┆ std ┆ voxel_count │
│ --- ┆ --- ┆ --- ┆ --- ┆ --- ┆ --- │
│ str ┆ str ┆ str ┆ f64 ┆ f64 ┆ i64 │
╞════════════╪═════════╪══════════╪═══════╪══════╪═════════════╡
│ 4213-a3_bl ┆ wm_mask ┆ null ┆ null ┆ null ┆ 602620 │
│ 4225-a3_bl ┆ wm_mask ┆ null ┆ null ┆ null ┆ 718758 │
└────────────┴─────────┴──────────┴───────┴──────┴─────────────┘

Теперь я хочу добавить новые значения voxel_count, соответствующие правому sid, что должно дать что-то вроде
┌────────────┬─────────┬───────────┬───────┬───────┬─────────────┬────────────────┐
│ sid ┆ roi ┆ endpoint ┆ value ┆ std ┆ voxel_count ┆ wm_mask__count │
│ --- ┆ --- ┆ --- ┆ --- ┆ --- ┆ --- ┆ --- │
│ str ┆ str ┆ str ┆ f64 ┆ f64 ┆ i64 ┆ i64 │
╞════════════╪═════════╪═══════════╪═══════╪═══════╪═════════════╪════════════════╡
│ 4213-a3_bl ┆ AF_L ┆ afd_along ┆ 0.4 ┆ 0.21 ┆ 57334 ┆ 602620 │
│ 4213-a3_bl ┆ AF_L ┆ radfODF ┆ 0.08 ┆ 0.045 ┆ 57334 ┆ 602620 │
│ 4213-a3_bl ┆ AF_R ┆ afd_along ┆ 0.42 ┆ 0.22 ┆ 53916 ┆ 602620 │
│ 4213-a3_bl ┆ AF_R ┆ radfODF ┆ 0.08 ┆ 0.04 ┆ 53916 ┆ 602620 │
│ 4213-a3_bl ┆ CC_1 ┆ afd_along ┆ null ┆ null ┆ null ┆ 602620 │
│ 4213-a3_bl ┆ CC_2a ┆ afd_along ┆ 0.54 ┆ 0.3 ┆ 3264 ┆ 602620 │
│ 4213-a3_bl ┆ wm_mask ┆ null ┆ null ┆ null ┆ 602620 ┆ 602620 │
│ 4225-a3_bl ┆ CC_2a ┆ radfODF ┆ 0.06 ┆ 0.04 ┆ 3264 ┆ 718758 │
│ 4225-a3_bl ┆ CC_2b ┆ afd_along ┆ 0.47 ┆ 0.24 ┆ 18833 ┆ 718758 │
│ 4225-a3_bl ┆ wm_mask ┆ null ┆ null ┆ null ┆ 718758 ┆ 718758 │
└────────────┴─────────┴───────────┴───────┴───────┴─────────────┴────────────────┘

Можете ли вы рассказать мне, как это выразить в полярных выражениях?
Если это может помочь, я бы сделал это в пандах следующим образом:
pd_df = df.to_pandas()
pd_df = pd_df.set_index("sid", drop=True)

pd_df_wm_volumes = pd_df[pd_df.roi == "wm_mask"].groupby("sid", as_index=True).first()
pd_df["wm_mask__count"] = pd_df_wm_volumes["voxel_count"]

pd_df = pd_df.reset_index(drop=False)


Подробнее здесь: https://stackoverflow.com/questions/725 ... -dataframe