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..64aee9652bb85 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -268,6 +268,13 @@ codegen-units = 16 lto = false strip = false # Retain debug info for flamegraphs +# 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" } + [profile.profiling] inherits = "release" debug = true diff --git a/datafusion-cli/src/exec.rs b/datafusion-cli/src/exec.rs index f43854821b2d5..4cfde75e447ad 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; @@ -513,8 +514,18 @@ 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 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(()) } @@ -591,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-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..f9c03cf5ee06c 100644 --- a/datafusion-cli/src/object_storage/stdin.rs +++ b/datafusion-cli/src/object_storage/stdin.rs @@ -103,9 +103,10 @@ impl StdinUtils { /// 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. @@ -119,6 +120,9 @@ impl StdinUtils { state: &SessionState, url: &Url, ) -> Result> { + // 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; @@ -252,7 +256,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(&Url::parse("stdin:///").unwrap(), store); ctx.sql(&format!( "CREATE EXTERNAL TABLE t STORED AS {stored_as} LOCATION '{location}' {options}" )) @@ -281,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(&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!( @@ -304,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(&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) diff --git a/datafusion/catalog-listing/src/options.rs b/datafusion/catalog-listing/src/options.rs index 44337e52a1e05..61e67222f0278 100644 --- a/datafusion/catalog-listing/src/options.rs +++ b/datafusion/catalog-listing/src/options.rs @@ -277,10 +277,13 @@ 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 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?; @@ -290,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 ); } @@ -300,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) } @@ -369,20 +375,24 @@ 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 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 632b829b161a0..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; @@ -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,17 @@ use object_store::ObjectStore; use std::collections::HashMap; use std::sync::Arc; +/// 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)] pub struct ListFilesResult { @@ -627,16 +639,14 @@ impl TableProvider for ListingTable { None }; - let Some(object_store_url) = - self.table_paths.first().map(ListingTableUrl::object_store) - 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) @@ -704,13 +714,24 @@ 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(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( state, store.as_ref(), - table_path, + rebased, &[], &self.options.file_extension, &self.options.table_partition_cols, @@ -724,8 +745,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); } @@ -733,8 +754,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(), @@ -785,15 +806,55 @@ 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. + /// + /// 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, + ) -> datafusion_common::Result> { + let Some(first) = self.table_paths.first() else { + return Ok(None); + }; + let env = ctx.runtime_env(); + let resolved = first.resolve(env)?; + 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>( &'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 +901,7 @@ impl ListingTable { ); } - let store = if let Some(url) = self.table_paths.first() { - ctx.runtime_env().object_store(url)? - } else { + let Some(resolved) = self.resolve_object_store(ctx)? else { return Ok(ListFilesResult { file_groups: vec![], statistics: Statistics::new_unknown(&self.file_schema), @@ -850,7 +909,13 @@ impl ListingTable { }); }; let (file_group, inexact_stats) = self - .collect_files_for_scan(ctx, &store, 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 @@ -902,17 +967,16 @@ impl ListingTable { ); } - let store = if let Some(url) = self.table_paths.first() { - ctx.runtime_env().object_store(url)? - } else { + let Some(resolved) = 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, &resolved.store, &resolved.table_urls, &[], 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..968a4dc62cfe4 100644 --- a/datafusion/core/src/datasource/listing/table.rs +++ b/datafusion/core/src/datasource/listing/table.rs @@ -51,17 +51,16 @@ 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 resolved = first.resolve(state.runtime_env())?; - let file = self - .table_paths - .first() - .unwrap() - .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/core/tests/datasource/object_store_access.rs b/datafusion/core/tests/datasource/object_store_access.rs index 2503de862e06a..0d074208b7c05 100644 --- a/datafusion/core/tests/datasource/object_store_access.rs +++ b/datafusion/core/tests/datasource/object_store_access.rs @@ -920,6 +920,222 @@ 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" + ); +} + +/// 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}" + ); +} + +/// 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-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..7a7f83b1b8242 100644 --- a/datafusion/datasource/src/mod.rs +++ b/datafusion/datasource/src/mod.rs @@ -569,6 +569,7 @@ mod tests { } #[test] + #[expect(deprecated)] fn test_object_store_listing_url() { let listing = ListingTableUrl::parse("file:///").unwrap(); let store = listing.object_store(); diff --git a/datafusion/datasource/src/url.rs b/datafusion/datasource/src/url.rs index 7985a29e4fd94..e486838bafeed 100644 --- a/datafusion/datasource/src/url.rs +++ b/datafusion/datasource/src/url.rs @@ -21,6 +21,7 @@ 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::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)] @@ -320,11 +339,55 @@ impl ListingTableUrl { } /// 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`]. + /// + /// 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 = `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: 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 = + ObjectStoreUrl::parse(format!("{authority}/{}", segments[..kept].join("/")))?; + + 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 pub fn is_folder(&self) -> bool { self.url.scheme() == "file" && self.is_collection() @@ -515,7 +578,9 @@ mod tests { use datafusion_common::config::TableOptions; use datafusion_execution::TaskContext; use datafusion_execution::config::SessionConfig; - 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::{ @@ -629,6 +694,74 @@ mod tests { assert!(url.strip_prefix(&path).is_none()); } + 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()), + ); + 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 env = env_with("s3://bucket"); + let url = ListingTableUrl::parse("s3://bucket/foo/bar/file.parquet").unwrap(); + 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 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 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 env = env_with("s3://bucket/user/repo"); + let url = ListingTableUrl::parse("s3://bucket/user/repo/data/").unwrap(); + 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 env = env_with("s3://bucket/user/repo"); + let url = ListingTableUrl::parse("s3://bucket/user/repo/").unwrap(); + 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_identity() { + // The retained (registered-prefix) segments keep their original encoding. + let env = env_with("s3://bucket/us%20er"); + let url = ListingTableUrl::parse("s3://bucket/us%20er/data.parquet").unwrap(); + 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] 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..599caf8991c42 100644 --- a/datafusion/execution/src/object_store.rs +++ b/datafusion/execution/src/object_store.rs @@ -19,14 +19,16 @@ //! 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 std::sync::Arc; +use object_store::path::Path; +use object_store::registry::ObjectStoreRegistry as UpstreamObjectStoreRegistry; +use std::collections::HashSet; +use std::sync::{Arc, RwLock}; use url::Url; /// A parsed URL identifying a particular [`ObjectStore`] instance @@ -59,15 +61,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,27 +182,61 @@ 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>, + /// 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 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. + 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 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() - } +/// 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`) +/// +/// 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. + /// + /// `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)>>, } impl Default for DefaultObjectStoreRegistry { @@ -204,73 +245,100 @@ impl Default for DefaultObjectStoreRegistry { } } +/// 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 object_stores: DashMap> = DashMap::new(); - object_stores.insert("file://".to_string(), Arc::new(LocalFileSystem::new())); - Self { object_stores } + 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()), + ); + this } /// Default without any backend registered. #[cfg(target_arch = "wasm32")] pub fn new() -> Self { - let object_stores: DashMap> = DashMap::new(); - Self { object_stores } + 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) + }) } } -/// -/// 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.registered + .write() + .unwrap() + .insert(registration_key(url)); + 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) + 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> { - 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)> { + // 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))) + } } #[cfg(test)] @@ -291,44 +359,151 @@ 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 from the returned path. + 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://"); + 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)); + } - let url = ObjectStoreUrl::parse("s3://bucket").unwrap(); - let key = get_url_key(&url.url); - assert_eq!(key.as_str(), "s3://bucket"); + #[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() + ); - 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"); + // 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() + ); } } 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 =