From 55f1e1efc57ecd8c2a4ec7c3ddf8a71cad1057be Mon Sep 17 00:00:00 2001 From: Chris Lu Date: Sun, 1 Nov 2026 12:59:47 -0800 Subject: [PATCH] vector: cache reads of write-once Lance files across requests (#12303) * lance worker: let dataset::open take an object store wrapper open_with is open with the dataset's object store wrapped, so a caller such as weed-vector can put a read cache in front of the table's files. open keeps its signature and opens unwrapped, as before. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_011UpvafvK8JpagF94MHWqwK * vector: cache reads of write-once Lance files across requests Every VectorService call opens the table afresh, and Lance keeps no data pages between opens, so each QueryVectors re-read the whole vector column through the S3 gateway: 83 MB and about 228 range GETs per query at 50k x 384. The gateway and volume server spent their CPU streaming it and throughput stayed at 6-8 queries/s whatever the client concurrency. weed-vector now wraps each table's object store with a bounded in-memory cache (--read-cache-mb, 512 MiB by default, 0 disables). It only serves data files, deletion files and index files: Lance writes each once under a fresh random name, so a cached read cannot go stale, even when an index is dropped and recreated at the same location. Manifests, tags and every other path still go to the store, and each call still describes the table, loads the latest version and takes the query lease, so reads see every write committed before them. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_011UpvafvK8JpagF94MHWqwK * vector: cache an owned copy of each fetched read A range from get_ranges can be a slice of a larger response buffer, such as a coalesced fetch, and caching the slice kept the whole buffer alive while the weigher counted only the slice. The cache now copies each fetched range and whole-object read into its own buffer, so --read-cache-mb bounds what it actually holds. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_011UpvafvK8JpagF94MHWqwK --------- Co-authored-by: Claude Opus 5.5 --- seaweed-vector/Cargo.lock | 4 + seaweed-vector/crates/server/Cargo.toml | 4 + seaweed-vector/crates/server/src/lib.rs | 1 + seaweed-vector/crates/server/src/main.rs | 9 +- .../crates/server/src/read_cache.rs | 565 ++++++++++++++++++ seaweed-vector/crates/server/src/service.rs | 23 +- seaweed-worker/crates/lance/src/dataset.rs | 31 +- 7 files changed, 630 insertions(+), 7 deletions(-) create mode 100644 seaweed-vector/crates/server/src/read_cache.rs diff --git a/seaweed-vector/Cargo.lock b/seaweed-vector/Cargo.lock index e1a3225934..da09df9035 100644 --- a/seaweed-vector/Cargo.lock +++ b/seaweed-vector/Cargo.lock @@ -6486,11 +6486,15 @@ dependencies = [ "arrow-array", "arrow-cast", "arrow-schema", + "async-trait", + "bytes", "clap", "futures", "lance", "lance-arrow", "lance-linalg", + "moka", + "object_store", "prost", "prost-types", "protoc-bin-vendored", diff --git a/seaweed-vector/crates/server/Cargo.toml b/seaweed-vector/crates/server/Cargo.toml index 5f938f66a4..8d65d590d1 100644 --- a/seaweed-vector/crates/server/Cargo.toml +++ b/seaweed-vector/crates/server/Cargo.toml @@ -27,8 +27,12 @@ reqwest = { version = "0.12", default-features = false } tonic = { version = "0.14", features = ["tls-aws-lc"] } tonic-prost = "0.14" anyhow.workspace = true +async-trait = "0.1" +bytes = "1" clap = { version = "4", features = ["derive", "env"] } futures = "0.3" +moka = { version = "0.12", features = ["sync"] } +object_store = "0.13" serde = { version = "1", features = ["derive"] } serde_json = "1" tokio.workspace = true diff --git a/seaweed-vector/crates/server/src/lib.rs b/seaweed-vector/crates/server/src/lib.rs index 322ec27dcc..84db0b59cd 100644 --- a/seaweed-vector/crates/server/src/lib.rs +++ b/seaweed-vector/crates/server/src/lib.rs @@ -1,5 +1,6 @@ //! The VectorService gRPC server the S3 gateway talks to. +pub mod read_cache; pub mod service; pub mod remote_pb { diff --git a/seaweed-vector/crates/server/src/main.rs b/seaweed-vector/crates/server/src/main.rs index 453aa8ed5f..10d46f372d 100644 --- a/seaweed-vector/crates/server/src/main.rs +++ b/seaweed-vector/crates/server/src/main.rs @@ -4,6 +4,7 @@ use tonic::transport::Server; use weed_lance_worker::catalog::NamespaceClient; use weed_lance_worker::dataset::FallbackOptions; use weed_vector_server::pb::vector::vector_service_server::VectorServiceServer; +use weed_vector_server::read_cache::ReadCache; use weed_vector_server::service::VectorServer; /// The VectorService the S3 gateway fans vector calls out to. @@ -27,6 +28,9 @@ struct Args { /// gateway discovers it via ListClusterNodes(client_type=vector). #[arg(long, alias = "master", env = "WEED_VECTOR_MASTERS")] masters: Option, + /// 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, } fn parse_option(s: &str) -> Result<(String, String), String> { @@ -108,7 +112,10 @@ async fn main() -> Result<()> { } let fallback: FallbackOptions = args.storage_options.into_iter().collect(); let client = NamespaceClient::new(args.namespace.clone()).with_fallback(&fallback); - let service = VectorServer::new(client, fallback); + let mut service = VectorServer::new(client, fallback); + if args.read_cache_mb > 0 { + service = service.with_read_cache(ReadCache::new(args.read_cache_mb << 20)); + } tracing::info!(namespace = %args.namespace, listen = %args.listen, "weed-vector"); Server::builder() // Match the gateway's 1 GiB cap — put batches over the tonic 4 MiB diff --git a/seaweed-vector/crates/server/src/read_cache.rs b/seaweed-vector/crates/server/src/read_cache.rs new file mode 100644 index 0000000000..748bd4b66d --- /dev/null +++ b/seaweed-vector/crates/server/src/read_cache.rs @@ -0,0 +1,565 @@ +//! A bounded in-memory cache of reads from the files Lance never rewrites. +//! +//! Data, index and deletion files are written once under fresh random names, +//! so a cached read cannot go stale, even when an index is dropped and +//! recreated at the same location. Manifests, tags and every other path are +//! read from the store each time, which keeps read-your-writes intact. + +use std::fmt; +use std::ops::Range; +use std::sync::Arc; + +use bytes::Bytes; +use futures::StreamExt; +use futures::stream::BoxStream; +use lance::io::WrappingObjectStore; +use object_store::path::Path; +use object_store::{ + Attributes, CopyOptions, GetOptions, GetResult, GetResultPayload, ListResult, MultipartUpload, + ObjectMeta, ObjectStore, PutMultipartOptions, PutOptions, PutPayload, PutResult, RenameOptions, + Result, +}; + +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +struct Key { + store: Arc, + path: Path, + /// `None` for a whole-object read. + range: Option>, +} + +#[derive(Debug, Clone)] +struct Entry { + bytes: Bytes, + /// Set for a whole-object read, which answers with its metadata. + object: Option>, +} + +/// Wraps each object store Lance opens; clones share one cache. +#[derive(Clone)] +pub struct ReadCache(moka::sync::Cache); + +impl ReadCache { + pub fn new(capacity_bytes: u64) -> Self { + Self( + moka::sync::Cache::builder() + .max_capacity(capacity_bytes) + .weigher(|key: &Key, entry: &Entry| { + u32::try_from(key.path.as_ref().len() + entry.bytes.len()).unwrap_or(u32::MAX) + }) + .build(), + ) + } +} + +impl fmt::Debug for ReadCache { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("ReadCache") + .field("entries", &self.0.entry_count()) + .field("bytes", &self.0.weighted_size()) + .finish() + } +} + +impl WrappingObjectStore for ReadCache { + fn wrap(&self, store_prefix: &str, original: Arc) -> Arc { + Arc::new(CachedStore { + inner: original, + store: store_prefix.into(), + cache: self.0.clone(), + }) + } +} + +/// Lance names data files `data/.lance`, deletion files +/// `_deletions/--.` and index files +/// `_indices//`; nothing else under a table is write-once. +fn immutable(path: &Path) -> bool { + let parts: Vec<_> = path.parts().collect(); + let Some((name, dirs)) = parts.split_last() else { + return false; + }; + match dirs { + [.., dir] if dir.as_ref() == "data" => name.as_ref().ends_with(".lance"), + [.., dir] if dir.as_ref() == "_deletions" => true, + [.., dir, id] if dir.as_ref() == "_indices" => !id.as_ref().starts_with('_'), + _ => false, + } +} + +/// Lance reads files up to its block size whole; larger whole reads stream +/// through uncached. +const MAX_WHOLE_OBJECT: u64 = 4 << 20; + +/// A read with no conditions, so its answer depends only on the object. +fn plain(options: &GetOptions) -> bool { + options.if_match.is_none() + && options.if_none_match.is_none() + && options.if_modified_since.is_none() + && options.if_unmodified_since.is_none() + && options.range.is_none() + && options.version.is_none() + && !options.head +} + +/// A fetched range can be a slice of a larger response buffer; the cache +/// keeps its own copy so it holds only the bytes it counts. +fn owned(bytes: Bytes) -> Bytes { + Bytes::copy_from_slice(&bytes) +} + +fn whole(bytes: Bytes, object: &(ObjectMeta, Attributes)) -> GetResult { + GetResult { + range: 0..bytes.len() as u64, + payload: GetResultPayload::Stream(futures::stream::once(async { Ok(bytes) }).boxed()), + meta: object.0.clone(), + attributes: object.1.clone(), + } +} + +#[derive(Debug)] +struct CachedStore { + inner: Arc, + store: Arc, + cache: moka::sync::Cache, +} + +impl CachedStore { + fn key(&self, location: &Path, range: Option>) -> Key { + Key { + store: self.store.clone(), + path: location.clone(), + range, + } + } +} + +impl fmt::Display for CachedStore { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "ReadCache({})", self.inner) + } +} + +#[async_trait::async_trait] +#[deny(clippy::missing_trait_methods)] +impl ObjectStore for CachedStore { + async fn put_opts( + &self, + location: &Path, + payload: PutPayload, + opts: PutOptions, + ) -> Result { + self.inner.put_opts(location, payload, opts).await + } + + async fn put_multipart_opts( + &self, + location: &Path, + opts: PutMultipartOptions, + ) -> Result> { + self.inner.put_multipart_opts(location, opts).await + } + + async fn get_opts(&self, location: &Path, options: GetOptions) -> Result { + if !immutable(location) || !plain(&options) { + return self.inner.get_opts(location, options).await; + } + let key = self.key(location, None); + if let Some(Entry { + bytes, + object: Some(object), + }) = self.cache.get(&key) + { + return Ok(whole(bytes, &object)); + } + let result = self.inner.get_opts(location, options).await?; + if result.meta.size > MAX_WHOLE_OBJECT { + return Ok(result); + } + let object = Arc::new((result.meta.clone(), result.attributes.clone())); + let bytes = owned(result.bytes().await?); + self.cache.insert( + key, + Entry { + bytes: bytes.clone(), + object: Some(object.clone()), + }, + ); + Ok(whole(bytes, &object)) + } + + async fn get_ranges(&self, location: &Path, ranges: &[Range]) -> Result> { + if !immutable(location) { + return self.inner.get_ranges(location, ranges).await; + } + let mut out: Vec> = ranges + .iter() + .map(|range| { + self.cache + .get(&self.key(location, Some(range.clone()))) + .map(|entry| entry.bytes) + }) + .collect(); + let missing: Vec> = ranges + .iter() + .zip(&out) + .filter(|(_, hit)| hit.is_none()) + .map(|(range, _)| range.clone()) + .collect(); + if !missing.is_empty() { + let mut fetched = self.inner.get_ranges(location, &missing).await?.into_iter(); + for (range, slot) in ranges.iter().zip(out.iter_mut()) { + if slot.is_some() { + continue; + } + let Some(bytes) = fetched.next() else { + break; + }; + let bytes = owned(bytes); + self.cache.insert( + self.key(location, Some(range.clone())), + Entry { + bytes: bytes.clone(), + object: None, + }, + ); + *slot = Some(bytes); + } + } + out.into_iter() + .map(|bytes| { + bytes.ok_or_else(|| object_store::Error::Generic { + store: "ReadCache", + source: "get_ranges returned fewer ranges than requested".into(), + }) + }) + .collect() + } + + fn delete_stream( + &self, + locations: BoxStream<'static, Result>, + ) -> BoxStream<'static, Result> { + self.inner.delete_stream(locations) + } + + fn list(&self, prefix: Option<&Path>) -> BoxStream<'static, Result> { + self.inner.list(prefix) + } + + fn list_with_offset( + &self, + prefix: Option<&Path>, + offset: &Path, + ) -> BoxStream<'static, Result> { + self.inner.list_with_offset(prefix, offset) + } + + async fn list_with_delimiter(&self, prefix: Option<&Path>) -> Result { + self.inner.list_with_delimiter(prefix).await + } + + async fn copy_opts(&self, from: &Path, to: &Path, options: CopyOptions) -> Result<()> { + self.inner.copy_opts(from, to, options).await + } + + async fn rename_opts(&self, from: &Path, to: &Path, options: RenameOptions) -> Result<()> { + self.inner.rename_opts(from, to, options).await + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::pb::vector as pb; + use crate::pb::vector::vector_service_server::VectorService; + use crate::service::VectorServer; + use object_store::ObjectStoreExt; + use object_store::memory::InMemory; + use tonic::Request; + use weed_lance_worker::catalog::NamespaceClient; + use weed_lance_worker::dataset::FallbackOptions; + + fn cached(inner: Arc) -> Arc { + ReadCache::new(1 << 20).wrap("memory", inner) + } + + #[test] + fn only_write_once_files_are_immutable() { + for path in [ + "vector/x/data/0101abcd.lance", + "vector/data/data/0101abcd.lance", + "vector/x/_deletions/3-7-12345.arrow", + "vector/x/_indices/1b6f/index.idx", + ] { + assert!(immutable(&Path::from(path)), "{path}"); + } + for path in [ + "vector/data/_versions/1.manifest", + "vector/data/_latest.manifest", + "vector/x/_versions/18446744073709551614.manifest", + "vector/x/_refs/tags/weed-vector-query-1.json", + "vector/x/_transactions/0-abc.txn", + "vector/x/data/notes.txt", + "vector/_indices/_versions/1.manifest", + "data/x.lance/_versions/1.manifest", + ] { + assert!(!immutable(&Path::from(path)), "{path}"); + } + } + + async fn read(store: &Arc, path: &Path, range: Range) -> Bytes { + let mut got = store + .get_ranges(path, std::slice::from_ref(&range)) + .await + .unwrap(); + got.remove(0) + } + + #[tokio::test] + async fn write_once_reads_come_from_the_cache() { + let inner = Arc::new(InMemory::new()); + let store = cached(inner.clone()); + let data = Path::from("vector/x/data/a.lance"); + let manifest = Path::from("vector/x/_versions/1.manifest"); + for path in [&data, &manifest] { + inner + .put(path, Bytes::from_static(b"first").into()) + .await + .unwrap(); + assert_eq!(read(&store, path, 0..3).await, &b"fir"[..]); + assert_eq!( + store.get(path).await.unwrap().bytes().await.unwrap(), + &b"first"[..] + ); + inner + .put(path, Bytes::from_static(b"second").into()) + .await + .unwrap(); + } + // The overwrite stands in for a store answering differently: the data + // file keeps its cached bytes, the manifest is read again. + assert_eq!( + store.get_ranges(&data, &[0..3, 1..4]).await.unwrap(), + [&b"fir"[..], &b"eco"[..]] + ); + let got = store.get(&data).await.unwrap(); + assert_eq!(got.meta.size, 5); + assert_eq!(got.bytes().await.unwrap(), &b"first"[..]); + assert_eq!(read(&store, &manifest, 0..3).await, &b"sec"[..]); + assert_eq!( + store.get(&manifest).await.unwrap().bytes().await.unwrap(), + &b"second"[..] + ); + // A conditional or ranged get is the store's to answer. + let ranged = GetOptions { + range: Some((0..2).into()), + ..Default::default() + }; + let got = store.get_opts(&data, ranged).await.unwrap(); + assert_eq!(got.bytes().await.unwrap(), &b"se"[..]); + } + + #[tokio::test] + async fn the_cache_stays_within_its_capacity() { + let cache = ReadCache::new(64 << 10); + let inner = Arc::new(InMemory::new()); + let store = cache.wrap("memory", inner.clone()); + for i in 0..64 { + let path = Path::from(format!("vector/x/data/{i}.lance")); + inner + .put(&path, Bytes::from(vec![0u8; 8 << 10]).into()) + .await + .unwrap(); + read(&store, &path, 0..(8 << 10)).await; + } + cache.0.run_pending_tasks(); + assert!(cache.0.weighted_size() <= 64 << 10, "{cache:?}"); + } + + #[tokio::test] + async fn a_cached_range_does_not_pin_its_response_buffer() { + let cache = ReadCache::new(1 << 20); + let inner = Arc::new(InMemory::new()); + let store = cache.wrap("memory", inner.clone()); + let path = Path::from("vector/x/data/big.lance"); + inner + .put(&path, Bytes::from(vec![7u8; 256 << 10]).into()) + .await + .unwrap(); + // InMemory answers a range with a slice of the whole object. + let whole = inner.get(&path).await.unwrap().bytes().await.unwrap(); + let backing = whole.as_ptr_range(); + read(&store, &path, 10..20).await; + let key = Key { + store: "memory".into(), + path, + range: Some(10..20), + }; + let cached = cache.0.get(&key).unwrap().bytes; + assert_eq!(cached.len(), 10); + assert!(!backing.contains(&cached.as_ptr())); + } + + /// A catalog that answers every request with `location`. + async fn catalog_at(location: String) -> String { + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let url = format!("http://{}", listener.local_addr().unwrap()); + let body = serde_json::json!({ "location": location }).to_string(); + tokio::spawn(async move { + loop { + let (mut conn, _) = listener.accept().await.unwrap(); + let body = body.clone(); + tokio::spawn(async move { + let mut request = [0u8; 8192]; + let _ = conn.read(&mut request).await; + let response = format!( + "HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{body}", + body.len() + ); + let _ = conn.write_all(response.as_bytes()).await; + let _ = conn.shutdown().await; + }); + } + }); + url + } + + fn vector(key: &str, data: Vec) -> pb::Vector { + pb::Vector { + key: key.to_string(), + data, + metadata: None, + } + } + + async fn get(server: &VectorServer, table: &pb::Table, keys: &[&str]) -> Vec { + let mut vectors = server + .get_vectors(Request::new(pb::GetVectorsRequest { + table: Some(table.clone()), + keys: keys.iter().map(|k| k.to_string()).collect(), + return_data: true, + return_metadata: false, + })) + .await + .unwrap() + .into_inner() + .vectors; + vectors.sort_by(|a, b| a.key.cmp(&b.key)); + vectors + } + + async fn query(server: &VectorServer, table: &pb::Table, v: Vec) -> Vec { + server + .query_vectors(Request::new(pb::QueryVectorsRequest { + table: Some(table.clone()), + query_vector: v, + top_k: 10, + ..Default::default() + })) + .await + .unwrap() + .into_inner() + .vectors + .into_iter() + .map(|r| r.key) + .collect() + } + + async fn create(server: &VectorServer, table: &pb::Table, dimension: u32) { + server + .create_index(Request::new(pb::CreateIndexRequest { + table: Some(table.clone()), + dimension, + ..Default::default() + })) + .await + .unwrap(); + } + + async fn put(server: &VectorServer, table: &pb::Table, vectors: Vec) { + server + .put_vectors(Request::new(pb::PutVectorsRequest { + table: Some(table.clone()), + vectors, + })) + .await + .unwrap(); + } + + /// Reads through the cache see every committed write, and an index + /// dropped and recreated at the same location answers with its new rows. + #[tokio::test] + async fn cached_reads_follow_writes_and_a_recreated_index() { + let location = format!("shared-memory://cache-{}/vector/x", uuid_like(),); + let server = VectorServer::new( + NamespaceClient::new(catalog_at(location.clone()).await), + FallbackOptions::new(), + ) + .with_read_cache(ReadCache::new(16 << 20)); + let table = pb::Table { + id: vec!["bkt".into(), "vector".into(), "x".into()], + version: 0, + }; + + create(&server, &table, 4).await; + put( + &server, + &table, + vec![ + vector("a", vec![1.0, 0.0, 0.0, 0.0]), + vector("b", vec![0.0, 1.0, 0.0, 0.0]), + ], + ) + .await; + assert_eq!( + query(&server, &table, vec![1.0, 0.0, 0.0, 0.0]).await, + ["a", "b"] + ); + assert_eq!( + get(&server, &table, &["a"]).await[0].data, + [1.0, 0.0, 0.0, 0.0] + ); + + put( + &server, + &table, + vec![ + vector("a", vec![0.0, 0.0, 2.0, 0.0]), + vector("c", vec![1.0, 0.1, 0.0, 0.0]), + ], + ) + .await; + assert_eq!( + query(&server, &table, vec![1.0, 0.0, 0.0, 0.0]).await, + ["c", "b", "a"] + ); + assert_eq!( + get(&server, &table, &["a"]).await[0].data, + [0.0, 0.0, 2.0, 0.0] + ); + + let (store, path) = lance::io::ObjectStore::from_uri(&location).await.unwrap(); + store.remove_dir_all(path).await.unwrap(); + create(&server, &table, 2).await; + put(&server, &table, vec![vector("z", vec![0.5, 0.5])]).await; + let vectors = get(&server, &table, &["a", "b", "c", "z"]).await; + assert_eq!(vectors.len(), 1); + assert_eq!( + (vectors[0].key.as_str(), vectors[0].data.as_slice()), + ("z", &[0.5, 0.5][..]) + ); + assert_eq!(query(&server, &table, vec![1.0, 0.0]).await, ["z"]); + } + + fn uuid_like() -> String { + format!( + "{:x}", + std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap() + .as_nanos() + ) + } +} diff --git a/seaweed-vector/crates/server/src/service.rs b/seaweed-vector/crates/server/src/service.rs index fb8ab22264..cb9f51b1c4 100644 --- a/seaweed-vector/crates/server/src/service.rs +++ b/seaweed-vector/crates/server/src/service.rs @@ -10,6 +10,7 @@ use lance::dataset::WriteParams; use lance::dataset::{MergeInsertBuilder, WhenMatched, WhenNotMatched}; use lance::io::ObjectStoreParams; use lance::io::StorageOptionsAccessor; +use lance::io::WrappingObjectStore; use lance::table::format::Fragment; use lance_linalg::distance::DistanceType; use prost_types::{ListValue, Struct, Value, value::Kind}; @@ -22,6 +23,7 @@ use weed_vector_core::query::{self, Query}; use weed_vector_core::table::VectorTable; use crate::pb::vector::{self as pb, vector_service_server::VectorService}; +use crate::read_cache::ReadCache; fn internal(err: impl Into) -> Status { Status::internal(format!("{:#}", err.into())) @@ -371,15 +373,32 @@ fn rows_as_vectors( pub struct VectorServer { client: NamespaceClient, fallback: FallbackOptions, + read_cache: Option>, } impl VectorServer { pub fn new(client: NamespaceClient, fallback: FallbackOptions) -> Self { - Self { client, fallback } + Self { + client, + fallback, + read_cache: None, + } + } + + /// Serve repeated reads of write-once table files from `cache`. + pub fn with_read_cache(mut self, cache: ReadCache) -> Self { + self.read_cache = Some(Arc::new(cache)); + self } async fn open(&self, table: &pb::Table) -> Result { - let opened = dataset::open(&self.client, &table.id, &self.fallback).await?; + let opened = dataset::open_with( + &self.client, + &table.id, + &self.fallback, + self.read_cache.clone(), + ) + .await?; let mut dataset = opened.dataset; if table.version != 0 { dataset = dataset.checkout_version(table.version).await?; diff --git a/seaweed-worker/crates/lance/src/dataset.rs b/seaweed-worker/crates/lance/src/dataset.rs index 19eb1b82e4..d28bcf2d4c 100644 --- a/seaweed-worker/crates/lance/src/dataset.rs +++ b/seaweed-worker/crates/lance/src/dataset.rs @@ -11,7 +11,10 @@ use std::sync::Arc; use anyhow::{Context, Result}; use lance::dataset::Dataset; use lance::dataset::builder::DatasetBuilder; -use lance::io::{ObjectStore, ObjectStoreParams, ObjectStoreRegistry, StorageOptionsAccessor}; +use lance::io::{ + ObjectStore, ObjectStoreParams, ObjectStoreRegistry, StorageOptionsAccessor, + WrappingObjectStore, +}; use crate::catalog::{NamespaceClient, TableDescription}; @@ -42,9 +45,19 @@ pub async fn open( client: &NamespaceClient, id: &[String], fallback: &FallbackOptions, +) -> Result { + open_with(client, id, fallback, None).await +} + +/// [`open`] with the dataset's object store wrapped, e.g. by a read cache. +pub async fn open_with( + client: &NamespaceClient, + id: &[String], + fallback: &FallbackOptions, + wrapper: Option>, ) -> Result { let description = client.describe_table(id).await?; - let dataset = open_at(&description, fallback).await?; + let dataset = open_at(&description, fallback, wrapper).await?; Ok(OpenTable { id: id.to_vec(), location: description.location, @@ -56,11 +69,21 @@ pub async fn open( /// through. What it vends always wins: the fallback exists because a deployment /// without STS vends no credentials at all, and then the worker has no other way /// to reach the data. -async fn open_at(description: &TableDescription, fallback: &FallbackOptions) -> Result { +async fn open_at( + description: &TableDescription, + fallback: &FallbackOptions, + wrapper: Option>, +) -> Result { let mut options: HashMap = fallback.clone(); options.extend(description.storage_options.clone()); DatasetBuilder::from_uri(&description.location) - .with_storage_options(options) + .with_store_params(ObjectStoreParams { + storage_options_accessor: Some(Arc::new(StorageOptionsAccessor::with_static_options( + options, + ))), + object_store_wrapper: wrapper, + ..Default::default() + }) .load() .await .with_context(|| format!("open lance dataset at {}", description.location))