主页 / user-28176357

armeya vaidya's questions

Martin Hope
armeya vaidya
Asked: 2024-11-07 22:52:47 +0800 CST
  • 8

我想根据极坐标数据框中的某个条件增加一列值，同时考虑满足该条件的次数。

示例数据。

import polars as pl

df = pl.DataFrame({
    "before": [0, 0, 0, 0, 0, 0, 0, 0, 0],
    "cdl_type": ["REC", "REC", "GEC", None, None, "GEC", None, "REC", "GEC"],
})

当前的方法。

df = df.with_columns(
    a=(
        pl.when(pl.col("cdl_type").is_in(["GEC", "REC"])).then(
            pl.int_ranges(
                pl.col("cdl_type")
                .is_in(["REC", "GEC"])
                .rle()
                .struct.field("len")
            ).flatten()
        )
        .when(pl.col('cdl_type').is_null().and_(pl.col('cdl_type').shift(1).is_not_null()))
        .then(pl.lit(1))
        .otherwise(0)
    )
)

预期输出。

┌────────┬──────────┬───────┐
│ before ┆ cdl_type ┆ after │
│ ---    ┆ ---      ┆ ---   │
│ i64    ┆ str      ┆ i64   │
╞════════╪══════════╪═══════╡
│ 0      ┆ REC      ┆ 0     │
│ 0      ┆ REC      ┆ 1     │
│ 0      ┆ GEC      ┆ 2     │
│ 0      ┆ null     ┆ 3     │
│ 0      ┆ null     ┆ 0     │
│ 0      ┆ GEC      ┆ 0     │
│ 0      ┆ null     ┆ 1     │
│ 0      ┆ REC      ┆ 0     │
│ 0      ┆ GEC      ┆ 1     │
└────────┴──────────┴───────┘
python