Polars でセルの書き換えを頑張る

セルの書き換え、私の苦手な言葉です

Pandasでよくやる df[idx, col] = value のように値を書き換える機能は、PolarsのDataFrameでは削除されたり、再実装されたりしています。また、LazyFrameには実装されていません。
こういったセルの書き換えは、あまりPolarsっぽくないのですが、値の書き換えをPolarsで多用せざるを得ないとき、どうしたらいいかを考えてみます。

なお、この記事でいう「書き換え」とは、元データの直接的な変更ではなく、新たなDataFrame (もしくはLazyFrame) を得ることを指しています。
(そのため、元データを直接変更できる方法と比べ、オーバーヘッドは大きくなります)

方法1: 別の形式に変換する

pd.DataFrame, np.ndarray, dictなど、なんでも好きな形式に変換して書き換えてから、Polarsに戻しましょう。

import polars as pl

df = pl.DataFrame(...)

# PandasのDataFrameに変換してから戻す
tmp = df.to_pandas()
tmp[idx, col] = value
df = pl.from_pandas(tmp)

# numpyのndarrayに変換してから戻す
arr = df.to_numpy()
schema = df.schema
arr[i, j] = value
df = pl.DataFrame(arr, schema)

# dictに変換してから戻す
tmp = df.to_dict()
tmp[col][idx] = value
df = pl.DataFrame(tmp)

書き換えるセルが多数な場合は、このように別の形式に変換してから行った方がいいのかもしれません。

方法2: (列内の同じ値を全部書き換えていい場合は) Expr.replace(...) を用いる

Expr.replaceは列内の値を一括で置き換えるエクスプレッションです。

df = pl.DataFrame(
    {
        "service": ["twitter", "instagram", "twitter"],
        "user_id": ["foo_official", "foo_official", "foo_campaign"],
    }
)

df = df.with_columns(pl.col("service").replace("twitter", "X"))

列内の該当する値を全て書き換えたい場合や、該当する値が1つしかないことが分かっている場合には、この方法が便利です。

方法3: (条件に当てはまる行を全部書き換えていい場合は) pl.when(...) を用いる

方法2は列が特定の値ならば別の値に置き換えるという単純なものです。より複雑な条件を指定するには pl.when(<condition>).then(<value>).otherwise(<value>) を用います。

df = pl.DataFrame(
    {
        "name": names,
        "email": emails,
    }
)

df = df.with_columns(
        pl.when(
            pl.col("name") == name,
        ).then(
            pl.lit(new_email),
        ).otherwise(
            pl.col("email"),
        ).alias("email")
    )

ただし、name が一意でない(同じ name が複数行ありえる)場合、全ての行が書き換わることに注意して下さい。これを避けるための方法が、次に示すものです。

方法4: インデックス列を追加して pl.when(...) を用いる

IDのような一意的なキーを持つ列があれば、方法3で行を特定して値を書き換えることができます。
そのようなキーがない場合や、キーを引くのが面倒な場合は、インデックス列を追加して pl.when(...) を用いることができます。

df = pl.DataFrame(
    {
        "name": names,
        "email": emails,
    }
)

# idx行目(0始まり)を書き換える
df = df.with_row_index().with_columns(
        pl.when(
            pl.col("index") == idx,
        ).then(
            pl.lit(new_email),
        ).otherwise(
            pl.col("email"),
        ).alias("email")
    ).drop("index")

インデックス列は、DataFrame.with_row_index()のほか、pl.int_range(pl.len())によっても作ることができます。

df = pl.DataFrame(
    {
        "name": names,
        "email": emails,
    }
)

# idx行目(0始まり)を書き換える
df = df.with_columns(
        pl.when(
            pl.int_range(pl.len()) == idx,
        ).then(
            pl.lit(new_email),
        ).otherwise(
            pl.col("email"),
        ).alias("email")
    )

方法5: 書き換えがたくさん発生するなら join を用いる

書き換えがたくさんあるなら、書き換え用のDataFrameを用意して、それをjoinする方法もあります。

df = pl.DataFrame(
    {
        "name": names,
        "email": emails,
    }
).with_row_index()

# 更新データのDataFrame
# indexを元のdfと同じ型にしないといけないため、schema_overrideを入れている
df_update = pl.DataFrame(
    {
        "index": indices,
        "email": new_emails,
    },
    schema_overrides={"index": df.schema["index"]},
)

# df_updateは、indexに重複がないことを想定しているが、もし重複がありうる場合は重複削除を行う
# (例えば、以下では、一番最後の行のみを抽出している)
df_update = df_update.group_by("index").last()

df = df.join(
        df_update,
        on="index",
        how="left",
        coalesce=True,
    ).with_columns(
        pl.coalesce(["email_right", "email"]).alias("email"),
    ).drop(
        "index",
        "email_right",
    )

方法6: [unstable] join のかわりに DataFrame.update(...) を用いる

この方法はversion 1.17.1現在はunstableですが、上記と同等のことをDataFrame.update(...)で行うことができます。

df = pl.DataFrame(
    {
        "name": names,
        "email": emails,
    }
).with_row_index()

# 更新データのDataFrame
# indexを元のdfと同じ型にしないといけないため、schema_overrideを入れている
df_update = pl.DataFrame(
    {
        "index": indices,
        "email": new_emails,
    },
    schema_overrides={"index": df.schema["index"]},
)

# df_updateは、indexに重複がないことを想定しているが、もし重複がありうる場合は重複削除を行う
# (例えば、以下では、一番最後の行のみを抽出している)
df_update = df_update.group_by("index").last()

df = df.update(df_update, on="index").drop("index")

この記事について

記事の初版執筆当初、筆者は、 pl.DataFrame.__setitem__ のドキュメントが見当たらなかったため、セルの書き換えは行えないと思って記事を書いておりました。しかし、コメントで書き換えが行えることを指摘いただいたため、記事の構成を修正しました。