sqlite-vec一个SQLite3高效向量搜索扩展--JDBC环境使用

最近要用SQLite3，之前放出来了SQLiteUtile工具，方便操作。今天发现AIGC方面，RAG知识库需要使用向量数据库，来存储知识信息。一般呢都是用mysql，但无奈的是mysql就是不让用。突然又发现SQLite3有向量库扩展组件，索性直接搞下来。用了一下。还可以。

SQLite3的向量库扩展extension，是个开源项目，名字叫sqlite-vec。目前我用到最新版本是0.1.5，配套使用的JDBC是SQLite3.47.0

使用需要注意的是，根据操作系统的不同，下载不同的Release版本库，一般linux要so的，windows要dll的。另外还要注意，下载64位版本的话，JDK、操作系统都得是配套的64位，否则会出现找不到模块的问题。

windows开发环境中，下载sqlite-vec-0.1.5-loadable-windows-x86_64.tar.gz。解压缩后得到vec0.dll。

在工程路径下创建一个extension文件夹，将vec0.dll放进去，便于程序运行时指定相对路径，访问到dll或者so。

demo代码如下：

package org.superx.demo.sqltools;

import org.sqlite.SQLiteConfig;

import java.sql.Connection;
import java.sql.DriverManager;
import java.sql.ResultSet;
import java.sql.Statement;

/***
 *@title DemoSQLiteVec
 *@description JDBC环境下SQLite扩展sqlite-vec的使用，向量数据库支持。这个示例工程，只能运行一次。二次运行请把./data/sqlite_vec.db删掉
 *@author superX
 *@version 1.0.0
 *@create 2024/11/20 下午3:51
 **/
public class DemoSQLiteVec {
    public static void main(String[] args) {
        // SQLite 连接字符串，建立一个向量数据库
        String url = "jdbc:sqlite:.\\data\\sqlite_vec.db";
        // 创建向量表的SQL语句
        String createTableSQL = "create virtual table IF NOT EXISTS vec_examples using vec0(sample_embedding float[8])";


        String insertDataSQL1 = "insert into vec_examples(rowid, sample_embedding) " +
                "values" +
                "(1, '[-0.200, 0.250, 0.341, -0.211, 0.645, 0.935, -0.316, -0.924]')," +
                "(2, '[0.443, -0.501, 0.355, -0.771, 0.707, -0.708, -0.185, 0.362]')," +
                "(3, '[0.716, -0.927, 0.134, 0.052, -0.669, 0.793, -0.634, -0.162]')," +
                "(4, '[-0.710, 0.330, 0.656, 0.041, -0.990, 0.726, 0.385, -0.958]')";

        String selectSQL = "select rowid,distance " +
                "from vec_examples " +
                "where sample_embedding match '[0.890, 0.544, 0.825, 0.961, 0.358, 0.0196, 0.521, 0.175]' " +
                "order by distance limit 2";

        //数据内容表，只能通过rowid来进行关联
        String createTableSQL2 = "CREATE TABLE IF NOT EXISTS vec_metadata (rowid INTEGER PRIMARY KEY, describe TEXT, label TEXT)";
        String insertDataSQL2 = "insert into vec_metadata(rowid, describe, label) " +
                "values" +
                "(1,'数据描述1','数据标签1')," +
                "(2,'数据描述2','数据标签2')," +
                "(3,'数据描述3','数据标签3')," +
                "(4,'数据描述4','数据标签4')";

        //数据内容表，只能通过rowid来进行关联查询。而且要注意，vec_examples的查询必须是独立的子查询，否则总是会报错
        //[SQLITE_ERROR] SQL error or missing database (A LIMIT or 'k = ?' constraint is required on vec0 knn queries.)
        String selectSQL2 = "SELECT ve.rowid, ve.sample_embedding, vm.describe, vm.label " +
                "FROM  (SELECT rowid, sample_embedding, distance " +
                "     FROM vec_examples " +
                "     WHERE sample_embedding MATCH '[0.890, 0.544, 0.825, 0.961, 0.358, 0.0196, 0.521, 0.175]' " +
                "     ORDER BY distance " +
                "     LIMIT 2) ve " +
                "JOIN vec_metadata vm ON ve.rowid = vm.rowid ";


        // 创建sqlite配置对象，启用加载扩展功能
        SQLiteConfig config = new SQLiteConfig();
        config.enableLoadExtension(true);

        // 使用配置初始化数据库连接
        try (Connection conn = DriverManager.getConnection(url, config.toProperties());
             Statement stmt = conn.createStatement()) {
            // 加载sqlite-vec扩展库，注意这里dll只能在windows下使用，如果是linux应该是so
            stmt.execute("SELECT load_extension('./extension/vec0.dll')");

            // 建表
            stmt.execute(createTableSQL);


            // 插入数据
            stmt.execute(insertDataSQL1);

            // 查询数据
            ResultSet rs = stmt.executeQuery(selectSQL);

            // 打印结果
            while (rs.next()) {
                int id = rs.getInt("rowid");
                float a = rs.getFloat("distance");
                System.out.println("Row ID: " + id + " distance: " + a);
            }

            //创建关联信息表、插入数据并查询
            stmt.execute(createTableSQL2);
            stmt.execute(insertDataSQL2);
            ResultSet rs2 = stmt.executeQuery(selectSQL2);
            while (rs2.next()) {
                int id = rs2.getInt("rowid");
                String describe = rs2.getString("describe");
                String label = rs2.getString("label");
                System.out.println("Row ID: " + id + " describe: " + describe + " label: " + label);
            }

/* CMD打印信息结果：代表成功
Row ID: 2 distance: 2.3868737
Row ID: 1 distance: 2.389785
Row ID: 2 describe: 数据描述2 label: 数据标签2
Row ID: 1 describe: 数据描述1 label: 数据标签1
 */


        } catch (Exception e) {
            e.printStackTrace();
        }

    }
}

sqlite-vec创建的向量表限制还是比较多的，virtual table ，using vec0，里面不能随便增加字段。增加字段会报错。所以虚拟向量表是依靠rowid与其他表关联的。这一点要注意！！！

即一般我们做RAG应用时，会embeding文本成向量，然后把向量、文本成对儿存储，再用向量检索机制来寻找最相似的向量对应的文本。使用sqlite-vec的话，需要最少创建2张表。

1张虚拟向量表，只存储rowid和向量信息。另一张表，存储rowid和文本信息，或其它标签信息。检索时，需要进行双表关联检索才能得到想要的信息。如demo

另外需要注意的是，双表关联的语法也有要求。因为sqlite-vec实现是用KNN进行相似搜索，所以查询虚拟向量表时，必须是单表查询，且必须指定limit记录数。所以，关联操作必须以子查询方式进行关联。其它方式都会报错：
[SQLITE_ERROR] SQL error or missing database (A LIMIT or 'k = ?' constraint is required on vec0 knn queries.)

上面的坑为各位踩过了。demo只能运行一遍，因为第二遍运行insert时会报主键冲突，所以多次运行的话，运行前最好把./data/sqlite_vec.db删掉。

SQLite向量扩展开源项目链接，需要其它系统的链接库，自己下载即可：

GitHub - asg017/sqlite-vec: A vector search SQLite extension that runs anywhere!