PostgreSQL で集合を返すユーザー定義関数の書き方

作成日:2016.12.04

この記事はPostgreSQL Advent Calendar 2016の4日目の記事である。

このページは PostgreSQL の集合を返すユーザー定義関数を C 言語で簡単に実装する方法を紹介する。 PostgreSQL 9.5 検証しているが現在 EOL を迎えていないバージョンであれば適用できると思われる。

PostgreSQL の他の記事へのインデックスはここ


更新履歴
(2016.12.04) 作成。


目次

1. 集合を返す関数

Oracle Database をはじめとする RDBMS 製品はユーザーが独自の関数を定義できるのが普通である。 このような機構を ユーザー定義関数(User Defined Function; UDF) と呼ぶことが多い。

ユーザー定義関数は、関数なので複数の引数をとり、1 つだけ結果を返す。 結果は整数・浮動小数点・文字列のようなスカラー値を返すのが一般的だが、製品によっては配列(array)や複合型(composite type)を返すことができるものもある。

PostgreSQL は集合を返す関数(Set Returning Function; SRF)が存在する。 少し毛色が異なり、1 回の呼び出しに対して複数の結果行を返す関数である。

例えば generate_series() は集合を返す組み込み関数である。 指定された範囲の整数を1行づつ返す。 下の SQL の場合、T1 には 3 行しかないので SELECT * FROM T1 は 3 行を返すはずだが、集合を返す関数が含まれることにによりそれよりも多い 12 行を返すことになった。

SQL-1: 集合を返す関数のサンプル
CREATE TABLE T1 (C1 INT);

INSERT INTO T1 (C1) VALUES (1), (2), (3);

SELECT C1, generate_series(1, C1) AS G1, generate_series(1, C1 * 2) AS G2 FROM T1;

 c1 | g1 | g2
----+----+----
  1 |  1 |  1
  1 |  1 |  2
  2 |  1 |  1
  2 |  2 |  2
  2 |  1 |  3
  2 |  2 |  4
  3 |  1 |  1
  3 |  2 |  2
  3 |  3 |  3
  3 |  1 |  4
  3 |  2 |  5
  3 |  3 |  6
(12 rows)

つまり「集合を返す関数」とは「テーブルを返す(ように見える)関数」である。

ユーザー定義関数として集合を返す関数を作れると、色々と便利なことが多い。 例えば、

集合を返すユーザー定義は FROM 句の後に指定することで、上記のようなユーザー定義関数を作っておいて WHERE 句の絞込みをかけたり、集約演算を行うこともできる。

2. 集合を返す関数の一般的な書き方

集合を返す関数を C 言語で記述する公式のやり方は 35.4.8. 集合を返すSQL関数35.9.9. 集合を返す に記載されている。

例として Linux の /proc/diskstats を読み込むユーザー定義関数を作成する。 関数の名前を show_diskstats1 とする。

/proc/diskstats は Linux システム内のディスクとそのアクセス統計情報を記録している。

   8      16 sdb 58721 30112 1645690 68805 700 10529 89832 126220 0 58921 195021
   8      17 sdb1 58665 30108 1645210 68696 700 10529 89832 126220 0 58812 194912
   8       0 sda 74696 37688 2461892 203698 1462510 2120160 28664504 5535643 0 1210183 5739814
   8       1 sda1 578 17 4754 569 7 1 64 4 0 572 572
   8       2 sda2 73612 37639 2452834 201727 1462503 2120159 28664440 5535639 0 1209602 5737840
   8       3 sda3 354 32 3088 1335 0 0 0 0 0 1335 1335

/proc/diskstats の 1 つの行は 14 列あるので、show_diskstats1 の定義は以下のように書ける。 引数(入力パラメータ)は 0 個である。

CREATE FUNCTION public.show_diskstats1()
RETURNS TABLE(major int, minor int, diskname text,
              read_completed bigint, read_merged bigint, read_sectors bigint, read_time int,
              write_completed bigint, write_merged bigint, write_sectors bigint, write_time int,
              io int, io_time int, weighted_io_time int)
AS 'MODULE_PATHNAME'
LANGUAGE C VOLATILE STRICT;

あるいは出力パラメータを用いて以下のように定義することもできる。

CREATE FUNCTION public.show_diskstats1(
              OUT major int, OUT minor int, OUT diskname text,
              OUT read_completed bigint, OUT read_merged bigint, OUT read_sectors bigint, OUT read_time int,
              OUT write_completed bigint, OUT write_merged bigint, OUT write_sectors bigint, OUT write_time int,
              OUT io int, OUT io_time int, OUT weighted_io_time int)
RETURNS SETOF record
AS 'MODULE_PATHNAME'
LANGUAGE C VOLATILE STRICT;

show_diskstats1 の実際のコードは以下のようになる。 SQL のレベルで show_diskstats1 が 1 回呼ばれている間に、C 言語の show_diskstats1() は複数回呼ばれることになる。

