Расчленение и повторение строк на основе числа, содержащегося в другом столбце Polars. ⇐ Python

[Anonymous]

У меня есть этот DataFrame:

Код: Выделить всё

import polars as pl

df = pl.from_repr("""
┌────────────┬────────────┬───────┐
│ numberType ┆ NumberInst ┆ Type  │
│ ---        ┆ ---        ┆ ---   │
│ i64        ┆ str        ┆ str   │
╞════════════╪════════════╪═══════╡
│ 1          ┆ None       ┆ Car   │
│ 2          ┆ 1          ┆ Bus   │
│ 3          ┆ 1          ┆ Plane │
└────────────┴────────────┴───────┘
""")

Я сделал это, чтобы разнести и дублировать строки:

Код: Выделить всё

df = df.with_columns(pl.col("NumberInst").alias("numberInstExclude"))
df = df.select(pl.exclude("numberInstExclude").repeat_by("numberType").explode())

Что дает мне:

Код: Выделить всё

┌────────────┬────────────┬───────┐
│ numberType ┆ NumberInst ┆ Type  │
│ ---        ┆ ---        ┆ ---   │
│ i64        ┆ str        ┆ str   │
╞════════════╪════════════╪═══════╡
│ 1          ┆ None       ┆ Car   │
│ 2          ┆ 1          ┆ Bus   │
│ 2          ┆ 1          ┆ Bus   │
│ 3          ┆ 1          ┆ Plane │
│ 3          ┆ 1          ┆ Plane │
│ 3          ┆ 1          ┆ Plane │
└────────────┴────────────┴───────┘

Результат, который я хочу получить:

Код: Выделить всё

┌────────────┬────────────┬───────┬───────────────┐
│ numberType ┆ NumberInst ┆ Type  ┆ NumberInstRep │
│ ---        ┆ ---        ┆ ---   ┆ ---           │
│ i64        ┆ str        ┆ str   ┆ i64           │
╞════════════╪════════════╪═══════╪═══════════════╡
│ 1          ┆ None       ┆ Car   ┆ 1             │
│ 2          ┆ 1          ┆ Bus   ┆ 1             │
│ 2          ┆ 1          ┆ Bus   ┆ 2             │
│ 3          ┆ 1          ┆ Plane ┆ 1             │
│ 3          ┆ 1          ┆ Plane ┆ 2             │
│ 3          ┆ 1          ┆ Plane ┆ 3             │
└────────────┴────────────┴───────┴───────────────┘

Как создать столбец NumberInstRep?
Что я хочу сделать:

Код: Выделить всё

if NumberInst is None then NumberInstRep = 1
else use numberType to have NumberInstRep

Спасибо.

Подробнее здесь: https://stackoverflow.com/questions/796 ... ars-column