From 63da8ba05a7364e125537f0a820ae2dd12089163 Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Sun, 12 Jul 2026 19:07:42 +0100 Subject: [PATCH 1/7] Support reporting statistics from write api Expose currently written byte size to writers and through vortex-jni Signed-off-by: Robert Kruszewski --- .../vortex/api/VortexColumnStatistics.java | 82 ++++++ .../dev/vortex/api/VortexWriteSummary.java | 40 +++ .../java/dev/vortex/api/VortexWriter.java | 39 ++- .../java/dev/vortex/jni/NativeWriter.java | 7 + .../java/dev/vortex/jni/JNIWriterTest.java | 19 ++ vortex-file/src/counting.rs | 14 +- vortex-file/src/lib.rs | 1 + vortex-file/src/writer.rs | 58 ++++ vortex-jni/src/writer.rs | 276 +++++++++++++++++- 9 files changed, 523 insertions(+), 13 deletions(-) create mode 100644 java/vortex-jni/src/main/java/dev/vortex/api/VortexColumnStatistics.java create mode 100644 java/vortex-jni/src/main/java/dev/vortex/api/VortexWriteSummary.java diff --git a/java/vortex-jni/src/main/java/dev/vortex/api/VortexColumnStatistics.java b/java/vortex-jni/src/main/java/dev/vortex/api/VortexColumnStatistics.java new file mode 100644 index 00000000000..308ca64a0df --- /dev/null +++ b/java/vortex-jni/src/main/java/dev/vortex/api/VortexColumnStatistics.java @@ -0,0 +1,82 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +package dev.vortex.api; + +import java.util.Optional; +import java.util.OptionalLong; + +/** Statistics and physical size reported for one top-level column after a Vortex write. */ +public final class VortexColumnStatistics { + private final int columnIndex; + private final long compressedSize; + private final long valueCount; + private final long nullValueCount; + private final long nanValueCount; + private final Object lowerBound; + private final Object upperBound; + + /** + * Construct native write statistics. + * + *

This constructor is public so the JNI implementation can instantiate the value without reflective access. + * Applications should obtain instances from {@link VortexWriteSummary#columnStatistics()}. + */ + public VortexColumnStatistics( + int columnIndex, + long compressedSize, + long valueCount, + long nullValueCount, + long nanValueCount, + Object lowerBound, + Object upperBound) { + this.columnIndex = columnIndex; + this.compressedSize = compressedSize; + this.valueCount = valueCount; + this.nullValueCount = nullValueCount; + this.nanValueCount = nanValueCount; + this.lowerBound = lowerBound; + this.upperBound = upperBound; + } + + /** Zero-based position of this column in the writer's Arrow schema. */ + public int columnIndex() { + return columnIndex; + } + + /** Compressed bytes referenced by this column's physical layout. */ + public long compressedSize() { + return compressedSize; + } + + /** Number of values in this top-level column. */ + public long valueCount() { + return valueCount; + } + + /** Exact null count, or empty when Vortex did not compute one for this data type. */ + public OptionalLong nullValueCount() { + return nullValueCount >= 0 ? OptionalLong.of(nullValueCount) : OptionalLong.empty(); + } + + /** Exact NaN count, or empty for non-floating-point columns. */ + public OptionalLong nanValueCount() { + return nanValueCount >= 0 ? OptionalLong.of(nanValueCount) : OptionalLong.empty(); + } + + /** + * Lower bound represented using the corresponding Arrow scalar's Java type. + * + *

Integers use {@link Integer}, {@link Long}, or {@link java.math.BigInteger}; floating-point values use + * {@link Float} or {@link Double}; decimal values use {@link java.math.BigDecimal}; strings use {@link String}; and + * binary values use {@code byte[]}. + */ + public Optional lowerBound() { + return Optional.ofNullable(lowerBound); + } + + /** Upper bound represented using the corresponding Arrow scalar's Java type. */ + public Optional upperBound() { + return Optional.ofNullable(upperBound); + } +} diff --git a/java/vortex-jni/src/main/java/dev/vortex/api/VortexWriteSummary.java b/java/vortex-jni/src/main/java/dev/vortex/api/VortexWriteSummary.java new file mode 100644 index 00000000000..69f042d2ff4 --- /dev/null +++ b/java/vortex-jni/src/main/java/dev/vortex/api/VortexWriteSummary.java @@ -0,0 +1,40 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +package dev.vortex.api; + +import java.util.List; + +/** Immutable metadata returned after a Vortex writer has finalized its file. */ +public final class VortexWriteSummary { + private final long fileSize; + private final long rowCount; + private final List columnStatistics; + + /** + * Construct a native write summary. + * + *

This constructor is public so the JNI implementation can instantiate the value without reflective access. + * Applications should obtain summaries from {@link VortexWriter#finish()}. + */ + public VortexWriteSummary(long fileSize, long rowCount, VortexColumnStatistics[] columnStatistics) { + this.fileSize = fileSize; + this.rowCount = rowCount; + this.columnStatistics = List.of(columnStatistics.clone()); + } + + /** Exact size of the completed file in bytes. */ + public long fileSize() { + return fileSize; + } + + /** Exact number of rows written to the file. */ + public long rowCount() { + return rowCount; + } + + /** Per-column statistics in Arrow schema order. */ + public List columnStatistics() { + return columnStatistics; + } +} 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 7761be59cdc..e883e7911f7 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 @@ -29,6 +29,7 @@ public final class VortexWriter implements AutoCloseable { private final long pointer; private final AtomicBoolean closed = new AtomicBoolean(false); + private volatile VortexWriteSummary summary; private VortexWriter(long pointer) { Preconditions.checkArgument(pointer != 0, "invalid writer pointer"); @@ -79,15 +80,45 @@ public void writeBatch(long arrowArrayAddr, long arrowSchemaAddr) throws IOExcep } } - /** Flush any pending batches and finalize the file. Idempotent. */ - @Override - public void close() throws IOException { + /** + * Return the number of bytes successfully written to the underlying sink so far. + * + *

This count does not include queued batches or data still buffered by layout strategies, so it may lag the + * amount of input accepted by {@link #writeBatch(long, long)}. After {@link #finish()}, it is the exact completed + * file size and is equal to {@link VortexWriteSummary#fileSize()}. + */ + public synchronized long bytesWritten() { + if (summary != null) { + return summary.fileSize(); + } + Preconditions.checkState(!closed.get(), "writer closed without a write summary"); + long bytesWritten = NativeWriter.bytesWritten(pointer); + Preconditions.checkState(bytesWritten >= 0, "native writer returned an invalid byte count"); + return bytesWritten; + } + + /** + * Flush pending batches, finalize the file, and return its statistics and physical sizes. + * + *

This method is idempotent. Later calls return the same immutable summary. + */ + public synchronized VortexWriteSummary finish() throws IOException { if (closed.compareAndSet(false, true)) { try { - NativeWriter.close(pointer); + summary = NativeWriter.finish(pointer); } catch (RuntimeException e) { throw new IOException("failed to close writer", e); } } + if (summary == null) { + throw new IOException("writer was closed without retaining its write summary"); + } + return summary; + } + + /** Flush any pending batches and finalize the file. Idempotent. */ + @Override + public void close() throws IOException { + finish(); } } 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 d1f2e725f8c..7536712665a 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 @@ -3,6 +3,7 @@ package dev.vortex.jni; +import dev.vortex.api.VortexWriteSummary; import java.util.Map; /** JNI boundary for {@link dev.vortex.api.VortexWriter}. */ @@ -27,6 +28,12 @@ public static native long create( */ public static native boolean writeBatch(long writerPointer, long arrowArrayAddress, long arrowSchemaAddress); + /** Number of bytes successfully written to the underlying sink so far. */ + public static native long bytesWritten(long writerPointer); + /** Flush and close the writer. Must be called exactly once. */ public static native void close(long writerPointer); + + /** Flush and close the writer, returning its completed-file summary. Must be called exactly once. */ + public static native VortexWriteSummary finish(long writerPointer); } diff --git a/java/vortex-jni/src/test/java/dev/vortex/jni/JNIWriterTest.java b/java/vortex-jni/src/test/java/dev/vortex/jni/JNIWriterTest.java index 03869ae8598..ed84111bdf8 100644 --- a/java/vortex-jni/src/test/java/dev/vortex/jni/JNIWriterTest.java +++ b/java/vortex-jni/src/test/java/dev/vortex/jni/JNIWriterTest.java @@ -14,6 +14,7 @@ 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 java.io.IOException; @@ -157,6 +158,8 @@ public void testWriteBatch() throws IOException { Schema schema = personSchema(); Session session = Session.create(); + VortexWriteSummary summary; + long bytesWhileOpen; try (VortexWriter writer = VortexWriter.create(session, writePath, schema, new HashMap<>(), allocator); VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator)) { VarCharVector nameVec = (VarCharVector) root.getVector("name"); @@ -179,9 +182,25 @@ public void testWriteBatch() throws IOException { Data.exportVectorSchemaRoot(allocator, root, null, arrowArray, arrowSchemaFfi); writer.writeBatch(arrowArray.memoryAddress(), arrowSchemaFfi.memoryAddress()); } + bytesWhileOpen = writer.bytesWritten(); + summary = writer.finish(); + assertEquals(summary.fileSize(), writer.bytesWritten()); } assertTrue(Files.exists(outputPath), "output file should exist"); + assertTrue(bytesWhileOpen >= 0); + assertTrue(bytesWhileOpen <= summary.fileSize()); + assertEquals(Files.size(outputPath), summary.fileSize()); + assertEquals(3L, summary.rowCount()); + assertEquals(2, summary.columnStatistics().size()); + assertEquals(0, summary.columnStatistics().get(0).columnIndex()); + assertTrue(summary.columnStatistics().get(0).compressedSize() > 0); + assertEquals(3L, summary.columnStatistics().get(0).valueCount()); + assertEquals(0L, summary.columnStatistics().get(0).nullValueCount().orElseThrow()); + assertEquals("Alice", summary.columnStatistics().get(0).lowerBound().orElseThrow()); + assertEquals("Carol", summary.columnStatistics().get(0).upperBound().orElseThrow()); + assertEquals(25, summary.columnStatistics().get(1).lowerBound().orElseThrow()); + assertEquals(40, summary.columnStatistics().get(1).upperBound().orElseThrow()); DataSource ds = DataSource.open(session, writePath); assertEquals(new DataSource.RowCount.Exact(3L), ds.rowCount()); diff --git a/vortex-file/src/counting.rs b/vortex-file/src/counting.rs index 78f8a8741a9..bb5c029da68 100644 --- a/vortex-file/src/counting.rs +++ b/vortex-file/src/counting.rs @@ -9,13 +9,14 @@ use std::sync::atomic::Ordering; use vortex_io::IoBuf; use vortex_io::VortexWrite; -/// A wrapper around an `VortexWrite` that counts the number of bytes written. -pub(crate) struct CountingVortexWrite { +/// A wrapper around a [`VortexWrite`] that counts the number of bytes written. +pub struct CountingVortexWrite { inner: W, bytes_written: Arc, } impl CountingVortexWrite { + /// Wrap a writer with a new byte counter. pub fn new(inner: W) -> Self { Self { inner, @@ -23,6 +24,15 @@ impl CountingVortexWrite { } } + /// Wrap a writer with an existing shared byte counter. + pub fn with_counter(inner: W, bytes_written: Arc) -> Self { + Self { + inner, + bytes_written, + } + } + + /// Returns the shared byte counter updated by this writer. pub fn counter(&self) -> Arc { Arc::clone(&self.bytes_written) } diff --git a/vortex-file/src/lib.rs b/vortex-file/src/lib.rs index e9c7741af93..261ba7191d6 100644 --- a/vortex-file/src/lib.rs +++ b/vortex-file/src/lib.rs @@ -113,6 +113,7 @@ mod tests; pub mod v2; mod writer; +pub use counting::CountingVortexWrite; pub use file::*; pub use footer::*; pub use forever_constant::*; diff --git a/vortex-file/src/writer.rs b/vortex-file/src/writer.rs index 2afd61f7b23..d6fda0dad22 100644 --- a/vortex-file/src/writer.rs +++ b/vortex-file/src/writer.rs @@ -39,14 +39,18 @@ use vortex_io::VortexWrite; use vortex_io::kanal_ext::KanalExt; use vortex_io::runtime::BlockingRuntime; use vortex_io::session::RuntimeSessionExt; +use vortex_layout::LayoutChildType; +use vortex_layout::LayoutRef; use vortex_layout::LayoutStrategy; use vortex_layout::layouts::file_stats::accumulate_stats; +use vortex_layout::segments::SegmentId; use vortex_layout::sequence::SequenceId; use vortex_layout::sequence::SequentialStreamAdapter; use vortex_layout::sequence::SequentialStreamExt; use vortex_session::SessionExt; use vortex_session::VortexSession; use vortex_session::registry::ReadContext; +use vortex_utils::aliases::hash_set::HashSet; use crate::ALLOWED_ENCODINGS; use crate::Footer; @@ -505,4 +509,58 @@ impl WriteSummary { pub fn row_count(&self) -> u64 { self.footer.row_count() } + + /// Returns the compressed size in bytes of each top-level column in schema order. + /// + /// A column's size includes every physical segment referenced by its layout subtree, + /// including auxiliary segments such as zone maps and dictionaries. A segment is counted at + /// most once per column even if more than one layout node references it. File-level metadata, + /// alignment padding, and the footer are not attributed to any column. + /// + /// For a non-struct file, the returned vector contains a single entry for the root column. + pub fn compressed_column_sizes(&self) -> VortexResult> { + let root = self.footer.layout(); + if !root.dtype().is_struct() { + return Ok(vec![compressed_layout_size( + root, + self.footer.segment_map(), + )?]); + } + + let mut sizes = Vec::with_capacity( + root.dtype() + .as_struct_fields_opt() + .map_or(1, |f| f.nfields()), + ); + for idx in 0..root.nchildren() { + if matches!(root.child_type(idx), LayoutChildType::Field(_)) { + sizes.push(compressed_layout_size( + &root.child(idx)?, + self.footer.segment_map(), + )?); + } + } + Ok(sizes) + } +} + +fn compressed_layout_size( + layout: &LayoutRef, + segments: &[crate::footer::SegmentSpec], +) -> VortexResult { + let mut segment_ids = HashSet::::default(); + for node in layout.depth_first_traversal() { + segment_ids.extend(node?.segment_ids()); + } + + segment_ids.into_iter().try_fold(0u64, |size, segment_id| { + let segment = segments.get(*segment_id as usize).ok_or_else(|| { + vortex_err!( + "layout references missing segment {} (segment count: {})", + *segment_id, + segments.len() + ) + })?; + Ok(size + u64::from(segment.length)) + }) } diff --git a/vortex-jni/src/writer.rs b/vortex-jni/src/writer.rs index 5b098b4c320..b44d4b53d4d 100644 --- a/vortex-jni/src/writer.rs +++ b/vortex-jni/src/writer.rs @@ -8,6 +8,8 @@ use std::path::PathBuf; use std::sync::Arc; +use std::sync::atomic::AtomicU64; +use std::sync::atomic::Ordering; use arrow_array::RecordBatch; use arrow_array::StructArray; @@ -17,24 +19,33 @@ use arrow_schema::SchemaRef; use async_fs::File; use futures::SinkExt; use futures::channel::mpsc; +use jni::Env; use jni::EnvUnowned; use jni::objects::JClass; use jni::objects::JObject; use jni::objects::JString; +use jni::objects::JValue; use jni::sys::JNI_FALSE; use jni::sys::JNI_TRUE; use jni::sys::jboolean; use jni::sys::jlong; +use jni::sys::jobject; use object_store::ObjectStore; use object_store::path::Path as ObjectStorePath; use vortex::array::ArrayRef; use vortex::array::VTable; +use vortex::array::scalar::PValue; +use vortex::array::scalar::Scalar; +use vortex::array::scalar::ScalarValue; +use vortex::array::stats::StatsSet; use vortex::array::stream::ArrayStreamAdapter; use vortex::dtype::DType; use vortex::dtype::Field as DTypeField; use vortex::dtype::FieldPath; use vortex::error::VortexResult; use vortex::error::vortex_err; +use vortex::expr::stats::Stat; +use vortex::file::CountingVortexWrite; use vortex::file::WriteOptionsSessionExt; use vortex::file::WriteStrategyBuilder; use vortex::file::WriteSummary; @@ -130,6 +141,7 @@ pub struct NativeWriter { session: VortexSession, arrow_schema: SchemaRef, write_schema: DType, + bytes_written: Arc, sender: mpsc::Sender>, } @@ -138,6 +150,7 @@ impl NativeWriter { session: VortexSession, arrow_schema: SchemaRef, write_schema: DType, + bytes_written: Arc, handle: Task>, sender: mpsc::Sender>, ) -> Self { @@ -146,6 +159,7 @@ impl NativeWriter { session, arrow_schema, write_schema, + bytes_written, sender, } } @@ -183,19 +197,224 @@ impl NativeWriter { .map_err(|e| vortex_err!("failed to send batch: {e}")) } - fn close(mut self) -> VortexResult<()> { + fn bytes_written(&self) -> u64 { + self.bytes_written.load(Ordering::Relaxed) + } + + fn close(mut self) -> VortexResult { self.sender.disconnect(); let handle = self .handle .take() .ok_or_else(|| vortex_err!("writer already closed"))?; - RUNTIME.block_on(async { - handle.await?; - VortexResult::Ok(()) - }) + RUNTIME.block_on(handle) + } +} + +fn exact_stat_scalar(stats: &StatsSet, dtype: &DType, stat: Stat) -> Option { + let value = stats.get(stat).as_exact()?; + Scalar::try_new(stat.dtype(dtype)?, Some(value)).ok() +} + +fn exact_count(stats: &StatsSet, dtype: &DType, stat: Stat) -> Option { + u64::try_from(&exact_stat_scalar(stats, dtype, stat)?).ok() +} + +fn checked_jlong(value: u64, name: &str) -> VortexResult { + jlong::try_from(value).map_err(|_| vortex_err!("{name} exceeds Java long range: {value}")) +} + +fn exact_count_jlong( + stats: Option<&StatsSet>, + dtype: Option<&DType>, + stat: Stat, +) -> VortexResult { + stats + .zip(dtype) + .and_then(|(stats, dtype)| exact_count(stats, dtype, stat)) + .map(|value| checked_jlong(value, stat.name())) + .transpose() + .map(|value| value.unwrap_or(-1)) +} + +fn boxed_integer<'local>(env: &mut Env<'local>, value: i32) -> Result, JNIError> { + Ok(env.new_object( + jni::jni_str!("java/lang/Integer"), + jni::jni_sig!("(I)V"), + &[JValue::Int(value)], + )?) +} + +fn boxed_long<'local>(env: &mut Env<'local>, value: i64) -> Result, JNIError> { + Ok(env.new_object( + jni::jni_str!("java/lang/Long"), + jni::jni_sig!("(J)V"), + &[JValue::Long(value)], + )?) +} + +fn boxed_float<'local>(env: &mut Env<'local>, value: f32) -> Result, JNIError> { + Ok(env.new_object( + jni::jni_str!("java/lang/Float"), + jni::jni_sig!("(F)V"), + &[JValue::Float(value)], + )?) +} + +fn boxed_double<'local>(env: &mut Env<'local>, value: f64) -> Result, JNIError> { + Ok(env.new_object( + jni::jni_str!("java/lang/Double"), + jni::jni_sig!("(D)V"), + &[JValue::Double(value)], + )?) +} + +fn boxed_boolean<'local>(env: &mut Env<'local>, value: bool) -> Result, JNIError> { + Ok(env.new_object( + jni::jni_str!("java/lang/Boolean"), + jni::jni_sig!("(Z)V"), + &[JValue::Bool(if value { JNI_TRUE } else { JNI_FALSE })], + )?) +} + +fn big_integer<'local>( + env: &mut Env<'local>, + value: impl ToString, +) -> Result, JNIError> { + let string = env.new_string(value.to_string())?; + Ok(env.new_object( + jni::jni_str!("java/math/BigInteger"), + jni::jni_sig!("(Ljava/lang/String;)V"), + &[JValue::Object(string.as_ref())], + )?) +} + +fn scalar_to_java<'local>( + env: &mut Env<'local>, + scalar: &Scalar, +) -> Result, JNIError> { + if scalar.is_null() { + return Ok(JObject::null()); + } + if scalar.dtype().is_extension() { + return scalar_to_java(env, &scalar.as_extension().to_storage_scalar()); + } + + let Some(value) = scalar.value() else { + return Ok(JObject::null()); + }; + match value { + ScalarValue::Bool(value) => boxed_boolean(env, *value), + ScalarValue::Primitive(value) => match value { + PValue::U8(value) => boxed_integer(env, i32::from(*value)), + PValue::U16(value) => boxed_integer(env, i32::from(*value)), + PValue::U32(value) => boxed_long(env, i64::from(*value)), + PValue::U64(value) => big_integer(env, value), + PValue::I8(value) => boxed_integer(env, i32::from(*value)), + PValue::I16(value) => boxed_integer(env, i32::from(*value)), + PValue::I32(value) => boxed_integer(env, *value), + PValue::I64(value) => boxed_long(env, *value), + PValue::F16(value) => boxed_float(env, f32::from(*value)), + PValue::F32(value) => boxed_float(env, *value), + PValue::F64(value) => boxed_double(env, *value), + }, + ScalarValue::Decimal(value) => { + let DType::Decimal(decimal_dtype, _) = scalar.dtype() else { + return Err(JNIError::Vortex(vortex_err!( + "decimal statistic has non-decimal dtype {}", + scalar.dtype() + ))); + }; + let unscaled = big_integer(env, value.as_i256())?; + Ok(env.new_object( + jni::jni_str!("java/math/BigDecimal"), + jni::jni_sig!("(Ljava/math/BigInteger;I)V"), + &[ + JValue::Object(&unscaled), + JValue::Int(i32::from(decimal_dtype.scale())), + ], + )?) + } + ScalarValue::Utf8(value) => Ok(env.new_string(value.as_str())?.into()), + ScalarValue::Binary(value) => Ok(env.byte_array_from_slice(value.as_slice())?.into()), + ScalarValue::Tuple(_) | ScalarValue::Variant(_) => Err(JNIError::Vortex(vortex_err!( + "cannot return nested scalar write statistic with dtype {} to Java", + scalar.dtype() + ))), } } +fn write_summary_to_java<'local>( + env: &mut Env<'local>, + summary: &WriteSummary, +) -> Result, JNIError> { + let column_sizes = summary.compressed_column_sizes()?; + let file_stats = summary.footer().statistics(); + let columns = env.new_object_array( + i32::try_from(column_sizes.len()) + .map_err(|_| vortex_err!("column count exceeds Java array range"))?, + jni::jni_str!("dev/vortex/api/VortexColumnStatistics"), + JObject::null(), + )?; + + for (column_index, compressed_size) in column_sizes.into_iter().enumerate() { + let (stats, dtype) = file_stats + .and_then(|all_stats| { + all_stats + .stats_sets() + .get(column_index) + .zip(all_stats.dtypes().get(column_index)) + }) + .map_or((None, None), |(stats, dtype)| (Some(stats), Some(dtype))); + let null_count = exact_count_jlong(stats, dtype, Stat::NullCount)?; + let nan_count = exact_count_jlong(stats, dtype, Stat::NaNCount)?; + let lower_bound = stats + .zip(dtype) + .and_then(|(stats, dtype)| exact_stat_scalar(stats, dtype, Stat::Min)); + let upper_bound = stats + .zip(dtype) + .and_then(|(stats, dtype)| exact_stat_scalar(stats, dtype, Stat::Max)); + let column = env.with_local_frame_returning_local::<_, JObject, JNIError>(16, |env| { + let lower_bound = match &lower_bound { + Some(value) => scalar_to_java(env, value)?, + None => JObject::null(), + }; + let upper_bound = match &upper_bound { + Some(value) => scalar_to_java(env, value)?, + None => JObject::null(), + }; + Ok(env.new_object( + jni::jni_str!("dev/vortex/api/VortexColumnStatistics"), + jni::jni_sig!("(IJJJJLjava/lang/Object;Ljava/lang/Object;)V"), + &[ + JValue::Int( + i32::try_from(column_index) + .map_err(|_| vortex_err!("column index exceeds Java int range"))?, + ), + JValue::Long(checked_jlong(compressed_size, "compressed column size")?), + JValue::Long(checked_jlong(summary.row_count(), "row count")?), + JValue::Long(null_count), + JValue::Long(nan_count), + JValue::Object(&lower_bound), + JValue::Object(&upper_bound), + ], + )?) + })?; + columns.set_element(env, column_index, &column)?; + env.delete_local_ref(column); + } + + Ok(env.new_object( + jni::jni_str!("dev/vortex/api/VortexWriteSummary"), + jni::jni_sig!("(JJ[Ldev/vortex/api/VortexColumnStatistics;)V"), + &[ + JValue::Long(checked_jlong(summary.size(), "file size")?), + JValue::Long(checked_jlong(summary.row_count(), "row count")?), + JValue::Object(columns.as_ref()), + ], + )?) +} + #[unsafe(no_mangle)] pub extern "system" fn Java_dev_vortex_jni_NativeWriter_create( mut env: EnvUnowned, @@ -223,6 +442,8 @@ 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 write_options = write_options_for_schema(session, &write_schema); + let bytes_written = Arc::new(AtomicU64::new(0)); + let task_bytes_written = Arc::clone(&bytes_written); let handle = session.handle().spawn(async move { match resolved { @@ -231,14 +452,22 @@ pub extern "system" fn Java_dev_vortex_jni_NativeWriter_create( async_fs::create_dir_all(parent).await?; } let mut file = File::create(path).await?; - let summary = write_options.write(&mut file, stream).await?; + let write = CountingVortexWrite::with_counter( + &mut file, + Arc::clone(&task_bytes_written), + ); + let summary = write_options.write(write, stream).await?; file.shutdown().await?; Ok(summary) } ResolvedStore::ObjectStore(store, path) => { let mut write = ObjectStoreWrite::new(Arc::new(Compat::new(store)), &path).await?; - let summary = write_options.write(&mut write, stream).await?; + let counted_write = CountingVortexWrite::with_counter( + &mut write, + Arc::clone(&task_bytes_written), + ); + let summary = write_options.write(counted_write, stream).await?; write.shutdown().await?; Ok(summary) } @@ -249,6 +478,7 @@ pub extern "system" fn Java_dev_vortex_jni_NativeWriter_create( session.clone(), arrow_schema, write_schema, + bytes_written, handle, tx, )) @@ -285,6 +515,38 @@ pub extern "system" fn Java_dev_vortex_jni_NativeWriter_writeBatch( }) } +#[unsafe(no_mangle)] +pub extern "system" fn Java_dev_vortex_jni_NativeWriter_bytesWritten( + 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.bytes_written(), "bytes written")?) + }) +} + +#[unsafe(no_mangle)] +pub extern "system" fn Java_dev_vortex_jni_NativeWriter_finish( + mut env: EnvUnowned, + _class: JClass, + writer_ptr: jlong, +) -> jobject { + if writer_ptr <= 0 { + return JObject::null().into_raw(); + } + let writer = unsafe { NativeWriter::from_raw(writer_ptr) }; + try_or_throw(&mut env, |env| { + let summary = writer.close()?; + Ok(write_summary_to_java(env, &summary)?.into_raw()) + }) +} + #[unsafe(no_mangle)] pub extern "system" fn Java_dev_vortex_jni_NativeWriter_close( mut env: EnvUnowned, From a7a04b9d8306bdec491e9202bbded4903045a632 Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Mon, 13 Jul 2026 20:16:29 +0100 Subject: [PATCH 2/7] Compute compressed sizes for arbitrarily nested fields Replace the top-level-only column size computation with a single walk of the layout tree that attributes each segment to the deepest field path referencing it, exposed as Footer::compressed_field_sizes. WriteSummary::compressed_column_sizes is now a thin schema-order wrapper over it. Signed-off-by: "Robert Kruszewski" Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01JE2aKEzyogTh1dP126gudc --- vortex-file/src/footer/field_sizes.rs | 243 ++++++++++++++++++++++++++ vortex-file/src/footer/mod.rs | 11 ++ vortex-file/src/writer.rs | 70 ++------ 3 files changed, 273 insertions(+), 51 deletions(-) create mode 100644 vortex-file/src/footer/field_sizes.rs diff --git a/vortex-file/src/footer/field_sizes.rs b/vortex-file/src/footer/field_sizes.rs new file mode 100644 index 00000000000..5d7d2ce7c02 --- /dev/null +++ b/vortex-file/src/footer/field_sizes.rs @@ -0,0 +1,243 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_array::dtype::Field; +use vortex_array::dtype::FieldPath; +use vortex_error::VortexResult; +use vortex_error::vortex_err; +use vortex_layout::LayoutChildType; +use vortex_layout::LayoutRef; +use vortex_layout::segments::SegmentId; +use vortex_utils::aliases::hash_map::Entry; +use vortex_utils::aliases::hash_map::HashMap; + +use crate::footer::SegmentSpec; + +/// Compressed sizes of a file's fields, keyed by field path. +/// +/// Sizes are derived from the file's layout tree and segment map: walking the tree from the root, +/// each physical segment is attributed to the deepest [`FieldPath`] whose layout subtree +/// references it. A segment referenced by more than one field is attributed to their longest +/// common ancestor, so a segment is counted exactly once and the size of a field always equals +/// the sum of its child field sizes plus the bytes attributed directly to it (e.g. struct +/// validity, or auxiliary segments such as zone maps and dictionaries written above the field +/// split). +/// +/// Fields nested inside non-struct types (e.g. structs within lists) are not split into separate +/// layout subtrees by the writer, so their bytes are attributed to the enclosing field. +/// +/// File-level metadata, alignment padding, and the footer are not part of any segment and are +/// never attributed. +#[derive(Debug, Clone)] +pub struct CompressedFieldSizes { + sizes: HashMap, +} + +impl CompressedFieldSizes { + pub(crate) fn try_new(root: &LayoutRef, segments: &[SegmentSpec]) -> VortexResult { + // Every segment is attributed to the deepest field path referencing it, collapsing to the + // longest common ancestor if multiple fields reference the same segment. + let mut attribution = HashMap::::default(); + // Seed every field path present in the layout tree so empty fields report a zero size. + let mut sizes = HashMap::::default(); + sizes.insert(FieldPath::root(), 0); + + let mut stack = vec![(LayoutRef::clone(root), FieldPath::root())]; + while let Some((layout, path)) = stack.pop() { + for segment_id in layout.segment_ids() { + match attribution.entry(segment_id) { + Entry::Occupied(mut entry) => { + let ancestor = common_prefix(entry.get(), &path); + entry.insert(ancestor); + } + Entry::Vacant(entry) => { + entry.insert(path.clone()); + } + } + } + + for idx in 0..layout.nchildren() { + let child_path = match layout.child_type(idx) { + LayoutChildType::Field(name) => { + let child_path = path.clone().push(name); + sizes.entry(child_path.clone()).or_insert(0); + child_path + } + _ => path.clone(), + }; + stack.push((layout.child(idx)?, child_path)); + } + } + + for (segment_id, path) in attribution { + let segment = segments.get(*segment_id as usize).ok_or_else(|| { + vortex_err!( + "layout references missing segment {} (segment count: {})", + *segment_id, + segments.len() + ) + })?; + // A segment counts towards its attributed field and every enclosing field. + for len in 0..=path.parts().len() { + *sizes + .entry(FieldPath::from(path.parts()[..len].to_vec())) + .or_insert(0) += u64::from(segment.length); + } + } + + Ok(Self { sizes }) + } + + /// Returns the compressed size in bytes of the field at the given path, including all of its + /// descendants, or `None` if the layout tree contains no such field. + /// + /// [`FieldPath::root`] returns the same value as [`total`][Self::total]. + pub fn get(&self, path: &FieldPath) -> Option { + self.sizes.get(path).copied() + } + + /// Returns the total compressed size in bytes of all segments referenced by the layout tree. + pub fn total(&self) -> u64 { + self.get(&FieldPath::root()).unwrap_or_default() + } + + /// Iterates over all field paths in the layout tree and their compressed sizes. + pub fn iter(&self) -> impl Iterator { + self.sizes.iter().map(|(path, size)| (path, *size)) + } +} + +fn common_prefix(left: &FieldPath, right: &FieldPath) -> FieldPath { + left.parts() + .iter() + .zip(right.parts()) + .take_while(|(l, r)| l == r) + .map(|(l, _)| l.clone()) + .collect::>() + .into() +} + +#[cfg(test)] +mod tests { + use std::sync::LazyLock; + + use vortex_array::IntoArray; + use vortex_array::array_session; + use vortex_array::arrays::StructArray; + use vortex_array::dtype::FieldPath; + use vortex_array::field_path; + use vortex_buffer::ByteBufferMut; + use vortex_buffer::buffer; + use vortex_error::VortexResult; + use vortex_io::session::RuntimeSession; + use vortex_layout::session::LayoutSession; + use vortex_session::VortexSession; + + use crate::WriteOptionsSessionExt; + use crate::writer::WriteSummary; + + static SESSION: LazyLock = LazyLock::new(|| { + let session = array_session() + .with::() + .with::(); + crate::register_default_encodings(&session); + session + }); + + async fn write_nested() -> VortexResult { + let inner = StructArray::from_fields(&[ + ("c", buffer![10i64, 20, 30, 40].into_array()), + ( + "d", + StructArray::from_fields(&[("e", buffer![1.0f64, 2.0, 3.0, 4.0].into_array())])? + .into_array(), + ), + ])?; + let root = StructArray::from_fields(&[ + ("a", buffer![1u32, 2, 3, 4].into_array()), + ("b", inner.into_array()), + ])?; + + let mut buf = ByteBufferMut::empty(); + SESSION + .write_options() + .write(&mut buf, root.into_array().to_array_stream()) + .await + } + + #[tokio::test] + async fn nested_field_sizes() -> VortexResult<()> { + let summary = write_nested().await?; + let sizes = summary.footer().compressed_field_sizes()?; + + for path in [ + field_path!(a), + field_path!(b), + field_path!(b.c), + field_path!(b.d), + field_path!(b.d.e), + ] { + assert!( + sizes.get(&path).is_some_and(|size| size > 0), + "expected non-zero size for {path}, got {:?}", + sizes.get(&path) + ); + } + assert_eq!(sizes.get(&field_path!(missing)), None); + assert_eq!(sizes.get(&field_path!(b.d.e.missing)), None); + + // Attribution partitions segments: a field's size is the sum of its children plus its own + // directly-attributed segments (none here, as the structs are non-nullable). + let size = |path: FieldPath| sizes.get(&path).unwrap_or_default(); + assert_eq!( + size(field_path!(b)), + size(field_path!(b.c)) + size(field_path!(b.d)) + ); + assert_eq!(size(field_path!(b.d)), size(field_path!(b.d.e))); + assert_eq!(sizes.total(), size(field_path!(a)) + size(field_path!(b))); + + // Every segment in the file is attributed exactly once. + let segment_total: u64 = summary + .footer() + .segment_map() + .iter() + .map(|segment| u64::from(segment.length)) + .sum(); + assert_eq!(sizes.total(), segment_total); + + Ok(()) + } + + #[tokio::test] + async fn column_sizes_in_schema_order() -> VortexResult<()> { + let summary = write_nested().await?; + let sizes = summary.footer().compressed_field_sizes()?; + + assert_eq!( + summary.compressed_column_sizes()?, + vec![ + sizes.get(&field_path!(a)).unwrap_or_default(), + sizes.get(&field_path!(b)).unwrap_or_default(), + ] + ); + Ok(()) + } + + #[tokio::test] + async fn non_struct_file_reports_single_column() -> VortexResult<()> { + let mut buf = ByteBufferMut::empty(); + let summary = SESSION + .write_options() + .write( + &mut buf, + buffer![1i32, 2, 3, 4].into_array().to_array_stream(), + ) + .await?; + + let sizes = summary.footer().compressed_field_sizes()?; + assert!(sizes.total() > 0); + assert_eq!(sizes.get(&FieldPath::root()), Some(sizes.total())); + assert_eq!(summary.compressed_column_sizes()?, vec![sizes.total()]); + Ok(()) + } +} diff --git a/vortex-file/src/footer/mod.rs b/vortex-file/src/footer/mod.rs index f4bf7e19760..fda02182e1b 100644 --- a/vortex-file/src/footer/mod.rs +++ b/vortex-file/src/footer/mod.rs @@ -8,6 +8,7 @@ //! //! The byte-level footer and postscript layout is part of the file-format spec; this module exposes //! the structured Rust representation and serializer/deserializer state machine. +mod field_sizes; mod file_layout; mod file_statistics; mod postscript; @@ -19,6 +20,7 @@ mod serializer; pub use serializer::*; mod deserializer; pub use deserializer::*; +pub use field_sizes::CompressedFieldSizes; pub use file_statistics::FileStatistics; use flatbuffers::root; use itertools::Itertools; @@ -146,6 +148,15 @@ impl Footer { self.statistics.as_ref() } + /// Computes the compressed size in bytes of every field in the file, keyed by field path. + /// + /// Sizes are derived by attributing each segment in the [segment map][Self::segment_map] to a + /// field in the [layout tree][Self::layout]; see [`CompressedFieldSizes`] for the exact + /// attribution semantics. No IO is performed. + pub fn compressed_field_sizes(&self) -> VortexResult { + CompressedFieldSizes::try_new(&self.root_layout, &self.segments) + } + /// Returns the [`DType`] of the file. pub fn dtype(&self) -> &DType { self.root_layout.dtype() diff --git a/vortex-file/src/writer.rs b/vortex-file/src/writer.rs index d6fda0dad22..80da9f4fb89 100644 --- a/vortex-file/src/writer.rs +++ b/vortex-file/src/writer.rs @@ -19,6 +19,7 @@ use itertools::Itertools; use vortex_array::ArrayContext; use vortex_array::ArrayRef; use vortex_array::dtype::DType; +use vortex_array::dtype::FieldPath; use vortex_array::expr::stats::Stat; use vortex_array::iter::ArrayIterator; use vortex_array::iter::ArrayIteratorExt; @@ -39,18 +40,14 @@ use vortex_io::VortexWrite; use vortex_io::kanal_ext::KanalExt; use vortex_io::runtime::BlockingRuntime; use vortex_io::session::RuntimeSessionExt; -use vortex_layout::LayoutChildType; -use vortex_layout::LayoutRef; use vortex_layout::LayoutStrategy; use vortex_layout::layouts::file_stats::accumulate_stats; -use vortex_layout::segments::SegmentId; use vortex_layout::sequence::SequenceId; use vortex_layout::sequence::SequentialStreamAdapter; use vortex_layout::sequence::SequentialStreamExt; use vortex_session::SessionExt; use vortex_session::VortexSession; use vortex_session::registry::ReadContext; -use vortex_utils::aliases::hash_set::HashSet; use crate::ALLOWED_ENCODINGS; use crate::Footer; @@ -512,55 +509,26 @@ impl WriteSummary { /// Returns the compressed size in bytes of each top-level column in schema order. /// - /// A column's size includes every physical segment referenced by its layout subtree, - /// including auxiliary segments such as zone maps and dictionaries. A segment is counted at - /// most once per column even if more than one layout node references it. File-level metadata, - /// alignment padding, and the footer are not attributed to any column. + /// A column's size includes every physical segment attributed to its layout subtree, + /// including auxiliary segments such as zone maps and dictionaries; see + /// [`Footer::compressed_field_sizes`] for the exact attribution semantics and for sizes of + /// nested fields. Bytes not attributable to a specific column (e.g. top-level struct + /// validity) are not included in any column's size. /// /// For a non-struct file, the returned vector contains a single entry for the root column. pub fn compressed_column_sizes(&self) -> VortexResult> { - let root = self.footer.layout(); - if !root.dtype().is_struct() { - return Ok(vec![compressed_layout_size( - root, - self.footer.segment_map(), - )?]); - } - - let mut sizes = Vec::with_capacity( - root.dtype() - .as_struct_fields_opt() - .map_or(1, |f| f.nfields()), - ); - for idx in 0..root.nchildren() { - if matches!(root.child_type(idx), LayoutChildType::Field(_)) { - sizes.push(compressed_layout_size( - &root.child(idx)?, - self.footer.segment_map(), - )?); - } - } - Ok(sizes) + let sizes = self.footer.compressed_field_sizes()?; + let Some(fields) = self.footer.dtype().as_struct_fields_opt() else { + return Ok(vec![sizes.total()]); + }; + Ok(fields + .names() + .iter() + .map(|name| { + sizes + .get(&FieldPath::from_name(name.clone())) + .unwrap_or_default() + }) + .collect()) } } - -fn compressed_layout_size( - layout: &LayoutRef, - segments: &[crate::footer::SegmentSpec], -) -> VortexResult { - let mut segment_ids = HashSet::::default(); - for node in layout.depth_first_traversal() { - segment_ids.extend(node?.segment_ids()); - } - - segment_ids.into_iter().try_fold(0u64, |size, segment_id| { - let segment = segments.get(*segment_id as usize).ok_or_else(|| { - vortex_err!( - "layout references missing segment {} (segment count: {})", - *segment_id, - segments.len() - ) - })?; - Ok(size + u64::from(segment.length)) - }) -} From 9211c32892c76970f091d87a567820cc40e8b613 Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Mon, 13 Jul 2026 20:23:49 +0100 Subject: [PATCH 3/7] Remove CountingVortexWrite::with_counter Construct the counting write before spawning the JNI write task and retrieve the shared counter directly via counter(), instead of threading a pre-allocated counter through with_counter. Sink creation now happens synchronously in NativeWriter.create, surfacing IO errors there instead of at close. Signed-off-by: "Robert Kruszewski" Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01JE2aKEzyogTh1dP126gudc --- vortex-file/src/counting.rs | 8 ------- vortex-jni/src/writer.rs | 47 +++++++++++++++++++------------------ 2 files changed, 24 insertions(+), 31 deletions(-) diff --git a/vortex-file/src/counting.rs b/vortex-file/src/counting.rs index bb5c029da68..003e3eb9465 100644 --- a/vortex-file/src/counting.rs +++ b/vortex-file/src/counting.rs @@ -24,14 +24,6 @@ impl CountingVortexWrite { } } - /// Wrap a writer with an existing shared byte counter. - pub fn with_counter(inner: W, bytes_written: Arc) -> Self { - Self { - inner, - bytes_written, - } - } - /// Returns the shared byte counter updated by this writer. pub fn counter(&self) -> Arc { Arc::clone(&self.bytes_written) diff --git a/vortex-jni/src/writer.rs b/vortex-jni/src/writer.rs index b44d4b53d4d..b75f8983efa 100644 --- a/vortex-jni/src/writer.rs +++ b/vortex-jni/src/writer.rs @@ -42,6 +42,7 @@ use vortex::array::stream::ArrayStreamAdapter; use vortex::dtype::DType; use vortex::dtype::Field as DTypeField; use vortex::dtype::FieldPath; +use vortex::error::VortexError; use vortex::error::VortexResult; use vortex::error::vortex_err; use vortex::expr::stats::Stat; @@ -442,37 +443,37 @@ 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 write_options = write_options_for_schema(session, &write_schema); - let bytes_written = Arc::new(AtomicU64::new(0)); - let task_bytes_written = Arc::clone(&bytes_written); - let handle = session.handle().spawn(async move { - match resolved { - ResolvedStore::Path(path) => { + let (bytes_written, handle) = match resolved { + ResolvedStore::Path(path) => { + let file = RUNTIME.block_on(async { if let Some(parent) = path.parent().filter(|p| !p.as_os_str().is_empty()) { async_fs::create_dir_all(parent).await?; } - let mut file = File::create(path).await?; - let write = CountingVortexWrite::with_counter( - &mut file, - Arc::clone(&task_bytes_written), - ); - let summary = write_options.write(write, stream).await?; - file.shutdown().await?; + Ok::<_, VortexError>(File::create(path).await?) + })?; + let mut write = CountingVortexWrite::new(file); + 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) - } - ResolvedStore::ObjectStore(store, path) => { - let mut write = - ObjectStoreWrite::new(Arc::new(Compat::new(store)), &path).await?; - let counted_write = CountingVortexWrite::with_counter( - &mut write, - Arc::clone(&task_bytes_written), - ); - let summary = write_options.write(counted_write, stream).await?; + }); + (bytes_written, handle) + } + ResolvedStore::ObjectStore(store, path) => { + let object_write = + RUNTIME.block_on(ObjectStoreWrite::new(Arc::new(Compat::new(store)), &path))?; + let mut write = CountingVortexWrite::new(object_write); + 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) - } + }); + (bytes_written, handle) } - }); + }; Ok(Box::new(NativeWriter::new( session.clone(), From 711f8e2a8ed32fb601d034301dcd882b816dd887 Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Wed, 15 Jul 2026 17:18:02 +0100 Subject: [PATCH 4/7] simplify Signed-off-by: Robert Kruszewski --- vortex-jni/src/writer.rs | 109 ++++++++++++++------------------------- 1 file changed, 38 insertions(+), 71 deletions(-) diff --git a/vortex-jni/src/writer.rs b/vortex-jni/src/writer.rs index b75f8983efa..7fa58ce7b91 100644 --- a/vortex-jni/src/writer.rs +++ b/vortex-jni/src/writer.rs @@ -212,15 +212,6 @@ impl NativeWriter { } } -fn exact_stat_scalar(stats: &StatsSet, dtype: &DType, stat: Stat) -> Option { - let value = stats.get(stat).as_exact()?; - Scalar::try_new(stat.dtype(dtype)?, Some(value)).ok() -} - -fn exact_count(stats: &StatsSet, dtype: &DType, stat: Stat) -> Option { - u64::try_from(&exact_stat_scalar(stats, dtype, stat)?).ok() -} - fn checked_jlong(value: u64, name: &str) -> VortexResult { jlong::try_from(value).map_err(|_| vortex_err!("{name} exceeds Java long range: {value}")) } @@ -232,52 +223,12 @@ fn exact_count_jlong( ) -> VortexResult { stats .zip(dtype) - .and_then(|(stats, dtype)| exact_count(stats, dtype, stat)) + .and_then(|(stats, dtype)| stats.get_as::(stat, dtype).as_exact()) .map(|value| checked_jlong(value, stat.name())) .transpose() .map(|value| value.unwrap_or(-1)) } -fn boxed_integer<'local>(env: &mut Env<'local>, value: i32) -> Result, JNIError> { - Ok(env.new_object( - jni::jni_str!("java/lang/Integer"), - jni::jni_sig!("(I)V"), - &[JValue::Int(value)], - )?) -} - -fn boxed_long<'local>(env: &mut Env<'local>, value: i64) -> Result, JNIError> { - Ok(env.new_object( - jni::jni_str!("java/lang/Long"), - jni::jni_sig!("(J)V"), - &[JValue::Long(value)], - )?) -} - -fn boxed_float<'local>(env: &mut Env<'local>, value: f32) -> Result, JNIError> { - Ok(env.new_object( - jni::jni_str!("java/lang/Float"), - jni::jni_sig!("(F)V"), - &[JValue::Float(value)], - )?) -} - -fn boxed_double<'local>(env: &mut Env<'local>, value: f64) -> Result, JNIError> { - Ok(env.new_object( - jni::jni_str!("java/lang/Double"), - jni::jni_sig!("(D)V"), - &[JValue::Double(value)], - )?) -} - -fn boxed_boolean<'local>(env: &mut Env<'local>, value: bool) -> Result, JNIError> { - Ok(env.new_object( - jni::jni_str!("java/lang/Boolean"), - jni::jni_sig!("(Z)V"), - &[JValue::Bool(if value { JNI_TRUE } else { JNI_FALSE })], - )?) -} - fn big_integer<'local>( env: &mut Env<'local>, value: impl ToString, @@ -292,32 +243,48 @@ fn big_integer<'local>( fn scalar_to_java<'local>( env: &mut Env<'local>, - scalar: &Scalar, + scalar: Scalar, ) -> Result, JNIError> { if scalar.is_null() { return Ok(JObject::null()); } if scalar.dtype().is_extension() { - return scalar_to_java(env, &scalar.as_extension().to_storage_scalar()); + return scalar_to_java(env, scalar.as_extension().to_storage_scalar()); } let Some(value) = scalar.value() else { return Ok(JObject::null()); }; match value { - ScalarValue::Bool(value) => boxed_boolean(env, *value), + ScalarValue::Bool(value) => Ok(env.new_object( + jni::jni_str!("java/lang/Boolean"), + jni::jni_sig!("(Z)V"), + &[JValue::Bool(if *value { JNI_TRUE } else { JNI_FALSE })], + )?), ScalarValue::Primitive(value) => match value { - PValue::U8(value) => boxed_integer(env, i32::from(*value)), - PValue::U16(value) => boxed_integer(env, i32::from(*value)), - PValue::U32(value) => boxed_long(env, i64::from(*value)), + PValue::U8(_) | PValue::U16(_) | PValue::I8(_) | PValue::I16(_) | PValue::I32(_) => { + Ok(env.new_object( + jni::jni_str!("java/lang/Integer"), + jni::jni_sig!("(I)V"), + &[JValue::Int(value.cast::()?)], + )?) + } + PValue::U32(_) | PValue::I64(_) => Ok(env.new_object( + jni::jni_str!("java/lang/Long"), + jni::jni_sig!("(J)V"), + &[JValue::Long(value.cast::()?)], + )?), PValue::U64(value) => big_integer(env, value), - PValue::I8(value) => boxed_integer(env, i32::from(*value)), - PValue::I16(value) => boxed_integer(env, i32::from(*value)), - PValue::I32(value) => boxed_integer(env, *value), - PValue::I64(value) => boxed_long(env, *value), - PValue::F16(value) => boxed_float(env, f32::from(*value)), - PValue::F32(value) => boxed_float(env, *value), - PValue::F64(value) => boxed_double(env, *value), + PValue::F16(_) | PValue::F32(_) => Ok(env.new_object( + jni::jni_str!("java/lang/Float"), + jni::jni_sig!("(F)V"), + &[JValue::Float(value.cast::()?)], + )?), + PValue::F64(value) => Ok(env.new_object( + jni::jni_str!("java/lang/Double"), + jni::jni_sig!("(D)V"), + &[JValue::Double(*value)], + )?), }, ScalarValue::Decimal(value) => { let DType::Decimal(decimal_dtype, _) = scalar.dtype() else { @@ -369,18 +336,18 @@ fn write_summary_to_java<'local>( .map_or((None, None), |(stats, dtype)| (Some(stats), Some(dtype))); let null_count = exact_count_jlong(stats, dtype, Stat::NullCount)?; let nan_count = exact_count_jlong(stats, dtype, Stat::NaNCount)?; - let lower_bound = stats - .zip(dtype) - .and_then(|(stats, dtype)| exact_stat_scalar(stats, dtype, Stat::Min)); - let upper_bound = stats - .zip(dtype) - .and_then(|(stats, dtype)| exact_stat_scalar(stats, dtype, Stat::Max)); + let lower_bound = stats.zip(dtype).map(|(stats, dtype)| unsafe { + Scalar::new_unchecked(dtype.clone(), stats.get(Stat::Min).as_exact()) + }); + let upper_bound = stats.zip(dtype).map(|(stats, dtype)| unsafe { + Scalar::new_unchecked(dtype.clone(), stats.get(Stat::Max).as_exact()) + }); let column = env.with_local_frame_returning_local::<_, JObject, JNIError>(16, |env| { - let lower_bound = match &lower_bound { + let lower_bound = match lower_bound { Some(value) => scalar_to_java(env, value)?, None => JObject::null(), }; - let upper_bound = match &upper_bound { + let upper_bound = match upper_bound { Some(value) => scalar_to_java(env, value)?, None => JObject::null(), }; From 92a8257217679ac8ee811d86dc89a403ea4fcd37 Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Wed, 15 Jul 2026 17:53:08 +0100 Subject: [PATCH 5/7] less 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 7fa58ce7b91..64ff255137d 100644 --- a/vortex-jni/src/writer.rs +++ b/vortex-jni/src/writer.rs @@ -222,8 +222,8 @@ fn exact_count_jlong( stat: Stat, ) -> VortexResult { stats - .zip(dtype) - .and_then(|(stats, dtype)| stats.get_as::(stat, dtype).as_exact()) + .zip(dtype.and_then(|dt| stat.dtype(dt))) + .and_then(|(stats, dt)| stats.get_as::(stat, &dt).as_exact()) .map(|value| checked_jlong(value, stat.name())) .transpose() .map(|value| value.unwrap_or(-1)) From fff5e012ed073e509353aea17b409bcaf30f36d0 Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Wed, 15 Jul 2026 18:13:50 +0100 Subject: [PATCH 6/7] fixmore Signed-off-by: Robert Kruszewski --- vortex-jni/src/writer.rs | 18 ++++++++++++------ 1 file changed, 12 insertions(+), 6 deletions(-) diff --git a/vortex-jni/src/writer.rs b/vortex-jni/src/writer.rs index 64ff255137d..5ac99416811 100644 --- a/vortex-jni/src/writer.rs +++ b/vortex-jni/src/writer.rs @@ -336,12 +336,18 @@ fn write_summary_to_java<'local>( .map_or((None, None), |(stats, dtype)| (Some(stats), Some(dtype))); let null_count = exact_count_jlong(stats, dtype, Stat::NullCount)?; let nan_count = exact_count_jlong(stats, dtype, Stat::NaNCount)?; - let lower_bound = stats.zip(dtype).map(|(stats, dtype)| unsafe { - Scalar::new_unchecked(dtype.clone(), stats.get(Stat::Min).as_exact()) - }); - let upper_bound = stats.zip(dtype).map(|(stats, dtype)| unsafe { - Scalar::new_unchecked(dtype.clone(), stats.get(Stat::Max).as_exact()) - }); + let lower_bound = + stats + .zip(dtype.and_then(|dt| Stat::Min.dtype(dt))) + .map(|(stats, dtype)| unsafe { + Scalar::new_unchecked(dtype, stats.get(Stat::Min).as_exact()) + }); + let upper_bound = + stats + .zip(dtype.and_then(|dt| Stat::Max.dtype(dt))) + .map(|(stats, dtype)| unsafe { + Scalar::new_unchecked(dtype, stats.get(Stat::Max).as_exact()) + }); let column = env.with_local_frame_returning_local::<_, JObject, JNIError>(16, |env| { let lower_bound = match lower_bound { Some(value) => scalar_to_java(env, value)?, From 25dcaa8ec9a70cbb4da9c27069591ee4770acf1e Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Wed, 15 Jul 2026 18:55:52 +0100 Subject: [PATCH 7/7] fix Signed-off-by: Robert Kruszewski --- vortex-array/src/stats/stats_set.rs | 8 ++------ vortex-jni/src/writer.rs | 19 +++++++------------ 2 files changed, 9 insertions(+), 18 deletions(-) diff --git a/vortex-array/src/stats/stats_set.rs b/vortex-array/src/stats/stats_set.rs index 6025afed2c6..cb4af0b2c7f 100644 --- a/vortex-array/src/stats/stats_set.rs +++ b/vortex-array/src/stats/stats_set.rs @@ -1,6 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::any::type_name; use std::fmt::Debug; use enum_iterator::all; @@ -133,12 +134,7 @@ impl StatsSet { .vortex_expect("failed to construct a scalar statistic"), ) .unwrap_or_else(|err| { - vortex_panic!( - err, - "Failed to get stat {} as {}", - stat, - std::any::type_name::() - ) + vortex_panic!(err, "Failed to get stat {} as {}", stat, type_name::()) }) }) } diff --git a/vortex-jni/src/writer.rs b/vortex-jni/src/writer.rs index 5ac99416811..309248fd73c 100644 --- a/vortex-jni/src/writer.rs +++ b/vortex-jni/src/writer.rs @@ -46,6 +46,7 @@ use vortex::error::VortexError; 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; @@ -336,18 +337,12 @@ fn write_summary_to_java<'local>( .map_or((None, None), |(stats, dtype)| (Some(stats), Some(dtype))); let null_count = exact_count_jlong(stats, dtype, Stat::NullCount)?; let nan_count = exact_count_jlong(stats, dtype, Stat::NaNCount)?; - let lower_bound = - stats - .zip(dtype.and_then(|dt| Stat::Min.dtype(dt))) - .map(|(stats, dtype)| unsafe { - Scalar::new_unchecked(dtype, stats.get(Stat::Min).as_exact()) - }); - let upper_bound = - stats - .zip(dtype.and_then(|dt| Stat::Max.dtype(dt))) - .map(|(stats, dtype)| unsafe { - Scalar::new_unchecked(dtype, stats.get(Stat::Max).as_exact()) - }); + let lower_bound = stats + .zip(dtype) + .and_then(|(stats, dtype)| stats.as_typed_ref(dtype).get(Stat::Min).into_inner()); + let upper_bound = stats + .zip(dtype) + .and_then(|(stats, dtype)| stats.as_typed_ref(dtype).get(Stat::Max).into_inner()); let column = env.with_local_frame_returning_local::<_, JObject, JNIError>(16, |env| { let lower_bound = match lower_bound { Some(value) => scalar_to_java(env, value)?,