Search

table = db.create_table(
    "documents",
    [
        {
            "id": 1,
            "text": "A puppy playing in the park",
            "category": "animals",
            "vector": [1.0, 0.0],
        },
        {
            "id": 2,
            "text": "A dog resting at home",
            "category": "animals",
            "vector": [0.9, 0.1],
        },
        {
            "id": 3,
            "text": "A guide to growing tomatoes",
            "category": "gardening",
            "vector": [0.0, 1.0],
        },
    ],
)
rows = (
    table.search([1.0, 0.0])
    .distance_type("l2")
    .select(["id", "text", "_distance"])
    .limit(2)
    .to_list()
)
print(rows)

const lancedb = await import("@lancedb/lancedb");
const { mkdtemp, rm } = await import("node:fs/promises");
const { tmpdir } = await import("node:os");
const { join } = await import("node:path");
const uri = await mkdtemp(join(tmpdir(), "lancedb-search-"));
try {
  const db = await lancedb.connect(uri);
  const table = await db.createTable("documents", [
    {
      id: 1,
      text: "A puppy playing in the park",
      category: "animals",
      vector: [1.0, 0.0],
    },
    {
      id: 2,
      text: "A dog resting at home",
      category: "animals",
      vector: [0.9, 0.1],
    },
    {
      id: 3,
      text: "A guide to growing tomatoes",
      category: "gardening",
      vector: [0.0, 1.0],
    },
  ]);
  const rows = await table
    .vectorSearch([1.0, 0.0])
    .distanceType("l2")
    .select(["id", "text", "_distance"])
    .limit(2)
    .toArray();
  console.log(rows);

} finally {
  await rm(uri, { recursive: true, force: true });
}

use arrow_array::types::Float32Type;
use arrow_array::{FixedSizeListArray, Int64Array, RecordBatch, RecordBatchIterator, StringArray};
use arrow_schema::{DataType, Field, Schema};
use futures::TryStreamExt;
use lancedb::query::{ExecutableQuery, QueryBase, Select};
use std::sync::Arc;

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    let directory = tempfile::tempdir()?;
    let db = lancedb::connect(directory.path().to_str().unwrap())
        .execute()
        .await?;
    let schema = Arc::new(Schema::new(vec![
        Field::new("id", DataType::Int64, false),
        Field::new("text", DataType::Utf8, false),
        Field::new("category", DataType::Utf8, false),
        Field::new(
            "vector",
            DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float32, true)), 2),
            true,
        ),
    ]));
    let batch = RecordBatch::try_new(
        schema.clone(),
        vec![
            Arc::new(Int64Array::from(vec![1, 2, 3])),
            Arc::new(StringArray::from(vec![
                "A puppy playing in the park",
                "A dog resting at home",
                "A guide to growing tomatoes",
            ])),
            Arc::new(StringArray::from(vec!["animals", "animals", "gardening"])),
            Arc::new(
                FixedSizeListArray::from_iter_primitive::<Float32Type, _, _>(
                    [[1.0, 0.0], [0.9, 0.1], [0.0, 1.0]]
                        .into_iter()
                        .map(|v| Some(v.into_iter().map(Some).collect::<Vec<_>>())),
                    2,
                ),
            ),
        ],
    )?;
    let reader = RecordBatchIterator::new(vec![Ok(batch)].into_iter(), schema);
    let table = db
        .create_table(
            "documents",
            Box::new(reader) as Box<dyn arrow_array::RecordBatchReader + Send>,
        )
        .execute()
        .await?;
    let query = table
        .vector_search(&[1.0_f32, 0.0])?
        .distance_type(lancedb::DistanceType::L2)
        .select(Select::columns(&["id", "text", "_distance"]))
        .limit(2);
    let rows: Vec<RecordBatch> = query.execute().await?.try_collect().await?;
    println!("{rows:?}");

    Ok(())
}