プランナサポート関数を作って自作演算子でインデックスが使えるようにする（続・LIKE 演算子とインデックスの謎を探る）

# はじめに

この記事は[PostgreSQL Advent Calendar 2024](https://qiita.com/advent-calendar/2024/postgresql) 2 枚目の23日目（+ [SRA Advent Calendar 2024](https://qiita.com/advent-calendar/2024/sra)の7日目）で、[LIKE 演算子とインデックスの謎を探る](https://qiita.com/yugo-n/items/8c4c05fea9d0d5ec6e67)の続編です。

### これまでのあらすじ

[前の記事](https://qiita.com/yugo-n/items/8c4c05fea9d0d5ec6e67)では PostgreSQL の [`LIKE` 演算子](https://www.postgresql.jp/document/16/html/functions-matching.html#FUNCTIONS-LIKE)を使った文字列（`text`型）の前方一致検索で、インデックスが使われるしくみについて説明しました。内容を簡単にまとめると以下のようなことを書きました。

- `LIKE` 演算子は構文解析の段階で、それと等価な `~~` 演算子に変換される
- しかし、`~~` 演算子それ自体に対しては直接インデックスを使用することができない[^opc]
- そのため、実行計画を作る段階で、インデックスが使用可能な等価演算子（`=`）、比較演算子（`>=`, `<`, `~>=~`, `~<~`）などを使った別の条件式に変換される[^opc2]
- それを行っているのは、`~~` 演算子を実装する二項論理関数 `textlike` の「[プランナサポート関数](https://www.postgresql.jp/document/16/html/xfunc-optimization.html)」である `textlike_support` である

[^opc]:より正確には、`text` 型に対して作られるインデックスの演算子クラスに `~~` 演算子が含まれていない、ということです。[演算子クラス](https://www.postgresql.jp/document/16/html/indexes-opclass.html)については、よろしければ[こちら](https://www.postgresql.jp/sites/default/files/2024-12/T3_JPUG_2024_index_final.pdf)もどうぞ。
[^opc2]:text型上のインデックスの演算子クラスはこれらの演算子を含んでいるので、インデックスが使用可能です。

### 今回は

この記事ではプランナサポート関数についてさらに理解を深めるため、新しい演算子とプランナサポート関数を自作してみます。

# 新しい演算子を作る

さて、ではどんな演算子を作りましょうか。

前回取り上げた`LIKE`演算子（`~~`演算子）の特徴は「それ自体ではインデックスが使えないが、インデックスが使用可能な別な条件式に変換可能」という点でした。そこで、今回はその単純な例として「整数の絶対値が等しい」ことを表す演算子`|=|`を作ってみましょう。

この新しい演算子でインデックスが使えなくとも、`x |=| y` 、すなわち整数 `x` と `y` の絶対値が等しい場合には `x = ANY(ARRAY[y, -y])` （あるいは `x IN (y,-y)` ）といった条件式が成り立ちます[^any]。整数同士の等号演算 `=` は普通にインデックスが使用可能ですので、この式だったらインデックスが使用可能になるはずです。これを実現してみよう、というわけです。

[^any]: 素直に考えると `x = y OR x = -y` と `OR` を使った式になるのですが、プランナサポート関数を用いたインデックス条件生成が現時点では `OR` 式に対応していないため `ANY` を使うことにしました。

では手始めに「絶対値が等しい」ときに真を返す二項論理関数 `int_abs_eq` を作りましょう。もう細かいことは考えないでちゃちゃっと `PL/pgSQL` 中で `abs` 関数を使って書いてしまいます。[^plpgsql]。なお、今回は簡単のため `int4` 型だけを考慮することにします。


[^plpgsql]:SQL関数でも作れるのですが、そうすると実行計画作成時にプランナによって関数が「インライン展開」されてしまい想定された通りにインデックスが使用されないので、今回は `PL/pgSQL` にしました。

```sql
test=# CREATE FUNCTION int_abs_eq(int, int) RETURNS bool AS
         $$ BEGIN RETURN abs($1) = abs($2); END; $$ LANGUAGE plpgsql;
CREATE FUNCTION

test=# SELECT int_abs_eq (-1,1);
 int_abs_eq 
------------
 t
(1 row)
```

次に、演算子 `|=|` を作ります。ここも細かい説明は抜きに、`CREATE OPERATOR` を使って、上で実装した`int_abs_eq`関数と左右の型（`int4`）を指定してあげるだけです。

```sql
test=# CREATE OPERATOR |=| (FUNCTION = int_abs_eq, LEFTARG = int4, RIGHTARG = int4);
CREATE OPERATOR

test=# SELECT -1 |=| 1;
 ?column? 
----------
 t
(1 row)
```

はい。とりあえず、演算子の準備はこれでいいでしょう。

# インデックスが使えないことの確認

`int4` 型のカラムを持つ適当なテーブルとインデックスを作ります。

```sql
test=# CREATE TABLE tbl10 (i int);
CREATE INDEX ON tbl10 (i);

test=# INSERT INTO tbl10 SELECT generate_series(-10000,10000);
INSERT 0 20001

test=# CREATE TABLE
CREATE INDEX

test=# EXPLAIN (COSTS off) SELECT * FROM tbl10 WHERE i = 1;
                 QUERY PLAN                 
--------------------------------------------
 Index Only Scan using tbl10_i_idx on tbl10
   Index Cond: (i = 1)
(2 rows)
```

`|=|` 演算子を使った検索もできますが、当然、その場合には上で作ったインデックスが使えず `SeqScan` となります。

```sql
test=# SELECT * FROM tbl10 WHERE i |=| 1;
 i  
----
 -1
  1
(2 rows)

test=# EXPLAIN (COSTS off) SELECT * FROM tbl10 WHERE i |=| 1;
     QUERY PLAN      
---------------------
 Seq Scan on tbl10
   Filter: (i |=| 1)
(2 rows)
```

# プランナサポート関数を作る

それでは、本題です。`|=|` 演算子でインデックスを使えるようにするため、プランナサポート関数を作りましょう。

[ドキュメント](https://www.postgresql.jp/document/16/html/xfunc-optimization.html)にも説明がある通り、「条件式をインデックス可能な演算子句に変換」するサポート関数は、`SupportRequestIndexCondition`リクエスト型を受け取り、変換後の式を返します。今回は、`x |=| y` という式を `x = ANY(ARRAY[y,-y])`に変換して返します。 関数名は `int_abs_eq_support` としましょう。

これを実装する[C言語関数](https://www.postgresql.jp/document/16/html/xfunc-c.html)は以下のとおりです。内容は PostgreSQL の実装を知らないと難しいと思いますが、少しでも流れが想像できるように日本語でコメントを書いてみました。

```c:int_abs_eq_support.c
#include "postgres.h"
#include "fmgr.h"

#include "catalog/namespace.h"
#include "catalog/pg_operator_d.h"
#include "catalog/pg_opfamily.h"
#include "nodes/makefuncs.h"
#include "nodes/nodeFuncs.h"
#include "nodes/supportnodes.h"
#include "utils/lsyscache.h"

PG_MODULE_MAGIC;
PG_FUNCTION_INFO_V1(int_abs_eq_support);

Datum int_abs_eq_support(PG_FUNCTION_ARGS)
{
	Node	*rawreq = (Node *) PG_GETARG_POINTER(0);
	List	*result = NIL;

    /* SupportRequestIndexCondition リクエスト型の場合に対応 */
	if (IsA(rawreq, SupportRequestIndexCondition))
	{
		SupportRequestIndexCondition *req = (SupportRequestIndexCondition *) rawreq;

		/* 入力が演算子式であり、インデックスが整数型、btree の場合に対応 */
		if (is_opclause(req->node) && req->opfamily == INTEGER_BTREE_FAM_OID)
		{
			Expr *neg_rightop;
			ScalarArrayOpExpr *saopexpr;
			ArrayExpr  *arrexpr;

			/* オリジナルの条件式の左辺と右辺を取得 */
			OpExpr *clause = (OpExpr *) req->node;
			Node *leftop = linitial(clause->args);
			Node *rightop = lsecond(clause->args);

			/* 整数の負の演算子(-) を取得 */
			Oid ngopr = OpernameGetOprid(list_make2(makeString("pg_catalog"),
										 makeString("-")), InvalidOid, INT4OID);

			/*
			 * インデックスカラムが右辺ならば左辺と右辺を交換して、
			 * 左辺がインデックスカラム col、右辺が定数 c になるようにする
			 *
			 *   c |=| col    ->    col |=| c
			 */
			if (req->indexarg != 0)
			{
				Node *swap = leftop;
				leftop = rightop;
				rightop = swap;
			}

			/* 右辺の値を負値にする式 -c を作る */
			neg_rightop = make_opclause(ngopr, BOOLOID, false,
								   (Expr *) rightop, NULL,
								   InvalidOid, InvalidOid);

			/* 右辺と右辺の負値を含む配列 {c, -c} を作る */
			arrexpr = makeNode(ArrayExpr);
			arrexpr->array_typeid = get_array_type(INT4OID);
			arrexpr->array_collid = InvalidOid;
			arrexpr->element_typeid = INT4OID;
			arrexpr->elements = list_make2(rightop, neg_rightop);
			arrexpr->multidims = false;
			arrexpr->location = -1;

			/*
			 * 左辺は右辺か右辺の負値のどちらかに等しい：以下の式を作る
			 *
			 * col = ANY(ARRAY[c,-c])
			 */
			saopexpr = makeNode(ScalarArrayOpExpr);
			saopexpr->opno = Int4EqualOperator;
			saopexpr->opfuncid = get_opcode(Int4EqualOperator);
			saopexpr->hashfuncid = InvalidOid;
			saopexpr->negfuncid = InvalidOid;
			saopexpr->useOr = true;
			saopexpr->inputcollid = InvalidOid;
			saopexpr->args = list_make2(leftop, arrexpr);
			saopexpr->location = -1;

			result = list_make1(saopexpr);
		}
	}

	PG_RETURN_POINTER(result);
}
```

上記がC言語で書かれたプランナサポート関数の中身ですが、これを PostgreSQL で使えるようにするためには関数オブジェクトを作る必要があります。そのために、C言語関数を[コンパイル](https://www.postgresql.jp/document/16/html/xfunc-c.html#DFUNC)して共有ライブラリを作り、PostgreSQL の`$LIBDIR`にコピーします。`pg_config` にパスが通っている状態ならば、以下を実行すればOKです。

```shell
$ gcc -Fpic -I`pg_config --includedir-server` -o int_abs_eq_support.o -c int_abs_eq_support.c 
$ gcc -Fpic -shared -I`pg_config --includedir-server` -o int_abs_eq_support.so int_abs_eq_support.o
$ cp int_abs_eq_support.so `pg_config --libdir`
```

次に`CREATE FUNCION`でサポート関数オブジェクトを作って、`ALTER FUNCION ... SUPPORT`でそれを`|=|`演算子の実装関数`int_abs_eq`のプランナサポート関数として登録しましょう。

```sql
test=# CREATE FUNCTION int_abs_eq_support(internal) RETURNS internal AS 'int_abs_eq_support' LANGUAGE C;
CREATE FUNCTION

test=# ALTER FUNCTION int_abs_eq ( integer, integer) SUPPORT int_abs_eq_support;
ALTER FUNCTION
```

# インデックスが使えるか確かめる

最後に `|=|` 演算子でインデックスが使えるようになったか確認してみましょう。

```sql
test=# EXPLAIN (COSTS off) SELECT * FROM tbl10 WHERE i |=| 1;
                 QUERY PLAN                 
--------------------------------------------
 Index Only Scan using tbl10_i_idx on tbl10
   Index Cond: (i = ANY (ARRAY[1, (- 1)]))
   Filter: (i |=| 1)
(3 rows)
```

めでたくインデックスが使用されました。成功です。

```sql
test=# SELECT * FROM tbl10 WHERE i |=| 1;
 i  
----
 -1
  1
(2 rows)
```

もちろん検索も正しくできています。

# おわりに

今回は PostgreSQL で新しい演算子を自作して、さらにプランナサポート関数を作成することでインデックスが使えるようにしてみました。

なお、この記事の目的は、以前の[LIKE演算子とインデックスの記事]()で紹介したプランナサポート関数についての理解を深めることで、実際に新しい演算子でインデックスを使えるようにするのは、[演算子クラスを作成](https://www.postgresql.jp/document/16/html/xindex.html)するのが一般的かと思います[^sup]。あくまでPostgreSQL の内部動作を理解するための一端として少しでも貢献できれば幸いです。

[^sup]:さらに今回のような例だったら、`SupportRequestIndexCondition` リクエストではなく、`SupportRequestSimplify` リクエスト型プランナサポート関数で式を単純変換するだけでもよかったかもしれません。リクエスト型については[こちら](https://www.postgresql.jp/document/16/html/xfunc-optimization.html)。