PostgreSQLでCTE内の関数呼び出しが、思った回数の数倍呼ばれる場合

はじめに

PostgreSQLのユーザー関数をCTE内で実行するクエリを書いていて、予期しないパフォーマンスの低下を発見したので説明します。
（厳密にはパフォーマンスが低下したわけではなく、データ量の増大によりSlow Queryレポートの閾値を超えてしまって報告されたということです）

PostgreSQLのユーザー関数の揮発性

PostgreSQLのユーザー関数の揮発性にはIMMUTABLE, STABLE, VOLATILEがあり、「なんでもVOLATILE」よりは「できるならSTABLE、IMMUTABLEにすべし」と思っていました。

CTE(共通テーブル式)内で関数を呼び出す

あるユーザー関数の戻り値をクエリ内で複数回参照したい場合、CTE内で関数を実行して後から参照するというクエリを書くことがあると思います。

WITH cte AS (
    SELECT foo(user_id) AS user_id
      FROM users
)
SELECT user_id,
       user_id || 'def' as user_id_2
  FROM cte
;

foo()が思った回数の二倍実行されている！

usersの行数が440行だった場合、foo()が実行されるのは何回でしょうか？
440回？
まず最初にCTE内のクエリが実行されメモリ上に展開された後、後続のクエリで参照されるから440回だろうと思っていましたが、ここで880回実行されるという現象に遭遇しました。

これは、PostgreSQL 12の新機能「CTEの高速化」によるもので、少なくともPostgreSQL 13はこの挙動のままでした。（14でどうなるかは未確認）。

今回の原因は、foo()の揮発性がSTABLEだったことで、VOLATILEだと発生しません。

回避策

関数の揮発性をVOLATILEにするか（これは本質的ではない気がする）、CTEを使うときにMATERIALIZEするように指示します。

WITH cte AS MATERIALIZED (
    SELECT foo(user_id) AS user_id
      FROM users
)
SELECT user_id,
       user_id || 'def' as user_id_2
  FROM cte
;

実行計画

参考までに、PostgreSQL 13.6における実行計画を載せておきます。

MATERIALIZEしない場合

Index Only Scan using users_index on public.users  (cost=0.27..233.47 rows=440 width=64) (actual time=0.059..1.381 rows=440 loops=1)
  Output: foo((users.user_id)::text), (foo((users.user_id)::text) || 'def'::text)
  Heap Fetches: 0
  Buffers: shared hit=5
Planning Time: 0.172 ms
Execution Time: 1.460 ms

MATERIALIZEする場合

CTE Scan on cte  (cost=122.37..132.27 rows=440 width=64) (actual time=0.067..1.565 rows=440 loops=1)
  Output: cte.user_id, (cte.user_id || 'def'::text)
  Buffers: shared hit=5
  CTE cte
    ->  Index Only Scan using users_index on public.users  (cost=0.27..122.37 rows=440 width=32) (actual time=0.060..1.236 rows=440 loops=1)
          Output: foo((users.user_id)::text)
          Heap Fetches: 0
          Buffers: shared hit=5
Planning Time: 0.146 ms
Execution Time: 1.674 ms

ユーザー関数のコスト指定は無力（かも）

PostgreSQLが実行計画を作成するにあたり、関数のCOSTも参照すると思いますが、今回のケースではCOST 1000000にしても揮発性がSTABLEだと自動的にはMATERIALIZEしてくれませんでした。

COST 1000000にしたときの実行計画

Index Only Scan using users_index on public.users  (cost=0.27..2200013.47 rows=440 width=64) (actual time=0.071..1.775 rows=440 loops=1)
  Output: foo((users.user_id)::text), (foo((users.user_id)::text) || 'def'::text)
  Heap Fetches: 0
  Buffers: shared hit=5
Planning Time: 0.151 ms
Execution Time: 1.869 ms

おまけ

ユーザー関数が何回呼ばれているかは、以下のクエリで確認できます。
※ postgresql.confでtrack_functionsにplかallを指定する必要があります。

SELECT * FROM pg_stat_user_functions;

以上です。