diff --git a/rust/lancedb/src/database/listing.rs b/rust/lancedb/src/database/listing.rs index ed217e632..6ba3cb029 100644 --- a/rust/lancedb/src/database/listing.rs +++ b/rust/lancedb/src/database/listing.rs @@ -25,7 +25,7 @@ use crate::database::namespace::LanceNamespaceDatabase; use crate::error::{CreateDirSnafu, Error, Result}; use crate::io::object_store::MirroringObjectStoreWrapper; use crate::table::NativeTable; -use crate::utils::validate_table_name; +use crate::utils::{PatchStoreParam, validate_table_name}; use lance_namespace::models::{ CreateNamespaceRequest, CreateNamespaceResponse, DescribeNamespaceRequest, @@ -355,6 +355,13 @@ impl ListingDatabase { url.to_string() } + fn uses_local_file_provider(object_store: &ObjectStore) -> bool { + matches!( + object_store.scheme(), + "file" | "file-object-store" | "file+uring" + ) + } + async fn prepare_namespace_root( uri: &str, storage_options: &HashMap, @@ -582,7 +589,7 @@ impl ListingDatabase { None => None, }; #[cfg(windows)] - let write_store_wrapper = if object_store.is_local() { + let write_store_wrapper = if Self::uses_local_file_provider(&object_store) { // Local manifest commits need create-only rename semantics, // including on filesystems that do not support hard links. Some( @@ -657,7 +664,7 @@ impl ListingDatabase { .await?; #[cfg(windows)] - let write_store_wrapper = object_store.is_local().then(|| { + let write_store_wrapper = Self::uses_local_file_provider(&object_store).then(|| { // Local manifest commits need create-only rename semantics, // including on filesystems that do not support hard links. Arc::new(crate::io::object_store::windows::WindowsLocalFileSystemWrapper) @@ -1133,6 +1140,12 @@ impl Database for ListingDatabase { }, ..Default::default() }; + let storage_params = match self.store_wrapper.clone() { + Some(wrapper) => Some(storage_params) + .patch_with_store_wrapper(wrapper)? + .expect("patching store params always returns parameters"), + None => storage_params, + }; let read_params = ReadParams { store_options: Some(storage_params.clone()), session: Some(self.session.clone()), @@ -1315,6 +1328,7 @@ mod tests { use crate::connection::ConnectRequest; use crate::data::scannable::Scannable; use crate::database::{CreateTableMode, CreateTableRequest}; + use crate::io::object_store::io_tracking::IoStatsHolder; use crate::query::QueryRequest; use crate::table::{AnyQuery, WriteOptions}; use arrow_array::{Int32Array, RecordBatch, StringArray}; @@ -1414,6 +1428,25 @@ mod tests { assert!(!tempdir.path().join("__manifest").exists()); } + #[tokio::test] + async fn test_file_object_store_uses_local_file_provider() { + let tempdir = tempdir().unwrap(); + let path = tempdir.path().to_string_lossy().replace('\\', "/"); + let uri = if path.starts_with('/') { + format!("file-object-store://{path}") + } else { + format!("file-object-store:///{path}") + }; + let registry = Arc::new(lance_io::object_store::ObjectStoreRegistry::default()); + let (store, _) = + ObjectStore::from_uri_and_params(registry, &uri, &ObjectStoreParams::default()) + .await + .unwrap(); + + assert_eq!(store.scheme(), "file-object-store"); + assert!(ListingDatabase::uses_local_file_provider(&store)); + } + /// Regression test for https://github.com/lancedb/lancedb/issues/1600. /// /// Opening a table used to create a separate object-store client instead of @@ -1623,6 +1656,46 @@ mod tests { ); } + #[tokio::test] + async fn test_clone_table_uses_connection_store_wrapper() { + let (_tempdir, mut db) = setup_database().await; + let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)])); + db.create_table(CreateTableRequest { + name: "source_table".to_string(), + namespace_path: vec![], + data: Box::new(RecordBatch::new_empty(schema)) as Box, + mode: CreateTableMode::Create, + write_options: Default::default(), + location: None, + namespace_client: None, + }) + .await + .unwrap(); + + let source_uri = db.table_uri("source_table").unwrap(); + let tracker = IoStatsHolder::default(); + db.store_wrapper = Some(Arc::new(tracker.clone())); + let _ = tracker.incremental_stats(); + + db.clone_table(CloneTableRequest { + target_table_name: "cloned_table".to_string(), + target_namespace_path: vec![], + source_uri, + source_version: None, + source_tag: None, + is_shallow: true, + namespace_client: None, + }) + .await + .unwrap(); + + let stats = tracker.incremental_stats(); + assert!( + stats.write_iops > 0, + "clone bypassed the wrapper: {stats:?}" + ); + } + #[tokio::test] async fn test_clone_table_with_data() { let (_tempdir, db) = setup_database().await; diff --git a/rust/lancedb/src/table.rs b/rust/lancedb/src/table.rs index 3e0137314..a4c5ea196 100644 --- a/rust/lancedb/src/table.rs +++ b/rust/lancedb/src/table.rs @@ -60,7 +60,7 @@ use crate::index::{IndexConfig, IndexStatisticsImpl, IndexType}; use crate::job::Job; use crate::query::{IntoQueryVector, Query, QueryExecutionOptions, TakeQuery, VectorQuery}; use crate::table::datafusion::insert::InsertExec; -use crate::utils::{PatchWriteParam, resolve_arrow_field_path}; +use crate::utils::{PatchReadParam, PatchWriteParam, resolve_arrow_field_path}; use self::dataset::DatasetConsistencyWrapper; use self::merge::MergeInsertBuilder; @@ -2385,6 +2385,20 @@ impl NativeTable { managed_versioning: Option, ) -> Result { let params = params.unwrap_or_default(); + let has_caller_store_wrapper = params + .store_options + .as_ref() + .and_then(|options| options.object_store_wrapper.as_ref()) + .is_some(); + // A caller wrapper must remain outside connection-level compatibility + // behavior. When there is no caller wrapper, apply the compatibility + // layer after loading so the session's registered store can be reused. + let (params, wrapper_after_load) = match write_store_wrapper { + Some(wrapper) if has_caller_store_wrapper => { + (params.patch_with_store_wrapper(wrapper)?, None) + } + wrapper => (params, wrapper), + }; // Build table_id from namespace + name let mut table_id = namespace.clone(); @@ -2443,7 +2457,7 @@ impl NativeTable { // connection-level write wrapper. Wrapper identity is part of the // registry key, so including it in ReadParams prevents reuse when the // opened table (and its wrapped store) is short-lived. - let dataset = match write_store_wrapper { + let dataset = match wrapper_after_load { Some(wrapper) => dataset.with_object_store_wrappers([wrapper]), None => dataset, }; @@ -2548,6 +2562,17 @@ impl NativeTable { if let Some(sess) = session { params.session(sess); } + let has_caller_store_wrapper = params + .store_options + .as_ref() + .and_then(|options| options.object_store_wrapper.as_ref()) + .is_some(); + let (params, wrapper_after_load) = match write_store_wrapper { + Some(wrapper) if has_caller_store_wrapper => { + (params.patch_with_store_wrapper(wrapper)?, None) + } + wrapper => (params, wrapper), + }; // Build table_id from namespace + name let mut table_id = namespace.clone(); @@ -2573,7 +2598,7 @@ impl NativeTable { // Apply the write wrapper after the session registry has resolved the // shared store. The cloned dataset retains the wrapper for subsequent // reads, manifest commits, and any additional base stores. - let dataset = match write_store_wrapper { + let dataset = match wrapper_after_load { Some(wrapper) => dataset.with_object_store_wrappers([wrapper]), None => dataset, }; @@ -3711,8 +3736,8 @@ pub struct FragmentSummaryStats { #[cfg(test)] #[allow(deprecated)] mod tests { - use std::sync::Arc; use std::sync::atomic::{AtomicBool, Ordering}; + use std::sync::{Arc, Mutex}; use std::time::Duration; use arrow_array::{ @@ -3942,6 +3967,66 @@ mod tests { } } + #[derive(Debug)] + struct OrderedStoreWrapper { + name: &'static str, + order: Arc>>, + } + + impl WrappingObjectStore for OrderedStoreWrapper { + fn wrap( + &self, + _store_prefix: &str, + original: Arc, + ) -> Arc { + self.order.lock().unwrap().push(self.name); + original + } + } + + #[tokio::test] + async fn test_open_with_params_keeps_caller_store_wrapper_outermost() { + let tmp_dir = tempdir().unwrap(); + let dataset_path = tmp_dir.path().join("test.lance"); + let uri = dataset_path.to_str().unwrap(); + let batch = make_test_batches(); + let reader = RecordBatchIterator::new(vec![Ok(batch.clone())], batch.schema()); + Dataset::write(reader, uri, None).await.unwrap(); + + let order = Arc::new(Mutex::new(Vec::new())); + let caller_wrapper = Arc::new(OrderedStoreWrapper { + name: "caller", + order: order.clone(), + }); + let compatibility_wrapper = Arc::new(OrderedStoreWrapper { + name: "compatibility", + order: order.clone(), + }); + let params = ReadParams { + store_options: Some(ObjectStoreParams { + object_store_wrapper: Some(caller_wrapper), + ..Default::default() + }), + ..Default::default() + }; + + NativeTable::open_with_params( + uri, + "test", + vec![], + Some(compatibility_wrapper), + Some(params), + None, + None, + HashSet::new(), + None, + ) + .await + .unwrap(); + + assert_eq!(*order.lock().unwrap(), vec!["compatibility", "caller"]); + } + #[tokio::test] async fn test_open_table_options() { let tmp_dir = tempdir().unwrap(); diff --git a/rust/lancedb/src/utils/mod.rs b/rust/lancedb/src/utils/mod.rs index c7bc95c9e..b5e5a6f4d 100644 --- a/rust/lancedb/src/utils/mod.rs +++ b/rust/lancedb/src/utils/mod.rs @@ -39,7 +39,10 @@ impl PatchStoreParam for Option { ) -> Result> { let mut params = self.unwrap_or_default(); params.object_store_wrapper = Some(match params.object_store_wrapper.take() { - Some(existing) => Arc::new(ChainedWrappingObjectStore::new(vec![existing, wrapper])), + // The wrapper being patched in is connection-level compatibility + // behavior. Keep it closest to the target store so an existing + // caller wrapper remains outermost and can observe every operation. + Some(existing) => Arc::new(ChainedWrappingObjectStore::new(vec![wrapper, existing])), None => wrapper, }); @@ -470,7 +473,7 @@ impl Stream for MaxBatchLengthStream { #[cfg(test)] mod tests { - use std::sync::atomic::{AtomicUsize, Ordering}; + use std::sync::Mutex; use arrow_array::Int32Array; use arrow_schema::Field; @@ -482,30 +485,38 @@ mod tests { use super::*; #[derive(Debug)] - struct CountingStoreWrapper(Arc); + struct OrderedStoreWrapper { + name: &'static str, + order: Arc>>, + } - impl WrappingObjectStore for CountingStoreWrapper { + impl WrappingObjectStore for OrderedStoreWrapper { fn wrap( &self, _store_prefix: &str, original: Arc, ) -> Arc { - self.0.fetch_add(1, Ordering::Relaxed); + self.order.lock().unwrap().push(self.name); original } } #[test] - fn test_patch_store_param_chains_wrappers() { - let existing_count = Arc::new(AtomicUsize::new(0)); - let added_count = Arc::new(AtomicUsize::new(0)); + fn test_patch_store_param_keeps_caller_wrapper_outermost() { + let order = Arc::new(Mutex::new(Vec::new())); let params = Some(ObjectStoreParams { - object_store_wrapper: Some(Arc::new(CountingStoreWrapper(existing_count.clone()))), + object_store_wrapper: Some(Arc::new(OrderedStoreWrapper { + name: "caller", + order: order.clone(), + })), ..Default::default() }); let params = params - .patch_with_store_wrapper(Arc::new(CountingStoreWrapper(added_count.clone()))) + .patch_with_store_wrapper(Arc::new(OrderedStoreWrapper { + name: "compatibility", + order: order.clone(), + })) .unwrap() .unwrap(); params @@ -513,8 +524,7 @@ mod tests { .unwrap() .wrap("memory", Arc::new(InMemory::new()) as Arc); - assert_eq!(existing_count.load(Ordering::Relaxed), 1); - assert_eq!(added_count.load(Ordering::Relaxed), 1); + assert_eq!(*order.lock().unwrap(), vec!["compatibility", "caller"]); } #[test]