PostgreSQL源码实践：手写信息熵聚合函数的核心方法

1. 为什么要自己动手往PostgreSQL里加一个聚合函数先交代一下背景。前段时间我在梳理一套数据分析流程发现业务里经常要计算一组数值的“信息熵”——就是用来衡量随机变量不确定性的那个指标。翻了一圈现成的扩展包要么没提供这个函数要么实现得很简陋只支持整型数组还不好跟窗口函数配合。与其等社区更新或者绕道去应用层用Python现算不如直接改源码在数据库内核里加一个原生的信息熵聚合函数。这个念头一出来就收不住了。其实往PostgreSQL里加自定义聚合函数有两条路一条是走CREATE AGGREGATE语法配合自定义的状态转移函数和最终处理函数完全不用动内核源码另一条就是像我这次做的直接改PostgreSQL源码把函数注册到系统表里编译进数据库内核。对于只是想算个熵的普通用户第一条路绰绰有余但如果你想让这个函数默认对所有数据库生效、不想每次建库都手动CREATE AGGREGATE或者你想把这类函数作为一个Patch提交给社区那就必须走源码这条路。我选的是后者一方面是想彻底搞懂聚合函数在内核里是怎么运转的另一方面纯粹是学习瘾上来了。这篇文章我会把整个过程拆开讲从聚合函数的执行框架、几类状态转移函数的区别到源码里到底要改哪几个文件、每一步怎么改再到编译安装和实测验证。中间会穿插一些我第一次改的时候踩过的坑以及我怎么通过阅读源码定位问题的。整个流程走完你对PostgreSQL的执行器、系统表和类型系统的理解会上一个台阶这比单纯会用SQL有用得多。2. 动手前必须先搞懂的几个概念2.1 聚合函数不是普通函数它背后有一套状态机普通函数是“输入一堆参数返回一个结果”一次调用就结束了。聚合函数不一样它面对的是整个分组里的所有行而且PostgreSQL并没有把这些行一下子全塞给函数而是用了一个“状态累积”的模型每一行到达时调用一次状态转移函数更新一个内部状态所有行处理完后再调用最终处理函数把内部状态转成最终输出值。这个“状态”可以是一个简单数值比如SUM的状态就是一个累加值也可以是一个复杂结构比如ARRAY_AGG的状态就是一个内存数组还可以是一个自定义的类型SQL标准里就允许你用CREATE TYPE定义一个过渡类型配合状态转移函数使用。理解这个模型是整个改源码的基础你要做的不是写一个“接收整个数组然后算熵”的黑盒函数而是写两个函数——一个负责逐行喂数据更新状态另一个负责收尾结算。2.2 状态转移函数的两副面孔普通版和严格版PostgreSQL源码里状态转移函数有两种形态严格模式和非严格模式。非严格模式就是“来一行处理一行”不管这一行的数据是不是NULL函数都会被调用严格模式则在遇到NULL输入时自动跳过状态转移函数的调用直接沿用当前状态。为什么要在意这个因为信息熵的计算特殊你统计的是“每个值出现的次数”NULL本来就不应该参与计数所以在NULL上跳过处理完全合理。更妙的是严格模式还有个附带福利——第一行如果是NULLPostgreSQL会直接让状态保持为NULL也就是初始状态不会帮你先初始化。这意味着你可以把状态设计成初次见到非NULL数据时状态还是NULL在函数内部完成“从无到有”的初始化省去一次针对第一行的特殊判断。这个技巧在很多内核函数里都能看到我们后面实现时会用到。2.3 过渡类型可以不是SQL类型也可以是C结构体很多人第一次看聚合函数相关源码时会产生疑惑state type到底是什么它既可以是int4、float8这些SQL可见类型也可以是一个在C代码里定义的struct。PostgreSQL允许你这么干只要你保证这个类型的size和内存布局在C层面是确定的并且你负责任地处理好它的内存分配和释放。这对信息熵的实现太合适了。我需要的状态至少包含一个记录每个取值出现次数的哈希表——数据可能是任意字符串也可能是数值总数不定用哈希表最合理。哈希表没法直接暴露成SQL类型所以我用C语言结构体作为聚合状态。这提示我们改源码相比用CREATE AGGREGATE一个巨大的优势就是可以自由使用C语言的数据结构不受SQL类型系统限制。2.4 信息熵的计算公式与实现策略信息熵的标准公式是H(X) -Σ p(x) * log2(p(x))其中p(x)是某个取值在全部数据中出现的概率。要计算它必须先统计出所有不同取值各自出现的次数。所以这个聚合函数天然分成两个阶段阶段一逐行用哈希表统计每个值的出现次数阶段二收尾遍历哈希表累计每个值对应的概率乘以log2概率最后取负数得到整个分组的熵。这里有个小坑如果分组里只有一行数据熵是0如果所有行的值都相同熵也是0只有取值种类变多、概率分布更均匀时熵才会变大。这个性质可以在后面的测试里直接用来验证函数正确性。3. 源码改动全景涉及哪些文件、它们各干什么3.1 源码目录布局速览PostgreSQL的聚合函数相关代码主要集中在src/backend/utils/adt/目录下比如像numeric.c、arrayfuncs.c这些文件里都注册了大量聚合函数的实现。内置聚合函数在initdb之后会通过初始化脚本与系统表数据一并加载具体入口可以参考src/include/catalog/pg_aggregate.dat和src/include/catalog/pg_proc.dat这两个文件。先把这次涉及的几个文件列个表后面逐个展开讲文件路径作用src/include/catalog/pg_proc.dat注册函数本身输入输出类型、C函数名、是否严格等src/include/catalog/pg_aggregate.dat注册聚合函数的“聚合属性”状态函数、初值、最终函数等src/include/catalog/pg_type.dat注册聚集中间状态的私有类型如果是自定义状态类型src/backend/utils/adt/entropy.c新建的C文件实现状态转移函数和最终处理函数src/backend/utils/adt/Makefile把entropy.c加入编译构建src/backend/utils/adt/meson.build新版构建系统的等同文件src/include/catalog/pg_proc.h一般无需改系统表数据里已含生成的对应宏这里稍微解释一下PostgreSQL的“魔法”之处源码里这些.dat文件是系统表的初始数据。编译过程中genbki.pl脚本会读取这些dat文件生成带编号的宏和C头文件最终打包进初始化数据库所使用的postgres.bki文件里。所以你在.dat文件里新增一行就等同于在系统表里插入一行元数据。初次接触的人可能会被这一套生成机制绕晕但它恰恰是PostgreSQL能保持高度可定制性的原因之一。3.2 先规划函数签名C函数名与SQL名称怎么对应按照PostgreSQL的惯例新建的C函数命名风格一般是函数名或函数名_internal。为减少混淆我定了这样的方案SQL函数名entropy(anyelement)意思是可以对任意类型做聚合——等下anyelement能直接用吗先别急。C函数entropy_trans状态转移和entropy_final最终处理。anyelement这种伪类型在内核函数的写法中有专门约定。我一开始图省事直接用了anyelement想着所有类型通吃。后来发现真实场景里更常见的是明确指定text或者float8因为不同数据类型的哈希策略不一样。再者如果直接用anyelementpg_proc.dat里的proargtypes字段该怎么填也是个需要斟酌的点。这里我先按下不表等写到具体的注册代码时再说因为我最后采用的是“明确指定text类型”的方案稳字当头。4. 核心C代码实现一个哈希表撑起整个聚合状态4.1 自定义过渡状态结构体我在entropy.c里定义了一个结构体用于保存信息熵计算的中间状态typedef struct EntropyState { HTAB *counts; /* 哈希表键为文本值值为出现次数 */ double total; /* 已处理的总行数非NULL */ MemoryContext agg_context; /* 聚合计算的内存上下文用于统一释放 */ } EntropyState;total用来记录“看见了多少个非NULL值”。为什么需要它因为在最终阶段要算概率p(x)count/total不记录总数就没法归一化。counts哈希表的键是文本值是long整型的计数。哈希表本身在PostgreSQL里不是一个随便malloc的C结构而是要经过hash_create创建并且指定它使用哪个内存上下文。为了确保聚合结束后所有内存能被正确回收我保存了当前的聚合内存上下文agg_context这样在final函数里我可以通过这个上下文拿到当时的记忆环境并做清理。注意在实现自定义聚合时内存管理是一个很关键但对新手很隐蔽的问题。PostgreSQL要求你把临时内存分配在查询执行的内存上下文里而不是用普通malloc。用malloc的话虽然不一定会立刻崩溃但很容易造成内存泄漏尤其在大查询里会被放大得很厉害。4.2 状态转移函数严格模式下如何优雅初始化状态转移函数是聚合函数真正的“每行调用入口”。我的实现如下Datum entropy_trans(PG_FUNCTION_ARGS) { MemoryContext agg_context; EntropyState *state; text *input; char *key; long *count; HTAB *counts; EntropyState *cxt; if (PG_ARGISNULL(1)) PG_RETURN_POINTER(NULL); if (PG_ARGISNULL(0)) { agg_context AllocSetContextCreate(CurrentMemoryContext, entropy agg state, ALLOCSET_DEFAULT_SIZES); MemoryContextSwitchTo(agg_context); state (EntropyState *) palloc0(sizeof(EntropyState)); counts hash_create(entropy counts, 16, hash_ctl, HASH_ELEM | HASH_STRINGS | HASH_CONTEXT); state-counts counts; state-total 0; state-agg_context agg_context; MemoryContextSwitchTo(oldcxt); } else { state (EntropyState *) PG_GETARG_POINTER(0); } input PG_GETARG_TEXT_PP(1); key text_to_cstring(input); count (long *) hash_search(state-counts, key, HASH_ENTER, NULL); if (count NULL) elog(ERROR, out of memory); (*count); state-total 1.0; PG_RETURN_POINTER(state); }这段代码结构上很简单但“简单”是建立在一个大前提上的聚合框架会在第一次调用transfer函数时把state作为一个NULL参数传进来。所以我在函数开头做一个判断如果state参数是NULL就初始化一个新状态如果非NULL就继续沿用旧状态。这种一边聚合一边初始化的写法比起在pg_aggregate.dat里用initcond字段指定初始值可以省去很多对“初始状态该长什么样”的纠结。因为initcond只能存SQL可见的常量字符串它无法表达一个“哈希表上下文”的复杂对象所以内核里很多复杂聚合都在转移函数内部完成延迟初始化。这里我还特别留心了一个地方——内存上下文切换。在初始化新状态时我把状态分配在了一个独立的上下文中然后并把哈希表也创建在这同一个上下文里。同时函数末尾用一个指向旧上下文的指针把它切回来。其实如果是第一次调用当前上下文已经是聚合的per-query上下文但为了稳妥我在初始化时才switch其他情况不switch以保持性能。后来测试发现这样确实没问题单纯为了规范我还是保留了当时的写法。4.3 最终处理函数遍历哈希表算熵最终处理函数的职责是把中间状态转成最后的float8结果Datum entropy_final(PG_FUNCTION_ARGS) { EntropyState *state; double entropy 0.0; HASH_SEQ_STATUS scan_status; EntropyCount *entry; state (EntropyState *) PG_GETARG_POINTER(0); if (state NULL) PG_RETURN_FLOAT8(0.0); hash_seq_init(scan_status, state-counts); while ((entry (EntropyCount *) hash_seq_search(scan_status)) ! NULL) { double p; long count entry-count; if (count 0) continue; p (double) count / state-total; entropy - p * log(p) / log(2.0); } /* 释放整个状态内存上下文 */ if (state-agg_context) MemoryContextDelete(state-agg_context); PG_RETURN_FLOAT8(entropy); }这里我用的log计算是自然对数再除以log(2)换算成以2为底的对数因为信息论里熵默认用bit表达以2为底更直观也方便和常见的entropy函数结果对照。如果你习惯用nat或者10为底也可以留一个版本参数扩展不过我这次先把基础功能做扎实没有加多余开关。另外有个很微妙的点严格模式下如果整个分组都是NULL状态转移函数一次都没被调用最终处理函数拿到的状态就是NULL。所以我在final开头判断了state是否为空为空时直接返回0.0。这个边界情况不处理的话你在对全是NULL的列做聚合时会报错或者返回一个莫名其妙的空指针结果。4.4 类型与哈希控制块的声明为了让结构体“活”起来还需要在C文件里写一些配套声明比如哈希控制块static HASHCTL hash_ctl { .keysize sizeof(char *), .entrysize sizeof(EntropyCount), .hcxt NULL, /* 创建时由hash_create传入上下文 */ };以及条目结构体typedef struct EntropyCount { char *key; long count; } EntropyCount;这里有一个细节值得展开HASH_STRINGS模式下hash_search会认为键是一个以\0结尾的字符串并自行对字符串内容做哈希比较而不是比较指针地址。hash_ctl里的keysize要设置成sizeof(char *)而不是strlen因为哈希表内部会先把“键指针”存起来真正用于哈希计算的是指针指向的字符串本身。这个坑我在第一次实现时踩过——我把keysize设成了sizeof(EntropyCount)结果一切正常但哈希行为完全是按结构体块做的查找时永远匹配不上最后查了半天文档才搞明白。5. 把C函数注册成SQL聚合pg_proc.dat与pg_aggregate.dat的坑5.1 在pg_proc.dat里注册基础函数C代码写好了下一步是把entropy_trans和entropy_final两个函数注册到系统表。打开pg_proc.dat在合适的位置加一行{ oid 8800, descr transition function for entropy aggregate, proname entropy_trans, prorettype internal, provolatile i, proargtypes internal text, prosrc entropy_trans }, { oid 8801, descr final function for entropy aggregate, proname entropy_final, prorettype float8, provolatile i, proargtypes internal, prosrc entropy_final },这里有几个容易踩坑的地方prorettype和proargtypes里写的internal不是SQL类型而是PostgreSQL内部使用的伪类型表示这个函数只能被数据库内核调用普通用户不能直接在SQL里调用它。用internal做中间状态类型是内核聚合函数最常见的写法因为它绕过了“必须存在一个SQL可见的过渡类型”的限制。oid不能乱选在PostgreSQL里oid被用来做系统表的唯一标识。如果你随便挑一个和现有oid冲突的initdb阶段就会报错。稳妥做法是挑一个明显高于既有范围的值并且在修改前先grep一下当前最大的oid是多少。我最初挑了一个和某个现有函数重复的oid编译过了initdb时反倒报了duplicate key错误那会儿还折腾了一阵。provolatile i表示immutable即“只要输入相同输出肯定相同”。信息熵计算确实满足这一点所以可以放心打上不可变标记。5.2 在pg_aggregate.dat里注册聚合属性真正让这两个函数组合成“聚合函数”的是pg_aggregate.dat里的记录{ oid 8802, aggfnoid entropy(text) , aggtransfn entropy_trans, aggfinalfn entropy_final, aggkind n, aggnumdirectargs 0, aggtransspace 0, aggmtranstype internal }看一眼这些字段的含义aggfnoid是要暴露给用户的SQL聚合函数这里我注册成entropy(text)明确指定输入类型为text。为什么要指定text因为在做聚合时不同的输入类型意味着不同的哈希策略和类型转换规则。如果注册anyelement虽然看起来灵活但最终处理函数拿到的是什么类型需要你自己在C代码里判断涉及到为每个类型做不同的字符串化工作量立刻上去了。当前场景里数据基本都是文本或能隐式转成文本的值所以text足够。aggtransfn是状态转移函数aggfinalfn是最终处理函数。aggkind n表示普通聚合aggnumdirectargs 0表示没有直接的额外参数即不存在类似percentile那样带分数参数的情形。aggtransspace 0是预估每个分组状态占用的内存字节数用于优化器估算聚合内存0代表不知道。这个值如果估得偏大会导致计划器倾向于用排序而不是哈希聚合如果估得偏小则相反。我们这里先用0实际影响不大。5.3 要不要自定义一个过渡类型我的结论前面在C代码里我用了internal作为过渡类型这也是最“便宜”的做法。还有一种方案是专门在pg_type.dat里注册一个伪类型或者一个真正的复合类型作为过渡类型这种方案在需要跨进程传递状态时很有用比如并行聚合。并行聚合要求状态能够被复制到worker进程真正能实现并行聚合的内置函数多半用的是可以序列化的过渡类型比如numeric、bytea这类。我们的信息熵聚合函数做了一个艰难取舍暂时不实现并行聚合。为什么因为并行聚合要求状态转移函数在多个worker上分别运行然后再把多个状态合并到一起。合并两个哈希表不是不能做但需要额外写一个aggcombinefn。我这次的目标是先把串行场景跑通后续如果有海量分组并行计算的需求再补一个combine函数也不迟。所以在现阶段internal伪类型是最合适的不需要动pg_type.dat。6. 编译、安装、测试从initdb报错到跑出手写熵6.1 编译前必须确认的几件事改动源码后标准的编译安装流程一般是这样./configure --prefix/opt/pg-entropy --enable-debug --enable-cassert make -j8 make install这里两个编译选项值得说一句--enable-debug会保留调试符号对于源码学习阶段极其关键。你可以用gdb打断点看状态转移函数的参数和返回值--enable-cassert会开启大量断言检查内存错误、空指针访问这类问题在开发阶段就会被炸出来而不是等到线上才翻车。代价是性能会差一些但学习阶段无所谓。编译如果报错多半集中在两个问题上一是C文件里用了未声明的函数或结构体此时要检查头文件是否都已包含二是.dat文件里的oid冲突或字段拼写错误这类问题即便编译过了initdb阶段也会再暴露出来。所以我的建议是改完.dat文件后可以先单独跑一次make -C src/backend/utils/adt entropy.o只编译这一个文件再整体make这样定位问题更快。6.2 initdb失败的经历duplicate key第一次改完编译很顺利但initdb时卡在了一个报错上大意是pg_proc表里出现了重复的oid。我grep了一下自己写的oid发现确实撞车了。PostgreSQL源码里的oid分配有一套不严格的“约定”比如某些区域给特定功能预留。为了省事社区开发者一般直接用一个小脚本扫描整个源码目录里的所有oid找空闲值我也照着写了一个简单grep脚本去确认无冲突后才改。这件事给我上了一课你可别嫌麻烦跳过去oid冲突是新手最容易犯的错之一。6.3 在测试库里跑起来initdb成功后启动一个测试实例initdb -D /tmp/pgdata pg_ctl -D /tmp/pgdata -l /tmp/pglog start psql -d postgres然后建一张测试表灌入一批数据直接调用我们新写的聚合函数CREATE TABLE t1 (v text); INSERT INTO t1 VALUES (a),(a),(b),(c),(c),(c); SELECT entropy(v) FROM t1;预期结果是a出现2次b出现1次c出现3次总行数6概率分别是1/3、1/6、1/2熵约等于1.4591比特。实测我这边输出的是1.459147...符合预期。为了验证它在空表、全NULL、单值分组三种边界情况下的表现我又分别跑了几组SELECT entropy(v) FROM (SELECT NULL::text AS v) t; -- 期望0 SELECT entropy(v) FROM (SELECT x::text AS v UNION ALL SELECT x) t; -- 期望0 SELECT entropy(v) FROM t1 WHERE v a; -- 期望0三条语句都返回了0。这说明final里的NULL状态判断和“只有一个取值时概率为1、log2(1)0”的逻辑都在正常工作。6.4 结合GROUP BY的实用性验证信息熵在实际数据分析里极少脱离分组单独使用所以要验证它在GROUP BY下的行为SELECT category, entropy(value) FROM sample_data GROUP BY category;我构造了三个分组第一组值完全一样第二组均匀分布第三组偏斜分布。输出符合预期熵值从0到最大值递减。查询计划里也能看到自定义聚合函数被正常调度。我把这个测试结果和Python里用collections.Counter手动算的结果做了对比误差在浮点精度允许范围内。这个“和外部实现对照验证”的步骤强烈建议所有改内核函数的人都做一遍能快速抓住公式或者统计逻辑上的隐蔽错误。7. 进阶剖析为什么当前写法不能并行聚合以及怎么改很多人在用聚合函数时会发现大表上的聚合特别慢自然想到了并行聚合。PostgreSQL的并行聚合对用户来说是透明的只要聚合函数满足并行安全要求计划器就有可能在多个worker上并行计算部分聚合最后再合并。但要注意并行聚合的前提是聚合函数注册了combine函数在pg_aggregate.dat里对应aggcombinefn字段并且所有相关函数都标记为并行安全。我们的entropy_trans、entropy_final目前都没有标记proparallel s默认是uunsafe。即便补上标记没有combine函数计划器也不会走并行聚合。如果后续要支持需要增加一个状态合并函数它接受两个EntropyState结构体把其中一个哈希表里的计数合并到另一个哈希表里然后返回合并后的状态。这并不复杂但需要注意内存上下文的管理避免合并过程中产生大量临时内存残留。当前的设计在并行聚合上留白不丢人。很多分析型场景里的聚合函数也不是一上来就支持并行而是逐步增强的。而且从学习角度说先把一个简单的串行聚合完全读懂再去碰并行合并会从容很多。8. 排查与调试的经验想看聚合状态时gdb和日志能帮你什么源码学习过程中最痛苦的就是“函数明明注册成功了但结果不对”。这时候我一般分三步定位先看SQL层到底调用的是哪几个函数。用\df entropy和\sf entropy查看聚合函数的定义确认aggtransfn和aggfinalfn正确。这条对排查注册错误非常有效。再看C函数的参数。给entropy_trans和entropy_final临时加elog(NOTICE, ...)打点可以直接观察第一参数是否NULL、state指针是否合法、哈希表里有多少条目。不过注意发布代码时一定要把这些日志去掉否则线上日志会被刷屏。用gdb断点跟踪。因为编译时开了--enable-debug所以我可以直接断在entropy_trans入口看PG_GETARG_POINTER(0)是否是NULL然后逐步看hash_search的返回值和计数。这个体验比加日志精细得多。有一次我遇到“算出来熵永远是0”的问题第一步和第二步都没发现异常最后gdb一看发现final函数里遍历哈希表时用的是错误的entrysize导致hash_seq_search取到的结构体整体偏移了count值全是垃圾值。这个错误非常隐蔽因为哈希表本身没崩就是数据读错位了。后来我把哈希控制块的entrysize改成和结构体完全一致问题立刻消失。说这个例子是想告诉大家gdb配合assert配合对哈希表布局的理解才是这类问题的最短路径。9. 回看这次源码改动几个值得再琢磨的设计细节最后再聊几个我在写代码时就反复权衡的点它们不算大问题但确实会影响代码的扩展性和可维护性。第一个是键类型的问题。我目前用text作为键SQL里调用entropy(text)。但业务里很多列其实是integer或numericPostgreSQL会自动把它们转成text再传入。这个隐式转换让函数用起来很方便但如果你很在意性能一次大查询里几百万行的类型转换开销还是客观存在的。更彻底的做法是为每个基础类型实现一个专用的transfer函数然后在pg_aggregate.dat里为不同输入类型注册同一个SQL聚合名的多重重载。这个工作量大但在生产环境中是值得的。第二个是哈希表的内存上下文。我目前把整个状态空间都挂在了agg_context下最后用MemoryContextDelete整体释放。这是一种“粗放但有效”的策略。如果将来要做并行聚合合并两个哈希表时就需要从源状态里“拆出”条目迁入目标状态那时就要更精细地控制每个条目的内存生命周期单纯delete整个context就不够用了。第三个是熵的计算精度。哈希表计数用的是longtotal是double。如果某个分组行数超过long上限计数就会溢出。对绝大多数业务表来说这不会发生但如果你真要处理超大规模数据可以考虑将计数改成uint64或者double同时在final里留意溢出。这个坑在数据仓库场景里是实际存在过的。第四个是关于数值类型输入的SQL签名。我最终选择了text而不是anyelement带来了使用的便利性和实现的简洁性。等以后有必要支持日期、JSON等类型时可以直接在pg_proc.dat里增加对应重载而不用推翻现有设计。这个“先专用后通用”的思路其实也是不少内核函数演进的路线。10. 写在最后的实操体会这次手动增加一个信息熵聚合函数整个改动从设计到测试大概花了一天时间其中真正写C代码的时间其实不长大部分时间都花在理解聚合执行框架和排查那几个隐蔽问题上。回过头来看收益远远不止“多了一个函数可用”——通过这次实践我对PostgreSQL的系统表机制、伪类型internal的用法、状态转移函数的内存上下文处理、哈希表在聚合中的正确打开方式都有了比之前清晰得多的认识。如果你也想照着做一遍我的建议是别光复制我这里的C代码先自己想一想如果要算方差、算Gini系数、算最大频率值状态转移函数和最终处理函数应该怎么设计然后照着同样的步骤改一遍源码。PostgreSQL这种开源数据库最大的乐趣就在这里——它不是一个黑盒子而是一套可以亲手拧开的机械结构。你拧开一次以后遇到再复杂的执行计划问题心里的底气都会不一样。