Support prefix-scoped object stores in the registry

Cargo.toml

@@ -267,6 +267,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

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::<Int64Array>()
+            .unwrap()
+            .value(0);
+        assert_eq!(n, 25);
+
+        Ok(())
+    }
     #[tokio::test]
     async fn copy_to_external_object_store_test() -> Result<()> {
         let aws_envs = vec![

datafusion-cli/src/object_storage/instrumented.rs

@@ -826,6 +826,13 @@ impl ObjectStoreRegistry for InstrumentedObjectStoreRegistry {
     fn get_store(&self, url: &Url) -> datafusion::common::Result<Arc<dyn ObjectStore>> {
         self.inner.get_store(url)
     }
+
+    fn resolve(
+        &self,
+        url: &Url,
+    ) -> datafusion::common::Result<(Arc<dyn ObjectStore>, Path)> {
+        self.inner.resolve(url)
+    }
 }
 
 #[cfg(test)]

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<Arc<dyn ObjectStore>> {
+        // 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)

datafusion/catalog-listing/src/options.rs

@@ -277,10 +277,13 @@ impl ListingOptions {
         state: &dyn Session,
         table_path: &'a ListingTableUrl,
     ) -> datafusion_common::Result<SchemaRef> {
-        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<Vec<String>> {
-        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()