🛠️

NumPy構造体配列をDuckDBのテーブルに格納する

2024/12/20に公開

SQL

はじめに

以前から気になっていたDuckDBをやっと使い始めました。端的に言ってとても素晴らしいですね。
そんなDuckDBを触っていて、早速「NumPyの構造体配列をDuckDBのテーブルに追加するには」ということで悩んだので、メモがてら記事を書いています。

DuckDBとは？

以下の記事によれば「DuckDBはOLAP(オンライン分散処理)分析に特化したデータベースです」とのことです。
個人的な印象は「めっちゃ便利なSQLite」ですかね。

環境

本記事は、Python 3.11.6に以下のライブラリがインストールされた環境に基づいて書かれています。

requirements.txt

base58==2.1.1
duckdb==1.1.3
numpy==2.2.0
polars==1.17.1
pyarrow==18.1.0

また、コードの簡略化のためimport文を省略しています。実際には以下のimport文が必要です。

import os

import base58
import duckdb
import numpy as np
import polars as pl

前提知識: DuckDBでNumPy配列を扱うには？

1次元配列

DuckDBはもとよりNumPy配列をシームレスに扱えるようになっており、両者は簡単に連携できます。
なお、これ以降の実行例は、PythonのREPL環境にて実行したコードとその結果です。

>>> numpy_array = np.random.rand(4).astype(np.float32)
>>> numpy_array
array([0.98205835, 0.11887381, 0.21006313, 0.14625888], dtype=float32)

>>> duckdb.sql("SELECT * FROM numpy_array")
┌────────────┐
│  column0   │
│   float    │
├────────────┤
│ 0.98205835 │
│ 0.11887381 │
│ 0.21006313 │
│ 0.14625888 │
└────────────┘

実に簡単ですね。でもこの例を最初に見たとき、一瞬何が起こったのか分かりませんでした。
DuckDBには「Replacement Scan」という仕組みがあり、データベース中に存在しないテーブル名がSQL文中に見つかった時に、外側の環境を参照してくれます。
そのおかげで、SQL文中のnumpy_arrayという識別子とnumpy_array変数をわざわざ結びつけることなく、numpy_array変数を参照することができます。
かなり雑な説明なので、詳しくは公式サイトなどをご参照ください。

2次元配列

続いて2次元配列の例を見てみます。

>>> numpy_array = np.random.rand(4, 2).astype(np.float32)
>>> numpy_array
array([[0.4122309 , 0.11685179],
       [0.78982043, 0.81812525],
       [0.08608095, 0.31977254],
       [0.6476003 , 0.11576548]], dtype=float32)

>>> duckdb.sql("SELECT * FROM numpy_array")
┌────────────┬────────────┬────────────┬────────────┐
│  column0   │  column1   │  column2   │  column3   │
│   float    │   float    │   float    │   float    │
├────────────┼────────────┼────────────┼────────────┤
│  0.4122309 │ 0.78982043 │ 0.08608095 │  0.6476003 │
│ 0.11685179 │ 0.81812525 │ 0.31977254 │ 0.11576548 │
└────────────┴────────────┴────────────┴────────────┘

あれれ、気持ち的には2列4行のテーブルとして扱ってほしいところですが、4列2行になってしまいました。
ひとまず現時点では、事前に転置（.T）することで2列4行として扱っています。

>>> numpy_array_t = numpy_array.T
>>> duckdb.sql("SELECT * FROM numpy_array_t")
┌────────────┬────────────┐
│  column0   │  column1   │
│   float    │   float    │
├────────────┼────────────┤
│  0.4122309 │ 0.11685179 │
│ 0.78982043 │ 0.81812525 │
│ 0.08608095 │ 0.31977254 │
│  0.6476003 │ 0.11576548 │
└────────────┴────────────┘

本題: DuckDBでNumPy構造体配列を扱うには？

ここからが本題です。実際にDuckDBを使いたかったシチュエーションでは、NumPy構造体配列（Structured array）を使っており、これをDuckDBのテーブルにまとめて追加する方法について悩みました。

