This is an automated email from the ASF dual-hosted git repository.
adamreeve pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-dotnet.git
The following commit(s) were added to refs/heads/main by this push:
new c964702 fix: Dispose record batches exported through the C stream
interface (#483)
c964702 is described below
commit c96470238a99a4b8542ed582c2fd5007b78c4365
Author: ump45nose <[email protected]>
AuthorDate: Fri Oct 23 05:08:21 2026 +0800
fix: Dispose record batches exported through the C stream interface (#483)
- `ExportedAllocationOwner` gets `TakeOwnership(IDisposable)`. Owned
objects are disposed in `Dispose`, after the exported buffers have been
unpinned and their shared handles released, so memory managers that
throw when disposed while pinned are not affected. This also covers
buffers exported by pinning rather than through a shared handle.
- `CArrowArrayExporter` and `CArrowDeviceArrayExporter` get internal
`ExportRecordBatch` overloads with a `transferOwnership` flag. The
stream exporters pass `true`, so a batch is disposed when the consumer
releases the exported array (or right away if the export fails). The
public `ExportRecordBatch` methods still leave the batch to the caller.
Batches returned by an `IArrowArrayStream` are owned by the caller, so
streams that keep using a batch after handing it out would now see it
disposed once the consumer releases it.
---
src/Apache.Arrow/C/CArrowArrayExporter.cs | 15 +++
src/Apache.Arrow/C/CArrowArrayStreamExporter.cs | 4 +-
src/Apache.Arrow/C/CArrowDeviceArrayExporter.cs | 13 ++-
.../C/CArrowDeviceArrayStreamExporter.cs | 4 +-
src/Apache.Arrow/Memory/ExportedAllocationOwner.cs | 18 +++
test/Apache.Arrow.Tests/CDataInterfaceDataTests.cs | 129 +++++++++++++++++++++
.../CDeviceDataInterfaceTests.cs | 30 +++++
7 files changed, 210 insertions(+), 3 deletions(-)
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 @@ namespace Apache.Arrow.C
/// </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 @@ namespace Apache.Arrow.C
}
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 @@ namespace Apache.Arrow.C
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 @@ namespace Apache.Arrow.C
/// </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 @@ namespace Apache.Arrow.C
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 @@ namespace Apache.Arrow.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 @@ namespace Apache.Arrow.Memory
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 @@ namespace Apache.Arrow.Memory
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 @@ namespace Apache.Arrow.Memory
_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 @@ namespace Apache.Arrow.Tests
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;
+ }
+ }
}
}
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 @@ namespace Apache.Arrow.Tests
}
}
+ [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()
{
