multicornを使ってPythonでfdwを実装する - そこはかとなく書くよ。

今の会社に入ってからPostgreSQLを使い始めました。というわけで、初心者ですがPostgreSQL アドベントカレンダー 2011に参加してみます。

さて、PostgreSQL 9.1からSQL/MED規格の一部である外部データラッパ (FDW:Foreign Data Wrapper) がサポートされました。すでに PostgreSQL アドベントカレンダー 12月3日分にて、板垣さんが WebAPIをfdwから叩いてます。この記事では www_fdw という既存のfdwを利用する方法でしたが、今回は実際に自分でfdwを作ってみます。

fdwはAPIを6個程度書くだけで良いのでそれだけでもかなり書きやすいのですが、いかんせんCで書く必要があります。そこで、 multicorn を使ってpythonでfdwを書いてみました。

なにを題材に書こうかなと思ったのですが、Pure Pythonで作られた全文検索エンジンの Whoosh を使ってみました。

ソースコードは、こちらに置いてあります。

準備

まず、以下の二つをインストールする必要があります。

whoosh
multicorn

下にubuntuで叩いたコマンドを書いておきます。

  % sudo atp-get install postgresql-9.1  # postgresql-9.1を入れる
  % sudo apt-get install python-dev  # multicornのcompileにヘッダが必要
  % sudo apt-get install python-pip  # pythonのパッケージ管理ツールpipを入れる

  % sudo pip install pgxnclient  # pgxnを扱うためのクライアント
  % sudo pip install whoosh  # 全文検索エンジン
  % sudo pip install -e "hg+http://bitbucket.org/r_rudi/whooshfdw#egg=whooshfdw"  # 今回実装したパッケージをbitbucketからインストールする

  % sudo pgxn install multicorn --testing  # multicornを入れる

なお、今回はPostgreSQLサーバがpythonを叩くため、パッケージはvirtualenv環境ではなく、システム自体にインストールしてください。

whoosh index作成

全文検索するためにはまずインデックスを作成する必要があります。

registerというクラスを用意しておきましたのでそれを使います。

今回は、wikipediaの全ページのタイトルだけが含まれたファイルを検索してみます。

http://dumps.wikimedia.org/jawiki/20111203/ から jawiki-20111203-all-titles-in-ns0.gz をダウンロードし、/tmpに展開しておいてください。

from whooshfdw import register

register.register(file='/tmp/jawiki-20111203-all-titles-in-ns0', indexdir='/tmp/indexdir')

手元のマシンではだいたい40分ほどかかりました。お試しの場合はhead -n 1000とかしてください。

なお、インデックスですが決め打ちで2-gramで作成してあります。このあたりはコードを読んでくださいね。

SQL

準備が整ったところでいよいよfdwです。まずは Server と Tableを作成します。

CREATE EXTENSION multicorn;  -- multicornをDBにいれる
CREATE SERVER whoosh_srv FOREIGN DATA WRAPPER multicorn options (
        wrapper 'whooshfdw.whooshfdw.WhooshFDW' -- Server作成
);
CREATE FOREIGN TABLE whooshtable( -- FOREIGN TABLE作成
        id numeric,
        title character varying
) server whoosh_srv options (
        indexdir '/tmp/indexdir'  -- インデックスの場所をオプションで指定
);

できたところで SELECT 実行です。

% psql -c "SELECT * from whooshtable where title LIKE '%どこでも%';"
 id |       title        
----+--------------------
  1 | どこでもまど
  2 | どこでもドア
  3 | どこでもパス
  4 | どこでも大砲
  5 | どこでもCHU
  6 | どこでもドアー
  7 | どこでもいっしょ
  8 | どこでもチョコボ
  9 | どこでもWiiの間
 10 | どこでもドラえもん
(10 rows)

ちゃんと返ってきてますね。

実装の解説

multicornを使ったfdwの実装は ForeignDataWrapper クラスを継承したクラスを作成し、コンストラクタである __init__ と実行メソッドである execute を実装するだけです。

__init__ではsuperを呼び出すだけです。また、実際の検索を行うexecuteの中では yield を使って値を次々に返して行くだけです。

下はfdwの機能に絞ったコードの抜粋です。(全部はbitbucketを見てください)

class WhooshFDW(ForeignDataWrapper):
    def __init__(self, options, columns):
        super(WhooshFDW, self).__init__(options, columns)
        self.columns = columns
        self.indexdir = options["indexdir"]  # CREATE TABLE で指定したオプションを持ってくる

    def execute(self, quals, columns):
        for query in quals:  # qualsはWHERE句の中身が入っている
            q = parser.parse(query.value)
            with ix.searcher() as searcher:
                count = 1
                for r in searcher.search(q):
                    res = {'id': count, 'title': r["title"].encode('utf-8')}
                    yield res  # 値を辞書形式で返してあげる
                    count += 1

まとめると、

__ini__の中でsuperを呼ぶ
executeの中でyieldを使って値を次々に返せるようにする

これだけ。

注意点

なお、qualsというのはWHERE句の中身です。そしてwhere句は*必ず*適用されます。例えば、

def execute(self, quals, columns):
    for i in range(1,20):
        yield(i, "hoge")

としておいて、

SELECT * FROM whooshtable WHERE id = 10;

というSQLを実行した場合、「executeでqualはなにも使ってないんだから1から20まで表示されるだろう」と思いますよね。でもそうではなく、ここでは 10,hoge しか出てきません。

速度は？

実用を求めるならCで書けよな、と思いますが一応。

本体	Mac Mini (mid 2011)
OS	Ubuntu on Virtualbox
CPU	2.4GHz(Core 2 Duo)
メモリ	8G
postgres	9.1

PostgreSQLの設定はubuntuのそのまんまです。っていうか、VirtualBoxという時点で評価の意味がないような気がしますが…。

1247130行のwikipediaのタイトルだけを抜き出したものでindexを作成するのに

  % time python register.py
  2571.88s user 289.96s system 91% cpu 51:54.10 total

2571秒はだいたい42分。これをfdw経由で検索すると、

 Foreign Scan on whooshtable  (cost=10.00..15.00 rows=9999999 width=64) (actual time=619.132..620.073 rows=10 loops=1)
   Filter: ((title)::text ~~ '%どこでも%'::text)
   Foreign multicorn: multicorn
   Foreign multicorn cost: 10
 Total runtime: 681.000 ms

でした。ちなみに普通にpythonで検索してみました。timeitを使って、100回実行して平均を出してみると、

  % python search.py
  15.917749 msec

おやまあ、やっぱりPostgreSQLからpythonを呼び出すコストが結構かかるようですね。pypyでも試してみたくなりますが止めておきます。

まとめ

今回は PostgreSQL 9.1から搭載されたfdwをPythonで実装できるmulticornを使って、pure pythonの全文検索エンジンwhooshをSQLから叩けるようにしてみました。

…字で書くとなんかむちゃくちゃですね。でも、やったことは非常に少しなのでPostgreSQLの拡張性すごい！ということで。

明日は DaiMotoh さん、お願いします。