速览:语义搜索系统轻松搭建在Rust的世界里,想打造强大的语义搜索系统吗?别愁啦!借助Rig和LanceDB,这一切变得超简单。我们会手把手带你实操,从创建向量嵌入,到存储和搜索,每一步都清晰明了。不管是构建RAG系统,还是打造语义搜索引擎,这套方法都能让你事半功倍。完整源代码已放
在 Rust 的世界里,想打造强大的语义搜索系统吗?别愁啦!借助 Rig 和 LanceDB,这一切变得超简单。我们会手把手带你实操,从创建向量嵌入,到存储和搜索,每一步都清晰明了。不管是构建 RAG 系统,还是打造语义搜索引擎,这套方法都能让你事半功倍。
完整源代码已放在 GitHub 仓库,赶紧去看看吧!
传统的关键词搜索已经 OUT 啦!语义搜索才是当下的潮流。它能精准捕捉查询背后的真实意图,让信息检索更加细致入微。不过,构建这样的系统,听起来就头大,什么复杂的嵌入、向量数据库,还有相似性搜索算法,让人望而却步。
别急,LanceDB 来救场!
LanceDB 可是开源向量数据库界的明星,专为 AI 应用和向量搜索量身定制。它的优势简直爆棚:
再结合 Rig 的嵌入和 LLM 能力,用最少的代码就能打造出强大又高效的语义搜索解决方案,是不是超心动?
在动手之前,你得先确保这些条件都满足:
一切就绪后,就可以开始创建新的 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"这些依赖项都有各自的作用:
rig-core 和 rig-lancedb:生成嵌入和进行向量搜索的核心库。lancedb:嵌入式向量数据库。tokio:提供异步运行时支持。arrow-array:处理 Arrow 的列式格式,这可是 LanceDB 内部用的。最后,创建一个 .env 文件,把你的 OpenAI API 密钥存进去:
echo "OPENAI_API_KEY=your_key_here" > .env我们把这个过程拆分成几个简单的步骤,先创建一个实用函数,处理 Rig 的嵌入和 LanceDB 格式之间的数据转换。
src/utils.rsuse 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表的模式
pub fn 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 as i32,
            ),
            false,
        ),
    ]))
}这个模式函数定义了表的结构:
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 as i32,
    );
    RecordBatch::try_from_iter(vec![
        ("id", Arc::new(id) as ArrayRef),
        ("content", Arc::new(content) as ArrayRef),
        ("embedding", Arc::new(embedding) as ArrayRef),
    ])
}这个函数超重要,它把 Rust 数据结构转换成 Arrow 的列式格式,也就是 LanceDB 内部用的格式:
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};这些导入包括:
创建一个简单的结构体,用来表示搜索结果:
#[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提供了零成本抽象和内存安全,无需垃圾回收。",
        )
        .simple_document(
            "doc2",
            "Python通过显著的空白强调代码可读性。",
        )
        // 使用枚举生成唯一ID的虚拟文档以满足最低要求
        .simple_documents(
            dummy_docs
                .into_iter()
                .enumerate()
                .map(|(i, doc)| (format!("doc{}", i + 3), doc))
                .collect(),
        )
        .build()
        .await?;
    Ok(embeddings)
}这个函数负责:
现在,设置 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?)
}这个设置函数:
最后,主函数把整个过程协调起来:
#[tokio::main]
async fn main() -> Result<()> {
    // 初始化OpenAI客户端
    let openai_api_key = env::var("OPENAI_API_KEY").expect("OPENAI_API_KEY未设置");
    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!("已为{}个文档创建嵌入", embeddings.len());
    // 设置向量存储
    let store = setup_vector_store(embeddings, model).await?;
    println!("向量存储初始化成功");
    // 执行语义搜索
    let query = "告诉我关于安全的编程语言";
    let results = store.top_n::<SearchResult>(query, 2).await?;
    println!("\n搜索结果为:{}\n", query);
    for (score, id, result) in results {
        println!(
            "得分:{:.4}\nID:{}\n内容:{}\n",
            score, id, result.content
        );
    }
    Ok(())
}向量搜索系统得在准确性和性能之间找到平衡,尤其是数据集越来越大的时候。LanceDB 提供了两种方法来应对:精确最近邻(ENN)和近似最近邻(ANN)搜索。
精确最近邻(ENN):
近似最近邻(ANN):
选 ENN 的情况:
选 ANN 的情况:
在这个教程里,我们用 ANN 来保证可扩展性。要是数据集小,ENN 会更合适。
小提示:开发的时候先从 ENN 开始,等数据和性能需求增加了,再换成 ANN。可以看看 ENN 示例。
要运行项目,只需执行:
cargo run预期输出如下:
已为258个文档创建嵌入
向量存储初始化成功
搜索结果为:告诉我关于安全的编程语言
得分:0.3982
ID:doc2-0
内容:Python通过显著的空白强调代码可读性。
得分:0.4369
ID:doc1-0
内容:Rust提供了零成本抽象和内存安全,无需垃圾回收。如果你想用 Rig 构建更多功能,这里有一些实用示例:
想让你的 LLM 访问自定义知识?看看这个教程 使用 Rig 在 100 行代码内构建 RAG 系统。
准备好打造更具交互性的 AI 应用程序?参考 代理示例。
现在,就用 Rig 和 LanceDB 开启你的 Rust 向量搜索应用之旅吧!
 
                如果觉得我的文章对您有用,请随意打赏。你的支持将鼓励我继续创作!