PG_FUNCTION_INFO_V1(show_diskstats1);
Datum
show_diskstats1(PG_FUNCTION_ARGS)
{
    FuncCallContext  *funcctx;
    Datum             result;
    struct
    {
        FILE   *file;
    } *diskstat_ctx;

    if (SRF_IS_FIRSTCALL())
    {
        /* 最初の 1 回はここを通る */
        MemoryContext oldcontext;
	TupleDesc     tupleDesc;

        oldcontext = MemoryContextSwitchTo(funcctx->multi_call_memory_ctx);

        /* このユーザー定義関数の戻り値となる複合型を tupleDesc として取り出す */
        if (get_call_result_type(fcinfo, NULL, &tupleDesc) != TYPEFUNC_COMPOSITE)
            elog(ERROR, "...");

        /* /proc/diskstats をオープンし、そのディスクリプタの情報を funcctx->user_fctx にリンクする */
        diskstat_ctx = palloc0(sizeof(*diskstat_ctx));
        diskstat_ctx->file = fopen("/proc/diskstats", "r");

        if (diskstat_ctx->file == NULL)
            elog(ERROR, "...");

        funcctx->user_fctx = diskstat_ctx;

        /* この関数の戻り値の型と tuple table slot を作成し保存する */
        funcctx->tuple_desc = tupleDesc;
        funcctx->slot = MakeSingleTupleTableSlot(tupleDesc);

        MemoryContextSwitchTo(oldcontext);
    }

    funcctx = SRF_PERCALL_SETUP();
    diskstat_ctx = funcctx->user_fctx;

    /* 全ての行を返し終わったかどうかを判定する。
     * この判定用に funcctx->call_cntr と funcctx->max_calls を使ってもよいが、
     * /proc/diskstats の場合は読み込んで EOF が出た時点が終了になる。
     */
    if (condition)
    {
        /* 返す行がある場合はこちら
         * この時点で funcctx->slot に結果内容を格納しておく。
         */
        HeapTuple tuple;

        /* 最終的な結果は Heap Tuple 形式なので virtual tuple から変換 */
        tuple = ExecCopySlotTuple(funcctx->slot);
        result = HeapTupleGetDatum(tuple);

        SRF_RETURN_NEXT(funcctx, result);
    }
    else
    {
        /* 全ての行を返し終わったのでクリーンナップを行う。
         * この回は結果行を返さない
         */
        fclose(diskstat_ctx->file);

        SRF_RETURN_DONE(funcctx);
    }
}

コード全体は github 上の diskstats.c を参照のこと。 Heap tuple や tuple table slot についてはPostgreSQL の基本データ型とタプルの扱いを参照して欲しい。

このコードは初回のコール、途中のコール、最後のコールを 1 つの関数の中で呼び分ける必要があるので、/proc/diskstats を読み込むという処理も 3 つのパーツに分断されてしまい、直線的にならない。 またこの関数が PostgreSQL の実行フレームワークの中から何回呼び出されるかも分かり辛い(例えば SQL-1 ならこの関数は 15 回呼び出されることになる)。

3. 集合を返す関数の効率的な書き方

2 章 の集合を返すユーザー定義関数の書き方とは別に、返すべき複数の結果を行を tuplestore に詰めて一度に返すという方法が存在する。 この方法は PostgreSQL 文書では紹介されていないが、集合を返すユーザー定義関数を書く場合、こちらの方がより簡単に記述できる。 また実行方法も直感的である(SQL-1 なら C 関数は 3 回しか呼び出されない)。

関数の名前を show_diskstats2 とする。 CREATE FUNCTION による関数の定義は show_diskstats1 と同様である。 show_diskstats2 の C 言語コードの概略は以下のようになる。 特に重要なコードは赤字でマークしている。 実際に動作するコードは github 上の diskstats.c に配置した。

PG_FUNCTION_INFO_V1(show_diskstats2);
Datum
show_diskstats2(PG_FUNCTION_ARGS)
{
    ReturnSetInfo  *rsi;
    TupleDesc       tupdesc;
    TupleTableSlot *tts;
    MemoryContext   oldcontext;

    rsi = (ReturnSetInfo *) fcinfo->resultinfo;

    /* 本当は rsi が有効かチェックした方がよい。やり方は github コードには載せている */  

    /* fcinfo->resultinfo->returnMode に SFRM_Materialize を指定すると tuplestore を使うことを宣言できる */
    rsi->returnMode = SFRM_Materialize;

    /* このユーザー定義関数の戻り値となる複合型を tupleDesc として取り出す */
    if (get_call_result_type(fcinfo, NULL, &tupleDesc) != TYPEFUNC_COMPOSITE)
        elog(ERROR, "...");

    oldcontext = MemoryContextSwitchTo(funcctx->multi_call_memory_ctx);

    /* この関数の戻り値の型と tuple table slot を作成し保存する */    
    rsi->setDesc = CreateTupleDescCopy(tupdesc);
    BlessTupleDesc(rsi->setDesc);
    tts = MakeSingleTupleTableSlot(rsi->setDesc);

    /* 結果行を格納する tuplestore を作成する */
    rsi->setResult = tuplestore_begin_heap(rsi->allowedModes &SFRM_Materialize_Random, false, work_mem);

    /* /proc/diskstats を一気に読み込む */
    {
        FILE *file;

        file = fopen("/proc/diskstats", "r");
        if (file == NULL)
            ereport(ERROR,
                    (errcode_for_file_access(),
                     errmsg("could not open file \"/proc/diskstats\" for writing: %m")));

        /* /proc/diskstats を EOF になるまで読み込む。
         * 読み込んだ結果は tts に格納しておく。
         */
        while (condition)
            tuplestore_puttupleslot(rsi->setResult, tts);

        fclose(file);
    }

    MemoryContextSwitchTo(oldcontext);

    /* ダミーの戻り値として NULL を返す。*/
    PG_RETURN_NULL();
}

show_diskstats2 は結果行を一気に作成することができるので、プログラムの制御フローが簡単になり書き易い。 ただし SQL の実行側から見ると、最初の結果行が返ってくるまでに時間がかかるという問題がある。 またネットワーク経由で受信したデータを逐次返すようなユーザー定義関数は作成できない。

コメント

トラックバック   [Trackback URL: http://www.nminoru.jp/cgi-bin/tb.cgi/postgresql__pg-write-set-returning-udf]
コメントを書き込む

TOP    掲示板    戻る
Written by NAKAMURA Minoru, Email: nminoru atmark nminoru dot jp, Twitter:@nminoru_jp