Use async_fs instead of LocalFileSystem in java bindings for local files

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

Author: robert3005

Cargo.lock

@@ -42,8 +42,8 @@ private VortexWriter(long pointer) {
     }
 
     /**
-     * Create a writer that streams records into the file at {@code uri}. The Arrow schema describes the exact layout of
-     * every batch written.
+     * Create a writer that streams records into the file at {@code uri}. The path may be a full URI or a plain local
+     * filesystem path. The Arrow schema describes the exact layout of every batch written.
      */
     public static VortexWriter create(
             Session session, String uri, Schema arrowSchema, Map<String, String> options, BufferAllocator allocator)

java/vortex-jni/src/main/java/dev/vortex/api/VortexWriter.java

@@ -68,6 +68,22 @@ public void testCreateWriter() throws IOException {
         assertTrue(Files.exists(outputPath), "output file should exist");
     }
 
+    @Test
+    public void testCreateWriterPlainLocalPath() throws IOException {
+        Path outputPath = tempDir.resolve("test_create_plain_path.vortex");
+        String writePath = outputPath.toAbsolutePath().toString();
+
+        BufferAllocator allocator = ArrowAllocation.rootAllocator();
+        Map<String, String> options = new HashMap<>();
+
+        Session session = Session.create();
+        try (VortexWriter writer = VortexWriter.create(session, writePath, personSchema(), options, allocator)) {
+            assertNotNull(writer);
+        }
+
+        assertTrue(Files.exists(outputPath), "output file should exist");
+    }
+
     @Test
     public void testWriteBatch() throws IOException {
         Path outputPath = tempDir.resolve("test_ffi.vortex");

java/vortex-jni/src/test/java/dev/vortex/jni/JNIWriterTest.java

@@ -19,6 +19,7 @@ categories = { workspace = true }
 [dependencies]
 arrow-array = { workspace = true, features = ["ffi"] }
 arrow-schema = { workspace = true }
+async-fs = { workspace = true }
 futures = { workspace = true }
 jni = { workspace = true }
 object_store = { workspace = true, features = ["aws", "azure", "gcp"] }

vortex-jni/Cargo.toml

@@ -6,12 +6,14 @@
 //! Writes go through an in-flight queue of at most [`WRITE_CHANNEL_CAPACITY`] pending
 //! batches on the same thread that drives the current-thread runtime.
 
+use std::path::PathBuf;
 use std::sync::Arc;
 
 use arrow_array::RecordBatch;
 use arrow_array::StructArray;
 use arrow_array::ffi::FFI_ArrowArray;
 use arrow_array::ffi::FFI_ArrowSchema;
+use async_fs::File;
 use futures::SinkExt;
 use futures::channel::mpsc;
 use jni::EnvUnowned;
@@ -22,7 +24,8 @@ use jni::sys::JNI_FALSE;
 use jni::sys::JNI_TRUE;
 use jni::sys::jboolean;
 use jni::sys::jlong;
-use object_store::path::Path;
+use object_store::ObjectStore;
+use object_store::path::Path as ObjectStorePath;
 use url::Url;
 use vortex::array::ArrayRef;
 use vortex::array::arrow::FromArrowArray;
@@ -52,6 +55,40 @@ 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,
+    },
+}
+
+fn resolve_write_target(
+    uri: &str,
+    properties: &HashMap<String, String>,
+) -> VortexResult<WriteTarget> {
+    match Url::parse(uri) {
+        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));
+        }
+        Ok(url) if uri.contains("://") => {
+            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 });
+        }
+        Err(err) if uri.contains("://") => {
+            return Err(vortex_err!("invalid URL {uri}: {err}"));
+        }
+        _ => {}
+    }
+
+    Ok(WriteTarget::Local(PathBuf::from(uri)))
+}
+
 /// Native writer holding a write-task handle and a sender that Java pushes batches into.
 pub struct NativeWriter {
     handle: Option<Task<VortexResult<WriteSummary>>>,
@@ -136,21 +173,27 @@ pub extern "system" fn Java_dev_vortex_jni_NativeWriter_create(
         let write_schema = import_dtype_from_arrow(arrow_schema_addr)?;
 
         let file_path: String = uri.try_to_string(env)?;
-        let url = Url::parse(&file_path)
-            .map_err(|e| JNIError::Vortex(vortex_err!("invalid URL {file_path}: {e}")))?;
         let properties: HashMap<String, String> = extract_properties(env, &options)?;
-        let path = Path::from_url_path(url.path())
-            .map_err(|_| vortex_err!("invalid object_store path: {}", url.path()))?;
-
-        let store = make_object_store(&url, &properties)?;
+        let target = resolve_write_target(&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 {
-            let mut write = ObjectStoreWrite::new(Arc::new(Compat::new(store)), &path).await?;
-            let summary = session.write_options().write(&mut write, stream).await?;
-            write.shutdown().await?;
-            Ok(summary)
+            match target {
+                WriteTarget::Local(path) => {
+                    let mut file = File::create(path).await?;
+                    let summary = session.write_options().write(&mut file, stream).await?;
+                    file.shutdown().await?;
+                    Ok(summary)
+                }
+                WriteTarget::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?;
+                    write.shutdown().await?;
+                    Ok(summary)
+                }
+            }
         });
 
         Ok(Box::new(NativeWriter::new(write_schema, handle, tx)).into_raw())