From 71b18b5f14f0a06c81c8bf0fc0145a5b6f09acc5 Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Thu, 23 Jul 2026 10:40:49 +0100 Subject: [PATCH 1/4] Expose LayoutStrategy::buffered_bytes through the jni boundary Signed-off-by: Robert Kruszewski --- .../java/dev/vortex/api/VortexWriter.java | 17 ++++++ .../java/dev/vortex/jni/NativeWriter.java | 3 ++ vortex-jni/src/writer.rs | 52 ++++++++++++++----- 3 files changed, 58 insertions(+), 14 deletions(-) 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 a62243aab07..8ebe8405df4 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 @@ -125,6 +125,23 @@ public synchronized long bytesWritten() { return bytesWritten; } + /** + * Return the number of uncompressed bytes accepted by the writer but not yet written to the sink. + * + *

Together with {@link #bytesWritten()}, this lets callers estimate the in-progress file size: bytes that + * reached the sink are already compressed, while buffered bytes are still uncompressed and will shrink by + * roughly the file's observed compression ratio once flushed. After {@link #finish()}, this is zero. + */ + public synchronized long bufferedBytes() { + if (summary != null) { + return 0; + } + Preconditions.checkState(!closed.get(), "writer closed without a write summary"); + long bufferedBytes = NativeWriter.bufferedBytes(pointer); + Preconditions.checkState(bufferedBytes >= 0, "native writer returned an invalid buffered byte count"); + return bufferedBytes; + } + /** * Flush pending batches, finalize the file, and return its statistics and physical sizes. * 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 4f61a381438..f0c0a7471ad 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 @@ -37,6 +37,9 @@ public static native long create( /** Number of bytes successfully written to the underlying sink so far. */ public static native long bytesWritten(long writerPointer); + /** Number of uncompressed bytes buffered by the native writer that have not yet reached the sink. */ + public static native long bufferedBytes(long writerPointer); + /** Flush and close the writer. Must be called exactly once. */ public static native void close(long writerPointer); diff --git a/vortex-jni/src/writer.rs b/vortex-jni/src/writer.rs index 53649eef4ab..34b96704b49 100644 --- a/vortex-jni/src/writer.rs +++ b/vortex-jni/src/writer.rs @@ -47,17 +47,17 @@ use vortex::error::VortexResult; use vortex::error::vortex_err; use vortex::expr::stats::Stat; use vortex::expr::stats::StatsProvider; -use vortex::file::CountingVortexWrite; -use vortex::file::WriteOptionsSessionExt; -use vortex::file::WriteStrategyBuilder; use vortex::file::WriteSummary; use vortex::file::multi::parse_uri_or_path; +use vortex::file::{ALLOWED_ENCODINGS, CountingVortexWrite}; +use vortex::file::{WriteOptionsSessionExt, WriteStrategyBuilder}; use vortex::io::VortexWrite; use vortex::io::compat::Compat; use vortex::io::object_store::ObjectStoreWrite; use vortex::io::runtime::BlockingRuntime; use vortex::io::runtime::Task; use vortex::io::session::RuntimeSessionExt; +use vortex::layout::LayoutStrategy; use vortex::session::VortexSession; use vortex::utils::aliases::hash_map::HashMap; use vortex_arrow::ArrowSessionExt; @@ -99,21 +99,18 @@ fn resolve_store( } } -fn write_options_for_schema( - session: &VortexSession, - write_schema: &DType, -) -> vortex::file::VortexWriteOptions { +fn write_strategy_for_schema(write_schema: &DType) -> Arc { let variant_paths = variant_field_paths(write_schema); if variant_paths.is_empty() { - return session.write_options(); + return WriteStrategyBuilder::default().build(); } - let mut allowed = vortex::file::ALLOWED_ENCODINGS.clone(); + let mut allowed = ALLOWED_ENCODINGS.clone(); allowed.insert(ParquetVariant.id()); - let strategy = WriteStrategyBuilder::default().with_allow_encodings(allowed); - - session.write_options().with_strategy(strategy.build()) + WriteStrategyBuilder::default() + .with_allow_encodings(allowed) + .build() } fn variant_field_paths(dtype: &DType) -> Vec { @@ -145,6 +142,7 @@ pub struct NativeWriter { arrow_schema: SchemaRef, write_schema: DType, bytes_written: Arc, + strategy: Arc, sender: mpsc::Sender>, } @@ -154,6 +152,7 @@ impl NativeWriter { arrow_schema: SchemaRef, write_schema: DType, bytes_written: Arc, + strategy: Arc, handle: Task>, sender: mpsc::Sender>, ) -> Self { @@ -163,6 +162,7 @@ impl NativeWriter { arrow_schema, write_schema, bytes_written, + strategy, sender, } } @@ -204,6 +204,10 @@ impl NativeWriter { self.bytes_written.load(Ordering::Relaxed) } + fn buffered_bytes(&self) -> u64 { + self.strategy.buffered_bytes() + } + fn close(mut self) -> VortexResult { self.sender.disconnect(); let handle = self @@ -413,7 +417,8 @@ pub extern "system" fn Java_dev_vortex_jni_NativeWriter_create( let resolved = resolve_store(&file_path, &properties)?; 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 strategy = write_strategy_for_schema(&write_schema); + let write_options = session.write_options().with_strategy(strategy.clone()); let (bytes_written, handle) = match resolved { ResolvedStore::Path(path) => { @@ -451,6 +456,7 @@ pub extern "system" fn Java_dev_vortex_jni_NativeWriter_create( arrow_schema, write_schema, bytes_written, + strategy, handle, tx, )) @@ -491,7 +497,8 @@ pub extern "system" fn Java_dev_vortex_jni_NativeWriter_createStream( 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 strategy = write_strategy_for_schema(&write_schema); + let write_options = session.write_options().with_strategy(strategy.clone()); let mut write = CountingVortexWrite::new(JavaWrite::new(vm, writable)); let bytes_written = write.counter(); @@ -506,6 +513,7 @@ pub extern "system" fn Java_dev_vortex_jni_NativeWriter_createStream( arrow_schema, write_schema, bytes_written, + strategy, handle, tx, )) @@ -558,6 +566,22 @@ pub extern "system" fn Java_dev_vortex_jni_NativeWriter_bytesWritten( }) } +#[unsafe(no_mangle)] +pub extern "system" fn Java_dev_vortex_jni_NativeWriter_bufferedBytes( + mut env: EnvUnowned, + _class: JClass, + writer_ptr: jlong, +) -> jlong { + if writer_ptr <= 0 { + return -1; + } + + try_or_throw(&mut env, |_env| { + let writer = unsafe { NativeWriter::from_ptr(writer_ptr) }; + Ok(checked_jlong(writer.buffered_bytes(), "buffered bytes")?) + }) +} + #[unsafe(no_mangle)] pub extern "system" fn Java_dev_vortex_jni_NativeWriter_finish( mut env: EnvUnowned, From f71d0905d28a01f1ec2dc13f943ddb6b553f9e99 Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Thu, 23 Jul 2026 11:53:51 +0100 Subject: [PATCH 2/4] more Signed-off-by: Robert Kruszewski --- vortex-jni/src/writer.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/vortex-jni/src/writer.rs b/vortex-jni/src/writer.rs index 34b96704b49..459aa751414 100644 --- a/vortex-jni/src/writer.rs +++ b/vortex-jni/src/writer.rs @@ -418,7 +418,7 @@ pub extern "system" fn Java_dev_vortex_jni_NativeWriter_create( let (tx, rx) = mpsc::channel(WRITE_CHANNEL_CAPACITY); let stream = ArrayStreamAdapter::new(write_schema.clone(), rx); let strategy = write_strategy_for_schema(&write_schema); - let write_options = session.write_options().with_strategy(strategy.clone()); + let write_options = session.write_options().with_strategy(Arc::clone(&strategy)); let (bytes_written, handle) = match resolved { ResolvedStore::Path(path) => { @@ -498,7 +498,7 @@ pub extern "system" fn Java_dev_vortex_jni_NativeWriter_createStream( let (tx, rx) = mpsc::channel(WRITE_CHANNEL_CAPACITY); let stream = ArrayStreamAdapter::new(write_schema.clone(), rx); let strategy = write_strategy_for_schema(&write_schema); - let write_options = session.write_options().with_strategy(strategy.clone()); + let write_options = session.write_options().with_strategy(Arc::clone(&strategy)); let mut write = CountingVortexWrite::new(JavaWrite::new(vm, writable)); let bytes_written = write.counter(); From 56b4bdf59d6f8d785013cb1766aa8c1fc0e5e24d Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Thu, 23 Jul 2026 11:54:00 +0100 Subject: [PATCH 3/4] format Signed-off-by: Robert Kruszewski --- vortex-jni/src/writer.rs | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/vortex-jni/src/writer.rs b/vortex-jni/src/writer.rs index 459aa751414..ebfadc4727a 100644 --- a/vortex-jni/src/writer.rs +++ b/vortex-jni/src/writer.rs @@ -47,10 +47,12 @@ use vortex::error::VortexResult; use vortex::error::vortex_err; use vortex::expr::stats::Stat; use vortex::expr::stats::StatsProvider; +use vortex::file::ALLOWED_ENCODINGS; +use vortex::file::CountingVortexWrite; +use vortex::file::WriteOptionsSessionExt; +use vortex::file::WriteStrategyBuilder; use vortex::file::WriteSummary; use vortex::file::multi::parse_uri_or_path; -use vortex::file::{ALLOWED_ENCODINGS, CountingVortexWrite}; -use vortex::file::{WriteOptionsSessionExt, WriteStrategyBuilder}; use vortex::io::VortexWrite; use vortex::io::compat::Compat; use vortex::io::object_store::ObjectStoreWrite; From af5a1e5e76277c302d1d634be50b6435ab628056 Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Thu, 23 Jul 2026 11:57:16 +0100 Subject: [PATCH 4/4] format Signed-off-by: Robert Kruszewski --- .../vortex-jni/src/main/java/dev/vortex/api/VortexWriter.java | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) 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 8ebe8405df4..e59a7972587 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 @@ -129,8 +129,8 @@ public synchronized long bytesWritten() { * Return the number of uncompressed bytes accepted by the writer but not yet written to the sink. * *

Together with {@link #bytesWritten()}, this lets callers estimate the in-progress file size: bytes that - * reached the sink are already compressed, while buffered bytes are still uncompressed and will shrink by - * roughly the file's observed compression ratio once flushed. After {@link #finish()}, this is zero. + * reached the sink are already compressed, while buffered bytes are still uncompressed and will shrink by roughly + * the file's observed compression ratio once flushed. After {@link #finish()}, this is zero. */ public synchronized long bufferedBytes() { if (summary != null) {