Polarsのsortはどこに実装されているのか ~デバッグの仕方を添えて~
はじめに
この記事はPolars Adbent Calendar 2023 22日目の記事です。
本記事では、Python版のPolarsについて話します。
Polarsは動作が早いです。Polarsが高速に動作する要因は複数ありますが、その一つに高速に動作するRustで書かれているから、というものがあります。
Python上でPolarsを扱う際にもこのRustの恩恵にあずかれるということは、Rustの実装をPython上で呼び出せるような仕組みがどこかにあるはずです。
この、別の言語間で相互で呼び出せるようにする仕組みのことをバインディングと呼びます。今回はPythonとRust間のため、Python↔︎Rustバインディングの仕組みが存在することになります。
PolarsではPython↔︎Rustバインディングとして、PyO3とmaturinが利用されています。
Polarsの実装を追っていると、このバインディングの仕組みを理解しないと実装を追えない場面が多々出てきます。今回はsortを例に、このPyO3/maturinがどのように用いられているか理解し、Polarsの実装を可能な限り追ってみることにします。
注: 私のRustの知識不足もあり、Rust側の実装を追うパートは怪しい内容になっています。わかる方いたら追記しますので是非教えてください。
PyO3とmaturinとは何か？
PyO3はRustとPythonのバインディング機能を提供するライブラリです。
maturinはRustをPythonパッケージとしてビルドし、publishするためのライブラリです。以下の図がその関係性を分かりやすく表しています。引用元
実装が追えなくなる場面
今回は、polars.DataFrameのsort機能を使う場面で説明します。以下のコードはPolarsのバージョン0.20.2です。
例えば、以下のようなコードがあるとします。
import polars as pl df = pl.DataFrame("../input/titanic.csv") df.sort("fare")
titanic.csvは、Kaggleなどで用いられるタイタニックの乗船客のデータです。
タイタニックの乗船客のデータをfare（料金）ごとにsortする非常にシンプルなコードです。
では、df.sort("fare")の中ではどのような実装をしているかを見ていきましょう。
VSCodeでのデバッグを行なっていきます。df.sort("fare")の行にbreakpointを打ち、↓（Step Into）を行います。
Step Intoをすると、Polarsライブラリの実装を見に行くことになります。ただし、VSCodeのデバッグは、デフォルトでは外部ライブラリの実装を見に行かないため、VSCodeのデバッグの設定ファイルであるlaunch.jsonの設定を変更します。
まず、下記の歯車ボタンを押します。
すると、launch.jsonが立ち上がるので、justMyCodeの部分をtrueからfalseに変更しましょう。
無事、Polarsの実装を見に行けると、以下のような画面になると思います。
sortの中では
(
self.lazy()
.sort(by, *more_by, descending=descending, nulls_last=nulls_last)
.collect(_eager=True)
)
というようなコードが行われていたんですね。.lazy()と.collect(_eager=True)はlazy APIのためのコードのため、sortは.sort(by, *more_by, descending=descending, nulls_last=nulls_last)で行われていることになります。
では.sortの中身を見てみましょう。Step Over(左から２番目、時計回りの矢印）を一度クリックし、.sortがハイライトされた状態にします。
上記画面になったら、再度↓（Step Into）を行います。今回は5-6回↓ボタンを押す必要があると思います。すると以下のような画面にいきます。
おそらく、sortの処理は、self._ldf.sort(by, descending, nulls_last, maintain_order)で行われていそうだと分かります。sortはself._ldfのメソッドのようですが、self._ldfとはなんでしょうか？こういう時はVSCodeの「Go To Definition」が便利です。_ldfをダブルクリックし、右クリックで表示されるメニューから「Go To Definition」を選択します。
すると、_ldf: PyLazyFrameという実装が見つかります。
今度はPyLazyFrameが何かという話になりますね。再度「Go To Definition」を行います。
あれ？追えません。仕方ないのでCtrl + FでPyLazyFrameを検索しましょう。
すると、以下のようなコードが見つかりますが、VSCodeのシンタックスハイライトが機能していません。
そう、このfrom polars.polars import PyLazyFrameこそが、PolarsがRustで書かれた実装をPyO3/maturinを用いてPythonモジュールとして呼び出されている箇所なのです。
つまり、sortの実装の本体を追いたければ、from polars.polars import PyLazyFrameがどこから呼び出されているのかを知り、PyLazyFrameに対してsortメソッドを用いた場合にどのコードが呼ばれるかを知る必要があるわけです。これは大変だ。
Polarsにおいて、PyO3/maturinが使われている箇所、およびその読み方
というわけで、課題感としては「Python側の実装で呼び出されているRustの実装が追えない」ということにありました。 ここからはRust側の実装を追ってみましょう。
まず前提知識なのですが、Rustはパッケージを作る際、Cargo.tomlというそれらのクレート（バイナリかライブラリ）をどのようにビルドするかを説明するファイルを持ちます。また、src/main.rs が、パッケージと同じ名前を持つバイナリクレートのクレートルートであるという慣習があります。同じようにCargoはパッケージディレクトリにsrc/lib.rs が含まれていたら、パッケージにはパッケージと同じ名前のライブラリクレートが含まれており、src/lib.rs がそのクレートルートなのだと判断します。
参考: doc.rust-jp.rs
src配下にsrc/main.rsかsrc/lib.rsのどちらかがあれば、そこからパッケージが呼び出されているはず、ということになります。
では、GitHubの画面で見てみましょう。Linkはこちら
src/main.rsはなく、src/lib.rsがあるこが分かります。
では次に、src/lib.rsの中身を確認します。上から辿っていくと、fn polarsというそれっぽい実装がされている箇所が見つかります。
ここが、PyO3によってPython↔︎Rustバインディングを行なっている箇所になります。#[pymodule]で指定されたモジュールがPythonにエクスポートされます。m.add_class::<PyLazyFrame>().unwrap();とあるため、PyLazyFrameもここでpolarsのクラスとして追加されていることが分かります。つまり、このPyLazyFrameが、先ほどPython側でも利用したモジュールになるはずです。
参考:
PyLazyFrameがどこから呼び出されているかを再度検索すると、use crate::lazyframe::PyLazyFrame;と呼び出されていました。
Rustでcrateから呼び出されている場合は、クレートルート（この場合はsrc/lib.rs）からの呼び出しになるため、crate::lazyframeはsrc/lazyframe.rsで実装されていることになります。
そのため、今度はsrc/lazyframe.rsを見ていきます。Linkはこちら
コードを上から追っていくと、sortメソッドらしきものが見つかります。
「これが我々が追い求めていたsortの実装！？」と喜ぶのは束の間、このsortはself.ldfのメソッドを呼んでいるだけのメソッドだと分かります。
このself.ldfを検索を行うと、以下のself.ldf = LazyFrame::from(lp);という実装で呼び出されていそうです。そしてこのlpはLogicalPlanという型で実装されています。
（ここからはRustの知識不足に伴い、かなり怪しいです。間違っていたらぜひ教えてください。）
このLogicalPlanがどこで実装されているかを調べていくと以下のようなコードが見つかります。Linkはこちら
このコードの意味をChatGPTに聞くと、
LogicalPlanBuilder という名前の公開構造体（pub struct）を定義しています。この構造体は、LogicalPlan（論理計画）を公開フィールドとして持っており、これにより他のコードからアクセス可能になります。
と回答がありました。つまり、LogicalPlangはLogicalPlanBuilderによって構築されると考えられます。
このLogicalPlanBuilderの実装の中から、sortメソッドらしきものを探すと、以下が見つかります。コードはこちら
このコードをChatGPTに説明してもらうと
このコードは、Polarsデータフレームライブラリの一部で、LogicalPlanBuilderのsortメソッドを定義しています。このメソッドはソート操作の論理計画を作成するためのものです。パラメータとして列のベクトル（by_column）、各列が降順かどうかを示すブール値のベクトル（descending）、null値を最後にするかどうかのブール値（null_last）、そしてソートした際に元の順序を維持するかどうかのブール値（maintain_order）を取ります。これらのパラメータを使用して、ソートの論理計画を表すLogicalPlan::Sort構造体を生成し、それを元にデータフレームに対するソート操作の実行計画が立てられます。
とあるので、どうやらこれが我々が追い求めていたsortの実装・・・なんですかね？Sort構造体を作る、とあるので、まだ実装は別の部分にある気がします。
が、これ以上は追い方が分からないため、今回はここまでとします！
最後に
もしsortがどこで実装されているか分かる方がいらっしゃったらぜひ教えてください🙇♂️
