目录
一、概述
最近公司有个任务是向opengauss移植pgvector扩展，opengauss的可能是基于PG内核9.1，内核中许多基础设施，与最新的PG并不相同，而pgvector是基于至少PG12以上的版本。
这个过程我学到很多，编写PG扩展，到底可以多大程度的改变PG的运行？都有哪些API（基础设施/机制）可用？加载到内核中的扩展，是否可以为所欲为？
答案是，差不多可以为所欲为吧。
加载到Linux内核中的驱动程序运行在内核态，也可以为所欲为，修改内核数据结构、直接改硬件寄存器，也会使用Linux内核提供的基础设施。
加载到PG内核中的扩展和Linux驱动程序不同，虽然也可以修改PG内核数据结构，破坏PG内核的运行机制，但是只能运行在用户态。
虽然通过自己实现 _PG_init() 可以在so加载时，调用so中的函数调用，可以在里面为所欲为，但是，一般我们还是按照PG内部基础设施提供的API，按照安全的编程规范实现扩展功能，关于这些编程规范和API，有些在PG文档和代码里的README有详细描述，有些缺乏文档，只能看其他人写的，使用这些API的代码作为例子。
以自定义索引为例（围绕自定义索引所需要的知识，正是pgvector实现的核心），PG支持不同程度的自定义索引，例如，文档中所描述的，先自定义complex类型，然后再自定义基于它的（btree）索引
（自定义基础类型complex见文档PostgreSQL: Documentation: 9.1: User-defined Types）
PostgreSQL: Documentation: 9.1: Interfacing Extensions To Indexes
你可以自定义一种类型，如果还希望在此类型上建立（ btree） 索引，可以定义operator class：
CREATE OPERATOR CLASS complex_abs_ops
DEFAULT FOR TYPE complex USING btree AS
OPERATOR 1 < ,
OPERATOR 2 <= ,
OPERATOR 3 = ,
OPERATOR 4 >= ,
OPERATOR 5 > ,
FUNCTION 1 complex_abs_cmp(complex, complex);
它告诉PG内核，对complex类型创建索引时，使用btree的索引访问方法，为此，必须实现一组operator（>, <, = 等等），每个operator对应一个编号（strategy number），每个编号代表指定的语义，与之对应的operator必须实现其语义，btree的代码中以编号调用operator，操作自定义的complex数据。
具体到我们的向量类型和索引，如果只是二维向量，直接就可以用PG内置的索引和类型，但是这里我们要定义可变长的向量类型，涉及到自定义TOAST基础类型，这部分知识倒也不难，参考 PostgreSQL: Documentation: 9.1: User-defined Types 就够了。
然后定义向量的各种操作符，有了这些，PG就支持变长变量的存储、查询和距离计算了，只是要想加速，还需要自定义索引算法。
我们可以基于btree索引，定义一个operator class，就可以在自定义向量上建立btree索引了。
但是btree索引，并不能计算在多维向量数据库中求最近n个向量的问题，为此必须自定义索引访问方法（Access Method），这就是pgvector的价值所在，pgvector-0.6.0 中定义了两种索引访问方法，分别是ivfflat和hnsw，它们功能是一样的，只是用了不同的索引算法。
PG提供了一套自定义索引的编程机制和API（Index Access Method Interface），它给用户提供了足够的发挥空间，又帮助用户完成了许多工作。想要使用这套机制和API，用户需要编写一个extension，在其中调用API或实现回调函数（接口）来实现自己的索引算法。
PostgreSQL: Documentation: 12: Chapter 61. Index Access Method Interface Definition简单说一下PG extension的知识，一般来说一个编译好的extension包括control文件、一个sql文件，一个so文件，在PG命令行输入create extension xxx 时，数据会找相应的control文件，加载so文件，然后执行sql文件，可能细节上这一描述有点误差，但这是我理解pgvector插件的一个最简单的模型，而且是有效的，这里我不想深究extension的其它机制。
总之，当执行 create extension vector 时，vector.sql会被执行，可以看到里面创建了自定义的类型和各种函数、操作符，其中许多函数都是C语言实现的，因此这些逻辑应该预先存在于PG内核中，也就是应该先加载 vector.so，我猜测只要sql文件中不涉及C语言实现的自定义函数或者基础类型，extension也可以没有C代码写的so。
下面就来解读一下pgvector-0.6.0的代码，代码来自，请对比阅读pgvector/pgvector: Open-source vector similarity search for Postgres (github.com)
前面说过，pgvector-0.6.0实现了两种Index access method，分别是 ivfflat 和 hnsw，在编程上这两者可以互相独立存在，没有互相依赖，所有可以用SQL编写的逻辑都在 vector.sql 中，包括ivfflat 和 hnsw的访问方法的定义、各种函数和运算符的定义、operator class的定义等等。
src目录下的.c文件是自定义索引的逻辑代码，也包括一些在vector.sql中定义的函数的C语言实现，它们最终都会编译到so中。
二、IVFFLAT
vector.sql
定义了 vector 基础类型，定义基础类型需要定义in、out、send、receive函数和C语言结构体，细节可以看 PostgreSQL: Documentation: 9.1: User-defined Types 不难理解。
然后是对于这种类型的函数和操作符，这些对于 ivfflat 和 hnsw 是共用的。
其中，operator 是通过函数实现的，这些函数在 sql 文件中也有定义，实现则是C语言代码，RESTRICT 和 JOIN 参数则是指定一些方法，跟优化器有关，这里用的是 PG 内核内建的函数，没有必要深入研究。
然后是 ACCESS METHOD，注意 OpenGauss 原来是不支持创建ACCESS METHOD，这部分内核代码由 向量数据库功能内核相关(pgvector插件依赖) · Pull Request !4845 · openGauss/openGauss-server - Gitee.com 修改后可以支持。
对于 ivfflat 和 hnsw，分别创建两套 ACCESS METHOD ，注意两个信息：
1、ACCESS METHOD 的名称（紫色部分）
2、handler函数（黄色部分）
CREATE ACCESS METHOD ivfflat TYPE INDEX HANDLER ivfflathandler;
ACCESS METHOD 的名称是创建索引时需要指定的，例如：
CREATE INDEX ON items USING ivfflat (embedding vector_l2_ops) WITH (lists = 100);
（红色部分是operator class）
handler函数是执行CREATE ACCESS METHOD时PG内核调用的函数，这个函数也要自己定义，在这个函数里，要填充实现函数的指针，填充到 IndexAmRoutine 结构体里，PG内核会根据这个结构体创建一条 pg_am 的记录。一个访问方法本质上是一条 pg_am 记录，之所以这里要按照这个规则填写，因为自定义索引，要按照PG内核提供框架的来编写的，不是用户完全自由发挥。
然后是operator class的定义，一套访问方法，可能对应多于一个operator class，例如，对于ivfflat索引，计算向量距离有L2、IP、CONSINE多种方法，当我们创建ivfflat索引时，指定不同的距离计算方法，就是选择了不同的operator class，下面图中包含 DEFAULT FOR TYPE 的operator class，是不指定operator class时，创建ivfflat索引选择的默认operator class：
1、如果创建ivfflat索引时指定 vector_l2_ops，则索引代码使用vector_l2_ops operator class提供的函数，按欧氏计算距离。注意，欧氏距离越小两个向量越相似！
欧氏距离的计算公式：
2、如果创建ivfflat索引时指定 vector_ip_ops，则索引代码使用vector_ip_ops operator class提供的函数，按内积计算两个向量的相似度。注意，内积越大两个向量越相似！
内积计算公式：
把两个n维向量想象成二维空间中两个都以（0，0）为起点的向量，内积的中蕴含的信息包括了两个n维向量在长度和夹角上的差别。
对于IP距离，可以在计算IP距离前对两个向量进行normalize，把它们的模都变成1，然后在计算IP距离，如果这样的话，算出的IP距离等价于cosine距离。
3、如果创建ivffat索引时指定 vector_consine_ops，索引代码使用vector_consine_ops operator class提供的函数。注意，余弦相似度越大两个向量越相似！
把两个n维向量想象成二维空间中两个都以（0，0）为起点的向量，但是与内积不同的是，cos需要把两个向量的模长度都normalize为1，然后计算内积，cos距离只关注两个向量“夹角”属性的差异，并不关心“模长度”属性的差异，而内积既关心“夹角”的差异也关心“模长度”的差异。
关于向量距离的选择，见这篇文章，除了解释了三种距离的计算方法，作者提出，比较向量相似度时选择哪种距离算法，与模型训练时选择哪种距离有关：
Vector Similarity Explained | Pinecone
vector.c
在vector.sql中定义的存储过程的实现，基本上都在这个文件里。
还定义了这个extension的初始化函数：_PG_init
定义了vector基础类型所必须的vector_in/out 等函数，注意这些函数在分配内存时使用的时palloc，那么用的是哪个上下文呢？
这个文件还定义了各种距离计算函数，这些函数不仅仅是C函数，而是定义为PG的存储过程，这样是为了在operator class里引用和在index method代码里调用：
vector_l2_squared_distance -- 计算两个向量的欧氏距离，但是不取平方根，对于比较两个向量的距离大小足以。
l2_distance -- 计算两个向量的欧氏距离，取平方根。
inner_product -- 计算两个向量的内积。
vector_negative_inner_product -- 计算两个向量的内积，再取负数，因为ip越大越相似，与目标向量最近的前k个向量，其实与之计算ip最大的k个向量，如果用order by <ip相似度> 把它选出来们，需要对结果做降序排序，order by降序排序在PG中不能用索引，若想用索引，只有order by x按x升序，如果内积取负，ip最大的k个向量，取负以后就是最小的k个向量，order by -<ip形似度>就既可以是升序，又可以选出ip最大的前k个向量，这样也可以用索引了。
vector_norm -- 计算一个向量的模长度，即各分量平方的和的开方。
vector_spherical_distance -- 计算两个向量的球体距离，其实就是圆的弧长，但是圆的弧长和圆的半径有关，那这个半径是多少呢？这个函数假设两个向量都是单位向量（模为1的向量），如果两个向量不是单位向量，就需要先对两个向量进行normalize，再由这个函数计算弧长，这样弧长其实是单位圆的弧长。
上面几个函数都是 operator class 的支持函数，索引代码里，计算距离时会引用到这些函数。
对于向量的其它操作符（函数）和转换函数，作为C实现的存储过程，也在 vector.c 中定义。
ivfflat.c
这里定义整个 ivfflat 索引公用的一些函数和数据结构，定义了handler函数的实现 ivfflathandler，当执行CREATE ACCESS METHOD 语句时，ivfflathandler 函数被调用，填充索引相关的实现函数，在pg_am中加入一条记录。
对于OpenGauss，由于没有这种机制，需要将每个实现函数创建为存储过程，然后在pg_am中记录存储过程名，然后通过存储过程名来调用实现函数，OG对这些函数进行了封装，例如，存储过程 ivfflatbuild 内部调用了 ivfflatbuild_internal 等等。
这些实现函数中，有一些对于我们自定索引非常重要，且必须自己定义，有一些不那么重要，可以使用内建的标准定义，例如：
这些函数对于自定义索引非常重要，必须按照自定义索引的逻辑自己实现：
ivfflatbuildivfflatbuildempty
ivfflatinsert
ivfflatbulkdelete
ivfflatvacuumcleanup
ivfflatbeginscan
ivfflatrescan
ivfflatgettuple
ivfflatendscan
这些函数可以照搬PG的标准实现，在pgvector的代码中，这几个函数我没有去详细理解，这里不打算解读：
ivfflatcostestimate
ivfflatoptions
ivfflatvalidate
参考 PostgreSQL: Documentation: 16: 64.2. Index Access Method Functions
IvfflatInit 函数是 _PG_init 调用的，这里定义了创建和使用 ivfflat 索引时，可以使用的参数，例如，建立索引时，把所有的向量分成几个“区域”，即lists：
CREATE INDEX ON items USING ivfflat (embedding vector_cosine_ops) WITH (lists = 100);
</
最低0.47元/天 解锁文章
扫一扫
3583
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
