2016年7月13日

TF-IDFでデータベース内の類似テキストを検索する Part 3 (性能改善編)

PostgreSQL 感動巨編 TF-IDF 3部作の最終回、「性能改善編」です。 前回の最後で、今回作成した UDF である euclidean_distance の処理に時間がかかってそうだ、ということが分かってきました。

そのため、本エントリではこの UDF をもう少し詳細に見ながら、パフォーマンスを改善する方法を探ります。

■euclidean_distanceの性能分析


処理時間がかかっていることが分かった euclidean_distance 関数ですが、改めて処理を詳細に見ていきます。
CREATE OR REPLACE FUNCTION euclidean_distance(tfidf_a jsonb, tfidf_b jsonb)
  RETURNS float8
AS $$
    from sklearn.metrics.pairwise import euclidean_distances
    import json

    aa = json.loads(tfidf_a)
    bb = json.loads(tfidf_b)
    w = list(set(aa.keys()).union(set(bb.keys())))
    vec_a = []
    vec_b = []
    for t in w:
        if t in aa:
            vec_a.append(aa[t])
        else:
            vec_a.append(0)
        if t in bb:
            vec_b.append(bb[t])
        else:
            vec_b.append(0)
    distance = euclidean_distances([vec_a], [vec_b])
    return distance[0][0]
$$
LANGUAGE 'plpython2u';
まず、ボトルネックの確認のため、少しずつこの関数のコードをコメントアウトして実行時間を比較してみると、

    aa = json.loads(tfidf_a)
    bb = json.loads(tfidf_b)
の処理をコメントアウトすると実行時間が大幅に短縮されることが分かります。つまり、JSONB から Python のオブジェクト、今回の場合は dict 型ですが、ここへの変換に時間がかかっているわけです。

ちなみに、今回の TF-IDF の JSONB 型のデータのサイズは、平均 130kB 以上もある大きな JSON ドキュメントです。そのため、これをロードするにはそれなりに時間がかかっているというのは、それなりに理解できる状況ではあります。
testdb=> SELECT avg(pg_column_size(tfidf)) FROM pgsql_doc;
         avg
---------------------
 137897.235429447853
(1 行)

testdb=>
つまり、この JSON のロード処理を無くすことができれば、かなり処理時間を短縮できると思われます。

というわけで、どうやったらここで JSON のロードをせずに済ませられるか、ということを考えてみます。

■euclidean_distanceの処理詳細


ここで、改めて euclidean_distance 関数の処理の詳細を見てみます。
CREATE OR REPLACE FUNCTION euclidean_distance(tfidf_a jsonb, tfidf_b jsonb)
まず、引数として受け取る2つの JSONB は TF-IDF の情報であり、キーがトークン、値は TF-IDF の値のペアを持つ Key-Value 構造になっています。
testdb=> SELECT substring(tfidf::text from 300500 for 100) FROM pgsql_doc LIMIT 1;
                                                            substring
----------------------------------------------------------------------------------------------------------------------------------
 ン": 0.0, "ミレニアム": 0.0, "メカニズム": 0.0, "メガバイト": 0.0025927042824036564, "メタデータ": 0.0, "メッセージ": 0.0, "モジ
(1 行)

testdb=>
このように、「すべてのドキュメントに出てくるトークン」をキーとして持つ、非常に大きな Key-Value 構造になっているわけです。

そして、このような大きな2つの JSONB ドキュメントを受け取った後、
    w = list(set(aa.keys()).union(set(bb.keys())))
で、2つの JSONB ドキュメントのキーをマージします。(実際には、ここでマージしなくても、2つの JSONB ドキュメントですでにキーはすべて同じになっているはずなのですが、念のための処理です)

そして、キー(トークン)をソートした後、各 dict からトークンに対応する TF-IDF 値を取り出し、リストに追加します。
    vec_a = []
    vec_b = []
    for t in w:
        if t in aa:
            vec_a.append(aa[t])
        else:
            vec_a.append(0)
        if t in bb:
            vec_b.append(bb[t])
        else:
            vec_b.append(0)
この処理によって最後に得られるのは、ソートされたトークンの順番に並べられた TF-IDF 値のリスト、つまり float8 型のベクトルになります。

