PostgreSQL 17からEXPLAINにSERIALIZEオプションが追加され、クエリをシリアライズするコストも取得できるようになりました

PostgreSQL 17からはテーブルがTOAST化されていてもEXPLAIN (ANALYZE,SERIALIZE)でより精緻な実行計画を取得できるようになりました
2024.06.03

PostgreSQL 17から EXPLAINSERIALIZE オプションが追加され、クエリに伴うシリアライズ処理も考慮したデータ量とwireフォーマットに変換する処理時間を取得できるようになりました。 例えば、TOAST(The Oversized-Attribute Storage Technique)で行外格納されているようなケースにおいて、顕著な効果を確認できます。

TOAST化されたレコードを1万行含むテーブルに対するSELECT結果を比較します。

test=> EXPLAIN ANALYZE SELECT * FROM vec;
                                              QUERY PLAN
-------------------------------------------------------------------------------------------------------
 Seq Scan on vec  (cost=0.00..173.98 rows=9998 width=26) (actual time=0.006..0.656 rows=10000 loops=1)
 Planning Time: 0.045 ms
 Execution Time: 1.029 ms
(3 rows)

test=> EXPLAIN (ANALYZE,SERIALIZE) SELECT * FROM vec;
                                              QUERY PLAN
-------------------------------------------------------------------------------------------------------
 Seq Scan on vec  (cost=0.00..173.98 rows=9998 width=26) (actual time=0.006..3.049 rows=10000 loops=1)
 Planning Time: 0.037 ms
 Serialization: time=752.544 ms  output=111420kB  format=text
 Execution Time: 757.661 ms
(4 rows)

SERIALIZE オプションの有無により、データサイズは 0.25MBから108.8MBへ約440倍、処理時間は 1.029msから758ms へと約740倍の開きがあり、SERIALIZE の方がより正確な値です。

EXPLAIN方法 ANALYZE ANALYZE,SERIALIZE 係数(倍)
サイズ(MB) 0.25 108.8 439
処理時間(ms) 1.029 757.661 736

EXPLAIN (ANALYZE)の弱点を克服

PostgreSQLではEXPLAINコマンドで実行計画を取得できます。特に、ANALYZE オプションを追加すると、実際にクエリを発行して実行計画と実行処理時間の両方を取得できます。

ただし、このANALYZEオプションにも欠点があります。

PostgreSQLでは2kBを超える *1データをテーブルに格納する際には、TOAST(The Oversized-Attribute Storage Technique)化され、特殊領域にデータ配置されます。

クエリ実行時にシリアライズを伴うようなケースにおいて、シリアライズ処理も含めた処理時間やシリアライズ後のワイアフォーマットでのデータサイズは、実際にクエリを実行しないと取得できませんでした。

EXPLAIN (ANALYZE, SERIALIZE) というように SERIALIZE オプションも含めることで、この課題は解消されます。

ドキュメントから引用します。

Include information on the cost of serializing the query's output data, that is converting it to text or binary format to send to the client. This can be a significant part of the time required for regular execution of the query, if the datatype output functions are expensive or if TOASTed values must be fetched from out-of-line storage. EXPLAIN's default behavior, SERIALIZE NONE, does not perform these conversions. If SERIALIZE TEXT or SERIALIZE BINARY is specified, the appropriate conversions are performed, and the time spent doing so is measured (unless TIMING OFF is specified). If the BUFFERS option is also specified, then any buffer accesses involved in the conversions are counted too. In no case, however, will EXPLAIN actually send the resulting data to the client; hence network transmission costs cannot be investigated this way. Serialization may only be enabled when ANALYZE is also enabled. If SERIALIZE is written without an argument, TEXT is assumed.

PostgreSQL: Documentation: 17: EXPLAIN

やってみた

RAGのバックエンドには一般にベクトルデータベースが利用され、PostgreSQLは pgvector Extension を利用すると、ベクトルデータベースとして利用できます。埋め込みデータは大きな次元のfloat型で管理するため、簡単にTOAST化される2kBのしきい値を突破します。一般的な4バイトのfloat(fp32)を1024次元で持つと、1レコードあたり4kB要し、TOASTのしきい値の2倍もあります。

テキスト埋め込みモデルにAmazon Bedrockで提供される1024次元のAmazon Titan Text Embedding V2を利用し、TOAST化されたデータに対して EXPLAIN (ANALYZE, SERIALIZE) を実行して効果を確認します。

環境

  • PostgreSQL 17 beta 1(Amazon RDSプレビュー環境を利用)
  • pgvector : 0.7.0
  • 埋め込みモデル : Amazon Titan Text Embedding V2 # Amazon Bedrock
  • 次元 : 1024

pgvectorの有効化とテーブル作成

PostgreSQL 17 beta 1 を用意し、pgvector(0.7.0) を有効化します。

test=> CREATE EXTENSION vector;
CREATE EXTENSION

test=> \dx
                              List of installed extensions
   Name   | Version |   Schema   |                     Description
----------+---------+------------+------------------------------------------------------
 plpgsql  | 1.0     | pg_catalog | PL/pgSQL procedural language
 vector   | 0.7.0   | public     | vector data type and ivfflat and hnsw access methods
(2 rows)

次に、1024次元のベクトルカラムを含んだテーブルを用意します

CREATE TABLE vec (
    id bigserial PRIMARY KEY,
    embedding vector(1024)
);

テーブルを確認すると、1024次元のベクトル型からなる embedding カラムの Storage は external となっています。

