DuckDBを使ってTiDBにspatial機能を拡張する

はじめに

TiDBで、まれにspatial(GIS)関連機能の取り扱いについて質問されることがあります。残念ながら現在はサポートされていません。

ところで DuckDBにはspatial機能があります。これはspatial関連の関数群 ST_** を使えるようにするものですが、これはmysqlデータソースのデータに対しても利用できます。これを使ってTiDBにGEOMETRYデータを保存したり、関数を使うことが可能になります。

恐らく真剣に使う場合は一度DuckDB側にデータをインポートした方が速いと思います。こういうやり方もあるよくらいに思ってください。

準備

DuckDB Spatial リポジトリをクローンします。これは単にテストデータが利用したいだけですが、サンプル等もありDuckDBのspatial機能を試すのにもちょうどいいので丸ごとクローンしておきます。

git clone https://github.com/duckdb/duckdb_spatial.git

本記事のサンプルもDuckDB SpatialのExample をベースにしたものです。

Spatial機能を使う

TiDBへの接続

DuckDBを立ち上げて、TiDBに接続します。詳しい方法については先日書いたblogを参照してください。

-- DuckDBで実行
INSTALL mysql;
LOAD mysql;

-- TiDBの接続情報
CREATE SECRET  (
    TYPE MYSQL,
    HOST 'gateway01.ap-northeast-1.prod.aws.tidbcloud.com',
    PORT 4000,
    DATABASE test,
    USER '<TiDBのユーザー>',
    PASSWORD '<TiDBのパスワード>',
    SSL_MODE 'verify_identity',
    SSL_CA '/etc/ssl/cert.pem'
);

ATTACH '' AS tidb (TYPE MYSQL);

テーブルの作成 (TiDB)

TiDB側にデータインポート用のテーブルを作っておきます。DuckDBでは CREATE TABLE ... AS SELECT ... (CTAS) 構文が使えますが、残念ながら今回は利用できません。理由は:

• TiDBはGEOMETRY型をサポートしていないので、GEOMETRYで作ってもIMPORTできない
  • なおMySQLで実行した場合もMySQLのGEOMETRYとDuckDBのGEOMETRYは中身が違うのでエラーになりました
• テキストやBLOBに変換した場合、大きいGEOMETRYがエラーになる
  • これはDuckDBのMySQLマッピングでTEXT, BLOBになってしまうためです(MySQLだと2^16 までしか入りません）

そのためここでは、TiDB側で事前にテーブルを作り、INSERT ... SELECT ... でデータ投入することにします。Geometryをマップする型はサイズが大きければなんでもいいですが、ここではJSON型にしました。

-- TiDB側で実行
CREATE TABLE test.zones (zone text, locationId int, borough text, geom json);

Spatial機能を使ったデータのインポート

さてデータをインポートしていきます。DuckDBから実行します。元のgeomカラムはGEOMETRY型なので、ST_AsGeoJSON() を使ってJSONに変換しています。

-- DuckDBから実行
-- tidbはattach済み
INSERT INTO tidb.zones (zone, locationId, borough, geom) SELECT zone, LocationId, borough, ST_AsGeoJSON(geom) as geom
FROM ST_Read('./duckdb_spatial/test/data/nyc_taxi/taxi_zones/taxi_zones.shx');

TiDB側にもちゃんと入っていることが確認できます。

Spatial関数を使ったクエリ

例として、ある点が含まれる区画を探してみます。

-- DuckDBで実行
SELECT zone FROM tidb.zones 
WHERE ST_Within(
  ST_Transform(ST_POINT(40.722223,-73.98385299999998), 'EPSG:4326', 'ESRI:102718'),
  ST_GeomFromGeoJSON(geom));
┌──────────────┐
│     zone     │
│   varchar    │
├──────────────┤
│ East Village │
└──────────────┘

geomはJSON型なので、ST_GeomFromGeoJSON() でGEOMETRY型にしなければいけないことに注意してください。ただ、変換後は普通にSpatial関数の引数とすることができます。

おわりに

私自身あまりSpatial機能に詳しくないので、単にサンプルを動かしてみた程度のものですが、DuckDBのいわば外部表としてTiDBを使えるというのはちょっとおもしろいと思います。

これ別にTiDBである必要なくて、他のSpatialをサポートしてないストレージ形式、例えばParquetやCSVでも同じことできるんですよね。結構便利じゃないかなと思います。