diff --git a/seaweed-vector/Cargo.lock b/seaweed-vector/Cargo.lock index c01416032f..012645c40c 100644 --- a/seaweed-vector/Cargo.lock +++ b/seaweed-vector/Cargo.lock @@ -7256,6 +7256,7 @@ dependencies = [ "tonic-prost-build", "tracing", "tracing-subscriber", + "uuid", "weed-lance-worker", "weed-vector-core", ] diff --git a/seaweed-vector/crates/core/src/main.rs b/seaweed-vector/crates/core/src/main.rs index 687320d5ce..da92d2bbde 100644 --- a/seaweed-vector/crates/core/src/main.rs +++ b/seaweed-vector/crates/core/src/main.rs @@ -170,7 +170,7 @@ async fn open_dataset( storage_options: Default::default(), managed_versioning: false, }; - return dataset::open_at(&description, storage.wrapper(id).await?).await; + return dataset::open_at(&description, storage.wrapper(id).await?, None).await; } let table = dataset::open(client, &parse_id(id), storage).await?; Ok(table.dataset) diff --git a/seaweed-vector/crates/core/src/maintenance.rs b/seaweed-vector/crates/core/src/maintenance.rs index 048015c7f2..546cc62756 100644 --- a/seaweed-vector/crates/core/src/maintenance.rs +++ b/seaweed-vector/crates/core/src/maintenance.rs @@ -2,8 +2,8 @@ //! and version tags act as the lease maintenance must respect. use std::collections::{HashMap, HashSet}; -use std::sync::{LazyLock, Mutex}; -use std::time::{Duration, SystemTime, UNIX_EPOCH}; +use std::sync::{Arc, LazyLock, Mutex, Weak}; +use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH}; use anyhow::Result; use lance::Dataset; @@ -21,6 +21,11 @@ const FOREIGN_TTL: Duration = Duration::from_secs(24 * 3600); /// Expired-lease scans run at most once per interval per dataset, not per /// request. const REAP_INTERVAL: Duration = Duration::from_secs(60); +/// How long a shared lease stays after the last read that took it. +const SHARED_LINGER: Duration = Duration::from_secs(10); +/// A shared lease is handed out for this long after it was taken, then +/// replaced, so a busy version's tag never ages into a reaper's TTL. +const SHARED_FOR: Duration = Duration::from_secs(300); /// A process-unique owner id embedded in lease tags so a reaper only applies /// the short TTL to its own leases. static INSTANCE: LazyLock = LazyLock::new(|| Uuid::new_v4().simple().to_string()); @@ -138,6 +143,104 @@ fn untrack(tag: &str) { } } +/// Leases that reads of one version share: the first read tags the version, +/// reads that follow reuse the tag, and the tag goes once no read has taken it +/// for [`SHARED_LINGER`]. Spares a tag write and delete per read. +pub struct SharedLeases { + /// By location, version and manifest ETag: a table dropped and created + /// again at its location reuses version numbers, not manifests. + held: Mutex, + linger: Duration, + share_for: Duration, +} + +type LeaseKey = (String, u64, Option); + +#[derive(Default)] +struct Held { + slots: HashMap, + /// Slots are pruned once they reach this many, so a lookup does not + /// scan them all. + prune_at: usize, +} +/// Reads of a version that find no lease wait here for the first to take it. +type Slot = Arc>>; + +impl Default for SharedLeases { + fn default() -> Self { + Self { + held: Mutex::default(), + linger: SHARED_LINGER, + share_for: SHARED_FOR, + } + } +} + +/// A lease held for as long as any read, or the linger after the last one, +/// keeps it. +pub struct SharedLease { + _lease: Lease, + taken: Instant, +} + +impl SharedLeases { + /// [`try_pin`], reusing a lease another read of this version holds. + pub async fn try_pin(&self, dataset: &Dataset) -> Result>> { + let key = ( + dataset.uri().to_string(), + dataset.version().version, + dataset.manifest_location().e_tag.clone(), + ); + let slot = match self.held.lock() { + Ok(mut held) => { + if held.slots.len() >= held.prune_at { + // A slot another read has taken but not yet locked stays. + held.slots.retain(|_, slot| { + Arc::strong_count(slot) > 1 + || slot + .try_lock() + .map_or(true, |lease| lease.strong_count() > 0) + }); + held.prune_at = (2 * held.slots.len()).max(64); + // A burst's allocation would otherwise make every later + // prune scan it. + let capacity = held.prune_at; + held.slots.shrink_to(capacity); + } + held.slots.entry(key).or_default().clone() + } + Err(_) => Slot::default(), + }; + let mut slot = slot.lock().await; + let lease = match slot + .upgrade() + .filter(|lease| lease.taken.elapsed() < self.share_for) + { + Some(lease) => lease, + None => { + let Some(lease) = try_pin(dataset).await? else { + return Ok(None); + }; + let lease = Arc::new(SharedLease { + _lease: lease, + taken: Instant::now(), + }); + *slot = Arc::downgrade(&lease); + lease + } + }; + drop(slot); + if let Ok(handle) = tokio::runtime::Handle::try_current() { + let (lingering, linger) = (lease.clone(), self.linger); + handle.spawn(async move { + tokio::time::sleep(linger).await; + drop(lingering); + }); + } + Ok(Some(lease)) + } +} + /// Drop lease tags whose owner never released them, so an abandoned lease /// cannot keep a version out of cleanup forever. async fn reap_expired(dataset: &Dataset, now: u128) { @@ -324,6 +427,131 @@ mod tests { b.release().await.unwrap(); } + /// Reads of one version share a tag, another version gets its own, and + /// the tags go once the linger after the last read has passed. + #[tokio::test] + async fn reads_of_a_version_share_one_lease() { + let dir = TempDir::new().unwrap(); + let mut dataset = + testing::write_dataset(&dir, &testing::vectors(4, 4), &["a", "b", "c", "d"]).await; + let leases = SharedLeases { + linger: Duration::from_millis(50), + ..Default::default() + }; + let a = leases.try_pin(&dataset).await.unwrap().unwrap(); + let b = leases.try_pin(&dataset).await.unwrap().unwrap(); + assert!(Arc::ptr_eq(&a, &b)); + assert_eq!(dataset.tags().list().await.unwrap().len(), 1); + + testing::append_dataset(&mut dataset, &testing::vectors(1, 4), &["e"]).await; + let c = leases.try_pin(&dataset).await.unwrap().unwrap(); + assert!(!Arc::ptr_eq(&a, &c)); + assert_eq!(dataset.tags().list().await.unwrap().len(), 2); + + drop((a, b, c)); + tokio::time::sleep(leases.linger).await; + let mut left = usize::MAX; + for _ in 0..1000 { + // A listing that meets a tag as it is deleted fails; try again. + if let Ok(tags) = dataset.tags().list().await { + left = tags.len(); + } + if left == 0 { + break; + } + tokio::time::sleep(Duration::from_millis(10)).await; + } + assert_eq!(left, 0); + } + + /// Reads that arrive together, before any lease exists, share the one the + /// first takes. + #[tokio::test] + async fn concurrent_reads_take_one_lease() { + let dir = TempDir::new().unwrap(); + let dataset = + testing::write_dataset(&dir, &testing::vectors(4, 4), &["a", "b", "c", "d"]).await; + let leases = SharedLeases::default(); + let pins: Vec<_> = futures::future::join_all((0..8).map(|_| leases.try_pin(&dataset))) + .await + .into_iter() + .map(|lease| lease.unwrap().unwrap()) + .collect(); + assert!(pins.iter().all(|lease| Arc::ptr_eq(lease, &pins[0]))); + assert_eq!(dataset.tags().list().await.unwrap().len(), 1); + } + + /// A table dropped and written again at its location gets its own lease + /// for a version number the old table had. + #[tokio::test] + async fn a_recreated_table_takes_its_own_lease() { + let dir = TempDir::new().unwrap(); + let rows = testing::vectors(4, 4); + let old = testing::write_dataset(&dir, &rows, &["a", "b", "c", "d"]).await; + let leases = SharedLeases::default(); + let first = leases.try_pin(&old).await.unwrap().unwrap(); + std::fs::remove_dir_all(dir.path()).unwrap(); + let new = testing::write_dataset(&dir, &rows, &["e", "f", "g", "h"]).await; + assert_eq!(new.version().version, old.version().version); + let second = leases.try_pin(&new).await.unwrap().unwrap(); + assert!(!Arc::ptr_eq(&first, &second)); + assert_eq!(new.tags().list().await.unwrap().len(), 1); + } + + /// Slots of released leases are pruned as new versions come in. + #[tokio::test] + async fn released_lease_slots_are_pruned() { + let dir = TempDir::new().unwrap(); + let mut dataset = testing::write_dataset(&dir, &testing::vectors(1, 4), &["a"]).await; + let leases = SharedLeases { + linger: Duration::ZERO, + ..Default::default() + }; + for _ in 0..200 { + drop(leases.try_pin(&dataset).await.unwrap()); + tokio::task::yield_now().await; + testing::append_dataset(&mut dataset, &testing::vectors(1, 4), &["b"]).await; + } + assert!(leases.held.lock().unwrap().slots.len() <= 128); + } + + /// A slot a read has taken but not yet locked survives a prune, so that + /// read and the next share the lease. + #[tokio::test] + async fn a_taken_slot_survives_a_prune() { + let dir = TempDir::new().unwrap(); + let dataset = testing::write_dataset(&dir, &testing::vectors(1, 4), &["a"]).await; + let leases = SharedLeases::default(); + let key: LeaseKey = ("elsewhere".into(), 1, None); + let taken = { + let mut held = leases.held.lock().unwrap(); + held.prune_at = 0; + held.slots.entry(key.clone()).or_default().clone() + }; + leases.try_pin(&dataset).await.unwrap().unwrap(); + assert!(Arc::ptr_eq( + &leases.held.lock().unwrap().slots[&key], + &taken + )); + } + + /// A shared lease is not handed out past SHARED_FOR, so a busy version + /// moves to a fresh tag. + #[tokio::test] + async fn a_shared_lease_is_replaced_when_old() { + let dir = TempDir::new().unwrap(); + let dataset = + testing::write_dataset(&dir, &testing::vectors(4, 4), &["a", "b", "c", "d"]).await; + let leases = SharedLeases { + share_for: Duration::from_millis(50), + ..Default::default() + }; + let a = leases.try_pin(&dataset).await.unwrap().unwrap(); + tokio::time::sleep(leases.share_for).await; + let b = leases.try_pin(&dataset).await.unwrap().unwrap(); + assert!(!Arc::ptr_eq(&a, &b)); + } + /// A stale lease tag left behind is reaped by the next pin; a stale-named /// tag this process still holds is not. #[tokio::test] diff --git a/seaweed-vector/crates/server/Cargo.toml b/seaweed-vector/crates/server/Cargo.toml index cbcf163291..c92e020168 100644 --- a/seaweed-vector/crates/server/Cargo.toml +++ b/seaweed-vector/crates/server/Cargo.toml @@ -36,6 +36,7 @@ bytes = "1" clap = { version = "4", features = ["derive", "env"] } futures = "0.3" moka = { version = "0.12", features = ["sync"] } +uuid = { version = "1", features = ["v4"] } object_store = { version = "0.13", features = ["aws"] } serde = { version = "1", features = ["derive"] } serde_json = "1" diff --git a/seaweed-vector/crates/server/src/main.rs b/seaweed-vector/crates/server/src/main.rs index f181220f20..fe0e482902 100644 --- a/seaweed-vector/crates/server/src/main.rs +++ b/seaweed-vector/crates/server/src/main.rs @@ -41,6 +41,13 @@ struct Args { /// Memory for caching reads of write-once table files, in MiB; 0 disables. #[arg(long, env = "WEED_VECTOR_READ_CACHE_MB", default_value_t = 512)] read_cache_mb: u64, + /// Memory for indexes loaded by queries, kept for the next query, in MiB. + #[arg(long, env = "WEED_VECTOR_INDEX_CACHE_MB", default_value_t = 1024)] + index_cache_mb: usize, + /// Memory for table metadata (manifests, file footers) kept across + /// queries, in MiB. + #[arg(long, env = "WEED_VECTOR_METADATA_CACHE_MB", default_value_t = 256)] + metadata_cache_mb: usize, /// Partitions each query probes in a table's vector index. Default: a /// tenth of the index's partitions, at least 8; more under a filter. #[arg(long, env = "WEED_VECTOR_NPROBES")] @@ -163,6 +170,7 @@ async fn main() -> Result<()> { .then(|| ReadCache::new(args.read_cache_mb << 20).with_metrics(metrics.clone())); let mut service = VectorServer::new() .with_nprobes(args.nprobes.map(std::num::NonZeroUsize::get)) + .with_lance_caches(args.index_cache_mb << 20, args.metadata_cache_mb << 20) .with_storage(reads); if let Some(cache) = &cache { service = service.with_read_cache(cache.clone()); diff --git a/seaweed-vector/crates/server/src/read_cache.rs b/seaweed-vector/crates/server/src/read_cache.rs index 71e881714f..563760b823 100644 --- a/seaweed-vector/crates/server/src/read_cache.rs +++ b/seaweed-vector/crates/server/src/read_cache.rs @@ -353,7 +353,7 @@ mod tests { use super::*; use crate::pb::vector as pb; use crate::pb::vector::vector_service_server::VectorService; - use crate::service::VectorServer; + use crate::service::{VectorServer, schema_for, table_id}; use object_store::ObjectStoreExt; use object_store::memory::InMemory; use tonic::Request; @@ -616,6 +616,117 @@ mod tests { assert_eq!(query(&server, &table, vec![1.0, 0.0]).await, ["z"]); } + /// A query loads the index once; the next one finds it in the session + /// every read shares. + #[tokio::test] + async fn queries_reuse_the_loaded_index() { + let bucket = format!("session-{}", uuid_like()); + let location = format!("shared-memory://{bucket}/vector/x"); + let server = VectorServer::new().with_schemes(&["shared-memory"]); + let table = pb::Table { + id: vec![bucket, "vector".into(), "x".into()], + location: location.clone(), + ..Default::default() + }; + create(&server, &table, 2).await; + put( + &server, + &table, + (0..300) + .map(|i| vector(&format!("k{i}"), vec![i as f32, 1.0])) + .collect(), + ) + .await; + let mut dataset = lance::Dataset::open(&location).await.unwrap(); + weed_lance_worker::vector_index::build(&mut dataset, "vector", "idx", 2, false) + .await + .unwrap(); + + assert_eq!(query(&server, &table, vec![5.0, 1.0]).await[0], "k5"); + let session = server.reads.lock().unwrap().session.clone(); + let loaded = session.index_cache_stats().await; + assert!(session.approx_num_items() > 0); + assert_eq!(query(&server, &table, vec![7.0, 1.0]).await[0], "k7"); + let again = session.index_cache_stats().await; + assert!(server.handles.contains_key(&(location, false))); + assert!(again.hits > loaded.hits, "{loaded:?} then {again:?}"); + assert_eq!(again.misses, loaded.misses, "{loaded:?} then {again:?}"); + } + + /// An index dropped and created again at the same location reaches the + /// version numbers of the old one; the shared session must not answer + /// for them with the old index. + #[tokio::test] + async fn a_recreated_table_does_not_see_the_old_index() { + recreate_after(false).await; + recreate_after(true).await; + } + + /// `unmarked`: the first table is written without a table id, as indexes + /// created before ids were are. + async fn recreate_after(unmarked: bool) { + let bucket = format!("recreate-{}", uuid_like()); + let location = format!("shared-memory://{bucket}/vector/x"); + let server = VectorServer::new().with_schemes(&["shared-memory"]); + let table = pb::Table { + id: vec![bucket, "vector".into(), "x".into()], + location: location.clone(), + ..Default::default() + }; + let rows = |offset: usize| { + (0..300) + .map(|i| vector(&format!("k{}", i + offset), vec![(i + offset) as f32, 1.0])) + .collect::>() + }; + if unmarked { + let schema = schema_for(2, pb::DistanceMetric::L2, &[]); + let empty = arrow_array::RecordBatch::new_empty(schema.clone()); + let reader = arrow_array::RecordBatchIterator::new(vec![Ok(empty)], schema); + lance::Dataset::write(reader, &location, None) + .await + .unwrap(); + } else { + create(&server, &table, 2).await; + } + put(&server, &table, rows(0)).await; + let mut dataset = lance::Dataset::open(&location).await.unwrap(); + assert_eq!(table_id(&dataset).is_empty(), unmarked); + weed_lance_worker::vector_index::build(&mut dataset, "vector", "idx", 2, false) + .await + .unwrap(); + let indexed = dataset.version().version; + assert_eq!(query(&server, &table, vec![5.0, 1.0]).await[0], "k5"); + let other = pb::Table { + id: vec![table.id[0].clone(), "vector".into(), "y".into()], + location: location.replace("/vector/x", "/vector/y"), + ..Default::default() + }; + create(&server, &other, 2).await; + put(&server, &other, rows(0)).await; + assert_eq!(query(&server, &other, vec![5.0, 1.0]).await[0], "k5"); + + let (store, path) = lance::io::ObjectStore::from_uri(&location).await.unwrap(); + store.remove_dir_all(path).await.unwrap(); + create(&server, &table, 2).await; + let mut offset = 1000; + while lance::Dataset::open(&location) + .await + .unwrap() + .version() + .version + < indexed + { + put(&server, &table, rows(offset)).await; + offset += 1000; + } + assert_eq!(query(&server, &table, vec![1005.0, 1.0]).await[0], "k1005"); + server.handles.run_pending_tasks(); + let session = server.reads.lock().unwrap().session.clone(); + for (_, (_, dataset)) in server.handles.iter() { + assert!(Arc::ptr_eq(&dataset.session(), &session)); + } + } + fn uuid_like() -> String { format!( "{:x}", diff --git a/seaweed-vector/crates/server/src/service.rs b/seaweed-vector/crates/server/src/service.rs index df7c9ae0e9..e35e830620 100644 --- a/seaweed-vector/crates/server/src/service.rs +++ b/seaweed-vector/crates/server/src/service.rs @@ -9,16 +9,18 @@ use lance::dataset::ProjectionRequest; use lance::dataset::WriteParams; use lance::dataset::{MergeInsertBuilder, WhenMatched, WhenNotMatched}; use lance::io::WrappingObjectStore; +use lance::session::Session; use lance::table::format::Fragment; use prost_types::{ListValue, Struct, Value, value::Kind}; use std::collections::{BTreeMap, HashMap}; -use std::sync::Arc; +use std::sync::{Arc, Mutex}; use tonic::{Request, Response, Status}; use weed_lance_worker::catalog::TableDescription; use weed_lance_worker::dataset; use weed_lance_worker::direct_read::{DirectReads, Wrappers}; use weed_lance_worker::jobs::vector_index::is_vector_table; use weed_vector_core::index; +use weed_vector_core::maintenance::SharedLeases; use weed_vector_core::query::{self, Query}; use weed_vector_core::table::VectorTable; @@ -55,7 +57,9 @@ const KEY_COLUMN: &str = "key"; const METADATA_COLUMN: &str = "metadata"; const VECTOR_COLUMN: &str = "vector"; const ROW_ID_COLUMN: &str = "_rowid"; -fn schema_for( +const DEFAULT_INDEX_CACHE_BYTES: usize = 1 << 30; +const DEFAULT_METADATA_CACHE_BYTES: usize = 256 << 20; +pub(crate) fn schema_for( dimension: usize, metric: pb::DistanceMetric, non_filterable: &[String], @@ -444,9 +448,85 @@ fn resolved(table: &pb::Table, schemes: &[&str]) -> Option { }) } +/// Marks the table an index was created as, in its schema metadata, which +/// every later version keeps. +const TABLE_ID_KEY: &str = "weed_vector_table_id"; + +/// The table a dataset was created as; empty for one created before tables +/// were marked. +pub(crate) fn table_id(dataset: &Dataset) -> &str { + dataset + .schema() + .metadata + .get(TABLE_ID_KEY) + .map_or("", String::as_str) +} + +/// The session reads share, and the table each location held when this +/// session first read it. Lance keys some of its caches by location and +/// version alone, and an index dropped and created again at its location +/// numbers its versions from 1 again, so a session that sees a location hold +/// another table starts over. +pub(crate) struct Reads { + pub(crate) session: Arc, + generation: u64, + tables: HashMap, +} + +/// Opened tables kept for the next read of them. +const MAX_HANDLES: u64 = 1024; + +fn handle_key(table: &pb::Table) -> (String, bool) { + (table.location.clone(), table.sse_s3) +} + +/// Locations a session remembers before it starts over. +const MAX_READ_TABLES: usize = 100_000; + +impl Reads { + fn new(index_bytes: usize, metadata_bytes: usize) -> Self { + Self { + session: Arc::new(Session::new( + index_bytes, + metadata_bytes, + Default::default(), + )), + generation: 0, + tables: HashMap::new(), + } + } + + /// Whether a dataset of `table_id` at `location`, opened in this session, + /// may be read. Otherwise the session was replaced and the caller opens + /// the table again. + fn admit(&mut self, location: &str, table_id: &str, (index, metadata): (usize, usize)) -> bool { + match self.tables.get(location) { + Some(seen) if seen == table_id => return true, + None if self.tables.len() < MAX_READ_TABLES => { + self.tables + .insert(location.to_string(), table_id.to_string()); + return true; + } + _ => {} + } + let generation = self.generation + 1; + *self = Self::new(index, metadata); + self.generation = generation; + false + } +} + pub struct VectorServer { storage: Option, read_cache: Option, + /// Lance's index and metadata caches, shared by every request that only + /// reads. Writes open in a session of their own. + pub(crate) reads: Mutex, + /// A table opened by an earlier read, by location and encryption, with + /// the generation of the session it was opened in. + pub(crate) handles: moka::sync::Cache<(String, bool), (u64, Dataset)>, + lance_cache_bytes: (usize, usize), + leases: SharedLeases, nprobes: Option, /// URI schemes a table's location may have. schemes: &'static [&'static str], @@ -463,6 +543,13 @@ impl VectorServer { Self { storage: None, read_cache: None, + reads: Mutex::new(Reads::new( + DEFAULT_INDEX_CACHE_BYTES, + DEFAULT_METADATA_CACHE_BYTES, + )), + lance_cache_bytes: (DEFAULT_INDEX_CACHE_BYTES, DEFAULT_METADATA_CACHE_BYTES), + handles: moka::sync::Cache::new(MAX_HANDLES), + leases: SharedLeases::default(), nprobes: None, schemes: &["s3"], } @@ -488,6 +575,13 @@ impl VectorServer { self } + /// Bytes of loaded indexes and of table metadata kept across reads. + pub fn with_lance_caches(mut self, index_bytes: usize, metadata_bytes: usize) -> Self { + self.reads = Mutex::new(Reads::new(index_bytes, metadata_bytes)); + self.lance_cache_bytes = (index_bytes, metadata_bytes); + self + } + /// Serve repeated reads of write-once table files from `cache`. pub fn with_read_cache(mut self, cache: ReadCache) -> Self { self.read_cache = Some(cache); @@ -529,9 +623,72 @@ impl VectorServer { } async fn open(&self, table: &pb::Table) -> Result { + self.open_in(table, None).await + } + + /// [`Self::open`] in the shared session, for a request that only reads. + async fn open_to_read(&self, table: &pb::Table) -> Result { + loop { + let (session, generation) = { + let reads = self.reads.lock().unwrap(); + (reads.session.clone(), reads.generation) + }; + let dataset = match self.handle(table, generation).await? { + Some(dataset) => dataset, + None => { + let dataset = self.open_in(table, Some(session)).await?; + // Under the lock, so no handle of a replaced session + // lands after the replacement cleared them. + let reads = self.reads.lock().unwrap(); + if table.version == 0 && reads.generation == generation { + self.handles + .insert(handle_key(table), (generation, dataset.clone())); + } + dataset + } + }; + let mut reads = self.reads.lock().unwrap(); + if reads.generation != generation { + continue; + } + if reads.admit(&table.location, table_id(&dataset), self.lance_cache_bytes) { + return Ok(dataset); + } + // The handles hold the replaced session and its caches. + self.handles.invalidate_all(); + } + } + + /// The table from the handle an earlier read of it left in this session, + /// moved to the version asked for: the latest costs a listing of its + /// versions, not a new store and manifest read. + async fn handle(&self, table: &pb::Table, generation: u64) -> Result, Status> { + self.describe(table)?; + let key = handle_key(table); + let Some((at, mut dataset)) = self.handles.get(&key) else { + return Ok(None); + }; + if at != generation { + return Ok(None); + } + let moved = match table.version { + 0 => dataset.checkout_latest().await.map(|()| dataset), + version => dataset.checkout_version(version).await, + }; + moved.map(Some).map_err(|err| { + self.handles.invalidate(&key); + open_status(err.into()) + }) + } + + async fn open_in( + &self, + table: &pb::Table, + session: Option>, + ) -> Result { let description = self.describe(table)?; let wrapper = self.wrapper(table); - let mut dataset = dataset::open_at(&description, wrapper) + let mut dataset = dataset::open_at(&description, wrapper, session) .await .map_err(open_status)?; if table.version != 0 { @@ -564,6 +721,13 @@ impl VectorService for VectorServer { // The gateway has already created the table through the catalog. let description = self.describe(&table)?; let schema = schema_for(dimension, metric, &req.non_filterable_metadata_keys); + let schema = Arc::new(Schema::new_with_metadata( + schema.fields().clone(), + HashMap::from([( + TABLE_ID_KEY.to_string(), + uuid::Uuid::new_v4().simple().to_string(), + )]), + )); let empty = RecordBatch::new_empty(schema.clone()); let reader = arrow_array::RecordBatchIterator::new(vec![Ok(empty)].into_iter(), schema); let dataset = Dataset::write( @@ -680,12 +844,10 @@ impl VectorService for VectorServer { let table = req .table .ok_or_else(|| Status::invalid_argument("missing table"))?; - let dataset = self.open(&table).await?; + let dataset = self.open_to_read(&table).await?; // Lease the resolved version so maintenance cannot reclaim its files // while this query runs against it. - let _lease = weed_vector_core::maintenance::try_pin(&dataset) - .await - .map_err(internal)?; + let _lease = self.leases.try_pin(&dataset).await.map_err(internal)?; let resolved = VectorTable::resolve(&dataset, None, Some(KEY_COLUMN)).map_err(invalid)?; let columns = projection(req.return_data, req.return_metadata); let batches = query::get(&dataset, &resolved, &req.keys, Some(&columns)) @@ -730,7 +892,7 @@ impl VectorService for VectorServer { "segment_index must be below segment_count, which is at most {MAX_SEGMENTS}" ))); } - let dataset = self.open(&table).await?; + let dataset = self.open_to_read(&table).await?; let offset: u64 = if req.next_token.is_empty() { 0 } else { @@ -750,9 +912,7 @@ impl VectorService for VectorServer { // Data and metadata come from a second read of this version, so // lease it as get_vectors does. let _lease = if columns.len() > 1 { - weed_vector_core::maintenance::try_pin(&dataset) - .await - .map_err(internal)? + self.leases.try_pin(&dataset).await.map_err(internal)? } else { None }; @@ -797,12 +957,10 @@ impl VectorService for VectorServer { let table = req .table .ok_or_else(|| Status::invalid_argument("missing table"))?; - let dataset = self.open(&table).await?; + let dataset = self.open_to_read(&table).await?; // Lease the resolved version so maintenance cannot reclaim its files // while this query runs against it. - let _lease = weed_vector_core::maintenance::try_pin(&dataset) - .await - .map_err(internal)?; + let _lease = self.leases.try_pin(&dataset).await.map_err(internal)?; let resolved = VectorTable::resolve(&dataset, None, Some(KEY_COLUMN)).map_err(invalid)?; if req.query_vector.len() != resolved.dimension { return Err(Status::invalid_argument(format!( diff --git a/seaweed-worker/crates/lance/src/dataset.rs b/seaweed-worker/crates/lance/src/dataset.rs index 93056056c1..5935f55620 100644 --- a/seaweed-worker/crates/lance/src/dataset.rs +++ b/seaweed-worker/crates/lance/src/dataset.rs @@ -15,6 +15,7 @@ use lance::io::{ ObjectStore, ObjectStoreParams, ObjectStoreRegistry, StorageOptionsAccessor, WrappingObjectStore, }; +use lance::session::Session; use crate::catalog::{NamespaceClient, TableDescription}; use crate::direct_read::DirectReads; @@ -74,7 +75,7 @@ pub fn unsigned_options() -> HashMap { pub async fn open(client: &NamespaceClient, id: &[String], storage: &Storage) -> Result { let description = client.describe_table(id).await?; let wrapper = storage.wrapper(&description.location).await?; - let dataset = open_at(&description, wrapper).await?; + let dataset = open_at(&description, wrapper, None).await?; Ok(OpenTable { id: id.to_vec(), location: description.location, @@ -82,13 +83,19 @@ pub async fn open(client: &NamespaceClient, id: &[String], storage: &Storage) -> }) } -/// Open a table the caller has already resolved, with the store wrapped. +/// Open a table the caller has already resolved, with the store wrapped, in +/// `session` when given, otherwise in a session of its own. pub async fn open_at( description: &TableDescription, wrapper: Option>, + session: Option>, ) -> Result { - DatasetBuilder::from_uri(&description.location) - .with_store_params(store_params(wrapper)) + let mut builder = + DatasetBuilder::from_uri(&description.location).with_store_params(store_params(wrapper)); + if let Some(session) = session { + builder = builder.with_session(session); + } + builder .load() .await .with_context(|| format!("open lance dataset at {}", description.location))