RAG实战|向量数据库LanceDB指南

LanceDB介绍LanceDB是一个开源的用 Rust 实现的向量数据库(),它的主要特点是:提供单机服务,可以直接嵌入到应用程序中支持多种向量索引算法,包括Flat、HNSW、IVF等。 支持全文检索,包括BM25、TF-IDF等。 支持多种向量相似度算法,包括Cosine、L2等。

RAG实战|向量数据库LanceDB指南

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/

相关资讯

低成本快速定制大模型,这次我们来深度探讨下RAG 和向量数据库

当今人工智能领域,最受关注的毋庸置疑是大模型。然而,高昂的训练成本、漫长的训练时间等都成为了制约大多数企业入局大模型的关键瓶颈。这种背景下,向量数据库凭借其独特的优势,成为解决低成本快速定制大模型问题的关键所在。向量数据库是一种专门用于存储和处理高维向量数据的技术。它采用高效的索引和查询算法,实现了海量数据的快速检索和分析。如此优秀的性能之外,向量数据库还可以为特定领域和任务提供定制化的解决方案。科技巨头诸如腾讯、阿里等公司纷纷布局向量数据库研发,力求在大模型领域实现突破。大量中小型公司也借助向量数据库的能力快速进

为什么AI需要向量数据库?

大模型火遍全球,DeepSeek、OpenAI、谷歌、百度、抖音等科技巨头争相发布自家产品。 多数人会想当然认为,大模型越大越强大,参数量越多就越聪明。  现实呢?

低成本快速入局大模型,你需要学下向量数据库

在 ChatGPT 等大语言模型(LLM)盛行的当下,一直不温不火的向量数据库(Vector Data Base ,VectorDB)开始受到大家的关注。一般而言,大模型在回答具有普遍性的问题上游刃有余,但在回答垂直专业的问题上,就显得不那么出色,例如,大模型在医疗行业会存在回答不准确的情况。这时,为大模型配备一个「超级大脑」变得尤为重要,这个「超级大脑」可以存储一些专有知识,这样一来,大模型就能从海量的数据中快速检索出最合适的答案,提高它们的准确性和效率,而向量数据库就充当了「超级大脑」角色。其实,早在 2013