そして、このベクトルを scikit-learn の euclidean_distances 関数に与えてユークリッド距離を計算して、それを返却しています。
    distance = euclidean_distances([vec_a], [vec_b])
    return distance[0][0]
ここまで見てきたように、入力として渡された JSONB のデータは、最終的には float8 のベクトルに変換される必要があります。

逆に言えば、この UDF の入力として float8 のベクトルが渡されて来れば、JSONB 型でやりとりする必要性はそもそも無くなるわけです。

先のパフォーマンスの分析で、JSONB のロードに時間がかかっていることが分かっていますので、この JSONB 型からの脱却を試みてみます。

■TF-IDF のベクトル化


まず、テーブルに TF-IDF をベクトルとして保存するカラムを追加します。ここでは、float8 の配列型のカラムを使います。
testdb=> ALTER TABLE pgsql_doc ADD COLUMN tfidf_vec float8[];
ALTER TABLE
testdb=> \d pgsql_doc
                                テーブル "public.pgsql_doc"
    列     |         型         |                          修飾語
-----------+--------------------+-----------------------------------------------------------
 docid     | integer            | not null default nextval('pgsql_doc_docid_seq'::regclass)
 filename  | text               | not null
 html      | text               | not null
 plain     | text               |
 tf        | jsonb              |
 tfidf     | jsonb              |
 tfidf_vec | double precision[] |
インデックス:
    "pgsql_doc_pkey" PRIMARY KEY, btree (docid)

testdb=>
次に、TF-IDF を JSONB で受け取って、キー(トークン)でソートした後に、TF-IDF 値を float8 配列で返却する UDF を作成します。
CREATE OR REPLACE FUNCTION tfidf_jsonb2vec(tfidf jsonb)
  RETURNS float8[]
AS $$
    import json

    a = json.loads(tfidf)
    vec = []
    for t in sorted(a):
        vec.append(a[t])
    return vec
$$
LANGUAGE 'plpython2u';
そして、この UDF を使って、 tfidf カラムの内容を float8 配列に変換して tfidf_vec カラムに保存します。
testdb=> UPDATE pgsql_doc SET tfidf_vec = tfidf_jsonb2vec(tfidf);
UPDATE 1304
testdb=> SELECT count(tfidf_vec) FROM pgsql_doc;
 count
-------
  1304
(1 行)

testdb=>
以上で、ベクトル化は完了です。

以下のように、数値のみの配列となり、カラムのサイズも平均 4kB まで小さくなっていることが分かります。
testdb=> SELECT substring(tfidf_vec::text from 0 for 40) FROM pgsql_doc LIMIT 1;
                substring
