diff --git a/src/Apache.Arrow/C/CArrowArrayExporter.cs b/src/Apache.Arrow/C/CArrowArrayExporter.cs
index 882533c..12761da 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 17b320f..aec5ddb 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 05dfabb..8f70b37 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 d695efb..5aa235c 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 63d0ce1..8434a0e 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 5fcabb2..d265cad 100644
--- a/test/Apache.Arrow.Tests/CDataInterfaceDataTests.cs
+++ b/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, 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()
+ .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
+ {
+ 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;
+
+ 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 3954d9c..a197a42 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()
{