CUBE SUGAR CONTAINER

技術系のこと書きます。

Python: Polars の DataFrame をゼロコピーで Pandas の DataFrame に変換する

Polars pandas Python

Polars の DataFrame は to_pandas() メソッドを使うことで Pandas の DataFrame に変換できる。 このとき、デフォルトではメモリのコピーが生じる。

pola-rs.github.io

ただし、オプションとして use_pyarrow_extension_array=True を渡すとゼロコピーで変換できる。 これは、メモリ上のデータ表現として Apache Arrow を Polars と Pandas の DataFrame で共有できるため。 今回は、この機能について実際にメモリの使用量をプロファイラで確認しながら使ってみよう。

なお、この機能には注意点もあるようだ。 具体的には、変換後の DataFrame に対して PyArrow でサポートされていない処理を実施した場合、NumPy 配列に変換した上で実行される。 つまり、オンデマンドでメモリのコピーを含むオーバーヘッドが生じることになる。

使った環境は次のとおり。

$ cat /etc/lsb-release 
DISTRIB_ID=Ubuntu
DISTRIB_RELEASE=20.04
DISTRIB_CODENAME=focal
DISTRIB_DESCRIPTION="Ubuntu 20.04.5 LTS"
$ uname -srm
Linux 5.15.0-60-generic x86_64
$ python -V
Python 3.10.10
$ pip list | egrep "(polars|pandas)"
pandas            1.5.3
polars            0.16.9

もくじ

下準備

下準備として pandas と pyarrow のオプションを有効にして Polars をインストールする。 また、memory_profiler もインストールしておこう。

$ pip install "polars[pandas,pyarrow]" memory_profiler

大きめのデータセットとして Higgs データセットをダウンロードしておく。

$ wget https://archive.ics.uci.edu/ml/machine-learning-databases/00280/HIGGS.csv.gz
$ gunzip HIGGS.csv.gz

Polars の DataFrame を Pandas の DataFrame に変換する

まずはオプションを指定せずに pl.DataFrame#to_pandas() を実行してみよう。 実行する main() 関数は memory_profiler のプロファイル対象にする。

#!/usr/bin/env python
# -*- coding: utf-8 -*-

import polars as pl
from memory_profiler import profile


@profile
def main():
    # CSV を読み込んで Polars のデータフレームにする
    pl_df = pl.read_csv("HIGGS.csv")
    # Polars のデータフレームを Pandas のデータフレームに変換する
    pd_df = pl_df.to_pandas()


if __name__ == '__main__':
    main()

上記を実行する。 すると、CSV を Polars の DataFrame に読み込むタイミングでメモリの使用量が 3,412MiB 増加する。 そして、Pandas の DataFrame に変換するタイミングで、さらにメモリの使用量が 2,436MiB 増えている。

$ python withcopy.py 
Filename: /home/amedama/Documents/temporary/withcopy.py

Line #    Mem usage    Increment  Occurrences   Line Contents
=============================================================
     8    100.1 MiB    100.1 MiB           1   @profile
     9                                         def main():
    10                                             # CSV を読み込んで Polars のデータフレームにする
    11   3512.4 MiB   3412.3 MiB           1       pl_df = pl.read_csv("HIGGS.csv")
    12                                             # Polars のデータフレームを Pandas のデータフレームに変換する
    13   5948.5 MiB   2436.1 MiB           1       pd_df = pl_df.to_pandas()

続いては use_pyarrow_extension_array=True を指定する。

#!/usr/bin/env python
# -*- coding: utf-8 -*-

import polars as pl
from memory_profiler import profile


@profile
def main():
    pl_df = pl.read_csv("HIGGS.csv")
    # PyArrow を使ってゼロコピーで Pandas の DataFrame に変換する
    pd_df = pl_df.to_pandas(use_pyarrow_extension_array=True)
    # NOTE: 後続の処理が PyArrow でサポートされていない場合は NumPy 配列に変換されることもある
    ...


if __name__ == '__main__':
    main()

上記を実行する。 すると、今度は Pandas の DataFrame に変換するタイミングでは 2.6MiB しかメモリの使用量が増えていない。

$ python zerocopy.py 
Filename: /home/amedama/Documents/temporary/zerocopy.py

