From 0728051d55d0c3912fd87abb62eb5ac5f1cd9002 Mon Sep 17 00:00:00 2001 From: ump45nose <52391318+ump45nose@users.noreply.github.com> Date: Tue, 20 Oct 2026 15:06:44 +0800 Subject: [PATCH 1/2] fix: Dispose record batches exported through the C stream interface CArrowArrayStreamExporter and CArrowDeviceArrayStreamExporter read each batch from the managed stream, export it and then drop it without disposing it. The exported array holds its own references to the buffers, so when the consumer releases it only those references go away; the batch's references are left for the finalizer, which in a NativeAOT library may never run. The exported array now takes ownership of the batch: the ExportedAllocationOwner disposes it when the array is released, after it has unpinned the exported buffers. Exporting a batch directly with ExportRecordBatch is unchanged and still leaves the batch to the caller. Closes #464. --- src/Apache.Arrow/C/CArrowArrayExporter.cs | 15 ++++ .../C/CArrowArrayStreamExporter.cs | 4 +- .../C/CArrowDeviceArrayExporter.cs | 13 +++- .../C/CArrowDeviceArrayStreamExporter.cs | 4 +- .../Memory/ExportedAllocationOwner.cs | 18 +++++ .../CDataInterfaceDataTests.cs | 74 +++++++++++++++++++ .../CDeviceDataInterfaceTests.cs | 30 ++++++++ 7 files changed, 155 insertions(+), 3 deletions(-) diff --git a/src/Apache.Arrow/C/CArrowArrayExporter.cs b/src/Apache.Arrow/C/CArrowArrayExporter.cs index 882533c0..12761daa 100644 --- a/src/Apache.Arrow/C/CArrowArrayExporter.cs +++ b/src/Apache.Arrow/C/CArrowArrayExporter.cs @@ -91,6 +91,17 @@ public static unsafe void ExportArray(IArrowArray array, CArrowArray* cArray) /// /// public static unsafe void ExportRecordBatch(RecordBatch batch, CArrowArray* cArray) + { + ExportRecordBatch(batch, cArray, transferOwnership: false); + } + + /// + /// Export a to a . When + /// is true, the exported array takes over the caller's + /// reference to and disposes it when the array is released (or + /// immediately if the export fails). + /// + internal static unsafe void ExportRecordBatch(RecordBatch batch, CArrowArray* cArray, bool transferOwnership) { if (batch == null) { @@ -106,6 +117,10 @@ public static unsafe void ExportRecordBatch(RecordBatch batch, CArrowArray* cArr } ExportedAllocationOwner allocationOwner = new ExportedAllocationOwner(); + if (transferOwnership) + { + allocationOwner.TakeOwnership(batch); + } try { ConvertRecordBatch(allocationOwner, batch, cArray); diff --git a/src/Apache.Arrow/C/CArrowArrayStreamExporter.cs b/src/Apache.Arrow/C/CArrowArrayStreamExporter.cs index 17b320f8..aec5ddb7 100644 --- a/src/Apache.Arrow/C/CArrowArrayStreamExporter.cs +++ b/src/Apache.Arrow/C/CArrowArrayStreamExporter.cs @@ -104,7 +104,9 @@ private unsafe static int GetNext(CArrowArrayStream* cArrayStream, CArrowArray* RecordBatch recordBatch = arrayStream.ArrowArrayStream.ReadNextRecordBatchAsync().Result; if (recordBatch != null) { - CArrowArrayExporter.ExportRecordBatch(recordBatch, cArray); + // The stream hands us ownership of each batch it returns. The exported array keeps + // it alive and disposes it on release, rather than leaving it to the finalizer. + CArrowArrayExporter.ExportRecordBatch(recordBatch, cArray, transferOwnership: true); } return arrayStream.ClearError(); } diff --git a/src/Apache.Arrow/C/CArrowDeviceArrayExporter.cs b/src/Apache.Arrow/C/CArrowDeviceArrayExporter.cs index 05dfabbd..8f70b378 100644 --- a/src/Apache.Arrow/C/CArrowDeviceArrayExporter.cs +++ b/src/Apache.Arrow/C/CArrowDeviceArrayExporter.cs @@ -67,6 +67,17 @@ public static unsafe void ExportArray(IArrowArray array, CArrowDeviceArray* devi /// /// public static unsafe void ExportRecordBatch(RecordBatch batch, CArrowDeviceArray* deviceArray) + { + ExportRecordBatch(batch, deviceArray, transferOwnership: false); + } + + /// + /// Export a to a . When + /// is true, the exported array takes over the caller's + /// reference to and disposes it when the array is released (or + /// immediately if the export fails). + /// + internal static unsafe void ExportRecordBatch(RecordBatch batch, CArrowDeviceArray* deviceArray, bool transferOwnership) { if (batch == null) { @@ -77,7 +88,7 @@ public static unsafe void ExportRecordBatch(RecordBatch batch, CArrowDeviceArray throw new ArgumentNullException(nameof(deviceArray)); } - CArrowArrayExporter.ExportRecordBatch(batch, &deviceArray->array); + CArrowArrayExporter.ExportRecordBatch(batch, &deviceArray->array, transferOwnership); deviceArray->device_type = ArrowDeviceType.Cpu; deviceArray->device_id = -1; deviceArray->sync_event = null; diff --git a/src/Apache.Arrow/C/CArrowDeviceArrayStreamExporter.cs b/src/Apache.Arrow/C/CArrowDeviceArrayStreamExporter.cs index d695efb8..5aa235c0 100644 --- a/src/Apache.Arrow/C/CArrowDeviceArrayStreamExporter.cs +++ b/src/Apache.Arrow/C/CArrowDeviceArrayStreamExporter.cs @@ -105,7 +105,9 @@ private unsafe static int GetNext(CArrowDeviceArrayStream* cDeviceArrayStream, C RecordBatch recordBatch = stream.ArrowArrayStream.ReadNextRecordBatchAsync().Result; if (recordBatch != null) { - CArrowDeviceArrayExporter.ExportRecordBatch(recordBatch, cDeviceArray); + // The stream hands us ownership of each batch it returns. The exported array keeps + // it alive and disposes it on release, rather than leaving it to the finalizer. + CArrowDeviceArrayExporter.ExportRecordBatch(recordBatch, cDeviceArray, transferOwnership: true); } return stream.ClearError(); } diff --git a/src/Apache.Arrow/Memory/ExportedAllocationOwner.cs b/src/Apache.Arrow/Memory/ExportedAllocationOwner.cs index 63d0ce10..8434a0e1 100644 --- a/src/Apache.Arrow/Memory/ExportedAllocationOwner.cs +++ b/src/Apache.Arrow/Memory/ExportedAllocationOwner.cs @@ -26,6 +26,7 @@ internal sealed class ExportedAllocationOwner : IDisposable private readonly List _pointers = new List(); private readonly List _handles = new List(); private readonly List _sharedHandles = new List(); + private readonly List _ownedObjects = new List(); private long _allocationSize; private long _referenceCount; private bool _disposed; @@ -62,6 +63,15 @@ public unsafe IntPtr Acquire(SharedMemoryHandle sharedHandle) return pointer; } + /// + /// Makes this owner responsible for disposing once the exported + /// memory is released. + /// + public void TakeOwnership(IDisposable obj) + { + _ownedObjects.Add(obj); + } + public void IncRef() { Interlocked.Increment(ref _referenceCount); @@ -103,6 +113,14 @@ public void Dispose() _sharedHandles[i] = default; } + // Owned objects are disposed only after the pins above have been released, because + // some memory managers throw when they are disposed while still pinned. + for (int i = 0; i < _ownedObjects.Count; i++) + { + _ownedObjects[i]?.Dispose(); + _ownedObjects[i] = null; + } + // Nothing is allocated when exporting an array without buffers or children (e.g. a NullArray), // and RemoveMemoryPressure throws for zero. if (_allocationSize > 0) diff --git a/test/Apache.Arrow.Tests/CDataInterfaceDataTests.cs b/test/Apache.Arrow.Tests/CDataInterfaceDataTests.cs index 5fcabb2e..c9747e6e 100644 --- a/test/Apache.Arrow.Tests/CDataInterfaceDataTests.cs +++ b/test/Apache.Arrow.Tests/CDataInterfaceDataTests.cs @@ -17,7 +17,11 @@ using System.Collections.Generic; using System.Linq; using System.Runtime.InteropServices; +using System.Threading; +using System.Threading.Tasks; using Apache.Arrow.C; +using Apache.Arrow.Ipc; +using Apache.Arrow.Memory; using Apache.Arrow.Types; using Xunit; @@ -298,5 +302,75 @@ RecordBatch GetTestRecordBatch() CArrowArray.Free(cArray); } + [Fact] + public async Task ExportedArrayStreamDisposesBatchesOnRelease() + { + var allocator = new TestMemoryAllocator(); + RecordBatch batch = GetAllocatorBackedRecordBatch(allocator); + Assert.True(allocator.Rented > 0); + + IArrowArrayStream imported; + unsafe + { + CArrowArrayStream* cArrayStream = CArrowArrayStream.Create(); + CArrowArrayStreamExporter.ExportArrayStream(new SingleBatchStream(batch), cArrayStream); + imported = CArrowArrayStreamImporter.ImportArrayStream(cArrayStream); + CArrowArrayStream.Free(cArrayStream); + } + + using (imported) + { + RecordBatch importedBatch = await imported.ReadNextRecordBatchAsync(); + Assert.Equal(new int?[] { 1, 2, null, 4 }, (Int32Array)importedBatch.Column(0)); + Assert.True(allocator.Rented > 0); + + // Releasing the exported array should release the batch the stream handed over. + importedBatch.Dispose(); + Assert.Equal(0, allocator.Rented); + + Assert.Null(await imported.ReadNextRecordBatchAsync()); + } + } + + internal static RecordBatch GetAllocatorBackedRecordBatch(MemoryAllocator allocator) + { + Int32Array array = new Int32Array.Builder() + .Append(1) + .Append(2) + .AppendNull() + .Append(4) + .Build(allocator); + Schema schema = new Schema.Builder() + .Field(f => f.Name("ints").DataType(Int32Type.Default).Nullable(true)) + .Build(); + return new RecordBatch(schema, new IArrowArray[] { array }, array.Length); + } + + private sealed class SingleBatchStream : IArrowArrayStream + { + private RecordBatch _batch; + + public SingleBatchStream(RecordBatch batch) + { + Schema = batch.Schema; + _batch = batch; + } + + public Schema Schema { get; } + + public ValueTask ReadNextRecordBatchAsync(CancellationToken cancellationToken = default) + { + // Ownership of the returned batch passes to the caller. + RecordBatch result = _batch; + _batch = null; + return new ValueTask(result); + } + + public void Dispose() + { + _batch?.Dispose(); + _batch = null; + } + } } } diff --git a/test/Apache.Arrow.Tests/CDeviceDataInterfaceTests.cs b/test/Apache.Arrow.Tests/CDeviceDataInterfaceTests.cs index 3954d9c6..a197a42b 100644 --- a/test/Apache.Arrow.Tests/CDeviceDataInterfaceTests.cs +++ b/test/Apache.Arrow.Tests/CDeviceDataInterfaceTests.cs @@ -317,6 +317,36 @@ public async Task ExportImportDeviceStreamRoundTrip() } } + [Fact] + public async Task ExportedDeviceStreamDisposesBatchesOnRelease() + { + var allocator = new TestMemoryAllocator(); + RecordBatch batch = CDataInterfaceDataTests.GetAllocatorBackedRecordBatch(allocator); + Assert.True(allocator.Rented > 0); + + IArrowArrayStream imported; + unsafe + { + CArrowDeviceArrayStream* ptr = CArrowDeviceArrayStream.Create(); + CArrowDeviceArrayStreamExporter.ExportArrayStream(new TestArrayStream(batch.Schema, batch), ptr); + imported = CArrowDeviceArrayStreamImporter.ImportDeviceArrayStream(ptr); + CArrowDeviceArrayStream.Free(ptr); + } + + using (imported) + { + RecordBatch importedBatch = await imported.ReadNextRecordBatchAsync(); + Assert.Equal(new int?[] { 1, 2, null, 4 }, (Int32Array)importedBatch.Column(0)); + Assert.True(allocator.Rented > 0); + + // Releasing the exported array should release the batch the stream handed over. + importedBatch.Dispose(); + Assert.Equal(0, allocator.Rented); + + Assert.Null(await imported.ReadNextRecordBatchAsync()); + } + } + [Fact] public unsafe void ImportNonCpuDeviceStreamThrows() { From ee38eb449ad5061c343796bcbd3b8eb4a6574820 Mon Sep 17 00:00:00 2001 From: ump45nose Date: Wed, 21 Oct 2026 21:16:15 +0800 Subject: [PATCH 2/2] test: Cover exporting a stream whose batch pins ReadOnlyMemory buffers --- .../CDataInterfaceDataTests.cs | 55 +++++++++++++++++++ 1 file changed, 55 insertions(+) diff --git a/test/Apache.Arrow.Tests/CDataInterfaceDataTests.cs b/test/Apache.Arrow.Tests/CDataInterfaceDataTests.cs index c9747e6e..d265cad9 100644 --- a/test/Apache.Arrow.Tests/CDataInterfaceDataTests.cs +++ b/test/Apache.Arrow.Tests/CDataInterfaceDataTests.cs @@ -14,6 +14,7 @@ // limitations under the License. using System; +using System.Buffers; using System.Collections.Generic; using System.Linq; using System.Runtime.InteropServices; @@ -332,6 +333,40 @@ public async Task ExportedArrayStreamDisposesBatchesOnRelease() } } + [Fact] + public async Task ExportedArrayStreamUnpinsMemoryBeforeDisposingBatch() + { + // The values buffer wraps a ReadOnlyMemory, so exporting it pins the memory. + // The native memory manager throws if it is disposed while still pinned. + var owner = new DisposeTrackingMemoryOwner(new NativeMemoryAllocator().Allocate(4 * sizeof(int))); + new int[] { 1, 2, 3, 4 }.CopyTo(MemoryMarshal.Cast(owner.Memory.Span)); + var values = new ArrowBuffer((ReadOnlyMemory)owner.Memory.Slice(0, 4 * sizeof(int))); + var array = new Int32Array(values, ArrowBuffer.Empty, 4, 0, 0); + Schema schema = new Schema.Builder() + .Field(f => f.Name("ints").DataType(Int32Type.Default).Nullable(false)) + .Build(); + var batch = new RecordBatch(schema, owner, new List { array }, array.Length); + + IArrowArrayStream imported; + unsafe + { + CArrowArrayStream* cArrayStream = CArrowArrayStream.Create(); + CArrowArrayStreamExporter.ExportArrayStream(new SingleBatchStream(batch), cArrayStream); + imported = CArrowArrayStreamImporter.ImportArrayStream(cArrayStream); + CArrowArrayStream.Free(cArrayStream); + } + + using (imported) + { + RecordBatch importedBatch = await imported.ReadNextRecordBatchAsync(); + Assert.Equal(new int?[] { 1, 2, 3, 4 }, (Int32Array)importedBatch.Column(0)); + Assert.False(owner.Disposed); + + importedBatch.Dispose(); + Assert.True(owner.Disposed); + } + } + internal static RecordBatch GetAllocatorBackedRecordBatch(MemoryAllocator allocator) { Int32Array array = new Int32Array.Builder() @@ -346,6 +381,26 @@ internal static RecordBatch GetAllocatorBackedRecordBatch(MemoryAllocator alloca return new RecordBatch(schema, new IArrowArray[] { array }, array.Length); } + private sealed class DisposeTrackingMemoryOwner : IMemoryOwner + { + private readonly IMemoryOwner _inner; + + public DisposeTrackingMemoryOwner(IMemoryOwner inner) + { + _inner = inner; + } + + public Memory Memory => _inner.Memory; + + public bool Disposed { get; private set; } + + public void Dispose() + { + _inner.Dispose(); + Disposed = true; + } + } + private sealed class SingleBatchStream : IArrowArrayStream { private RecordBatch _batch;