diff --git a/vortex-cloud/src/hf/mod.rs b/vortex-cloud/src/hf/mod.rs index 4d25978537c..9a8caff8246 100644 --- a/vortex-cloud/src/hf/mod.rs +++ b/vortex-cloud/src/hf/mod.rs @@ -8,7 +8,7 @@ //! `resolve` prefix, carrying a bearer token when one is available, is the whole implementation. //! Reads therefore keep every [`object_store::ClientOptions`] setting a caller passes — //! connect/request timeouts, retries, proxy configuration, `allow_http` — unlike the OpenDAL-backed -//! schemes in [`crate::opendal`], whose bridge owns its own HTTP client. +//! schemes in the crate's `opendal` module, whose bridge owns its own HTTP client. //! //! # URL grammar //! diff --git a/vortex-duckdb/Cargo.toml b/vortex-duckdb/Cargo.toml index 35e2f904c49..f7512861cff 100644 --- a/vortex-duckdb/Cargo.toml +++ b/vortex-duckdb/Cargo.toml @@ -40,6 +40,7 @@ tracing-subscriber = { workspace = true } url = { workspace = true } vortex = { workspace = true, features = [ "files", + "hf", "tokio", "object_store", "object_store_registry", diff --git a/vortex-duckdb/src/multi_file.rs b/vortex-duckdb/src/multi_file.rs index e118dcaa8f4..b8ddc076b8e 100644 --- a/vortex-duckdb/src/multi_file.rs +++ b/vortex-duckdb/src/multi_file.rs @@ -18,7 +18,6 @@ use vortex::io::filesystem::FileSystemRef; use vortex::io::object_store::ObjectStoreFileSystem; use vortex::io::runtime::BlockingRuntime; use vortex::layout::scan::multi::MultiLayoutDataSource; -use vortex_utils::aliases::hash_map::HashMap; use crate::RUNTIME; use crate::SESSION; @@ -28,23 +27,34 @@ use crate::duckdb::ExtractedValue; /// Process-wide registry, so repeated scans against the same bucket share one client. static REGISTRY: LazyLock = LazyLock::new(Registry::new); -fn resolve_filesystem(base_url: &Url) -> VortexResult { +fn resolve_filesystem(glob_url: &Url) -> VortexResult<(FileSystemRef, String)> { // Compat makes us use tokio which is very bad for local reads on // high-core machines because reads go into blocking pool - if base_url.scheme() == "file" { - return Ok(Arc::new(ObjectStoreFileSystem::local(RUNTIME.handle()))); + if glob_url.scheme() == "file" { + return Ok(( + Arc::new(ObjectStoreFileSystem::local(RUNTIME.handle())), + glob_url.path().to_string(), + )); } - // `base_url` has its path cleared by the caller, so the resolved path is empty and only the - // store matters here. Going through the shared registry means DuckDB resolves the same set of - // schemes as the Python and Java bindings, including the OpenDAL-backed ones when the - // `opendal` feature is on. - let (object_store, _) = REGISTRY.resolve(base_url)?; + // The full URL goes through the shared registry, which reports the glob as a path *within* + // the store it returns. For most schemes the store is mounted at the URL authority, so the + // path is the whole URL path — but not for all of them: an `hf://` store is rooted at a + // repository and revision, which occupy path segments. Only the registry knows how deep the + // store is mounted, so globbing anything other than the path it reports would address the + // wrong keys. Going through the registry also means DuckDB resolves the same set of schemes + // as the Python and Java bindings, including the OpenDAL-backed ones when the `opendal` + // feature is on. The registry caches one client per store prefix, so repeated scans against + // the same bucket or repository share a client even though the filesystem wrapper is rebuilt. + let (object_store, path) = REGISTRY.resolve(glob_url)?; - Ok(Arc::new(ObjectStoreFileSystem::new( - Arc::new(Compat::new(object_store)), - RUNTIME.handle(), - ))) + Ok(( + Arc::new(ObjectStoreFileSystem::new( + Arc::new(Compat::new(object_store)), + RUNTIME.handle(), + )), + path.to_string(), + )) } /// Shared bind logic for both single-glob and multi-glob variants. @@ -77,28 +87,16 @@ pub fn bind_multi_file_scan(input: &BindInputRef) -> VortexResult = HashMap::new(); - for glob_url in &glob_urls { - let mut base_url = glob_url.clone(); - base_url.set_path(""); - if !fs_cache.contains_key(&base_url) { - let fs = resolve_filesystem(&base_url)?; - fs_cache.insert(base_url, fs); - } - } + let resolved = glob_urls + .iter() + .map(resolve_filesystem) + .collect::>>()?; RUNTIME.block_on(async { let mut builder = MultiFileDataSource::new(SESSION.clone()); - for glob_url in &glob_urls { - let mut base_url = glob_url.clone(); - base_url.set_path(""); - let fs = fs_cache - .get(&base_url) - .map(Arc::clone) - .unwrap_or_else(|| unreachable!("fs should be cached for all base URLs")); - builder = builder.with_glob(glob_url.path(), Some(fs)); + for (fs, glob) in resolved { + builder = builder.with_glob(&glob, Some(fs)); } builder.build().await