https://blog.mocobeta.dev/posts/feed.xml

LanceDBのRust APIで全文検索，ベクトル検索，ハイブリッド検索をする

2025-08-09

LanceDBのRust APIを使って，全文検索，ベクトル検索，ハイブリッド検索を試します。

LanceDBのPythonSDKで全文検索，ベクトル検索，ハイブリッド検索をするのRust版です。

公式ドキュメントにRustのサンプルは少ないですが，GitHubリポジトリのexamplesがあります。

準備

Cargo.tomlに依存関係を追加します。arrow-arrayのバージョンは，LanceDBが依存しているバージョンに合わせます。また，embeddingにSentence Transformersを使用するために，lancedbのsentence-transformers featureを有効にします。

[dependencies]
arrow-array = "55.2.0"
arrow-schema = "55.2.0"
futures = "0.3.31"
lance-index = "0.32.0"
lancedb = { version = "0.21.2", features = ["sentence-transformers"] }
tokio = { version = "1", features = ["full"] }

テーブルを作成して，初期データを投入します。テーブルスキーマは，Apache Arrow Schemaで定義します。

#[tokio::main]
async fn main() -> Result<()> {
    // データベースに接続（または新規作成）
    let uri = "data/testdb";
    let db = connect(uri).execute().await?;

    // sentence-transformersモデルをデータベースに登録
    let embedding = Arc::new(SentenceTransformersEmbeddings::builder().build()?);
    db.embedding_registry()
        .register("sentence-transformers", embedding.clone())?;

    // テーブルスキーマ定義
    let schema = Arc::new(Schema::new(vec![Field::new(
        "movie",
        DataType::Utf8,
        false,
    )]));

    // 初期データ
    let data = StringArray::from_iter_values(vec![
        "The Godfather is Francis Ford Coppola's epic crime drama often cited as one of the greatest films ever made.",
        "The Shawshank Redemption is a powerful and uplifting prison story that consistently tops audience polls.",
        "Citizen Kane, directed by Orson Welles, is celebrated for its revolutionary cinematic techniques.",
        "Pulp Fiction is Quentin Tarantino's non-linear crime film and a landmark of independent cinema.",
        "2001: A Space Odyssey is Stanley Kubrick's visually stunning and philosophical science fiction epic.",
        "The Dark Knight, directed by Christopher Nolan, is hailed for its dark themes and iconic villain.",
        "Schindler's List is Steven Spielberg's poignant and powerful historical drama about the Holocaust.",
        "Seven Samurai is Akira Kurosawa's influential epic that set the template for the modern action film.",
        "Casablanca is a timeless romantic drama known for its iconic lines and performances.",
        "Parasite is Bong Joon-ho's satirical thriller that made history at the Academy Awards.",
    ]);
    let batch = RecordBatchIterator::new(
        RecordBatch::try_new(schema.clone(), vec![Arc::new(data)])
            .into_iter()
            .map(Ok),
        schema.clone(),
    );

    // "movie"カラムに対してembeddingを定義。embeddingの格納先は"embeddings"カラム
    let embdef = EmbeddingDefinition::new("movie", "sentence-transformers", Some("embeddings"));
    // テーブル作成
    let table = db
        .create_table("best_movies", batch)
        .add_embedding(embdef)?
        .execute()
        .await?;

    Ok(())
}

全文検索用のインデックス作成

Table.create_indexメソッドで，movieカラムに対して全文検索インデックスを作成します。index引数にIndex::FTSを指定すると全文検索インデックスが作成されます。

    table
        .create_index(&["movie"], Index::FTS(FtsIndexBuilder::default()))
        .execute()
        .await?;

(オプショナル) ANNインデックス作成

Table.create_indexメソッドで，embeddingsカラムに対してANNインデックスを作成します。index引数にIndex::IvfFlatを指定するとInverted File Flatインデックスが作成されます。ここでは，ドキュメント数が少ないのでコメントアウト。

//    table
//        .create_index(
//            &["embeddings"],
//            Index::IvfFlat(IvfFlatIndexBuilder::default()),
//        )
//        .execute()
//        .await?;

検索する

Table.queryメソッドを使用して検索を行います。結果は，ArrowのRecordBatchで返されます。

全文検索

Query.full_text_searchメソッドに全文検索クエリFullTextSearchQueryを指定します。

    println!("Full-text search results:");
    let fts_query = FullTextSearchQuery::new("history".to_string());
    let results = table
        .query()
        .full_text_search(fts_query.clone())
        .limit(3)
        .execute()
        .await?
        .try_collect::<Vec<_>>()
        .await?;
    for batch in results {
        let column = batch
            .column_by_name("movie")
            .unwrap()
            .as_any()
            .downcast_ref::<StringArray>()
            .unwrap();
        for v in column.iter() {
            println!("{}", v.unwrap());
        }
    }

Full-text search results:
Parasite is Bong Joon-ho's satirical thriller that made history at the Academy Awards.

"history"という検索キーワードを含む映画だけがヒットする。

ベクトル検索

Query.nearest_toメソッドにクエリembeddingを指定するとベクトル検索になります。Python版の場合，クエリembeddingは，カラムに紐づけたembedding functionで暗黙的に計算されますが，Rust版だと自分で計算する必要があります。

    println!("\nVector search results:");
    // クエリembeddingの計算
    let vector_query = embedding
        .compute_query_embeddings(Arc::new(StringArray::from_iter_values(once("history"))))?;

    let results = table
        .query()
        .nearest_to(vector_query.clone())?
        .limit(3)
        .execute()
        .await?
        .try_collect::<Vec<_>>()
        .await?;
    for batch in results {
        let column = batch
            .column_by_name("movie")
            .unwrap()
            .as_any()
            .downcast_ref::<StringArray>()
            .unwrap();
        for v in column.iter() {
            println!("{}", v.unwrap());
        }
    }

Vector search results:
Casablanca is a timeless romantic drama known for its iconic lines and performances.
The Dark Knight, directed by Christopher Nolan, is hailed for its dark themes and iconic villain.
Parasite is Bong Joon-ho's satirical thriller that made history at the Academy Awards.

ふんわり，"history"に関連しそうな映画がヒットする。なぜか，Python版と検索結果が違う・・・。

ハイブリッド検索

ハイブリッド検索のサンプルはドキュメントにもAPIリファレンスにも見つからなかったので，ソースコードを見ながら実装してみました。

Query.nearest_toメソッドとQuery.full_text_searchメソッドでそれぞれクエリembeddingと全文検索クエリを渡して，executeの代わりにexecute_hybridメソッドを呼び出すことでハイブリッド検索ができるようです。APIが少し落ち着かない感じ。

    println!("\nHybrid search results:");
    let results = table
        .query()
        .nearest_to(vector_query)?
        .full_text_search(fts_query)
        .limit(3)
        .execute_hybrid(QueryExecutionOptions::default())
        .await?
        .try_collect::<Vec<_>>()
        .await?;
    for batch in results {
        let column = batch
            .column_by_name("movie")
            .unwrap()
            .as_any()
            .downcast_ref::<StringArray>()
            .unwrap();
        for v in column.iter() {
            println!("{}", v.unwrap());
        }
    }

Hybrid search results:
Parasite is Bong Joon-ho's satirical thriller that made history at the Academy Awards.
Casablanca is a timeless romantic drama known for its iconic lines and performances.
The Dark Knight, directed by Christopher Nolan, is hailed for its dark themes and iconic villain.

返ってくる結果セットはベクトル検索と一緒だけど，"history"という検索キーワードを含む映画が1位にヒットしている。