Postgresql 编写自定义 C 函数

前言

postgresql 支持自定义函数,并且还支持多种语言进行编写, 极大的提高了可扩展性。postgresql 支持使用 pgSQL(postgresql提供的类sql语言),python 和 c 语言来编写函数,本篇主要讲解 c 语言,因为它的性能是最好的。

示例

编写 c 函数分为三部分:

  1. 编写 c 文件,定义函数实现
  2. 编译程序,需要编译成共享库
  3. 在 postgresql 中执行 CREATE FUNCTION注册函数

下面展示一个简单的函数,用于两数相加。通过这个简单的例子,来看看操作流程。

编写函数

创建一个my_add_func.c 文件,内容如下

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
#include "postgres.h"
#include "fmgr.h"

PG_MODULE_MAGIC;

PG_FUNCTION_INFO_V1(my_add_func);
Datum my_add_func(PG_FUNCTION_ARGS)
{
    int32 a = PG_GETARG_INT32(0);
    int32 b = PG_GETARG_INT32(1);
    int64 result = a + b;
    PG_RETURN_INT64(result);
}

这里面的函数很简单,只不过调用了一些莫名其妙的宏,所以会觉得疑惑。这些宏的原理在后面会讲到,这里先简单的介绍下作用。

PG_MODULE_MAGIC宏必须要使用,后面编译生成的库才可以被 postgresql 加载。

PG_FUNCTION_INFO_V1宏声明了函数的版本,这个也必须要使用,目前postgresql 只支持 v1 版本的函数。

PG_FUNCTION_ARGS宏定义了参数类型和名称,等于FunctionCallInfo fcinfo,这个函数名称会在取参数值会用到。

PG_GETARG_INT32宏表示取出指定位置的参数,并且转换为 int32 类型。

PG_RETURN_INT64宏表示将 int64 类型的数值,转换为Datum并且返回。

编译函数

上述的 c 文件,包含了 postgresql 里的头文件,所以在编译的时候需要指定头文件的位置。通过pg_config --includedir-server命令就可以找到。我是通过 yum 在 Centos 上装的 postgresql,执行情况如下:

1
2
[root@host1]# /usr/pgsql-9.6/bin/pg_config --includedir-server
/usr/pgsql-9.6/include/server

然后执行编译命令,记住这个生成共享库(so 结尾的文件)的路径,在创建函数时会用到

1
2
gcc -fPIC -c my_add_func.c -I/usr/pgsql-9.6/include/server/  # 生成my_add_func.o文件
gcc -shared -o my_add_func.so my_add_func.o            # 生成共享库

创建函数

然后在 postgresql 命令行执行CREATE FUNCTION命令

1
2
CREATE FUNCTION my_add(a integer, b integer) RETURNS integer
     AS '/var/lib/pgsql/my_add_func.so', 'my_add_func' LANGUAGE C STRICT;

接下来我们尝试使用 my_add 函数

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
test=# select my_add(1, 2);
 my_add 
--------
      3
(1 row)

test=# select my_add(id, price), id, price from orders ;
 my_add | id | price 
--------+----+-------
    124 |  1 |   123
(1 row)

数据类型

通过上面的例子,可以看到创建一个 c 函数并不难,不过还有些细节需要注意到,就是数据是如何在 postgresql 和函数之间传递的。

postgresql 支持的数据类型有多种,基本类型 int,char,double 等,还有复杂类型,比如字符串,结构体等。这些数据传递时都是由Datum类型表示,也就是指针类型。指针长度根据平台不同而不一样,有32位或64位。那么它是如何能够表示这么多类型的数据呢。下面分为两种情形来分析,基本数据类型和复杂数据类型。

基本数据传递

基本数据类型包含 int8,int16,int32,int64,float等数值型。因为Datum的长度根据系统的不同,可能是 4个字节 或 8个字节。如果基本类型的长度小于等于 Datum,那么将它的值直接存储在Datum。如果大于 Datum,就需要分配堆内存,然后将内存地址存储在Datum

我们以 int16 类型的数值 -1 为例,它的二进制位1111 1111 1111 1111。这里假设Datum的长度为 4 个字节,将其转换为Datum之后的二进制位0000 0000 0000 0000 1111 1111 1111 1111。这里只是在高位填充0。当从Datum转换 int16 类型时,直接取对应的低位。这里转换规则和我们平时接触的不太一样,在 c 语言中如果是负数,会在高位填充1,正数才会填充0,目前还不太清楚原因。

再来看看 int64 数据类型,Int64GetDatum定义了转换规则。如果Datum的长度等于 8 个字节,那么就定义了USE_FLOAT8_BYVAL宏,也就是直接存储到Datum里。如果Datum的长度小于 8 个字节,那么则分配堆内存,然后将内存地址存储在Datum

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
#ifdef USE_FLOAT8_BYVAL
#define Int64GetDatum(X) ((Datum) (X))
#else
extern Datum Int64GetDatum(int64 X);
#endif

