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..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 @@ -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..ebfadc4727a 100644 --- a/vortex-jni/src/writer.rs +++ b/vortex-jni/src/writer.rs @@ -47,6 +47,7 @@ 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; @@ -58,6 +59,7 @@ 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 +101,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 +144,7 @@ pub struct NativeWriter { arrow_schema: SchemaRef, write_schema: DType, bytes_written: Arc, + strategy: Arc, sender: mpsc::Sender>, } @@ -154,6 +154,7 @@ impl NativeWriter { arrow_schema: SchemaRef, write_schema: DType, bytes_written: Arc, + strategy: Arc, handle: Task>, sender: mpsc::Sender>, ) -> Self { @@ -163,6 +164,7 @@ impl NativeWriter { arrow_schema, write_schema, bytes_written, + strategy, sender, } } @@ -204,6 +206,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 +419,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(Arc::clone(&strategy)); let (bytes_written, handle) = match resolved { ResolvedStore::Path(path) => { @@ -451,6 +458,7 @@ pub extern "system" fn Java_dev_vortex_jni_NativeWriter_create( arrow_schema, write_schema, bytes_written, + strategy, handle, tx, )) @@ -491,7 +499,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(Arc::clone(&strategy)); let mut write = CountingVortexWrite::new(JavaWrite::new(vm, writable)); let bytes_written = write.counter(); @@ -506,6 +515,7 @@ pub extern "system" fn Java_dev_vortex_jni_NativeWriter_createStream( arrow_schema, write_schema, bytes_written, + strategy, handle, tx, )) @@ -558,6 +568,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,