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-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-file/src/counting.rs b/vortex-file/src/counting.rs index 78f8a8741a9..003e3eb9465 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,7 @@ impl CountingVortexWrite { } } + /// 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/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/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..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; @@ -505,4 +506,29 @@ 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 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 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()) + } } diff --git a/vortex-jni/src/writer.rs b/vortex-jni/src/writer.rs index 5b098b4c320..309248fd73c 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,35 @@ 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::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; use vortex::file::WriteSummary; @@ -130,6 +143,7 @@ pub struct NativeWriter { session: VortexSession, arrow_schema: SchemaRef, write_schema: DType, + bytes_written: Arc, sender: mpsc::Sender>, } @@ -138,6 +152,7 @@ impl NativeWriter { session: VortexSession, arrow_schema: SchemaRef, write_schema: DType, + bytes_written: Arc, handle: Task>, sender: mpsc::Sender>, ) -> Self { @@ -146,6 +161,7 @@ impl NativeWriter { session, arrow_schema, write_schema, + bytes_written, sender, } } @@ -183,17 +199,189 @@ 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 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(|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)) +} + +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) => 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(_) | 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::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 { + 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)| 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)?, + 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)] @@ -224,31 +412,42 @@ pub extern "system" fn Java_dev_vortex_jni_NativeWriter_create( let stream = ArrayStreamAdapter::new(write_schema.clone(), rx); let write_options = write_options_for_schema(session, &write_schema); - 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 summary = write_options.write(&mut file, 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?; + }); + (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(), arrow_schema, write_schema, + bytes_written, handle, tx, )) @@ -285,6 +484,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,