CUBE SUGAR CONTAINER

技術系のこと書きます。

Python: DuckDB と Polars を組み合わせて使う

DuckDB のバージョン 0.7 で Polars とのインテグレーションが強化された。 たとえば Polars の DataFrame に対して DuckDB の Python API で直接クエリを発行できるようになったらしい。 また、DuckDB で実行したクエリの結果を Polars の DataFrame に変換することもできる。 ユースケースとしては、たとえばすべてインメモリで扱うのが辛くなったときに読み書きする先を DuckDB に切り替えるなどが考えられる。 あるいは、データの加工をなるべく SQL でやりたい場合にもマッチしそうだ。 今回は、その機能を実際に試してみることにする。

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

$ sw_vers                        
ProductName:        macOS
ProductVersion:     13.2.1
BuildVersion:       22D68
$ python -V
Python 3.10.10
$ pip list | egrep -i "(polars|duckdb)"
duckdb            0.7.0
polars            0.16.5

もくじ

下準備

あらかじめ DuckDB の Python API と Polars をインストールしておく。 Polars のインストール時はオプションとして PyArrow を選択する。

$ pip install duckdb "polars[pyarrow]"

サンプルに使う Diamonds データセットの CSV をダウンロードしておく。

$ wget https://raw.githubusercontent.com/mwaskom/seaborn-data/master/diamonds.csv

中身はこんな感じ。

$ head diamonds.csv 
"carat","cut","color","clarity","depth","table","price","x","y","z"
0.23,"Ideal","E","SI2",61.5,55,326,3.95,3.98,2.43
0.21,"Premium","E","SI1",59.8,61,326,3.89,3.84,2.31
0.23,"Good","E","VS1",56.9,65,327,4.05,4.07,2.31
0.29,"Premium","I","VS2",62.4,58,334,4.2,4.23,2.63
0.31,"Good","J","SI2",63.3,58,335,4.34,4.35,2.75
0.24,"Very Good","J","VVS2",62.8,57,336,3.94,3.96,2.48
0.24,"Very Good","I","VVS1",62.3,57,336,3.95,3.98,2.47
0.26,"Very Good","H","SI1",61.9,55,337,4.07,4.11,2.53
0.22,"Fair","E","VS2",65.1,61,337,3.87,3.78,2.49

続いて Python のインタプリタを起動する。

$ python

Diamonds データセットの CSV ファイルを Polars の DataFrame に読み込む。

>>> import polars as pl
>>> df = pl.read_csv("diamonds.csv")
>>> df.head()
shape: (5, 10)
┌───────┬─────────┬───────┬─────────┬─────┬───────┬──────┬──────┬──────┐
│ carat ┆ cut     ┆ color ┆ clarity ┆ ... ┆ price ┆ x    ┆ y    ┆ z    │
│ ---   ┆ ---     ┆ ---   ┆ ---     ┆     ┆ ---   ┆ ---  ┆ ---  ┆ ---  │
│ f64   ┆ strstrstr     ┆     ┆ i64   ┆ f64  ┆ f64  ┆ f64  │
╞═══════╪═════════╪═══════╪═════════╪═════╪═══════╪══════╪══════╪══════╡
│ 0.23  ┆ Ideal   ┆ E     ┆ SI2     ┆ ... ┆ 3263.953.982.43 │
│ 0.21  ┆ Premium ┆ E     ┆ SI1     ┆ ... ┆ 3263.893.842.31 │
│ 0.23  ┆ Good    ┆ E     ┆ VS1     ┆ ... ┆ 3274.054.072.31 │
│ 0.29  ┆ Premium ┆ I     ┆ VS2     ┆ ... ┆ 3344.24.232.63 │
│ 0.31  ┆ Good    ┆ J     ┆ SI2     ┆ ... ┆ 3354.344.352.75 │
└───────┴─────────┴───────┴─────────┴─────┴───────┴──────┴──────┴──────┘

DuckDB で Polars の DataFrame に対してクエリを発行する

では、DuckDB の Python API を使って Polars の DataFrame にクエリを発行してみよう。

まずは duckdb パッケージをインポートする。

>>> import duckdb

あとは ducbdb.sql() 関数を使ってクエリを発行できる。 ここで SQL の from 句に、現在のスコープにある Polars の DataFrame の名前を指定できる。

>>> duckdb.sql("select * from df")

