more

Signed-off-by: Robert Kruszewski <github@robertk.io>

Author: robert3005

vortex-jni/src/writer.rs

@@ -55,38 +55,30 @@ use crate::session::session_ref;
 /// the writer is felt on the Java thread producing batches.
 const WRITE_CHANNEL_CAPACITY: usize = 4;
 
-enum WriteTarget {
-    Local(PathBuf),
-    ObjectStore {
-        store: Arc<dyn ObjectStore>,
-        path: ObjectStorePath,
-    },
+enum ResolvedStore {
+    ObjectStore(Arc<dyn ObjectStore>, ObjectStorePath),
+    Path(PathBuf),
 }
 
-fn resolve_write_target(
-    uri: &str,
+fn resolve_store(
+    url_or_path: &str,
     properties: &HashMap<String, String>,
-) -> VortexResult<WriteTarget> {
-    match Url::parse(uri) {
+) -> VortexResult<ResolvedStore> {
+    match Url::parse(url_or_path) {
         Ok(url) if url.scheme() == "file" => {
             let path = url
                 .to_file_path()
-                .map_err(|_| vortex_err!("invalid file URL: {uri}"))?;
-            return Ok(WriteTarget::Local(path));
+                .map_err(|_| vortex_err!("invalid file URL: {url_or_path}"))?;
+            Ok(ResolvedStore::Path(path))
         }
-        Ok(url) if uri.contains("://") => {
+        Ok(url) => {
             let path = ObjectStorePath::from_url_path(url.path())
                 .map_err(|_| vortex_err!("invalid object_store path: {}", url.path()))?;
             let store = make_object_store(&url, properties)?;
-            return Ok(WriteTarget::ObjectStore { store, path });
+            Ok(ResolvedStore::ObjectStore(store, path))
         }
-        Err(err) if uri.contains("://") => {
-            return Err(vortex_err!("invalid URL {uri}: {err}"));
-        }
-        _ => {}
+        Err(_) => Ok(ResolvedStore::Path(PathBuf::from(url_or_path))),
     }
-
-    Ok(WriteTarget::Local(PathBuf::from(uri)))
 }
 
 /// Native writer holding a write-task handle and a sender that Java pushes batches into.
@@ -174,13 +166,13 @@ pub extern "system" fn Java_dev_vortex_jni_NativeWriter_create(
 
         let file_path: String = uri.try_to_string(env)?;
         let properties: HashMap<String, String> = extract_properties(env, &options)?;
-        let target = resolve_write_target(&file_path, &properties)?;
+        let resolved = resolve_store(&file_path, &properties)?;
         let (tx, rx) = mpsc::channel(WRITE_CHANNEL_CAPACITY);
         let stream = ArrayStreamAdapter::new(write_schema.clone(), rx);
 
         let handle = session.handle().spawn(async move {
-            match target {
-                WriteTarget::Local(path) => {
+            match resolved {
+                ResolvedStore::Path(path) => {
                     if let Some(parent) = path.parent().filter(|p| !p.as_os_str().is_empty()) {
                         async_fs::create_dir_all(parent).await?;
                     }
@@ -189,7 +181,7 @@ pub extern "system" fn Java_dev_vortex_jni_NativeWriter_create(
                     file.shutdown().await?;
                     Ok(summary)
                 }
-                WriteTarget::ObjectStore { store, path } => {
+                ResolvedStore::ObjectStore(store, path) => {
                     let mut write =
                         ObjectStoreWrite::new(Arc::new(Compat::new(store)), &path).await?;
                     let summary = session.write_options().write(&mut write, stream).await?;

vortex-python/src/object_store/resolve.rs

@@ -10,6 +10,7 @@ use object_store::path::Path;
 use object_store::registry::ObjectStoreRegistry;
 use url::Url;
 use vortex::error::VortexResult;
+use vortex::error::vortex_err;
 
 use crate::object_store::registry::Registry;
 
@@ -32,6 +33,12 @@ pub(crate) fn resolve_store(
         None => {
             // If the URL does not parse
             match Url::parse(url_or_path) {
+                Ok(url) if url.scheme() == "file" => {
+                    let path = url
+                        .to_file_path()
+                        .map_err(|_| vortex_err!("invalid file URL: {url_or_path}"))?;
+                    Ok(ResolvedStore::Path(path))
+                }
                 Ok(url) => {
                     let (store, path) = REGISTRY.resolve(&url)?;
                     Ok(ResolvedStore::ObjectStore(store, path))
@@ -92,6 +99,13 @@ mod test {
             PathBuf::from("/my/absolute/path")
         );
 
+        assert_eq!(
+            resolve_store("file:///my/absolute/path", None)
+                .unwrap()
+                .unwrap_path(),
+            PathBuf::from("/my/absolute/path")
+        );
+
         let (_store, path) = resolve_store("s3://my-bucket/first/second/third/", None)
             .unwrap()
             .unwrap_store();