This is an automated email from the ASF dual-hosted git repository.
JingsongLi pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/paimon-mosaic.git
The following commit(s) were added to refs/heads/main by this push:
new b613dd2 [java] Support Arrow batch writes across allocator roots (#73)
b613dd2 is described below
commit b613dd25eba45afbc860171e109df19495ba0e3c
Author: jianguotian <[email protected]>
AuthorDate: Sun Aug 16 09:41:03 2026 +0800
[java] Support Arrow batch writes across allocator roots (#73)
---
core/src/reader_tests.rs | 55 ++
core/src/spec.rs | 2 +-
core/src/writer.rs | 221 ++++++-
docs/java-api.html | 2 +-
.../org/apache/paimon/mosaic/MosaicWriter.java | 196 +++++-
.../apache/paimon/mosaic/MosaicRoundtripTest.java | 725 +++++++++++++++++++++
jni/src/lib.rs | 354 ++++++----
python/mosaic/mosaic.py | 65 +-
python/tests/test_mosaic.py | 108 +++
9 files changed, 1583 insertions(+), 145 deletions(-)
diff --git a/core/src/reader_tests.rs b/core/src/reader_tests.rs
index 44e936e..d89c2e7 100644
--- a/core/src/reader_tests.rs
+++ b/core/src/reader_tests.rs
@@ -504,6 +504,61 @@ fn test_roundtrip_basic() {
}
}
+#[test]
+fn test_writer_recovers_from_batch_type_mismatch() {
+ let arrow_schema = Schema::new(vec![Field::new("id", DataType::Int32,
false)]);
+ let mut writer = MosaicWriter::new(
+ MemOutputFile::new(),
+ &arrow_schema,
+ WriterOptions {
+ compression: COMPRESSION_NONE,
+ num_buckets: 1,
+ row_group_max_size: u64::MAX,
+ ..Default::default()
+ },
+ )
+ .unwrap();
+
+ let first_batch = RecordBatch::try_new(
+ Arc::new(arrow_schema.clone()),
+ vec![Arc::new(Int32Array::from(vec![1]))],
+ )
+ .unwrap();
+ writer.write_batch(&first_batch).unwrap();
+
+ let wrong_type_batch = RecordBatch::try_new(
+ Arc::new(Schema::new(vec![Field::new("id", DataType::Int64, false)])),
+ vec![Arc::new(Int64Array::from(vec![2]))],
+ )
+ .unwrap();
+ let error = writer.write_batch(&wrong_type_batch).unwrap_err();
+ assert_eq!(error.kind(), io::ErrorKind::InvalidInput);
+ assert_eq!(
+ "column 'id' type mismatch: expected Int32, got Int64",
+ error.to_string()
+ );
+
+ let second_batch = RecordBatch::try_new(
+ Arc::new(arrow_schema),
+ vec![Arc::new(Int32Array::from(vec![3]))],
+ )
+ .unwrap();
+ writer.write_batch(&second_batch).unwrap();
+ writer.close().unwrap();
+
+ let data = writer.output().buf.clone();
+ let len = data.len() as u64;
+ let reader = MosaicReader::new(ByteArrayInputFile::new(data),
len).unwrap();
+ assert_eq!(reader.num_row_groups(), 1);
+
+ let mut row_group = reader.row_group_reader(0).unwrap();
+ let batch = row_group.read_columns().unwrap();
+ let ids = batch_col_i32(&batch, "id");
+ assert_eq!(ids.len(), 2);
+ assert_eq!(ids.value(0), 1);
+ assert_eq!(ids.value(1), 3);
+}
+
#[test]
fn test_roundtrip_with_nulls() {
let columns = vec![
diff --git a/core/src/spec.rs b/core/src/spec.rs
index 6d180bf..37bc782 100644
--- a/core/src/spec.rs
+++ b/core/src/spec.rs
@@ -15,7 +15,7 @@
// specific language governing permissions and limitations
// under the License.
-pub const MAGIC: [u8; 4] = [b'M', b'O', b'S', b'A'];
+pub const MAGIC: [u8; 4] = *b"MOSA";
pub const VERSION: u8 = 1;
pub const FOOTER_SIZE: usize = 32;
diff --git a/core/src/writer.rs b/core/src/writer.rs
index 4c1e353..3a3f922 100644
--- a/core/src/writer.rs
+++ b/core/src/writer.rs
@@ -120,6 +120,13 @@ struct RowGroupMeta {
stats: Vec<ColumnStats>,
}
+#[derive(Clone, Copy, Debug, Eq, PartialEq)]
+enum WriterState {
+ Open,
+ Aborted,
+ Closed,
+}
+
pub struct MosaicWriter<S: OutputFile> {
out: S,
schema: MosaicSchema,
@@ -139,7 +146,7 @@ pub struct MosaicWriter<S: OutputFile> {
total_uncompressed: u64,
total_compressed: u64,
stats_collector: Option<StatsCollector>,
- closed: bool,
+ state: WriterState,
}
impl<S: OutputFile> MosaicWriter<S> {
@@ -251,7 +258,7 @@ impl<S: OutputFile> MosaicWriter<S> {
total_uncompressed: 0,
total_compressed: 0,
stats_collector,
- closed: false,
+ state: WriterState::Open,
})
}
@@ -282,12 +289,36 @@ impl<S: OutputFile> MosaicWriter<S> {
}
pub fn write_batch(&mut self, batch: &RecordBatch) -> io::Result<()> {
- if self.closed {
- return Err(io::Error::new(
- io::ErrorKind::InvalidInput,
- "writer is already closed",
- ));
+ match self.state {
+ WriterState::Open => {}
+ WriterState::Aborted => {
+ return Err(io::Error::new(
+ io::ErrorKind::InvalidInput,
+ "writer is aborted after a previous failure",
+ ));
+ }
+ WriterState::Closed => {
+ return Err(io::Error::new(
+ io::ErrorKind::InvalidInput,
+ "writer is already closed",
+ ));
+ }
}
+
+ self.validate_batch(batch)?;
+
+ // From this point on, an error or panic can leave bucket, row-group,
or output state
+ // partially advanced. Keep the writer aborted until the whole
operation succeeds so
+ // retry, close, and Drop cannot flush a batch whose write was
reported as failed.
+ self.state = WriterState::Aborted;
+ let result = self.write_batch_mutating(batch);
+ if result.is_ok() {
+ self.state = WriterState::Open;
+ }
+ result
+ }
+
+ fn validate_batch(&self, batch: &RecordBatch) -> io::Result<()> {
let num_cols = self.schema.columns.len();
if batch.num_columns() != num_cols {
return Err(io::Error::new(
@@ -301,18 +332,33 @@ impl<S: OutputFile> MosaicWriter<S> {
}
for (i, col) in self.schema.columns.iter().enumerate() {
- if !col.nullable &&
batch.column(self.batch_col_map[i]).null_count() > 0 {
+ let array = batch.column(self.batch_col_map[i]);
+ if array.data_type() != &col.data_type {
+ return Err(io::Error::new(
+ io::ErrorKind::InvalidInput,
+ format!(
+ "column '{}' type mismatch: expected {:?}, got {:?}",
+ col.name,
+ col.data_type,
+ array.data_type()
+ ),
+ ));
+ }
+ if !col.nullable && array.null_count() > 0 {
return Err(io::Error::new(
io::ErrorKind::InvalidInput,
format!(
"non-nullable column '{}' has {} nulls in batch",
col.name,
- batch.column(self.batch_col_map[i]).null_count()
+ array.null_count()
),
));
}
}
+ Ok(())
+ }
+ fn write_batch_mutating(&mut self, batch: &RecordBatch) -> io::Result<()> {
let mut size = 0u64;
for &b in &self.active_buckets {
let global_indices = &self.schema.bucket_to_global[b];
@@ -542,11 +588,27 @@ impl<S: OutputFile> MosaicWriter<S> {
}
pub fn close(&mut self) -> io::Result<()> {
- if self.closed {
- return Ok(());
+ match self.state {
+ WriterState::Closed => return Ok(()),
+ WriterState::Aborted => {
+ return Err(io::Error::new(
+ io::ErrorKind::InvalidInput,
+ "writer is aborted after a previous failure",
+ ));
+ }
+ WriterState::Open => {}
}
- self.closed = true;
+ // Closing mutates row-group metadata and the output stream. Keep
failures sticky so a
+ // retry cannot report success for a file whose footer or final flush
was not completed.
+ self.state = WriterState::Aborted;
+ let result = self.close_inner();
+ if result.is_ok() {
+ self.state = WriterState::Closed;
+ }
+ result
+ }
+ fn close_inner(&mut self) -> io::Result<()> {
self.flush_row_group()?;
// Write schema block
@@ -626,7 +688,7 @@ impl<S: OutputFile> MosaicWriter<S> {
impl<S: OutputFile> Drop for MosaicWriter<S> {
fn drop(&mut self) {
- if !self.closed {
+ if self.state == WriterState::Open {
if let Err(e) = self.close() {
eprintln!("MosaicWriter::drop: close failed: {}", e);
}
@@ -638,7 +700,7 @@ impl<S: OutputFile> Drop for MosaicWriter<S> {
mod tests {
use super::*;
use arrow_schema::{DataType, Field, Schema};
- use std::sync::Arc;
+ use std::sync::{Arc, Mutex};
struct MemOutputFile {
buf: Vec<u8>,
@@ -663,6 +725,137 @@ mod tests {
}
}
+ #[derive(Default)]
+ struct FailingOutputState {
+ write_calls: usize,
+ flush_calls: usize,
+ bytes: Vec<u8>,
+ }
+
+ struct FailOnceOutputFile {
+ state: Arc<Mutex<FailingOutputState>>,
+ fail: bool,
+ }
+
+ impl OutputFile for FailOnceOutputFile {
+ fn write(&mut self, data: &[u8]) -> io::Result<()> {
+ let mut state = self.state.lock().unwrap();
+ state.write_calls += 1;
+ if self.fail {
+ self.fail = false;
+ return Err(io::Error::other("sentinel output failure"));
+ }
+ state.bytes.extend_from_slice(data);
+ Ok(())
+ }
+
+ fn flush(&mut self) -> io::Result<()> {
+ self.state.lock().unwrap().flush_calls += 1;
+ Ok(())
+ }
+
+ fn pos(&self) -> u64 {
+ self.state.lock().unwrap().bytes.len() as u64
+ }
+ }
+
+ #[test]
+ fn test_write_failure_aborts_writer_without_retry_or_drop_flush() {
+ let arrow_schema = Schema::new(vec![Field::new("id", DataType::Int32,
false)]);
+ let batch = RecordBatch::try_new(
+ Arc::new(arrow_schema.clone()),
+ vec![Arc::new(Int32Array::from(vec![7]))],
+ )
+ .unwrap();
+ let state = Arc::new(Mutex::new(FailingOutputState::default()));
+
+ {
+ let out = FailOnceOutputFile {
+ state: Arc::clone(&state),
+ fail: true,
+ };
+ let mut writer = MosaicWriter::new(
+ out,
+ &arrow_schema,
+ WriterOptions {
+ compression: COMPRESSION_NONE,
+ num_buckets: 1,
+ row_group_max_size: 1,
+ ..Default::default()
+ },
+ )
+ .unwrap();
+
+ let error = writer.write_batch(&batch).unwrap_err();
+ assert_eq!("sentinel output failure", error.to_string());
+ let calls_after_failure = state.lock().unwrap().write_calls;
+ assert_eq!(1, calls_after_failure);
+
+ let retry_error = writer.write_batch(&batch).unwrap_err();
+ assert!(retry_error
+ .to_string()
+ .contains("writer is aborted after a previous failure"));
+ assert_eq!(calls_after_failure, state.lock().unwrap().write_calls);
+
+ let close_error = writer.close().unwrap_err();
+ assert!(close_error
+ .to_string()
+ .contains("writer is aborted after a previous failure"));
+ assert_eq!(calls_after_failure, state.lock().unwrap().write_calls);
+ }
+
+ let state = state.lock().unwrap();
+ assert_eq!(1, state.write_calls);
+ assert_eq!(0, state.flush_calls);
+ assert!(state.bytes.is_empty());
+ }
+
+ #[test]
+ fn test_close_failure_aborts_writer_without_retry_or_drop_flush() {
+ let arrow_schema = Schema::new(vec![Field::new("id", DataType::Int32,
false)]);
+ let batch = RecordBatch::try_new(
+ Arc::new(arrow_schema.clone()),
+ vec![Arc::new(Int32Array::from(vec![7]))],
+ )
+ .unwrap();
+ let state = Arc::new(Mutex::new(FailingOutputState::default()));
+
+ {
+ let out = FailOnceOutputFile {
+ state: Arc::clone(&state),
+ fail: true,
+ };
+ let mut writer = MosaicWriter::new(
+ out,
+ &arrow_schema,
+ WriterOptions {
+ compression: COMPRESSION_NONE,
+ num_buckets: 1,
+ row_group_max_size: u64::MAX,
+ ..Default::default()
+ },
+ )
+ .unwrap();
+
+ writer.write_batch(&batch).unwrap();
+ let error = writer.close().unwrap_err();
+ assert_eq!("sentinel output failure", error.to_string());
+ let calls_after_failure = state.lock().unwrap().write_calls;
+ assert_eq!(1, calls_after_failure);
+
+ let retry_error = writer.close().unwrap_err();
+ assert!(retry_error
+ .to_string()
+ .contains("writer is aborted after a previous failure"));
+ assert_eq!(calls_after_failure, state.lock().unwrap().write_calls);
+ }
+
+ let state = state.lock().unwrap();
+ assert_eq!(1, state.write_calls);
+ assert_eq!(0, state.flush_calls);
+ assert!(state.bytes.is_empty());
+ }
+
#[test]
fn test_write_simple_file() {
let arrow_schema = Schema::new(vec![
diff --git a/docs/java-api.html b/docs/java-api.html
index 9b9796d..ffc95de 100644
--- a/docs/java-api.html
+++ b/docs/java-api.html
@@ -238,7 +238,7 @@ w.endMap();</code></pre>
<tr><th>Method</th><th>Return</th><th>Description</th></tr>
</thead>
<tbody>
-
<tr><td><code>write(VectorSchemaRoot)</code></td><td><code>void</code></td><td>Write
an Arrow batch (zero-copy via C Data Interface)</td></tr>
+
<tr><td><code>write(VectorSchemaRoot)</code></td><td><code>void</code></td><td>Write
an Arrow batch synchronously (zero-copy via C Data Interface). All top-level
and nested vectors must share one allocator root, which may differ from the
writer allocator root; temporary export metadata is charged to the writer
allocator.</td></tr>
<tr><td><code>estimatedFileSize()</code></td><td><code>long</code></td><td>Estimated
output file size in bytes (for file rolling)</td></tr>
<tr><td><code>close()</code></td><td><code>void</code></td><td>Flush remaining
data and write footer</td></tr>
<tr><td><code>numRowGroups()</code></td><td><code>int</code></td><td>Number of
row groups written (available after close)</td></tr>
diff --git a/java/src/main/java/org/apache/paimon/mosaic/MosaicWriter.java
b/java/src/main/java/org/apache/paimon/mosaic/MosaicWriter.java
index 461a694..bd6818c 100644
--- a/java/src/main/java/org/apache/paimon/mosaic/MosaicWriter.java
+++ b/java/src/main/java/org/apache/paimon/mosaic/MosaicWriter.java
@@ -29,15 +29,30 @@ import java.util.Map;
import org.apache.arrow.c.ArrowArray;
import org.apache.arrow.c.ArrowSchema;
import org.apache.arrow.c.Data;
+import org.apache.arrow.c.jni.JniWrapper;
+import org.apache.arrow.c.jni.PrivateData;
+import org.apache.arrow.memory.ArrowBuf;
import org.apache.arrow.memory.BufferAllocator;
+import org.apache.arrow.vector.FieldVector;
import org.apache.arrow.vector.VectorSchemaRoot;
import org.apache.arrow.vector.types.pojo.Schema;
public class MosaicWriter implements AutoCloseable {
+ @FunctionalInterface
+ interface RootArrayExporter {
+
+ void export(long address, PrivateData privateData);
+ }
+
+ private static final RootArrayExporter JNI_ROOT_ARRAY_EXPORTER =
+ (address, privateData) ->
+ JniWrapper.get().exportArray(address, privateData);
+
private long handle;
private boolean closed;
private final BufferAllocator allocator;
+ private final RootArrayExporter rootArrayExporter;
private List<Map<String, ColumnStatistics>> rowGroupStats;
public MosaicWriter(OutputStream outputStream, Schema arrowSchema,
BufferAllocator allocator) {
@@ -45,7 +60,17 @@ public class MosaicWriter implements AutoCloseable {
}
public MosaicWriter(OutputStream outputStream, Schema arrowSchema,
WriterOptions options, BufferAllocator allocator) {
+ this(outputStream, arrowSchema, options, allocator,
JNI_ROOT_ARRAY_EXPORTER);
+ }
+
+ MosaicWriter(
+ OutputStream outputStream,
+ Schema arrowSchema,
+ WriterOptions options,
+ BufferAllocator allocator,
+ RootArrayExporter rootArrayExporter) {
this.allocator = allocator;
+ this.rootArrayExporter = rootArrayExporter;
try (ArrowSchema cSchema = ArrowSchema.allocateNew(allocator)) {
try {
Data.exportSchema(allocator, arrowSchema, null, cSchema);
@@ -69,19 +94,175 @@ public class MosaicWriter implements AutoCloseable {
}
}
+ /**
+ * Writes an Arrow batch synchronously.
+ *
+ * <p>All top-level and nested field vectors must share one allocator
root. That root may be
+ * independent from the writer allocator root; temporary Arrow C Data
metadata remains charged
+ * to the writer allocator supplied at construction time. The caller
retains ownership of the
+ * batch and must keep it open until this method returns.
+ *
+ * @param root batch to write
+ * @throws IllegalArgumentException if field vectors use different
allocator roots
+ * @throws IllegalStateException if the writer is closed
+ */
public void write(VectorSchemaRoot root) {
if (closed || handle == 0) {
throw new IllegalStateException("writer is closed");
}
+ boolean sameAllocatorRoot = sharesWriterAllocatorRoot(root);
try (ArrowArray arrowArray = ArrowArray.allocateNew(allocator);
ArrowSchema arrowSchema = ArrowSchema.allocateNew(allocator)) {
try {
- Data.exportVectorSchemaRoot(allocator, root, null, arrowArray,
arrowSchema);
+ if (sameAllocatorRoot) {
+ Data.exportVectorSchemaRoot(
+ allocator, root, null, arrowArray, arrowSchema);
+ } else {
+ Data.exportSchema(allocator, root.getSchema(), null,
arrowSchema);
+ exportCrossRootArray(
+ allocator, root, arrowArray, rootArrayExporter);
+ }
NativeLib.nativeWriterWriteBatch(handle,
arrowArray.memoryAddress(), arrowSchema.memoryAddress());
} finally {
releaseExported(arrowArray);
releaseExported(arrowSchema);
}
+ } catch (Throwable failure) {
+ throw propagateNativeFailure("write batch failed", failure);
+ }
+ }
+
+ private static void exportCrossRootArray(
+ BufferAllocator exportAllocator,
+ VectorSchemaRoot root,
+ ArrowArray arrowArray,
+ RootArrayExporter rootArrayExporter) {
+ // Data.exportVectorSchemaRoot reloads every field into a temporary
StructVector. If that
+ // reload fails before the root ArrowArray owns a release callback,
Arrow 15 can retain
+ // already-associated input buffers. Export each child directly
instead: input buffers are
+ // retained without being associated with the writer allocator, while
temporary C Data
+ // metadata remains charged to that allocator.
+ RootArrayPrivateData privateData = new RootArrayPrivateData();
+ try {
+ privateData.bufferPointers = exportAllocator.buffer(Long.BYTES);
+ privateData.bufferPointers.writeLong(0L);
+
+ List<FieldVector> vectors = root.getFieldVectors();
+ if (!vectors.isEmpty()) {
+ privateData.childPointers =
+ exportAllocator.buffer((long) vectors.size() *
Long.BYTES);
+ for (int i = 0; i < vectors.size(); i++) {
+ ArrowArray child = ArrowArray.allocateNew(exportAllocator);
+ privateData.children.add(child);
+ privateData.childPointers.writeLong(child.memoryAddress());
+ }
+ for (int i = 0; i < vectors.size(); i++) {
+ Data.exportVector(
+ exportAllocator, vectors.get(i), null,
privateData.children.get(i));
+ }
+ }
+
+ ArrowArray.Snapshot snapshot = new ArrowArray.Snapshot();
+ snapshot.length = root.getRowCount();
+ snapshot.null_count = 0;
+ snapshot.offset = 0;
+ snapshot.n_buffers = 1;
+ snapshot.n_children = vectors.size();
+ snapshot.buffers = privateData.bufferPointers.memoryAddress();
+ snapshot.children =
+ privateData.childPointers == null
+ ? 0
+ : privateData.childPointers.memoryAddress();
+ snapshot.dictionary = 0;
+ snapshot.release = 0;
+ arrowArray.save(snapshot);
+ rootArrayExporter.export(arrowArray.memoryAddress(), privateData);
+ } catch (RuntimeException | Error failure) {
+ if (arrowArray.snapshot().release != 0) {
+ // Arrow 15 may install the root callback even when
NewGlobalRef leaves a pending
+ // OutOfMemoryError. Once installed, that callback owns child
traversal, so keep
+ // the child pointer table alive until the callback has run.
+ try {
+ arrowArray.release();
+ } catch (RuntimeException | Error cleanupFailure) {
+ failure.addSuppressed(cleanupFailure);
+ }
+ if (arrowArray.snapshot().release == 0) {
+ try {
+ privateData.close();
+ } catch (RuntimeException | Error cleanupFailure) {
+ failure.addSuppressed(cleanupFailure);
+ }
+ }
+ } else {
+ privateData.abort(failure);
+ }
+ throw failure;
+ }
+ }
+
+ private static final class RootArrayPrivateData implements PrivateData {
+
+ private ArrowBuf bufferPointers;
+ private ArrowBuf childPointers;
+ private final List<ArrowArray> children = new ArrayList<>();
+
+ private void abort(Throwable failure) {
+ for (ArrowArray child : children) {
+ try {
+ releaseExported(child);
+ } catch (RuntimeException | Error cleanupFailure) {
+ failure.addSuppressed(cleanupFailure);
+ }
+ }
+ try {
+ close();
+ } catch (RuntimeException | Error cleanupFailure) {
+ failure.addSuppressed(cleanupFailure);
+ }
+ }
+
+ @Override
+ public void close() {
+ if (bufferPointers != null) {
+ bufferPointers.close();
+ bufferPointers = null;
+ }
+ if (childPointers != null) {
+ childPointers.close();
+ childPointers = null;
+ }
+ for (ArrowArray child : children) {
+ child.close();
+ }
+ children.clear();
+ }
+ }
+
+ private boolean sharesWriterAllocatorRoot(VectorSchemaRoot root) {
+ List<FieldVector> vectors = root.getFieldVectors();
+ if (vectors.isEmpty()) {
+ return true;
+ }
+
+ BufferAllocator inputRoot = vectors.get(0).getAllocator().getRoot();
+ for (FieldVector vector : vectors) {
+ validateAllocatorRoot(vector, inputRoot,
vector.getField().getName());
+ }
+ return inputRoot == allocator.getRoot();
+ }
+
+ private static void validateAllocatorRoot(
+ FieldVector vector, BufferAllocator expectedRoot, String
fieldPath) {
+ if (vector.getAllocator().getRoot() != expectedRoot) {
+ throw new IllegalArgumentException(
+ "All field vectors must share the same allocator root;
field '"
+ + fieldPath
+ + "' uses a different root");
+ }
+ for (FieldVector child : vector.getChildrenFromFields()) {
+ validateAllocatorRoot(
+ child, expectedRoot, fieldPath + "." +
child.getField().getName());
}
}
@@ -97,6 +278,17 @@ public class MosaicWriter implements AutoCloseable {
}
}
+ private static RuntimeException propagateNativeFailure(
+ String message, Throwable failure) {
+ if (failure instanceof RuntimeException) {
+ return (RuntimeException) failure;
+ }
+ if (failure instanceof Error) {
+ throw (Error) failure;
+ }
+ return new RuntimeException(message, failure);
+ }
+
public long estimatedFileSize() {
return NativeLib.nativeWriterEstimatedSize(handle);
}
@@ -125,6 +317,8 @@ public class MosaicWriter implements AutoCloseable {
try {
NativeLib.nativeWriterClose(handle);
collectStatistics();
+ } catch (Throwable failure) {
+ throw propagateNativeFailure("close failed", failure);
} finally {
NativeLib.nativeWriterFree(handle);
handle = 0;
diff --git
a/java/src/test/java/org/apache/paimon/mosaic/MosaicRoundtripTest.java
b/java/src/test/java/org/apache/paimon/mosaic/MosaicRoundtripTest.java
index 4e7aea5..437cd5d 100644
--- a/java/src/test/java/org/apache/paimon/mosaic/MosaicRoundtripTest.java
+++ b/java/src/test/java/org/apache/paimon/mosaic/MosaicRoundtripTest.java
@@ -20,15 +20,23 @@
package org.apache.paimon.mosaic;
import java.io.ByteArrayOutputStream;
+import java.io.IOException;
+import java.io.OutputStream;
import java.lang.ref.WeakReference;
import java.nio.ByteBuffer;
import java.nio.ByteOrder;
+import java.util.ArrayList;
import java.util.Arrays;
+import java.util.List;
+import org.apache.arrow.c.jni.JniWrapper;
+import org.apache.arrow.memory.ArrowBuf;
import org.apache.arrow.memory.BufferAllocator;
+import org.apache.arrow.memory.OutOfMemoryException;
import org.apache.arrow.memory.RootAllocator;
import org.apache.arrow.vector.BigIntVector;
import org.apache.arrow.vector.BitVector;
+import org.apache.arrow.vector.FieldVector;
import org.apache.arrow.vector.Float4Vector;
import org.apache.arrow.vector.Float8Vector;
import org.apache.arrow.vector.IntVector;
@@ -59,6 +67,88 @@ public class MosaicRoundtripTest {
private BufferAllocator allocator;
+ private static final class InjectableListVector extends ListVector {
+
+ private InjectableListVector(String name, BufferAllocator allocator) {
+ super(name, allocator,
FieldType.nullable(ArrowType.List.INSTANCE), null);
+ }
+
+ private void setDataVector(FieldVector vector) {
+ replaceDataVector(vector);
+ }
+ }
+
+ private static final class InjectedExportException extends
RuntimeException {}
+
+ private static final class FailOnceListVector extends ListVector {
+
+ private boolean fail = true;
+
+ private FailOnceListVector(String name, BufferAllocator allocator) {
+ super(name, allocator,
FieldType.nullable(ArrowType.List.INSTANCE), null);
+ replaceDataVector(new IntVector(ListVector.DATA_VECTOR_NAME,
allocator));
+ }
+
+ @Override
+ public List<ArrowBuf> getFieldBuffers() {
+ if (fail) {
+ fail = false;
+ throw new InjectedExportException();
+ }
+ return super.getFieldBuffers();
+ }
+ }
+
+ private enum FailurePoint {
+ WRITE,
+ FLUSH
+ }
+
+ private static final class FailOnceOutputStream extends OutputStream {
+
+ private final ByteArrayOutputStream delegate = new
ByteArrayOutputStream();
+ private final FailurePoint failurePoint;
+ private final IllegalStateException failureCause;
+ private final IOException failure;
+ private boolean failed;
+ private int writeCalls;
+ private int flushCalls;
+
+ private FailOnceOutputStream(FailurePoint failurePoint, String
message) {
+ this.failurePoint = failurePoint;
+ this.failureCause = new IllegalStateException(message + "-cause");
+ this.failure = new IOException(message, failureCause);
+ }
+
+ @Override
+ public void write(int value) throws IOException {
+ write(new byte[] {(byte) value}, 0, 1);
+ }
+
+ @Override
+ public void write(byte[] bytes, int offset, int length) throws
IOException {
+ writeCalls++;
+ if (!failed && failurePoint == FailurePoint.WRITE) {
+ failed = true;
+ throw failure;
+ }
+ delegate.write(bytes, offset, length);
+ }
+
+ @Override
+ public void flush() throws IOException {
+ flushCalls++;
+ if (!failed && failurePoint == FailurePoint.FLUSH) {
+ failed = true;
+ throw failure;
+ }
+ }
+
+ private int size() {
+ return delegate.size();
+ }
+ }
+
@Before
public void setUp() {
allocator = new RootAllocator();
@@ -88,6 +178,14 @@ public class MosaicRoundtripTest {
return MosaicReader.open(inputFile, data.length, allocator);
}
+ private static Schema wideIntSchema(int width) {
+ List<Field> fields = new ArrayList<>(width);
+ for (int i = 0; i < width; i++) {
+ fields.add(Field.nullable("c" + i, new ArrowType.Int(32, true)));
+ }
+ return new Schema(fields);
+ }
+
private static void awaitGarbageCollection(WeakReference<?> reference)
throws InterruptedException {
for (int i = 0; i < 20 && reference.get() != null; i++) {
System.gc();
@@ -181,6 +279,633 @@ public class MosaicRoundtripTest {
}
}
+ @Test
+ public void testWriteFromIndependentRootAllocator() {
+ Schema arrowSchema = new Schema(Arrays.asList(
+ Field.notNullable("id", new ArrowType.Int(32, true)),
+ Field.nullable("name", ArrowType.Utf8.INSTANCE)
+ ));
+
+ byte[] data;
+ try (BufferAllocator inputAllocator = new RootAllocator();
+ VectorSchemaRoot root = VectorSchemaRoot.create(arrowSchema,
inputAllocator)) {
+ IntVector ids = (IntVector) root.getVector("id");
+ VarCharVector names = (VarCharVector) root.getVector("name");
+
+ ids.allocateNew(3);
+ names.allocateNew(3);
+ for (int i = 0; i < 3; i++) {
+ ids.set(i, i + 1);
+ names.setSafe(i, ("input_" + i).getBytes());
+ }
+ root.setRowCount(3);
+
+ data = writeToBytes(arrowSchema, writer -> writer.write(root));
+ }
+
+ try (MosaicReader reader = readerFromBytes(data);
+ VectorSchemaRoot batch = reader.readRowGroup(0, allocator)) {
+ assertEquals(3, batch.getRowCount());
+ IntVector ids = (IntVector) batch.getVector("id");
+ VarCharVector names = (VarCharVector) batch.getVector("name");
+ for (int i = 0; i < 3; i++) {
+ assertEquals(i + 1, ids.get(i));
+ assertEquals("input_" + i, new String(names.get(i)));
+ }
+ }
+ }
+
+ @Test
+ public void testWriteFromLimitedChildAllocator() {
+ Schema arrowSchema = new Schema(Arrays.asList(
+ Field.notNullable("id", new ArrowType.Int(32, true))
+ ));
+
+ byte[] data;
+ try (BufferAllocator inputRoot = new RootAllocator();
+ BufferAllocator limitedAllocator =
+ inputRoot.newChildAllocator("limited-input", 0, 512);
+ VectorSchemaRoot root =
+ VectorSchemaRoot.create(arrowSchema, limitedAllocator)) {
+ IntVector ids = (IntVector) root.getVector("id");
+ ids.allocateNew(1);
+ ids.set(0, 7);
+ root.setRowCount(1);
+
+ long inputBytes = inputRoot.getAllocatedMemory();
+ long inputPeak = limitedAllocator.getPeakMemoryAllocation();
+ int inputChildren = inputRoot.getChildAllocators().size();
+ ByteArrayOutputStream output = new ByteArrayOutputStream();
+ try (MosaicWriter writer =
+ new MosaicWriter(output, arrowSchema, allocator)) {
+ writer.write(root);
+ assertEquals(inputBytes, inputRoot.getAllocatedMemory());
+ assertEquals(inputPeak,
limitedAllocator.getPeakMemoryAllocation());
+ assertEquals(inputChildren,
inputRoot.getChildAllocators().size());
+ }
+ data = output.toByteArray();
+ }
+
+ try (MosaicReader reader = readerFromBytes(data);
+ VectorSchemaRoot batch = reader.readRowGroup(0, allocator)) {
+ assertEquals(1, batch.getRowCount());
+ assertEquals(7, ((IntVector) batch.getVector("id")).get(0));
+ }
+ }
+
+ @Test
+ public void testCrossRootExportUsesWriterAllocatorAndReleasesMetadata() {
+ Schema arrowSchema = wideIntSchema(5_000);
+
+ byte[] data;
+ try (RootAllocator inputRoot = new RootAllocator(16L * 1024 * 1024);
+ BufferAllocator limitedAllocator =
+ inputRoot.newChildAllocator("limited-input", 0, 512);
+ VectorSchemaRoot root =
+ VectorSchemaRoot.create(arrowSchema, limitedAllocator)) {
+ root.setRowCount(0);
+ long inputBytes = inputRoot.getAllocatedMemory();
+ long inputPeak = inputRoot.getPeakMemoryAllocation();
+ int inputChildren = inputRoot.getChildAllocators().size();
+ long writerBytes = allocator.getAllocatedMemory();
+
+ ByteArrayOutputStream output = new ByteArrayOutputStream();
+ try (MosaicWriter writer =
+ new MosaicWriter(output, arrowSchema, allocator)) {
+ long writerPeak = allocator.getPeakMemoryAllocation();
+ writer.write(root);
+ assertEquals(inputBytes, inputRoot.getAllocatedMemory());
+ assertEquals(inputPeak, inputRoot.getPeakMemoryAllocation());
+ assertEquals(inputChildren,
inputRoot.getChildAllocators().size());
+ assertEquals(writerBytes, allocator.getAllocatedMemory());
+ assertTrue(
+ "expected Arrow C Data metadata on the writer
allocator",
+ allocator.getPeakMemoryAllocation() > writerPeak);
+ }
+ data = output.toByteArray();
+ }
+
+ assertTrue(data.length > 32);
+ }
+
+ @Test
+ public void testSameRootExportKeepsWriterAllocatorAccounting() {
+ Schema arrowSchema = new Schema(Arrays.asList(
+ Field.notNullable("id", new ArrowType.Int(32, true))
+ ));
+
+ byte[] data;
+ try (RootAllocator sharedRoot = new RootAllocator(16L * 1024 * 1024);
+ BufferAllocator writerAllocator =
+ sharedRoot.newChildAllocator("writer", 0, 16L * 1024 *
1024);
+ BufferAllocator inputAllocator =
+ sharedRoot.newChildAllocator("input", 0, 16L * 1024 *
1024);
+ VectorSchemaRoot root =
+ VectorSchemaRoot.create(arrowSchema, inputAllocator)) {
+ IntVector ids = (IntVector) root.getVector("id");
+ ids.allocateNew(1);
+ ids.set(0, 7);
+ root.setRowCount(1);
+
+ long inputBytes = inputAllocator.getAllocatedMemory();
+ long inputPeak = inputAllocator.getPeakMemoryAllocation();
+
+ ByteArrayOutputStream output = new ByteArrayOutputStream();
+ try (MosaicWriter writer =
+ new MosaicWriter(output, arrowSchema, writerAllocator)) {
+ long writerPeak = writerAllocator.getPeakMemoryAllocation();
+ writer.write(root);
+ assertTrue(
+ "expected Arrow C Data metadata on the writer
allocator",
+ writerAllocator.getPeakMemoryAllocation() >
writerPeak);
+ assertEquals(inputBytes, inputAllocator.getAllocatedMemory());
+ assertEquals(inputPeak,
inputAllocator.getPeakMemoryAllocation());
+ }
+ data = output.toByteArray();
+ }
+
+ try (MosaicReader reader = readerFromBytes(data);
+ VectorSchemaRoot batch = reader.readRowGroup(0, allocator)) {
+ assertEquals(1, batch.getRowCount());
+ assertEquals(7, ((IntVector) batch.getVector("id")).get(0));
+ }
+ }
+
+ @Test
+ public void testCrossRootWriterAllocatorOutOfMemoryCanRetryWithoutLeak() {
+ Schema arrowSchema = new Schema(Arrays.asList(
+ Field.nullable("id", new ArrowType.Int(32, true))
+ ));
+
+ byte[] data;
+ try (RootAllocator writerRoot = new RootAllocator(16L * 1024 * 1024);
+ RootAllocator inputRoot = new RootAllocator(16L * 1024 * 1024)) {
+ try (BufferAllocator writerAllocator =
+ writerRoot.newChildAllocator("writer", 0, 16L * 1024
* 1024);
+ BufferAllocator inputAllocator =
+ inputRoot.newChildAllocator("limited-input", 0, 16L *
1024 * 1024);
+ VectorSchemaRoot root =
+ VectorSchemaRoot.create(arrowSchema, inputAllocator))
{
+ IntVector ids = (IntVector) root.getVector("id");
+ int rowCount = 65_536;
+ ids.allocateNew(rowCount);
+ for (int i = 0; i < rowCount; i++) {
+ ids.set(i, i);
+ }
+ root.setRowCount(rowCount);
+
+ long inputBytes = inputRoot.getAllocatedMemory();
+ ByteArrayOutputStream output = new ByteArrayOutputStream();
+ try (MosaicWriter writer =
+ new MosaicWriter(output, arrowSchema,
writerAllocator)) {
+ long writerBytes = writerAllocator.getAllocatedMemory();
+ writerAllocator.setLimit(writerBytes + 512);
+ assertThrows(OutOfMemoryException.class, () ->
writer.write(root));
+ assertEquals(inputBytes, inputRoot.getAllocatedMemory());
+ assertEquals(inputBytes,
inputAllocator.getAllocatedMemory());
+ assertEquals(writerBytes,
writerAllocator.getAllocatedMemory());
+ assertEquals(rowCount - 1, ids.get(rowCount - 1));
+
+ writerAllocator.setLimit(16L * 1024 * 1024);
+ writer.write(root);
+ }
+ data = output.toByteArray();
+ }
+ assertEquals(0, inputRoot.getAllocatedMemory());
+ assertEquals(0, writerRoot.getAllocatedMemory());
+ }
+
+ try (MosaicReader reader = readerFromBytes(data);
+ VectorSchemaRoot batch = reader.readRowGroup(0, allocator)) {
+ assertEquals(65_536, batch.getRowCount());
+ assertEquals(65_535, ((IntVector)
batch.getVector("id")).get(65_535));
+ }
+ }
+
+ @Test
+ public void testCrossRootFailureAfterRootRegistrationCanRetryWithoutLeak()
{
+ Schema arrowSchema = new Schema(Arrays.asList(
+ Field.notNullable("id", new ArrowType.Int(32, true))
+ ));
+ OutOfMemoryError injected =
+ new OutOfMemoryError("injected after root callback
registration");
+ boolean[] failOnce = {true};
+ MosaicWriter.RootArrayExporter rootArrayExporter =
+ (address, privateData) -> {
+ JniWrapper.get().exportArray(address, privateData);
+ if (failOnce[0]) {
+ failOnce[0] = false;
+ throw injected;
+ }
+ };
+
+ byte[] data;
+ try (RootAllocator writerRoot = new RootAllocator(16L * 1024 * 1024);
+ RootAllocator inputRoot = new RootAllocator(16L * 1024 * 1024);
+ BufferAllocator writerAllocator =
+ writerRoot.newChildAllocator("writer", 0, 16L * 1024 *
1024);
+ BufferAllocator inputAllocator =
+ inputRoot.newChildAllocator("input", 0, 16L * 1024 *
1024);
+ VectorSchemaRoot root =
+ VectorSchemaRoot.create(arrowSchema, inputAllocator)) {
+ IntVector ids = (IntVector) root.getVector("id");
+ ids.allocateNew(1);
+ ids.set(0, 7);
+ root.setRowCount(1);
+
+ long inputBytes = inputAllocator.getAllocatedMemory();
+ List<ArrowBuf> fieldBuffers = ids.getFieldBuffers();
+ int[] refCounts = new int[fieldBuffers.size()];
+ for (int i = 0; i < fieldBuffers.size(); i++) {
+ refCounts[i] = fieldBuffers.get(i).refCnt();
+ }
+
+ ByteArrayOutputStream output = new ByteArrayOutputStream();
+ try (MosaicWriter writer =
+ new MosaicWriter(
+ output,
+ arrowSchema,
+ new WriterOptions(),
+ writerAllocator,
+ rootArrayExporter)) {
+ long writerBytes = writerAllocator.getAllocatedMemory();
+ assertSame(injected, assertThrows(OutOfMemoryError.class, ()
-> writer.write(root)));
+ assertEquals(inputBytes, inputAllocator.getAllocatedMemory());
+ assertEquals(writerBytes,
writerAllocator.getAllocatedMemory());
+ for (int i = 0; i < fieldBuffers.size(); i++) {
+ assertEquals(refCounts[i], fieldBuffers.get(i).refCnt());
+ }
+
+ writer.write(root);
+ assertEquals(inputBytes, inputAllocator.getAllocatedMemory());
+ assertEquals(writerBytes,
writerAllocator.getAllocatedMemory());
+ for (int i = 0; i < fieldBuffers.size(); i++) {
+ assertEquals(refCounts[i], fieldBuffers.get(i).refCnt());
+ }
+ }
+ data = output.toByteArray();
+ }
+
+ try (MosaicReader reader = readerFromBytes(data);
+ VectorSchemaRoot batch = reader.readRowGroup(0, allocator)) {
+ assertEquals(1, batch.getRowCount());
+ assertEquals(7, ((IntVector) batch.getVector("id")).get(0));
+ }
+ }
+
+ @Test
+ public void testCrossRootPreflightValidationFailureCanRetryWithoutLeak() {
+ Schema arrowSchema = new Schema(Arrays.asList(
+ Field.notNullable("id", new ArrowType.Int(32, true))
+ ));
+
+ byte[] data;
+ try (RootAllocator writerRoot = new RootAllocator(16L * 1024 * 1024);
+ RootAllocator inputRoot = new RootAllocator(16L * 1024 * 1024);
+ BufferAllocator writerAllocator =
+ writerRoot.newChildAllocator("writer", 0, 16L * 1024 *
1024);
+ BufferAllocator inputAllocator =
+ inputRoot.newChildAllocator("input", 0, 16L * 1024 *
1024);
+ VectorSchemaRoot root =
+ VectorSchemaRoot.create(arrowSchema, inputAllocator)) {
+ IntVector ids = (IntVector) root.getVector("id");
+ ids.allocateNew(1);
+ root.setRowCount(1);
+
+ long inputBytes = inputAllocator.getAllocatedMemory();
+ List<ArrowBuf> fieldBuffers = ids.getFieldBuffers();
+ int[] refCounts = new int[fieldBuffers.size()];
+ for (int i = 0; i < fieldBuffers.size(); i++) {
+ refCounts[i] = fieldBuffers.get(i).refCnt();
+ }
+
+ ByteArrayOutputStream output = new ByteArrayOutputStream();
+ try (MosaicWriter writer =
+ new MosaicWriter(output, arrowSchema, writerAllocator)) {
+ long writerBytes = writerAllocator.getAllocatedMemory();
+ RuntimeException error =
+ assertThrows(RuntimeException.class, () ->
writer.write(root));
+ assertTrue(error.getMessage().contains("non-nullable column
'id' has 1 nulls"));
+ assertEquals(inputBytes, inputAllocator.getAllocatedMemory());
+ assertEquals(writerBytes,
writerAllocator.getAllocatedMemory());
+ for (int i = 0; i < fieldBuffers.size(); i++) {
+ assertEquals(refCounts[i], fieldBuffers.get(i).refCnt());
+ }
+
+ ids.set(0, 7);
+ writer.write(root);
+ assertEquals(inputBytes, inputAllocator.getAllocatedMemory());
+ assertEquals(writerBytes,
writerAllocator.getAllocatedMemory());
+ for (int i = 0; i < fieldBuffers.size(); i++) {
+ assertEquals(refCounts[i], fieldBuffers.get(i).refCnt());
+ }
+ }
+ data = output.toByteArray();
+ }
+
+ int totalRows = 0;
+ try (MosaicReader reader = readerFromBytes(data)) {
+ for (int rg = 0; rg < reader.numRowGroups(); rg++) {
+ try (VectorSchemaRoot batch = reader.readRowGroup(rg,
allocator)) {
+ totalRows += batch.getRowCount();
+ assertEquals(7, ((IntVector)
batch.getVector("id")).get(0));
+ }
+ }
+ }
+ assertEquals(1, totalRows);
+ }
+
+ @Test
+ public void testOutputWriteFailureAbortsWriterAndPreservesThrowable() {
+ Schema arrowSchema = new Schema(Arrays.asList(
+ Field.notNullable("id", new ArrowType.Int(32, true))
+ ));
+ WriterOptions options =
+ new WriterOptions()
+ .compression(0)
+ .numBuckets(1)
+ .rowGroupMaxSize(1);
+
+ try (RootAllocator writerRoot = new RootAllocator(16L * 1024 * 1024);
+ RootAllocator inputRoot = new RootAllocator(16L * 1024 * 1024)) {
+ try (BufferAllocator writerAllocator =
+ writerRoot.newChildAllocator("writer", 0, 16L * 1024
* 1024);
+ BufferAllocator inputAllocator =
+ inputRoot.newChildAllocator("input", 0, 16L * 1024 *
1024);
+ VectorSchemaRoot root =
+ VectorSchemaRoot.create(arrowSchema, inputAllocator))
{
+ IntVector ids = (IntVector) root.getVector("id");
+ ids.allocateNew(1);
+ ids.set(0, 7);
+ root.setRowCount(1);
+
+ long inputBytes = inputAllocator.getAllocatedMemory();
+ long writerBytes = writerAllocator.getAllocatedMemory();
+ FailOnceOutputStream output =
+ new FailOnceOutputStream(
+ FailurePoint.WRITE, "sentinel-output-write");
+ MosaicWriter writer =
+ new MosaicWriter(output, arrowSchema, options,
writerAllocator);
+
+ RuntimeException error =
+ assertThrows(RuntimeException.class, () ->
writer.write(root));
+ assertEquals("write batch failed", error.getMessage());
+ assertSame(output.failure, error.getCause());
+ assertEquals("sentinel-output-write",
error.getCause().getMessage());
+ assertSame(output.failureCause, error.getCause().getCause());
+ assertEquals(inputBytes, inputAllocator.getAllocatedMemory());
+ assertEquals(writerBytes,
writerAllocator.getAllocatedMemory());
+ assertEquals(1, output.writeCalls);
+ assertEquals(0, output.size());
+
+ RuntimeException retryError =
+ assertThrows(RuntimeException.class, () ->
writer.write(root));
+ assertTrue(retryError
+ .getMessage()
+ .contains("writer is aborted after a previous
failure"));
+ assertEquals(1, output.writeCalls);
+
+ RuntimeException closeError =
+ assertThrows(RuntimeException.class, writer::close);
+ assertTrue(closeError
+ .getMessage()
+ .contains("writer is aborted after a previous
failure"));
+ assertEquals(1, output.writeCalls);
+ assertEquals(0, output.flushCalls);
+ assertEquals(0, output.size());
+ assertEquals(inputBytes, inputAllocator.getAllocatedMemory());
+ assertEquals(writerBytes,
writerAllocator.getAllocatedMemory());
+ }
+ assertEquals(0, inputRoot.getAllocatedMemory());
+ assertEquals(0, writerRoot.getAllocatedMemory());
+ }
+ }
+
+ @Test
+ public void testOutputFlushFailurePreservesThrowableWithoutFreeRetry() {
+ Schema arrowSchema = new Schema(Arrays.asList(
+ Field.notNullable("id", new ArrowType.Int(32, true))
+ ));
+ WriterOptions options =
+ new WriterOptions()
+ .compression(0)
+ .numBuckets(1)
+ .rowGroupMaxSize(1);
+
+ try (VectorSchemaRoot root = VectorSchemaRoot.create(arrowSchema,
allocator)) {
+ IntVector ids = (IntVector) root.getVector("id");
+ ids.allocateNew(1);
+ ids.set(0, 7);
+ root.setRowCount(1);
+
+ FailOnceOutputStream output =
+ new FailOnceOutputStream(
+ FailurePoint.FLUSH, "sentinel-output-flush");
+ MosaicWriter writer =
+ new MosaicWriter(output, arrowSchema, options, allocator);
+ writer.write(root);
+ int writesBeforeClose = output.writeCalls;
+
+ RuntimeException error = assertThrows(RuntimeException.class,
writer::close);
+ assertEquals("close failed", error.getMessage());
+ assertSame(output.failure, error.getCause());
+ assertEquals("sentinel-output-flush",
error.getCause().getMessage());
+ assertSame(output.failureCause, error.getCause().getCause());
+ assertTrue(output.writeCalls > writesBeforeClose);
+ assertEquals(1, output.flushCalls);
+ assertTrue(output.size() > 0);
+
+ int writesAfterClose = output.writeCalls;
+ writer.close();
+ assertEquals(writesAfterClose, output.writeCalls);
+ assertEquals(1, output.flushCalls);
+ }
+ }
+
+ @Test
+ public void testCrossRootPartialExportFailureCanRetryWithoutLeak() {
+ byte[] data;
+ try (RootAllocator writerRoot = new RootAllocator(16L * 1024 * 1024);
+ RootAllocator inputRoot = new RootAllocator(16L * 1024 * 1024)) {
+ try (BufferAllocator writerAllocator =
+ writerRoot.newChildAllocator("writer", 0, 16L * 1024
* 1024);
+ BufferAllocator inputAllocator =
+ inputRoot.newChildAllocator("input", 0, 16L * 1024 *
1024)) {
+ IntVector first = new IntVector("first", inputAllocator);
+ FailOnceListVector second = new FailOnceListVector("second",
inputAllocator);
+ try (VectorSchemaRoot root = VectorSchemaRoot.of(first,
second)) {
+ first.allocateNew(2);
+ second.allocateNew();
+ first.set(0, 1);
+ first.set(1, 2);
+ UnionListWriter listWriter = second.getWriter();
+ listWriter.setPosition(0);
+ listWriter.startList();
+ listWriter.writeInt(3);
+ listWriter.endList();
+ listWriter.setPosition(1);
+ listWriter.startList();
+ listWriter.writeInt(4);
+ listWriter.endList();
+ root.setRowCount(2);
+
+ long inputBytes = inputRoot.getAllocatedMemory();
+ ByteArrayOutputStream output = new ByteArrayOutputStream();
+ try (MosaicWriter writer =
+ new MosaicWriter(output, root.getSchema(),
writerAllocator)) {
+ long writerBytes =
writerAllocator.getAllocatedMemory();
+ assertThrows(InjectedExportException.class, () ->
writer.write(root));
+ assertEquals(inputBytes,
inputRoot.getAllocatedMemory());
+ assertEquals(inputBytes,
inputAllocator.getAllocatedMemory());
+ assertEquals(writerBytes,
writerAllocator.getAllocatedMemory());
+ assertEquals(2, first.get(1));
+ assertEquals("[4]", second.getObject(1).toString());
+
+ writer.write(root);
+ }
+ data = output.toByteArray();
+ }
+ }
+ assertEquals(0, inputRoot.getAllocatedMemory());
+ assertEquals(0, writerRoot.getAllocatedMemory());
+ }
+
+ try (MosaicReader reader = readerFromBytes(data);
+ VectorSchemaRoot batch = reader.readRowGroup(0, allocator)) {
+ assertEquals(2, batch.getRowCount());
+ assertEquals(1, ((IntVector) batch.getVector("first")).get(0));
+ assertEquals(2, ((IntVector) batch.getVector("first")).get(1));
+ assertEquals("[3]",
batch.getVector("second").getObject(0).toString());
+ assertEquals("[4]",
batch.getVector("second").getObject(1).toString());
+ }
+ }
+
+ @Test
+ public void testWriteSequentialBundlesFromDifferentRootAllocators() {
+ Schema arrowSchema = new Schema(Arrays.asList(
+ Field.notNullable("id", new ArrowType.Int(32, true))
+ ));
+
+ byte[] data = writeToBytes(arrowSchema, writer -> {
+ for (int batch = 0; batch < 2; batch++) {
+ try (BufferAllocator inputAllocator = new RootAllocator();
+ VectorSchemaRoot root =
+ VectorSchemaRoot.create(arrowSchema,
inputAllocator)) {
+ IntVector ids = (IntVector) root.getVector("id");
+ ids.allocateNew(2);
+ ids.set(0, batch * 2);
+ ids.set(1, batch * 2 + 1);
+ root.setRowCount(2);
+ writer.write(root);
+ }
+ }
+ });
+
+ boolean[] seen = new boolean[4];
+ int totalRows = 0;
+ try (MosaicReader reader = readerFromBytes(data)) {
+ for (int rg = 0; rg < reader.numRowGroups(); rg++) {
+ try (VectorSchemaRoot batch = reader.readRowGroup(rg,
allocator)) {
+ IntVector ids = (IntVector) batch.getVector("id");
+ for (int i = 0; i < batch.getRowCount(); i++) {
+ seen[ids.get(i)] = true;
+ totalRows++;
+ }
+ }
+ }
+ }
+ assertEquals(4, totalRows);
+ assertArrayEquals(new boolean[]{true, true, true, true}, seen);
+ }
+
+ @Test
+ public void testRejectsFieldVectorsFromDifferentAllocatorRoots() {
+ Schema arrowSchema = new Schema(Arrays.asList(
+ Field.notNullable("id", new ArrowType.Int(32, true)),
+ Field.nullable("name", ArrowType.Utf8.INSTANCE)
+ ));
+
+ try (BufferAllocator idAllocator = new RootAllocator();
+ BufferAllocator nameAllocator = new RootAllocator()) {
+ IntVector ids = new IntVector("id", idAllocator);
+ VarCharVector names = new VarCharVector("name", nameAllocator);
+ try (VectorSchemaRoot root = VectorSchemaRoot.of(ids, names)) {
+ ids.allocateNew(1);
+ names.allocateNew(1);
+ ids.set(0, 1);
+ names.setSafe(0, "one".getBytes());
+ root.setRowCount(1);
+
+ IllegalArgumentException error =
+ assertThrows(
+ IllegalArgumentException.class,
+ () -> writeToBytes(arrowSchema, writer ->
writer.write(root)));
+ assertTrue(error.getMessage().contains("same allocator root"));
+ assertTrue(error.getMessage().contains("name"));
+ }
+ }
+ }
+
+ @Test
+ public void
testRejectsNestedFieldVectorFromDifferentAllocatorRootWithoutLeak() {
+ try (RootAllocator parentAllocator = new RootAllocator(16L * 1024 *
1024);
+ RootAllocator nestedAllocator = new RootAllocator(16L * 1024 *
1024)) {
+ InjectableListVector list = new InjectableListVector("items",
parentAllocator);
+ IntVector data = new IntVector(ListVector.DATA_VECTOR_NAME,
nestedAllocator);
+ list.setDataVector(data);
+
+ try (VectorSchemaRoot root = VectorSchemaRoot.of(list)) {
+ list.allocateNew();
+ list.startNewValue(0);
+ data.set(0, 7);
+ list.endValue(0, 1);
+ list.setValueCount(1);
+ root.setRowCount(1);
+
+ long parentBefore = parentAllocator.getAllocatedMemory();
+ long nestedBefore = nestedAllocator.getAllocatedMemory();
+
+ IllegalArgumentException error =
+ assertThrows(
+ IllegalArgumentException.class,
+ () -> writeToBytes(root.getSchema(), writer ->
writer.write(root)));
+ assertTrue(error.getMessage().contains("same allocator root"));
+ assertTrue(error.getMessage().contains("items." +
ListVector.DATA_VECTOR_NAME));
+ assertEquals(parentBefore,
parentAllocator.getAllocatedMemory());
+ assertEquals(nestedBefore,
nestedAllocator.getAllocatedMemory());
+ } finally {
+ data.close();
+ }
+
+ assertEquals(0, parentAllocator.getAllocatedMemory());
+ assertEquals(0, nestedAllocator.getAllocatedMemory());
+ }
+ }
+
+ @Test
+ public void testWriterOpenFailurePreservesNativeMessage() {
+ Schema arrowSchema = new Schema(Arrays.asList(
+ Field.notNullable("id", new ArrowType.Int(32, true))
+ ));
+ WriterOptions options = new WriterOptions().statsColumns("missing");
+
+ RuntimeException error =
+ assertThrows(
+ RuntimeException.class,
+ () ->
+ new MosaicWriter(
+ new ByteArrayOutputStream(),
+ arrowSchema,
+ options,
+ allocator));
+
+ assertTrue(
+ error.getMessage(),
+ error.getMessage()
+ .contains(
+ "writer open failed: stats_columns: column
'missing' not found in schema"));
+ }
+
@Test
public void testNullValues() {
Schema arrowSchema = new Schema(Arrays.asList(
diff --git a/jni/src/lib.rs b/jni/src/lib.rs
index 98fe7c4..ac84808 100644
--- a/jni/src/lib.rs
+++ b/jni/src/lib.rs
@@ -20,8 +20,9 @@ use std::panic::{self, AssertUnwindSafe};
use std::ptr;
use std::sync::Arc;
+use jni::errors::Error as JniError;
use jni::objects::{
- GlobalRef, JByteArray, JClass, JMethodID, JObject, JObjectArray, JString,
JValue,
+ GlobalRef, JByteArray, JClass, JMethodID, JObject, JObjectArray, JString,
JThrowable, JValue,
};
use jni::sys::{jint, jlong, jlongArray};
use jni::JNIEnv;
@@ -53,14 +54,61 @@ struct JniOutputFile {
pos: u64,
cached_array: Option<GlobalRef>,
cached_array_len: usize,
+ pending_exception: Option<GlobalRef>,
}
unsafe impl Send for JniOutputFile {}
+impl JniOutputFile {
+ fn record_jni_error(&mut self, env: &mut JNIEnv, error: JniError) ->
io::Error {
+ if !matches!(error, JniError::JavaException) {
+ return io::Error::other(error.to_string());
+ }
+
+ let captured = (|| -> jni::errors::Result<Option<GlobalRef>> {
+ let exception = env.exception_occurred()?;
+ env.exception_clear()?;
+ if exception.is_null() || self.pending_exception.is_some() {
+ return Ok(None);
+ }
+ let global = env.new_global_ref(exception)?;
+ let exception_pending = env.exception_check()?;
+ if exception_pending {
+ env.exception_clear()?;
+ }
+ if global.as_obj().is_null() || exception_pending {
+ return Err(JniError::NullPtr("NewGlobalRef for pending Java
exception"));
+ }
+ Ok(Some(global))
+ })();
+
+ match captured {
+ Ok(Some(exception)) => {
+ self.pending_exception = Some(exception);
+ io::Error::other(error.to_string())
+ }
+ Ok(None) => io::Error::other(error.to_string()),
+ Err(capture_error) => {
+ // Cleanup must run without a pending Java exception. If
preserving the original
+ // throwable itself fails (for example due to OOM), report
both JNI failures.
+ let _ = env.exception_clear();
+ io::Error::other(format!(
+ "{} (failed to preserve Java exception: {})",
+ error, capture_error
+ ))
+ }
+ }
+ }
+
+ fn take_pending_exception(&mut self) -> Option<GlobalRef> {
+ self.pending_exception.take()
+ }
+}
+
impl OutputFile for JniOutputFile {
fn write(&mut self, data: &[u8]) -> io::Result<()> {
- let mut env = self
- .jvm
+ let jvm = Arc::clone(&self.jvm);
+ let mut env = jvm
.attach_current_thread()
.map_err(|e| io::Error::other(e.to_string()))?;
@@ -72,12 +120,26 @@ impl OutputFile for JniOutputFile {
};
if need_new {
- let byte_array = env
- .new_byte_array(len)
- .map_err(|e| io::Error::other(e.to_string()))?;
- let global = env
- .new_global_ref(&byte_array)
- .map_err(|e| io::Error::other(e.to_string()))?;
+ let byte_array = match env.new_byte_array(len) {
+ Ok(array) => array,
+ Err(error) => return Err(self.record_jni_error(&mut env,
error)),
+ };
+ let global = match env.new_global_ref(&byte_array) {
+ Ok(global) => global,
+ Err(error) => return Err(self.record_jni_error(&mut env,
error)),
+ };
+ match env.exception_check() {
+ Ok(true) => {
+ return Err(self.record_jni_error(&mut env,
JniError::JavaException));
+ }
+ Ok(false) => {}
+ Err(error) => return Err(io::Error::other(error.to_string())),
+ }
+ if global.as_obj().is_null() {
+ return Err(io::Error::other(
+ "failed to create global reference for output buffer",
+ ));
+ }
self.cached_array = Some(global);
self.cached_array_len = data.len();
}
@@ -85,10 +147,11 @@ impl OutputFile for JniOutputFile {
let raw = self.cached_array.as_ref().unwrap().as_raw();
let byte_array = unsafe { JByteArray::from_raw(raw) };
- env.set_byte_array_region(&byte_array, 0, bytemuck_cast(data))
- .map_err(|e| io::Error::other(e.to_string()))?;
+ if let Err(error) = env.set_byte_array_region(&byte_array, 0,
bytemuck_cast(data)) {
+ return Err(self.record_jni_error(&mut env, error));
+ }
- unsafe {
+ let call_result = unsafe {
env.call_method_unchecked(
&self.stream_ref,
self.write_mid,
@@ -99,7 +162,9 @@ impl OutputFile for JniOutputFile {
jni::sys::jvalue { i: len },
],
)
- .map_err(|e| io::Error::other(e.to_string()))?;
+ };
+ if let Err(error) = call_result {
+ return Err(self.record_jni_error(&mut env, error));
}
#[allow(clippy::forget_non_drop)]
std::mem::forget(byte_array);
@@ -108,18 +173,20 @@ impl OutputFile for JniOutputFile {
}
fn flush(&mut self) -> io::Result<()> {
- let mut env = self
- .jvm
+ let jvm = Arc::clone(&self.jvm);
+ let mut env = jvm
.attach_current_thread()
.map_err(|e| io::Error::other(e.to_string()))?;
- unsafe {
+ let call_result = unsafe {
env.call_method_unchecked(
&self.stream_ref,
self.flush_mid,
jni::signature::ReturnType::Primitive(jni::signature::Primitive::Void),
&[],
)
- .map_err(|e| io::Error::other(e.to_string()))?;
+ };
+ if let Err(error) = call_result {
+ return Err(self.record_jni_error(&mut env, error));
}
Ok(())
}
@@ -188,6 +255,35 @@ fn throw(env: &mut JNIEnv, msg: &str) {
let _ = env.throw_new("java/lang/RuntimeException", msg);
}
+fn rethrow(env: &mut JNIEnv, exception: &GlobalRef) {
+ if exception.as_obj().is_null() {
+ throw(env, "cannot rethrow a null Java exception reference");
+ return;
+ }
+ match env.new_local_ref(exception.as_obj()) {
+ Ok(local) if !local.is_null() => {
+ if let Err(error) = env.throw(JThrowable::from(local)) {
+ let _ = env.exception_clear();
+ throw(
+ env,
+ &format!("failed to rethrow preserved Java exception: {}",
error),
+ );
+ }
+ }
+ Ok(_) => {
+ let _ = env.exception_clear();
+ throw(env, "failed to create local reference for Java exception");
+ }
+ Err(error) => {
+ let _ = env.exception_clear();
+ throw(
+ env,
+ &format!("failed to rethrow preserved Java exception: {}",
error),
+ );
+ }
+ }
+}
+
struct WriterHandle {
inner: MosaicWriter<JniOutputFile>,
_stream_ref: GlobalRef,
@@ -210,53 +306,41 @@ pub extern "system" fn
Java_org_apache_paimon_mosaic_NativeLib_nativeWriterOpen(
stats_columns: JObjectArray<'_>,
page_size_threshold: jint,
) -> jlong {
- let raw_env = env.get_raw();
- let result = panic::catch_unwind(AssertUnwindSafe(|| {
+ let result = panic::catch_unwind(AssertUnwindSafe(|| -> Result<jlong,
String> {
if arrow_schema_addr == 0 {
- throw(&mut env, "null Arrow schema address");
- return 0;
+ return Err("null Arrow schema address".to_string());
}
let ffi_schema =
unsafe { FFI_ArrowSchema::from_raw(arrow_schema_addr as *mut
FFI_ArrowSchema) };
- let arrow_schema = match Schema::try_from(&ffi_schema) {
- Ok(s) => s,
- Err(e) => {
- throw(&mut env, &format!("Arrow schema import failed: {}", e));
- return 0;
- }
- };
-
- let stream_global = match env.new_global_ref(&stream) {
- Ok(g) => g,
- Err(e) => {
- throw(&mut env, &format!("failed to create global ref: {}",
e));
- return 0;
- }
- };
+ let arrow_schema = Schema::try_from(&ffi_schema)
+ .map_err(|e| format!("Arrow schema import failed: {}", e))?;
+ drop(ffi_schema);
+
+ let stream_global = env
+ .new_global_ref(&stream)
+ .map_err(|e| format!("failed to create global ref: {}", e))?;
+ if env
+ .exception_check()
+ .map_err(|e| format!("failed to check global ref exception: {}",
e))?
+ {
+ return Err("failed to create global ref: Java exception was
thrown".to_string());
+ }
+ if stream_global.as_obj().is_null() {
+ return Err("failed to create global ref: NewGlobalRef returned
null".to_string());
+ }
- let write_mid = match env.get_method_id("java/io/OutputStream",
"write", "([BII)V") {
- Ok(m) => m,
- Err(e) => {
- throw(&mut env, &format!("cannot find OutputStream.write: {}",
e));
- return 0;
- }
- };
- let flush_mid = match env.get_method_id("java/io/OutputStream",
"flush", "()V") {
- Ok(m) => m,
- Err(e) => {
- throw(&mut env, &format!("cannot find OutputStream.flush: {}",
e));
- return 0;
- }
- };
+ let write_mid = env
+ .get_method_id("java/io/OutputStream", "write", "([BII)V")
+ .map_err(|e| format!("cannot find OutputStream.write: {}", e))?;
+ let flush_mid = env
+ .get_method_id("java/io/OutputStream", "flush", "()V")
+ .map_err(|e| format!("cannot find OutputStream.flush: {}", e))?;
- let jvm = match env.get_java_vm() {
- Ok(vm) => Arc::new(vm),
- Err(e) => {
- throw(&mut env, &format!("cannot get JavaVM: {}", e));
- return 0;
- }
- };
+ let jvm = Arc::new(
+ env.get_java_vm()
+ .map_err(|e| format!("cannot get JavaVM: {}", e))?,
+ );
let jni_stream = JniOutputFile {
jvm,
@@ -266,35 +350,30 @@ pub extern "system" fn
Java_org_apache_paimon_mosaic_NativeLib_nativeWriterOpen(
pos: 0,
cached_array: None,
cached_array_len: 0,
+ pending_exception: None,
};
- let stats_cols: Vec<String> = match
env.get_array_length(&stats_columns) {
- Ok(len) if len > 0 => {
- let mut names = Vec::with_capacity(len as usize);
- for i in 0..len {
- let obj = match
env.get_object_array_element(&stats_columns, i) {
- Ok(o) => o,
- Err(_) => {
- throw(&mut env, "failed to read stats_columns
element");
- return 0;
- }
- };
- let jstr = JString::from(obj);
- let s: String = match env.get_string(&jstr) {
- Ok(s) => s.into(),
- Err(_) => {
- throw(
- &mut env,
- "failed to convert stats_columns element to
string",
- );
- return 0;
- }
- };
- names.push(s);
- }
- names
+ let stats_len = env
+ .get_array_length(&stats_columns)
+ .map_err(|e| format!("failed to read stats_columns length: {}",
e))?;
+ let stats_cols: Vec<String> = if stats_len > 0 {
+ let mut names = Vec::with_capacity(stats_len as usize);
+ for i in 0..stats_len {
+ let obj = env
+ .get_object_array_element(&stats_columns, i)
+ .map_err(|e| format!("failed to read stats_columns
element: {}", e))?;
+ let jstr = JString::from(obj);
+ let s: String = env
+ .get_string(&jstr)
+ .map_err(|e| {
+ format!("failed to convert stats_columns element to
string: {}", e)
+ })?
+ .into();
+ names.push(s);
}
- _ => Vec::new(),
+ names
+ } else {
+ Vec::new()
};
let buckets = if num_buckets <= 0 {
@@ -314,27 +393,28 @@ pub extern "system" fn
Java_org_apache_paimon_mosaic_NativeLib_nativeWriterOpen(
page_size_threshold: page_size_threshold as usize,
};
- let writer = match MosaicWriter::new(jni_stream, &arrow_schema, opts) {
- Ok(w) => w,
- Err(e) => {
- throw(&mut env, &format!("writer open failed: {}", e));
- return 0;
- }
- };
+ let writer = MosaicWriter::new(jni_stream, &arrow_schema, opts)
+ .map_err(|e| format!("writer open failed: {}", e))?;
let handle = Box::new(WriterHandle {
inner: writer,
_stream_ref: stream_global,
});
- Box::into_raw(handle) as jlong
+ Ok(Box::into_raw(handle) as jlong)
}));
- match result {
- Ok(val) => val,
- Err(e) => {
- let mut env = unsafe { JNIEnv::from_raw(raw_env).unwrap() };
- throw(&mut env, &panic_message(&e));
- 0
- }
+
+ let error = match result {
+ Ok(Ok(handle)) => return handle,
+ Ok(Err(error)) => error,
+ Err(error) => panic_message(&error),
+ };
+ // A failing JNI call can leave its original Java throwable pending. The
imported Arrow schema
+ // was released before any such call, so let that throwable propagate
instead of replacing it.
+ if env.exception_check().unwrap_or(false) {
+ return 0;
}
+ // Defer throwing until Rust-owned resources above have been dropped.
+ throw(&mut env, &error);
+ 0
}
#[no_mangle]
@@ -343,19 +423,32 @@ pub extern "system" fn
Java_org_apache_paimon_mosaic_NativeLib_nativeWriterClose
_class: JClass,
handle: jlong,
) {
- let raw_env = env.get_raw();
- let result = panic::catch_unwind(AssertUnwindSafe(|| {
+ let result = panic::catch_unwind(AssertUnwindSafe(|| -> Result<(), String>
{
if handle == 0 {
- return;
+ return Ok(());
}
let writer = unsafe { &mut *(handle as *mut WriterHandle) };
- if let Err(e) = writer.inner.close() {
- throw(&mut env, &format!("close failed: {}", e));
- }
+ writer
+ .inner
+ .close()
+ .map_err(|e| format!("close failed: {}", e))
}));
- if let Err(e) = result {
- let mut env = unsafe { JNIEnv::from_raw(raw_env).unwrap() };
- throw(&mut env, &panic_message(&e));
+
+ let pending_exception = if handle == 0 {
+ None
+ } else {
+ let writer = unsafe { &mut *(handle as *mut WriterHandle) };
+ writer.inner.output_mut().take_pending_exception()
+ };
+ if let Some(exception) = pending_exception {
+ rethrow(&mut env, &exception);
+ return;
+ }
+
+ match result {
+ Ok(Ok(())) => {}
+ Ok(Err(error)) => throw(&mut env, &error),
+ Err(error) => throw(&mut env, &panic_message(&error)),
}
}
@@ -537,15 +630,12 @@ pub extern "system" fn
Java_org_apache_paimon_mosaic_NativeLib_nativeWriterWrite
array_addr: jlong,
schema_addr: jlong,
) {
- let raw_env = env.get_raw();
- let result = panic::catch_unwind(AssertUnwindSafe(|| {
+ let result = panic::catch_unwind(AssertUnwindSafe(|| -> Result<(), String>
{
if writer_handle == 0 {
- throw(&mut env, "null writer handle");
- return;
+ return Err("null writer handle".to_string());
}
if array_addr == 0 || schema_addr == 0 {
- throw(&mut env, "null ArrowArray or ArrowSchema address");
- return;
+ return Err("null ArrowArray or ArrowSchema address".to_string());
}
let writer = unsafe { &mut *(writer_handle as *mut WriterHandle) };
@@ -554,24 +644,36 @@ pub extern "system" fn
Java_org_apache_paimon_mosaic_NativeLib_nativeWriterWrite
let arr_owned = unsafe { FFI_ArrowArray::from_raw(ffi_array) };
let schema_owned = unsafe { FFI_ArrowSchema::from_raw(ffi_schema) };
- let arr_data = match unsafe { arrow_array::ffi::from_ffi(arr_owned,
&schema_owned) } {
- Ok(d) => d,
- Err(e) => {
- throw(&mut env, &format!("Arrow import failed: {}", e));
- return;
- }
- };
+ let arr_data = unsafe { arrow_array::ffi::from_ffi(arr_owned,
&schema_owned) }
+ .map_err(|e| format!("Arrow import failed: {}", e))?;
let struct_array = StructArray::from(arr_data);
let batch = RecordBatch::from(struct_array);
- if let Err(e) = writer.inner.write_batch(&batch) {
- throw(&mut env, &format!("write_batch failed: {}", e));
- }
+ writer
+ .inner
+ .write_batch(&batch)
+ .map_err(|e| format!("write_batch failed: {}", e))
}));
- if let Err(e) = result {
- let mut env = unsafe { JNIEnv::from_raw(raw_env).unwrap() };
- throw(&mut env, &panic_message(&e));
+
+ let pending_exception = if writer_handle == 0 {
+ None
+ } else {
+ let writer = unsafe { &mut *(writer_handle as *mut WriterHandle) };
+ writer.inner.output_mut().take_pending_exception()
+ };
+ if let Some(exception) = pending_exception {
+ rethrow(&mut env, &exception);
+ return;
}
+
+ let error = match result {
+ Ok(Ok(())) => return,
+ Ok(Err(error)) => error,
+ Err(error) => panic_message(&error),
+ };
+ // Arrow's Java release callbacks may clear a pending JNI exception. Defer
throwing until all
+ // Rust-owned Arrow C Data objects above have been dropped and their
callbacks have completed.
+ throw(&mut env, &error);
}
// ======================== Reader ========================
diff --git a/python/mosaic/mosaic.py b/python/mosaic/mosaic.py
index b08096d..4796890 100644
--- a/python/mosaic/mosaic.py
+++ b/python/mosaic/mosaic.py
@@ -59,6 +59,60 @@ def _check_error(msg="operation failed"):
raise RuntimeError(msg)
+def _visible_exception_chain(root):
+ chain = []
+ seen = set()
+ current = root
+ while current is not None:
+ if id(current) in seen:
+ return chain, True
+ seen.add(id(current))
+ chain.append(current)
+ if current.__cause__ is not None:
+ current = current.__cause__
+ elif not current.__suppress_context__:
+ current = current.__context__
+ else:
+ current = None
+ return chain, False
+
+
+def _remove_exception_references(root, target_ids):
+ stack = [root]
+ seen = set()
+ while stack:
+ current = stack.pop()
+ if id(current) in seen:
+ continue
+ seen.add(id(current))
+ for attribute in ("__cause__", "__context__"):
+ linked = getattr(current, attribute)
+ if linked is not None and id(linked) in target_ids:
+ setattr(current, attribute, None)
+ elif linked is not None:
+ stack.append(linked)
+
+
+def _append_cleanup_exception(primary, secondary):
+ """Append a cleanup failure without replacing the primary exception."""
+ if primary is secondary:
+ return
+
+ primary_chain, has_cycle = _visible_exception_chain(primary)
+ if has_cycle or any(current is secondary for current in primary_chain):
+ return
+
+ # ``secondary`` was raised while ``primary`` was being handled, so Python
normally links it
+ # back to the primary chain. Remove every such edge, including shared
causes, before appending.
+ _remove_exception_references(secondary, {id(current) for current in
primary_chain})
+
+ tail = primary_chain[-1]
+ if tail.__suppress_context__:
+ tail.__cause__ = secondary
+ else:
+ tail.__context__ = secondary
+
+
def _fetch_rg_stats(num_stats_fn, stats_fn, handle, rg_index):
n_out = ctypes.c_uint32(0)
rc = num_stats_fn(handle, rg_index, ctypes.byref(n_out))
@@ -255,8 +309,15 @@ class MosaicWriter:
def __enter__(self):
return self
- def __exit__(self, *args):
- self.close()
+ def __exit__(self, exc_type, exc_value, _traceback):
+ try:
+ self.close()
+ except Exception as close_error:
+ if exc_type is None:
+ raise
+ # Keep the body exception primary while retaining the independent
+ # close failure for diagnostics.
+ _append_cleanup_exception(exc_value, close_error)
def __del__(self):
self.close()
diff --git a/python/tests/test_mosaic.py b/python/tests/test_mosaic.py
index 796e880..f88720c 100644
--- a/python/tests/test_mosaic.py
+++ b/python/tests/test_mosaic.py
@@ -17,6 +17,7 @@
import io
import struct
+import traceback
import pyarrow as pa
import pytest
@@ -29,6 +30,7 @@ from mosaic import (
read_table,
write_table,
)
+from mosaic.mosaic import _append_cleanup_exception
def _write_to_bytes(pa_schema, data, options=None):
@@ -781,6 +783,112 @@ class TestWriter:
with pytest.raises(RuntimeError, match="writer is closed"):
writer.write(batch)
+ def test_context_manager_preserves_write_failure(self):
+ class FailOnceOutput(io.BytesIO):
+ def __init__(self):
+ super().__init__()
+ self.fail = True
+
+ def write(self, data):
+ if self.fail:
+ self.fail = False
+ raise OSError("sentinel output failure")
+ return super().write(data)
+
+ pa_schema = pa.schema([pa.field("x", pa.int32(), nullable=False)])
+ batch = pa.record_batch(
+ [pa.array([1], type=pa.int32())], schema=pa_schema
+ )
+ options = WriterOptions(
+ compression=WriterOptions.COMPRESSION_NONE,
+ num_buckets=1,
+ row_group_max_size=1,
+ )
+
+ with pytest.raises(
+ RuntimeError, match="write_batch failed: write callback failed"
+ ) as exc_info:
+ with MosaicWriter(FailOnceOutput(), pa_schema, options) as writer:
+ writer.write(batch)
+ assert writer._closed
+ assert writer._handle is None
+ assert isinstance(exc_info.value.__context__, RuntimeError)
+ assert "writer is aborted after a previous failure" in str(
+ exc_info.value.__context__
+ )
+
+ def test_context_manager_preserves_body_and_close_failures(self):
+ class FlushFailOutput(io.BytesIO):
+ def flush(self):
+ raise OSError("sentinel flush failure")
+
+ pa_schema = pa.schema([pa.field("x", pa.int32(), nullable=False)])
+ batch = pa.record_batch(
+ [pa.array([1], type=pa.int32())], schema=pa_schema
+ )
+ options = WriterOptions(
+ compression=WriterOptions.COMPRESSION_NONE,
+ num_buckets=1,
+ )
+
+ with pytest.raises(ValueError, match="body failure") as exc_info:
+ with MosaicWriter(FlushFailOutput(), pa_schema, options) as writer:
+ writer.write(batch)
+ raise ValueError("body failure")
+
+ assert writer._closed
+ assert writer._handle is None
+ assert isinstance(exc_info.value.__context__, RuntimeError)
+ assert "close failed: flush callback failed" in str(
+ exc_info.value.__context__
+ )
+
+ @pytest.mark.parametrize("cause", [KeyError("explicit cause"), None])
+ def
test_context_manager_keeps_close_failure_visible_with_suppressed_context(
+ self, cause
+ ):
+ class FlushFailOutput(io.BytesIO):
+ def flush(self):
+ raise OSError("sentinel flush failure")
+
+ pa_schema = pa.schema([pa.field("x", pa.int32(), nullable=False)])
+ batch = pa.record_batch(
+ [pa.array([1], type=pa.int32())], schema=pa_schema
+ )
+
+ try:
+ raise LookupError("hidden body context")
+ except LookupError:
+ with pytest.raises(ValueError, match="body failure") as exc_info:
+ with MosaicWriter(FlushFailOutput(), pa_schema) as writer:
+ writer.write(batch)
+ raise ValueError("body failure") from cause
+
+ if cause is not None:
+ assert exc_info.value.__cause__ is cause
+ rendered = "".join(
+ traceback.format_exception(
+ type(exc_info.value),
+ exc_info.value,
+ exc_info.value.__traceback__,
+ )
+ )
+ assert "close failed: flush callback failed" in rendered
+ assert "hidden body context" not in rendered
+
+ def test_cleanup_exception_chain_avoids_shared_cause_cycle(self):
+ shared_cause = KeyError("shared cause")
+ primary = ValueError("body failure")
+ primary.__cause__ = shared_cause
+ secondary = RuntimeError("close failure")
+ secondary.__cause__ = shared_cause
+
+ _append_cleanup_exception(primary, secondary)
+
+ assert primary.__cause__ is shared_cause
+ assert shared_cause.__context__ is secondary
+ assert secondary.__cause__ is None
+
def test_writer_stats_basic(self):
pa_schema = pa.schema(
[
