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