┌────────┬───────────┬─────────┬─────────┬───┬────────┬────────┬────────┐
│ carat  │    cut    │  color  │ clarity │ … │   x    │   y    │   z    │
│ double │  varchar  │ varchar │ varchar │   │ double │ double │ double │
├────────┼───────────┼─────────┼─────────┼───┼────────┼────────┼────────┤
│   0.23 │ Ideal     │ E       │ SI2     │ … │   3.953.982.43 │
│   0.21 │ Premium   │ E       │ SI1     │ … │   3.893.842.31 │
│   0.23 │ Good      │ E       │ VS1     │ … │   4.054.072.31 │
│   0.29 │ Premium   │ I       │ VS2     │ … │    4.24.232.63 │
│   0.31 │ Good      │ J       │ SI2     │ … │   4.344.352.75 │
│   0.24 │ Very Good │ J       │ VVS2    │ … │   3.943.962.48 │
│   0.24 │ Very Good │ I       │ VVS1    │ … │   3.953.982.47 │
│   0.26 │ Very Good │ H       │ SI1     │ … │   4.074.112.53 │
│   0.22 │ Fair      │ E       │ VS2     │ … │   3.873.782.49 │
│   0.23 │ Very Good │ H       │ VS1     │ … │    4.04.052.39 │
│     ·  │    ·      │ ·       │  ·      │ · │     ·  │     ·  │     ·  │
│     ·  │    ·      │ ·       │  ·      │ · │     ·  │     ·  │     ·  │
│     ·  │    ·      │ ·       │  ·      │ · │     ·  │     ·  │     ·  │
│    1.2 │ Premium   │ J       │ VS1     │ … │   6.776.724.18 │
│   1.12 │ Premium   │ H       │ SI2     │ … │   6.746.664.15 │
│    1.0 │ Premium   │ F       │ SI1     │ … │   6.436.383.86 │
│    1.0 │ Premium   │ E       │ SI2     │ … │   6.396.343.93 │
│    1.0 │ Very Good │ D       │ SI2     │ … │   6.386.354.03 │
│    1.0 │ Very Good │ E       │ SI2     │ … │   6.386.314.03 │
│    1.0 │ Premium   │ E       │ SI2     │ … │   6.416.363.92 │
│    1.0 │ Premium   │ E       │ SI2     │ … │   6.486.443.95 │
│    1.0 │ Premium   │ D       │ SI1     │ … │   6.416.293.94 │
│    1.0 │ Fair      │ D       │ SI1     │ … │    6.26.134.0 │
├────────┴───────────┴─────────┴─────────┴───┴────────┴────────┴────────┤
│ ? rows (>9999 rows, 20 shown)                    10 columns (7 shown) │
└───────────────────────────────────────────────────────────────────────┘

もちろん、通常の SQL のように色々な条件が書ける。 たとえば 1 カラット以上のレコードについて caratcutcolor カラムだけ取り出すとか。

>>> duckdb.sql("select carat, cut, color from df where carat > 1.0")
┌────────┬───────────┬─────────┐
│ carat  │    cut    │  color  │
│ double │  varchar  │ varchar │
├────────┼───────────┼─────────┤
│   1.17 │ Very Good │ J       │
│   1.01 │ Premium   │ F       │
│   1.01 │ Fair      │ E       │
│   1.01 │ Premium   │ H       │
│   1.05 │ Very Good │ J       │
│   1.05 │ Fair      │ J       │
│   1.01 │ Fair      │ E       │
│   1.04 │ Premium   │ G       │
│    1.2 │ Fair      │ F       │
│   1.02 │ Premium   │ G       │
│     ·  │   ·       │ ·       │
│     ·  │   ·       │ ·       │
│     ·  │   ·       │ ·       │
│   1.02 │ Ideal     │ G       │
│    1.5 │ Good      │ G       │
│   1.06 │ Premium   │ E       │
│   1.06 │ Ideal     │ G       │
│    1.5 │ Fair      │ H       │
│    1.7 │ Premium   │ I       │
│    3.0 │ Fair      │ H       │
│   1.56 │ Good      │ I       │
│   1.88 │ Good      │ J       │
│   1.51 │ Very Good │ J       │
├────────┴───────────┴─────────┤
│ ? rows             3 columns │
└──────────────────────────────┘

ちなみに DuckDB は CSV ファイルなどに対して直接クエリを発行することもできたりする。

