1. LanceDB介绍
LanceDB是一个开源的用 Rust 实现的向量数据库(https://github.com/lancedb/lancedb),它的主要特点是:
- 提供单机服务,可以直接嵌入到应用程序中
- 支持多种向量索引算法,包括Flat、HNSW、IVF等。
- 支持全文检索,包括BM25、TF-IDF等。
- 支持多种向量相似度算法,包括Cosine、L2等。
- 与Arrow生态系统紧密集成,允许通过 SIMD 和 GPU 加速在共享内存中实现真正的零拷贝访问。
2. LanceDB安装
复制pip install lancedb
预览版本:
复制pip install --pre --extra-index-url https://pypi.fury.io/lancedb/ lancedb
3. 快速入门
3.1 连接或者打开数据库
创建数据库:
复制import lancedb db = lancedb.connect("./test") # 如果数据库不存在,会自动创建
打开数据库:
复制db = lancedb.open("./test") # 如果数据库不存在,会报错
3.2 创建表
复制data = [ {"vector": [1, 2], "text": "hello"}, {"vector": [3, 4], "text": "world"}, ] table = db.create_table("my_table", data=data, mode="overwrite") df = pd.DataFrame(data) table = db.create_table("my_table", data=df, mode="overwrite")
3.3 查看当前 db 中的表
复制print(db.table_names())
3.4 插入数据
复制data = [ {"vector": [1, 2], "text": "hello"}, {"vector": [3, 4], "text": "world"}, ] table.add(data)
3.5 查询数据
通过向量查询数据:
复制query = [1, 2] results = table.search(query).limit(1).to_pandas()
通过文本查询数据:
复制query = "hello" results = table.search(query).limit(1).to_pandas()
3.6 创建索引
复制table.create_index()
LanceDB 不会自动创建索引,对于数据量较大的情况下,建议手动创建,否则会走全文检索(速度会比较慢)。
3.7 删除数据
复制table.delete(f'text = "hello"')
删除数据当然也支持 SQL 语法,具体参考官方文档(https://lancedb.github.io/lancedb/sql/#pre-and-post-filtering)。
3.8 删除表
复制db.drop_table("my_table")
注意:如果表不存在,会报错,可以通过传入参数忽略错误 ignore_missing=True。
4. 向量搜索
4.1 什么是向量搜索
向量搜索是一种在高维空间中搜索向量的方法,主要是将原始数据通过嵌入模型得到向量,然后通过向量相似度算法计算向量之间的距离,从而找到最相似的向量。
4.2 embedding
embedding 是将原始数据通过嵌入模型得到向量的过程,嵌入模型可以是预训练的模型,也可以是自己训练的模型,是一种将文本、图像、音频等数据投影到二维空间的方法。
4.3 索引
和关系型数据库一样,向量数据库也需要索引来加速查询,索引是一种数据结构,用于快速查找数据,LanceDB 使用基于磁盘的索引:IVF-PQ,是倒排索引的一种变体,使用PQ 来做压缩嵌入。
PQ 原理主要分为以下步骤:
- 对向量进行分桶,将向量分为多个桶,每个桶包含多个向量,比如 样本的维度 D=1024,桶的数量 M=64,每个桶的维度 16;
- 对样本的每个段分别进行聚类,聚成 k=256(其中 K 可以定义) 个聚类中心,这样整个样本被分为 M*K 个聚类中心,每个聚类中心分配一个 ID(范围 0-K-1);
- 通过上面的聚类和 ID 分配,每个样本就会变成量化后的向量,例如 [28, 100, 99, 255 ...];
- 对于新加入样本按照第一步的分桶方式切分,然后再聚类的分段里面找到最近的类中心,然后将聚类中心的 ID 作为量化后的向量;
通过以上的处理,原来 1024 维度向量*float类型(1024 * 4 字节)被压缩到 64 个字节,大大减少了存储空间和计算量,当然量化是有损的,所以对于数据量不大的情况,可以不使用索引,直接暴力搜索。
4.4 暴力搜索和 ANN 搜索
- 如果要搜索的准确度,执行暴力搜索是一种好的选择,基本上就是对所有的向量进行相似度计算,然后返回最相似的向量,相当于 kNN 搜索。
- kNN 和每个向量都做距离计算,计算量比较大,所以需要使用 ANN 搜索,ANN 搜索是一种基于树的搜索方法,使用树结构来存储向量,然后通过树的搜索来找到最相似的向量。
4.5 HNSW
HNSW 是一种基于图的搜索方法,使用图结构来存储向量,然后通过图的搜索来找到最相似的向量,原理类似跳跃表,通过分层的 k-ANN 图来实现递归搜索。
LanceDB 创建 HNSW 索引样例:
复制data = [ {"vector": row, "item": f"item {i}"} for i, row in enumerate(np.random.random((10_000, 1536)).astype('float32')) ] tbl = db.create_table("my_vectors", data=data) tbl.create_index(index_type=IVF_HNSW_SQ)
5. LanceDB 使用指南
5.1 通过 Pandas DataFrame 插入数据
复制import pandas as pd import lancedb import numpy as np db = lancedb.connect("./test") table = db.create_table("my_table", data=[], mode="overwrite") df = pd.DataFrame({ "vector": [np.random.rand(100) for _ in range(100)], "text": [f"hello {i}" for i in range(100)], }) table.add(df)
5.2 通过 Arrow Table 插入数据
复制import pyarrow as pa import lancedb import numpy as np db = lancedb.connect("./test") table = db.create_table("my_table", data=[], mode="overwrite") table = db.create_table("my_table", data=pa.Table.from_pandas(df), mode="overwrite")
5.3 通过 Model 插入数据
复制import lancedb from lancedb.pydantic import LanceModel class MyModel(LanceModel): vector: list[float] text: str db = lancedb.connect("./test") table = db.create_table("my_table", schema=MyModel, mode="overwrite") model = MyModel(vector=[1, 2], text="hello") table.add(model)
5.4 通过迭代器写入大规模数据
复制import lancedb import pyarrow as pa def make_batches(): for i in range(1000): yield pa.Table.from_pandas(pd.DataFrame({ "vector": [np.random.rand(100) for _ in range(100)], "text": [f"hello {i}"for i in range(100)], })) schema = pa.schema( [ pa.field("vector", pa.list_(pa.float32(), 4)), pa.field("item", pa.utf8()), ] ) db = lancedb.connect("./test") table = db.create_table("my_table", make_batches(), schema=schema, mode="overwrite")
或者通过迭代器写入数据:
复制import lancedb import pyarrow as pa def make_batches(): for i in range(1000): yield pa.Table.from_pandas(pd.DataFrame({ "vector": [np.random.rand(100) for _ in range(100)], "text": [f"hello {i}" for i in range(100)], })) db = lancedb.connect("./test") table = db.create_table("my_table", data=[], mode="overwrite") table.add(make_batches())
5.5 删除指定的数据
复制db = lancedb.connect("./test") data = [ {"x": 1, "vector": [1, 2]}, {"x": 2, "vector": [3, 4]}, {"x": 3, "vector": [5, 6]}, ] # Synchronous client table = db.create_table("delete_row", data) table.to_pandas() # x vector # 0 1 [1.0, 2.0] # 1 2 [3.0, 4.0] # 2 3 [5.0, 6.0] table.delete("x = 2") table.to_pandas() # x vector # 0 1 [1.0, 2.0] # 1 3 [5.0, 6.0]
5.6 更新数据
复制db = lancedb.connect("./test") data = [ {"x": 1, "vector": [1, 2]}, {"x": 2, "vector": [3, 4]}, {"x": 3, "vector": [5, 6]}, ] # Synchronous client table = db.create_table("update_row", data) table.update(where="x = 2", values={"vector": [10, 10]})
5.7 一致性
由于 lancedb 是嵌入到各个应用中,所以数据更新并不能保持一致,可以通过设置 read_consistency_interval 参数来保证数据更新的一致性。
read_consistency_interval 是一个时间间隔,单位是秒。
- 不设置,数据库不检查其他进程对表所做的更新。这提供了最佳查询性能,但意味着客户端可能无法看到最新的数据,此设置适用于在表引用的生命周期内数据不会发生变化的应用程序。
- 如果设置为 0,数据库在每次读取时检查更新。这提供了最强的一致性保证,确保所有客户端都看到最新提交的数据,但是,它的开销最大。当一致性比高 QPS 更重要时,此设置是合适的。
- 自定义间隔时间,数据库以自定义间隔(例如每 5 秒)检查更新。这提供了最终一致性,允许写入和读取操作之间有一些滞后,从性能方面来看,这是强一致性和无一致性检查之间的中间地带,此设置适用于即时一致性并不重要但客户端最终应该看到更新数据的应用程序。
from datetime import timedelta uri = "data/sample-lancedb" # 保障数据最终一致性 db = lancedb.connect(uri, read_consistency_interval=timedelta(secnotallow=5)) tbl = db.open_table("test_table")
5.8 构建 ANN 索引
复制import lancedb import numpy as np db = lancedb.connect("./test") data = [ {"vector": row, "item": f"item {i}"} for i, row in enumerate(np.random.random((10_000, 1536)).astype('float32')) ] tbl = db.create_table("my_vectors", data=data) tbl.create_index(distance_type='l2', num_partitinotallow=2, num_sub_vectors=4)
- distance_type:距离度量算法,可以参考(cosine, l2)等;
- num_partitions: 分区数量;
- num_sub_vectors:子向量数量,PQ 的子向量数量;
- num_bits:用于编码的子向量的位数,支持 4 和 8;
如果需要通过 CUDA 加速,可以增加参数:
复制tbl.create_index(distance_type='l2', num_partitinotallow=2, num_sub_vectors=4, accelerator='cuda')
5.9 搜索数据
1)kNN 搜索:不建立索引,就会扫描全表,计算每个向量的距离,然后返回最相似的 k 个向量,也可以指定距离度量算法。
复制query = np.random.random(1536).astype('float32') results = tbl.search(query).limit(10).distance_type("cosine").to_pandas()
2)ANN 搜索:通过索引搜索,支持 nprobes 和 refine_factor 参数。
- nprobes 数字越大,搜索越准确,但是速度越慢;
- refine_factor 对召回的进行重排优化;
query = np.random.random(1536).astype('float32') tbl.search(query).limit(2).nprobes(20).refine_factor( 10 ).to_pandas()
3)基于距离范围搜索:主要用于不通过topk 查询,而是通过距离范围查询。
复制query = np.random.random(1536).astype('float32') tbl.search(query).distance_range(0.1, 0.5).to_pandas()
4)全文搜索:如果需要对字符串进行索引,并通过关键字搜索进行查询,可以通过创建 FTS 索引。
复制from lancedb.index import FTS tbl = db.create_table("my_vectors", data=[ {"vector": np.rand.random(10), "item": f"this item {i}"} {"vector": np.rand.random(10), "item": f"this item {i + 100}"} ]) tbl.create_fts_index("text", use_tantivy=False) tbl.search("this item 10").limit(10).select(["item"]).to_pandas()
5)过滤搜索:通过 SQL 语法进行过滤搜索。
复制tbl.search("this item 10").limit(10).where("item='this'", prefilter=True).to_pandas()
5.10 SQL 语法
LanceDB 支持 SQL 语法如下:
复制>, <, >=, <= AND, OR, NOT IS NULL, IS NOT NULL IS TRUE, IS FALSE IN LIKE, NOT LIKE CAST regexp_match(column, pattern)
样例如下:
复制table.search("this item 10").where( "(item IN ('item 0', 'item 2')) AND (id > 10)" ).to_arrow()
6. LanceDB 结合 embedding
6.1 注册 embedding 模型
LanceDB 支持结合 embedding 模型进行搜索。
复制import lancedb from lancedb.pydantic import LanceModel, Vector from lancedb.embeddings import get_registry db = lancedb.connect("./test") func = get_registry().get("openai").create(name="text-embedding-ada-002") class Words(LanceModel): text: str = func.SourceField() vector: Vector(func.ndims()) = func.VectorField() table = db.create_table("words", schema=Words, mode="overwrite") table.add( [ {"text": "hello world"}, {"text": "goodbye world"} ] ) query = "greetings" actual = table.search(query).limit(1).to_pydantic(Words)[0] print(actual.text)
通过 get_registry() 注册不同的模型参数,其中支持的代码如下:
- get_registry().get("sentence-transformers").create(name="BAAI/bge-small-en-v1.5", device="cpu")
- get_registry().get("huggingface").create(name='facebook/bart-base')
- get_registry().get("ollama").create(name="nomic-embed-text")
- get_registry().get("openai").create(name="text-embedding-ada-002")
- et_registry().get("instructor").create(source_instructinotallow="represent the docuement for retreival", query_instructinotallow="represent the document for retreiving the most similar documents")
- get_registry().get("gemini-text").create()
- get_registry().get("open-clip").create()
- get_registry().get("imagebind").create()...
6.2 完整的使用样例
1)注册 embedding 函数
复制from lancedb.embeddings import get_registry registry = get_registry() clip = registry.get("open-clip").create()
2)定义数据模型
复制from lancedb.pydantic import LanceModel, Vector from typing import List class Document(LanceModel): id: str vector: Vector(clip.ndims()) = clip.VectorField() image_uri: str = clip.SourceField()
3)创建表并添加数据
复制db = lancedb.connect("~/lancedb") table = db.create_table("pets", schema=Pets) table.add([{"image_uri": u} for u in uris])
4)查询数据
复制results = ( table.search("dog") .limit(10) .to_pandas() )
参考
(1)https://lancedb.github.io/lancedb/
(2)https://excalidraw-phi-woad.vercel.app/