From af4fa4a62d024cdac7ad7503af15efc50c397982 Mon Sep 17 00:00:00 2001 From: Krisztian Szucs Date: Thu, 2 Jul 2026 12:24:07 +0200 Subject: [PATCH 1/7] Support prefix-scoped object stores in the registry Register object stores under a path prefix (e.g. `hf://bucket/user/repo`) so multiple stores can coexist under one scheme+authority, and hand each store paths with its registered prefix stripped. This fixes double-prefixing for stores rooted at a prefix (OpenDAL #7786). - Wrap object_store's path-aware registry (register/resolve/deregister) and add `ObjectStoreRegistry::resolve` returning the store plus a store-relative path. - Add `ListingTableUrl::resolve` and rewire the read and schema/partition inference paths to list against the rebased, store-relative url. - Allow `ObjectStoreUrl` to carry a path prefix; key the datafusion-cli stdin store by authority. --- Cargo.lock | 11 +- Cargo.toml | 6 + datafusion-cli/src/exec.rs | 10 +- .../src/object_storage/instrumented.rs | 7 + datafusion-cli/src/object_storage/stdin.rs | 23 +- datafusion/catalog-listing/src/options.rs | 10 +- datafusion/catalog-listing/src/table.rs | 52 +++- .../core/src/datasource/listing/table.rs | 13 +- .../tests/datasource/object_store_access.rs | 143 +++++++++++ datafusion/datasource/src/url.rs | 131 +++++++++- datafusion/execution/src/object_store.rs | 234 +++++++++++------- 11 files changed, 512 insertions(+), 128 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 5b43435ec0a7b..9576fa501e7dc 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4333,8 +4333,7 @@ dependencies = [ [[package]] name = "object_store" version = "0.13.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "622acbc9100d3c10e2ee15804b0caa40e55c933d5aa53814cd520805b7958a49" +source = "git+https://github.com/kszucs/arrow-rs-object-store?branch=registry-deregister-0.13#609e65ab0181a55a429c1f16238fb2273a9b91bf" dependencies = [ "async-trait", "base64 0.22.1", @@ -4834,7 +4833,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "03da047801ff44bb6a4d407d4860c05fd70bb81714e6b2f3812603d5b145b042" dependencies = [ "heck", - "itertools 0.14.0", + "itertools 0.13.0", "log", "multimap", "petgraph", @@ -4853,7 +4852,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b570b25f7617e43d59005d0990ccb79e950a423952cea19671b7a876da390adf" dependencies = [ "anyhow", - "itertools 0.14.0", + "itertools 0.13.0", "proc-macro2", "quote", "syn", @@ -4889,9 +4888,9 @@ dependencies = [ [[package]] name = "quick-xml" -version = "0.39.2" +version = "0.39.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "958f21e8e7ceb5a1aa7fa87fab28e7c75976e0bfe7e23ff069e0a260f894067d" +checksum = "cdcc8dd4e2f670d309a5f0e83fe36dfdc05af317008fea29144da1a2ac858e5e" dependencies = [ "memchr", "serde", diff --git a/Cargo.toml b/Cargo.toml index 0bfaad9a68b3e..bc2148a4c4160 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -268,6 +268,12 @@ codegen-units = 16 lto = false strip = false # Retain debug info for flamegraphs +# Pin object_store to arrow-rs-object-store PR #784 (adds a `deregister` method to +# the registry). Uses the 0.13-compatible branch so the patch applies to `parquet` +# (which requires object_store ^0.13). Drop once the PR is released. +[patch.crates-io] +object_store = { git = "https://github.com/kszucs/arrow-rs-object-store", branch = "registry-deregister-0.13" } + [profile.profiling] inherits = "release" debug = true diff --git a/datafusion-cli/src/exec.rs b/datafusion-cli/src/exec.rs index f43854821b2d5..a7c936c56b91b 100644 --- a/datafusion-cli/src/exec.rs +++ b/datafusion-cli/src/exec.rs @@ -513,8 +513,14 @@ pub(crate) async fn register_object_store_and_config_extensions( ) .await?; - // Register the retrieved object store in the session context's runtime environment - ctx.register_object_store(url, store); + // Register the retrieved object store in the session context's runtime + // environment. The stdin store is shared across all `stdin://` object paths, + // so it is keyed by scheme/authority rather than by the object path. + if scheme == StdinUtils::SCHEME { + ctx.register_object_store(&StdinUtils::store_url(url), store); + } else { + ctx.register_object_store(url, store); + } Ok(()) } diff --git a/datafusion-cli/src/object_storage/instrumented.rs b/datafusion-cli/src/object_storage/instrumented.rs index a0321cacb374b..7e4c31c55a4b4 100644 --- a/datafusion-cli/src/object_storage/instrumented.rs +++ b/datafusion-cli/src/object_storage/instrumented.rs @@ -826,6 +826,13 @@ impl ObjectStoreRegistry for InstrumentedObjectStoreRegistry { fn get_store(&self, url: &Url) -> datafusion::common::Result> { self.inner.get_store(url) } + + fn resolve( + &self, + url: &Url, + ) -> datafusion::common::Result<(Arc, Path)> { + self.inner.resolve(url) + } } #[cfg(test)] diff --git a/datafusion-cli/src/object_storage/stdin.rs b/datafusion-cli/src/object_storage/stdin.rs index 9e63068f4f32f..8f28214fca124 100644 --- a/datafusion-cli/src/object_storage/stdin.rs +++ b/datafusion-cli/src/object_storage/stdin.rs @@ -99,6 +99,17 @@ impl StdinUtils { format!("{}:///{object_name}", Self::SCHEME) } + /// The URL under which the shared stdin store is registered: scheme + + /// authority only (no object path). All `stdin://` object paths resolve to + /// this single store. + pub(crate) fn store_url(url: &Url) -> Url { + let mut base = url.clone(); + base.set_path("/"); + base.set_query(None); + base.set_fragment(None); + base + } + /// Returns the object store backing the `stdin://` scheme, buffering all of /// standard input when the store is first constructed and reusing that /// buffer for any subsequent `stdin://` table created in the same session. @@ -119,7 +130,11 @@ impl StdinUtils { state: &SessionState, url: &Url, ) -> Result> { - let Ok(existing) = state.runtime_env().object_store_registry.get_store(url) + // Every `stdin://` URL shares a single buffered store, keyed by + // scheme/authority (the object-store registry matches on path prefixes, + // so the store must be registered at the authority, not the object path). + let base = Self::store_url(url); + let Ok(existing) = state.runtime_env().object_store_registry.get_store(&base) else { return Self::object_store(state, url).await; }; @@ -252,7 +267,7 @@ mod tests { let store = StdinUtils::in_memory_object_store(&url, data).await?; let ctx = SessionContext::new(); - ctx.register_object_store(&url, store); + ctx.register_object_store(&StdinUtils::store_url(&url), store); ctx.sql(&format!( "CREATE EXTERNAL TABLE t STORED AS {stored_as} LOCATION '{location}' {options}" )) @@ -281,7 +296,7 @@ mod tests { buffered.put(&path, b"a\n1\n2\n".to_vec().into()).await?; let ctx = SessionContext::new(); - ctx.register_object_store(&url, Arc::clone(&buffered)); + ctx.register_object_store(&StdinUtils::store_url(&url), Arc::clone(&buffered)); let reused = StdinUtils::get_or_create(&ctx.state(), &url).await?; assert!( @@ -304,7 +319,7 @@ mod tests { StdinUtils::in_memory_object_store(&csv_url, b"a\n1\n".to_vec()).await?; let ctx = SessionContext::new(); - ctx.register_object_store(&csv_url, store); + ctx.register_object_store(&StdinUtils::store_url(&csv_url), store); let json_url = Url::parse("stdin:///stdin.json").unwrap(); let err = StdinUtils::get_or_create(&ctx.state(), &json_url) diff --git a/datafusion/catalog-listing/src/options.rs b/datafusion/catalog-listing/src/options.rs index 44337e52a1e05..3634e1e78fc9e 100644 --- a/datafusion/catalog-listing/src/options.rs +++ b/datafusion/catalog-listing/src/options.rs @@ -277,7 +277,10 @@ impl ListingOptions { state: &dyn Session, table_path: &'a ListingTableUrl, ) -> datafusion_common::Result { - let store = state.runtime_env().object_store(table_path)?; + // Resolve through the registry so a store registered under a path prefix + // receives store-relative paths (see `ListingTableUrl::resolve`). + let (store, _, table_path) = + table_path.resolve(state.runtime_env().object_store_registry.as_ref())?; let all_files: Vec<_> = table_path .list_all_files(state, store.as_ref(), &self.file_extension) @@ -369,7 +372,10 @@ impl ListingOptions { state: &dyn Session, table_path: &ListingTableUrl, ) -> datafusion_common::Result> { - let store = state.runtime_env().object_store(table_path)?; + // Resolve through the registry so a store registered under a path prefix + // receives store-relative paths (see `ListingTableUrl::resolve`). + let (store, _, table_path) = + table_path.resolve(state.runtime_env().object_store_registry.as_ref())?; // only use 10 files for inference // This can fail to detect inconsistent partition keys diff --git a/datafusion/catalog-listing/src/table.rs b/datafusion/catalog-listing/src/table.rs index 632b829b161a0..fb53ce46705fc 100644 --- a/datafusion/catalog-listing/src/table.rs +++ b/datafusion/catalog-listing/src/table.rs @@ -40,6 +40,7 @@ use datafusion_datasource::{ use datafusion_execution::cache::cache_manager::{ CachedFileMetadata, FileStatisticsCache, SchemaFingerprint, TableScopedPath, }; +use datafusion_execution::object_store::ObjectStoreUrl; use datafusion_expr::dml::InsertOp; use datafusion_expr::execution_props::ExecutionProps; use datafusion_expr::{ @@ -55,6 +56,12 @@ use object_store::ObjectStore; use std::collections::HashMap; use std::sync::Arc; +/// Result of resolving a table's paths against the object store registry: the +/// store serving the table, the base [`ObjectStoreUrl`] (including any registered +/// path prefix) to use as the scan key, and the table paths rebased into the +/// store's coordinate space. +type ResolvedTablePaths = (Arc, ObjectStoreUrl, Vec); + /// Result of a file listing operation from [`ListingTable::list_files_for_scan`]. #[derive(Debug)] pub struct ListFilesResult { @@ -627,9 +634,7 @@ impl TableProvider for ListingTable { None }; - let Some(object_store_url) = - self.table_paths.first().map(ListingTableUrl::object_store) - else { + let Some((_, object_store_url, _)) = self.resolve_object_store(state)? else { return Ok(ScanResult::new(Arc::new(EmptyExec::new(Arc::new( Schema::empty(), ))))); @@ -785,15 +790,39 @@ impl ListingTable { } } + /// Resolve the object store for this table's paths via the registry. + /// + /// Returns the [`ObjectStore`] serving the table, the [`ObjectStoreUrl`] (base + /// url including any registered path prefix) to use as the scan key, and the + /// table paths rebased into the store's coordinate space (registered prefix + /// stripped). Returns `None` when the table has no paths. + fn resolve_object_store( + &self, + ctx: &dyn Session, + ) -> datafusion_common::Result> { + let Some(first) = self.table_paths.first() else { + return Ok(None); + }; + let registry = Arc::clone(&ctx.runtime_env().object_store_registry); + let (store, object_store_url, _) = first.resolve(registry.as_ref())?; + let rebased = self + .table_paths + .iter() + .map(|p| Ok(p.resolve(registry.as_ref())?.2)) + .collect::>>()?; + Ok(Some((store, object_store_url, rebased))) + } + async fn collect_files_for_scan<'a>( &'a self, ctx: &'a dyn Session, store: &'a Arc, + table_paths: &'a [ListingTableUrl], listing_time_filters: &'a [Expr], file_limit: Option, ) -> datafusion_common::Result<(FileGroup, bool)> { // list files (with partitions) - let file_list = future::try_join_all(self.table_paths.iter().map(|table_path| { + let file_list = future::try_join_all(table_paths.iter().map(|table_path| { pruned_partition_list( ctx, store.as_ref(), @@ -840,9 +869,7 @@ impl ListingTable { ); } - let store = if let Some(url) = self.table_paths.first() { - ctx.runtime_env().object_store(url)? - } else { + let Some((store, _, rebased_paths)) = self.resolve_object_store(ctx)? else { return Ok(ListFilesResult { file_groups: vec![], statistics: Statistics::new_unknown(&self.file_schema), @@ -850,7 +877,7 @@ impl ListingTable { }); }; let (file_group, inexact_stats) = self - .collect_files_for_scan(ctx, &store, filters, limit) + .collect_files_for_scan(ctx, &store, &rebased_paths, filters, limit) .await?; // Threshold: 0 = disabled, N > 0 = enabled when distinct_keys >= N @@ -902,17 +929,16 @@ impl ListingTable { ); } - let store = if let Some(url) = self.table_paths.first() { - ctx.runtime_env().object_store(url)? - } else { + let Some((store, _, rebased_paths)) = self.resolve_object_store(ctx)? else { return Ok(ListFilesResult { file_groups: vec![], statistics: Statistics::new_unknown(&self.file_schema), grouped_by_partition: false, }); }; - let (file_group, inexact_stats) = - self.collect_files_for_scan(ctx, &store, &[], None).await?; + let (file_group, inexact_stats) = self + .collect_files_for_scan(ctx, &store, &rebased_paths, &[], None) + .await?; let mut file_groups = file_group.split_files(file_group_count); if !file_groups.is_empty() { file_groups.resize_with(file_group_count, || FileGroup::new(vec![])); diff --git a/datafusion/core/src/datasource/listing/table.rs b/datafusion/core/src/datasource/listing/table.rs index 39c20f9b786c2..4bf05049519fb 100644 --- a/datafusion/core/src/datasource/listing/table.rs +++ b/datafusion/core/src/datasource/listing/table.rs @@ -51,16 +51,15 @@ impl ListingTableConfigExt for ListingTableConfig { self, state: &dyn Session, ) -> datafusion_common::Result { - let store = if let Some(url) = self.table_paths.first() { - state.runtime_env().object_store(url)? - } else { + let Some(first) = self.table_paths.first() else { return Ok(self); }; + // Resolve through the registry so a store registered under a path prefix + // receives store-relative paths (see `ListingTableUrl::resolve`). + let (store, _, table_path) = + first.resolve(state.runtime_env().object_store_registry.as_ref())?; - let file = self - .table_paths - .first() - .unwrap() + let file = table_path .list_all_files(state, store.as_ref(), "") .await? .next() diff --git a/datafusion/core/tests/datasource/object_store_access.rs b/datafusion/core/tests/datasource/object_store_access.rs index 2503de862e06a..a4f3d3bbb16dc 100644 --- a/datafusion/core/tests/datasource/object_store_access.rs +++ b/datafusion/core/tests/datasource/object_store_access.rs @@ -920,6 +920,149 @@ async fn query_single_parquet_file_multi_row_groups_multiple_predicates() { ); } +/// Build a single-file parquet buffer with two columns `a` and `b`. +fn single_file_parquet_bytes() -> Vec { + let a: ArrayRef = Arc::new(Int32Array::from_iter_values(0..200)); + let b: ArrayRef = Arc::new(Int32Array::from_iter_values(1000..1200)); + let batch = RecordBatch::try_from_iter([("a", a), ("b", b)]).unwrap(); + let mut buffer = vec![]; + let mut writer = + parquet::arrow::ArrowWriter::try_new(&mut buffer, batch.schema(), None).unwrap(); + writer.write(&batch).unwrap(); + writer.close().unwrap(); + buffer +} + +/// Registering a store under a *path prefix* (as an OpenDAL store rooted at a +/// specific repo would be) must not double-prefix the object paths handed to it. +/// The store is registered at `mem://bucket/user/repo` and, like an +/// OpenDAL operator rooted at the repo, expects/serves *store-relative* paths +/// (`data/f.parquet`), never the registered prefix. +/// +/// Regression test for . +#[tokio::test] +async fn prefixed_store_is_not_double_prefixed() { + let store = Arc::new(RequestCountingObjectStore::new()); + + // The store serves store-relative paths: the file lives at `data/f.parquet`, + // NOT at `user/repo/data/f.parquet`. + store + .inner + .put( + &Path::from("data/f.parquet"), + PutPayload::from(single_file_parquet_bytes()), + ) + .await + .unwrap(); + + let ctx = SessionContext::new(); + ctx.runtime_env().register_object_store( + &Url::parse("mem://bucket/user/repo").unwrap(), + Arc::clone(&store) as Arc, + ); + + // Register and query through the fully-qualified, prefixed URL. + ctx.register_parquet( + "t", + "mem://bucket/user/repo/data/f.parquet", + ParquetReadOptions::new(), + ) + .await + .unwrap(); + + let batches = ctx + .sql("SELECT count(*) as n FROM t") + .await + .unwrap() + .collect() + .await + .unwrap(); + let count = batches[0] + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .value(0); + assert_eq!(count, 200, "query did not read the file successfully"); + + // The store must only ever see store-relative paths. If double-prefixing + // regressed, it would see `user/repo/data/f.parquet` instead. + let requests = store.recorded_requests(); + assert!(!requests.is_empty(), "store received no requests"); + for req in &requests { + let rendered = format!("{req}"); + assert!( + !rendered.contains("user/repo"), + "store received a double-prefixed path: {rendered}" + ); + } +} + +/// Two stores registered under the *same authority* at different path prefixes +/// must coexist, each resolving to its own store. +#[tokio::test] +async fn two_prefixed_stores_coexist_under_one_authority() { + async fn make_repo() -> Arc { + let store = Arc::new(RequestCountingObjectStore::new()); + store + .inner + .put( + &Path::from("data/f.parquet"), + PutPayload::from(single_file_parquet_bytes()), + ) + .await + .unwrap(); + store + } + + let repo1 = make_repo().await; + let repo2 = make_repo().await; + + let ctx = SessionContext::new(); + ctx.runtime_env().register_object_store( + &Url::parse("mem://bucket/userA/repo1").unwrap(), + Arc::clone(&repo1) as Arc, + ); + ctx.runtime_env().register_object_store( + &Url::parse("mem://bucket/userB/repo2").unwrap(), + Arc::clone(&repo2) as Arc, + ); + + ctx.register_parquet( + "t1", + "mem://bucket/userA/repo1/data/f.parquet", + ParquetReadOptions::new(), + ) + .await + .unwrap(); + ctx.register_parquet( + "t2", + "mem://bucket/userB/repo2/data/f.parquet", + ParquetReadOptions::new(), + ) + .await + .unwrap(); + + // Each query must have gone to its own store, and neither should be + // double-prefixed. + repo1.clear_requests(); + repo2.clear_requests(); + ctx.sql("SELECT count(*) FROM t1") + .await + .unwrap() + .collect() + .await + .unwrap(); + assert!( + !repo1.recorded_requests().is_empty(), + "repo1 was not queried" + ); + assert!( + repo2.recorded_requests().is_empty(), + "repo2 was queried for a repo1 table" + ); +} + /// Runs tests with a request counting object store struct Test { object_store: Arc, diff --git a/datafusion/datasource/src/url.rs b/datafusion/datasource/src/url.rs index 7985a29e4fd94..a0c6b6e5d9fdf 100644 --- a/datafusion/datasource/src/url.rs +++ b/datafusion/datasource/src/url.rs @@ -20,7 +20,7 @@ use std::sync::Arc; use datafusion_common::{DataFusionError, Result, TableReference}; use datafusion_execution::cache::cache_manager::CachedFileList; use datafusion_execution::cache::cache_manager::TableScopedPath; -use datafusion_execution::object_store::ObjectStoreUrl; +use datafusion_execution::object_store::{ObjectStoreRegistry, ObjectStoreUrl}; use datafusion_session::Session; use futures::stream::BoxStream; @@ -320,11 +320,75 @@ impl ListingTableUrl { } /// Return the [`ObjectStoreUrl`] for this [`ListingTableUrl`] + /// + /// This returns the scheme + authority only, and does not account for any + /// store registered under a path prefix. For reading, prefer + /// [`Self::resolve`], which consults the registry and strips the registered + /// prefix from the path handed to the store. pub fn object_store(&self) -> ObjectStoreUrl { let url = &self.url[url::Position::BeforeScheme..url::Position::BeforePath]; ObjectStoreUrl::parse(url).unwrap() } + /// Resolve this url against the object store `registry`. + /// + /// Returns: + /// - the [`ObjectStore`] serving this url, + /// - the [`ObjectStoreUrl`] (scheme + authority + any registered path prefix) + /// to use as the scan key so the same store can be re-resolved at execution + /// time, and + /// - a copy of this [`ListingTableUrl`] rebased into the resolved store's + /// coordinate space: its [`prefix`](Self::prefix) has the registered prefix + /// stripped, so listing/`head`/partition parsing operate on paths relative + /// to the store (avoiding double-prefixing for stores rooted at a prefix). + /// + /// For a store registered by scheme + authority only (e.g. `s3://bucket`), the + /// rebased prefix is identical to [`Self::prefix`] and the behavior is + /// unchanged. + pub fn resolve( + &self, + registry: &dyn ObjectStoreRegistry, + ) -> Result<(Arc, ObjectStoreUrl, ListingTableUrl)> { + let (store, store_relative_prefix) = registry.resolve(&self.url)?; + + // The registered prefix is the leading path segments that the registry + // stripped: full prefix segments minus the store-relative ones. + let stripped = self + .prefix + .parts() + .count() + .saturating_sub(store_relative_prefix.parts().count()); + + // Rebuild the base URL (scheme + authority + registered prefix) preserving + // the original percent-encoding of the retained segments. + let authority = &self.url[url::Position::BeforeScheme..url::Position::BeforePath]; + let kept: Vec<&str> = self + .url + .path() + .split('/') + .filter(|s| !s.is_empty()) + .take(stripped) + .collect(); + let base_str = if kept.is_empty() { + authority.to_string() + } else { + format!("{authority}/{}", kept.join("/")) + }; + let object_store_url = ObjectStoreUrl::parse(&base_str)?; + + // The rebased url keeps the original `url` (used only for scheme / + // is_collection / display); only `prefix` becomes store-relative, which is + // what all listing and partition-parsing logic keys off. + let rebased = ListingTableUrl { + url: self.url.clone(), + prefix: store_relative_prefix, + glob: self.glob.clone(), + table_ref: self.table_ref.clone(), + }; + + Ok((store, object_store_url, rebased)) + } + /// Returns true if the [`ListingTableUrl`] points to the folder pub fn is_folder(&self) -> bool { self.url.scheme() == "file" && self.is_collection() @@ -515,6 +579,7 @@ mod tests { use datafusion_common::config::TableOptions; use datafusion_execution::TaskContext; use datafusion_execution::config::SessionConfig; + use datafusion_execution::object_store::DefaultObjectStoreRegistry; use datafusion_execution::runtime_env::RuntimeEnv; use datafusion_expr::execution_props::ExecutionProps; use datafusion_expr::registry::ExtensionTypeRegistryRef; @@ -629,6 +694,70 @@ mod tests { assert!(url.strip_prefix(&path).is_none()); } + fn registry_with(base: &str) -> DefaultObjectStoreRegistry { + let registry = DefaultObjectStoreRegistry::new(); + registry.register_store( + &Url::parse(base).unwrap(), + Arc::new(object_store::memory::InMemory::new()), + ); + registry + } + + #[test] + fn test_resolve_authority_only_store() { + // A store registered by scheme+authority only leaves the prefix unchanged. + let registry = registry_with("s3://bucket"); + let url = ListingTableUrl::parse("s3://bucket/foo/bar/file.parquet").unwrap(); + let (_, base, rebased) = url.resolve(®istry).unwrap(); + assert_eq!(base.as_str(), "s3://bucket/"); + assert_eq!(rebased.prefix().as_ref(), "foo/bar/file.parquet"); + assert_eq!(rebased.prefix(), url.prefix()); + } + + #[test] + fn test_resolve_prefixed_store_strips_prefix() { + // A store registered under a path prefix: the prefix is stripped from the + // rebased url and retained in the base ObjectStoreUrl (the scan key). + let registry = registry_with("s3://bucket/user/repo"); + let url = + ListingTableUrl::parse("s3://bucket/user/repo/data/file.parquet").unwrap(); + let (_, base, rebased) = url.resolve(®istry).unwrap(); + assert_eq!(base.as_str(), "s3://bucket/user/repo"); + assert_eq!(rebased.prefix().as_ref(), "data/file.parquet"); + assert!(!rebased.is_collection()); + } + + #[test] + fn test_resolve_prefixed_collection_preserves_trailing_slash() { + // A directory url keeps its collection-ness after rebasing. + let registry = registry_with("s3://bucket/user/repo"); + let url = ListingTableUrl::parse("s3://bucket/user/repo/data/").unwrap(); + let (_, base, rebased) = url.resolve(®istry).unwrap(); + assert_eq!(base.as_str(), "s3://bucket/user/repo"); + assert_eq!(rebased.prefix().as_ref(), "data"); + assert!(rebased.is_collection()); + } + + #[test] + fn test_resolve_query_at_registered_root() { + // Querying exactly the registered prefix yields an empty rebased prefix. + let registry = registry_with("s3://bucket/user/repo"); + let url = ListingTableUrl::parse("s3://bucket/user/repo/").unwrap(); + let (_, base, rebased) = url.resolve(®istry).unwrap(); + assert_eq!(base.as_str(), "s3://bucket/user/repo"); + assert_eq!(rebased.prefix().as_ref(), ""); + } + + #[test] + fn test_resolve_preserves_percent_encoding_in_base() { + // The retained (registered-prefix) segments keep their original encoding. + let registry = registry_with("s3://bucket/us%20er"); + let url = ListingTableUrl::parse("s3://bucket/us%20er/data.parquet").unwrap(); + let (_, base, rebased) = url.resolve(®istry).unwrap(); + assert_eq!(base.as_str(), "s3://bucket/us%20er"); + assert_eq!(rebased.prefix().as_ref(), "data.parquet"); + } + #[test] fn test_split_glob() { fn test(input: &str, expected: Option<(&str, &str)>) { diff --git a/datafusion/execution/src/object_store.rs b/datafusion/execution/src/object_store.rs index 22ce1f0cf2bbf..28f866134e188 100644 --- a/datafusion/execution/src/object_store.rs +++ b/datafusion/execution/src/object_store.rs @@ -19,13 +19,14 @@ //! This allows the user to extend DataFusion with different storage systems such as S3 or HDFS //! and query data inside these systems. -use dashmap::DashMap; use datafusion_common::{ DataFusionError, Result, exec_err, internal_datafusion_err, not_impl_err, }; use object_store::ObjectStore; #[cfg(not(target_arch = "wasm32"))] use object_store::local::LocalFileSystem; +use object_store::path::Path; +use object_store::registry::ObjectStoreRegistry as UpstreamObjectStoreRegistry; use std::sync::Arc; use url::Url; @@ -59,15 +60,20 @@ impl ObjectStoreUrl { let mut parsed = Url::parse(s.as_ref()).map_err(|e| DataFusionError::External(Box::new(e)))?; - let remaining = &parsed[url::Position::BeforePath..]; - if !remaining.is_empty() && remaining != "/" { + // An ObjectStoreUrl identifies a store, optionally including a registered + // path prefix (e.g. `hf://bucket/user/repo`). Query strings and fragments + // are not part of a store key and are rejected. + if parsed.query().is_some() || parsed.fragment().is_some() { + let remaining = &parsed[url::Position::AfterPath..]; return exec_err!( - "ObjectStoreUrl must only contain scheme and authority, got: {remaining}" + "ObjectStoreUrl must not contain a query or fragment, got: {remaining}" ); } - // Always set path for consistency - parsed.set_path("/"); + // Normalize an empty path to "/" for consistency. + if parsed.path().is_empty() { + parsed.set_path("/"); + } Ok(Self { url: parsed }) } @@ -175,102 +181,89 @@ pub trait ObjectStoreRegistry: Send + Sync + std::fmt::Debug + 'static { /// the `url` and [`ObjectStoreRegistry`] implementation. An [`ObjectStore`] may be lazily /// created and registered. fn get_store(&self, url: &Url) -> Result>; -} - -/// The default [`ObjectStoreRegistry`] -pub struct DefaultObjectStoreRegistry { - /// A map from scheme to object store that serve list / read operations for the store - object_stores: DashMap>, -} -impl std::fmt::Debug for DefaultObjectStoreRegistry { - fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { - f.debug_struct("DefaultObjectStoreRegistry") - .field( - "schemes", - &self - .object_stores - .iter() - .map(|o| o.key().clone()) - .collect::>(), - ) - .finish() + /// Resolve `url` to an [`ObjectStore`] together with the object [`Path`] + /// *relative to that store*. + /// + /// Unlike [`Self::get_store`], this allows a store to be registered under a + /// path prefix (e.g. `hf://bucket/user/repo`). When resolving a longer URL + /// (e.g. `hf://bucket/user/repo/data/f.parquet`) the registered prefix is + /// stripped from the returned path (`data/f.parquet`), so a store rooted at + /// that prefix is not prefixed twice. The store with the longest matching + /// path prefix is returned. + /// + /// The default implementation preserves the legacy behavior of matching on + /// scheme + authority only, returning the full URL path unchanged. + fn resolve(&self, url: &Url) -> Result<(Arc, Path)> { + let store = self.get_store(url)?; + let path = Path::from_url_path(url.path()) + .map_err(|e| internal_datafusion_err!("Invalid object store path: {e}"))?; + Ok((store, path)) } } -impl Default for DefaultObjectStoreRegistry { - fn default() -> Self { - Self::new() - } +/// The default [`ObjectStoreRegistry`]. +/// +/// This is a thin adapter over [`object_store::registry::DefaultObjectStoreRegistry`], +/// which performs path-segment based longest-prefix matching. This means stores can +/// be registered under a path prefix and coexist under a single scheme + authority: +/// +/// - `file:///my_path` returns the local filesystem store +/// - `s3://bucket/path` returns a store registered with `s3://bucket` (authority only) +/// - `hf://bucket/user/repo/data/f.parquet` returns a store registered with +/// `hf://bucket/user/repo`, and [`ObjectStoreRegistry::resolve`] strips the +/// registered prefix from the returned path (`data/f.parquet`) +/// +/// Stores for known schemes (`file`, `s3`, `gs`, `az`, `http`, ...) are created +/// lazily on first access via [`object_store::parse_url_opts`]. +#[derive(Debug, Default)] +pub struct DefaultObjectStoreRegistry { + inner: object_store::registry::DefaultObjectStoreRegistry, } impl DefaultObjectStoreRegistry { /// This will register [`LocalFileSystem`] to handle `file://` paths #[cfg(not(target_arch = "wasm32"))] pub fn new() -> Self { - let object_stores: DashMap> = DashMap::new(); - object_stores.insert("file://".to_string(), Arc::new(LocalFileSystem::new())); - Self { object_stores } + let inner = object_store::registry::DefaultObjectStoreRegistry::new(); + inner.register( + Url::parse("file://").unwrap(), + Arc::new(LocalFileSystem::new()), + ); + Self { inner } } /// Default without any backend registered. #[cfg(target_arch = "wasm32")] pub fn new() -> Self { - let object_stores: DashMap> = DashMap::new(); - Self { object_stores } + Self::default() } } -/// -/// Stores are registered based on the scheme, host and port of the provided URL -/// with a [`LocalFileSystem::new`] automatically registered for `file://` (if the -/// target arch is not `wasm32`). -/// -/// For example: -/// -/// - `file:///my_path` will return the default LocalFS store -/// - `s3://bucket/path` will return a store registered with `s3://bucket` if any -/// - `hdfs://host:port/path` will return a store registered with `hdfs://host:port` if any impl ObjectStoreRegistry for DefaultObjectStoreRegistry { fn register_store( &self, url: &Url, store: Arc, ) -> Option> { - let s = get_url_key(url); - self.object_stores.insert(s, store) + self.inner.register(url.clone(), store) } fn deregister_store(&self, url: &Url) -> Result> { - let s = get_url_key(url); - let (_, object_store) = self.object_stores - .remove(&s) - .ok_or_else(|| { - internal_datafusion_err!("Failed to deregister object store. No suitable object store found for {url}. See `RuntimeEnv::register_object_store`") - })?; - - Ok(object_store) + self.inner.deregister(url).ok_or_else(|| { + internal_datafusion_err!("Failed to deregister object store. No suitable object store found for {url}. See `RuntimeEnv::register_object_store`") + }) } fn get_store(&self, url: &Url) -> Result> { - let s = get_url_key(url); - self.object_stores - .get(&s) - .map(|o| Arc::clone(o.value())) - .ok_or_else(|| { - internal_datafusion_err!("No suitable object store found for {url}. See `RuntimeEnv::register_object_store`") - }) + self.resolve(url).map(|(store, _)| store) } -} -/// Get the key of a url for object store registration. -/// The credential info will be removed -fn get_url_key(url: &Url) -> String { - format!( - "{}://{}", - url.scheme(), - &url[url::Position::BeforeHost..url::Position::AfterPort], - ) + fn resolve(&self, url: &Url) -> Result<(Arc, Path)> { + self.inner + .resolve(url) + .map_err(|e| DataFusionError::External(Box::new(e))) + } } #[cfg(test)] @@ -291,44 +284,99 @@ mod tests { let err = ObjectStoreUrl::parse("s3://bucket:invalid").unwrap_err(); assert_eq!(err.strip_backtrace(), "External error: invalid port number"); + // A path prefix is now allowed: it identifies a store registered under + // that prefix (e.g. an OpenDAL store rooted at a specific repo). + let url = ObjectStoreUrl::parse("s3://host:123/foo").unwrap(); + assert_eq!(url.as_str(), "s3://host:123/foo"); + + let url = + ObjectStoreUrl::parse("s3://username:password@host:123/foo/bar").unwrap(); + assert_eq!(url.as_str(), "s3://username:password@host:123/foo/bar"); + + // Query strings and fragments are still rejected. let err = ObjectStoreUrl::parse("s3://bucket?").unwrap_err(); assert_eq!( err.strip_backtrace(), - "Execution error: ObjectStoreUrl must only contain scheme and authority, got: ?" + "Execution error: ObjectStoreUrl must not contain a query or fragment, got: ?" ); let err = ObjectStoreUrl::parse("s3://bucket?foo=bar").unwrap_err(); assert_eq!( err.strip_backtrace(), - "Execution error: ObjectStoreUrl must only contain scheme and authority, got: ?foo=bar" + "Execution error: ObjectStoreUrl must not contain a query or fragment, got: ?foo=bar" ); + } - let err = ObjectStoreUrl::parse("s3://host:123/foo").unwrap_err(); - assert_eq!( - err.strip_backtrace(), - "Execution error: ObjectStoreUrl must only contain scheme and authority, got: /foo" + #[test] + fn test_registry_resolves_prefixed_store() { + use object_store::memory::InMemory; + use object_store::prefix::PrefixStore; + + let registry = DefaultObjectStoreRegistry::new(); + + // Register two stores under the SAME authority at different prefixes. + let repo1 = Arc::new(PrefixStore::new(InMemory::new(), "userA/repo1")) + as Arc; + let repo2 = Arc::new(PrefixStore::new(InMemory::new(), "userB/repo2")) + as Arc; + registry.register_store( + &Url::parse("hf://bucket/userA/repo1").unwrap(), + Arc::clone(&repo1), + ); + registry.register_store( + &Url::parse("hf://bucket/userB/repo2").unwrap(), + Arc::clone(&repo2), ); - let err = - ObjectStoreUrl::parse("s3://username:password@host:123/foo").unwrap_err(); - assert_eq!( - err.strip_backtrace(), - "Execution error: ObjectStoreUrl must only contain scheme and authority, got: /foo" + // Each query resolves to its own store, with the registered prefix stripped. + let (store, path) = registry + .resolve(&Url::parse("hf://bucket/userA/repo1/data/f.parquet").unwrap()) + .unwrap(); + assert_eq!(path.as_ref(), "data/f.parquet"); + assert!(Arc::ptr_eq(&store, &repo1)); + + let (store, path) = registry + .resolve(&Url::parse("hf://bucket/userB/repo2/x/y.parquet").unwrap()) + .unwrap(); + assert_eq!(path.as_ref(), "x/y.parquet"); + assert!(Arc::ptr_eq(&store, &repo2)); + + // Deregister repo1 and confirm it is gone while repo2 survives. + registry + .deregister_store(&Url::parse("hf://bucket/userA/repo1").unwrap()) + .unwrap(); + assert!( + registry + .resolve(&Url::parse("hf://bucket/userA/repo1/data/f.parquet").unwrap()) + .is_err() + ); + let (store, _) = registry + .resolve(&Url::parse("hf://bucket/userB/repo2/x/y.parquet").unwrap()) + .unwrap(); + assert!(Arc::ptr_eq(&store, &repo2)); + + // Deregistering a URL that was never registered is an error. + assert!( + registry + .deregister_store(&Url::parse("hf://bucket/nope").unwrap()) + .is_err() ); } #[test] - fn test_get_url_key() { - let file = ObjectStoreUrl::parse("file://").unwrap(); - let key = get_url_key(&file.url); - assert_eq!(key.as_str(), "file://"); - - let url = ObjectStoreUrl::parse("s3://bucket").unwrap(); - let key = get_url_key(&url.url); - assert_eq!(key.as_str(), "s3://bucket"); - - let url = ObjectStoreUrl::parse("s3://username:password@host:123").unwrap(); - let key = get_url_key(&url.url); - assert_eq!(key.as_str(), "s3://host:123"); + fn test_registry_authority_only_store() { + use object_store::memory::InMemory; + + // A store registered by scheme+authority only serves the whole authority; + // the full path is returned unchanged (the legacy behavior). + let registry = DefaultObjectStoreRegistry::new(); + let store = Arc::new(InMemory::new()) as Arc; + registry.register_store(&Url::parse("s3://bucket").unwrap(), Arc::clone(&store)); + + let (resolved, path) = registry + .resolve(&Url::parse("s3://bucket/a/b/c.parquet").unwrap()) + .unwrap(); + assert_eq!(path.as_ref(), "a/b/c.parquet"); + assert!(Arc::ptr_eq(&resolved, &store)); } } From 09daae8b6916fab82965991eddcdc63f636118d3 Mon Sep 17 00:00:00 2001 From: Krisztian Szucs Date: Thu, 2 Jul 2026 15:39:13 +0200 Subject: [PATCH 2/7] Key the stdin store by its object_store() identity Remove the bespoke StdinUtils::store_url helper. The stdin store is a single per-session store identified by the stdin:// authority, which is exactly what ListingTableUrl::object_store() returns, so register under that identity and let the registry resolve object paths to it. --- datafusion-cli/src/exec.rs | 5 +-- datafusion-cli/src/object_storage/stdin.rs | 36 +++++++++------------- 2 files changed, 17 insertions(+), 24 deletions(-) diff --git a/datafusion-cli/src/exec.rs b/datafusion-cli/src/exec.rs index a7c936c56b91b..e225e12d50a34 100644 --- a/datafusion-cli/src/exec.rs +++ b/datafusion-cli/src/exec.rs @@ -515,9 +515,10 @@ pub(crate) async fn register_object_store_and_config_extensions( // Register the retrieved object store in the session context's runtime // environment. The stdin store is shared across all `stdin://` object paths, - // so it is keyed by scheme/authority rather than by the object path. + // so it is keyed by scheme/authority (its `object_store()` identity) rather + // than by the object path. if scheme == StdinUtils::SCHEME { - ctx.register_object_store(&StdinUtils::store_url(url), store); + ctx.register_object_store(table_path.object_store().as_ref(), store); } else { ctx.register_object_store(url, store); } diff --git a/datafusion-cli/src/object_storage/stdin.rs b/datafusion-cli/src/object_storage/stdin.rs index 8f28214fca124..f9c03cf5ee06c 100644 --- a/datafusion-cli/src/object_storage/stdin.rs +++ b/datafusion-cli/src/object_storage/stdin.rs @@ -99,24 +99,14 @@ impl StdinUtils { format!("{}:///{object_name}", Self::SCHEME) } - /// The URL under which the shared stdin store is registered: scheme + - /// authority only (no object path). All `stdin://` object paths resolve to - /// this single store. - pub(crate) fn store_url(url: &Url) -> Url { - let mut base = url.clone(); - base.set_path("/"); - base.set_query(None); - base.set_fragment(None); - base - } - /// Returns the object store backing the `stdin://` scheme, buffering all of /// standard input when the store is first constructed and reusing that /// buffer for any subsequent `stdin://` table created in the same session. /// - /// stdin is a one-shot stream: it can only be read once. The object store - /// registry keys by scheme/authority, so every `stdin://` URL maps to the - /// same store. Without this guard, a second `CREATE EXTERNAL TABLE ... + /// stdin is a one-shot stream: it can only be read once. The stdin store is + /// registered at the `stdin://` authority, so every `stdin://` object path + /// resolves to the same store. Without this guard, a second `CREATE EXTERNAL + /// TABLE ... /// LOCATION '/dev/stdin'` would re-read (now-EOF) stdin, build an empty /// store, and overwrite the populated one, silently emptying the earlier /// table. Reusing the already-registered store avoids that. @@ -130,11 +120,10 @@ impl StdinUtils { state: &SessionState, url: &Url, ) -> Result> { - // Every `stdin://` URL shares a single buffered store, keyed by - // scheme/authority (the object-store registry matches on path prefixes, - // so the store must be registered at the authority, not the object path). - let base = Self::store_url(url); - let Ok(existing) = state.runtime_env().object_store_registry.get_store(&base) + // Every `stdin://` object path is served by the single per-session stdin + // store, registered at the `stdin://` authority; the registry resolves + // this url to it. + let Ok(existing) = state.runtime_env().object_store_registry.get_store(url) else { return Self::object_store(state, url).await; }; @@ -267,7 +256,7 @@ mod tests { let store = StdinUtils::in_memory_object_store(&url, data).await?; let ctx = SessionContext::new(); - ctx.register_object_store(&StdinUtils::store_url(&url), store); + ctx.register_object_store(&Url::parse("stdin:///").unwrap(), store); ctx.sql(&format!( "CREATE EXTERNAL TABLE t STORED AS {stored_as} LOCATION '{location}' {options}" )) @@ -296,7 +285,10 @@ mod tests { buffered.put(&path, b"a\n1\n2\n".to_vec().into()).await?; let ctx = SessionContext::new(); - ctx.register_object_store(&StdinUtils::store_url(&url), Arc::clone(&buffered)); + ctx.register_object_store( + &Url::parse("stdin:///").unwrap(), + Arc::clone(&buffered), + ); let reused = StdinUtils::get_or_create(&ctx.state(), &url).await?; assert!( @@ -319,7 +311,7 @@ mod tests { StdinUtils::in_memory_object_store(&csv_url, b"a\n1\n".to_vec()).await?; let ctx = SessionContext::new(); - ctx.register_object_store(&StdinUtils::store_url(&csv_url), store); + ctx.register_object_store(&Url::parse("stdin:///").unwrap(), store); let json_url = Url::parse("stdin:///stdin.json").unwrap(); let err = StdinUtils::get_or_create(&ctx.state(), &json_url) From 1c2cbe7eb99744aa81dd7d65e102ffa58c755856 Mon Sep 17 00:00:00 2001 From: Krisztian Szucs Date: Fri, 3 Jul 2026 12:42:37 +0200 Subject: [PATCH 3/7] Resolve object stores via ListingTableUrl::resolve Replace the prefix-blind ListingTableUrl::object_store() with a single registry-aware ListingTableUrl::resolve() that returns the serving store, the store identity (scheme + authority + registered prefix), and the url rebased into the store's coordinate space (ResolvedTableUrl). The registry resolve() returns (store, store-relative path); the identity is derived from the url and that path. This makes both read and write paths correct for stores registered under a path prefix. --- datafusion-cli/src/exec.rs | 10 +- datafusion/catalog-listing/src/options.rs | 26 +-- datafusion/catalog-listing/src/table.rs | 30 +-- .../core/src/datasource/listing/table.rs | 8 +- datafusion/core/src/physical_planner.rs | 9 +- datafusion/core/src/test_util/parquet.rs | 7 +- datafusion/datasource-csv/src/source.rs | 8 +- datafusion/datasource-json/src/source.rs | 8 +- datafusion/datasource-parquet/src/writer.rs | 8 +- datafusion/datasource/src/mod.rs | 11 -- datafusion/datasource/src/url.rs | 177 +++++++++--------- datafusion/execution/src/object_store.rs | 8 +- 12 files changed, 153 insertions(+), 157 deletions(-) diff --git a/datafusion-cli/src/exec.rs b/datafusion-cli/src/exec.rs index e225e12d50a34..c36d3dba51113 100644 --- a/datafusion-cli/src/exec.rs +++ b/datafusion-cli/src/exec.rs @@ -32,6 +32,7 @@ use datafusion::config::{ConfigFileType, Dialect}; use datafusion::datasource::listing::ListingTableUrl; use datafusion::error::{DataFusionError, Result}; use datafusion::execution::memory_pool::MemoryConsumer; +use datafusion::execution::object_store::ObjectStoreUrl; use datafusion::logical_expr::{DdlStatement, LogicalPlan}; use datafusion::physical_plan::execution_plan::EmissionType; use datafusion::physical_plan::spill::get_record_batch_memory_size; @@ -515,10 +516,13 @@ pub(crate) async fn register_object_store_and_config_extensions( // Register the retrieved object store in the session context's runtime // environment. The stdin store is shared across all `stdin://` object paths, - // so it is keyed by scheme/authority (its `object_store()` identity) rather - // than by the object path. + // so it is registered at its scheme/authority identity rather than at the + // object path (the registry then resolves every `stdin://` path to it). if scheme == StdinUtils::SCHEME { - ctx.register_object_store(table_path.object_store().as_ref(), store); + let identity = ObjectStoreUrl::parse( + &url[::url::Position::BeforeScheme..::url::Position::BeforePath], + )?; + ctx.register_object_store(identity.as_ref(), store); } else { ctx.register_object_store(url, store); } diff --git a/datafusion/catalog-listing/src/options.rs b/datafusion/catalog-listing/src/options.rs index 3634e1e78fc9e..61e67222f0278 100644 --- a/datafusion/catalog-listing/src/options.rs +++ b/datafusion/catalog-listing/src/options.rs @@ -279,11 +279,11 @@ impl ListingOptions { ) -> datafusion_common::Result { // Resolve through the registry so a store registered under a path prefix // receives store-relative paths (see `ListingTableUrl::resolve`). - let (store, _, table_path) = - table_path.resolve(state.runtime_env().object_store_registry.as_ref())?; + let resolved = table_path.resolve(state.runtime_env())?; - let all_files: Vec<_> = table_path - .list_all_files(state, store.as_ref(), &self.file_extension) + let all_files: Vec<_> = resolved + .table_url + .list_all_files(state, resolved.store.as_ref(), &self.file_extension) .await? .try_collect() .await?; @@ -293,7 +293,7 @@ impl ListingOptions { "No files found at {}. \ Cannot infer schema from an empty location; either add data files \ or declare an explicit schema for the table.", - table_path + resolved.table_url ); } @@ -303,7 +303,10 @@ impl ListingOptions { .filter(|object_meta| object_meta.size > 0) .collect(); - let schema = self.format.infer_schema(state, &store, &files).await?; + let schema = self + .format + .infer_schema(state, &resolved.store, &files) + .await?; Ok(schema) } @@ -374,21 +377,22 @@ impl ListingOptions { ) -> datafusion_common::Result> { // Resolve through the registry so a store registered under a path prefix // receives store-relative paths (see `ListingTableUrl::resolve`). - let (store, _, table_path) = - table_path.resolve(state.runtime_env().object_store_registry.as_ref())?; + let resolved = table_path.resolve(state.runtime_env())?; // only use 10 files for inference // This can fail to detect inconsistent partition keys // A DFS traversal approach of the store can help here - let files: Vec<_> = table_path - .list_all_files(state, store.as_ref(), &self.file_extension) + let files: Vec<_> = resolved + .table_url + .list_all_files(state, resolved.store.as_ref(), &self.file_extension) .await? .take(10) .try_collect() .await?; let stripped_path_parts = files.iter().map(|file| { - table_path + resolved + .table_url .strip_prefix(&file.location) .unwrap() .collect_vec() diff --git a/datafusion/catalog-listing/src/table.rs b/datafusion/catalog-listing/src/table.rs index fb53ce46705fc..d9d8a9814c3a1 100644 --- a/datafusion/catalog-listing/src/table.rs +++ b/datafusion/catalog-listing/src/table.rs @@ -709,13 +709,21 @@ impl TableProvider for ListingTable { ); } - // Get the object store for the table path. - let store = state.runtime_env().object_store(table_path)?; + // Resolve through the registry so a store registered under a path prefix + // gets store-relative paths (see `ListingTableUrl::with_prefix`). The + // resolved identity and rebased paths are embedded in the sink config so + // the store re-resolves and writes correctly at execution time. + let Some((store, object_store_url, rebased_paths)) = + self.resolve_object_store(state)? + else { + return Err(internal_datafusion_err!("ListingTable has no table paths")); + }; + let rebased = &rebased_paths[0]; let file_list_stream = pruned_partition_list( state, store.as_ref(), - table_path, + rebased, &[], &self.options.file_extension, &self.options.table_partition_cols, @@ -729,8 +737,8 @@ impl TableProvider for ListingTable { // Invalidate cache entries for this table if they exist if let Some(lfc) = state.runtime_env().cache_manager.get_list_files_cache() { let key = TableScopedPath { - table: table_path.get_table_ref().clone(), - path: table_path.prefix().clone(), + table: rebased.get_table_ref().clone(), + path: rebased.prefix().clone(), }; let _ = lfc.remove(&key); } @@ -738,8 +746,8 @@ impl TableProvider for ListingTable { // Sink related option, apart from format let config = FileSinkConfig { original_url: String::default(), - object_store_url: self.table_paths()[0].object_store(), - table_paths: self.table_paths().clone(), + object_store_url, + table_paths: rebased_paths, file_group, output_schema: self.schema(), table_partition_cols: self.options.table_partition_cols.clone(), @@ -803,14 +811,14 @@ impl ListingTable { let Some(first) = self.table_paths.first() else { return Ok(None); }; - let registry = Arc::clone(&ctx.runtime_env().object_store_registry); - let (store, object_store_url, _) = first.resolve(registry.as_ref())?; + let env = ctx.runtime_env(); + let resolved = first.resolve(env)?; let rebased = self .table_paths .iter() - .map(|p| Ok(p.resolve(registry.as_ref())?.2)) + .map(|p| Ok(p.resolve(env)?.table_url)) .collect::>>()?; - Ok(Some((store, object_store_url, rebased))) + Ok(Some((resolved.store, resolved.identity, rebased))) } async fn collect_files_for_scan<'a>( diff --git a/datafusion/core/src/datasource/listing/table.rs b/datafusion/core/src/datasource/listing/table.rs index 4bf05049519fb..968a4dc62cfe4 100644 --- a/datafusion/core/src/datasource/listing/table.rs +++ b/datafusion/core/src/datasource/listing/table.rs @@ -56,11 +56,11 @@ impl ListingTableConfigExt for ListingTableConfig { }; // Resolve through the registry so a store registered under a path prefix // receives store-relative paths (see `ListingTableUrl::resolve`). - let (store, _, table_path) = - first.resolve(state.runtime_env().object_store_registry.as_ref())?; + let resolved = first.resolve(state.runtime_env())?; - let file = table_path - .list_all_files(state, store.as_ref(), "") + let file = resolved + .table_url + .list_all_files(state, resolved.store.as_ref(), "") .await? .next() .await diff --git a/datafusion/core/src/physical_planner.rs b/datafusion/core/src/physical_planner.rs index b6d28e7b21c79..7a7b0be5a1046 100644 --- a/datafusion/core/src/physical_planner.rs +++ b/datafusion/core/src/physical_planner.rs @@ -749,8 +749,13 @@ impl DefaultPhysicalPlanner { }) => { let original_url = output_url.clone(); let input_exec = children.one()?; - let parsed_url = ListingTableUrl::parse(output_url)?; - let object_store_url = parsed_url.object_store(); + // Resolve through the registry so a store registered under a path + // prefix gets store-relative paths; the identity and rebased url + // are embedded in the sink config for execution-time resolution. + let resolved = ListingTableUrl::parse(output_url)? + .resolve(session_state.runtime_env())?; + let object_store_url = resolved.identity; + let parsed_url = resolved.table_url; let schema = Arc::clone(input.schema().inner()); diff --git a/datafusion/core/src/test_util/parquet.rs b/datafusion/core/src/test_util/parquet.rs index c53495421307b..ce5a5752be890 100644 --- a/datafusion/core/src/test_util/parquet.rs +++ b/datafusion/core/src/test_util/parquet.rs @@ -24,7 +24,7 @@ use std::sync::Arc; use crate::arrow::{datatypes::SchemaRef, record_batch::RecordBatch}; use crate::common::ToDFSchema; use crate::config::ConfigOptions; -use crate::datasource::listing::{ListingTableUrl, PartitionedFile}; +use crate::datasource::listing::PartitionedFile; use crate::datasource::object_store::ObjectStoreUrl; use crate::datasource::physical_plan::ParquetSource; use crate::error::Result; @@ -115,9 +115,8 @@ impl TestParquetFile { .into(); }; - let object_store_url = - ListingTableUrl::parse(canonical_path.to_str().unwrap_or_default())? - .object_store(); + // The generated dataset is always a local file. + let object_store_url = ObjectStoreUrl::local_filesystem(); let object_meta = ObjectMeta { location: Path::parse(canonical_path.to_str().unwrap_or_default())?, diff --git a/datafusion/datasource-csv/src/source.rs b/datafusion/datasource-csv/src/source.rs index 25ec311880405..cca27402eebb1 100644 --- a/datafusion/datasource-csv/src/source.rs +++ b/datafusion/datasource-csv/src/source.rs @@ -449,9 +449,7 @@ pub async fn plan_to_csv( path: impl AsRef, ) -> Result<()> { let path = path.as_ref(); - let parsed = ListingTableUrl::parse(path)?; - let object_store_url = parsed.object_store(); - let store = task_ctx.runtime_env().object_store(&object_store_url)?; + let resolved = ListingTableUrl::parse(path)?.resolve(&task_ctx.runtime_env())?; let writer_buffer_size = task_ctx .session_config() .options() @@ -459,9 +457,9 @@ pub async fn plan_to_csv( .objectstore_writer_buffer_size; let mut join_set = JoinSet::new(); for i in 0..plan.output_partitioning().partition_count() { - let storeref = Arc::clone(&store); + let storeref = Arc::clone(&resolved.store); let plan: Arc = Arc::clone(&plan); - let filename = format!("{}/part-{i}.csv", parsed.prefix()); + let filename = format!("{}/part-{i}.csv", resolved.table_url.prefix()); let file = object_store::path::Path::parse(filename)?; let mut stream = plan.execute(i, Arc::clone(&task_ctx))?; diff --git a/datafusion/datasource-json/src/source.rs b/datafusion/datasource-json/src/source.rs index 8632d6b942bc1..74e7ccb28e47b 100644 --- a/datafusion/datasource-json/src/source.rs +++ b/datafusion/datasource-json/src/source.rs @@ -464,9 +464,7 @@ pub async fn plan_to_json( path: impl AsRef, ) -> Result<()> { let path = path.as_ref(); - let parsed = ListingTableUrl::parse(path)?; - let object_store_url = parsed.object_store(); - let store = task_ctx.runtime_env().object_store(&object_store_url)?; + let resolved = ListingTableUrl::parse(path)?.resolve(&task_ctx.runtime_env())?; let writer_buffer_size = task_ctx .session_config() .options() @@ -474,9 +472,9 @@ pub async fn plan_to_json( .objectstore_writer_buffer_size; let mut join_set = JoinSet::new(); for i in 0..plan.output_partitioning().partition_count() { - let storeref = Arc::clone(&store); + let storeref = Arc::clone(&resolved.store); let plan: Arc = Arc::clone(&plan); - let filename = format!("{}/part-{i}.json", parsed.prefix()); + let filename = format!("{}/part-{i}.json", resolved.table_url.prefix()); let file = object_store::path::Path::parse(filename)?; let mut stream = plan.execute(i, Arc::clone(&task_ctx))?; diff --git a/datafusion/datasource-parquet/src/writer.rs b/datafusion/datasource-parquet/src/writer.rs index d37b6e26a7536..7ed3a67ec118e 100644 --- a/datafusion/datasource-parquet/src/writer.rs +++ b/datafusion/datasource-parquet/src/writer.rs @@ -35,17 +35,15 @@ pub async fn plan_to_parquet( writer_properties: Option, ) -> datafusion_common::Result<()> { let path = path.as_ref(); - let parsed = ListingTableUrl::parse(path)?; - let object_store_url = parsed.object_store(); - let store = task_ctx.runtime_env().object_store(&object_store_url)?; + let resolved = ListingTableUrl::parse(path)?.resolve(&task_ctx.runtime_env())?; let mut join_set = JoinSet::new(); for i in 0..plan.output_partitioning().partition_count() { let plan: Arc = Arc::clone(&plan); - let filename = format!("{}/part-{i}.parquet", parsed.prefix()); + let filename = format!("{}/part-{i}.parquet", resolved.table_url.prefix()); let file = Path::parse(filename)?; let propclone = writer_properties.clone(); - let storeref = Arc::clone(&store); + let storeref = Arc::clone(&resolved.store); let buf_writer = BufWriter::with_capacity( storeref, file.clone(), diff --git a/datafusion/datasource/src/mod.rs b/datafusion/datasource/src/mod.rs index 7c8cae337f1eb..7458a32e4ac37 100644 --- a/datafusion/datasource/src/mod.rs +++ b/datafusion/datasource/src/mod.rs @@ -568,17 +568,6 @@ mod tests { Arc::new(schema) } - #[test] - fn test_object_store_listing_url() { - let listing = ListingTableUrl::parse("file:///").unwrap(); - let store = listing.object_store(); - assert_eq!(store.as_str(), "file:///"); - - let listing = ListingTableUrl::parse("s3://bucket/").unwrap(); - let store = listing.object_store(); - assert_eq!(store.as_str(), "s3://bucket/"); - } - #[test] fn test_get_store_hdfs() { let sut = DefaultObjectStoreRegistry::default(); diff --git a/datafusion/datasource/src/url.rs b/datafusion/datasource/src/url.rs index a0c6b6e5d9fdf..df3bbe818aa9c 100644 --- a/datafusion/datasource/src/url.rs +++ b/datafusion/datasource/src/url.rs @@ -20,7 +20,8 @@ use std::sync::Arc; use datafusion_common::{DataFusionError, Result, TableReference}; use datafusion_execution::cache::cache_manager::CachedFileList; use datafusion_execution::cache::cache_manager::TableScopedPath; -use datafusion_execution::object_store::{ObjectStoreRegistry, ObjectStoreUrl}; +use datafusion_execution::object_store::ObjectStoreUrl; +use datafusion_execution::runtime_env::RuntimeEnv; use datafusion_session::Session; use futures::stream::BoxStream; @@ -33,6 +34,24 @@ use object_store::path::Path; use object_store::{ObjectMeta, ObjectStore, ObjectStoreExt}; use url::Url; +/// The outcome of resolving a [`ListingTableUrl`] against the object store +/// registry (see [`ListingTableUrl::resolve`]). +/// +/// A short-lived, local value: use its fields directly, taking only the ones a +/// given call site needs. +#[derive(Debug, Clone)] +pub struct ResolvedTableUrl { + /// The object store serving the url. + pub store: Arc, + /// The store's identity (scheme + authority + any registered path prefix). + /// Use as the scan/sink key so the store re-resolves at execution time. + pub identity: ObjectStoreUrl, + /// This url rebased into the store's coordinate space: its + /// [`prefix`](ListingTableUrl::prefix) is store-relative, so listing and + /// partition parsing operate on paths relative to the store. + pub table_url: ListingTableUrl, +} + /// A parsed URL identifying files for a listing table, see [`ListingTableUrl::parse`] /// for more information on the supported expressions #[derive(Debug, Clone, Eq, PartialEq, Hash)] @@ -319,74 +338,41 @@ impl ListingTableUrl { self.as_ref() } - /// Return the [`ObjectStoreUrl`] for this [`ListingTableUrl`] - /// - /// This returns the scheme + authority only, and does not account for any - /// store registered under a path prefix. For reading, prefer - /// [`Self::resolve`], which consults the registry and strips the registered - /// prefix from the path handed to the store. - pub fn object_store(&self) -> ObjectStoreUrl { - let url = &self.url[url::Position::BeforeScheme..url::Position::BeforePath]; - ObjectStoreUrl::parse(url).unwrap() - } - - /// Resolve this url against the object store `registry`. + /// Resolves this url against the object store registry in `env`. See + /// [`ResolvedTableUrl`]. /// - /// Returns: - /// - the [`ObjectStore`] serving this url, - /// - the [`ObjectStoreUrl`] (scheme + authority + any registered path prefix) - /// to use as the scan key so the same store can be re-resolved at execution - /// time, and - /// - a copy of this [`ListingTableUrl`] rebased into the resolved store's - /// coordinate space: its [`prefix`](Self::prefix) has the registered prefix - /// stripped, so listing/`head`/partition parsing operate on paths relative - /// to the store (avoiding double-prefixing for stores rooted at a prefix). - /// - /// For a store registered by scheme + authority only (e.g. `s3://bucket`), the - /// rebased prefix is identical to [`Self::prefix`] and the behavior is - /// unchanged. - pub fn resolve( - &self, - registry: &dyn ObjectStoreRegistry, - ) -> Result<(Arc, ObjectStoreUrl, ListingTableUrl)> { - let (store, store_relative_prefix) = registry.resolve(&self.url)?; - - // The registered prefix is the leading path segments that the registry - // stripped: full prefix segments minus the store-relative ones. - let stripped = self - .prefix - .parts() - .count() - .saturating_sub(store_relative_prefix.parts().count()); - - // Rebuild the base URL (scheme + authority + registered prefix) preserving - // the original percent-encoding of the retained segments. + /// The returned [`table_url`](ResolvedTableUrl::table_url) is this url rebased + /// into the store's coordinate space: its [`prefix`](Self::prefix) is replaced + /// by the store-relative path the registry stripped, so listing and partition + /// parsing operate on paths relative to the store (avoiding double-prefixing + /// for a store rooted at a path prefix). Only `prefix` changes; `url` / `glob` + /// / `table_ref` are preserved. + pub fn resolve(&self, env: &RuntimeEnv) -> Result { + let (store, path) = env.object_store_registry.resolve(&self.url)?; + + // Store identity = scheme + authority + registered prefix, i.e. the + // leading segments of `url` the registry stripped to produce `path`, + // keeping their original percent-encoding. + let segments = || self.url.path().split('/').filter(|s| !s.is_empty()); + let kept = segments().count().saturating_sub(path.parts().count()); let authority = &self.url[url::Position::BeforeScheme..url::Position::BeforePath]; - let kept: Vec<&str> = self - .url - .path() - .split('/') - .filter(|s| !s.is_empty()) - .take(stripped) - .collect(); - let base_str = if kept.is_empty() { - authority.to_string() + let identity = if kept == 0 { + ObjectStoreUrl::parse(authority)? } else { - format!("{authority}/{}", kept.join("/")) - }; - let object_store_url = ObjectStoreUrl::parse(&base_str)?; - - // The rebased url keeps the original `url` (used only for scheme / - // is_collection / display); only `prefix` becomes store-relative, which is - // what all listing and partition-parsing logic keys off. - let rebased = ListingTableUrl { - url: self.url.clone(), - prefix: store_relative_prefix, - glob: self.glob.clone(), - table_ref: self.table_ref.clone(), + let prefix = segments().take(kept).collect::>().join("/"); + ObjectStoreUrl::parse(format!("{authority}/{prefix}"))? }; - Ok((store, object_store_url, rebased)) + Ok(ResolvedTableUrl { + store, + identity, + table_url: Self { + url: self.url.clone(), + prefix: path, + glob: self.glob.clone(), + table_ref: self.table_ref.clone(), + }, + }) } /// Returns true if the [`ListingTableUrl`] points to the folder @@ -579,8 +565,9 @@ mod tests { use datafusion_common::config::TableOptions; use datafusion_execution::TaskContext; use datafusion_execution::config::SessionConfig; - use datafusion_execution::object_store::DefaultObjectStoreRegistry; - use datafusion_execution::runtime_env::RuntimeEnv; + use datafusion_execution::object_store::{ + DefaultObjectStoreRegistry, ObjectStoreRegistry, + }; use datafusion_expr::execution_props::ExecutionProps; use datafusion_expr::registry::ExtensionTypeRegistryRef; use datafusion_expr::{ @@ -694,68 +681,72 @@ mod tests { assert!(url.strip_prefix(&path).is_none()); } - fn registry_with(base: &str) -> DefaultObjectStoreRegistry { + fn env_with(base: &str) -> Arc { + use datafusion_execution::runtime_env::RuntimeEnvBuilder; let registry = DefaultObjectStoreRegistry::new(); registry.register_store( &Url::parse(base).unwrap(), Arc::new(object_store::memory::InMemory::new()), ); - registry + RuntimeEnvBuilder::new() + .with_object_store_registry(Arc::new(registry)) + .build_arc() + .unwrap() } #[test] fn test_resolve_authority_only_store() { // A store registered by scheme+authority only leaves the prefix unchanged. - let registry = registry_with("s3://bucket"); + let env = env_with("s3://bucket"); let url = ListingTableUrl::parse("s3://bucket/foo/bar/file.parquet").unwrap(); - let (_, base, rebased) = url.resolve(®istry).unwrap(); - assert_eq!(base.as_str(), "s3://bucket/"); - assert_eq!(rebased.prefix().as_ref(), "foo/bar/file.parquet"); - assert_eq!(rebased.prefix(), url.prefix()); + let resolved = url.resolve(&env).unwrap(); + assert_eq!(resolved.identity.as_str(), "s3://bucket/"); + assert_eq!(resolved.table_url.prefix().as_ref(), "foo/bar/file.parquet"); + assert_eq!(resolved.table_url.prefix(), url.prefix()); } #[test] fn test_resolve_prefixed_store_strips_prefix() { // A store registered under a path prefix: the prefix is stripped from the - // rebased url and retained in the base ObjectStoreUrl (the scan key). - let registry = registry_with("s3://bucket/user/repo"); + // rebased url and retained in the store identity (the scan key). + let env = env_with("s3://bucket/user/repo"); let url = ListingTableUrl::parse("s3://bucket/user/repo/data/file.parquet").unwrap(); - let (_, base, rebased) = url.resolve(®istry).unwrap(); - assert_eq!(base.as_str(), "s3://bucket/user/repo"); - assert_eq!(rebased.prefix().as_ref(), "data/file.parquet"); - assert!(!rebased.is_collection()); + let resolved = url.resolve(&env).unwrap(); + assert_eq!(resolved.identity.as_str(), "s3://bucket/user/repo"); + assert_eq!(resolved.table_url.prefix().as_ref(), "data/file.parquet"); + assert!(!resolved.table_url.is_collection()); } #[test] fn test_resolve_prefixed_collection_preserves_trailing_slash() { // A directory url keeps its collection-ness after rebasing. - let registry = registry_with("s3://bucket/user/repo"); + let env = env_with("s3://bucket/user/repo"); let url = ListingTableUrl::parse("s3://bucket/user/repo/data/").unwrap(); - let (_, base, rebased) = url.resolve(®istry).unwrap(); - assert_eq!(base.as_str(), "s3://bucket/user/repo"); - assert_eq!(rebased.prefix().as_ref(), "data"); - assert!(rebased.is_collection()); + let resolved = url.resolve(&env).unwrap(); + assert_eq!(resolved.identity.as_str(), "s3://bucket/user/repo"); + assert_eq!(resolved.table_url.prefix().as_ref(), "data"); + assert!(resolved.table_url.is_collection()); } #[test] fn test_resolve_query_at_registered_root() { // Querying exactly the registered prefix yields an empty rebased prefix. - let registry = registry_with("s3://bucket/user/repo"); + let env = env_with("s3://bucket/user/repo"); let url = ListingTableUrl::parse("s3://bucket/user/repo/").unwrap(); - let (_, base, rebased) = url.resolve(®istry).unwrap(); - assert_eq!(base.as_str(), "s3://bucket/user/repo"); - assert_eq!(rebased.prefix().as_ref(), ""); + let resolved = url.resolve(&env).unwrap(); + assert_eq!(resolved.identity.as_str(), "s3://bucket/user/repo"); + assert_eq!(resolved.table_url.prefix().as_ref(), ""); } #[test] - fn test_resolve_preserves_percent_encoding_in_base() { + fn test_resolve_preserves_percent_encoding_in_identity() { // The retained (registered-prefix) segments keep their original encoding. - let registry = registry_with("s3://bucket/us%20er"); + let env = env_with("s3://bucket/us%20er"); let url = ListingTableUrl::parse("s3://bucket/us%20er/data.parquet").unwrap(); - let (_, base, rebased) = url.resolve(®istry).unwrap(); - assert_eq!(base.as_str(), "s3://bucket/us%20er"); - assert_eq!(rebased.prefix().as_ref(), "data.parquet"); + let resolved = url.resolve(&env).unwrap(); + assert_eq!(resolved.identity.as_str(), "s3://bucket/us%20er"); + assert_eq!(resolved.table_url.prefix().as_ref(), "data.parquet"); } #[test] diff --git a/datafusion/execution/src/object_store.rs b/datafusion/execution/src/object_store.rs index 28f866134e188..5579888ca58ae 100644 --- a/datafusion/execution/src/object_store.rs +++ b/datafusion/execution/src/object_store.rs @@ -189,8 +189,9 @@ pub trait ObjectStoreRegistry: Send + Sync + std::fmt::Debug + 'static { /// path prefix (e.g. `hf://bucket/user/repo`). When resolving a longer URL /// (e.g. `hf://bucket/user/repo/data/f.parquet`) the registered prefix is /// stripped from the returned path (`data/f.parquet`), so a store rooted at - /// that prefix is not prefixed twice. The store with the longest matching - /// path prefix is returned. + /// that prefix is not prefixed twice. The store with the longest matching path + /// prefix is returned. The store identity (scheme + authority + registered + /// prefix) can be reconstructed from `url` and the returned path. /// /// The default implementation preserves the legacy behavior of matching on /// scheme + authority only, returning the full URL path unchanged. @@ -328,7 +329,8 @@ mod tests { Arc::clone(&repo2), ); - // Each query resolves to its own store, with the registered prefix stripped. + // Each query resolves to its own store, with the registered prefix + // stripped from the returned path. let (store, path) = registry .resolve(&Url::parse("hf://bucket/userA/repo1/data/f.parquet").unwrap()) .unwrap(); From c14e075407294a097e213ffcfa58e8fef684a052 Mon Sep 17 00:00:00 2001 From: Krisztian Szucs Date: Fri, 3 Jul 2026 16:33:29 +0200 Subject: [PATCH 4/7] Simplify store identity derivation in ListingTableUrl::resolve --- datafusion/datasource/src/url.rs | 21 +++++++++++---------- 1 file changed, 11 insertions(+), 10 deletions(-) diff --git a/datafusion/datasource/src/url.rs b/datafusion/datasource/src/url.rs index df3bbe818aa9c..d8b74c33283d3 100644 --- a/datafusion/datasource/src/url.rs +++ b/datafusion/datasource/src/url.rs @@ -350,18 +350,19 @@ impl ListingTableUrl { pub fn resolve(&self, env: &RuntimeEnv) -> Result { let (store, path) = env.object_store_registry.resolve(&self.url)?; - // Store identity = scheme + authority + registered prefix, i.e. the - // leading segments of `url` the registry stripped to produce `path`, + // Store identity = `url` with the store-relative `path` stripped off its + // tail: scheme + authority + the leading segments the registry consumed, // keeping their original percent-encoding. - let segments = || self.url.path().split('/').filter(|s| !s.is_empty()); - let kept = segments().count().saturating_sub(path.parts().count()); + let segments: Vec<&str> = self + .url + .path() + .split('/') + .filter(|s| !s.is_empty()) + .collect(); + let kept = segments.len().saturating_sub(path.parts().count()); let authority = &self.url[url::Position::BeforeScheme..url::Position::BeforePath]; - let identity = if kept == 0 { - ObjectStoreUrl::parse(authority)? - } else { - let prefix = segments().take(kept).collect::>().join("/"); - ObjectStoreUrl::parse(format!("{authority}/{prefix}"))? - }; + let identity = + ObjectStoreUrl::parse(format!("{authority}/{}", segments[..kept].join("/")))?; Ok(ResolvedTableUrl { store, From f5b5aed72ce7c172d9a538a33a59bef0d3f04b87 Mon Sep 17 00:00:00 2001 From: Krisztian Szucs Date: Fri, 3 Jul 2026 21:38:39 +0200 Subject: [PATCH 5/7] Fix object store registration, multi-path scans, and strict resolution - datafusion-cli: register object stores at their scheme/authority identity instead of the full object path, so a store rooted at the authority is not scoped to a deep prefix and handed an empty relative path (broke every CLI CREATE EXTERNAL TABLE / COPY with a path component). - catalog-listing: reject a ListingTable whose paths resolve to different object stores (a scan carries a single store), replace the positional ResolvedTablePaths tuple with a named struct, and resolve the first path once. - execution: preserve the strict "No suitable object store found" error by gating resolve on explicit registration, so the registry never lazily creates a store from ambient environment credentials. Restore the file:// default registration for DefaultObjectStoreRegistry::default(). - Add regression tests for each fix. --- datafusion-cli/src/exec.rs | 60 ++++++-- datafusion/catalog-listing/src/table.rs | 74 ++++++--- .../tests/datasource/object_store_access.rs | 41 +++++ datafusion/execution/src/object_store.rs | 140 ++++++++++++++++-- 4 files changed, 271 insertions(+), 44 deletions(-) diff --git a/datafusion-cli/src/exec.rs b/datafusion-cli/src/exec.rs index c36d3dba51113..4cfde75e447ad 100644 --- a/datafusion-cli/src/exec.rs +++ b/datafusion-cli/src/exec.rs @@ -515,17 +515,17 @@ pub(crate) async fn register_object_store_and_config_extensions( .await?; // Register the retrieved object store in the session context's runtime - // environment. The stdin store is shared across all `stdin://` object paths, - // so it is registered at its scheme/authority identity rather than at the - // object path (the registry then resolves every `stdin://` path to it). - if scheme == StdinUtils::SCHEME { - let identity = ObjectStoreUrl::parse( - &url[::url::Position::BeforeScheme..::url::Position::BeforePath], - )?; - ctx.register_object_store(identity.as_ref(), store); - } else { - ctx.register_object_store(url, store); - } + // environment at its scheme/authority identity rather than at the full object + // path. `get_object_store` always returns an authority/root-rooted store (an + // S3 bucket, the local filesystem, an HTTP origin, or the shared per-session + // stdin store), so the registry resolves every object path under that + // authority to it. Registering at the full path would instead scope the store + // to that single prefix and hand it store-relative paths with the prefix + // stripped, breaking reads/writes. + let identity = ObjectStoreUrl::parse( + &url[::url::Position::BeforeScheme..::url::Position::BeforePath], + )?; + ctx.register_object_store(identity.as_ref(), store); Ok(()) } @@ -602,6 +602,44 @@ mod tests { Ok(()) } + + /// Regression test: `register_object_store_and_config_extensions` must + /// register the store at its scheme/authority identity, not at the full + /// object path. Otherwise the (root-rooted) store is scoped to a deep path + /// prefix and handed an empty store-relative path, so a `CREATE EXTERNAL + /// TABLE` over a local file can no longer be read back. + #[tokio::test] + async fn create_external_table_reads_local_file() -> Result<()> { + use datafusion::arrow::array::Int64Array; + + let ctx = SessionContext::new(); + // Relative to the datafusion-cli crate root (the test working dir). + let location = "../datafusion/core/tests/data/cars.csv"; + let sql = format!( + "CREATE EXTERNAL TABLE cars STORED AS CSV LOCATION '{location}' \ + OPTIONS ('has_header' 'TRUE')" + ); + + // Drive the full CLI planning path so the object store is registered. + let statement = DFParser::parse_sql(&sql)?.pop_front().unwrap(); + let plan = create_plan(&ctx, statement, false).await?; + ctx.execute_logical_plan(plan).await?; + + let batches = ctx + .sql("SELECT count(*) AS n FROM cars") + .await? + .collect() + .await?; + let n = batches[0] + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .value(0); + assert_eq!(n, 25); + + Ok(()) + } #[tokio::test] async fn copy_to_external_object_store_test() -> Result<()> { let aws_envs = vec![ diff --git a/datafusion/catalog-listing/src/table.rs b/datafusion/catalog-listing/src/table.rs index d9d8a9814c3a1..eaf55ded8db60 100644 --- a/datafusion/catalog-listing/src/table.rs +++ b/datafusion/catalog-listing/src/table.rs @@ -25,8 +25,8 @@ use async_trait::async_trait; use datafusion_catalog::{ScanArgs, ScanResult, Session, TableProvider}; use datafusion_common::stats::Precision; use datafusion_common::{ - Constraints, DFSchema, SchemaExt, Statistics, internal_datafusion_err, plan_err, - project_schema, + Constraints, DFSchema, SchemaExt, Statistics, exec_err, internal_datafusion_err, + plan_err, project_schema, }; use datafusion_datasource::file::FileSource; use datafusion_datasource::file_groups::FileGroup; @@ -56,11 +56,16 @@ use object_store::ObjectStore; use std::collections::HashMap; use std::sync::Arc; -/// Result of resolving a table's paths against the object store registry: the -/// store serving the table, the base [`ObjectStoreUrl`] (including any registered -/// path prefix) to use as the scan key, and the table paths rebased into the -/// store's coordinate space. -type ResolvedTablePaths = (Arc, ObjectStoreUrl, Vec); +/// Result of resolving a table's paths against the object store registry. +struct ResolvedTablePaths { + /// The object store serving the table (all `table_urls` resolve to it). + store: Arc, + /// The base [`ObjectStoreUrl`] (including any registered path prefix) to use + /// as the scan/sink key. + identity: ObjectStoreUrl, + /// The table paths rebased into the store's coordinate space. + table_urls: Vec, +} /// Result of a file listing operation from [`ListingTable::list_files_for_scan`]. #[derive(Debug)] @@ -634,14 +639,14 @@ impl TableProvider for ListingTable { None }; - let Some((_, object_store_url, _)) = self.resolve_object_store(state)? else { + let Some(resolved) = self.resolve_object_store(state)? else { return Ok(ScanResult::new(Arc::new(EmptyExec::new(Arc::new( Schema::empty(), ))))); }; let file_source = self.create_file_source(); - let scan_config = FileScanConfigBuilder::new(object_store_url, file_source) + let scan_config = FileScanConfigBuilder::new(resolved.identity, file_source) .with_file_groups(partitioned_file_lists) .with_constraints(self.constraints.clone()) .with_statistics(statistics) @@ -713,11 +718,14 @@ impl TableProvider for ListingTable { // gets store-relative paths (see `ListingTableUrl::with_prefix`). The // resolved identity and rebased paths are embedded in the sink config so // the store re-resolves and writes correctly at execution time. - let Some((store, object_store_url, rebased_paths)) = - self.resolve_object_store(state)? - else { + let Some(resolved) = self.resolve_object_store(state)? else { return Err(internal_datafusion_err!("ListingTable has no table paths")); }; + let ResolvedTablePaths { + store, + identity: object_store_url, + table_urls: rebased_paths, + } = resolved; let rebased = &rebased_paths[0]; let file_list_stream = pruned_partition_list( @@ -804,6 +812,9 @@ impl ListingTable { /// url including any registered path prefix) to use as the scan key, and the /// table paths rebased into the store's coordinate space (registered prefix /// stripped). Returns `None` when the table has no paths. + /// + /// A [`FileScanConfig`]/[`FileSinkConfig`] carries a single object store, so + /// every table path must resolve to the same store; this errors otherwise. fn resolve_object_store( &self, ctx: &dyn Session, @@ -813,12 +824,25 @@ impl ListingTable { }; let env = ctx.runtime_env(); let resolved = first.resolve(env)?; - let rebased = self - .table_paths - .iter() - .map(|p| Ok(p.resolve(env)?.table_url)) - .collect::>>()?; - Ok(Some((resolved.store, resolved.identity, rebased))) + let mut table_urls = Vec::with_capacity(self.table_paths.len()); + table_urls.push(resolved.table_url.clone()); + for path in &self.table_paths[1..] { + let other = path.resolve(env)?; + if other.identity != resolved.identity { + return exec_err!( + "ListingTable spans multiple object stores ({} and {}); \ + all table paths must resolve to the same object store", + resolved.identity.as_str(), + other.identity.as_str() + ); + } + table_urls.push(other.table_url); + } + Ok(Some(ResolvedTablePaths { + store: resolved.store, + identity: resolved.identity, + table_urls, + })) } async fn collect_files_for_scan<'a>( @@ -877,7 +901,7 @@ impl ListingTable { ); } - let Some((store, _, rebased_paths)) = self.resolve_object_store(ctx)? else { + let Some(resolved) = self.resolve_object_store(ctx)? else { return Ok(ListFilesResult { file_groups: vec![], statistics: Statistics::new_unknown(&self.file_schema), @@ -885,7 +909,13 @@ impl ListingTable { }); }; let (file_group, inexact_stats) = self - .collect_files_for_scan(ctx, &store, &rebased_paths, filters, limit) + .collect_files_for_scan( + ctx, + &resolved.store, + &resolved.table_urls, + filters, + limit, + ) .await?; // Threshold: 0 = disabled, N > 0 = enabled when distinct_keys >= N @@ -937,7 +967,7 @@ impl ListingTable { ); } - let Some((store, _, rebased_paths)) = self.resolve_object_store(ctx)? else { + let Some(resolved) = self.resolve_object_store(ctx)? else { return Ok(ListFilesResult { file_groups: vec![], statistics: Statistics::new_unknown(&self.file_schema), @@ -945,7 +975,7 @@ impl ListingTable { }); }; let (file_group, inexact_stats) = self - .collect_files_for_scan(ctx, &store, &rebased_paths, &[], None) + .collect_files_for_scan(ctx, &resolved.store, &resolved.table_urls, &[], None) .await?; let mut file_groups = file_group.split_files(file_group_count); if !file_groups.is_empty() { diff --git a/datafusion/core/tests/datasource/object_store_access.rs b/datafusion/core/tests/datasource/object_store_access.rs index a4f3d3bbb16dc..10cdac0e0a30e 100644 --- a/datafusion/core/tests/datasource/object_store_access.rs +++ b/datafusion/core/tests/datasource/object_store_access.rs @@ -1063,6 +1063,47 @@ async fn two_prefixed_stores_coexist_under_one_authority() { ); } +/// A single `ListingTable` whose paths resolve to *different* stores under the +/// same authority must be rejected: a scan carries a single object store, so +/// silently reading every path from the first store would drop/duplicate rows. +#[tokio::test] +async fn multi_path_table_spanning_two_stores_is_rejected() { + use datafusion::catalog::TableProvider; + use datafusion_catalog_listing::ListingTableConfig; + use object_store::prefix::PrefixStore; + + let ctx = SessionContext::new(); + ctx.runtime_env().register_object_store( + &Url::parse("mem://bucket/userA/repo1").unwrap(), + Arc::new(PrefixStore::new(InMemory::new(), "userA/repo1")) + as Arc, + ); + ctx.runtime_env().register_object_store( + &Url::parse("mem://bucket/userB/repo2").unwrap(), + Arc::new(PrefixStore::new(InMemory::new(), "userB/repo2")) + as Arc, + ); + + let options = ListingOptions::new(Arc::new(CsvFormat::default())); + let schema = Arc::new(arrow::datatypes::Schema::new(vec![ + arrow::datatypes::Field::new("a", arrow::datatypes::DataType::Int32, false), + ])); + let config = ListingTableConfig::new_with_multi_paths(vec![ + ListingTableUrl::parse("mem://bucket/userA/repo1/data/").unwrap(), + ListingTableUrl::parse("mem://bucket/userB/repo2/data/").unwrap(), + ]) + .with_listing_options(options) + .with_schema(schema); + let table = ListingTable::try_new(config).unwrap(); + + let state = ctx.state(); + let err = table.scan(&state, None, &[], None).await.unwrap_err(); + assert!( + err.to_string().contains("multiple object stores"), + "unexpected error: {err}" + ); +} + /// Runs tests with a request counting object store struct Test { object_store: Arc, diff --git a/datafusion/execution/src/object_store.rs b/datafusion/execution/src/object_store.rs index 5579888ca58ae..6a22b42aaf3e4 100644 --- a/datafusion/execution/src/object_store.rs +++ b/datafusion/execution/src/object_store.rs @@ -27,7 +27,8 @@ use object_store::ObjectStore; use object_store::local::LocalFileSystem; use object_store::path::Path; use object_store::registry::ObjectStoreRegistry as UpstreamObjectStoreRegistry; -use std::sync::Arc; +use std::collections::HashSet; +use std::sync::{Arc, RwLock}; use url::Url; /// A parsed URL identifying a particular [`ObjectStore`] instance @@ -215,29 +216,79 @@ pub trait ObjectStoreRegistry: Send + Sync + std::fmt::Debug + 'static { /// `hf://bucket/user/repo`, and [`ObjectStoreRegistry::resolve`] strips the /// registered prefix from the returned path (`data/f.parquet`) /// -/// Stores for known schemes (`file`, `s3`, `gs`, `az`, `http`, ...) are created -/// lazily on first access via [`object_store::parse_url_opts`]. -#[derive(Debug, Default)] +/// Only stores that have been explicitly registered (plus the default +/// `file://` store) are resolved: a URL with no matching registration returns an +/// error rather than lazily creating a store from the environment. This keeps +/// object store configuration explicit and avoids silently picking up ambient +/// (e.g. cloud) credentials. +#[derive(Debug)] pub struct DefaultObjectStoreRegistry { inner: object_store::registry::DefaultObjectStoreRegistry, + /// Scheme+authority and path segments of every explicitly registered store. + /// + /// The inner registry would otherwise lazily create a store from the + /// environment for any known scheme; tracking registrations lets `resolve` + /// reject unregistered URLs before that fallback runs. + registered: RwLock)>>, +} + +impl Default for DefaultObjectStoreRegistry { + fn default() -> Self { + Self::new() + } +} + +/// Splits `url` into its scheme+authority and non-empty path segments, the key +/// under which a store is registered / matched. +fn registration_key(url: &Url) -> (String, Vec) { + let authority = url[..url::Position::AfterPort].to_string(); + let segments = url + .path() + .split('/') + .filter(|s| !s.is_empty()) + .map(str::to_string) + .collect(); + (authority, segments) } impl DefaultObjectStoreRegistry { /// This will register [`LocalFileSystem`] to handle `file://` paths #[cfg(not(target_arch = "wasm32"))] pub fn new() -> Self { - let inner = object_store::registry::DefaultObjectStoreRegistry::new(); - inner.register( - Url::parse("file://").unwrap(), + let this = Self { + inner: object_store::registry::DefaultObjectStoreRegistry::new(), + registered: RwLock::new(HashSet::new()), + }; + this.register_store( + &Url::parse("file://").unwrap(), Arc::new(LocalFileSystem::new()), ); - Self { inner } + this } /// Default without any backend registered. #[cfg(target_arch = "wasm32")] pub fn new() -> Self { - Self::default() + Self { + inner: object_store::registry::DefaultObjectStoreRegistry::new(), + registered: RwLock::new(HashSet::new()), + } + } + + /// Returns `true` if some registered store's path is a prefix of `url`'s + /// (under the same scheme+authority), i.e. `url` resolves to a registered + /// store rather than requiring lazy creation. + fn has_registered_match(&self, url: &Url) -> bool { + let (authority, segments) = registration_key(url); + self.registered + .read() + .unwrap() + .iter() + .any(|(reg_authority, reg_segments)| { + *reg_authority == authority + && reg_segments.len() <= segments.len() + && reg_segments.iter().zip(&segments).all(|(a, b)| a == b) + }) } } @@ -247,13 +298,22 @@ impl ObjectStoreRegistry for DefaultObjectStoreRegistry { url: &Url, store: Arc, ) -> Option> { + self.registered + .write() + .unwrap() + .insert(registration_key(url)); self.inner.register(url.clone(), store) } fn deregister_store(&self, url: &Url) -> Result> { - self.inner.deregister(url).ok_or_else(|| { + let store = self.inner.deregister(url).ok_or_else(|| { internal_datafusion_err!("Failed to deregister object store. No suitable object store found for {url}. See `RuntimeEnv::register_object_store`") - }) + })?; + self.registered + .write() + .unwrap() + .remove(®istration_key(url)); + Ok(store) } fn get_store(&self, url: &Url) -> Result> { @@ -261,6 +321,13 @@ impl ObjectStoreRegistry for DefaultObjectStoreRegistry { } fn resolve(&self, url: &Url) -> Result<(Arc, Path)> { + // Reject unregistered URLs before the inner registry's lazy + // `parse_url_opts` fallback can create a store from the environment. + if !self.has_registered_match(url) { + return Err(internal_datafusion_err!( + "No suitable object store found for {url}. See `RuntimeEnv::register_object_store`" + )); + } self.inner .resolve(url) .map_err(|e| DataFusionError::External(Box::new(e))) @@ -381,4 +448,55 @@ mod tests { assert_eq!(path.as_ref(), "a/b/c.parquet"); assert!(Arc::ptr_eq(&resolved, &store)); } + + #[test] + fn test_registry_rejects_unregistered_store() { + use object_store::memory::InMemory; + + let registry = DefaultObjectStoreRegistry::new(); + + // An unregistered scheme+authority is a strict error: the registry must + // not lazily create a store from the environment. + let err = registry + .resolve(&Url::parse("s3://bucket/data/f.parquet").unwrap()) + .unwrap_err(); + assert!( + err.strip_backtrace() + .contains("No suitable object store found"), + "unexpected error: {}", + err.strip_backtrace() + ); + assert!( + registry + .get_store(&Url::parse("s3://bucket/x").unwrap()) + .is_err() + ); + + // The default `file://` store is registered and still resolves. + assert!( + registry + .get_store(&Url::parse("file:///tmp/x").unwrap()) + .is_ok() + ); + + // Once registered explicitly, the store resolves. + registry.register_store( + &Url::parse("s3://bucket").unwrap(), + Arc::new(InMemory::new()) as Arc, + ); + let (_, path) = registry + .resolve(&Url::parse("s3://bucket/data/f.parquet").unwrap()) + .unwrap(); + assert_eq!(path.as_ref(), "data/f.parquet"); + + // Deregistering restores the strict error. + registry + .deregister_store(&Url::parse("s3://bucket").unwrap()) + .unwrap(); + assert!( + registry + .resolve(&Url::parse("s3://bucket/data/f.parquet").unwrap()) + .is_err() + ); + } } From 75f189050a108a1f394330a8c0821e7eb266de97 Mon Sep 17 00:00:00 2001 From: Krisztian Szucs Date: Sun, 5 Jul 2026 20:44:17 +0200 Subject: [PATCH 6/7] Restore deprecated object_store(), add test coverage for prefixed registry - Restore ListingTableUrl::object_store() as #[deprecated] in favor of resolve(), which correctly handles stores registered under a path prefix. - Document the registered-set linear scan in DefaultObjectStoreRegistry as an intentional tradeoff at expected registration cardinalities. - Add a positive test for a multi-path ListingTable whose paths all resolve to the same prefixed store. - Add a proto round-trip test asserting a path-bearing ObjectStoreUrl survives ParquetExec serde. --- .../tests/datasource/object_store_access.rs | 32 +++++++++++++ datafusion/datasource/src/mod.rs | 12 +++++ datafusion/datasource/src/url.rs | 12 +++++ datafusion/execution/src/object_store.rs | 7 +++ .../tests/cases/roundtrip_physical_plan.rs | 48 +++++++++++++++++++ 5 files changed, 111 insertions(+) diff --git a/datafusion/core/tests/datasource/object_store_access.rs b/datafusion/core/tests/datasource/object_store_access.rs index 10cdac0e0a30e..0d074208b7c05 100644 --- a/datafusion/core/tests/datasource/object_store_access.rs +++ b/datafusion/core/tests/datasource/object_store_access.rs @@ -1104,6 +1104,38 @@ async fn multi_path_table_spanning_two_stores_is_rejected() { ); } +/// A single `ListingTable` whose paths all resolve to the *same* store under +/// one registered path prefix must be accepted, each path correctly rebased +/// to its store-relative prefix. +#[tokio::test] +async fn multi_path_table_spanning_one_prefixed_store_is_accepted() { + use datafusion::catalog::TableProvider; + use datafusion_catalog_listing::ListingTableConfig; + use object_store::prefix::PrefixStore; + + let ctx = SessionContext::new(); + ctx.runtime_env().register_object_store( + &Url::parse("mem://bucket/user/repo").unwrap(), + Arc::new(PrefixStore::new(InMemory::new(), "user/repo")) as Arc, + ); + + let options = ListingOptions::new(Arc::new(CsvFormat::default())); + let schema = Arc::new(arrow::datatypes::Schema::new(vec![ + arrow::datatypes::Field::new("a", arrow::datatypes::DataType::Int32, false), + ])); + let config = ListingTableConfig::new_with_multi_paths(vec![ + ListingTableUrl::parse("mem://bucket/user/repo/a/").unwrap(), + ListingTableUrl::parse("mem://bucket/user/repo/b/").unwrap(), + ]) + .with_listing_options(options) + .with_schema(schema); + let table = ListingTable::try_new(config).unwrap(); + + let state = ctx.state(); + // Must not error with "multiple object stores": both paths share one store. + table.scan(&state, None, &[], None).await.unwrap(); +} + /// Runs tests with a request counting object store struct Test { object_store: Arc, diff --git a/datafusion/datasource/src/mod.rs b/datafusion/datasource/src/mod.rs index 7458a32e4ac37..25e5c9aa2b9d6 100644 --- a/datafusion/datasource/src/mod.rs +++ b/datafusion/datasource/src/mod.rs @@ -568,6 +568,18 @@ mod tests { Arc::new(schema) } + #[test] + #[allow(deprecated)] + fn test_object_store_listing_url() { + let listing = ListingTableUrl::parse("file:///").unwrap(); + let store = listing.object_store(); + assert_eq!(store.as_str(), "file:///"); + + let listing = ListingTableUrl::parse("s3://bucket/").unwrap(); + let store = listing.object_store(); + assert_eq!(store.as_str(), "s3://bucket/"); + } + #[test] fn test_get_store_hdfs() { let sut = DefaultObjectStoreRegistry::default(); diff --git a/datafusion/datasource/src/url.rs b/datafusion/datasource/src/url.rs index d8b74c33283d3..e486838bafeed 100644 --- a/datafusion/datasource/src/url.rs +++ b/datafusion/datasource/src/url.rs @@ -338,6 +338,18 @@ impl ListingTableUrl { self.as_ref() } + /// Return the [`ObjectStoreUrl`] for this [`ListingTableUrl`] + #[deprecated( + since = "55.0.0", + note = "Use `resolve()` instead, which returns the resolved object \ + store and the store-relative path together; this method does \ + not account for stores registered under a path prefix" + )] + pub fn object_store(&self) -> ObjectStoreUrl { + let url = &self.url[url::Position::BeforeScheme..url::Position::BeforePath]; + ObjectStoreUrl::parse(url).unwrap() + } + /// Resolves this url against the object store registry in `env`. See /// [`ResolvedTableUrl`]. /// diff --git a/datafusion/execution/src/object_store.rs b/datafusion/execution/src/object_store.rs index 6a22b42aaf3e4..599caf8991c42 100644 --- a/datafusion/execution/src/object_store.rs +++ b/datafusion/execution/src/object_store.rs @@ -229,6 +229,13 @@ pub struct DefaultObjectStoreRegistry { /// The inner registry would otherwise lazily create a store from the /// environment for any known scheme; tracking registrations lets `resolve` /// reject unregistered URLs before that fallback runs. + /// + /// `has_registered_match` scans this set linearly on every `resolve`/ + /// `get_store` call. That's intentional: registration counts are expected + /// to stay in the range of a handful to low hundreds of stores (one per + /// mounted bucket/prefix, not per query or per file), so a linear scan is + /// cheaper and simpler than a trie/prefix-map. Revisit only if profiling + /// shows registration cardinality growing past that range. registered: RwLock)>>, } diff --git a/datafusion/proto/tests/cases/roundtrip_physical_plan.rs b/datafusion/proto/tests/cases/roundtrip_physical_plan.rs index 8acb891d6a94d..f1abef9123c3f 100644 --- a/datafusion/proto/tests/cases/roundtrip_physical_plan.rs +++ b/datafusion/proto/tests/cases/roundtrip_physical_plan.rs @@ -1124,6 +1124,54 @@ fn roundtrip_parquet_exec_attaches_cached_reader_factory_after_roundtrip() -> Re Ok(()) } +#[test] +fn roundtrip_parquet_exec_preserves_prefixed_object_store_url() -> Result<()> { + let file_schema = + Arc::new(Schema::new(vec![Field::new("col", DataType::Utf8, false)])); + let file_source = Arc::new(ParquetSource::new(Arc::clone(&file_schema))); + let object_store_url = ObjectStoreUrl::parse("s3://bucket/user/repo")?; + let scan_config = FileScanConfigBuilder::new(object_store_url.clone(), file_source) + .with_file_groups(vec![FileGroup::new(vec![PartitionedFile::new( + "/path/to/file.parquet".to_string(), + 1024, + )])]) + .with_statistics(Statistics { + num_rows: Precision::Inexact(100), + total_byte_size: Precision::Inexact(1024), + column_statistics: Statistics::unknown_column(&file_schema), + }) + .build(); + let exec_plan = DataSourceExec::from_data_source(scan_config); + + let ctx = SessionContext::new(); + // Decoding a `ParquetExec` resolves `object_store_url` to attach a cached + // reader factory, so the store must be registered at the same identity + // the scan config carries (including its path prefix). + ctx.runtime_env().register_object_store( + object_store_url.as_ref(), + Arc::new(object_store::memory::InMemory::new()), + ); + let codec = DefaultPhysicalExtensionCodec {}; + let proto_converter = DefaultPhysicalProtoConverter {}; + let roundtripped = + roundtrip_test_and_return(exec_plan, &ctx, &codec, &proto_converter)?; + + let data_source = roundtripped + .downcast_ref::() + .ok_or_else(|| { + internal_datafusion_err!("Expected DataSourceExec after roundtrip") + })?; + let file_scan = data_source + .data_source() + .downcast_ref::() + .ok_or_else(|| { + internal_datafusion_err!("Expected FileScanConfig after roundtrip") + })?; + + assert_eq!(file_scan.object_store_url, object_store_url); + Ok(()) +} + #[test] fn roundtrip_arrow_scan() -> Result<()> { let file_schema = From 9db3c0c1337a46975ed040ca0d141227042d83a5 Mon Sep 17 00:00:00 2001 From: Krisztian Szucs Date: Wed, 8 Jul 2026 21:42:25 +0200 Subject: [PATCH 7/7] fix: replace #[allow(deprecated)] with #[expect(deprecated)] for clippy clippy::allow_attributes denies #[allow], requiring #[expect] instead. Also clarify the object_store patch comment now that arrow-rs-object-store#784 (deregister) has merged upstream but is not yet released. --- Cargo.toml | 7 ++++--- datafusion/datasource/src/mod.rs | 2 +- 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index bc2148a4c4160..64aee9652bb85 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -268,9 +268,10 @@ codegen-units = 16 lto = false strip = false # Retain debug info for flamegraphs -# Pin object_store to arrow-rs-object-store PR #784 (adds a `deregister` method to -# the registry). Uses the 0.13-compatible branch so the patch applies to `parquet` -# (which requires object_store ^0.13). Drop once the PR is released. +# Pin object_store to a branch carrying the `deregister` method added to the +# registry by arrow-rs-object-store#784 (merged upstream, not yet released). +# Uses the 0.13-compatible branch so the patch applies to `parquet` (which +# requires object_store ^0.13). Drop once a release containing #784 ships. [patch.crates-io] object_store = { git = "https://github.com/kszucs/arrow-rs-object-store", branch = "registry-deregister-0.13" } diff --git a/datafusion/datasource/src/mod.rs b/datafusion/datasource/src/mod.rs index 25e5c9aa2b9d6..7a7f83b1b8242 100644 --- a/datafusion/datasource/src/mod.rs +++ b/datafusion/datasource/src/mod.rs @@ -569,7 +569,7 @@ mod tests { } #[test] - #[allow(deprecated)] + #[expect(deprecated)] fn test_object_store_listing_url() { let listing = ListingTableUrl::parse("file:///").unwrap(); let store = listing.object_store();