JingsongLi commented on code in PR #73:
URL: https://github.com/apache/paimon-mosaic/pull/73#discussion_r3789633890
##########
java/src/main/java/org/apache/paimon/mosaic/MosaicWriter.java:
##########
@@ -69,19 +73,153 @@ public MosaicWriter(OutputStream outputStream, Schema
arrowSchema, WriterOptions
}
}
+ /**
+ * 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);
+ }
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) {
+ // 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);
+ JniWrapper.get().exportArray(arrowArray.memoryAddress(),
privateData);
+ } catch (RuntimeException | Error failure) {
+ privateData.abort(failure);
Review Comment:
[P1] Release the registered root callback before aborting child metadata
`JniWrapper.exportArray` can return with an `OutOfMemoryError` after Arrow
15 has already installed the root `ArrowArray` release callback: its JNI code
does not check the result of `NewGlobalRef(privateData)` before assigning
`private_data` and `release`. In that state, `privateData.abort()` frees
`childPointers` and the child structs, but the outer `finally` then invokes the
still-registered root callback, whose first step is to walk `root.children`.
That dereferences freed native memory and can crash or corrupt the JVM under
memory pressure.
Please branch on root ownership here: when `arrowArray.snapshot().release !=
0`, release the root first while its child pointer table is still valid, then
call the idempotent `privateData.close()` as a fallback. Use the current manual
`abort()` path only while the root release callback is still zero. A
fault-injection test that throws after root callback registration would cover
this window.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