test=> \d+ v*
                         Sequence "public.vec_id_seq"
  Type  | Start | Minimum |       Maximum       | Increment | Cycles? | Cache
--------+-------+---------+---------------------+-----------+---------+-------
 bigint |     1 |       1 | 9223372036854775807 |         1 | no      |     1
Owned by: public.vec.id

                   Index "public.vec_pkey"
 Column |  Type  | Key? | Definition | Storage | Stats target
--------+--------+------+------------+---------+--------------
 id     | bigint | yes  | id         | plain   |
primary key, btree, for table "public.vec"

                                                           Table "public.vec"
  Column   |     Type     | Collation | Nullable |             Default             | Storage  | Compression | Stats target | Description
-----------+--------------+-----------+----------+---------------------------------+----------+-------------+--------------+-------------
 id        | bigint       |           | not null | nextval('vec_id_seq'::regclass) | plain    |             |              |
 embedding | vector(1024) |           |          |                                 | external |             |              |
Indexes:
    "vec_pkey" PRIMARY KEY, btree (id)
Access method: heap

初期データの投入

Amazon Bedrock の テキスト埋込モデル Amazon Titan Text Embedding V2を使うと、次の様にしてテキストデータを1024次元のベクトルにできます。

import boto3
import json

bedrock_runtime = boto3.client("bedrock-runtime")

inputText = """\
Your text string goes here"
"""

body = json.dumps(
    {
        "inputText": inputText,
        "dimensions": 1024,
    }
)

modelId = "amazon.titan-embed-text-v2:0"
accept = "application/json"
contentType = "application/json"

response = bedrock_runtime.invoke_model(
    body=body, modelId=modelId, accept=accept, contentType=contentType
)

response_body = json.loads(response["body"].read())
embedding = response_body.get("embedding")

print(len(embedding))  # 1024
print(embedding[:10])  # [-0.038241122, 0.022944672, -0.04564263, -0.008265017, 0.065133266, -0.009190205, 0.026028633, 0.03676082, -0.053044137, 0.009190205]

今回はベクトル検索の性能を評価したいわけではなく、PostgreSQLの EXPLAIN (ANALYZE, SERIALIZE) 機能を確認するのが狙いのため、BedrockのAPIを大量に呼び出さず、乱数で-1から1の範囲の1024次元からなる1万件のベクトルデータを作成し、テーブルに投入しました。

import random

import psycopg
from pgvector.psycopg import register_vector

DSN = "host=pg17.xxx.us-east-2.rds-preview.amazonaws.com dbname=test user=postgres password=123"
with psycopg.connect(DSN, autocommit=True) as conn:
    register_vector(conn)
    with conn.cursor() as cur:
        for _ in range(10000):
            embedding = [random.uniform(-1, 1) for _ in range(1024)]
            cur.execute(
                "INSERT INTO vec (embedding) VALUES (%s)", (embedding,)
            )

1万件登録されたことを確認します。

test=> select count(*) from vec;
 count
-------
 10000
(1 row)

SERIALIZEオプションの有無によるEXPLAIN ANALYZEの比較

1万レコードを含んだ TOAST化されたテーブルに EXPLAIN ANALYZE を実行すると、SERIALIZEオプションの有無によって次のテーブルのような大きな違いがありました。

EXPLAIN方法 ANALYZE ANALYZE,SERIALIZE 係数(倍)
サイズ(MB) 0.25 108.8 439
処理時間(ms) 1.029 757.661 736

まずは従来型の EXPLAIN ANALYZE です。

test=> EXPLAIN ANALYZE SELECT * FROM vec;
                                              QUERY PLAN
-------------------------------------------------------------------------------------------------------
 Seq Scan on vec  (cost=0.00..173.98 rows=9998 width=26) (actual time=0.006..0.656 rows=10000 loops=1)
 Planning Time: 0.045 ms
 Execution Time: 1.029 ms
(3 rows)

PostgreSQLのANALYZEでは 各行の平均サイズを width で取得でき、単位はバイトです。 width=26が9998行あるため、データサイズは 9998 * 26 = 259948 (byte) ~ 254kB ~ 0.25MB です。

また、実行時間(Execution Time)は 1.029 ms です。

次に PostgreSQL 17 から利用可能になる EXPLAIN (ANALYZE,SERIALIZE) です。

test=> EXPLAIN (ANALYZE,SERIALIZE) SELECT * FROM vec;
                                              QUERY PLAN
-------------------------------------------------------------------------------------------------------
 Seq Scan on vec  (cost=0.00..173.98 rows=9998 width=26) (actual time=0.006..3.049 rows=10000 loops=1)
 Planning Time: 0.037 ms
 Serialization: time=752.544 ms  output=111420kB  format=text
 Execution Time: 757.661 ms
(4 rows)

シリアライズにかかる情報は Serialization の行から取得できます。

シリアライズ後のサイズ(=output) は 111420kB = 108.8MBで、約440倍に増えています。

同様に、シリアライズ処理にかかる時間(752.544 ms)も取得できるようになり、処理全体の時間は 1.029 ms から 757.661 ms へと 736倍に伸びています。

より現実に即した情報を取得できるようになりました。

最後に

PostgreSQL 17から EXPLAINSERIALIZE オプションが追加され、クエリに伴うシリアライズ処理も考慮したデータサイズと処理時間を取得できるようになりました。TOAST化されたレコードを含んでいるようなケースで有用です。

PostgreSQL 17以降を利用できるなら、EXPLAIN (ANALYZE,SERIALIZE) としましょう。

参考

脚注

  1. 設定によります