Skip to content

Resolve every Java-binding URL through the vortex-cloud registry, adding hf:// #9273

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
robert3005 wants to merge 1 commit into
base: develop
Choose a base branch
from
Draft

Resolve every Java-binding URL through the vortex-cloud registry, adding hf:// #9273

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 6 additions & 10 deletions vortex-cloud/src/registry/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -86,17 +86,12 @@ pub struct Registry {
}

/// Source of the configuration variables consulted when building a store.
///
/// Tests construct a registry over a fixed set of variables rather than mutating the process
/// environment, which is unsound when tests run on multiple threads within one process (the
/// `std::env::set_var` block became `unsafe` in Rust 2024 for exactly this reason).
#[derive(Debug, Default)]
enum EnvSource {
/// Read from the process environment.
#[default]
Process,
/// A fixed set of variables.
#[cfg(test)]
Fixed(Vec<(String, String)>),
}

Expand All @@ -106,7 +101,6 @@ impl EnvSource {
fn lookup(&self, key: &str) -> Option<String> {
match self {
EnvSource::Process => std::env::var(key).ok(),
#[cfg(test)]
EnvSource::Fixed(vars) => vars
.iter()
.find(|(k, _)| k.eq_ignore_ascii_case(key))
Expand All @@ -120,7 +114,6 @@ impl EnvSource {
EnvSource::Process => std::env::vars()
.map(|(k, v)| (k.to_ascii_lowercase(), v))
.collect(),
#[cfg(test)]
EnvSource::Fixed(vars) => vars
.iter()
.map(|(k, v)| (k.to_ascii_lowercase(), v.clone()))
Expand Down Expand Up @@ -177,9 +170,12 @@ impl Registry {
Self::default()
}

/// Create a registry over a fixed set of configuration variables.
#[cfg(test)]
fn with_env<I>(vars: I) -> Self
/// Create a registry over a fixed set of configuration variables, consulted instead of the
/// process environment when building stores.
///
/// Lookups are case-insensitive. Pass each key at most once — the set's consumers disagree
/// on which duplicate wins.
pub fn with_vars<I>(vars: I) -> Self
where
I: IntoIterator<Item = (String, String)>,
{
Expand Down
2 changes: 1 addition & 1 deletion vortex-cloud/src/registry/tests.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,7 @@ use super::Registry;
/// A registry whose S3 configuration comes from a fixed map rather than the process environment,
/// so these tests neither read nor mutate global state.
fn registry() -> Registry {
Registry::with_env([("AWS_REGION".to_string(), "us-east-3".to_string())])
Registry::with_vars([("AWS_REGION".to_string(), "us-east-3".to_string())])
}

/// A percent-encoded segment (as HuggingFace dataset URLs use for `refs/convert/parquet`) decodes
Expand Down
7 changes: 1 addition & 6 deletions vortex-jni/Cargo.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -33,18 +33,13 @@ tracing-subscriber = { workspace = true, features = ["env-filter"] }
url = { workspace = true }
vortex = { workspace = true, features = ["object_store", "files"] }
vortex-arrow = { workspace = true }
vortex-cloud = { workspace = true, optional = true }
vortex-cloud = { workspace = true, features = ["hf", "opendal", "registry"] }
vortex-geo = { workspace = true }
vortex-parquet-variant = { workspace = true }

[dev-dependencies]
jni = { workspace = true, features = ["invocation"] }

[features]
# Enable OpenDAL-backed object stores (Tencent COS, Alibaba OSS) for `cos://` and `oss://` URLs.
# This pulls in the `opendal` dependency, so it is opt-in.
opendal = ["vortex-cloud/opendal"]

[lib]
crate-type = ["cdylib"]

Expand Down
40 changes: 3 additions & 37 deletions vortex-jni/src/data_source.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -30,7 +30,6 @@ use vortex::io::filesystem::FileSystemRef;
use vortex::io::runtime::BlockingRuntime;
use vortex::io::session::RuntimeSessionExt;
use vortex::scan::DataSourceRef;
use vortex::utils::aliases::hash_map::HashMap;
use vortex_arrow::ArrowSessionExt;

use crate::RUNTIME;
Expand Down Expand Up @@ -93,23 +92,11 @@ pub extern "system" fn Java_dev_vortex_jni_NativeDataSource_open(
.map(|g| parse_uri_or_path(g.as_str()))
.collect::<VortexResult<_>>()?;

let mut fs_cache: HashMap<Url, FileSystemRef> = HashMap::new();
for glob_url in &glob_urls {
let base = base_url(glob_url);
if !fs_cache.contains_key(&base) {
let fs = object_store_fs(glob_url, &properties, session.handle())?;
fs_cache.insert(base, fs);
}
}

// Glob by the path the resolver reports — only it knows how deep each store is mounted.
let mut builder = MultiFileDataSource::new(session.clone());
for glob_url in &glob_urls {
let base = base_url(glob_url);
let fs = fs_cache
.get(&base)
.cloned()
.unwrap_or_else(|| unreachable!("fs cached for every base url"));
builder = builder.with_glob(glob_url.path(), Some(fs));
let (fs, glob) = object_store_fs(glob_url, &properties, session.handle())?;
builder = builder.with_glob(glob, Some(fs));
}

let inner = RUNTIME
Expand Down Expand Up @@ -190,13 +177,6 @@ pub extern "system" fn Java_dev_vortex_jni_NativeDataSource_openFiles(
})
}

/// URL with the path cleared, used as a cache key for filesystem reuse.
fn base_url(url: &Url) -> Url {
let mut base = url.clone();
base.set_path("");
base
}

#[unsafe(no_mangle)]
pub extern "system" fn Java_dev_vortex_jni_NativeDataSource_free(
_env: EnvUnowned,
Expand Down Expand Up @@ -274,17 +254,3 @@ pub extern "system" fn Java_dev_vortex_jni_NativeDataSource_byteSize(
Ok(())
});
}

#[cfg(test)]
mod tests {
use super::*;

#[test]
fn test_base_url_strips_path() {
let url = Url::parse("s3://bucket/a/b/c").unwrap();
let base = base_url(&url);
assert_eq!(base.scheme(), "s3");
assert_eq!(base.host_str(), Some("bucket"));
assert_eq!(base.path(), "");
}
}
15 changes: 4 additions & 11 deletions vortex-jni/src/file.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,6 @@ use jni::objects::JObjectArray;
use jni::objects::JString;
use jni::sys::jlong;
use jni::sys::jobject;
use object_store::path::Path;
use vortex::buffer::ByteBuffer;
use vortex::error::VortexResult;
use vortex::error::vortex_err;
Expand Down Expand Up @@ -159,11 +158,7 @@ pub extern "system" fn Java_dev_vortex_jni_NativeFiles_readMetadata(
let url = parse_uri_or_path(&uri)?;
let properties = extract_properties(env, &options)?;

let fs = object_store_fs(&url, &properties, session.handle())?;
// `FileSystem` keys are literal, already-decoded paths, so decode as `listFiles` does.
let path = Path::from_url_path(url.path())
.map_err(|_| vortex_err!("cannot parse uri as object_store Path"))?
.to_string();
let (fs, path) = object_store_fs(&url, &properties, session.handle())?;
let source = RUNTIME.block_on(async { fs.open_read(&path).await })?;

let segments = read_metadata_segments(session, source, &path, None)?;
Expand Down Expand Up @@ -226,11 +221,9 @@ pub extern "system" fn Java_dev_vortex_jni_NativeFiles_listFiles(

let properties = extract_properties(env, &options)?;

let fs = object_store_fs(&url, &properties, session.handle())?;
let prefix = Path::from_url_path(url.path())
.map_err(|_| vortex_err!("cannot parse root_path as object_store Path"))?;
let (fs, prefix) = object_store_fs(&url, &properties, session.handle())?;

let mut stream = fs.list(prefix.as_ref());
let mut stream = fs.list(&prefix);

let paths_vec = RUNTIME.block_on(async move {
let mut paths = Vec::new();
Expand Down Expand Up @@ -287,7 +280,7 @@ pub extern "system" fn Java_dev_vortex_jni_NativeFiles_delete(

let properties = extract_properties(env, &options)?;

let fs = object_store_fs(&store_url, &properties, session.handle())?;
let (fs, _path) = object_store_fs(&store_url, &properties, session.handle())?;

RUNTIME.block_on(async {
for uri in delete_uris {
Expand Down
Loading