This is an automated email from the ASF dual-hosted git repository.
kou pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-java.git
The following commit(s) were added to refs/heads/main by this push:
new 85d7ef30e GH-1239: Fix memory leak when C Data import hits allocator
limit mid-array (#1240)
85d7ef30e is described below
commit 85d7ef30ee5f623cf3442dd989df09c3c5754ad2
Author: Sandesh Kumar <[email protected]>
AuthorDate: Thu Aug 6 17:29:51 2026 -0700
GH-1239: Fix memory leak when C Data import hits allocator limit mid-array
(#1240)
`ReferenceCountedArrowArray.unsafeAssociateAllocation` calls `retain()`
before `wrapForeignAllocation`. If `wrapForeignAllocation` throws
(allocator over its limit), the retain is never balanced, the reference
count stays elevated, and the C Data release callback never fires —
leaking the producer's native memory.
**Fix:** call `retain()` after `wrapForeignAllocation` returns. Same
behavior on the success path; on failure the existing `finally` in
`ArrayImporter.importArray` drives the count to zero and fires the
release callback.
**Test:** `ImportOutOfMemoryTest` — exports a batch from a "producer"
allocator, tries to import it into a too-small consumer, and asserts the
producer drains to zero after the OOM. Fails on the original code,
passes with the fix.
Fixes: https://github.com/apache/arrow-java/issues/1239
---------
Signed-off-by: Sandesh Kumar <[email protected]>
Co-authored-by: Sandesh Kumar <[email protected]>
---
.../apache/arrow/c/ReferenceCountedArrowArray.java | 19 +--
.../org/apache/arrow/c/ImportOutOfMemoryTest.java | 139 +++++++++++++++++++++
2 files changed, 151 insertions(+), 7 deletions(-)
diff --git a/c/src/main/java/org/apache/arrow/c/ReferenceCountedArrowArray.java
b/c/src/main/java/org/apache/arrow/c/ReferenceCountedArrowArray.java
index cf50f9417..f51fb2510 100644
--- a/c/src/main/java/org/apache/arrow/c/ReferenceCountedArrowArray.java
+++ b/c/src/main/java/org/apache/arrow/c/ReferenceCountedArrowArray.java
@@ -64,13 +64,18 @@ final class ReferenceCountedArrowArray {
*/
ArrowBuf unsafeAssociateAllocation(
BufferAllocator trackingAllocator, long capacity, long memoryAddress) {
+ // Retain only after wrapForeignAllocation succeeds. On the
allocator-limit OOM path,
+ // wrapForeignAllocation throws before the ForeignAllocation is
associated, so release0()
+ // is not called; retaining first would leave the count elevated with no
matching release0().
+ ArrowBuf buf =
+ trackingAllocator.wrapForeignAllocation(
+ new ForeignAllocation(capacity, memoryAddress) {
+ @Override
+ protected void release0() {
+ ReferenceCountedArrowArray.this.release();
+ }
+ });
retain();
- return trackingAllocator.wrapForeignAllocation(
- new ForeignAllocation(capacity, memoryAddress) {
- @Override
- protected void release0() {
- ReferenceCountedArrowArray.this.release();
- }
- });
+ return buf;
}
}
diff --git a/c/src/test/java/org/apache/arrow/c/ImportOutOfMemoryTest.java
b/c/src/test/java/org/apache/arrow/c/ImportOutOfMemoryTest.java
new file mode 100644
index 000000000..7c099f2ef
--- /dev/null
+++ b/c/src/test/java/org/apache/arrow/c/ImportOutOfMemoryTest.java
@@ -0,0 +1,139 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.arrow.c;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+import java.util.ArrayList;
+import java.util.List;
+import org.apache.arrow.memory.BufferAllocator;
+import org.apache.arrow.memory.OutOfMemoryException;
+import org.apache.arrow.memory.RootAllocator;
+import org.apache.arrow.vector.FieldVector;
+import org.apache.arrow.vector.VarCharVector;
+import org.apache.arrow.vector.VectorSchemaRoot;
+import org.apache.arrow.vector.types.pojo.Schema;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+
+/**
+ * Regression test: a mid-import {@link OutOfMemoryException} must not leak
the imported array.
+ *
+ * <p>A "producer" allocator owns the exported batch; if the C Data release
callback fires, the
+ * producer drains to zero. A too-small consumer allocator forces an OOM
part-way through the
+ * import. The test asserts the producer drains, confirming the release
callback fired despite the
+ * failure.
+ */
+final class ImportOutOfMemoryTest {
+ private static final int ROWS = 1024;
+ private static final int VALUE_BYTES = 256;
+ private static final int COLUMNS = 4;
+ // Far smaller than the exported batch, so the import OOMs part-way through
the buffers.
+ private static final long TINY_LIMIT = 16 * 1024;
+
+ private RootAllocator root;
+
+ @BeforeEach
+ public void setUp() {
+ root = new RootAllocator(Long.MAX_VALUE);
+ }
+
+ @AfterEach
+ public void tearDown() {
+ root.close();
+ }
+
+ @Test
+ public void importOomDoesNotLeakExportedArray() {
+ // "producer" owns only the exported batch buffers; the C Data struct
containers live on a
+ // separate allocator (they are consumed/closed by import, which would
otherwise muddy the
+ // producer's balance). So producer draining to zero is an exact signal
that the array's release
+ // callback fired.
+ try (BufferAllocator producer = root.newChildAllocator("producer", 0,
Long.MAX_VALUE);
+ BufferAllocator structs = root.newChildAllocator("structs", 0,
Long.MAX_VALUE)) {
+ try (ArrowArray array = ArrowArray.allocateNew(structs);
+ ArrowSchema schema = ArrowSchema.allocateNew(structs)) {
+ exportBatch(producer, array, schema);
+ assertTrue(
+ producer.getAllocatedMemory() > 0, "producer holds the exported
batch before import");
+
+ // A consumer allocator far too small to hold the batch: the import
throws part-way through.
+ try (BufferAllocator consumer = root.newChildAllocator("consumer", 0,
TINY_LIMIT);
+ CDataDictionaryProvider provider = new CDataDictionaryProvider()) {
+ Schema importSchema = Data.importSchema(consumer, schema, provider);
+ try (VectorSchemaRoot importRoot =
VectorSchemaRoot.create(importSchema, consumer)) {
+ Exception thrown =
+ assertThrows(
+ Exception.class,
+ () -> Data.importIntoVectorSchemaRoot(consumer, array,
importRoot, provider));
+ assertTrue(
+ hasOutOfMemoryCause(thrown),
+ "mid-import failure must be an allocator OOM: " + thrown);
+ }
+ }
+
+ // The array's release callback must have fired despite the mid-import
OOM, freeing the
+ // whole exported batch. On the unfixed retain-before-wrap code the
batch is stranded.
+ assertEquals(
+ 0L,
+ producer.getAllocatedMemory(),
+ "import OOM leaked the exported batch (producer not drained)");
+ }
+ }
+ }
+
+ /** True if {@code t} is, or is caused by, an Arrow {@link
OutOfMemoryException}. */
+ private static boolean hasOutOfMemoryCause(Throwable t) {
+ for (Throwable cause = t; cause != null; cause = cause.getCause()) {
+ if (cause instanceof OutOfMemoryException) {
+ return true;
+ }
+ }
+ return false;
+ }
+
+ /**
+ * Builds a wide multi-column VarChar batch on {@code alloc} and exports it
into the C structs.
+ */
+ private void exportBatch(BufferAllocator alloc, ArrowArray array,
ArrowSchema schema) {
+ byte[] value = new byte[VALUE_BYTES];
+ for (int i = 0; i < value.length; i++) {
+ value[i] = (byte) 'x';
+ }
+ List<FieldVector> vectors = new ArrayList<>(COLUMNS);
+ for (int c = 0; c < COLUMNS; c++) {
+ VarCharVector vector = new VarCharVector("col" + c, alloc);
+ vector.allocateNew((long) ROWS * VALUE_BYTES, ROWS);
+ for (int r = 0; r < ROWS; r++) {
+ vector.setSafe(r, value);
+ }
+ vector.setValueCount(ROWS);
+ vectors.add(vector);
+ }
+ try (VectorSchemaRoot source = new VectorSchemaRoot(vectors)) {
+ long total = 0;
+ for (FieldVector vector : source.getFieldVectors()) {
+ total += vector.getBufferSize();
+ }
+ assertTrue(total > TINY_LIMIT, "test setup: batch must exceed the
consumer limit");
+ Data.exportVectorSchemaRoot(alloc, source, null, array, schema);
+ }
+ }
+}