Datum Int64GetDatum(int64 X)
{
    // 使用palloc分配堆内存,palloc的作用同malloc一样,在它基础之上增加了内存管理
	int64	   *retval = (int64 *) palloc(sizeof(int64));
    // 设置堆内存值
	*retval = X;
    // 将内存地址存储在Datum
	return PointerGetDatum(retval);
}

相关宏

这里再来回顾下PG_FUNCTION_ARGS宏,它定义函数参数名称fcinfo,类型为FunctionCallInfoFunctionCallInfo是一个结构体,里面有个args成员,表示参数数组。

1
#define PG_FUNCTION_ARGS	FunctionCallInfo fcinfo

postgresql 为我们提供了提取参数的方法,下面列举其中一些

1
2
3
4
5
6
7
8
// 提取指定位置的参数
#define PG_GETARG_DATUM(n)	 (fcinfo->args[n].value)

// 将指定位置的参数转换为uint16
#define PG_GETARG_UINT16(n)  DatumGetUInt16(PG_GETARG_DATUM(n))

// 将指定位置的参数转换为int32
#define PG_GETARG_INT32(n)	 DatumGetInt32(PG_GETARG_DATUM(n))

同样 postgresql 也为我们提供了返回结果的方法,

1
2
3
4
// 将int32转换为Datum类型,并且返回
#define PG_RETURN_INT32(x)	 return Int32GetDatum(x)
// 将uint32转换为Datum类型,并且返回
#define PG_RETURN_UINT32(x)  return UInt32GetDatum(x)

复合数据传递

上述介绍的my_add_func例子,只是接收和返回基本类型。postgresql 还可以接收行数据,或者自定义的数据类型(比如postgis插件的 POINT 类型),这种情况叫做复合类型数据,它在 c 语言中 由 tuple 表示,下面演示一个实例:

c 文件内容

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
#include "postgres.h"
#include "funcapi.h"

PG_MODULE_MAGIC;

PG_FUNCTION_INFO_V1(my_tuple_func);
Datum my_tuple_func(PG_FUNCTION_ARGS)
{
    // 第一个参数为复合类型
	HeapTupleHeader tuple = PG_GETARG_HEAPTUPLEHEADER(0);
	
	// 取出返回数据的类型,由TupleDesc表示
	TupleDesc	tupdesc;
	if (get_call_result_type(fcinfo, NULL, &tupdesc) != TYPEFUNC_COMPOSITE)
		elog(ERROR, "return type must be a row type");

	// 构建返回值
	Datum values[2];
	values[0] = PointerGetDatum(cstring_to_text("hello"));
	values[1] = Int32GetDatum(1);
    // nulls 数组表示值是否为空
	bool		nulls[2];
	memset(nulls, 0, sizeof(nulls));
	// 构建返回数据,由HeapTuple表示
	HeapTuple result = heap_form_tuple(tupdesc, values, nulls);
	PG_RETURN_DATUM(HeapTupleGetDatum(result));
}

创建函数语句如下,定义了接收参数类型为 record, 返回结果有两列(text,int)

1
2
3
4
CREATE FUNCTION my_tuple_func(tuple record) 
     RETURNS TABLE (message text ,num int)
AS '/var/lib/pgsql/my_tuple_func', 'my_tuple_func'
LANGUAGE C STRICT;

执行语句,返回的数据为复合类型,可以通过.*符号展示所有列

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
test=# select my_tuple_func(orders) from orders where id = 1;
 my_tuple_func 
---------------
 (hello,1)
(1 row)

test=# select (my_tuple_func(orders)).* from orders where id = 1; 
 message | num 
---------+-----
 hello   |   1

通过上面的例子,可以看到可以接收行数据,参数为对应的表名即可。它由HeapTuple类表示,关于HeapTuple的原理,可以参考后续的文件。

在获取参数时,我们又看到了一个新的宏PG_GETARG_HEAPTUPLEHEADER,它负责提取参数并且转换为HeapTuple类。

在构建返回结果时,我们需要调用get_call_result_type函数获取返回结果的类型,也就是TupleDesc。返回结果的类型是在我们执行CREATE FUNCTION时候创建的,通过RETURNS TABLE 语句定义的。postgresql 在调用函数时,会将结果类型保存到参数FunctionCallInfo里。

还需要注意一点,上述函数返回的第一列是TEXT类型,它并不等于 C 语言的字符串,所以需要使用cstring_to_text函数转换。

返回多行数据

下面展示一个函数,用于生成指定数量的 斐波那契数列。每行数据包括两列(num int, value int)。

c 语言文件

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
#include "postgres.h"
#include "funcapi.h"

PG_MODULE_MAGIC;