結論から言えば、PolarsのDataFrameを経由することで、簡単にDuckDBで扱うことができます。
以下は、idという文字列型のフィールド、featureという3次元浮動小数点数型のフィールドを持つ構造体配列の例です。
（ここでは分かりやすさのために3次元としていますが、実際には512次元や1024次元など、もっと大きな次元のデータです）

>>> structured_numpy_array = np.zeros((4,), dtype=[("id", "U11"), ("feature", np.float32, (3,))])
>>> structured_numpy_array["id"] = np.array([base58.b58encode(os.urandom(8)).decode() for _ in range(4)])
>>> structured_numpy_array["feature"] = np.random.rand(4, 3).astype(np.float32)
>>> structured_numpy_array
array([('LujTePS1FFQ', [0.42245385, 0.2560287 , 0.9933379 ]),
       ('SAmyeNfPk3m', [0.29928076, 0.8344725 , 0.9878882 ]),
       ('4aGx62FXUJL', [0.7726692 , 0.62192047, 0.23283875]),
       ('Hzs6wudhNRz', [0.7106031 , 0.63772374, 0.6755866 ])],
      dtype=[('id', '<U11'), ('feature', '<f4', (3,))])

>>> polars_df = pl.DataFrame(structured_numpy_array)
>>> polars_df
shape: (4, 2)
┌─────────────┬────────────────────────────────┐
│ id          ┆ feature                        │
│ ---         ┆ ---                            │
│ str         ┆ array[f32, 3]                  │
╞═════════════╪════════════════════════════════╡
│ LujTePS1FFQ ┆ [0.422454, 0.256029, 0.993338] │
│ SAmyeNfPk3m ┆ [0.299281, 0.834472, 0.987888] │
│ 4aGx62FXUJL ┆ [0.772669, 0.62192, 0.232839]  │
│ Hzs6wudhNRz ┆ [0.710603, 0.637724, 0.675587] │
└─────────────┴────────────────────────────────┘

>>> duckdb.sql("SELECT * FROM polars_df")
┌─────────────┬─────────────────────────────────────┐
│     id      │               feature               │
│   varchar   │              float[3]               │
├─────────────┼─────────────────────────────────────┤
│ LujTePS1FFQ │ [0.42245385, 0.2560287, 0.9933379]  │
│ SAmyeNfPk3m │ [0.29928076, 0.8344725, 0.9878882]  │
│ 4aGx62FXUJL │ [0.7726692, 0.62192047, 0.23283875] │
│ Hzs6wudhNRz │ [0.7106031, 0.63772374, 0.6755866]  │
└─────────────┴─────────────────────────────────────┘

ここまでくればあとは簡単で、事前に定義したテーブルに対してデータを追加するだけです。

>>> con = duckdb.connect()
>>> con.sql("CREATE TABLE my_duckdb_table(id VARCHAR PRIMARY KEY, feature FLOAT[3])")
>>> con.sql("INSERT INTO my_duckdb_table SELECT * FROM polars_df")
>>> con.sql("SELECT * FROM my_duckdb_table")
┌─────────────┬─────────────────────────────────────┐
│     id      │               feature               │
│   varchar   │              float[3]               │
├─────────────┼─────────────────────────────────────┤
│ LujTePS1FFQ │ [0.42245385, 0.2560287, 0.9933379]  │
│ SAmyeNfPk3m │ [0.29928076, 0.8344725, 0.9878882]  │
│ 4aGx62FXUJL │ [0.7726692, 0.62192047, 0.23283875] │
│ Hzs6wudhNRz │ [0.7106031, 0.63772374, 0.6755866]  │
└─────────────┴─────────────────────────────────────┘

PolarsのDataFrameを経由することで、無事にNumPy構造体配列をDuckDBのテーブルに格納することができました。

おわりに

まだ触り始めたばかりのDuckDBですが、かなりの可能性を感じます。
今まではNumPy + Pandasでデータ加工を行ってきましたが、今後はNumPy + Polars + DuckDBで行っていこうと思っています。

本記事が何らかの参考になれば幸いです。