【Polars】polars.whenで文字列型(Utf8)のカラムを生成する時の注意

polarsでは、pl.whenを用いることで条件分岐によるカラム生成を行うことができる。SQLのcase when文に相当するものである。
特定のカラムの条件に応じて、新たに文字列型(pl.Utf8)のカラムを生成するケースにおいて、polarsのversionの違いでハマったことがあったので、備忘録として残しておく。

まずは polars==0.18.3 の場合。
以下のように書くことで、文字列の値を持つ colc が生成される。

import polars as pl

df = pl.DataFrame({
    'cola': [1, 2, 2, 3]
    , 'colb': [4, 5, 6, 7]
})

print((
    df
    .with_columns(
        pl.when(pl.col('cola') == 1).then('hoge')
        .otherwise('fuga')
        .alias('colc')
    )
))

# output
# shape: (4, 3)
# ┌──────┬──────┬──────┐
# │ cola ┆ colb ┆ colc │
# │ ---  ┆ ---  ┆ ---  │
# │ i64  ┆ i64  ┆ str  │
# ╞══════╪══════╪══════╡
# │ 1    ┆ 4    ┆ hoge │
# │ 2    ┆ 5    ┆ fuga │
# │ 2    ┆ 6    ┆ fuga │
# │ 3    ┆ 7    ┆ fuga │
# └──────┴──────┴──────┘

しかし、上記のコードをpolars==0.20.2で動かすと、以下のようなエラーが出た。

ColumnNotFoundError: hoge

Error originated just after this operation:
DF ["cola", "colb"]; PROJECT */2 COLUMNS; SELECTION: "None"

どうやらthen('hoge') が then(pl.col('hoge')) というふうに認識されているようである。当然データフレームにhogeというカラムは存在しないのでエラーとなる。

対処として、以下のようにpl.litを用いることで、どちらのversionでも期待通りの動作となる。

import polars as pl

df = pl.DataFrame({
    'cola': [1, 2, 2, 3]
    , 'colb': [4, 5, 6, 7]
})

print((
    df
    .with_columns(
        pl.when(pl.col('cola') == 1).then(pl.lit('hoge'))
        .otherwise(pl.lit('fuga'))
        .alias('colc')
    )
))

# output
# shape: (4, 3)
# ┌──────┬──────┬──────┐
# │ cola ┆ colb ┆ colc │
# │ ---  ┆ ---  ┆ ---  │
# │ i64  ┆ i64  ┆ str  │
# ╞══════╪══════╪══════╡
# │ 1    ┆ 4    ┆ hoge │
# │ 2    ┆ 5    ┆ fuga │
# │ 2    ┆ 6    ┆ fuga │
# │ 3    ┆ 7    ┆ fuga │
# └──────┴──────┴──────┘

polarsは破壊的変更がガンガン行われるので注意が必要である。