mirror of
https://github.com/lancedb/lancedb.git
synced 2026-10-13 15:22:30 +00:00
Brendan Clement
215f78de96
The TypeScript and Rust Quickstarts read the image from ../../docs/static, the layout before the documentation moved into this repository. The file is docs/web/static/assets/images/quickstart/sir-lancelot.jpg, so the TypeScript test failed with ENOENT and the Rust example panicked at the read. Both now resolve the image from the example's own directory (__dirname and CARGO_MANIFEST_DIR), and the two Quickstart exports are regenerated with docs/web-tests/mdx_snippets_gen.py.
348 lines
11 KiB
Rust
348 lines
11 KiB
Rust
// SPDX-License-Identifier: Apache-2.0
|
|
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
|
|
|
use std::sync::Arc;
|
|
|
|
use arrow_array::types::Float32Type;
|
|
use arrow_array::{
|
|
FixedSizeListArray, Int8Array, LargeStringArray, RecordBatch, RecordBatchIterator, StructArray,
|
|
};
|
|
use arrow_schema::{DataType, Field, FieldRef, Schema};
|
|
use lancedb::arrow::IntoPolars;
|
|
use lancedb::database::CreateTableMode;
|
|
use lancedb::query::{ExecutableQuery, QueryBase, Select};
|
|
use lancedb::{connect, table::NewColumnTransform};
|
|
use polars::prelude::DataFrame;
|
|
use serde::{Deserialize, Serialize};
|
|
|
|
// --8<-- [start:quickstart_define_struct]
|
|
// Define structs representing the data schema
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
struct Stats {
|
|
strength: i8,
|
|
magic: i8,
|
|
leadership: i8,
|
|
wisdom: i8,
|
|
}
|
|
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
struct Character {
|
|
id: String,
|
|
name: String,
|
|
role: String,
|
|
description: String,
|
|
stats: Stats,
|
|
vector: [f32; 4],
|
|
}
|
|
|
|
fn characters_schema() -> Arc<Schema> {
|
|
Arc::new(Schema::new(vec![
|
|
Field::new("id", DataType::LargeUtf8, false),
|
|
Field::new("name", DataType::LargeUtf8, false),
|
|
Field::new("role", DataType::LargeUtf8, false),
|
|
Field::new("description", DataType::LargeUtf8, false),
|
|
Field::new(
|
|
"stats",
|
|
DataType::Struct(arrow_schema::Fields::from(vec![
|
|
Arc::new(Field::new("strength", DataType::Int8, false)),
|
|
Arc::new(Field::new("magic", DataType::Int8, false)),
|
|
Arc::new(Field::new("leadership", DataType::Int8, false)),
|
|
Arc::new(Field::new("wisdom", DataType::Int8, false)),
|
|
])),
|
|
false,
|
|
),
|
|
Field::new(
|
|
"vector",
|
|
DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float32, true)), 4),
|
|
false,
|
|
),
|
|
]))
|
|
}
|
|
// --8<-- [end:quickstart_define_struct]
|
|
|
|
type BatchIter = Box<dyn arrow_array::RecordBatchReader + Send>;
|
|
|
|
fn characters_to_reader(schema: Arc<Schema>, rows: &[Character]) -> BatchIter {
|
|
let ids = LargeStringArray::from_iter_values(rows.iter().map(|row| row.id.as_str()));
|
|
let names = LargeStringArray::from_iter_values(rows.iter().map(|row| row.name.as_str()));
|
|
let roles = LargeStringArray::from_iter_values(rows.iter().map(|row| row.role.as_str()));
|
|
let descriptions =
|
|
LargeStringArray::from_iter_values(rows.iter().map(|row| row.description.as_str()));
|
|
|
|
let strength = Int8Array::from_iter_values(rows.iter().map(|row| row.stats.strength));
|
|
let magic = Int8Array::from_iter_values(rows.iter().map(|row| row.stats.magic));
|
|
let leadership = Int8Array::from_iter_values(rows.iter().map(|row| row.stats.leadership));
|
|
let wisdom = Int8Array::from_iter_values(rows.iter().map(|row| row.stats.wisdom));
|
|
let stats_fields: Vec<FieldRef> = vec![
|
|
Arc::new(Field::new("strength", DataType::Int8, false)),
|
|
Arc::new(Field::new("magic", DataType::Int8, false)),
|
|
Arc::new(Field::new("leadership", DataType::Int8, false)),
|
|
Arc::new(Field::new("wisdom", DataType::Int8, false)),
|
|
];
|
|
let stats = StructArray::new(
|
|
stats_fields.into(),
|
|
vec![
|
|
Arc::new(strength),
|
|
Arc::new(magic),
|
|
Arc::new(leadership),
|
|
Arc::new(wisdom),
|
|
],
|
|
None,
|
|
);
|
|
|
|
let vectors = FixedSizeListArray::from_iter_primitive::<Float32Type, _, _>(
|
|
rows.iter()
|
|
.map(|row| Some(row.vector.iter().copied().map(Some).collect::<Vec<_>>())),
|
|
4,
|
|
);
|
|
|
|
let batch = RecordBatch::try_new(
|
|
schema.clone(),
|
|
vec![
|
|
Arc::new(ids),
|
|
Arc::new(names),
|
|
Arc::new(roles),
|
|
Arc::new(descriptions),
|
|
Arc::new(stats),
|
|
Arc::new(vectors),
|
|
],
|
|
)
|
|
.unwrap();
|
|
|
|
Box::new(RecordBatchIterator::new(
|
|
vec![Ok(batch)].into_iter(),
|
|
schema,
|
|
))
|
|
}
|
|
|
|
#[tokio::main]
|
|
async fn main() {
|
|
let temp_dir = tempfile::tempdir().unwrap();
|
|
let uri = temp_dir.path().to_str().unwrap();
|
|
let db = connect(uri).execute().await.unwrap();
|
|
|
|
// --8<-- [start:quickstart_data]
|
|
let data = vec![
|
|
Character {
|
|
id: "1".to_string(),
|
|
name: "King Arthur".to_string(),
|
|
role: "King".to_string(),
|
|
description: "Leader of Camelot and wielder of Excalibur.".to_string(),
|
|
stats: Stats {
|
|
strength: 4,
|
|
magic: 1,
|
|
leadership: 5,
|
|
wisdom: 4,
|
|
},
|
|
vector: [0.7, 0.1, 0.9, 0.7],
|
|
},
|
|
Character {
|
|
id: "2".to_string(),
|
|
name: "Merlin".to_string(),
|
|
role: "Wizard".to_string(),
|
|
description: "Advisor and prophet with deep magical knowledge.".to_string(),
|
|
stats: Stats {
|
|
strength: 2,
|
|
magic: 5,
|
|
leadership: 4,
|
|
wisdom: 5,
|
|
},
|
|
vector: [0.2, 0.9, 0.4, 0.9],
|
|
},
|
|
Character {
|
|
id: "3".to_string(),
|
|
name: "Sir Lancelot".to_string(),
|
|
role: "Knight".to_string(),
|
|
description: "Legendary knight known for courage and combat skill.".to_string(),
|
|
stats: Stats {
|
|
strength: 5,
|
|
magic: 1,
|
|
leadership: 3,
|
|
wisdom: 3,
|
|
},
|
|
vector: [0.9, 0.1, 0.5, 0.4],
|
|
},
|
|
];
|
|
// --8<-- [end:quickstart_data]
|
|
|
|
// --8<-- [start:quickstart_create_table]
|
|
let schema = characters_schema();
|
|
let table = db
|
|
.create_table("characters", characters_to_reader(schema.clone(), &data))
|
|
.mode(CreateTableMode::Overwrite)
|
|
.execute()
|
|
.await
|
|
.unwrap();
|
|
// --8<-- [end:quickstart_create_table]
|
|
assert_eq!(table.count_rows(None).await.unwrap(), 3);
|
|
db.drop_table("characters", &[]).await.unwrap();
|
|
|
|
// --8<-- [start:quickstart_create_table_no_overwrite]
|
|
let table = db
|
|
.create_table("characters", characters_to_reader(schema.clone(), &data))
|
|
.execute()
|
|
.await
|
|
.unwrap();
|
|
// --8<-- [end:quickstart_create_table_no_overwrite]
|
|
assert_eq!(table.count_rows(None).await.unwrap(), 3);
|
|
|
|
// --8<-- [start:quickstart_vector_search_1]
|
|
// Search for examples similar to a "wise magical advisor"
|
|
let query_vector = [0.2, 0.8, 0.4, 0.9];
|
|
|
|
let result: DataFrame = table
|
|
.query()
|
|
.nearest_to(&query_vector)
|
|
.unwrap()
|
|
.select(Select::Columns(vec![
|
|
"name".to_string(),
|
|
"role".to_string(),
|
|
"description".to_string(),
|
|
"_distance".to_string(),
|
|
]))
|
|
.limit(2)
|
|
.execute()
|
|
.await
|
|
.unwrap()
|
|
.into_polars()
|
|
.await
|
|
.unwrap();
|
|
println!("{result:?}");
|
|
// --8<-- [end:quickstart_vector_search_1]
|
|
let name_col = result.column("name").unwrap().str().unwrap();
|
|
assert_eq!(name_col.get(0).unwrap(), "Merlin");
|
|
|
|
// --8<-- [start:quickstart_curate_with_metadata]
|
|
let curated: DataFrame = table
|
|
.query()
|
|
.nearest_to(&query_vector)
|
|
.unwrap()
|
|
.only_if("stats.magic >= 4")
|
|
.select(Select::Columns(vec![
|
|
"name".to_string(),
|
|
"role".to_string(),
|
|
"description".to_string(),
|
|
"_distance".to_string(),
|
|
]))
|
|
.limit(2)
|
|
.execute()
|
|
.await
|
|
.unwrap()
|
|
.into_polars()
|
|
.await
|
|
.unwrap();
|
|
println!("{curated:?}");
|
|
// --8<-- [end:quickstart_curate_with_metadata]
|
|
let curated_name_col = curated.column("name").unwrap().str().unwrap();
|
|
assert_eq!(curated_name_col.get(0).unwrap(), "Merlin");
|
|
|
|
// --8<-- [start:quickstart_output_array]
|
|
let result: DataFrame = table
|
|
.query()
|
|
.nearest_to(&query_vector)
|
|
.unwrap()
|
|
.select(Select::Columns(vec![
|
|
"name".to_string(),
|
|
"role".to_string(),
|
|
"description".to_string(),
|
|
"_distance".to_string(),
|
|
]))
|
|
.limit(2)
|
|
.execute()
|
|
.await
|
|
.unwrap()
|
|
.into_polars()
|
|
.await
|
|
.unwrap();
|
|
println!("{result:?}");
|
|
// --8<-- [end:quickstart_output_array]
|
|
let name_col = result.column("name").unwrap().str().unwrap();
|
|
assert_eq!(name_col.get(0).unwrap(), "Merlin");
|
|
|
|
// --8<-- [start:quickstart_add_feature]
|
|
table
|
|
.add_columns(
|
|
NewColumnTransform::SqlExpressions(vec![(
|
|
"power_score".to_string(),
|
|
"cast(((stats.strength + stats.magic + stats.leadership + stats.wisdom) / 4.0) as float)"
|
|
.to_string(),
|
|
)]),
|
|
None,
|
|
)
|
|
.await
|
|
.unwrap();
|
|
// --8<-- [end:quickstart_add_feature]
|
|
|
|
// --8<-- [start:quickstart_query_feature]
|
|
let features: DataFrame = table
|
|
.query()
|
|
.select(Select::Columns(vec![
|
|
"name".to_string(),
|
|
"role".to_string(),
|
|
"power_score".to_string(),
|
|
]))
|
|
.execute()
|
|
.await
|
|
.unwrap()
|
|
.into_polars()
|
|
.await
|
|
.unwrap();
|
|
println!("{features:?}");
|
|
// --8<-- [end:quickstart_query_feature]
|
|
assert!(features.column("power_score").is_ok());
|
|
|
|
// --8<-- [start:quickstart_multimodal_bytes]
|
|
use std::sync::Arc;
|
|
|
|
use arrow_array::{
|
|
BinaryArray, FixedSizeListArray, LargeStringArray, RecordBatch, RecordBatchIterator,
|
|
};
|
|
use arrow_schema::{DataType, Field, Schema};
|
|
|
|
let image_path = std::path::Path::new(env!("CARGO_MANIFEST_DIR"))
|
|
.join("../../web/static/assets/images/quickstart/sir-lancelot.jpg");
|
|
let image_bytes = std::fs::read(image_path).unwrap();
|
|
|
|
let image_schema = Arc::new(Schema::new(vec![
|
|
Field::new("id", DataType::LargeUtf8, false),
|
|
Field::new("description", DataType::LargeUtf8, false),
|
|
Field::new("image", DataType::Binary, false),
|
|
Field::new(
|
|
"vector",
|
|
DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float32, true)), 4),
|
|
false,
|
|
),
|
|
]));
|
|
let image_vectors = [[0.9_f32, 0.1, 0.5, 0.4]];
|
|
let image_batch = RecordBatch::try_new(
|
|
image_schema.clone(),
|
|
vec![
|
|
Arc::new(LargeStringArray::from_iter_values(["lancelot"])),
|
|
Arc::new(LargeStringArray::from_iter_values([
|
|
"Portrait of Sir Lancelot",
|
|
])),
|
|
Arc::new(BinaryArray::from_iter_values([image_bytes.as_slice()])),
|
|
Arc::new(
|
|
FixedSizeListArray::from_iter_primitive::<Float32Type, _, _>(
|
|
image_vectors
|
|
.iter()
|
|
.map(|vector| Some(vector.iter().copied().map(Some).collect::<Vec<_>>())),
|
|
4,
|
|
),
|
|
),
|
|
],
|
|
)
|
|
.unwrap();
|
|
let image_reader: Box<dyn arrow_array::RecordBatchReader + Send> = Box::new(
|
|
RecordBatchIterator::new(vec![Ok(image_batch)].into_iter(), image_schema),
|
|
);
|
|
let multimodal_table = db
|
|
.create_table("character_images", image_reader)
|
|
.mode(CreateTableMode::Overwrite)
|
|
.execute()
|
|
.await
|
|
.unwrap();
|
|
// --8<-- [end:quickstart_multimodal_bytes]
|
|
assert_eq!(multimodal_table.count_rows(None).await.unwrap(), 1);
|
|
}
|