diff --git a/Cargo.lock b/Cargo.lock index ca59b40e9dc..e23ac0d40f8 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -10326,6 +10326,8 @@ dependencies = [ "arrow-array 58.3.0", "arrow-schema 58.3.0", "async-fs", + "async-lock", + "async-trait", "futures", "jni", "object_store", diff --git a/java/vortex-jni/src/main/java/dev/vortex/api/DataSource.java b/java/vortex-jni/src/main/java/dev/vortex/api/DataSource.java index 93c5df2c61a..9a888f224f9 100644 --- a/java/vortex-jni/src/main/java/dev/vortex/api/DataSource.java +++ b/java/vortex-jni/src/main/java/dev/vortex/api/DataSource.java @@ -5,6 +5,7 @@ import com.google.common.base.Preconditions; import dev.vortex.VortexCleaner; +import dev.vortex.io.NativeReadable; import dev.vortex.jni.NativeDataSource; import java.util.Arrays; import java.util.Collections; @@ -71,6 +72,51 @@ public static DataSource open(Session session, List uris, Map readables) { + return open(session, readables, 0); + } + + /** + * Open one or more files through caller-provided byte sources instead of native storage clients. Every read the + * scan performs becomes an upcall into the corresponding {@link NativeReadable}, so this is the integration point + * for external I/O abstractions (for example Iceberg's {@code FileIO}). Each file is identified by + * {@link NativeReadable#name()}, which must be unique within {@code readables}. + * + *

The readables must remain open until this data source and all scans created from it are closed; native code + * never closes them. + * + * @param session open session + * @param readables byte sources + * @param readConcurrency maximum in-flight {@code readFully} calls across all files of this data source; {@code 0} + * selects the default. Each in-flight read occupies one native thread and typically one stream on its readable. + */ + public static DataSource open(Session session, List readables, int readConcurrency) { + Objects.requireNonNull(session, "session"); + Objects.requireNonNull(readables, "readables"); + Preconditions.checkArgument(!readables.isEmpty(), "at least one readable is required"); + Preconditions.checkArgument(readConcurrency >= 0, "readConcurrency must not be negative"); + Object[] readableArray = readables.toArray(); + String[] names = new String[readableArray.length]; + long[] lengths = new long[readableArray.length]; + for (int i = 0; i < readableArray.length; i++) { + Preconditions.checkArgument(readableArray[i] != null, "readables must not contain null values"); + NativeReadable readable = readables.get(i); + names[i] = readable.name(); + Preconditions.checkArgument(names[i] != null, "readable at index %s returned a null name", i); + lengths[i] = readable.length(); + Preconditions.checkArgument(lengths[i] >= 0, "readable for %s reported negative length", names[i]); + } + long pointer = + NativeDataSource.openFiles(session.nativePointer(), readableArray, names, lengths, readConcurrency); + return new DataSource(session, pointer); + } + /** Arrow schema of the data source (and of scans produced from it). */ public Schema arrowSchema(BufferAllocator allocator) { try (ArrowSchema schema = ArrowSchema.allocateNew(allocator)) { diff --git a/java/vortex-jni/src/main/java/dev/vortex/api/VortexWriter.java b/java/vortex-jni/src/main/java/dev/vortex/api/VortexWriter.java index e883e7911f7..a62243aab07 100644 --- a/java/vortex-jni/src/main/java/dev/vortex/api/VortexWriter.java +++ b/java/vortex-jni/src/main/java/dev/vortex/api/VortexWriter.java @@ -5,6 +5,7 @@ import com.google.common.base.Preconditions; import dev.vortex.VortexCleaner; +import dev.vortex.io.NativeWritable; import dev.vortex.jni.NativeWriter; import java.io.IOException; import java.util.Map; @@ -66,6 +67,33 @@ public static VortexWriter create( } } + /** + * Create a writer that streams the file into a caller-provided byte sink instead of a native storage client. This + * is the integration point for external I/O abstractions (for example Iceberg's {@code PositionOutputStream}). + * + *

The native side writes and flushes the sink but never closes it: after {@link #close()} returns, all bytes + * have been written and flushed, and the caller must close the sink to finalize the file. + */ + public static VortexWriter create( + Session session, NativeWritable writable, Schema arrowSchema, BufferAllocator allocator) + throws IOException { + Objects.requireNonNull(session, "session"); + Objects.requireNonNull(writable, "writable"); + Objects.requireNonNull(arrowSchema, "arrowSchema"); + Objects.requireNonNull(allocator, "allocator"); + ArrowSchema ffi = ArrowSchema.allocateNew(allocator); + try { + Data.exportSchema(allocator, arrowSchema, null, ffi); + long ptr = NativeWriter.createStream(session.nativePointer(), writable, ffi.memoryAddress()); + if (ptr <= 0) { + throw new IOException("failed to create stream writer (ptr=" + ptr + ")"); + } + return new VortexWriter(ptr); + } finally { + ffi.close(); + } + } + /** Write a batch directly from Arrow C Data Interface addresses. */ public void writeBatch(long arrowArrayAddr, long arrowSchemaAddr) throws IOException { Preconditions.checkState(!closed.get(), "writer already closed"); diff --git a/java/vortex-jni/src/main/java/dev/vortex/io/NativeReadable.java b/java/vortex-jni/src/main/java/dev/vortex/io/NativeReadable.java new file mode 100644 index 00000000000..84f264cc952 --- /dev/null +++ b/java/vortex-jni/src/main/java/dev/vortex/io/NativeReadable.java @@ -0,0 +1,47 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +package dev.vortex.io; + +import java.io.Closeable; +import java.io.IOException; +import java.nio.ByteBuffer; + +/** + * A random-access byte source supplied by the caller and read from native code. + * + *

Implementations bridge external I/O abstractions — for example Iceberg's {@code FileIO} input streams — into the + * Vortex native reader. The native side invokes {@link #readFully(long, ByteBuffer)} from its own worker threads, + * potentially many calls concurrently, so implementations must be safe for concurrent positional reads (typically by + * opening one underlying stream per concurrent call, or delegating to a positional-read API). + * + *

Lifecycle: the creator owns this object. Native code never calls {@link #close()}; close it only after every data + * source or scan built on top of it has been closed. + */ +public interface NativeReadable extends Closeable { + /** + * Unique name of this source, typically its original location (URI or file path). Used to key per-session caches + * and in native error messages, so it must be stable and unique across sources; it must not contain glob characters + * ({@code *?[}). + */ + String name(); + + /** Total length of the source in bytes. Must be cheap; called once when the source is registered. */ + long length(); + + /** + * Read exactly {@code buffer.remaining()} bytes starting at absolute {@code position} into {@code buffer}. + * + *

Unlike {@link java.io.InputStream#read}, short reads are not permitted: implementations must either fill the + * buffer completely ({@code buffer.remaining() == 0} on return) or throw. + * + *

{@code buffer} is a direct {@link ByteBuffer} wrapping native memory owned by the caller, so bytes + * written to it land in the native reader without an extra copy. It is valid only for the duration of this call: + * implementations must not retain a reference to it after returning. Implementations backed by {@code byte[]} APIs + * can bulk-{@link ByteBuffer#put(byte[], int, int) put} into it. + * + * @throws java.io.EOFException if the range extends past the end of the source + * @throws IOException if the underlying storage fails + */ + void readFully(long position, ByteBuffer buffer) throws IOException; +} diff --git a/java/vortex-jni/src/main/java/dev/vortex/io/NativeWritable.java b/java/vortex-jni/src/main/java/dev/vortex/io/NativeWritable.java new file mode 100644 index 00000000000..a86d127dad6 --- /dev/null +++ b/java/vortex-jni/src/main/java/dev/vortex/io/NativeWritable.java @@ -0,0 +1,26 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +package dev.vortex.io; + +import java.io.Closeable; +import java.io.IOException; + +/** + * A sequential byte sink supplied by the caller and written to from native code. + * + *

Implementations bridge external output abstractions — for example Iceberg's {@code PositionOutputStream} — into + * the Vortex native writer. Writes are sequential and single-threaded: the native side never issues concurrent + * {@code write} calls against the same instance. + * + *

Lifecycle: the creator owns this object. Native code calls {@link #write} and {@link #flush} but never + * {@link #close()}; after {@link dev.vortex.api.VortexWriter#close()} returns, all bytes have been written and flushed, + * and the creator must close the underlying stream to finalize it. + */ +public interface NativeWritable extends Closeable { + /** Append {@code length} bytes from {@code buffer} starting at {@code offset}. */ + void write(byte[] buffer, int offset, int length) throws IOException; + + /** Flush buffered bytes to the underlying storage. */ + void flush() throws IOException; +} diff --git a/java/vortex-jni/src/main/java/dev/vortex/jni/NativeDataSource.java b/java/vortex-jni/src/main/java/dev/vortex/jni/NativeDataSource.java index cc2aa163cee..76d09baf42e 100644 --- a/java/vortex-jni/src/main/java/dev/vortex/jni/NativeDataSource.java +++ b/java/vortex-jni/src/main/java/dev/vortex/jni/NativeDataSource.java @@ -22,6 +22,21 @@ private NativeDataSource() {} */ public static native long open(long sessionPointer, String[] uris, Map options); + /** + * Open a data source over caller-provided {@link dev.vortex.io.NativeReadable} objects. All reads become upcalls + * into the supplied readables; no native storage client is created. + * + * @param sessionPointer pointer from {@link NativeSession#newSession()} + * @param readables one {@link dev.vortex.io.NativeReadable} per file + * @param names unique name per file (from {@link dev.vortex.io.NativeReadable#name()}), parallel to + * {@code readables} + * @param lengths file sizes in bytes, parallel to {@code readables} + * @param readConcurrency maximum in-flight {@code readFully} upcalls across all files of this data source; + * {@code <= 0} selects the default + */ + public static native long openFiles( + long sessionPointer, Object[] readables, String[] names, long[] lengths, int readConcurrency); + /** Free a data source pointer. */ public static native void free(long pointer); diff --git a/java/vortex-jni/src/main/java/dev/vortex/jni/NativeWriter.java b/java/vortex-jni/src/main/java/dev/vortex/jni/NativeWriter.java index 7536712665a..4f61a381438 100644 --- a/java/vortex-jni/src/main/java/dev/vortex/jni/NativeWriter.java +++ b/java/vortex-jni/src/main/java/dev/vortex/jni/NativeWriter.java @@ -18,6 +18,12 @@ private NativeWriter() {} public static native long create( long sessionPointer, String uri, long arrowSchemaAddress, Map options); + /** + * Open a writer that streams the file into a caller-provided {@link dev.vortex.io.NativeWritable}. The native side + * writes and flushes but never closes the writable; the caller must close it after {@link #close}. + */ + public static native long createStream(long sessionPointer, Object writable, long arrowSchemaAddress); + /** * Write a batch directly from Arrow C Data Interface addresses. * diff --git a/java/vortex-jni/src/test/java/dev/vortex/io/NativeIOBridgeTest.java b/java/vortex-jni/src/test/java/dev/vortex/io/NativeIOBridgeTest.java new file mode 100644 index 00000000000..5a6d40e05db --- /dev/null +++ b/java/vortex-jni/src/test/java/dev/vortex/io/NativeIOBridgeTest.java @@ -0,0 +1,332 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +package dev.vortex.io; + +import static java.nio.charset.StandardCharsets.UTF_8; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import dev.vortex.api.DataSource; +import dev.vortex.api.Partition; +import dev.vortex.api.Scan; +import dev.vortex.api.ScanOptions; +import dev.vortex.api.Session; +import dev.vortex.api.VortexWriteSummary; +import dev.vortex.api.VortexWriter; +import dev.vortex.arrow.ArrowAllocation; +import dev.vortex.jni.NativeLoader; +import java.io.EOFException; +import java.io.IOException; +import java.io.OutputStream; +import java.nio.ByteBuffer; +import java.nio.channels.FileChannel; +import java.nio.file.Files; +import java.nio.file.Path; +import java.nio.file.StandardOpenOption; +import java.util.List; +import org.apache.arrow.c.ArrowArray; +import org.apache.arrow.c.ArrowSchema; +import org.apache.arrow.c.Data; +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.vector.IntVector; +import org.apache.arrow.vector.VarCharVector; +import org.apache.arrow.vector.VectorSchemaRoot; +import org.apache.arrow.vector.ViewVarCharVector; +import org.apache.arrow.vector.ipc.ArrowReader; +import org.apache.arrow.vector.types.pojo.ArrowType; +import org.apache.arrow.vector.types.pojo.Field; +import org.apache.arrow.vector.types.pojo.Schema; +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; + +/** Round-trip tests for the caller-provided I/O bridge ({@link NativeReadable} / {@link NativeWritable}). */ +public final class NativeIOBridgeTest { + @TempDir + Path tempDir; + + @BeforeAll + public static void loadLibrary() { + NativeLoader.loadJni(); + } + + private static Schema personSchema() { + return new Schema(List.of( + Field.notNullable("name", new ArrowType.Utf8()), + Field.notNullable("age", new ArrowType.Int(32, true)))); + } + + /** A {@link NativeReadable} over a local file, safe for concurrent positional reads. */ + private static final class FileChannelReadable implements NativeReadable { + private final String name; + private final FileChannel channel; + private final long length; + + FileChannelReadable(Path path) throws IOException { + this(path.toString(), path); + } + + FileChannelReadable(String name, Path path) throws IOException { + this.name = name; + this.channel = FileChannel.open(path, StandardOpenOption.READ); + this.length = channel.size(); + } + + @Override + public String name() { + return name; + } + + @Override + public long length() { + return length; + } + + @Override + public void readFully(long position, ByteBuffer buffer) throws IOException { + int len = buffer.remaining(); + long pos = position; + while (buffer.hasRemaining()) { + int read = channel.read(buffer, pos); + if (read < 0) { + throw new EOFException("EOF reading " + len + " bytes at position " + position); + } + pos += read; + } + } + + @Override + public void close() throws IOException { + channel.close(); + } + } + + /** A {@link NativeWritable} over a local file. */ + private static final class StreamWritable implements NativeWritable { + private final OutputStream out; + + StreamWritable(Path path) throws IOException { + this.out = Files.newOutputStream( + path, StandardOpenOption.CREATE, StandardOpenOption.TRUNCATE_EXISTING, StandardOpenOption.WRITE); + } + + @Override + public void write(byte[] buffer, int offset, int length) throws IOException { + out.write(buffer, offset, length); + } + + @Override + public void flush() throws IOException { + out.flush(); + } + + @Override + public void close() throws IOException { + out.close(); + } + } + + private void writePeopleFile(Session session, BufferAllocator allocator, Path path) throws IOException { + Schema schema = personSchema(); + VortexWriteSummary summary; + try (StreamWritable writable = new StreamWritable(path)) { + try (VortexWriter writer = VortexWriter.create(session, writable, schema, allocator); + VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator)) { + VarCharVector nameVec = (VarCharVector) root.getVector("name"); + IntVector ageVec = (IntVector) root.getVector("age"); + nameVec.allocateNew(3); + ageVec.allocateNew(3); + nameVec.setSafe(0, "Alice".getBytes(UTF_8)); + nameVec.setSafe(1, "Bob".getBytes(UTF_8)); + nameVec.setSafe(2, "Carol".getBytes(UTF_8)); + ageVec.setSafe(0, 30); + ageVec.setSafe(1, 25); + ageVec.setSafe(2, 40); + root.setRowCount(3); + + try (ArrowArray arrowArray = ArrowArray.allocateNew(allocator); + ArrowSchema arrowSchemaFfi = ArrowSchema.allocateNew(allocator)) { + Data.exportVectorSchemaRoot(allocator, root, null, arrowArray, arrowSchemaFfi); + writer.writeBatch(arrowArray.memoryAddress(), arrowSchemaFfi.memoryAddress()); + } + + // The byte counter must be live for stream writers too: bounded while + // the write task is still flushing, exact once finished. + long bytesWhileOpen = writer.bytesWritten(); + summary = writer.finish(); + assertTrue(bytesWhileOpen >= 0 && bytesWhileOpen <= summary.fileSize()); + assertEquals(summary.fileSize(), writer.bytesWritten()); + } + } + // Everything the writer counted must have reached the caller-provided sink. + assertEquals(Files.size(path), summary.fileSize()); + } + + @Test + public void testWritableThenReadableRoundTrip() throws IOException { + Path path = tempDir.resolve("bridge_roundtrip.vortex"); + BufferAllocator allocator = ArrowAllocation.rootAllocator(); + Session session = Session.create(); + + writePeopleFile(session, allocator, path); + assertTrue(Files.size(path) > 0, "stream writer should have produced bytes"); + + try (FileChannelReadable readable = new FileChannelReadable(path)) { + DataSource ds = DataSource.open(session, readable); + assertEquals(new DataSource.RowCount.Exact(3L), ds.rowCount()); + + Scan scan = ds.scan(ScanOptions.of()); + int rows = 0; + while (scan.hasNext()) { + Partition p = scan.next(); + try (ArrowReader reader = p.scanArrow(allocator)) { + while (reader.loadNextBatch()) { + VectorSchemaRoot root = reader.getVectorSchemaRoot(); + ViewVarCharVector nameOut = (ViewVarCharVector) root.getVector("name"); + IntVector ageOut = (IntVector) root.getVector("age"); + if (rows == 0) { + assertEquals("Alice", nameOut.getObject(0).toString()); + assertEquals(30, ageOut.get(0)); + } + rows += root.getRowCount(); + } + } + } + assertEquals(3, rows); + } + } + + @Test + public void testMultipleReadables() throws IOException { + Path first = tempDir.resolve("bridge_a.vortex"); + Path second = tempDir.resolve("bridge_b.vortex"); + BufferAllocator allocator = ArrowAllocation.rootAllocator(); + Session session = Session.create(); + + writePeopleFile(session, allocator, first); + writePeopleFile(session, allocator, second); + + try (FileChannelReadable firstReadable = new FileChannelReadable(first); + FileChannelReadable secondReadable = new FileChannelReadable(second)) { + DataSource ds = DataSource.open(session, List.of(firstReadable, secondReadable)); + // Only the first file is opened eagerly, so the count is an estimate until scanned. + assertEquals(6L, ds.rowCount().asOptional().orElseThrow()); + + Scan scan = ds.scan(ScanOptions.of()); + long rows = 0; + while (scan.hasNext()) { + Partition p = scan.next(); + try (ArrowReader reader = p.scanArrow(allocator)) { + while (reader.loadNextBatch()) { + rows += reader.getVectorSchemaRoot().getRowCount(); + } + } + } + assertEquals(6, rows); + } + } + + @Test + public void testDuplicateReadableNamesAreRejected() throws IOException { + Path path = tempDir.resolve("bridge_duplicate.vortex"); + BufferAllocator allocator = ArrowAllocation.rootAllocator(); + Session session = Session.create(); + + writePeopleFile(session, allocator, path); + + try (FileChannelReadable first = new FileChannelReadable("/bridge_duplicate.vortex", path); + FileChannelReadable second = new FileChannelReadable("bridge_duplicate.vortex", path)) { + RuntimeException thrown = + assertThrows(RuntimeException.class, () -> DataSource.open(session, List.of(first, second))); + assertTrue( + thrown.getMessage().contains("multiple Java readables normalize to path"), + "error should identify the normalized path collision, got: " + thrown.getMessage()); + } + } + + @Test + public void testLargeRoundTrip() throws IOException { + Path path = tempDir.resolve("bridge_large.vortex"); + BufferAllocator allocator = ArrowAllocation.rootAllocator(); + Session session = Session.create(); + Schema schema = personSchema(); + + int batches = 20; + int rowsPerBatch = 5_000; + try (StreamWritable writable = new StreamWritable(path)) { + try (VortexWriter writer = VortexWriter.create(session, writable, schema, allocator); + VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator)) { + VarCharVector nameVec = (VarCharVector) root.getVector("name"); + IntVector ageVec = (IntVector) root.getVector("age"); + for (int batch = 0; batch < batches; batch++) { + nameVec.allocateNew(rowsPerBatch); + ageVec.allocateNew(rowsPerBatch); + for (int row = 0; row < rowsPerBatch; row++) { + nameVec.setSafe(row, ("name-" + batch + "-" + row).getBytes(UTF_8)); + ageVec.setSafe(row, batch * rowsPerBatch + row); + } + root.setRowCount(rowsPerBatch); + try (ArrowArray arrowArray = ArrowArray.allocateNew(allocator); + ArrowSchema arrowSchemaFfi = ArrowSchema.allocateNew(allocator)) { + Data.exportVectorSchemaRoot(allocator, root, null, arrowArray, arrowSchemaFfi); + writer.writeBatch(arrowArray.memoryAddress(), arrowSchemaFfi.memoryAddress()); + } + } + } + } + + try (FileChannelReadable readable = new FileChannelReadable(path)) { + DataSource ds = DataSource.open(session, List.of(readable), 4); + Scan scan = ds.scan(ScanOptions.of()); + long rows = 0; + long ageSum = 0; + while (scan.hasNext()) { + Partition p = scan.next(); + try (ArrowReader reader = p.scanArrow(allocator)) { + while (reader.loadNextBatch()) { + VectorSchemaRoot root = reader.getVectorSchemaRoot(); + IntVector ageOut = (IntVector) root.getVector("age"); + for (int i = 0; i < root.getRowCount(); i++) { + ageSum += ageOut.get(i); + } + rows += root.getRowCount(); + } + } + } + long total = (long) batches * rowsPerBatch; + assertEquals(total, rows); + assertEquals(total * (total - 1) / 2, ageSum, "ages should be 0..N-1 exactly once"); + } + } + + @Test + public void testReadableExceptionPropagates() { + Session session = Session.create(); + NativeReadable failing = new NativeReadable() { + @Override + public String name() { + return "memory/failing.vortex"; + } + + @Override + public long length() { + return 1024; + } + + @Override + public void readFully(long position, ByteBuffer buffer) throws IOException { + throw new IOException("boom: injected read failure"); + } + + @Override + public void close() {} + }; + + RuntimeException thrown = assertThrows(RuntimeException.class, () -> DataSource.open(session, failing)); + assertTrue( + thrown.getMessage().contains("boom: injected read failure"), + "error should carry the Java exception message, got: " + thrown.getMessage()); + } +} diff --git a/vortex-jni/Cargo.toml b/vortex-jni/Cargo.toml index 767770db46f..6bb50cf097d 100644 --- a/vortex-jni/Cargo.toml +++ b/vortex-jni/Cargo.toml @@ -20,6 +20,8 @@ categories = { workspace = true } arrow-array = { workspace = true, features = ["ffi"] } arrow-schema = { workspace = true } async-fs = { workspace = true } +async-lock = { workspace = true } +async-trait = { workspace = true } futures = { workspace = true } jni = { workspace = true } object_store = { workspace = true, features = ["aws", "azure", "gcp"] } diff --git a/vortex-jni/src/data_source.rs b/vortex-jni/src/data_source.rs index ab86e78a056..47d9d83c577 100644 --- a/vortex-jni/src/data_source.rs +++ b/vortex-jni/src/data_source.rs @@ -16,6 +16,7 @@ use jni::objects::JLongArray; use jni::objects::JObject; use jni::objects::JObjectArray; use jni::objects::JString; +use jni::sys::jint; use jni::sys::jlong; use url::Url; use vortex::error::VortexResult; @@ -33,6 +34,7 @@ use crate::RUNTIME; use crate::dtype::export_dtype_to_arrow; use crate::errors::try_or_throw; use crate::file::extract_properties; +use crate::io::JavaFileSystem; use crate::object_store::object_store_fs; use crate::session::session_ref; @@ -115,6 +117,77 @@ pub extern "system" fn Java_dev_vortex_jni_NativeDataSource_open( }) } +/// Open a data source over caller-provided `dev.vortex.io.NativeReadable` objects. +/// +/// Unlike [`Java_dev_vortex_jni_NativeDataSource_open`], no storage client is created +/// on the native side: every read is an upcall into the corresponding Java object. +/// `paths` are opaque identifiers (typically the original file locations) used for +/// debugging and deduplication, `lengths` are the known file sizes in bytes. +/// `read_concurrency` caps in-flight `readFully` upcalls across *all* files of the +/// data source; values `<= 0` select the library default. +#[unsafe(no_mangle)] +pub extern "system" fn Java_dev_vortex_jni_NativeDataSource_openFiles( + mut env: EnvUnowned, + _class: JClass, + session_ptr: jlong, + readables: JObjectArray, + paths: JObjectArray, + lengths: JLongArray, + read_concurrency: jint, +) -> jlong { + try_or_throw(&mut env, |env| { + let session = unsafe { session_ref(session_ptr) }; + + let count = readables.len(env)?; + if count == 0 { + throw_runtime!("no readables provided"); + } + if paths.len(env)? != count || lengths.len(env)? != count { + throw_runtime!("readables, paths, and lengths must have equal length"); + } + + let mut sizes = vec![0 as jlong; count]; + lengths.get_region(env, 0, &mut sizes)?; + + let vm = env.get_java_vm()?; + let concurrency = usize::try_from(read_concurrency).ok().filter(|c| *c > 0); + let mut fs = JavaFileSystem::new(vm, session.handle(), concurrency); + let mut ordered_paths = Vec::with_capacity(count); + for idx in 0..count { + let path_obj = paths.get_element(env, idx)?; + let path: String = env.cast_local::(path_obj)?.try_to_string(env)?; + if path.contains(['*', '?', '[']) { + throw_runtime!("path '{path}' contains glob characters, which are unsupported"); + } + let size = u64::try_from(sizes[idx]) + .map_err(|_| vortex_err!("negative length for path '{path}'"))?; + + let readable = readables.get_element(env, idx)?; + if readable.is_null() { + throw_runtime!("null readable for path '{path}'"); + } + let readable = Arc::new(env.new_global_ref(&readable)?); + + // `MultiFileDataSource::with_glob` strips leading slashes (object-store paths + // are bucket-relative), so key the registry by the same normalized form. + let key = path.trim_start_matches('/').to_string(); + fs.insert(key.clone(), readable, size)?; + ordered_paths.push(key); + } + + let fs: FileSystemRef = Arc::new(fs); + let mut builder = MultiFileDataSource::new(session.clone()); + for path in ordered_paths { + builder = builder.with_glob(path, Some(Arc::clone(&fs))); + } + + let inner = RUNTIME + .block_on(builder.build()) + .map(|ds| Arc::new(ds) as DataSourceRef)?; + Ok(Box::new(NativeDataSource { inner }).into_raw()) + }) +} + /// URL with the path cleared, used as a cache key for filesystem reuse. fn base_url(url: &Url) -> Url { let mut base = url.clone(); diff --git a/vortex-jni/src/errors.rs b/vortex-jni/src/errors.rs index 4cf45a14667..0e903406ae0 100644 --- a/vortex-jni/src/errors.rs +++ b/vortex-jni/src/errors.rs @@ -11,6 +11,7 @@ use jni::sys::JNI_FALSE; use jni::sys::jboolean; use jni::sys::jobject; use vortex::error::VortexError; +use vortex::error::vortex_err; #[derive(Debug, thiserror::Error)] pub enum JNIError { @@ -38,6 +39,15 @@ impl From for JNIError { } } +impl From for VortexError { + fn from(error: JNIError) -> Self { + match error { + JNIError::Vortex(error) => error, + JNIError::Custom(error) => vortex_err!("JNI: {error}"), + } + } +} + /// Types that have a reasonable default value to use /// across the FFI. pub trait JNIDefault { diff --git a/vortex-jni/src/io/mod.rs b/vortex-jni/src/io/mod.rs new file mode 100644 index 00000000000..fa20a6a4685 --- /dev/null +++ b/vortex-jni/src/io/mod.rs @@ -0,0 +1,163 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Rust implementations of Vortex I/O traits backed by Java objects (upcalls). +//! +//! These bridges let Java callers supply their own I/O — e.g. Iceberg's `FileIO` +//! streams — instead of having the native side open storage itself. Java objects +//! are held as JNI global references and their methods are invoked from whichever +//! runtime thread executes the I/O. +//! +//! # Threading +//! +//! Vortex runtime threads (smol workers and the `blocking` pool used by +//! [`Handle::spawn_blocking`](vortex::io::runtime::Handle::spawn_blocking)) are not +//! JVM threads. Every upcall goes through [`with_jvm`], which attaches the current +//! thread on first use (detached automatically at thread exit) — re-attachment is a +//! cheap thread-local lookup. +//! +//! # JavaVM lifetime +//! +//! Each bridge struct stores the process-wide [`JavaVM`] pointer, captured once at +//! construction from the [`Env`] of the JNI entry point that created it. +//! This is safe with respect to serialization: only *Java* objects are serialized +//! (e.g. Iceberg's `FileIO` shipped to executors), and they reconstruct their native +//! state through fresh JNI calls after deserialization, at which point the entry +//! point's `Env` provides the VM again. Native bridge objects never outlive the +//! process, and JNI guarantees a single VM per process for its entire lifetime. + +mod read_at; +mod write; + +use jni::Env; +use jni::JavaVM; +pub(crate) use read_at::JavaFileSystem; +use vortex::error::VortexError; +use vortex::error::VortexResult; +pub(crate) use write::JavaWrite; + +use crate::errors::JNIError; + +/// Run `f` on the current thread with a JVM attachment. +/// +/// Attaching is a thread-local lookup when the thread is already attached; threads +/// attached here stay attached until they exit. A Java exception thrown inside `f` +/// is caught, cleared, and surfaced as a `VortexError` carrying the exception's +/// class, message, and stack trace. +pub(crate) fn with_jvm( + vm: &JavaVM, + f: impl FnOnce(&mut Env) -> Result, +) -> VortexResult { + vm.attach_current_thread(f).map_err(VortexError::from) +} + +#[cfg(test)] +pub(crate) mod tests { + use std::sync::Arc; + use std::sync::LazyLock; + use std::time::Duration; + + use jni::InitArgsBuilder; + use jni::JValue; + use jni::JavaVM; + use jni::errors::StartJvmError; + use jni::objects::JObject; + use jni::refs::Global; + use vortex::error::VortexResult; + use vortex::error::vortex_bail; + use vortex::io::filesystem::FileSystem; + use vortex::io::runtime::BlockingRuntime; + use vortex::io::runtime::current::CurrentThreadRuntime; + + use super::JavaFileSystem; + use super::with_jvm; + + /// Launch (or reuse) the process-wide JVM backing upcall tests. Returns `None` + /// when no JVM installation can be located, so tests skip on machines without a + /// JDK; any other launch failure panics. + pub(crate) fn test_vm() -> Option<&'static JavaVM> { + static VM: LazyLock> = LazyLock::new(|| { + let args = InitArgsBuilder::new() + .option("-Xcheck:jni") + .build() + .expect("valid JVM init args"); + match JavaVM::new(args) { + Ok(vm) => Some(vm), + Err(StartJvmError::NotFound(_)) => { + eprintln!("skipping JNI upcall test: no JVM found (set JAVA_HOME to run it)"); + None + } + Err(e) => panic!("failed to launch test JVM: {e}"), + } + }); + VM.as_ref() + } + + /// A fresh `java.lang.Object` global ref plus a `java.lang.ref.WeakReference` + /// observing the same object, so tests can detect when the global ref is deleted. + pub(crate) fn object_with_weak_ref( + vm: &JavaVM, + ) -> VortexResult<(Global>, Global>)> { + with_jvm(vm, |env| { + let obj = + env.new_object(jni::jni_str!("java/lang/Object"), jni::jni_sig!("()V"), &[])?; + let weak = env.new_object( + jni::jni_str!("java/lang/ref/WeakReference"), + jni::jni_sig!("(Ljava/lang/Object;)V"), + &[JValue::Object(&obj)], + )?; + Ok((env.new_global_ref(&obj)?, env.new_global_ref(&weak)?)) + }) + } + + /// Assert that the referent observed by `weak` becomes collectible, i.e. that no + /// JNI global reference pins it anymore. + pub(crate) fn assert_weak_ref_clears( + vm: &JavaVM, + weak: &Global>, + ) -> VortexResult<()> { + for _ in 0..100 { + let cleared = with_jvm(vm, |env| { + env.call_static_method( + jni::jni_str!("java/lang/System"), + jni::jni_str!("gc"), + jni::jni_sig!("()V"), + &[], + )?; + let referent = env + .call_method( + weak.as_ref(), + jni::jni_str!("get"), + jni::jni_sig!("()Ljava/lang/Object;"), + &[], + )? + .l()?; + Ok(referent.is_null()) + })?; + if cleared { + return Ok(()); + } + std::thread::sleep(Duration::from_millis(10)); + } + vortex_bail!("JNI global reference leaked: weak reference still has a referent") + } + + #[test] + fn test_file_system_drop_releases_global_refs() -> VortexResult<()> { + let Some(vm) = test_vm() else { + return Ok(()); + }; + let (readable, weak) = object_with_weak_ref(vm)?; + + let runtime = CurrentThreadRuntime::new(); + let mut fs = JavaFileSystem::new(vm.clone(), runtime.handle(), None); + fs.insert("data/file.vortex".to_string(), Arc::new(readable), 4)?; + // `open_read` clones the global ref into a `JavaReadable`; both it and the + // file system must let go of the Java object once dropped. + let reader = runtime.block_on(fs.open_read("data/file.vortex"))?; + drop(fs); + drop(reader); + + assert_weak_ref_clears(vm, &weak) + } +} diff --git a/vortex-jni/src/io/read_at.rs b/vortex-jni/src/io/read_at.rs new file mode 100644 index 00000000000..a08652c7cc1 --- /dev/null +++ b/vortex-jni/src/io/read_at.rs @@ -0,0 +1,293 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::fmt; +use std::fmt::Debug; +use std::sync::Arc; + +use async_lock::Semaphore; +use async_trait::async_trait; +use futures::FutureExt; +use futures::StreamExt; +use futures::future::BoxFuture; +use futures::stream; +use futures::stream::BoxStream; +use jni::JValue; +use jni::JavaVM; +use jni::objects::JObject; +use jni::refs::Global; +use vortex::array::buffer::BufferHandle; +use vortex::buffer::Alignment; +use vortex::buffer::ByteBufferMut; +use vortex::error::VortexResult; +use vortex::error::vortex_bail; +use vortex::error::vortex_err; +use vortex::io::CoalesceConfig; +use vortex::io::VortexReadAt; +use vortex::io::filesystem::FileListing; +use vortex::io::filesystem::FileSystem; +use vortex::io::runtime::Handle; +use vortex::utils::aliases::hash_map::EntryRef; +use vortex::utils::aliases::hash_map::HashMap; + +use crate::io::with_jvm; + +/// Default number of concurrent `readFully` upcalls to allow across all files of one +/// [`JavaFileSystem`]. Matches the object-store default since the backing storage is +/// typically remote. +const DEFAULT_CONCURRENCY: usize = 192; + +/// Shared cap on in-flight `readFully` upcalls across every file of one +/// [`JavaFileSystem`], so a wide scan cannot pin `files x concurrency` blocking +/// threads and Java streams. +#[derive(Clone)] +struct UpcallLimiter { + semaphore: Arc, + concurrency: usize, +} + +impl UpcallLimiter { + fn new(concurrency: usize) -> Self { + Self { + semaphore: Arc::new(Semaphore::new(concurrency)), + concurrency, + } + } +} + +/// A [`VortexReadAt`] backed by a Java object implementing `dev.vortex.io.NativeReadable`. +/// +/// Positional reads are forwarded as blocking `readFully(long, ByteBuffer)` upcalls +/// executed on the runtime's blocking pool. The destination is a direct `ByteBuffer` +/// wrapping the Rust-side allocation, so Java writes land in the buffer handed to the +/// scan without a copy through a heap `byte[]`. The file size is captured at +/// construction time (Java callers know it from metadata), so `size()` never crosses +/// the JNI boundary. +pub(crate) struct JavaReadable { + vm: JavaVM, + readable: Arc>>, + len: u64, + handle: Handle, + limiter: UpcallLimiter, +} + +impl JavaReadable { + fn new( + vm: JavaVM, + readable: Arc>>, + len: u64, + handle: Handle, + limiter: UpcallLimiter, + ) -> Self { + Self { + vm, + readable, + len, + handle, + limiter, + } + } +} + +impl VortexReadAt for JavaReadable { + fn coalesce_config(&self) -> Option { + // Upcalls have a fixed JNI overhead and the backing storage is usually + // remote, so favor fewer, larger reads. + Some(CoalesceConfig::object_storage()) + } + + fn concurrency(&self) -> usize { + self.limiter.concurrency + } + + fn size(&self) -> BoxFuture<'static, VortexResult> { + let len = self.len; + async move { Ok(len) }.boxed() + } + + fn read_at( + &self, + offset: u64, + length: usize, + alignment: Alignment, + ) -> BoxFuture<'static, VortexResult> { + let vm = self.vm.clone(); + let readable = Arc::clone(&self.readable); + let len = self.len; + let handle = self.handle.clone(); + let semaphore = Arc::clone(&self.limiter.semaphore); + + async move { + // Take a permit before occupying a blocking thread. The lock-free + // `try_acquire_arc` fast path deliberately barges ahead of queued + // waiters: slight unfairness is fine, the limiter must never become + // the bottleneck itself. + let permit = match semaphore.try_acquire_arc() { + Some(permit) => permit, + None => semaphore.acquire_arc().await, + }; + + handle + .spawn_blocking(move || { + // Keep the permit with the blocking work. Dropping the read future can + // cancel its task handle, but it cannot interrupt a `readFully` upcall + // that has already started. + let _permit = permit; + let end = offset + .checked_add(length as u64) + .ok_or_else(|| vortex_err!("read {offset}+{length} overflows u64"))?; + if end > len { + vortex_bail!("read {offset}..{end} out of bounds for file of length {len}"); + } + // `java.nio.Buffer` capacities are Java ints. + if i32::try_from(length).is_err() { + vortex_bail!("read length {length} exceeds ByteBuffer limit"); + } + let joffset = i64::try_from(offset) + .map_err(|_| vortex_err!("read offset {offset} exceeds i64"))?; + + let mut buffer = ByteBufferMut::with_capacity_aligned(length, alignment); + with_jvm(&vm, |env| { + // SAFETY: the pointer covers `length` bytes of live allocation, and + // the direct buffer wrapping it does not outlive this call — + // `readFully` is synchronous and the interface forbids retaining + // the buffer after it returns. + let dst = unsafe { + env.new_direct_byte_buffer( + buffer.spare_capacity_mut().as_mut_ptr().cast(), + length, + )? + }; + env.call_method( + readable.as_ref(), + jni::jni_str!("readFully"), + jni::jni_sig!("(JLjava/nio/ByteBuffer;)V"), + &[JValue::Long(joffset), JValue::Object(dst.as_ref())], + )?; + // Guard against implementations that return without filling the + // buffer, which would hand uninitialized memory to the scan. + let remaining = env + .call_method( + &dst, + jni::jni_str!("remaining"), + jni::jni_sig!("()I"), + &[], + )? + .i()?; + if remaining != 0 { + return Err(vortex_err!( + "readFully returned with {remaining} of {length} bytes unfilled" + ) + .into()); + } + Ok(()) + }) + .map_err(|e| e.with_context("readFully upcall failed"))?; + + // SAFETY: `readFully` wrote all `length` bytes through the direct + // buffer, verified by the `remaining()` check above. + unsafe { buffer.set_len(length) }; + Ok(BufferHandle::new_host(buffer.freeze())) + }) + .await + } + .boxed() + } +} + +struct JavaFileEntry { + readable: Arc>>, + size: u64, +} + +/// A [`FileSystem`] over a fixed set of Java-provided readables, keyed by path. +/// +/// Built by `NativeDataSource.openFiles`: every file the data source may touch is +/// registered up front together with its size, so `head` (and therefore exact-path +/// glob resolution) is answered without any upcall. `open_read` wraps the registered +/// Java object in a [`JavaReadable`]. +pub(crate) struct JavaFileSystem { + vm: JavaVM, + files: HashMap, + handle: Handle, + limiter: UpcallLimiter, +} + +impl JavaFileSystem { + pub(crate) fn new(vm: JavaVM, handle: Handle, concurrency: Option) -> Self { + Self { + vm, + files: HashMap::new(), + handle, + limiter: UpcallLimiter::new(concurrency.unwrap_or(DEFAULT_CONCURRENCY)), + } + } + + /// Register a Java readable for `path` with a known size. + pub(crate) fn insert( + &mut self, + path: String, + readable: Arc>>, + size: u64, + ) -> VortexResult<()> { + match self.files.entry_ref(&path) { + EntryRef::Occupied(_) => { + vortex_bail!("multiple Java readables normalize to path '{path}'"); + } + EntryRef::Vacant(v) => v.insert(JavaFileEntry { readable, size }), + }; + + Ok(()) + } +} + +impl Debug for JavaFileSystem { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("JavaFileSystem") + .field("files", &self.files.keys()) + .finish() + } +} + +#[async_trait] +impl FileSystem for JavaFileSystem { + fn list(&self, prefix: &str) -> BoxStream<'_, VortexResult> { + let listings: Vec> = self + .files + .iter() + .filter(|(path, _)| path.starts_with(prefix)) + .map(|(path, entry)| { + Ok(FileListing { + path: path.clone(), + size: Some(entry.size), + }) + }) + .collect(); + stream::iter(listings).boxed() + } + + async fn head(&self, path: &str) -> VortexResult> { + Ok(self.files.get(path).map(|entry| FileListing { + path: path.to_string(), + size: Some(entry.size), + })) + } + + async fn open_read(&self, path: &str) -> VortexResult> { + let entry = self + .files + .get(path) + .ok_or_else(|| vortex_err!("no Java readable registered for path '{path}'"))?; + Ok(Arc::new(JavaReadable::new( + self.vm.clone(), + Arc::clone(&entry.readable), + entry.size, + self.handle.clone(), + self.limiter.clone(), + ))) + } + + async fn delete(&self, path: &str) -> VortexResult<()> { + vortex_bail!("delete('{path}') is not supported by a Java-readable file system") + } +} diff --git a/vortex-jni/src/io/write.rs b/vortex-jni/src/io/write.rs new file mode 100644 index 00000000000..ffadf400e17 --- /dev/null +++ b/vortex-jni/src/io/write.rs @@ -0,0 +1,166 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::io; +use std::sync::Arc; + +use jni::JValue; +use jni::JavaVM; +use jni::objects::JObject; +use jni::refs::Global; +use vortex::io::IoBuf; +use vortex::io::VortexWrite; + +use crate::io::with_jvm; + +/// Largest chunk pushed through a single `write` upcall. Java arrays are indexed by +/// `int`, and keeping chunks bounded also keeps per-call array allocations modest. +const MAX_WRITE_CHUNK: usize = 1 << 30; + +/// A [`VortexWrite`] backed by a Java object implementing `dev.vortex.io.NativeWritable`. +/// +/// Bytes are forwarded as blocking `write(byte[], int, int)` upcalls, `flush` maps to +/// `flush()`, and `shutdown` maps to `flush()` as well: the Java caller created the +/// underlying stream and remains responsible for closing it once the writer finishes. +/// Writes run inline on the runtime thread driving the write task, which is attached +/// to the JVM on first use. +pub(crate) struct JavaWrite { + vm: JavaVM, + writable: Arc>>, +} + +impl JavaWrite { + pub(crate) fn new(vm: JavaVM, writable: Arc>>) -> Self { + Self { vm, writable } + } + + fn write_slice(&self, bytes: &[u8]) -> io::Result<()> { + for chunk in bytes.chunks(MAX_WRITE_CHUNK) { + let jlen = i32::try_from(chunk.len()).map_err(io::Error::other)?; + with_jvm(&self.vm, |env| { + let array = env.byte_array_from_slice(chunk)?; + env.call_method( + self.writable.as_ref(), + jni::jni_str!("write"), + jni::jni_sig!("([BII)V"), + &[ + JValue::Object(array.as_ref()), + JValue::Int(0), + JValue::Int(jlen), + ], + )?; + Ok(()) + }) + .map_err(io::Error::other)?; + } + Ok(()) + } + + fn flush_upcall(&self) -> io::Result<()> { + with_jvm(&self.vm, |env| { + env.call_method( + self.writable.as_ref(), + jni::jni_str!("flush"), + jni::jni_sig!("()V"), + &[], + )?; + Ok(()) + }) + .map_err(io::Error::other) + } +} + +impl VortexWrite for JavaWrite { + async fn write_all(&mut self, buffer: B) -> io::Result { + self.write_slice(buffer.as_slice())?; + Ok(buffer) + } + + async fn flush(&mut self) -> io::Result<()> { + self.flush_upcall() + } + + async fn shutdown(&mut self) -> io::Result<()> { + self.flush_upcall() + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use jni::JValue; + use jni::objects::JByteArray; + use vortex::error::VortexResult; + use vortex::error::vortex_err; + + use super::JavaWrite; + use crate::io::tests::assert_weak_ref_clears; + use crate::io::tests::test_vm; + use crate::io::with_jvm; + + /// Write upcalls from a fresh native thread must reach the Java object, and once + /// the writer (and every other native reference) is dropped, its JNI global ref + /// must be deleted so the Java object becomes collectible. + #[test] + fn test_write_upcalls_release_global_ref() -> VortexResult<()> { + let Some(vm) = test_vm() else { + return Ok(()); + }; + + // `ByteArrayOutputStream` has the same `write([BII)V`/`flush()V` shape as + // `dev.vortex.io.NativeWritable`, so it stands in for a caller-provided sink. + let (writable, content, weak) = with_jvm(vm, |env| { + let sink = env.new_object( + jni::jni_str!("java/io/ByteArrayOutputStream"), + jni::jni_sig!("()V"), + &[], + )?; + let weak = env.new_object( + jni::jni_str!("java/lang/ref/WeakReference"), + jni::jni_sig!("(Ljava/lang/Object;)V"), + &[JValue::Object(&sink)], + )?; + Ok(( + env.new_global_ref(&sink)?, + env.new_global_ref(&sink)?, + env.new_global_ref(&weak)?, + )) + })?; + + let thread_vm = vm.clone(); + std::thread::spawn(move || { + let write = JavaWrite::new(thread_vm.clone(), Arc::new(writable)); + write.write_slice(b"vortex")?; + write.flush_upcall()?; + // `with_jvm` attaches permanently: the thread stays attached until it + // exits, at which point the attachment is torn down automatically. + let attached = thread_vm + .is_thread_attached() + .map_err(|e| vortex_err!("is_thread_attached failed: {e}"))?; + assert!(attached, "write upcalls should leave the thread attached"); + VortexResult::Ok(()) + }) + .join() + .map_err(|_| vortex_err!("writer thread panicked"))??; + + let written = with_jvm(vm, |env| { + let array = env + .call_method( + content.as_ref(), + jni::jni_str!("toByteArray"), + jni::jni_sig!("()[B"), + &[], + )? + .l()?; + let array = env.cast_local::(array)?; + let mut bytes = vec![0i8; array.len(env)?]; + array.get_region(env, 0, &mut bytes)?; + Ok(bytes.into_iter().map(|b| b as u8).collect::>()) + })?; + assert_eq!(written, b"vortex"); + + drop(content); + assert_weak_ref_clears(vm, &weak) + } +} diff --git a/vortex-jni/src/lib.rs b/vortex-jni/src/lib.rs index a7d988c3c6a..c8369c28bff 100644 --- a/vortex-jni/src/lib.rs +++ b/vortex-jni/src/lib.rs @@ -23,6 +23,7 @@ mod dtype; mod errors; mod expression; mod file; +mod io; mod logging; mod object_store; mod runtime; diff --git a/vortex-jni/src/writer.rs b/vortex-jni/src/writer.rs index 309248fd73c..c1ffea2f14f 100644 --- a/vortex-jni/src/writer.rs +++ b/vortex-jni/src/writer.rs @@ -68,6 +68,7 @@ use crate::dtype::import_arrow_schema; use crate::errors::JNIError; use crate::errors::try_or_throw; use crate::file::extract_properties; +use crate::io::JavaWrite; use crate::object_store::make_object_store; use crate::session::session_ref; @@ -455,6 +456,61 @@ pub extern "system" fn Java_dev_vortex_jni_NativeWriter_create( }) } +/// Create a writer that streams the file into a caller-provided +/// `dev.vortex.io.NativeWritable` instead of a native storage client. +/// +/// Bytes are pushed through blocking `write`/`flush` upcalls on the runtime thread +/// driving the write task. The Java caller owns the underlying stream and must close +/// it after `NativeWriter.close` returns; the native side only flushes. +#[unsafe(no_mangle)] +pub extern "system" fn Java_dev_vortex_jni_NativeWriter_createStream( + mut env: EnvUnowned, + _class: JClass, + session_ptr: jlong, + writable: JObject, + arrow_schema_addr: jlong, +) -> jlong { + try_or_throw(&mut env, |env| { + if session_ptr == 0 { + throw_runtime!("null session pointer"); + } + if arrow_schema_addr == 0 { + throw_runtime!("null arrow schema address"); + } + if writable.is_null() { + throw_runtime!("null writable"); + } + let session = unsafe { session_ref(session_ptr) }; + + let arrow_schema = Arc::new(import_arrow_schema(arrow_schema_addr)?); + let write_schema = session.arrow().from_arrow_schema(arrow_schema.as_ref())?; + + let vm = env.get_java_vm()?; + let writable = Arc::new(env.new_global_ref(&writable)?); + let (tx, rx) = mpsc::channel(WRITE_CHANNEL_CAPACITY); + let stream = ArrayStreamAdapter::new(write_schema.clone(), rx); + let write_options = write_options_for_schema(session, &write_schema); + + let mut write = CountingVortexWrite::new(JavaWrite::new(vm, writable)); + let bytes_written = write.counter(); + let handle = session.handle().spawn(async move { + let summary = write_options.write(&mut write, stream).await?; + write.shutdown().await?; + Ok(summary) + }); + + Ok(Box::new(NativeWriter::new( + session.clone(), + arrow_schema, + write_schema, + bytes_written, + handle, + tx, + )) + .into_raw()) + }) +} + /// Write a batch to the Vortex file directly from Arrow C Data Interface pointers. #[unsafe(no_mangle)] pub extern "system" fn Java_dev_vortex_jni_NativeWriter_writeBatch( diff --git a/vortex-utils/src/aliases/hash_map.rs b/vortex-utils/src/aliases/hash_map.rs index 07047562b46..5937e2c1fc1 100644 --- a/vortex-utils/src/aliases/hash_map.rs +++ b/vortex-utils/src/aliases/hash_map.rs @@ -9,6 +9,8 @@ pub type RandomState = hashbrown::DefaultHashBuilder; pub type HashMap = hashbrown::HashMap; /// Entry type for HashMap. pub type Entry<'a, K, V, S> = hashbrown::hash_map::Entry<'a, K, V, S>; +/// Entry type for HashMap. +pub type EntryRef<'a, 'b, K, Q, V, S, A> = hashbrown::hash_map::EntryRef<'a, 'b, K, Q, V, S, A>; /// IntoIter type for HashMap. pub type IntoIter = hashbrown::hash_map::IntoIter; /// HashTable type alias.