fix: Dispose record batches exported through the C stream interface

src/Apache.Arrow/C/CArrowArrayExporter.cs

@@ -91,6 +91,17 @@ public static unsafe void ExportArray(IArrowArray array, CArrowArray* cArray)
         /// </code>
         /// </example>
         public static unsafe void ExportRecordBatch(RecordBatch batch, CArrowArray* cArray)
+        {
+            ExportRecordBatch(batch, cArray, transferOwnership: false);
+        }
+
+        /// <summary>
+        /// Export a <see cref="RecordBatch"/> to a <see cref="CArrowArray"/>. When
+        /// <paramref name="transferOwnership"/> is true, the exported array takes over the caller's
+        /// reference to <paramref name="batch"/> and disposes it when the array is released (or
+        /// immediately if the export fails).
+        /// </summary>
+        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);

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();
             }

src/Apache.Arrow/C/CArrowDeviceArrayExporter.cs

@@ -67,6 +67,17 @@ public static unsafe void ExportArray(IArrowArray array, CArrowDeviceArray* devi
         /// </code>
         /// </example>
         public static unsafe void ExportRecordBatch(RecordBatch batch, CArrowDeviceArray* deviceArray)
+        {
+            ExportRecordBatch(batch, deviceArray, transferOwnership: false);
+        }
+
+        /// <summary>
+        /// Export a <see cref="RecordBatch"/> to a <see cref="CArrowDeviceArray"/>. When
+        /// <paramref name="transferOwnership"/> is true, the exported array takes over the caller's
+        /// reference to <paramref name="batch"/> and disposes it when the array is released (or
+        /// immediately if the export fails).
+        /// </summary>
+        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;

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();
             }

src/Apache.Arrow/Memory/ExportedAllocationOwner.cs

@@ -26,6 +26,7 @@ internal sealed class ExportedAllocationOwner : IDisposable
         private readonly List<IntPtr> _pointers = new List<IntPtr>();
         private readonly List<MemoryHandle> _handles = new List<MemoryHandle>();
         private readonly List<SharedMemoryHandle> _sharedHandles = new List<SharedMemoryHandle>();
+        private readonly List<IDisposable> _ownedObjects = new List<IDisposable>();
         private long _allocationSize;
         private long _referenceCount;
         private bool _disposed;
@@ -62,6 +63,15 @@ public unsafe IntPtr Acquire(SharedMemoryHandle sharedHandle)
             return pointer;
         }
 
+        /// <summary>
+        /// Makes this owner responsible for disposing <paramref name="obj"/> once the exported
+        /// memory is released.
+        /// </summary>
+        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)

test/Apache.Arrow.Tests/CDataInterfaceDataTests.cs

@@ -14,10 +14,15 @@
 // limitations under the License.
 
 using System;
+using System.Buffers;
 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 +303,129 @@ 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());
+            }
+        }
+
+        [Fact]
+        public async Task ExportedArrayStreamUnpinsMemoryBeforeDisposingBatch()
+        {
+            // The values buffer wraps a ReadOnlyMemory<byte>, 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<byte, int>(owner.Memory.Span));
+            var values = new ArrowBuffer((ReadOnlyMemory<byte>)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<IArrowArray> { 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()
+                .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 DisposeTrackingMemoryOwner : IMemoryOwner<byte>
+        {
+            private readonly IMemoryOwner<byte> _inner;
+
+            public DisposeTrackingMemoryOwner(IMemoryOwner<byte> inner)
+            {
+                _inner = inner;
+            }
+
+            public Memory<byte> Memory => _inner.Memory;
+
+            public bool Disposed { get; private set; }
+
+            public void Dispose()
+            {
+                _inner.Dispose();
+                Disposed = true;
+            }
+        }
+
+        private sealed class SingleBatchStream : IArrowArrayStream
+        {
+            private RecordBatch _batch;
+
+            public SingleBatchStream(RecordBatch batch)
+            {
+                Schema = batch.Schema;
+                _batch = batch;
+            }
+
+            public Schema Schema { get; }
+
+            public ValueTask<RecordBatch> ReadNextRecordBatchAsync(CancellationToken cancellationToken = default)
+            {
+                // Ownership of the returned batch passes to the caller.
+                RecordBatch result = _batch;
+                _batch = null;
+                return new ValueTask<RecordBatch>(result);
+            }
+
+            public void Dispose()
+            {
+                _batch?.Dispose();
+                _batch = null;
+            }
+        }
     }
 }

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()
         {