>>> duckdb.sql("select carat, cut, color from 'diamonds.csv' where carat > 1.0")
┌────────┬───────────┬─────────┐
│ carat  │    cut    │  color  │
│ double │  varchar  │ varchar │
├────────┼───────────┼─────────┤
│   1.17 │ Very Good │ J       │
│   1.01 │ Premium   │ F       │
│   1.01 │ Fair      │ E       │
│   1.01 │ Premium   │ H       │
│   1.05 │ Very Good │ J       │
│   1.05 │ Fair      │ J       │
│   1.01 │ Fair      │ E       │
│   1.04 │ Premium   │ G       │
│    1.2 │ Fair      │ F       │
│   1.02 │ Premium   │ G       │
│     ·  │   ·       │ ·       │
│     ·  │   ·       │ ·       │
│     ·  │   ·       │ ·       │
│   1.02 │ Ideal     │ G       │
│    1.5 │ Good      │ G       │
│   1.06 │ Premium   │ E       │
│   1.06 │ Ideal     │ G       │
│    1.5 │ Fair      │ H       │
│    1.7 │ Premium   │ I       │
│    3.0 │ Fair      │ H       │
│   1.56 │ Good      │ I       │
│   1.88 │ Good      │ J       │
│   1.51 │ Very Good │ J       │
├────────┴───────────┴─────────┤
│ ? rows             3 columns │
└──────────────────────────────┘

ここまでの例では Polars の DataFrame にしておくとメモリにデータが乗るのでその分アクセスが速いというのはありそう。

DuckDB のクエリ結果を Polars の DataFrame に変換する

さて、DuckDB で実行したクエリの結果は、そのままだと DuckDB のオブジェクトになっている。

>>> query = "select carat, cut, color from df where carat > 1.0"
>>> type(duckdb.sql(query))
<class 'duckdb.DuckDBPyRelation'>

これは DuckDBPyRelation#pl() メソッドを使うことで Polars の DataFrame に変換できる。

>>> duckdb.sql(query).pl()
shape: (17502, 3)
┌───────┬───────────┬───────┐
│ carat ┆ cut       ┆ color │
│ ---   ┆ ---       ┆ ---   │
│ f64   ┆ strstr   │
╞═══════╪═══════════╪═══════╡
│ 1.17  ┆ Very Good ┆ J     │
│ 1.01  ┆ Premium   ┆ F     │
│ 1.01  ┆ Fair      ┆ E     │
│ 1.01  ┆ Premium   ┆ H     │
│ ...   ┆ ...       ┆ ...   │
│ 1.01  ┆ Very Good ┆ I     │
│ 1.04  ┆ Very Good ┆ I     │
│ 1.04  ┆ Fair      ┆ G     │
│ 1.02  ┆ Good      ┆ H     │
└───────┴───────────┴───────┘
>>> type(duckdb.sql(query).pl())
<class 'polars.internals.dataframe.frame.DataFrame'>

なお、これは Polars の DataFrame に対してクエリを発行したから、というわけではない。 DuckDB の Python API で実行したクエリの結果は pl() メソッドで Polars の DataFrame に変換できる。

試しに Diamonds データセットの CSV をデータベースに取り込んでからクエリを発行してみよう。 まずは CSV ファイルを元にテーブルを作成する。

>>> duckdb.sql("create table diamonds as select * from read_csv_auto('diamonds.csv')")
>>> duckdb.sql("show tables")
┌──────────┐
│   name   │
│ varchar  │
├──────────┤
│ diamonds │
└──────────┘

今度はテーブルに対してクエリを発行して、結果を pl() メソッドで Polars の DataFrame に変換してみる。

>>> query = "select carat, cut, color from diamonds where carat > 1.0"
>>> duckdb.sql(query).pl()
shape: (17502, 3)
┌───────┬───────────┬───────┐
│ carat ┆ cut       ┆ color │
│ ---   ┆ ---       ┆ ---   │
│ f64   ┆ strstr   │
╞═══════╪═══════════╪═══════╡
│ 1.17  ┆ Very Good ┆ J     │
│ 1.01  ┆ Premium   ┆ F     │
│ 1.01  ┆ Fair      ┆ E     │
│ 1.01  ┆ Premium   ┆ H     │
│ ...   ┆ ...       ┆ ...   │
│ 1.01  ┆ Very Good ┆ I     │
│ 1.04  ┆ Very Good ┆ I     │
│ 1.04  ┆ Fair      ┆ G     │
│ 1.02  ┆ Good      ┆ H     │
└───────┴───────────┴───────┘

ちゃんと変換できた。

まとめ

今回は DuckDB バージョン 0.7 で強化された Polars とのインテグレーションを試してみた。

なお、データベースに関しては、これまでも Polars は BigQuery などと連携して使うことはできた。 一方で、データセットのサイズが数十 GB 程度であれば DuckDB で事足りる状況も多いはず。 そこは、得られるスループットと、セットアップの手間や利用料金とのトレードオフで考えることになるだろう。

参考

duckdb.org