From 80ae1b0b0bc09462ba9959d938d19d938cf3a73f Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Tue, 4 Aug 2026 15:28:38 +0000 Subject: [PATCH 1/3] feat(cloud): serve Hugging Face Hub `hf://` URLs from the object store registry Vortex could not read from the Hugging Face Hub. `hf://` was understood only by `vortex.datasets`, which parsed the URI in Python and rewrote it into an authenticated `HTTPStore`, so `vx.open("hf://...")` failed and Polars, DuckDB and DataFusion could not open Hub URLs at all. A Hub repository is a set of files behind an HTTP endpoint that honours range requests, so serving one needs no cloud SDK: an `object_store::http::HttpStore` rooted at the repository's `resolve` prefix, carrying a bearer token when one is available, is the whole implementation. Reads therefore keep every `ClientOptions` setting a caller passes -- timeouts, retries, proxy configuration, `allow_http` -- unlike the OpenDAL-backed schemes, whose bridge owns its own HTTP client. Unlike every other scheme in the registry, the store is rooted at a (repository, revision) pair rather than at the URL authority, since both occupy URL path segments. It therefore reports the in-repository path itself and lets the registry mount it that deep; getting this wrong would send the repository name to the Hub as part of the file path, so it is covered from both the build-and-cache and the cached-store branch. Listing is not supported. The Hub does not implement WebDAV `PROPFIND`, which is how `object_store`'s HTTP store lists a prefix. Opening a known path -- `head` plus ranged `get`, both plain HTTP -- is what a Vortex scan needs, and callers that must expand a glob can list through the Hub's own API first. The `hf` feature adds no cloud SDK, only `http` and `percent-encoding`, so `vortex-python` enables it unconditionally rather than making it opt-in the way `opendal` is. Towards #5379. Signed-off-by: Robert Kruszewski Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01Lc5zw7Le2T3pakDEUdKTYd --- Cargo.lock | 4 + Cargo.toml | 2 + vortex-cloud/Cargo.toml | 9 + vortex-cloud/src/hf/mod.rs | 351 +++++++++++++++++++++++++++++ vortex-cloud/src/hf/tests.rs | 300 ++++++++++++++++++++++++ vortex-cloud/src/lib.rs | 9 +- vortex-cloud/src/registry/mod.rs | 17 +- vortex-cloud/src/registry/tests.rs | 42 ++++ vortex-python/Cargo.toml | 1 + vortex/Cargo.toml | 3 + 10 files changed, 734 insertions(+), 4 deletions(-) create mode 100644 vortex-cloud/src/hf/mod.rs create mode 100644 vortex-cloud/src/hf/tests.rs diff --git a/Cargo.lock b/Cargo.lock index 0d3665511c2..bf2d4c49ef4 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -9659,10 +9659,14 @@ dependencies = [ name = "vortex-cloud" version = "0.1.0" dependencies = [ + "http", "object_store", "object_store_opendal", "opendal", "parking_lot", + "percent-encoding", + "rstest", + "tempfile", "tracing", "url", "vortex-utils", diff --git a/Cargo.toml b/Cargo.toml index d096bce0c35..ebd88594ca9 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -170,6 +170,7 @@ glob = "0.3.2" goldenfile = "1" half = { version = "2.7.1", features = ["std", "num-traits"] } hashbrown = "0.17.1" +http = "1.5.0" humansize = "2.1.3" indicatif = "0.18.0" insta = "1.43" @@ -207,6 +208,7 @@ parquet-variant = "58.3" parquet-variant-compute = "58.3" paste = "1.0.15" pco = "1.0.1" +percent-encoding = "2.3.2" pin-project-lite = "0.2.15" primitive-types = { version = "0.14.0" } proc-macro2 = "1.0.95" diff --git a/vortex-cloud/Cargo.toml b/vortex-cloud/Cargo.toml index a10b9f1868f..ea8b5648a57 100644 --- a/vortex-cloud/Cargo.toml +++ b/vortex-cloud/Cargo.toml @@ -17,14 +17,20 @@ categories = { workspace = true } all-features = true [dependencies] +http = { workspace = true, optional = true } object_store = { workspace = true, features = ["fs"] } object_store_opendal = { workspace = true, optional = true } opendal = { workspace = true, optional = true } parking_lot = { workspace = true, optional = true } +percent-encoding = { workspace = true, optional = true } tracing = { workspace = true, optional = true } url = { workspace = true } vortex-utils = { workspace = true } +[dev-dependencies] +rstest = { workspace = true } +tempfile = { workspace = true } + [features] default = [] # The URL -> ObjectStore registry, plus the cloud backends it resolves URLs to. Kept optional so @@ -36,6 +42,9 @@ registry = [ "object_store/gcp", "object_store/http", ] +# The Hugging Face Hub, the `hf://` scheme. Served over `object_store`'s HTTP store, so it adds no +# cloud SDK of its own. +hf = ["dep:http", "dep:percent-encoding", "object_store/http"] # Tencent Cloud COS, the `cos://` scheme. cos = [ "dep:opendal", diff --git a/vortex-cloud/src/hf/mod.rs b/vortex-cloud/src/hf/mod.rs new file mode 100644 index 00000000000..e39a7045bbd --- /dev/null +++ b/vortex-cloud/src/hf/mod.rs @@ -0,0 +1,351 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! The Hugging Face Hub, served over [`object_store`]'s HTTP store. +//! +//! A Hub repository is a set of files behind an HTTP endpoint that honours range requests, so it +//! needs no cloud SDK of its own: a [`object_store::http::HttpStore`] rooted at the repository's +//! `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. +//! +//! # URL grammar +//! +//! ```text +//! hf://datasets//[@][/] +//! hf://spaces//[@][/] +//! hf:///[@][/] # a model repository +//! ``` +//! +//! matching the paths [`huggingface_hub`'s `HfFileSystem`][hffs] accepts. A revision containing `/` +//! (e.g. `refs/convert/parquet`) must be percent-encoded in the URL, as it must be there: +//! +//! ```text +//! hf://datasets/org/name@refs%2Fconvert%2Fparquet/data/train.vortex +//! ``` +//! +//! [hffs]: https://huggingface.co/docs/huggingface_hub/guides/hf_file_system +//! +//! # Configuration +//! +//! * `HF_TOKEN` — the API token. Falling back, as `huggingface_hub` does, to the token file at +//! `HF_TOKEN_PATH`, then `$HF_HOME/token`, then `$HOME/.cache/huggingface/token`. Without a token +//! the store reads anonymously, which is all a public repository needs. +//! * `HF_ENDPOINT` — the Hub endpoint, defaulting to `https://huggingface.co`. +//! +//! # Limitations +//! +//! The Hub does not implement WebDAV `PROPFIND`, which is how [`object_store`]'s HTTP store lists a +//! prefix, so [`object_store::ObjectStore::list`] fails against these URLs. Opening a file by name +//! works — that is `head` plus ranged `get`, both plain HTTP — so this serves reads of a known path. +//! Callers that need to expand a glob must list the repository through the Hub's own API first and +//! then open each returned path. + +use std::sync::Arc; + +use http::HeaderMap; +use http::HeaderValue; +use http::header::AUTHORIZATION; +use object_store::ClientOptions; +use object_store::ObjectStore; +use object_store::http::HttpBuilder; +use object_store::path::Path; +use percent_encoding::AsciiSet; +use percent_encoding::CONTROLS; +use percent_encoding::percent_decode_str; +use percent_encoding::utf8_percent_encode; +use url::Url; + +/// The URL scheme served by the Hugging Face Hub. +pub const HF_SCHEME: &str = "hf"; + +/// The Hub endpoint used when `HF_ENDPOINT` is unset. +const DEFAULT_ENDPOINT: &str = "https://huggingface.co"; + +/// The revision used when a URL does not name one. +const DEFAULT_REVISION: &str = "main"; + +const ENDPOINT_VAR: &str = "HF_ENDPOINT"; +const TOKEN_VAR: &str = "HF_TOKEN"; +const TOKEN_PATH_VAR: &str = "HF_TOKEN_PATH"; +const HF_HOME_VAR: &str = "HF_HOME"; +const HOME_VAR: &str = "HOME"; + +/// The URL authority that marks a dataset repository. +const DATASETS_HOST: &str = "datasets"; +/// The URL authority that marks a Space repository. +const SPACES_HOST: &str = "spaces"; + +/// Characters escaped when a revision is spliced into the `resolve` path. +/// +/// `/` is the point of this: the Hub routes `refs/convert/parquet` only in its `%2F` form, since an +/// unescaped `/` would read as another path segment. +const REVISION_ESCAPES: &AsciiSet = &CONTROLS.add(b'/').add(b'%').add(b'?').add(b'#').add(b' '); + +/// Whether `scheme` is served by this module. +/// +/// Callers dispatching on a URL scheme should ask this rather than comparing against [`HF_SCHEME`] +/// themselves, matching how [`crate::opendal::supports_scheme`] is used. +pub fn supports_scheme(scheme: &str) -> bool { + scheme == HF_SCHEME +} + +/// Which kind of Hub repository a URL addresses. +/// +/// The Hub routes each kind under its own path prefix, which is the only thing that differs between +/// them here. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum HfRepoType { + /// A dataset repository, `hf://datasets//`. + Dataset, + /// A model repository, `hf:///`. + Model, + /// A Space repository, `hf://spaces//`. + Space, +} + +impl HfRepoType { + /// The path prefix, including its trailing `/`, that the Hub routes this kind under. + fn url_prefix(self) -> &'static str { + match self { + HfRepoType::Dataset => "datasets/", + HfRepoType::Model => "", + HfRepoType::Space => "spaces/", + } + } +} + +/// Error type for building a Hugging Face Hub object store. +#[derive(Debug)] +pub enum HfStoreError { + /// The URL scheme is not one this module handles. + UnsupportedScheme(String), + /// The URL is not a well-formed `hf://` URL. + InvalidUrl(String), + /// The bearer token could not be used as an HTTP header value. + InvalidToken(http::header::InvalidHeaderValue), + /// The underlying HTTP store rejected the configuration. + Build(object_store::Error), +} + +impl std::fmt::Display for HfStoreError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + HfStoreError::UnsupportedScheme(s) => write!(f, "unsupported Hugging Face scheme: {s}"), + HfStoreError::InvalidUrl(url) => write!( + f, + "invalid Hugging Face URL {url}: expected hf://datasets//[@revision][/path], \ + hf://spaces//[@revision][/path] or hf:///[@revision][/path]" + ), + HfStoreError::InvalidToken(e) => write!(f, "invalid Hugging Face token: {e}"), + HfStoreError::Build(e) => write!(f, "failed to build Hugging Face store: {e}"), + } + } +} + +impl std::error::Error for HfStoreError { + fn source(&self) -> Option<&(dyn std::error::Error + 'static)> { + match self { + HfStoreError::InvalidToken(e) => Some(e), + HfStoreError::Build(e) => Some(e), + HfStoreError::UnsupportedScheme(_) | HfStoreError::InvalidUrl(_) => None, + } + } +} + +impl From for object_store::Error { + fn from(e: HfStoreError) -> Self { + object_store::Error::Generic { + store: "HuggingFace", + source: Box::new(e), + } + } +} + +/// Configuration for a store serving one Hub repository at one revision. +/// +/// A store covers a single `(repository, revision)` pair because that pair is what the Hub's +/// `resolve` prefix names; the object key within the store is then the in-repository path. +#[derive(Debug, Clone)] +pub struct HfConfig { + /// Which kind of repository to address. + pub repo_type: HfRepoType, + /// The repository, as `/`. + pub repo_id: String, + /// The revision to read: a branch, tag or commit. Held decoded, so `refs/convert/parquet` is + /// spelled with `/`; it is percent-encoded when the `resolve` URL is built. + pub revision: String, + /// Bearer token for private and gated repositories. `None` reads anonymously. + pub token: Option, + /// The Hub endpoint, e.g. `https://huggingface.co`. + pub endpoint: String, + /// HTTP client configuration, applied to every read. + pub client_options: ClientOptions, +} + +impl Default for HfConfig { + fn default() -> Self { + Self { + repo_type: HfRepoType::Dataset, + repo_id: String::new(), + revision: DEFAULT_REVISION.to_string(), + token: None, + endpoint: DEFAULT_ENDPOINT.to_string(), + client_options: ClientOptions::default(), + } + } +} + +impl HfConfig { + /// The URL this repository's files hang off, which is where the store is rooted. + fn resolve_prefix(&self) -> String { + let endpoint = self.endpoint.trim_end_matches('/'); + let prefix = self.repo_type.url_prefix(); + let revision = utf8_percent_encode(&self.revision, REVISION_ESCAPES); + format!("{endpoint}/{prefix}{}/resolve/{revision}", self.repo_id) + } +} + +/// Build an [`ObjectStore`] for a Hugging Face Hub repository from an [`HfConfig`]. +/// +/// This is the entry point for callers holding a strongly-typed configuration. The returned store is +/// rooted at the repository's `resolve` prefix, so its object keys are in-repository paths. +pub fn make_hf_store(config: HfConfig) -> Result, HfStoreError> { + if config.repo_id.is_empty() { + return Err(HfStoreError::InvalidUrl("".to_string())); + } + + let mut client_options = config.client_options.clone(); + if let Some(token) = config.token.as_deref() { + let mut headers = HeaderMap::new(); + let value = HeaderValue::from_str(&format!("Bearer {token}")) + .map_err(HfStoreError::InvalidToken)?; + // Sending the token as a default header, rather than per request, is what lets a plain + // `HttpStore` serve gated and private repositories. + headers.insert(AUTHORIZATION, value); + client_options = client_options.with_default_headers(headers); + } + + let store = HttpBuilder::new() + .with_url(config.resolve_prefix()) + .with_client_options(client_options) + .build() + .map_err(HfStoreError::Build)?; + + Ok(Arc::new(store)) +} + +/// Build an [`ObjectStore`] for an `hf://` URL, with the in-repository path of the addressed file. +/// +/// Configuration not carried by the URL is read from the process environment. The path is reported +/// the way [`object_store::parse_url_opts`] reports it — the key the returned store will see — so +/// that the caller can derive how deep into the URL the store is mounted. +pub fn make_hf_store_from_url(url: &Url) -> Result<(Arc, Path), HfStoreError> { + make_hf_store_from_url_with_env(url, |key| std::env::var(key).ok()) +} + +/// [`make_hf_store_from_url`], reading environment configuration through `env_lookup`. +/// +/// Tests pass a closure over a fixed map so they do not race against the process environment. +pub(crate) fn make_hf_store_from_url_with_env( + url: &Url, + env_lookup: F, +) -> Result<(Arc, Path), HfStoreError> +where + F: Fn(&str) -> Option, +{ + let (config, path) = url_to_config(url, env_lookup)?; + Ok((make_hf_store(config)?, path)) +} + +/// Translate an `hf://` URL into an [`HfConfig`] plus the in-repository path it addresses. +fn url_to_config(url: &Url, env_lookup: F) -> Result<(HfConfig, Path), HfStoreError> +where + F: Fn(&str) -> Option, +{ + if !supports_scheme(url.scheme()) { + return Err(HfStoreError::UnsupportedScheme(url.scheme().to_string())); + } + let invalid = || HfStoreError::InvalidUrl(url.to_string()); + + // `hf://datasets/org/name/...` parses with `datasets` as the authority and `/org/name/...` as + // the path, so the repository kind is the authority and a model repository's owner is too. + let host = url.host_str().ok_or_else(invalid)?; + let segments: Vec<&str> = url.path().split('/').filter(|s| !s.is_empty()).collect(); + + let (repo_type, owner, rest) = match host { + DATASETS_HOST | SPACES_HOST => { + let repo_type = if host == DATASETS_HOST { + HfRepoType::Dataset + } else { + HfRepoType::Space + }; + let (owner, rest) = segments.split_first().ok_or_else(invalid)?; + (repo_type, *owner, rest) + } + owner => (HfRepoType::Model, owner, segments.as_slice()), + }; + + let (name, rest) = rest.split_first().ok_or_else(invalid)?; + if owner.is_empty() || name.is_empty() { + return Err(invalid()); + } + + // A revision is appended to the repository name with `@`, and arrives percent-encoded when it + // contains `/`. Hold it decoded; `resolve_prefix` re-encodes it on the way back out. + let (name, revision) = match name.split_once('@') { + Some((name, revision)) => { + if name.is_empty() || revision.is_empty() { + return Err(invalid()); + } + let revision = percent_decode_str(revision) + .decode_utf8() + .map_err(|_| invalid())? + .into_owned(); + (name, revision) + } + None => (*name, DEFAULT_REVISION.to_string()), + }; + + let config = HfConfig { + repo_type, + repo_id: format!("{owner}/{name}"), + revision, + token: resolve_token(&env_lookup), + endpoint: env_lookup(ENDPOINT_VAR) + .filter(|endpoint| !endpoint.is_empty()) + .unwrap_or_else(|| DEFAULT_ENDPOINT.to_string()), + client_options: ClientOptions::default(), + }; + + // The segments still carry their URL escapes, so decode them into the object key rather than + // joining them raw. + let path = Path::from_url_path(rest.join("/")).map_err(|_| invalid())?; + + Ok((config, path)) +} + +/// The bearer token to read with, following the same precedence as `huggingface_hub.get_token()`: +/// `HF_TOKEN`, then the token file at `HF_TOKEN_PATH`, `$HF_HOME/token` or +/// `$HOME/.cache/huggingface/token`. +fn resolve_token(env_lookup: &F) -> Option +where + F: Fn(&str) -> Option, +{ + if let Some(token) = env_lookup(TOKEN_VAR).filter(|token| !token.is_empty()) { + return Some(token); + } + + let token_path = env_lookup(TOKEN_PATH_VAR) + .or_else(|| env_lookup(HF_HOME_VAR).map(|home| format!("{home}/token"))) + .or_else(|| env_lookup(HOME_VAR).map(|home| format!("{home}/.cache/huggingface/token")))?; + + // A missing token file is the ordinary anonymous case, not an error. + let token = std::fs::read_to_string(token_path).ok()?; + let token = token.trim(); + (!token.is_empty()).then(|| token.to_string()) +} + +#[cfg(test)] +mod tests; diff --git a/vortex-cloud/src/hf/tests.rs b/vortex-cloud/src/hf/tests.rs new file mode 100644 index 00000000000..cc3b9c86c5a --- /dev/null +++ b/vortex-cloud/src/hf/tests.rs @@ -0,0 +1,300 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use rstest::rstest; +use url::Url; + +use super::DEFAULT_ENDPOINT; +use super::HfConfig; +use super::HfRepoType; +use super::HfStoreError; +use super::url_to_config; + +/// Configuration lookups over a fixed map, so these tests neither read nor mutate the process +/// environment. +fn env(vars: &[(&str, &str)]) -> impl Fn(&str) -> Option + use<> { + let vars: Vec<(String, String)> = vars + .iter() + .map(|(k, v)| ((*k).to_string(), (*v).to_string())) + .collect(); + move |key| { + vars.iter() + .find(|(k, _)| k == key) + .map(|(_, v)| v.to_string()) + } +} + +/// Parse a URL with no environment configuration at all, so no token is found and the default +/// endpoint applies. +fn parse(url: &str) -> Result<(HfConfig, String), Box> { + let url = Url::parse(url)?; + let (config, path) = url_to_config(&url, env(&[]))?; + Ok((config, path.as_ref().to_string())) +} + +#[rstest] +// A dataset repository, with and without an explicit revision. +#[case("hf://datasets/org/name", HfRepoType::Dataset, "org/name", "main", "")] +#[case( + "hf://datasets/org/name@main", + HfRepoType::Dataset, + "org/name", + "main", + "" +)] +#[case( + "hf://datasets/org/name/train.vortex", + HfRepoType::Dataset, + "org/name", + "main", + "train.vortex" +)] +#[case( + "hf://datasets/org/name/data/nested/train.vortex", + HfRepoType::Dataset, + "org/name", + "main", + "data/nested/train.vortex" +)] +#[case( + "hf://datasets/org/name@v1.0/train.vortex", + HfRepoType::Dataset, + "org/name", + "v1.0", + "train.vortex" +)] +// A revision containing `/` arrives percent-encoded and must be held decoded. +#[case( + "hf://datasets/org/name@refs%2Fconvert%2Fparquet/data/train.vortex", + HfRepoType::Dataset, + "org/name", + "refs/convert/parquet", + "data/train.vortex" +)] +// A bare owner/name is a model repository: the owner is the URL authority. +#[case( + "hf://org/name/model.vortex", + HfRepoType::Model, + "org/name", + "main", + "model.vortex" +)] +#[case( + "hf://org/name@abc123/model.vortex", + HfRepoType::Model, + "org/name", + "abc123", + "model.vortex" +)] +#[case( + "hf://spaces/org/name/app.vortex", + HfRepoType::Space, + "org/name", + "main", + "app.vortex" +)] +fn test_url_to_config( + #[case] url: &str, + #[case] repo_type: HfRepoType, + #[case] repo_id: &str, + #[case] revision: &str, + #[case] path: &str, +) -> Result<(), Box> { + let (config, parsed_path) = parse(url)?; + + assert_eq!(config.repo_type, repo_type); + assert_eq!(config.repo_id, repo_id); + assert_eq!(config.revision, revision); + assert_eq!(parsed_path, path); + Ok(()) +} + +#[rstest] +// A repository needs both an owner and a name. +#[case("hf://datasets/name-only")] +#[case("hf://org")] +// An `@` must have a name on the left and a revision on the right. +#[case("hf://datasets/org/@main")] +#[case("hf://datasets/org/name@")] +fn test_url_to_config_invalid(#[case] url: &str) -> Result<(), Box> { + let url = Url::parse(url)?; + assert!(matches!( + url_to_config(&url, env(&[])), + Err(HfStoreError::InvalidUrl(_)) + )); + Ok(()) +} + +#[test] +fn test_url_to_config_rejects_other_schemes() -> Result<(), Box> { + let url = Url::parse("s3://bucket/key.vortex")?; + assert!(matches!( + url_to_config(&url, env(&[])), + Err(HfStoreError::UnsupportedScheme(_)) + )); + Ok(()) +} + +#[rstest] +// Each repository kind hangs off its own prefix, and the store is rooted at the revision. +#[case( + HfRepoType::Dataset, + "main", + "https://huggingface.co/datasets/org/name/resolve/main" +)] +#[case( + HfRepoType::Model, + "main", + "https://huggingface.co/org/name/resolve/main" +)] +#[case( + HfRepoType::Space, + "main", + "https://huggingface.co/spaces/org/name/resolve/main" +)] +// The Hub only routes the escaped form of a revision containing `/`. +#[case( + HfRepoType::Dataset, + "refs/convert/parquet", + "https://huggingface.co/datasets/org/name/resolve/refs%2Fconvert%2Fparquet" +)] +fn test_resolve_prefix( + #[case] repo_type: HfRepoType, + #[case] revision: &str, + #[case] expected: &str, +) { + let config = HfConfig { + repo_type, + repo_id: "org/name".to_string(), + revision: revision.to_string(), + ..HfConfig::default() + }; + + assert_eq!(config.resolve_prefix(), expected); +} + +/// A percent-encoded revision must survive the URL -> config -> URL round trip, since that is how a +/// `refs/convert/parquet` read reaches the Hub. +#[test] +fn test_slash_revision_round_trips() -> Result<(), Box> { + let (config, path) = + parse("hf://datasets/org/name@refs%2Fconvert%2Fparquet/data/train.vortex")?; + + assert_eq!(config.revision, "refs/convert/parquet"); + assert_eq!( + config.resolve_prefix(), + "https://huggingface.co/datasets/org/name/resolve/refs%2Fconvert%2Fparquet" + ); + assert_eq!(path, "data/train.vortex"); + Ok(()) +} + +#[test] +fn test_endpoint_override() -> Result<(), Box> { + let url = Url::parse("hf://datasets/org/name/train.vortex")?; + let (config, _path) = url_to_config(&url, env(&[("HF_ENDPOINT", "https://hub.example.com/")]))?; + + assert_eq!(config.endpoint, "https://hub.example.com/"); + // The trailing slash on the endpoint must not double up in the resolve prefix. + assert_eq!( + config.resolve_prefix(), + "https://hub.example.com/datasets/org/name/resolve/main" + ); + Ok(()) +} + +#[test] +fn test_default_endpoint_when_unset_or_empty() -> Result<(), Box> { + let url = Url::parse("hf://datasets/org/name")?; + + let (unset, _) = url_to_config(&url, env(&[]))?; + assert_eq!(unset.endpoint, DEFAULT_ENDPOINT); + + let (empty, _) = url_to_config(&url, env(&[("HF_ENDPOINT", "")]))?; + assert_eq!(empty.endpoint, DEFAULT_ENDPOINT); + Ok(()) +} + +#[test] +fn test_token_from_env() -> Result<(), Box> { + let url = Url::parse("hf://datasets/org/name")?; + let (config, _path) = url_to_config(&url, env(&[("HF_TOKEN", "hf_from_env")]))?; + + assert_eq!(config.token.as_deref(), Some("hf_from_env")); + Ok(()) +} + +#[test] +fn test_token_from_token_file() -> Result<(), Box> { + let dir = tempfile::tempdir()?; + let token_file = dir.path().join("token"); + // The saved login file is newline-terminated, which must not reach the header value. + std::fs::write(&token_file, "hf_from_file\n")?; + + let url = Url::parse("hf://datasets/org/name")?; + let (config, _path) = url_to_config( + &url, + env(&[("HF_TOKEN_PATH", &token_file.to_string_lossy())]), + )?; + + assert_eq!(config.token.as_deref(), Some("hf_from_file")); + Ok(()) +} + +#[test] +fn test_token_env_wins_over_file() -> Result<(), Box> { + let dir = tempfile::tempdir()?; + let token_file = dir.path().join("token"); + std::fs::write(&token_file, "hf_from_file")?; + + let url = Url::parse("hf://datasets/org/name")?; + let (config, _path) = url_to_config( + &url, + env(&[ + ("HF_TOKEN", "hf_from_env"), + ("HF_TOKEN_PATH", &token_file.to_string_lossy()), + ]), + )?; + + assert_eq!(config.token.as_deref(), Some("hf_from_env")); + Ok(()) +} + +/// No token anywhere is the ordinary anonymous case, not an error. +#[test] +fn test_no_token_reads_anonymously() -> Result<(), Box> { + let url = Url::parse("hf://datasets/org/name")?; + let (config, _path) = url_to_config(&url, env(&[("HF_TOKEN_PATH", "/nonexistent/token")]))?; + + assert!(config.token.is_none()); + Ok(()) +} + +/// Building a store must succeed for both the anonymous and the credentialed path; a token that +/// cannot become a header value is reported rather than panicking. +#[test] +fn test_make_hf_store() -> Result<(), Box> { + let base = HfConfig { + repo_id: "org/name".to_string(), + ..HfConfig::default() + }; + + let _anonymous = super::make_hf_store(base.clone())?; + let _credentialed = super::make_hf_store(HfConfig { + token: Some("hf_token".to_string()), + ..base.clone() + })?; + + assert!(matches!( + super::make_hf_store(HfConfig { + token: Some("bad\nvalue".to_string()), + ..base + }), + Err(HfStoreError::InvalidToken(_)) + )); + assert!(matches!( + super::make_hf_store(HfConfig::default()), + Err(HfStoreError::InvalidUrl(_)) + )); + Ok(()) +} diff --git a/vortex-cloud/src/lib.rs b/vortex-cloud/src/lib.rs index 112b7e4a698..0799fe5bbf4 100644 --- a/vortex-cloud/src/lib.rs +++ b/vortex-cloud/src/lib.rs @@ -12,6 +12,8 @@ //! * `opendal` supplies stores for cloud services the `object_store` crate does not implement //! natively — Tencent Cloud COS and Alibaba Cloud OSS — bridged through //! `object_store_opendal`. +//! * `hf` serves Hugging Face Hub repositories over `object_store`'s HTTP store, adding no cloud +//! SDK of its own. //! //! Every Vortex language binding resolves URLs through this one crate, so a scheme added here is //! reachable from Python, Java and DuckDB alike. @@ -23,13 +25,16 @@ //! //! * `registry` — the `Registry`, plus the natively-supported cloud backends (S3, Azure, GCS, //! HTTP) it resolves URLs to. +//! * `hf` — the Hugging Face Hub, the `hf://` scheme. //! * `cos` — Tencent Cloud COS, the `cos://` scheme. //! * `oss` — Alibaba Cloud OSS, the `oss://` scheme. //! * `opendal` — every OpenDAL-backed service above. //! -//! The `registry` feature picks up whichever OpenDAL services are enabled, so a consumer that -//! turns on `oss` gets `oss://` resolution without touching its own scheme matching. +//! The `registry` feature picks up whichever services are enabled, so a consumer that turns on +//! `oss` gets `oss://` resolution without touching its own scheme matching. +#[cfg(feature = "hf")] +pub mod hf; #[cfg(any(feature = "cos", feature = "oss"))] pub mod opendal; #[cfg(feature = "registry")] diff --git a/vortex-cloud/src/registry/mod.rs b/vortex-cloud/src/registry/mod.rs index 4b44e3b5bda..cac2e8c985b 100644 --- a/vortex-cloud/src/registry/mod.rs +++ b/vortex-cloud/src/registry/mod.rs @@ -10,7 +10,8 @@ //! various `Store::from_env` builders behave (see //! ); //! 2. schemes that `object_store` does not recognize natively — the OpenDAL-backed `cos://` and -//! `oss://` — are served by the crate's `opendal` module under the matching service feature. +//! `oss://`, and the Hugging Face Hub's `hf://` — are served by the crate's `opendal` and `hf` +//! modules under the matching service feature. use std::sync::Arc; @@ -101,7 +102,7 @@ enum EnvSource { impl EnvSource { /// Case-insensitive lookup of a single configuration variable. - #[cfg(any(feature = "cos", feature = "oss"))] + #[cfg(any(feature = "cos", feature = "oss", feature = "hf"))] fn lookup(&self, key: &str) -> Option { match self { EnvSource::Process => std::env::var(key).ok(), @@ -231,6 +232,18 @@ impl Registry { return Ok((store, Path::from_url_path(to_resolve.path())?)); } + // The Hugging Face Hub is not recognized by `object_store` either. Unlike the OpenDAL + // schemes its store is not rooted at the URL authority — a repository and revision occupy + // path segments too — so it reports the in-repository path itself and the registry mounts + // the store as deep as that implies. + #[cfg(feature = "hf")] + if crate::hf::supports_scheme(to_resolve.scheme()) { + return Ok(crate::hf::make_hf_store_from_url_with_env( + to_resolve, + |key| self.env.lookup(key), + )?); + } + let (store, path) = parse_url_opts(to_resolve, self.env.normalized_vars())?; Ok((Arc::from(store), path)) } diff --git a/vortex-cloud/src/registry/tests.rs b/vortex-cloud/src/registry/tests.rs index b15815f76f9..fb0bc5177b6 100644 --- a/vortex-cloud/src/registry/tests.rs +++ b/vortex-cloud/src/registry/tests.rs @@ -147,6 +147,48 @@ fn test_registered_store_wins_over_build() -> Result<(), Box Result<(), Box> { + let registry = registry(); + let url = Url::parse(url)?; + + // First resolution builds and caches the store; the second takes the cached-store branch. + // Both must report the same key, since the cached branch recomputes it from the mount depth. + let (_store, path) = registry.resolve(&url)?; + assert_eq!(path, Path::from(expected)); + let (_store, path) = registry.resolve(&url)?; + assert_eq!(path, Path::from(expected)); + Ok(()) +} + +/// Two revisions of one repository are different stores, since the revision is part of the prefix +/// the store is rooted at. +#[cfg(feature = "hf")] +#[test] +fn test_hf_revisions_do_not_share_a_store() -> Result<(), Box> { + let registry = registry(); + + let (main, _) = registry.resolve(&Url::parse("hf://datasets/org/name/train.vortex")?)?; + let (tagged, _) = registry.resolve(&Url::parse("hf://datasets/org/name@v2/train.vortex")?)?; + + assert!(!Arc::ptr_eq(&main, &tagged)); + Ok(()) +} + /// `parse_url_opts`, which does not recognize them. Without an endpoint the build fails, but the /// error must come from the OpenDAL builder ("missing required OpenDAL store configuration"), not /// from `object_store`'s unrecognized-scheme path. diff --git a/vortex-python/Cargo.toml b/vortex-python/Cargo.toml index 81ee7742890..7747ec5e1d5 100644 --- a/vortex-python/Cargo.toml +++ b/vortex-python/Cargo.toml @@ -55,6 +55,7 @@ pyo3-object_store = { workspace = true } tokio = { workspace = true, features = ["rt-multi-thread"], optional = true } url = { workspace = true } vortex = { workspace = true, features = [ + "hf", "object_store", "object_store_registry", ] } diff --git a/vortex/Cargo.toml b/vortex/Cargo.toml index 5b26944c4ed..54aa3f1e515 100644 --- a/vortex/Cargo.toml +++ b/vortex/Cargo.toml @@ -76,6 +76,9 @@ memmap2 = ["vortex-buffer/memmap2"] object_store = ["vortex-file?/object_store", "vortex-io/object_store"] # The URL -> ObjectStore registry, plus the cloud backends it resolves URLs to. object_store_registry = ["object_store", "vortex-cloud/registry"] +# The Hugging Face Hub (`hf://`) in the object store registry. Served over HTTP, so unlike +# `opendal` it pulls in no cloud SDK. +hf = ["object_store_registry", "vortex-cloud/hf"] # OpenDAL-backed schemes (`cos://`, `oss://`) in the object store registry. opendal = ["object_store_registry", "vortex-cloud/opendal"] tokio = [ From 1307b0ea81aab694146d28e36257123ec85dc147 Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Tue, 4 Aug 2026 15:45:59 +0000 Subject: [PATCH 2/3] refactor(python): drop the Hub URL and auth plumbing now that Vortex reads `hf://` `vortex.datasets` built Hub URLs and authorization headers by hand because the reader could not resolve `hf://` itself. It can now, so the URL building (`hf_hub_url`), the saved-login lookup (`get_token`), the endpoint-relative path rewriting, and the separate prefix-rooted store for revisions containing `/` all collapse into emitting `hf://` URIs and letting the reader resolve them. Two cases an `hf://` URI cannot express keep a store of their own, both rooted at the encoded `resolve` prefix: * a `token` string, which the reader has no way to see; * `token=False`, which must suppress the credentials the reader would otherwise read from `HF_TOKEN` or the saved login. Routing this through `hf://` would have quietly authenticated a read the caller asked to be anonymous, so it is covered by its own test. `token=True` asks for exactly the saved login the reader already finds, so it joins the default on the `hf://` path. Listing stays here. The Hub serves no listing over the object-store protocol, so expanding a glob or a directory still goes through the Hub API, which is why `_parse_hf_uri` and the `HfApi` call remain. File paths are percent-encoded on their way into the URI, since they become URI segments; Vortex decodes them back into the object key. Signed-off-by: Robert Kruszewski Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01Lc5zw7Le2T3pakDEUdKTYd --- vortex-python/python/vortex/datasets.py | 67 +++++++++++------------ vortex-python/test/test_hf_datasets.py | 73 +++++++++++++++++-------- 2 files changed, 81 insertions(+), 59 deletions(-) diff --git a/vortex-python/python/vortex/datasets.py b/vortex-python/python/vortex/datasets.py index faab4d28d43..d10bc66d721 100644 --- a/vortex-python/python/vortex/datasets.py +++ b/vortex-python/python/vortex/datasets.py @@ -49,7 +49,7 @@ get_format_type_from_alias, ) from datasets.table import InMemoryTable - from huggingface_hub import HfApi, get_token, hf_hub_url, snapshot_download + from huggingface_hub import HfApi, snapshot_download from huggingface_hub import constants as hf_hub_constants except ImportError as e: # pragma: no cover - exercised only without optional deps. raise ImportError("Install vortex-data[hf] to use vortex.datasets.") from e @@ -62,6 +62,7 @@ _DEFAULT_SPLIT = "train" _DEFAULT_DATA_FILES = "**/*.vortex" +_DEFAULT_REVISION = "main" _ITERABLE_DATASET_HAS_SHUFFLING = "shuffling" in inspect.signature(hf_datasets.IterableDataset).parameters @@ -99,9 +100,9 @@ def load_dataset( keeps Vortex in charge of reading and pushes column selection, Vortex expressions, and row limits into each Vortex scan before examples are yielded to Hugging Face Datasets transforms. Hub repositories are streamed directly with HTTP range requests instead of being downloaded; - private and gated repositories authenticate with ``token`` or the locally saved login. Pass - ``streaming=False`` to eagerly materialize an in-memory ``datasets.Dataset``, which downloads - Hub files first (as does ``local_files_only=True``). + private and gated repositories authenticate with ``token``, ``HF_TOKEN`` or the locally saved + login. Pass ``streaming=False`` to eagerly materialize an in-memory ``datasets.Dataset``, which + downloads Hub files first (as does ``local_files_only=True``). """ split_to_files, store = _resolve_data_files( @@ -846,12 +847,17 @@ def _resolve_hub_files( ) -> tuple[dict[str, list[str]], ObjectStore | None]: """Resolve Hub repository files to locations Vortex can stream without downloading them. - Anonymous reads use plain ``resolve`` URLs. When credentials are available (an explicit - ``token`` or the saved login) the reads must carry an authorization header, so the files are - returned as endpoint-relative paths together with an authenticated HTTP store. Revisions - containing ``/`` also require a store, rooted at the percent-encoded ``resolve`` prefix, - because the Hub only routes the encoded form and percent escapes cannot round-trip through - inferred object store paths. + Vortex reads ``hf://`` URIs itself — resolving the revision, percent-encoding it when it + contains ``/``, and authenticating from ``HF_TOKEN`` or the saved login — so for the default + ``token`` (and ``token=True``, which asks for exactly that saved login) the matched files are + returned as ``hf://`` URIs and need no store. + + The two cases the reader cannot express are handled with a store of its own: a ``token`` string, + which the reader has no way to see, and ``token=False``, which must suppress the credentials the + reader would otherwise pick up from the environment. + + The Hub serves no listing over the object-store protocol, so the patterns are expanded here + through the Hub API either way. """ repo_files = HfApi(token=token).list_repo_files(repo_id, repo_type="dataset", revision=revision) @@ -865,35 +871,24 @@ def _resolve_hub_files( ) split_to_matches[split_name] = matches - headers = _hub_auth_headers(token) - client_options: ClientConfig | None = {"default_headers": headers} if headers else None - - if revision is not None and "/" in revision: - base = f"{hf_hub_constants.ENDPOINT}/datasets/{repo_id}/resolve/{quote(revision, safe='')}" + if token is False or isinstance(token, str): + # The Hub only routes the percent-encoded form of a revision containing `/`, so the store is + # rooted at the encoded `resolve` prefix and the files stay in-repository paths. + revision_path = quote(revision if revision is not None else _DEFAULT_REVISION, safe="") + base = f"{hf_hub_constants.ENDPOINT}/datasets/{repo_id}/resolve/{revision_path}" + client_options: ClientConfig | None = ( + {"default_headers": {"authorization": f"Bearer {token}"}} if isinstance(token, str) else None + ) return split_to_matches, HTTPStore(base, client_options=client_options) - split_to_urls = { - name: [hf_hub_url(repo_id, file, repo_type="dataset", revision=revision) for file in files] - for name, files in split_to_matches.items() - } - if client_options is None: - return split_to_urls, None - - endpoint = hf_hub_constants.ENDPOINT + prefix = f"hf://datasets/{repo_id}" + if revision is not None: + prefix = f"{prefix}@{quote(revision, safe='')}" + # The file paths become URI segments, so escape them (Vortex percent-decodes them back into the + # object key); `/` stays literal because it separates the segments. return { - name: [url.removeprefix(endpoint).lstrip("/") for url in urls] for name, urls in split_to_urls.items() - }, HTTPStore(endpoint, client_options=client_options) - - -def _hub_auth_headers(token: bool | str | None) -> dict[str, str]: - """Authorization headers for Hub reads: the explicit token, or the saved login unless - ``token=False`` opts out. Empty when no credentials are available (anonymous access).""" - if token is False: - return {} - resolved = token if isinstance(token, str) else get_token() - if resolved is None: - return {} - return {"authorization": f"Bearer {resolved}"} + name: [f"{prefix}/{quote(file, safe='/')}" for file in files] for name, files in split_to_matches.items() + }, None def _is_url(path: str) -> bool: diff --git a/vortex-python/test/test_hf_datasets.py b/vortex-python/test/test_hf_datasets.py index 6642ab5a14c..dbbc0e01442 100644 --- a/vortex-python/test/test_hf_datasets.py +++ b/vortex-python/test/test_hf_datasets.py @@ -16,7 +16,6 @@ import pyarrow as pa import pytest import vortex.datasets as vx_datasets -from huggingface_hub import hf_hub_url from typing_extensions import override from vortex.store import HTTPStore @@ -422,7 +421,8 @@ def list_repo_files(self, repo_id: str, *, repo_type: str | None = None, revisio return list(self.repo_files) -def test_hub_streaming_resolves_to_urls_without_download(monkeypatch: pytest.MonkeyPatch): +@pytest.mark.parametrize("token", [None, True]) +def test_hub_streaming_resolves_to_hf_uris_without_download(token: bool | None, monkeypatch: pytest.MonkeyPatch): monkeypatch.setattr(vx_datasets, "HfApi", _FakeHfApi) files, store = vx_datasets._resolve_data_files( # pyright: ignore[reportPrivateUsage] @@ -430,17 +430,42 @@ def test_hub_streaming_resolves_to_urls_without_download(monkeypatch: pytest.Mon data_files=None, split="train", revision=None, - token=False, + token=token, cache_dir=None, local_files_only=False, streaming=True, ) + # Vortex resolves `hf://` itself, including the saved login that `token=True` asks for, so + # there is no store and no URL building here. assert store is None - expected = [ - hf_hub_url("org/name", file, repo_type="dataset") for file in ["data/validation.vortex", "train.vortex"] - ] - assert files == {"train": expected} + assert files == { + "train": [ + "hf://datasets/org/name/data/validation.vortex", + "hf://datasets/org/name/train.vortex", + ] + } + + +def test_hub_streaming_with_token_false_forces_anonymous_store(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(vx_datasets, "HfApi", _FakeHfApi) + + files, store = vx_datasets._resolve_data_files( # pyright: ignore[reportPrivateUsage] + "org/name", + data_files=None, + split="train", + revision=None, + token=False, + cache_dir=None, + local_files_only=False, + streaming=True, + ) + + # `token=False` must suppress the credentials Vortex would otherwise read from the environment, + # which an `hf://` URI cannot express, so it gets an unauthenticated store instead. + assert isinstance(store, HTTPStore) + assert store.url.endswith("/datasets/org/name/resolve/main") + assert files == {"train": ["data/validation.vortex", "train.vortex"]} def test_hub_streaming_with_token_uses_authenticated_store(monkeypatch: pytest.MonkeyPatch): @@ -457,13 +482,11 @@ def test_hub_streaming_with_token_uses_authenticated_store(monkeypatch: pytest.M streaming=True, ) + # An explicitly passed token cannot reach the Vortex reader, so this is the one path that still + # builds a store; the files are then relative to it. assert isinstance(store, HTTPStore) - assert files == { - "train": [ - "datasets/org/name/resolve/main/data/validation.vortex", - "datasets/org/name/resolve/main/train.vortex", - ] - } + assert store.url.endswith("/datasets/org/name/resolve/main") + assert files == {"train": ["data/validation.vortex", "train.vortex"]} @pytest.mark.parametrize( @@ -498,7 +521,7 @@ def resolve(path: str): data_files=None, split="train", revision=None, - token=False, + token=None, cache_dir=None, local_files_only=False, streaming=True, @@ -506,14 +529,14 @@ def resolve(path: str): files, store = resolve("hf://datasets/org/name/train.vortex") assert store is None - assert files == {"train": [hf_hub_url("org/name", "train.vortex", repo_type="dataset")]} + assert files == {"train": ["hf://datasets/org/name/train.vortex"]} # A glob-free path naming a directory selects the default Vortex files beneath it. files, _store = resolve("hf://datasets/org/name/data") - assert files == {"train": [hf_hub_url("org/name", "data/validation.vortex", repo_type="dataset")]} + assert files == {"train": ["hf://datasets/org/name/data/validation.vortex"]} -def test_hf_uri_slash_revision_uses_prefix_rooted_store(monkeypatch: pytest.MonkeyPatch): +def test_hf_uri_slash_revision_stays_percent_encoded(monkeypatch: pytest.MonkeyPatch): monkeypatch.setattr(vx_datasets, "HfApi", _FakeHfApi) files, store = vx_datasets._resolve_data_files( # pyright: ignore[reportPrivateUsage] @@ -521,17 +544,21 @@ def test_hf_uri_slash_revision_uses_prefix_rooted_store(monkeypatch: pytest.Monk data_files=None, split="train", revision=None, - token=False, + token=None, cache_dir=None, local_files_only=False, streaming=True, ) - # The Hub only routes percent-encoded revisions, so the store is rooted at the encoded - # resolve prefix and the files are in-repo paths. - assert isinstance(store, HTTPStore) - assert store.url.endswith("/datasets/org/name/resolve/refs%2Fconvert%2Fparquet") - assert files == {"train": ["data/validation.vortex", "train.vortex"]} + # A revision containing `/` needs no store of its own any more: Vortex percent-encodes it when + # it builds the `resolve` URL, so it only has to survive round-tripping through the `hf://` URI. + assert store is None + assert files == { + "train": [ + "hf://datasets/org/name@refs%2Fconvert%2Fparquet/data/validation.vortex", + "hf://datasets/org/name@refs%2Fconvert%2Fparquet/train.vortex", + ] + } def test_hf_uri_revision_conflict_raises(): From 80f90fc0109092d8344b0597294a4f191a3325c2 Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Tue, 4 Aug 2026 18:46:46 +0000 Subject: [PATCH 3/3] feat(python): add `vortex.store.HfStore` and accept every store in `vx.open` `vortex.datasets` still hand-built a Hugging Face `resolve` URL and an authorization header for the two cases an `hf://` URI cannot express: a token held in a variable, and a read that must stay anonymous despite the environment offering credentials. Expose the store class those cases actually want. `HfStore(repo_id, *, repo_type, revision, token, endpoint)` wraps the `HfConfig` / `make_hf_store` pair the way `CosStore` wraps `CosConfig`. `token` follows `huggingface_hub`'s own convention, so a caller can pass its value straight through: `None`/`True` take whatever the environment offers, `False` forces an anonymous read, and a string is used directly. `revision` is passed literally -- unlike in a URL, the store percent-encodes a revision containing `/` itself. `vx.open` previously accepted only the `pyo3-object_store` classes, which is why `datasets.py` carried a comment that `CosStore` could not be read from. It now takes the same `AnyVortexStore` extraction `read_url`/`write` use, so `HfStore` works there -- and `CosStore` does too, removing that limitation. Two things this surfaced: * `HfConfig` built `ClientOptions::default()`, ignoring `ALLOW_HTTP`. Every other scheme picks it up because `parse_url_opts` reads its configuration from the environment, so a plain-HTTP `HF_ENDPOINT` -- a self-hosted Hub, or a test double -- behaved inconsistently. It is honoured now, which is what lets the new end-to-end read test exercise the real path against a local server. * `HfRepoType` gained `FromStr` so the binding can take `repo_type="dataset"`, accepting the plural the Hub spells in its URLs as well. With the store class in place, `datasets.py` no longer knows the Hub endpoint at all: the `huggingface_hub.constants` import and the last `HTTPStore` use are gone. Signed-off-by: Robert Kruszewski Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01Lc5zw7Le2T3pakDEUdKTYd --- docs/api/python/datasets.rst | 5 +- docs/api/python/store.rst | 1 + docs/api/python/store/huggingface.rst | 97 ++++++++++++++++ docs/conf.py | 4 + vortex-cloud/src/hf/mod.rs | 89 +++++++++++++-- vortex-cloud/src/hf/tests.rs | 64 +++++++++++ vortex-python/python/vortex/_lib/__init__.pyi | 20 ++++ vortex-python/python/vortex/_lib/file.pyi | 3 +- vortex-python/python/vortex/datasets.py | 26 ++--- vortex-python/python/vortex/file.py | 4 +- vortex-python/python/vortex/store/__init__.py | 5 +- vortex-python/python/vortex/store/_hf.py | 16 +++ vortex-python/src/file.rs | 4 +- vortex-python/src/hf_store.rs | 107 ++++++++++++++++++ vortex-python/src/io.rs | 21 ++-- vortex-python/src/lib.rs | 2 + vortex-python/test/test_hf_datasets.py | 55 +++++++-- 17 files changed, 475 insertions(+), 48 deletions(-) create mode 100644 docs/api/python/store/huggingface.rst create mode 100644 vortex-python/python/vortex/store/_hf.py create mode 100644 vortex-python/src/hf_store.rs diff --git a/docs/api/python/datasets.rst b/docs/api/python/datasets.rst index 6d7597a7b13..24e27e5fba4 100644 --- a/docs/api/python/datasets.rst +++ b/docs/api/python/datasets.rst @@ -19,8 +19,9 @@ transforms. Hub repositories are streamed in place: files are read with HTTP range requests, so only the projected columns and matching rows are ever transferred. Private and gated repositories -authenticate with the ``token`` argument or the locally saved login. Files are downloaded (with -the usual Hub caching) only when ``streaming=False`` or ``local_files_only=True``. +authenticate with the ``token`` argument, ``HF_TOKEN``, or the locally saved login — see +:doc:`store/huggingface` for the full precedence. Files are downloaded (with the usual Hub +caching) only when ``streaming=False`` or ``local_files_only=True``. .. code-block:: python diff --git a/docs/api/python/store.rst b/docs/api/python/store.rst index ce4a6f75b2d..59bd6661695 100644 --- a/docs/api/python/store.rst +++ b/docs/api/python/store.rst @@ -11,6 +11,7 @@ Vortex arrays support reading and writing to many object storage systems: store/gcs store/azure store/http + store/huggingface store/local store/memory store/opendal diff --git a/docs/api/python/store/huggingface.rst b/docs/api/python/store/huggingface.rst new file mode 100644 index 00000000000..fd6831ceeb4 --- /dev/null +++ b/docs/api/python/store/huggingface.rst @@ -0,0 +1,97 @@ +================== +Hugging Face Hub +================== + +Vortex reads Hugging Face Hub repositories over ``hf://`` URLs. A Hub repository is a set of files +behind an HTTP endpoint that honours range requests, so no cloud SDK is involved and no extra build +feature is needed. + +.. list-table:: + :header-rows: 1 + + * - URL + - Repository kind + * - ``hf://datasets//[@][/]`` + - Dataset + * - ``hf://spaces//[@][/]`` + - Space + * - ``hf:///[@][/]`` + - Model + +```` is a branch, tag or commit, defaulting to ``main``. A revision containing ``/`` must +be percent-encoded, e.g. ``hf://datasets/org/name@refs%2Fconvert%2Fparquet/data/train.vortex``. + +Configuration comes from the same environment variables ``huggingface_hub`` reads: + +.. list-table:: + :header-rows: 1 + + * - Variable + - Meaning + * - ``HF_TOKEN`` + - API token for private and gated repositories. Falls back to the token file at + ``HF_TOKEN_PATH``, then ``$HF_HOME/token``, then ``$HOME/.cache/huggingface/token``. + * - ``HF_ENDPOINT`` + - Hub endpoint, defaulting to ``https://huggingface.co``. + +Reading from the Hub +==================== + +Pass an ``hf://`` URL directly. Public repositories need no credentials; private and gated ones +authenticate from ``HF_TOKEN`` or the saved login: + +.. code-block:: python + + import vortex as vx + + vxf = vx.open("hf://datasets/org/name/data/train.vortex") + for batch in vxf.to_arrow(): + ... + +:class:`vortex.store.HfStore` +============================= + +.. py:class:: vortex.store.HfStore(repo_id, *, repo_type="dataset", revision=None, token=None, endpoint=None) + + A Hugging Face Hub object store, rooted at one repository and revision. + + A URL is enough for most reads, so reach for this class only for the two things a URL cannot + express: a token held in a variable rather than the environment, and a read that must stay + anonymous even though the environment offers credentials. + + Because the store is rooted at the repository and revision, the path passed alongside it is a + path *within* the repository. + + :param repo_id: The repository, as ``"/"``. + :param repo_type: ``"dataset"``, ``"model"`` or ``"space"``. Defaults to ``"dataset"``. + :param revision: A branch, tag or commit. Defaults to ``main``. Unlike in a URL, a revision + containing ``/`` is passed literally — the store percent-encodes it. + :param token: ``None`` (the default) or ``True`` authenticates from ``HF_TOKEN`` or the saved + login; ``False`` forces an anonymous read even when credentials are available; a string is + used as the token directly. + :param endpoint: Hub endpoint. Defaults to ``HF_ENDPOINT``, then ``https://huggingface.co``. + +.. code-block:: python + + import vortex as vx + from vortex.store import HfStore + + store = HfStore("org/name", revision="refs/convert/parquet", token="hf_...") + + # With `store=`, the path is a path within the repository. + vxf = vx.open("data/train.vortex", store=store) + +Listing +======= + +The Hub does not implement WebDAV ``PROPFIND``, which is how object-store HTTP listing works, so a +Hub store cannot list a prefix. Opening a known path works, since that is a ``HEAD`` plus ranged +``GET``. To expand a glob, list the repository through the Hub's own API first — which is what +``vortex.datasets.load_dataset`` does — and then open each path it returns. + +Hugging Face Datasets +===================== + +``vortex.datasets.load_dataset`` builds on this to load Vortex files from the Hub as Hugging Face +``Datasets`` objects, expanding globs and pushing projections, filters and row limits into each +scan. See :doc:`../datasets`. diff --git a/docs/conf.py b/docs/conf.py index b73ed99c763..d49bfad5d03 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -56,6 +56,10 @@ # and the `ObjectStore` type alias resolves to the private path. The public class is # fully documented in `opendal.rst`; the private path is intentionally not. ("py:class", "vortex.store._cos.CosStore"), + # `vortex.store.HfStore` is the native class re-exported through `vortex.store._hf`, so + # annotations resolve to its `vortex._lib` module path. The public class is fully documented + # in `huggingface.rst`; the native path is intentionally not. + ("py:class", "vortex._lib.HfStore"), ] doctest_global_setup = "import pyarrow; import vortex; import vortex as vx; import random; random.seed(a=0)" diff --git a/vortex-cloud/src/hf/mod.rs b/vortex-cloud/src/hf/mod.rs index e39a7045bbd..9b6ab374224 100644 --- a/vortex-cloud/src/hf/mod.rs +++ b/vortex-cloud/src/hf/mod.rs @@ -71,6 +71,7 @@ const TOKEN_VAR: &str = "HF_TOKEN"; const TOKEN_PATH_VAR: &str = "HF_TOKEN_PATH"; const HF_HOME_VAR: &str = "HF_HOME"; const HOME_VAR: &str = "HOME"; +const ALLOW_HTTP_VAR: &str = "ALLOW_HTTP"; /// The URL authority that marks a dataset repository. const DATASETS_HOST: &str = "datasets"; @@ -116,6 +117,21 @@ impl HfRepoType { } } +impl std::str::FromStr for HfRepoType { + type Err = HfStoreError; + + /// Parses either the singular name or the plural the Hub uses in its URLs, so a caller can pass + /// whichever it has. + fn from_str(s: &str) -> Result { + match s { + "dataset" | "datasets" => Ok(HfRepoType::Dataset), + "model" | "models" => Ok(HfRepoType::Model), + "space" | "spaces" => Ok(HfRepoType::Space), + other => Err(HfStoreError::UnknownRepoType(other.to_string())), + } + } +} + /// Error type for building a Hugging Face Hub object store. #[derive(Debug)] pub enum HfStoreError { @@ -123,6 +139,8 @@ pub enum HfStoreError { UnsupportedScheme(String), /// The URL is not a well-formed `hf://` URL. InvalidUrl(String), + /// The repository kind is not one the Hub serves. + UnknownRepoType(String), /// The bearer token could not be used as an HTTP header value. InvalidToken(http::header::InvalidHeaderValue), /// The underlying HTTP store rejected the configuration. @@ -138,6 +156,10 @@ impl std::fmt::Display for HfStoreError { "invalid Hugging Face URL {url}: expected hf://datasets//[@revision][/path], \ hf://spaces//[@revision][/path] or hf:///[@revision][/path]" ), + HfStoreError::UnknownRepoType(kind) => write!( + f, + "unknown Hugging Face repository type {kind}: expected one of dataset, model, space" + ), HfStoreError::InvalidToken(e) => write!(f, "invalid Hugging Face token: {e}"), HfStoreError::Build(e) => write!(f, "failed to build Hugging Face store: {e}"), } @@ -149,7 +171,9 @@ impl std::error::Error for HfStoreError { match self { HfStoreError::InvalidToken(e) => Some(e), HfStoreError::Build(e) => Some(e), - HfStoreError::UnsupportedScheme(_) | HfStoreError::InvalidUrl(_) => None, + HfStoreError::UnsupportedScheme(_) + | HfStoreError::InvalidUrl(_) + | HfStoreError::UnknownRepoType(_) => None, } } } @@ -198,6 +222,45 @@ impl Default for HfConfig { } impl HfConfig { + /// Configuration for one repository, taking the token and endpoint from the process environment + /// exactly as resolving an `hf://` URL does. + /// + /// `revision` defaults to `main`. Callers that must read anonymously should clear + /// [`HfConfig::token`] afterwards, since this picks up whatever credentials the environment + /// offers. + pub fn from_env( + repo_type: HfRepoType, + repo_id: impl Into, + revision: Option, + ) -> Self { + Self::from_env_lookup(repo_type, repo_id, revision, |key| std::env::var(key).ok()) + } + + /// [`HfConfig::from_env`], reading environment configuration through `env_lookup`. + fn from_env_lookup( + repo_type: HfRepoType, + repo_id: impl Into, + revision: Option, + env_lookup: F, + ) -> Self + where + F: Fn(&str) -> Option, + { + Self { + repo_type, + repo_id: repo_id.into(), + revision: revision.unwrap_or_else(|| DEFAULT_REVISION.to_string()), + token: resolve_token(&env_lookup), + endpoint: env_lookup(ENDPOINT_VAR) + .filter(|endpoint| !endpoint.is_empty()) + .unwrap_or_else(|| DEFAULT_ENDPOINT.to_string()), + // Every other scheme picks `allow_http` out of the environment, because that is how + // `parse_url_opts` reads its configuration. Honour it here too, so a plain-HTTP + // `HF_ENDPOINT` (a test double, or a self-hosted Hub) behaves the same way. + client_options: ClientOptions::default().with_allow_http(allow_http(&env_lookup)), + } + } + /// The URL this repository's files hang off, which is where the store is rooted. fn resolve_prefix(&self) -> String { let endpoint = self.endpoint.trim_end_matches('/'); @@ -308,16 +371,12 @@ where None => (*name, DEFAULT_REVISION.to_string()), }; - let config = HfConfig { + let config = HfConfig::from_env_lookup( repo_type, - repo_id: format!("{owner}/{name}"), - revision, - token: resolve_token(&env_lookup), - endpoint: env_lookup(ENDPOINT_VAR) - .filter(|endpoint| !endpoint.is_empty()) - .unwrap_or_else(|| DEFAULT_ENDPOINT.to_string()), - client_options: ClientOptions::default(), - }; + format!("{owner}/{name}"), + Some(revision), + &env_lookup, + ); // The segments still carry their URL escapes, so decode them into the object key rather than // joining them raw. @@ -326,6 +385,16 @@ where Ok((config, path)) } +/// Whether plain-HTTP endpoints are permitted, from the `ALLOW_HTTP` variable the `object_store` +/// builders read. Anything other than a `true`-ish value leaves HTTPS as the requirement. +fn allow_http(env_lookup: &F) -> bool +where + F: Fn(&str) -> Option, +{ + env_lookup(ALLOW_HTTP_VAR) + .is_some_and(|value| matches!(value.trim().to_ascii_lowercase().as_str(), "true" | "1")) +} + /// The bearer token to read with, following the same precedence as `huggingface_hub.get_token()`: /// `HF_TOKEN`, then the token file at `HF_TOKEN_PATH`, `$HF_HOME/token` or /// `$HOME/.cache/huggingface/token`. diff --git a/vortex-cloud/src/hf/tests.rs b/vortex-cloud/src/hf/tests.rs index cc3b9c86c5a..04b8fe40cb1 100644 --- a/vortex-cloud/src/hf/tests.rs +++ b/vortex-cloud/src/hf/tests.rs @@ -298,3 +298,67 @@ fn test_make_hf_store() -> Result<(), Box> { )); Ok(()) } + +#[rstest] +// Both the singular name and the plural the Hub spells in its URLs are accepted, so a caller can +// pass whichever it has. +#[case("dataset", HfRepoType::Dataset)] +#[case("datasets", HfRepoType::Dataset)] +#[case("model", HfRepoType::Model)] +#[case("models", HfRepoType::Model)] +#[case("space", HfRepoType::Space)] +#[case("spaces", HfRepoType::Space)] +fn test_repo_type_from_str( + #[case] spelling: &str, + #[case] expected: HfRepoType, +) -> Result<(), Box> { + assert_eq!(spelling.parse::()?, expected); + Ok(()) +} + +#[test] +fn test_repo_type_from_str_rejects_unknown() { + assert!(matches!( + "notarepo".parse::(), + Err(HfStoreError::UnknownRepoType(_)) + )); +} + +/// `from_env` must resolve the token and endpoint the same way an `hf://` URL does, since the store +/// class and the URL registry are two doors onto one configuration. +#[test] +fn test_from_env_matches_url_resolution() -> Result<(), Box> { + let vars = &[ + ("HF_TOKEN", "hf_shared"), + ("HF_ENDPOINT", "https://hub.example.com"), + ]; + let url = Url::parse("hf://datasets/org/name/train.vortex")?; + + let (from_url, _path) = url_to_config(&url, env(vars))?; + let direct = HfConfig::from_env_lookup(HfRepoType::Dataset, "org/name", None, env(vars)); + + assert_eq!(from_url.token, direct.token); + assert_eq!(from_url.endpoint, direct.endpoint); + assert_eq!(from_url.resolve_prefix(), direct.resolve_prefix()); + Ok(()) +} + +/// `ALLOW_HTTP` must reach the client options, since that is how every other scheme is told plain +/// HTTP is acceptable and a self-hosted or stubbed Hub endpoint needs the same door. `ClientOptions` +/// exposes no getter, so the observable assertion is an actual plain-HTTP read; that lives in +/// `vortex-python/test/test_hf_datasets.py::test_hf_store_reads_a_repository_relative_path`. +#[test] +fn test_allow_http_is_configured_from_env() { + let allowed = HfConfig::from_env_lookup( + HfRepoType::Dataset, + "org/name", + None, + env(&[("ALLOW_HTTP", "true")]), + ); + let denied = HfConfig::from_env_lookup(HfRepoType::Dataset, "org/name", None, env(&[])); + + assert_ne!( + format!("{:?}", allowed.client_options), + format!("{:?}", denied.client_options) + ); +} diff --git a/vortex-python/python/vortex/_lib/__init__.pyi b/vortex-python/python/vortex/_lib/__init__.pyi index 09c9a520406..22161cc2992 100644 --- a/vortex-python/python/vortex/_lib/__init__.pyi +++ b/vortex-python/python/vortex/_lib/__init__.pyi @@ -21,3 +21,23 @@ class CosStore: root: str | None = None, disable_config_load: bool = False, ) -> None: ... + +class HfStore: + """A Hugging Face Hub object store, rooted at one repository and revision. + + Reading an ``hf://`` URL needs no store; this covers a token held in a variable + rather than the environment, and reads that must stay anonymous regardless of it. + + Construct it and pass it to ``vortex.io.read_url(path, store=hf_store)``, where + ``path`` is a path within the repository. + """ + + def __init__( + self, + repo_id: str, + *, + repo_type: str = "dataset", + revision: str | None = None, + token: bool | str | None = None, + endpoint: str | None = None, + ) -> None: ... diff --git a/vortex-python/python/vortex/_lib/file.pyi b/vortex-python/python/vortex/_lib/file.pyi index ed8109210c4..f8c928be459 100644 --- a/vortex-python/python/vortex/_lib/file.pyi +++ b/vortex-python/python/vortex/_lib/file.pyi @@ -8,6 +8,7 @@ import pyarrow as pa from vortex.type_aliases import IntoProjection +from . import CosStore, HfStore from .arrays import Array from .dataset import VortexDataset from .dtype import DType @@ -55,6 +56,6 @@ class VortexFile: def open( path: str, *, - store: ObjectStore | None = None, + store: ObjectStore | CosStore | HfStore | None = None, without_segment_cache: bool = False, ) -> VortexFile: ... diff --git a/vortex-python/python/vortex/datasets.py b/vortex-python/python/vortex/datasets.py index d10bc66d721..3e5536fde1e 100644 --- a/vortex-python/python/vortex/datasets.py +++ b/vortex-python/python/vortex/datasets.py @@ -29,17 +29,17 @@ from vortex.expr import Expr, and_ from vortex.store import ( AzureStore, - ClientConfig, + CosStore, GCSStore, + HfStore, HTTPStore, LocalStore, MemoryStore, S3Store, ) -# The stores `vx.open` accepts. This is narrower than `vortex.store.ObjectStore`, which also -# covers the OpenDAL-backed `CosStore` that `vx.open` cannot read from. -ObjectStore: TypeAlias = AzureStore | GCSStore | HTTPStore | LocalStore | MemoryStore | S3Store +# The stores `vx.open` accepts. +ObjectStore: TypeAlias = AzureStore | CosStore | GCSStore | HfStore | HTTPStore | LocalStore | MemoryStore | S3Store try: import datasets as hf_datasets @@ -50,7 +50,6 @@ ) from datasets.table import InMemoryTable from huggingface_hub import HfApi, snapshot_download - from huggingface_hub import constants as hf_hub_constants except ImportError as e: # pragma: no cover - exercised only without optional deps. raise ImportError("Install vortex-data[hf] to use vortex.datasets.") from e @@ -62,7 +61,6 @@ _DEFAULT_SPLIT = "train" _DEFAULT_DATA_FILES = "**/*.vortex" -_DEFAULT_REVISION = "main" _ITERABLE_DATASET_HAS_SHUFFLING = "shuffling" in inspect.signature(hf_datasets.IterableDataset).parameters @@ -852,9 +850,9 @@ def _resolve_hub_files( ``token`` (and ``token=True``, which asks for exactly that saved login) the matched files are returned as ``hf://`` URIs and need no store. - The two cases the reader cannot express are handled with a store of its own: a ``token`` string, - which the reader has no way to see, and ``token=False``, which must suppress the credentials the - reader would otherwise pick up from the environment. + The two cases a URI cannot express go through an :class:`~vortex.store.HfStore` instead: a + ``token`` string, which the reader has no way to see, and ``token=False``, which must suppress + the credentials the reader would otherwise pick up from the environment. The Hub serves no listing over the object-store protocol, so the patterns are expanded here through the Hub API either way. @@ -872,14 +870,8 @@ def _resolve_hub_files( split_to_matches[split_name] = matches if token is False or isinstance(token, str): - # The Hub only routes the percent-encoded form of a revision containing `/`, so the store is - # rooted at the encoded `resolve` prefix and the files stay in-repository paths. - revision_path = quote(revision if revision is not None else _DEFAULT_REVISION, safe="") - base = f"{hf_hub_constants.ENDPOINT}/datasets/{repo_id}/resolve/{revision_path}" - client_options: ClientConfig | None = ( - {"default_headers": {"authorization": f"Bearer {token}"}} if isinstance(token, str) else None - ) - return split_to_matches, HTTPStore(base, client_options=client_options) + # An HfStore is rooted at the repository and revision, so the files stay in-repository paths. + return split_to_matches, HfStore(repo_id, revision=revision, token=token) prefix = f"hf://datasets/{repo_id}" if revision is not None: diff --git a/vortex-python/python/vortex/file.py b/vortex-python/python/vortex/file.py index 90336338e84..7ca666acf50 100644 --- a/vortex-python/python/vortex/file.py +++ b/vortex-python/python/vortex/file.py @@ -17,7 +17,9 @@ from .scan import RepeatedScan from .store import ( AzureStore, + CosStore, GCSStore, + HfStore, HTTPStore, LocalStore, MemoryStore, @@ -32,7 +34,7 @@ def open( path: str, *, - store: AzureStore | GCSStore | HTTPStore | LocalStore | MemoryStore | S3Store | None = None, + store: AzureStore | CosStore | GCSStore | HfStore | HTTPStore | LocalStore | MemoryStore | S3Store | None = None, without_segment_cache: bool = False, ) -> VortexFile: """ diff --git a/vortex-python/python/vortex/store/__init__.py b/vortex-python/python/vortex/store/__init__.py index c9127087bf0..690aa1a97d1 100644 --- a/vortex-python/python/vortex/store/__init__.py +++ b/vortex-python/python/vortex/store/__init__.py @@ -18,12 +18,13 @@ from ._client import ClientConfig from ._cos import CosStore from ._gcs import GCSConfig, GCSCredential, GCSCredentialProvider, GCSStore +from ._hf import HfStore from ._http import HTTPStore from ._local import LocalStore from ._memory import MemoryStore from ._retry import BackoffConfig, RetryConfig -ObjectStore: TypeAlias = AzureStore | CosStore | GCSStore | HTTPStore | S3Store | LocalStore | MemoryStore +ObjectStore: TypeAlias = AzureStore | CosStore | GCSStore | HfStore | HTTPStore | S3Store | LocalStore | MemoryStore """All supported ObjectStore implementations.""" @@ -152,6 +153,8 @@ def from_url( # type: ignore[misc] # docstring in pyi file "GCSStore", # HTTP "HTTPStore", + # Hugging Face Hub + "HfStore", # Local "LocalStore", "MemoryStore", diff --git a/vortex-python/python/vortex/store/_hf.py b/vortex-python/python/vortex/store/_hf.py new file mode 100644 index 00000000000..2e1832c01ca --- /dev/null +++ b/vortex-python/python/vortex/store/_hf.py @@ -0,0 +1,16 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright the Vortex contributors + +"""Hugging Face Hub object store. + +Reading an ``hf://`` URL needs no store: :func:`vortex.io.read_url` resolves it, taking credentials +from ``HF_TOKEN`` or the saved login. :class:`HfStore` covers the two cases a URL cannot express — a +token held in a variable rather than the environment, and a read that must stay anonymous even +though the environment offers credentials. +""" + +from __future__ import annotations + +from vortex._lib import HfStore as HfStore + +__all__ = ["HfStore"] diff --git a/vortex-python/src/file.rs b/vortex-python/src/file.rs index 5f8a39862b7..b3cdad9d153 100644 --- a/vortex-python/src/file.rs +++ b/vortex-python/src/file.rs @@ -8,7 +8,6 @@ use arrow_schema::Schema; use pyo3::exceptions::PyTypeError; use pyo3::prelude::*; use pyo3::types::PyList; -use pyo3_object_store::PyObjectStore; use vortex::array::ArrayRef; use vortex::array::ExecutionCtx; use vortex::array::VortexSessionExecute; @@ -39,6 +38,7 @@ use crate::dtype::PyDType; use crate::error::PyVortexResult; use crate::expr::PyExpr; use crate::install_module; +use crate::io::AnyVortexStore; use crate::iter::PyArrayIterator; use crate::object_store::resolve::ResolvedStore; use crate::object_store::resolve::resolve_store; @@ -65,7 +65,7 @@ pub(crate) fn init(py: Python, parent: &Bound) -> PyResult<()> { pub fn open( py: Python, path: &str, - store: Option, + store: Option, without_segment_cache: bool, ) -> PyVortexResult { let vxf = py.detach(move || { diff --git a/vortex-python/src/hf_store.rs b/vortex-python/src/hf_store.rs new file mode 100644 index 00000000000..612dda6d902 --- /dev/null +++ b/vortex-python/src/hf_store.rs @@ -0,0 +1,107 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Python-facing wrapper for the Hugging Face Hub object store. +//! +//! Reading an `hf://` URL needs no store at all — the URL registry resolves it, taking credentials +//! from `HF_TOKEN` or the saved login. This class exists for the two things a URL cannot say: a +//! token held in a variable rather than the environment, and a read that must stay anonymous even +//! though the environment offers credentials. + +use std::str::FromStr; +use std::sync::Arc; + +use pyo3::exceptions::PyValueError; +use pyo3::prelude::*; +use vortex::cloud::hf::HfConfig; +use vortex::cloud::hf::HfRepoType; +use vortex::cloud::hf::make_hf_store; + +/// How the `token` argument was spelled. +/// +/// Mirrors `huggingface_hub`'s own convention so that a caller can pass the same value through. +#[derive(Debug, Clone, FromPyObject)] +enum TokenArg { + /// `True` asks for the saved login, `False` forces an anonymous read. + Flag(bool), + /// An explicit token. + Value(String), +} + +/// A Hugging Face Hub object store, rooted at one repository and revision. +/// +/// Construct it with explicit configuration and pass it to +/// ``vortex.io.read_url(path, store=hf_store)``, where ``path`` is a path within the repository. +#[pyclass(name = "HfStore", module = "vortex._lib", frozen, from_py_object)] +#[derive(Clone, Debug)] +pub struct HfStore { + store: Arc, +} + +impl HfStore { + /// Clone the underlying object store as an `Arc`. + pub fn to_arc(&self) -> Arc { + Arc::clone(&self.store) + } +} + +#[pymethods] +impl HfStore { + #[new] + #[pyo3(signature = ( + repo_id, + *, + repo_type = "dataset", + revision = None, + token = None, + endpoint = None, + ))] + fn new( + repo_id: String, + repo_type: &str, + revision: Option, + token: Option, + endpoint: Option, + ) -> PyResult { + let repo_type = + HfRepoType::from_str(repo_type).map_err(|e| PyValueError::new_err(e.to_string()))?; + + let mut config = HfConfig::from_env(repo_type, repo_id, revision); + match token { + // The default and `token=True` both mean "whatever the environment offers", which + // `from_env` has already resolved. + None | Some(TokenArg::Flag(true)) => {} + Some(TokenArg::Flag(false)) => config.token = None, + Some(TokenArg::Value(token)) => config.token = Some(token), + } + if let Some(endpoint) = endpoint { + config.endpoint = endpoint; + } + + let store = make_hf_store(config).map_err(|e| PyValueError::new_err(e.to_string()))?; + Ok(Self { store }) + } +} + +/// Register the Hugging Face store class on the `vortex._lib` module. +pub(crate) fn init(_py: Python, parent: &Bound) -> PyResult<()> { + parent.add_class::()?; + Ok(()) +} + +#[cfg(test)] +mod tests { + use vortex::cloud::hf::HfConfig; + use vortex::cloud::hf::HfRepoType; + use vortex::cloud::hf::make_hf_store; + + /// `token=False` must reach `make_hf_store` with no token at all, which is the whole reason this + /// class exists alongside plain `hf://` URL resolution. + #[test] + fn anonymous_config_builds() { + let mut config = HfConfig::from_env(HfRepoType::Dataset, "org/name", None); + config.token = None; + + assert!(make_hf_store(config).is_ok()); + } +} diff --git a/vortex-python/src/io.rs b/vortex-python/src/io.rs index f99322673b7..3fa89b16380 100644 --- a/vortex-python/src/io.rs +++ b/vortex-python/src/io.rs @@ -39,6 +39,7 @@ use crate::classes::table_class; use crate::dataset::PyVortexDataset; use crate::error::PyVortexResult; use crate::expr::PyExpr; +use crate::hf_store::HfStore; use crate::install_module; use crate::iter::PyArrayIterator; use crate::object_store::resolve::ResolvedStore; @@ -66,7 +67,7 @@ pub(crate) fn init(py: Python, parent: &Bound) -> PyResult<()> { /// ---------- /// url : str /// The URL to read from. -/// store : vortex.store.AzureStore | vortex.store.CosStore | vortex.store.GCSStore | vortex.store.HTTPStore | vortex.store.LocalStore | vortex.store.MemoryStore | vortex.store.S3Store | None +/// store : vortex.store.AzureStore | vortex.store.CosStore | vortex.store.GCSStore | vortex.store.HfStore | vortex.store.HTTPStore | vortex.store.LocalStore | vortex.store.MemoryStore | vortex.store.S3Store | None /// Pre-configured object store with credentials and settings. /// If provided, uses this store's configuration. /// If None, checks session registry for matching URL pattern. @@ -144,11 +145,13 @@ pub fn read_url<'py>( /// A store object accepted by `read_url` / `write`. /// -/// This recognizes both the built-in `pyo3-object_store` classes (S3, Azure, GCS, HTTP, -/// Local, Memory) and Vortex's own OpenDAL-backed classes (`CosStore`). +/// This recognizes the built-in `pyo3-object_store` classes (S3, Azure, GCS, HTTP, Local, Memory) +/// and Vortex's own classes: the Hugging Face Hub store and the OpenDAL-backed `CosStore`. pub(crate) enum AnyVortexStore { /// A store extracted from one of the built-in `pyo3-object_store` classes. Builtin(PyObjectStore), + /// Vortex's Hugging Face Hub store. + Hf(HfStore), /// Vortex's OpenDAL-backed COS store. #[cfg(feature = "opendal")] Cos(CosStore), @@ -156,9 +159,10 @@ pub(crate) enum AnyVortexStore { impl AnyVortexStore { /// Consume self and return the underlying `Arc`. - fn into_inner(self) -> Arc { + pub(crate) fn into_inner(self) -> Arc { match self { AnyVortexStore::Builtin(s) => s.into_inner(), + AnyVortexStore::Hf(s) => s.to_arc(), #[cfg(feature = "opendal")] AnyVortexStore::Cos(s) => s.to_arc(), } @@ -172,12 +176,15 @@ impl<'py> FromPyObject<'_, 'py> for AnyVortexStore { if let Ok(builtin) = obj.extract::() { return Ok(AnyVortexStore::Builtin(builtin)); } + if let Ok(hf) = obj.extract::() { + return Ok(AnyVortexStore::Hf(hf)); + } #[cfg(feature = "opendal")] if let Ok(cos) = obj.extract::() { return Ok(AnyVortexStore::Cos(cos)); } Err(PyTypeError::new_err( - "Expected an object store instance (S3/Azure/GCS/HTTP/Local/Memory/COS/OSS store)", + "Expected an object store instance (S3/Azure/GCS/HTTP/Local/Memory/HF/COS/OSS store)", )) } } @@ -193,7 +200,7 @@ impl<'py> FromPyObject<'_, 'py> for AnyVortexStore { /// path : str /// The file path. /// -/// store : vortex.store.AzureStore | vortex.store.CosStore | vortex.store.GCSStore | vortex.store.HTTPStore | vortex.store.LocalStore | vortex.store.MemoryStore | vortex.store.S3Store | None +/// store : vortex.store.AzureStore | vortex.store.CosStore | vortex.store.GCSStore | vortex.store.HfStore | vortex.store.HTTPStore | vortex.store.LocalStore | vortex.store.MemoryStore | vortex.store.S3Store | None /// An optional object store configuration to use for writing the output. /// /// Examples @@ -335,7 +342,7 @@ impl PyVortexWriteOptions { /// path : str /// The file path. /// - /// store : vortex.store.AzureStore | vortex.store.CosStore | vortex.store.GCSStore | vortex.store.HTTPStore | vortex.store.LocalStore | vortex.store.MemoryStore | vortex.store.S3Store | None + /// store : vortex.store.AzureStore | vortex.store.CosStore | vortex.store.GCSStore | vortex.store.HfStore | vortex.store.HTTPStore | vortex.store.LocalStore | vortex.store.MemoryStore | vortex.store.S3Store | None /// An optional object store configuration to use for writing the output. /// /// Examples diff --git a/vortex-python/src/lib.rs b/vortex-python/src/lib.rs index 0912426d10b..b1ada751175 100644 --- a/vortex-python/src/lib.rs +++ b/vortex-python/src/lib.rs @@ -22,6 +22,7 @@ pub(crate) mod dtype; mod error; mod expr; mod file; +mod hf_store; mod io; mod iter; mod object_store; @@ -80,6 +81,7 @@ fn _lib(py: Python, m: &Bound) -> PyResult<()> { dtype::init(py, m)?; expr::init(py, m)?; file::init(py, m)?; + hf_store::init(py, m)?; io::init(py, m)?; iter::init(py, m)?; #[cfg(feature = "opendal")] diff --git a/vortex-python/test/test_hf_datasets.py b/vortex-python/test/test_hf_datasets.py index dbbc0e01442..6990c907f85 100644 --- a/vortex-python/test/test_hf_datasets.py +++ b/vortex-python/test/test_hf_datasets.py @@ -17,7 +17,7 @@ import pytest import vortex.datasets as vx_datasets from typing_extensions import override -from vortex.store import HTTPStore +from vortex.store import HfStore import vortex as vx import vortex.expr as ve @@ -462,9 +462,8 @@ def test_hub_streaming_with_token_false_forces_anonymous_store(monkeypatch: pyte ) # `token=False` must suppress the credentials Vortex would otherwise read from the environment, - # which an `hf://` URI cannot express, so it gets an unauthenticated store instead. - assert isinstance(store, HTTPStore) - assert store.url.endswith("/datasets/org/name/resolve/main") + # which an `hf://` URI cannot express, so it gets an anonymous store instead. + assert isinstance(store, HfStore) assert files == {"train": ["data/validation.vortex", "train.vortex"]} @@ -483,9 +482,8 @@ def test_hub_streaming_with_token_uses_authenticated_store(monkeypatch: pytest.M ) # An explicitly passed token cannot reach the Vortex reader, so this is the one path that still - # builds a store; the files are then relative to it. - assert isinstance(store, HTTPStore) - assert store.url.endswith("/datasets/org/name/resolve/main") + # builds a store; the files are then paths within the repository. + assert isinstance(store, HfStore) assert files == {"train": ["data/validation.vortex", "train.vortex"]} @@ -602,3 +600,46 @@ def test_local_directory_in_data_files(tmp_path: Path): ) def test_glob_match(pattern: str, path: str, expected: bool): assert vx_datasets._glob_match(path, pattern) is expected # pyright: ignore[reportPrivateUsage] + + +@pytest.mark.parametrize("repo_type", ["dataset", "datasets", "model", "space"]) +def test_hf_store_accepts_each_repo_type(repo_type: str): + assert isinstance(HfStore("org/name", repo_type=repo_type), HfStore) + + +def test_hf_store_rejects_unknown_repo_type(): + with pytest.raises(ValueError, match="repository type"): + _ = HfStore("org/name", repo_type="notarepo") + + +@pytest.mark.parametrize("token", [None, True, False, "hf_explicit_token"]) +def test_hf_store_accepts_each_token_spelling(token: bool | str | None): + # `token` follows huggingface_hub's convention: None/True use the environment, False forces an + # anonymous read, and a string is used directly. + assert isinstance(HfStore("org/name", token=token), HfStore) + + +def test_hf_store_rejects_unusable_token(): + with pytest.raises(ValueError, match="token"): + _ = HfStore("org/name", token="bad\nvalue") + + +def test_hf_store_reads_a_repository_relative_path(tmp_path: Path, monkeypatch: pytest.MonkeyPatch): + # Point the store at a local server standing in for the Hub, so this exercises the real read + # path: HfStore is rooted at the repository revision and the path is relative to it. + monkeypatch.setenv("ALLOW_HTTP", "true") + rows: list[dict[str, object]] = [{"idx": i} for i in range(4)] + resolve_dir = tmp_path / "datasets" / "org" / "name" / "resolve" / "main" + resolve_dir.mkdir(parents=True) + write_vortex(resolve_dir / "train.vortex", rows) + + handler = type("Handler", (_RangeRequestHandler,), {"directory": tmp_path}) + server = ThreadingHTTPServer(("127.0.0.1", 0), handler) + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + try: + endpoint = f"http://127.0.0.1:{server.server_address[1]}" + store = HfStore("org/name", endpoint=endpoint) + assert vx.open("train.vortex", store=store).to_arrow().read_all().to_pylist() == rows + finally: + server.shutdown()