[SQLite]CでSQLiteのユーザ定義関数を作る その３

前回はスカラー関数を作ったけど、SUMやAVGのような集約関数はどうなるだろうか。
今回は指定された列の行の値の積を返す関数を作ってみよう。

動作イメージ

> select 3 foo union all select 4 foo union all select 5 foo;
3
4
5
> select product(foo) from (select 3 foo union all select 4 foo union all select 5 foo);
60

productfunctions.c

#include <stddef.h>
#include "sqlite3ext.h"

SQLITE_EXTENSION_INIT1

typedef struct productCtx {
    sqlite_int64 product;
    sqlite_int64 cnt;
} productCtx;

// 1. 
static void productStep(sqlite3_context *context, int argc, sqlite3_value **argv) {
    productCtx *product_p;
    product_p = sqlite3_aggregate_context(context, sizeof(*product_p));
    int type = sqlite3_value_type(argv[0]);
    if (product_p && type != SQLITE_NULL) {
        if (product_p->cnt == 0) {
            product_p->product = 1;
        }
        product_p->product *= sqlite3_value_int64(argv[0]);
        (product_p->cnt)++;
    }
}

// 2.
static void productFinalize(sqlite3_context *context) {
    productCtx *product_p;
    product_p = sqlite3_aggregate_context(context, 0);
    sqlite3_result_int64(context, product_p->product);
}

// 3. 
int sqlite3_extension_init(
    sqlite3 *db, char **pzErrMsg, const sqlite3_api_routines *pApi){
    SQLITE_EXTENSION_INIT2(pApi);

    int res = sqlite3_create_function(
              db, "product", 1, SQLITE_ANY, NULL, NULL, productStep, productFinalize);
    if (res != SQLITE_OK) {
        *pzErrMsg = sqlite3_mprintf("Can't create product aggregate.");
        return 1;
    }

    return 0;
}

1. 集約関数の実装部分。
   レコード1行ごとにこの関数が呼ばれることになるので、状態を記憶しておく必要がある。そのためにはsqlite3_aggregate_context関数を使って、記憶領域を確保する。
   記憶領域は任意のサイズを指定できるので、第2引数に指定しよう。sqlite3_aggregate_contextは初回は領域を確保し、2回目以降は領域への参照を返してくれる。
   あとはスカラー関数と同じ要領で型に気をつけながら、計算結果を格納していこう。
2. 集約関数はもう１つ実装する関数が存在する。これはすべての行について、1.の関数が呼び出された後に呼び出される関数だ。
   ここで集約関数の戻り値を決定する。
   具体的にはsqlite3_aggregate_contextで結果領域を取得し、sqlite3_result_*で戻り値を設定する。
   また、sqlite3_aggregate_contextの第2引数には0を指定しておこう。
   これにより、後処理関数から戻った後に記憶領域が解放される。
3. 最後に関数の登録を行う。基本的にはスカラー関数の場合と同じだけど、関数は第6引数ではなく、第7引数に行ごとに実行される関数、第8引数に後処理関数を渡す。

これも同じようにビルドしよう。

$ gcc -I../sqlite-amalgamation-3150000gcc -I../sqlite-amalgamation-3150000 -L../sqlite-amalgamation-3150000 -fPIC -shared -Os productfunctions.c -o libproductfunctions.so -L../sqlite-amalgamation-3150000

$ LD_LIBRARY_PATH=./../sqlite-amalgamation-3150000/ ../sqlite-amalgamation-3150000/sqlite3
SQLite version 3.15.0 2016-10-14 10:20:30
Enter ".help" for usage hints.
Connected to a transient in-memory database.
Use ".open FILENAME" to reopen on a persistent database.
sqlite> .load libproductfunctions
sqlite> select product(foo) from (select 3 foo union all select 4 foo union all select 5 foo);
60
sqlite> 

集約関数の方もうまく動作していることが確認できたね。

前回のスカラー関数と今回の集約関数両方を含んだソースコードは以下となる。

productfunctions.c

#include <stddef.h>
#include "sqlite3ext.h"

SQLITE_EXTENSION_INIT1

static void productFunc(sqlite3_context *context, int argc, sqlite3_value **argv) {
    if (argc == 0) {
        sqlite3_result_int64(context, 0);
        return;
    }

    sqlite_int64 product = 1;
    int i;
    for (i = 0; i < argc; i++) {
        if (sqlite3_value_type(argv[i]) == SQLITE_NULL) return;
        product *= sqlite3_value_int64(argv[i]);
    }
    sqlite3_result_int64(context, product);
}

typedef struct productCtx {
    sqlite_int64 product;
    sqlite_int64 cnt;
} productCtx;

static void productStep(sqlite3_context *context, int argc, sqlite3_value **argv) {
    productCtx *product_p;
    product_p = sqlite3_aggregate_context(context, sizeof(*product_p));
    int type = sqlite3_value_type(argv[0]);
    if (product_p && type != SQLITE_NULL) {
        if (product_p->cnt == 0) {
            product_p->product = 1;
        }
        product_p->product *= sqlite3_value_int64(argv[0]);
        (product_p->cnt)++;
    }
}

static void productFinalize(sqlite3_context *context) {
    productCtx *product_p;
    product_p = sqlite3_aggregate_context(context, 0);
    sqlite3_result_int64(context, product_p->product);
}

// init
int sqlite3_extension_init(
    sqlite3 *db, char **pzErrMsg, const sqlite3_api_routines *pApi){
    SQLITE_EXTENSION_INIT2(pApi);

    int res = sqlite3_create_function(
              db, "product", -1, SQLITE_ANY, NULL, productFunc, NULL, NULL);
    if (res != SQLITE_OK) {
        *pzErrMsg = sqlite3_mprintf("Can't create product function.");
        return 1;
    }
    res = sqlite3_create_function(
              db, "product", 1, SQLITE_ANY, NULL, NULL, productStep, productFinalize);
    if (res != SQLITE_OK) {
        *pzErrMsg = sqlite3_mprintf("Can't create product aggregate.");
        return 1;
    }

    return 0;
}