-----------------------------------------
 {0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,
(1 行)

testdb=> SELECT avg(pg_column_size(tfidf_vec)) FROM pgsql_doc;
          avg
-----------------------
 4311.3703987730061350
(1 行)

testdb=>
最後に、ユークリッド距離を計算する UDF euclidean_distance を修正します。

データの受け渡しに JSONB を使わなくなったので、以下のようにシンプルになります。
CREATE OR REPLACE FUNCTION euclidean_distance(vec_a float8[], vec_b float8[])
  RETURNS float8
AS $$
    from sklearn.metrics.pairwise import euclidean_distances

    distance = euclidean_distances([vec_a], [vec_b])
    return distance[0][0]
$$
LANGUAGE 'plpython2u';

■TF-IDF のベクトルを使って類似性検索を実行する


では、前回と同じクエリを実行してみましょう。

変更されている点は JSONB 型の tfidf カラムではなく、 float8[] 型の tfidf_vec カラムを使うようにしたところだけです。
testdb=> SELECT
testdb->   filename,
testdb->   euclidean_distance(tfidf_vec, (SELECT tfidf_vec FROM pgsql_doc WHERE filename = 'wal.html') )
testdb-> FROM
testdb->   pgsql_doc
testdb-> ORDER BY
testdb->   2;
                     filename                      |  euclidean_distance
---------------------------------------------------+----------------------
 wal.html                                          | 1.49011611938477e-08
 wal-intro.html                                    |    0.670239012809462
 wal-internals.html                                |    0.696524116689475
 wal-reliability.html                              |    0.722138171836224
 admin.html                                        |    0.723666230739223
 wal-async-commit.html                             |    0.725507228362063
 runtime-config-wal.html                           |    0.725767186591778
 wal-configuration.html                            |    0.730449041608858
 release-7-4-29.html                               |    0.732007479060099
 disk-full.html                                    |    0.733419521902515
(...略...)
 brin.html                                         |      1.3933519134798
 nls.html                                          |     1.41006771577778
 gin.html                                          |     1.51897678394843
 event-trigger-matrix.html                         |     1.91950639453063
 sql-keywords-appendix.html                        |     2.23247656979142
(1304 行)

testdb=>
当然ながら、前回と同じ結果が返ってきています。

この時の実行時間を見ると約4秒強となり、前回の JSONB を使った実装で 100 秒近くかかったのと比べると、20倍程度の性能改善が行われたことが分かります。
testdb=> EXPLAIN ANALYZE SELECT
testdb->   filename,
testdb->   euclidean_distance(tfidf_vec, (SELECT tfidf_vec FROM pgsql_doc WHERE filename = 'wal.html') )
testdb-> FROM
testdb->   pgsql_doc
testdb-> ORDER BY
testdb->   2;
                                                        QUERY PLAN
--------------------------------------------------------------------------------------------------------------------------
 Sort  (cost=1104.81..1108.07 rows=1304 width=39) (actual time=4316.830..4318.665 rows=1304 loops=1)
   Sort Key: (euclidean_distance(pgsql_doc.tfidf_vec, $0))
   Sort Method: quicksort  Memory: 151kB
   InitPlan 1 (returns $0)
     ->  Seq Scan on pgsql_doc pgsql_doc_1  (cost=0.00..357.30 rows=1 width=18) (actual time=0.242..0.475 rows=1 loops=1)
           Filter: (filename = 'wal.html'::text)
           Rows Removed by Filter: 1303
   ->  Seq Scan on pgsql_doc  (cost=0.00..680.04 rows=1304 width=39) (actual time=4.590..4246.025 rows=1304 loops=1)
 Planning time: 0.064 ms
 Execution time: 4320.443 ms
(10 行)

testdb=>
これくらいまで性能が改善できれば、状況によっては使えるケースが増えてきそうです。

■その他の性能改善


今回は実施しませんでしたが、TF-IDF を事前に計算する処理も、もう少し性能改善できる可能性があります。

例えば、今回の処理では単純な Key-Value 構造しか扱わないため、 JSONB の代わりに hstore を使うことで性能が改善する可能性があります。
なお、hstore と Python の dict 型のマッピングは hstore_plpython2u という EXTENSION で実現することができます。
testdb=# CREATE EXTENSION hstore;
CREATE EXTENSION
testdb=# CREATE EXTENSION hstore_plpython2u;
CREATE EXTENSION
testdb=#
また、9.6 からはパラレルクエリが使えるようになりますので、CPU コアを増やすことで検索のレスポンスタイムを短縮させることができるようになります。

興味のある方は、こういった点にもぜひチャレンジしてみていただければと思います。

■まとめ


TF-IDF を用いた類似テキスト検索ということで、全3回に渡って PostgreSQL に保存したテキストデータに対してどのように類似検索を行うかをご紹介してきました。

PostgreSQL の UDF として実装したことによって、
-- TFの計算
UPDATE pgsql_doc SET tf = tf(mecab_tokenize(plain));
-- TF-IDFの計算
UPDATE pgsql_doc SET tfidf = tf_idf(tf, (SELECT idf(tf) FROM pgsql_doc));
-- TF-IDFのベクトル化
UPDATE pgsql_doc SET tfidf_vec = tfidf_jsonb2vec(tfidf);
という3ステップで、PostgreSQLの日本語マニュアルの類似検索を行えるようになりました。

また、最終回の今回は、データ構造を工夫することによって処理時間を大幅に短縮できることが分かりました。

PostgreSQLの強みはその拡張性であり、それはデータ分析というユースケースでも変わらないと思います。また、扱うデータが増えれば増えるほど、データベースから取り出さずに分析したい、というニーズは高まってくると思います。

ぜひ、PostgreSQLの拡張性を活かして、さまざまなデータ分析にPostgreSQLを使ってみていただければと思います。

では。

0 件のコメント:

コメントを投稿