// 计算斐波那契数列,需要记录前两个结果的值
typedef struct OlderValues {
    uint64 value1;
    uint64 value2;
} OlderValues;

PG_FUNCTION_INFO_V1(my_records_func);
Datum my_records_func(PG_FUNCTION_ARGS)
{
    FuncCallContext     *funcctx;

    // 判断是否第一次执行
    if (SRF_IS_FIRSTCALL())
    {
        // 进行此次函数第一次初始化
        funcctx = SRF_FIRSTCALL_INIT();
        // 切换内存分配器
        MemoryContext oldcontext = MemoryContextSwitchTo(funcctx->multi_call_memory_ctx);
        // 提取参数,保存到FuncCallContext的max_calls属性里
        funcctx->max_calls = PG_GETARG_UINT32(0);
        // 分配自定义上下文变量
        OlderValues * mydata = palloc(sizeof(OlderValues));
        mydata->value1 = 0;
        mydata->value2 = 1;
        funcctx->user_fctx = mydata;
        //切回内存分配器
        MemoryContextSwitchTo(oldcontext);
    }

    // 获取最新的FuncCallContext
    funcctx = SRF_PERCALL_SETUP();
    uint64 call_cntr = funcctx->call_cntr;
    uint64 max_calls = funcctx->max_calls;

    if (call_cntr < max_calls)
    {
        TupleDesc            tupdesc;
        if (get_call_result_type(fcinfo, NULL, &tupdesc) != TYPEFUNC_COMPOSITE)
            elog(ERROR, "return type must be a row type");

        // 构建返回值,使用调用次数作为第一列的值
        Datum values[2];
        values[0] = UInt64GetDatum(call_cntr);
        // 计算此次斐波那契的值
        uint64 value;
        if (call_cntr == 0) {
            value = 0;
        }
        if (call_cntr == 1) {
            value = 1;
        }
        if (call_cntr > 1) {
            // 取出上下文的变量,也就是前两个斐波那契的值
            OlderValues *mydata = (OlderValues *) funcctx->user_fctx;
            value = mydata->value1 + mydata->value2;
            mydata->value1 = mydata->value2;
            mydata->value2 = value;
        }
        values[1] = UInt64GetDatum(value);

        // nulls 数组表示值是否为空
        bool            nulls[2];
        memset(nulls, 0, sizeof(nulls));
        // 构建返回数据,由HeapTuple表示
        HeapTuple tuple = heap_form_tuple(tupdesc, values, nulls);
        Datum result = HeapTupleGetDatum(tuple);
        // 返回数据
        SRF_RETURN_NEXT(funcctx, result);
    }

    // 数据已经全部返回
    SRF_RETURN_DONE(funcctx);
}

创建函数

1
2
3
4
CREATE FUNCTION my_records_func(number_limit int) 
     RETURNS TABLE (num int ,value int)
AS '/var/lib/pgsql/my_records_func.so', 'my_records_func'
LANGUAGE C STRICT;

执行语句

1
2
3
4
5
6
7
8
9
test=# select * from my_records_func(5);
 num | value 
-------+-----
     0 |   0
     1 |   1
     2 |   1
     3 |   2
     4 |   3
(5 rows)

函数调用上下文

因为要返回多行数据,所以会执行函数多次。这里需要一个上下文变量,来保存需要在执行多次之间传递的变量,比如执行的次数等。FuncCallContext结构体,作为函数调用的内置上下文变量,下面只是列举常见的成员

1
2
3
4
5
6
7
typedef struct FuncCallContext
{
	uint64		call_cntr;      // 已经返回结果的次数
    void	   *user_fctx;      // 用于保存自定义上下文
	MemoryContext multi_call_memory_ctx;   // 用于分配内存
    // .......
} FuncCallContext;

相关宏

上述的程序调用了多个宏,接下来简单介绍下它们的用途:

SRF_IS_FIRSTCALL负责判断是否这是第一次执行。

SRF_FIRSTCALL_INIT负责初始化FuncCallContext,它会设置call_cntr为 0,然后在fn_mcxt内存管理器下,分配一个子管理器,负责这个函数执行时的内存申请。使用内存管理器有一个好处,它会记录所有的申请内存,在所有的行数据返回完之后,会释放掉执行期间申请的所有内存。

SRF_PERCALL_SETUP负责返回最新的FuncCallContext

SRF_RETURN_NEXT负责返回数据,会更新call_cntr自增,设置isDoneExprMultipleResult,表示还有更多的行数据。

SRF_RETURN_NEXT_NULL负责返回空数据,执行过程同SRF_RETURN_NEXT相同。

SRF_RETURN_DONE负责释放函数执行期间申请的所有内存,并且设置isDoneExprEndResult,表示所有的行数据已经返回完了。

updatedupdated2023-07-022023-07-02