CUBE SUGAR CONTAINER

技術系のこと書きます。

Python: Pandas 2 系ではデータ型のバックエンドを変更できる

Pandas の 2 系から、新たにデータ型のバックエンドという考え方が導入された。 これは、端的にいうと DataFrame のデータをどのような形式で持つかを表している。 たとえば Pandas 2.0.0 の時点では、次の 3 つからバックエンドを選ぶことができる。

  • NumPy (デフォルト)
  • NumPy Nullable
  • PyArrow

何も指定しないときに選ばれるデフォルトの NumPy は 1 系と使い勝手が変わらない。 このエントリでは、バックエンドを切り替えたときに何が起こるのかを解説する。 また Pandas の内部実装についても軽く紹介する。

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

$ sw_vers   
ProductName:        macOS
ProductVersion:     13.3
BuildVersion:       22E252
$ python -V         
Python 3.10.11
$ pip list | grep pandas
pandas          2.0.0

もくじ

下準備

あらかじめ Pandas と PyArrow をインストールしておく。 PyArrow についてはインストールしておかないとバックエンドに指定できない。

$ pip install -U pandas pyarrow

欠損値の扱いについて

分かりやすいところで欠損値の扱いについて紹介する。 以下のサンプルコードでは「fruits」と「price」という 2 つのカラムを持った CSV ファイルを読み込んでいる。 CSV ファイルの読み込みには pandas.read_csv() 関数を使う。 Pandas 2 系では、この関数に dtype_backend という引数が追加された。 この引数に "numpy_nullable" や "pyarrow" を指定することでバックエンドを変更できる。 ちなみに pandas.read_csv() 以外のデータを読み込む関数にも、同様に dtype_backend が追加された。 なお、既存の DataFrame のバックエンドを変更したいときは DataFrame#convert_dtypes() を使えば良い。

import io

import pandas as pd


def main():
    # CSV フォーマットのファイルライクオブジェクトを用意する
    csv_data = """
    fruits,price
    apple,198
    banana,
    cherry,398
    """
    csv_file_obj = io.StringIO(csv_data)

    # デフォルトの NumPy バックエンド
    default_df = pd.read_csv(
        csv_file_obj,
    )
    print("Default NumPy backend")
    print(default_df)
    print(default_df.dtypes)
    print("-" * 60)

    csv_file_obj.seek(0)
    # Nullable NumPy バックエンド
    # 浮動小数点以外でも欠損値を持てる
    nullable_df = pd.read_csv(
        csv_file_obj,
        dtype_backend="numpy_nullable",
    )
    print("Nullable NumPy backend")
    print(nullable_df)
    print(nullable_df.dtypes)
    print("-" * 60)

    csv_file_obj.seek(0)
    # PyArrow バックエンド
    # これも浮動小数点以外で欠損値を持てる
    pyarrow_df = pd.read_csv(
        csv_file_obj,
        dtype_backend="pyarrow",
    )
    print("PyArrow backend")
    print(pyarrow_df)
    print(pyarrow_df.dtypes)
    print("-" * 60)


if __name__ == "__main__":
    main()

「price」には欠損値が含まれるが、バックエンドによってどのような違いが出るだろうか。 上記の実行結果は次のとおり。

$ python backend.py 
Default NumPy backend
       fruits  price
0       apple  198.0
1      banana    NaN
2      cherry  398.0
    fruits     object
price         float64
dtype: object
------------------------------------------------------------
Nullable NumPy backend
       fruits  price
0       apple    198
1      banana   <NA>
2      cherry    398
    fruits    string[python]
price                  Int64
dtype: object
------------------------------------------------------------
PyArrow backend
       fruits  price
0       apple    198
1      banana   <NA>
2      cherry    398
    fruits    string[pyarrow]
price          int64[pyarrow]
dtype: object
------------------------------------------------------------

デフォルトの NumPy バックエンドでは欠損値があることでデータ型が暗黙に浮動小数点型 (float64) にキャストされてしまっている。 一方で NumPy Nullable や PyArrow では整数型 (Int64 / int64[pyarrow]) が維持されている。 このように、データ型のバックエンドを切り替えると振る舞いが変わることがある。

内部実装について

ここからは、データ型のバックエンドがどのように実現されているのか内部実装を見ていこう。 先ほどのサンプルコードで使った変数を Python の REPL で操作できるようにした。

NumPy バックエンド

まずはデフォルトの NumPy バックエンドを使った DataFrame から見ていく。

DataFrame のデータは、内部的に DataFrame#_data という変数に格納されている。 このオブジェクトは BlockManager と呼ばれる。 BlockManager には複数のブロックが登録されており、それぞれのブロックがカラムに対応している。

>>> default_df._data
BlockManager
Items: Index(['    fruits', 'price'], dtype='object')
Axis 1: RangeIndex(start=0, stop=3, step=1)
ObjectBlock: slice(0, 1, 1), 1 x 3, dtype: object
NumericBlock: slice(1, 2, 1), 1 x 3, dtype: float64