Line #    Mem usage    Increment  Occurrences   Line Contents
=============================================================
     8     99.2 MiB     99.2 MiB           1   @profile
     9                                         def main():
    10   3510.8 MiB   3411.6 MiB           1       pl_df = pl.read_csv("HIGGS.csv")
    11                                             # PyArrow を使ってゼロコピーで Pandas の DataFrame に変換する
    12   3513.4 MiB      2.6 MiB           1       pd_df = pl_df.to_pandas(use_pyarrow_extension_array=True)
    13                                             # 後続の処理が PyArrow でサポートされていない場合は NumPy 配列に変換されることもある
    14   3513.4 MiB      0.0 MiB           1       ...

GC 時のメモリ使用量について確認する

GC (Garbage Collection) が実行されたときの挙動も確認しておこう。 Pandas の DataFrame に変換した後で、Polars の DataFrame への参照を削除する。

まずはオプションを指定しない場合から。

#!/usr/bin/env python
# -*- coding: utf-8 -*-

import polars as pl
from memory_profiler import profile


@profile
def main():
    pl_df = pl.read_csv("HIGGS.csv")
    pd_df = pl_df.to_pandas()

    # Polars のデータフレームの参照を削除する
    # CPython であれば直後にオブジェクトが GC されるはず
    del pl_df


if __name__ == '__main__':
    main()

上記を実行する。 すると、Polars の DataFrame への参照を削除したタイミングでメモリの使用量が 2,433MiB 減少する。 減少したメモリの使用量は少し違和感がある。 というのも、Polars の DataFrame が最初に確保した 3,411MiB ではなく、Pandas の DataFrame に変換時する際に増加した 2,436MiB に近いため。 ここはどういった理屈なのか、まだよく分かっていない。

$ python withcopygc.py 
Filename: /home/amedama/Documents/temporary/withcopygc.py

Line #    Mem usage    Increment  Occurrences   Line Contents
=============================================================
     8     99.3 MiB     99.3 MiB           1   @profile
     9                                         def main():
    10   3510.9 MiB   3411.6 MiB           1       pl_df = pl.read_csv("HIGGS.csv")
    11   5946.9 MiB   2436.0 MiB           1       pd_df = pl_df.to_pandas()
    12                                         
    13                                             # Polars のデータフレームの参照を削除する
    14                                             # CPython であれば直後にオブジェクトが GC されるはず
    15   3513.3 MiB  -2433.6 MiB           1       del pl_df

ひとまず、オプションを付けた場合の検証に進む。 オプションを付けた以外にやっていることは先ほどと同じ。

#!/usr/bin/env python
# -*- coding: utf-8 -*-

import polars as pl
from memory_profiler import profile


@profile
def main():
    pl_df = pl.read_csv("HIGGS.csv")
    pd_df = pl_df.to_pandas(use_pyarrow_extension_array=True)

    # Polars のデータフレームの参照を削除する
    # メモリ内の Arrow 配列への参照は残るためメモリ使用量は変わらない想定
    del pl_df


if __name__ == '__main__':
    main()

上記を実行してみよう。 今度は Polars の DataFrame への参照を削除したタイミングで、メモリの使用量に変化がない。 これは、内部的にデータを共有している Pandas の DataFrame からの参照が残っているためにメモリが開放されなかったためと思われる。

$ python zerocopygc.py 
Filename: /home/amedama/Documents/temporary/zerocopygc.py

Line #    Mem usage    Increment  Occurrences   Line Contents
=============================================================
     8     99.0 MiB     99.0 MiB           1   @profile
     9                                         def main():
    10   3510.4 MiB   3411.4 MiB           1       pl_df = pl.read_csv("HIGGS.csv")
    11   3513.0 MiB      2.6 MiB           1       pd_df = pl_df.to_pandas(use_pyarrow_extension_array=True)
    12                                         
    13                                             # Polars のデータフレームの参照を削除する
    14   3513.0 MiB      0.0 MiB           1       del pl_df

まとめ

今回は Polars の DataFrame をゼロコピーで Pandas の DataFrame に変換する方法について確認した。 注意点もあるので、特性を理解した上で使う必要はあるものの、Pandas のエコシステムを流用できる可能性を秘めた便利な機能だと思う。

Linuxで動かしながら学ぶTCP/IPネットワーク入門

Linuxで動かしながら学ぶTCP/IPネットワーク入門

Amazon