为什么选择 LanceDB？
LanceDB 是一个专为 AI 应用和向量搜索设计的开源向量数据库。它提供了以下特性：
嵌入式数据库：直接在应用程序中工作，无需外部服务器。 高性能：利用 Arrow 格式实现高效的数据存储和检索。 可扩展性：高效处理 TB 级数据集。 向量索引：支持开箱即用的精确和近似最近邻搜索。
结合 Rig 的嵌入和 LLM 功能，您可以用最少的代码创建一个强大而高效的语义搜索解决方案。让我们开始吧！
您可以在我们的 GitHub 仓库中找到该项目的完整源代码。
前提条件
在开始之前，请确保您已具备以下条件：
已安装 Rust（rust-lang.org） 一个 OpenAI API 密钥（platform.openai.com） 基本的 Rust 和异步编程知识
项目设置
首先，创建一个新的 Rust 项目：
cargo new vector_search
cd vector_search
更新您的Cargo.toml 以添加必要的依赖项：
[dependencies] rig-core = "0.4.0" rig-lancedb = "0.1.1" lancedb = "0.10.0" tokio = { version = "1.40.0", features = ["full"] } anyhow = "1.0.89" futures = "0.3.30" serde = { version = "1.0.210", features = ["derive"] } serde_json = "1.0.128" arrow-array = "52.2.0"在项目根目录创建一个.env 文件以存储您的 OpenAI API 密钥：
echo "OPENAI_API_KEY=your_key_here" > .env
构建搜索系统
我们将其分解为可管理的步骤。首先，让我们创建一个实用函数来处理 Rig 的嵌入与 LanceDB 格式之间的数据转换。在src/utils.rs 中创建：
use std::sync::Arc;
use arrow_array::{
types::Float64Type, ArrayRef, FixedSizeListArray,
RecordBatch, StringArray
};
use lancedb::arrow::arrow_schema::{DataType, Field, Fields, Schema};
use rig::embeddings::DocumentEmbeddings;
// 定义 LanceDB 表的架构
pubfn schema(dims: usize) -> Schema {
Schema::new(Fields::from(vec![
Field::new("id", DataType::Utf8, false),
Field::new("content", DataType::Utf8, false),
Field::new(
"embedding",
DataType::FixedSizeList(
Arc::new(Field::new("item", DataType::Float64, true)),
dims asi32,
),
false,
),
]))
}
这个schema 函数定义了我们表的结构：
id：每个文档的唯一标识符。content：文档的文本内容。embedding：内容的向量表示。dims参数：表示嵌入向量的大小（例如，OpenAI 的 ada-002 模型为 1536）。
接下来，添加转换函数以将DocumentEmbeddings 转换为 LanceDB 的RecordBatch：
pub fn as_record_batch(
records: Vec<DocumentEmbeddings>,
dims: usize,
) -> Result<RecordBatch, lancedb::arrow::arrow_schema::ArrowError> {
let id = StringArray::from_iter_values(
records
.iter()
.flat_map(|record| (0..record.embeddings.len())
.map(|i| format!("{}-{i}", record.id)))
.collect::<Vec<_>>(),
);
let content = StringArray::from_iter_values(
records
.iter()
.flat_map(|record| {
record
.embeddings
.iter()
.map(|embedding| embedding.document.clone())
})
.collect::<Vec<_>>(),
);
let embedding = FixedSizeListArray::from_iter_primitive::<Float64Type, _, _>(
records
.into_iter()
.flat_map(|record| {
record
.embeddings
.into_iter()
.map(|embedding| embedding.vec.into_iter().map(Some).collect::<Vec<_>>())
.map(Some)
.collect::<Vec<_>>()
})
.collect::<Vec<_>>(),
dims asi32,
);
RecordBatch::try_from_iter(vec![
("id", Arc::new(id) as ArrayRef),
("content", Arc::new(content) as ArrayRef),
("embedding", Arc::new(embedding) as ArrayRef),
])
}
这个函数将我们的 Rust 数据结构转换为 LanceDB 内部使用的 Arrow 列格式：
为 ID 和内容创建字符串数组。 将嵌入转换为固定大小的列表。 将所有内容组装到一个 RecordBatch中。
准备好实用函数后，让我们在src/main.rs 中构建主要的搜索功能。我们将逐步实现这一过程，并解释每个部分。
设置依赖项
首先，让我们导入所需的库：
use anyhow::Result;
use arrow_array::RecordBatchIterator;
use lancedb::{index::vector::IvfPqIndexBuilder, DistanceType};
use rig::{
embeddings::{DocumentEmbeddings, EmbeddingModel, EmbeddingsBuilder},
providers::openai::{Client, TEXT_EMBEDDING_ADA_002},
vector_store::VectorStoreIndex,
};
use rig_lancedb::{LanceDbVectorStore, SearchParams};
use serde::Deserialize;
use std::{env, sync::Arc};
mod utils;
use utils::{as_record_batch, schema};
这些导入包括：
Rig 的嵌入和向量存储工具。 LanceDB 的数据库功能。 用于高效处理的 Arrow 数据结构。 用于序列化、错误处理和异步编程的实用工具。
定义数据结构
我们将创建一个简单的结构体来表示我们的搜索结果：
#[derive(Debug, Deserialize)]
struct SearchResult {
content: String,
}
这个结构体映射到数据库记录，表示我们想要检索的内容。
生成嵌入
生成文档嵌入是我们系统的核心部分。让我们实现这个功能：
async fn create_embeddings(client: &Client) -> Result<Vec<DocumentEmbeddings>> {
let model = client.embedding_model(TEXT_EMBEDDING_ADA_002);
// 设置虚拟数据以满足 IVF-PQ 索引的 256 行要求
let dummy_doc = "Let there be light".to_string();
let dummy_docs = vec![dummy_doc; 256];
// 为数据生成嵌入
let embeddings = EmbeddingsBuilder::new(model)
// 首先添加我们的真实文档
.simple_document(
"doc1",
"Rust provides zero-cost abstractions and memory safety without garbage collection.",
)
.simple_document(
"doc2",
"Python emphasizes code readability with significant whitespace.",
)
// 使用枚举生成唯一 ID，添加虚拟文档以满足最低要求
.simple_documents(
dummy_docs
.into_iter()
.enumerate()
.map(|(i, doc)| (format!("doc{}", i + 3), doc))
.collect(),
)
.build()
.await?;
Ok(embeddings)
}
这个函数处理：
初始化 OpenAI 嵌入模型。 为我们的真实文档创建嵌入。 添加虚拟数据以满足 LanceDB 的索引要求。
配置向量存储
现在，让我们设置 LanceDB 并配置适当的索引和搜索参数：
async fn setup_vector_store<M: EmbeddingModel>(
embeddings: Vec<DocumentEmbeddings>,
model: M,
) -> Result<LanceDbVectorStore<M>> {
// 初始化 LanceDB
let db = lancedb::connect("data/lancedb-store").execute().await?;
// 如果存在，删除现有表 - 对于开发很重要
if db
.table_names()
.execute()
.await?
.contains(&"documents".to_string())
{
db.drop_table("documents").await?;
}
// 使用嵌入创建表
let record_batch = as_record_batch(embeddings, model.ndims())?;
let table = db
.create_table(
"documents",
RecordBatchIterator::new(vec![Ok(record_batch)], Arc::new(schema(model.ndims()))),
)
.execute()
.await?;
// 使用 IVF-PQ 创建优化的向量索引
table
.create_index(
&["embedding"],
lancedb::index::Index::IvfPq(
IvfPqIndexBuilder::default().distance_type(DistanceType::Cosine),
),
)
.execute()
.await?;
// 配置搜索参数
let search_params = SearchParams::default().distance_type(DistanceType::Cosine);
// 创建并返回向量存储
Ok(LanceDbVectorStore::new(table, model, "id", search_params).await?)
}
这个设置函数：
连接到 LanceDB 数据库。 管理表的创建和删除。 设置向量索引以实现高效的相似性搜索。
整合所有功能
最后，主函数协调整个过程：
#[tokio::main]
asyncfn main() -> Result<()> {
// 初始化 OpenAI 客户端
let openai_api_key = env::var("OPENAI_API_KEY").expect("OPENAI_API_KEY not set");
let openai_client = Client::new(&openai_api_key);
let model = openai_client.embedding_model(TEXT_EMBEDDING_ADA_002);
// 创建嵌入（包括真实和虚拟文档）
let embeddings = create_embeddings(&openai_client).await?;
println!("Created embeddings for {} documents", embeddings.len());
// 设置向量存储
let store = setup_vector_store(embeddings, model).await?;
println!("Vector store initialized successfully");
// 执行语义搜索
let query = "Tell me about safe programming languages";
let results = store.top_n::<SearchResult>(query, 2).await?;
println!("\nSearch Results for: {}\n", query);
for (score, id, result) in results {
println!(
"Score: {:.4}\nID: {}\nContent: {}\n",
score, id, result.content
);
}
Ok(())
}
理解向量搜索方法
向量搜索系统需要在准确性和性能之间找到平衡，特别是当数据集增长时。LanceDB 提供了两种方法来处理这一问题：精确最近邻（ENN）和近似最近邻（ANN）搜索。
ENN 与 ANN
精确最近邻（ENN）：
在所有向量中进行详尽搜索。 保证找到真正的最近邻。 适用于小型数据集。 无最低数据要求。 较慢但更准确。
近似最近邻（ANN）：
使用索引加速搜索（如 IVF-PQ）。 返回近似结果。 适用于较大数据集。 更快但略微不准确。
选择合适的方法
使用 ENN 当：
数据集较小（< 1,000 个向量）。 精确匹配至关重要。 性能不是主要关注点。
使用 ANN 当：
数据集较大。 可以容忍轻微的近似。 需要快速搜索速度。
在我们的教程中，我们使用 ANN 以实现可扩展性。对于较小的数据集，ENN 更为合适。
提示：在开发过程中从 ENN 开始。随着数据和性能需求的增长，过渡到 ANN。查看 ENN 示例。
运行系统
运行项目：
cargo run
预期输出：
Created embeddings for 258 documents
Vector store initialized successfully
Search Results for: Tell me about safe programming languages
Score: 0.3982
ID: doc2-0
Content: Python emphasizes code readability with significant whitespace.
Score: 0.4369
ID: doc1-0
Content: Rust provides zero-cost abstractions and memory safety without garbage collection.