上記から、デフォルトの NumPy バックエンドで作った DataFrame の BlockManager には object 型の ObjectBlock と float64 型の NumericBlock が登録されていることがわかる。

それぞれのブロックには DataFrame#_data#blocks という配列を経由してアクセスできる。 さらに、データの実装には values という変数でアクセスできる。

>>> default_df._data.blocks[0].values
array([['    apple', '    banana', '    cherry']], dtype=object)
>>> default_df._data.blocks[1].values
array([[198.,  nan, 398.]])

上記から、デフォルトの NumPy バックエンドで作った場合、ここに NumPy 配列が直接入ることが分かる。

NumPy Nullable バックエンド

続いては NumPy Nullable バックエンドを見ていこう。 先ほどと同じように DataFrame#_data で BlockManager にアクセスする。

>>> nullable_df._data
BlockManager
Items: Index(['    fruits', 'price'], dtype='object')
Axis 1: RangeIndex(start=0, stop=3, step=1)
ExtensionBlock: slice(0, 1, 1), 1 x 3, dtype: string
ExtensionBlock: slice(1, 2, 1), 1 x 3, dtype: Int64

すると、今度は string 型の ExtensionBlock と Int64 型の ExtensionBlock が登録されていることがわかる。

先ほどと同じように DataFrame#_data#blocks の配列から values 変数にアクセスしてみる。 先ほどは NumPy 配列が直接入っていたところだ。

>>> nullable_df._data.blocks[0].values
<StringArray>
['    apple', '    banana', '    cherry']
Length: 3, dtype: string
>>> type(nullable_df._data.blocks[0].values)
<class 'pandas.core.arrays.string_.StringArray'>
>>> nullable_df._data.blocks[1].values
<IntegerArray>
[198, <NA>, 398]
Length: 3, dtype: Int64
>>> type(nullable_df._data.blocks[1].values)
<class 'pandas.core.arrays.integer.IntegerArray'>

今度は NumPy 配列が入っているのではなく StringArray や IntegerArray という型のオブジェクトが入っていた。

この中で IntegerArray について、さらに _data_mask という変数にアクセスしてみよう。

>>> nullable_df._data.blocks[1].values._data
array([                 198, -9223372036854775808,                  398])
>>> nullable_df._data.blocks[1].values._mask
array([False,  True, False])

すると、今度は NumPy 配列が入っている。 これは _data の方が整数値を入れておく配列で _mask の方が欠損値かどうかを表した真偽値の配列になっている。 つまり「欠損値かどうかを表した配列」をデータとは別に用意することで、暗黙のキャストをしなくても欠損値を表現できるようにしている。

PyArrow バックエンド

PyArrow バックエンドについても同様に見ていこう。

>>> pyarrow_df._data
BlockManager
Items: Index(['    fruits', 'price'], dtype='object')
Axis 1: RangeIndex(start=0, stop=3, step=1)
ExtensionBlock: slice(0, 1, 1), 1 x 3, dtype: string[pyarrow]
ExtensionBlock: slice(1, 2, 1), 1 x 3, dtype: int64[pyarrow]

今度は string[pyarrow] 型の ExtensionBlock と int64[pyarrow] 型の ExtensionBlock が BlockManager に登録されている。

それぞれのブロックの values 変数にアクセスすると ArrowExtensionArray という型のオブジェクトになっている。

>>> pyarrow_df._data.blocks[0].values
<ArrowExtensionArray>
['    apple', '    banana', '    cherry']
Length: 3, dtype: string[pyarrow]
>>> type(pyarrow_df._data.blocks[0].values)
<class 'pandas.core.arrays.arrow.array.ArrowExtensionArray'>
>>> pyarrow_df._data.blocks[1].values
<ArrowExtensionArray>
[198, <NA>, 398]
Length: 3, dtype: int64[pyarrow]
>>> type(pyarrow_df._data.blocks[1].values)
<class 'pandas.core.arrays.arrow.array.ArrowExtensionArray'>

さらに _data 変数にアクセスすると、ここに pyarrow.lib.ChunkedArray 型のオブジェクトが入っている。

>>> pyarrow_df._data.blocks[1].values._data
<pyarrow.lib.ChunkedArray object at 0x13fcb6f70>
[
  [
    198,
    null,
    398
  ]
]

上記は PyArrow パッケージで提供されている Arrow フォーマットの配列になる。 このオブジェクトは見て分かるとおり整数型であっても欠損値を null で表現できる。

まとめ

今回は Pandas 2 系で導入されたデータ型のバックエンドという概念を紹介した。 それぞれのバックエンドはデータの持ち方が違うことから、内部的な処理のやり方についても異なっている。 もちろん、処理にかかる時間もそれぞれで異なる。

ちなみに Pandas 2.0.0 の時点で PyArrow バックエンドは最適化の余地が多く残されているようだ。 要するに、デフォルトの NumPy バックエンドの方が高速に動作するケースも多い。 これは Pandas の GitHub リポジトリに登録された Issue から確認できる。 もちろん、この点は今後バージョンが進むごとに改善されていくはずだ。

いじょう。