This is an automated email from the ASF dual-hosted git repository.
curth 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 fb9199e Introduce a reference counted layer between ArrowBuffer and
memory (#297)
fb9199e is described below
commit fb9199e0a511b32f311ba2703c1f5a82bf973e8b
Author: Curt Hagenlocher <[email protected]>
AuthorDate: Mon Mar 30 13:24:08 2026 -0700
Introduce a reference counted layer between ArrowBuffer and memory (#297)
## What's Changed
Introduces a reference-counted layer between ArrowBuffer and the
underlying memory to allow buffers to be shared between multiple arrays.
Supports export of managed buffers.
Disables experimental workaround that was previously added to support
buffer export given that it is no longer necessary.
This is an alternative to #291 that's more flexible.
Closes #111.
---
src/Apache.Arrow/Arrays/Array.cs | 14 +
src/Apache.Arrow/Arrays/ArrayData.cs | 97 +++
src/Apache.Arrow/Arrays/ArrowArrayFactory.cs | 6 +
src/Apache.Arrow/ArrowBuffer.cs | 51 +-
src/Apache.Arrow/C/CArrowArrayExporter.cs | 15 +-
src/Apache.Arrow/ChunkedArray.cs | 71 ++-
src/Apache.Arrow/Column.cs | 39 +-
src/Apache.Arrow/Memory/ExportedAllocationOwner.cs | 15 +
src/Apache.Arrow/Memory/IOwnableAllocation.cs | 24 -
src/Apache.Arrow/Memory/NativeMemoryManager.cs | 23 +-
src/Apache.Arrow/Memory/SharedMemoryHandle.cs | 54 ++
src/Apache.Arrow/Memory/SharedMemoryOwner.cs | 62 ++
src/Apache.Arrow/RecordBatch.cs | 11 +
.../ArrayDataReferenceCountingTests.cs | 703 +++++++++++++++++++++
test/Apache.Arrow.Tests/ArrowBufferTests.cs | 319 +++++++++-
test/Apache.Arrow.Tests/CDataInterfaceDataTests.cs | 30 +
.../CDataInterfacePythonTests.cs | 95 ++-
17 files changed, 1485 insertions(+), 144 deletions(-)
diff --git a/src/Apache.Arrow/Arrays/Array.cs b/src/Apache.Arrow/Arrays/Array.cs
index 4abe63e..5083e75 100644
--- a/src/Apache.Arrow/Arrays/Array.cs
+++ b/src/Apache.Arrow/Arrays/Array.cs
@@ -60,11 +60,25 @@ namespace Apache.Arrow
}
}
+ /// <summary>
+ /// Slice this array without changing ownership. The returned slice may
+ /// become invalid if the original array is disposed.
+ /// </summary>
public Array Slice(int offset, int length)
{
return ArrowArrayFactory.Slice(this, offset, length) as Array;
}
+ /// <summary>
+ /// Slice this array with shared ownership. The returned slice keeps
the
+ /// underlying buffers alive via reference counting. The caller must
dispose
+ /// the returned array when done.
+ /// </summary>
+ public Array SliceShared(int offset, int length)
+ {
+ return ArrowArrayFactory.SliceShared(this, offset, length) as
Array;
+ }
+
public void Dispose()
{
Dispose(true);
diff --git a/src/Apache.Arrow/Arrays/ArrayData.cs
b/src/Apache.Arrow/Arrays/ArrayData.cs
index 25cda4d..87f3af8 100644
--- a/src/Apache.Arrow/Arrays/ArrayData.cs
+++ b/src/Apache.Arrow/Arrays/ArrayData.cs
@@ -118,6 +118,13 @@ namespace Apache.Arrow
Dictionary?.Dispose();
}
+ /// <summary>
+ /// Slice this ArrayData without ownership tracking. The returned
slice shares
+ /// the underlying buffers but does not keep them alive — the caller
must ensure
+ /// the original ArrayData outlives the slice.
+ /// Consider using <see cref="SliceShared"/> instead, which uses
reference counting
+ /// to keep the underlying buffers alive for the lifetime of the slice.
+ /// </summary>
public ArrayData Slice(int offset, int length)
{
if (offset > Length)
@@ -149,6 +156,96 @@ namespace Apache.Arrow
return new ArrayData(DataType, length, nullCount, offset, Buffers,
Children, Dictionary);
}
+ /// <summary>
+ /// Retain this ArrayData with shared ownership. The returned
ArrayData keeps the
+ /// underlying buffers alive via reference counting, and recursively
retains any
+ /// children and dictionary. The caller must dispose the returned
ArrayData when done.
+ /// </summary>
+ public ArrayData Retain()
+ {
+ return new ArrayData(
+ DataType,
+ Length,
+ NullCount,
+ Offset,
+ RetainBuffers(Buffers),
+ RetainChildren(Children),
+ Dictionary?.Retain());
+ }
+
+ /// <summary>
+ /// Slice this ArrayData with shared ownership. The returned slice
keeps the
+ /// underlying buffers alive via reference counting. The caller must
dispose the
+ /// returned ArrayData when done.
+ /// </summary>
+ public ArrayData SliceShared(int offset, int length)
+ {
+ if (offset > Length)
+ {
+ throw new ArgumentException($"Offset {offset} cannot be
greater than Length {Length} for Array.SliceShared");
+ }
+
+ length = Math.Min(Length - offset, length);
+ offset += Offset;
+
+ int nullCount;
+ if (NullCount == 0)
+ {
+ nullCount = 0;
+ }
+ else if (NullCount == Length)
+ {
+ nullCount = length;
+ }
+ else if (offset == Offset && length == Length)
+ {
+ nullCount = NullCount;
+ }
+ else
+ {
+ nullCount = RecalculateNullCount;
+ }
+
+ return new ArrayData(
+ DataType,
+ length,
+ nullCount,
+ offset,
+ RetainBuffers(Buffers),
+ RetainChildren(Children),
+ Dictionary?.Retain());
+ }
+
+ private static ArrowBuffer[] RetainBuffers(ArrowBuffer[] buffers)
+ {
+ if (buffers == null)
+ {
+ return null;
+ }
+
+ var retained = new ArrowBuffer[buffers.Length];
+ for (int i = 0; i < buffers.Length; i++)
+ {
+ retained[i] = buffers[i].Retain();
+ }
+ return retained;
+ }
+
+ private static ArrayData[] RetainChildren(ArrayData[] children)
+ {
+ if (children == null)
+ {
+ return null;
+ }
+
+ var retained = new ArrayData[children.Length];
+ for (int i = 0; i < children.Length; i++)
+ {
+ retained[i] = children[i]?.Retain();
+ }
+ return retained;
+ }
+
public ArrayData Clone(MemoryAllocator allocator = default)
{
return new ArrayData(
diff --git a/src/Apache.Arrow/Arrays/ArrowArrayFactory.cs
b/src/Apache.Arrow/Arrays/ArrowArrayFactory.cs
index 2518a32..ae5d7be 100644
--- a/src/Apache.Arrow/Arrays/ArrowArrayFactory.cs
+++ b/src/Apache.Arrow/Arrays/ArrowArrayFactory.cs
@@ -122,5 +122,11 @@ namespace Apache.Arrow
ArrayData newData = array.Data.Slice(offset, length);
return BuildArray(newData);
}
+
+ public static IArrowArray SliceShared(IArrowArray array, int offset,
int length)
+ {
+ ArrayData newData = array.Data.SliceShared(offset, length);
+ return BuildArray(newData);
+ }
}
}
diff --git a/src/Apache.Arrow/ArrowBuffer.cs b/src/Apache.Arrow/ArrowBuffer.cs
index 28a06c3..7b8ecc7 100644
--- a/src/Apache.Arrow/ArrowBuffer.cs
+++ b/src/Apache.Arrow/ArrowBuffer.cs
@@ -16,36 +16,37 @@
using System;
using System.Buffers;
using System.Runtime.CompilerServices;
-using Apache.Arrow.C;
using Apache.Arrow.Memory;
namespace Apache.Arrow
{
public readonly partial struct ArrowBuffer : IEquatable<ArrowBuffer>,
IDisposable
{
- private readonly IMemoryOwner<byte> _memoryOwner;
+ private readonly SharedMemoryHandle _handle;
private readonly ReadOnlyMemory<byte> _memory;
public static ArrowBuffer Empty => new ArrowBuffer(Memory<byte>.Empty);
public ArrowBuffer(ReadOnlyMemory<byte> data)
{
- _memoryOwner = null;
+ _handle = null;
_memory = data;
}
internal ArrowBuffer(IMemoryOwner<byte> memoryOwner)
{
- // When wrapping an IMemoryOwner, don't cache the Memory<byte>
- // since the owner may be disposed, and the cached Memory would
- // be invalid.
+ _handle = new SharedMemoryHandle(new
SharedMemoryOwner(memoryOwner));
+ _memory = Memory<byte>.Empty;
+ }
- _memoryOwner = memoryOwner;
+ private ArrowBuffer(SharedMemoryHandle handle)
+ {
+ _handle = handle;
_memory = Memory<byte>.Empty;
}
public ReadOnlyMemory<byte> Memory =>
- _memoryOwner != null ? _memoryOwner.Memory : _memory;
+ _handle != null ? _handle.Memory : _memory;
public bool IsEmpty => Memory.IsEmpty;
@@ -57,6 +58,20 @@ namespace Apache.Arrow
get => Memory.Span;
}
+ /// <summary>
+ /// Adds another reference to the memory used by this buffer. Allows a
single buffer
+ /// to be shared by multiple arrays.
+ /// </summary>
+ public ArrowBuffer Retain()
+ {
+ if (_handle != null)
+ {
+ return new ArrowBuffer(_handle.Retain());
+ }
+
+ return new ArrowBuffer(_memory);
+ }
+
public ArrowBuffer Clone(MemoryAllocator allocator = default)
{
return Span.Length == 0 ? Empty : new Builder<byte>(Span.Length)
@@ -71,34 +86,26 @@ namespace Apache.Arrow
public void Dispose()
{
- _memoryOwner?.Dispose();
+ _handle?.Dispose();
}
internal bool TryExport(ExportedAllocationOwner newOwner, out IntPtr
ptr)
{
if (IsEmpty)
{
- // _memoryOwner could be anything (for example null or a
NullMemoryOwner), but it doesn't matter here
ptr = IntPtr.Zero;
return true;
}
- if (_memoryOwner is IOwnableAllocation ownable &&
ownable.TryAcquire(out ptr, out int offset, out int length))
- {
- newOwner.Acquire(ptr, offset, length);
- ptr += offset;
- return true;
- }
-
- if (_memoryOwner == null &&
CArrowArrayExporter.EnableManagedMemoryExport)
+ if (_handle != null)
{
- var handle = _memory.Pin();
- ptr = newOwner.Reference(handle);
+ ptr = newOwner.Acquire(_handle.Retain());
return true;
}
- ptr = IntPtr.Zero;
- return false;
+ var pinHandle = _memory.Pin();
+ ptr = newOwner.Reference(pinHandle);
+ return true;
}
}
}
diff --git a/src/Apache.Arrow/C/CArrowArrayExporter.cs
b/src/Apache.Arrow/C/CArrowArrayExporter.cs
index 0e140b5..32daad3 100644
--- a/src/Apache.Arrow/C/CArrowArrayExporter.cs
+++ b/src/Apache.Arrow/C/CArrowArrayExporter.cs
@@ -27,8 +27,9 @@ namespace Apache.Arrow.C
public static class CArrowArrayExporter
{
/// <summary>
- /// Experimental feature to enable exporting managed memory to
CArrowArray. Use with caution.
+ /// Formerly-experimental feature to enable exporting managed memory
to CArrowArray. Now obsolete.
/// </summary>
+ [Obsolete("EnableManagedMemoryExport is obsolete and ignored; managed
memory export is now always enabled and this field will be removed in a future
release.")]
public static bool EnableManagedMemoryExport = false;
#if NET5_0_OR_GREATER
@@ -39,9 +40,9 @@ namespace Apache.Arrow.C
private static IntPtr ReleaseArrayPtr => s_releaseArray.Pointer;
#endif
/// <summary>
- /// Export an <see cref="IArrowArray"/> to a <see
cref="CArrowArray"/>. Whether or not the
- /// export succeeds, the original array becomes invalid. Clone an
array to continue using it
- /// after a copy has been exported.
+ /// Export an <see cref="IArrowArray"/> to a <see
cref="CArrowArray"/>. The exported array
+ /// shares the underlying buffers via reference counting, so the
original array remains valid
+ /// after export.
/// </summary>
/// <param name="array">The array to export</param>
/// <param name="cArray">An allocated but uninitialized CArrowArray
pointer.</param>
@@ -76,9 +77,9 @@ namespace Apache.Arrow.C
}
/// <summary>
- /// Export a <see cref="RecordBatch"/> to a <see cref="CArrowArray"/>.
Whether or not the
- /// export succeeds, the original record batch becomes invalid. Clone
the batch to continue using it
- /// after a copy has been exported.
+ /// Export a <see cref="RecordBatch"/> to a <see cref="CArrowArray"/>.
The exported record
+ /// batch shares the underlying buffers via reference counting, so the
original batch remains
+ /// valid after export.
/// </summary>
/// <param name="batch">The record batch to export</param>
/// <param name="cArray">An allocated but uninitialized CArrowArray
pointer.</param>
diff --git a/src/Apache.Arrow/ChunkedArray.cs b/src/Apache.Arrow/ChunkedArray.cs
index 2c238ac..10c457f 100644
--- a/src/Apache.Arrow/ChunkedArray.cs
+++ b/src/Apache.Arrow/ChunkedArray.cs
@@ -22,12 +22,13 @@ namespace Apache.Arrow
/// <summary>
/// A data structure to manage a list of primitive Array arrays logically
as one large array
/// </summary>
- public class ChunkedArray
+ public class ChunkedArray : IDisposable
{
private IList<IArrowArray> Arrays { get; }
public IArrowType DataType { get; }
public long Length { get; }
public long NullCount { get; }
+ private bool DisposeArrays { get; set; }
public int ArrayCount
{
@@ -39,11 +40,21 @@ namespace Apache.Arrow
public IArrowArray ArrowArray(int index) => Arrays[index];
public ChunkedArray(IList<Array> arrays)
- : this(Cast(arrays))
+ : this(Cast(arrays), disposeArrays: false)
{
}
public ChunkedArray(IList<IArrowArray> arrays)
+ : this(arrays, disposeArrays: false)
+ {
+ }
+
+ public ChunkedArray(Array array)
+ : this(new IArrowArray[] { array }, disposeArrays: false)
+ {
+ }
+
+ private ChunkedArray(IList<IArrowArray> arrays, bool disposeArrays)
{
Arrays = arrays ?? throw new ArgumentNullException(nameof(arrays));
if (arrays.Count < 1)
@@ -56,20 +67,19 @@ namespace Apache.Arrow
Length += array.Length;
NullCount += array.NullCount;
}
+ DisposeArrays = disposeArrays;
}
- public ChunkedArray(Array array) : this(new IArrowArray[] { array }) {
}
-
public ChunkedArray Slice(long offset, long length)
{
if (offset >= Length)
{
- throw new ArgumentException($"Index {offset} cannot be greater
than the Column's Length {Length}");
+ throw new ArgumentException($"Offset {offset} cannot be
greater than Length {Length} for ChunkedArray.Slice");
}
int curArrayIndex = 0;
int numArrays = Arrays.Count;
- while (curArrayIndex < numArrays && offset >
Arrays[curArrayIndex].Length)
+ while (curArrayIndex < numArrays && offset >=
Arrays[curArrayIndex].Length)
{
offset -= Arrays[curArrayIndex].Length;
curArrayIndex++;
@@ -92,6 +102,55 @@ namespace Apache.Arrow
return Slice(offset, Length - offset);
}
+ /// <summary>
+ /// Slice this chunked array with shared ownership. The returned slice
keeps the
+ /// underlying buffers alive via reference counting. The caller must
dispose
+ /// the returned chunked array when done.
+ /// </summary>
+ public ChunkedArray SliceShared(long offset, long length)
+ {
+ if (offset >= Length)
+ {
+ throw new ArgumentException($"Offset {offset} cannot be
greater than Length {Length} for ChunkedArray.SliceShared");
+ }
+
+ int curArrayIndex = 0;
+ int numArrays = Arrays.Count;
+ while (curArrayIndex < numArrays && offset >=
Arrays[curArrayIndex].Length)
+ {
+ offset -= Arrays[curArrayIndex].Length;
+ curArrayIndex++;
+ }
+
+ IList<IArrowArray> newArrays = new List<IArrowArray>();
+ while (curArrayIndex < numArrays && length > 0)
+ {
+
newArrays.Add(ArrowArrayFactory.SliceShared(Arrays[curArrayIndex], (int)offset,
+ length > Arrays[curArrayIndex].Length ?
Arrays[curArrayIndex].Length : (int)length));
+ length -= Arrays[curArrayIndex].Length - offset;
+ offset = 0;
+ curArrayIndex++;
+ }
+ return new ChunkedArray(newArrays, disposeArrays: true);
+ }
+
+ public ChunkedArray SliceShared(long offset)
+ {
+ return SliceShared(offset, Length - offset);
+ }
+
+ public void Dispose()
+ {
+ if (DisposeArrays)
+ {
+ DisposeArrays = false;
+ foreach (IArrowArray array in Arrays)
+ {
+ array.Dispose();
+ }
+ }
+ }
+
public override string ToString() => $"{nameof(ChunkedArray)}:
Length={Length}, DataType={DataType.Name}";
private static IArrowArray[] Cast(IList<Array> arrays)
diff --git a/src/Apache.Arrow/Column.cs b/src/Apache.Arrow/Column.cs
index 03be51c..5d4ddd7 100644
--- a/src/Apache.Arrow/Column.cs
+++ b/src/Apache.Arrow/Column.cs
@@ -22,22 +22,23 @@ namespace Apache.Arrow
/// <summary>
/// A Column data structure that logically represents a column in a dataset
/// </summary>
- public class Column
+ public class Column : IDisposable
{
public Field Field { get; }
public ChunkedArray Data { get; }
+ private bool DisposeArrayData { get; set; }
public Column(Field field, IList<Array> arrays)
- : this(field, new ChunkedArray(arrays), doValidation: true)
+ : this(field, new ChunkedArray(arrays), doValidation: true,
disposeArrayData: false)
{
}
public Column(Field field, IList<IArrowArray> arrays)
- : this(field, new ChunkedArray(arrays), doValidation: true)
+ : this(field, new ChunkedArray(arrays), doValidation: true,
disposeArrayData: false)
{
}
- private Column(Field field, ChunkedArray data, bool doValidation =
false)
+ private Column(Field field, ChunkedArray data, bool doValidation =
false, bool disposeArrayData = false)
{
Data = data;
Field = field;
@@ -45,6 +46,7 @@ namespace Apache.Arrow
{
throw new ArgumentException($"{Field.DataType} must match
{Data.DataType}");
}
+ DisposeArrayData = disposeArrayData;
}
public long Length => Data.Length;
@@ -62,6 +64,35 @@ namespace Apache.Arrow
return new Column(Field, Data.Slice(offset));
}
+ /// <summary>
+ /// Slice this column with shared ownership. The returned slice keeps
the
+ /// underlying buffers alive via reference counting. The caller must
+ /// dispose the returned column when done.
+ /// </summary>
+ public Column SliceShared(int offset, int length)
+ {
+ return new Column(Field, Data.SliceShared(offset, length),
disposeArrayData: true);
+ }
+
+ /// <summary>
+ /// Slice this column with shared ownership. The returned slice keeps
the
+ /// underlying buffers alive via reference counting. The caller must
+ /// dispose the returned column when done.
+ /// </summary>
+ public Column SliceShared(int offset)
+ {
+ return new Column(Field, Data.SliceShared(offset),
disposeArrayData: true);
+ }
+
+ public void Dispose()
+ {
+ if (DisposeArrayData)
+ {
+ DisposeArrayData = false;
+ Data?.Dispose();
+ }
+ }
+
private bool ValidateArrayDataTypes()
{
var dataTypeComparer = new ArrayDataTypeComparer(Field.DataType);
diff --git a/src/Apache.Arrow/Memory/ExportedAllocationOwner.cs
b/src/Apache.Arrow/Memory/ExportedAllocationOwner.cs
index cd52f6d..272841e 100644
--- a/src/Apache.Arrow/Memory/ExportedAllocationOwner.cs
+++ b/src/Apache.Arrow/Memory/ExportedAllocationOwner.cs
@@ -25,6 +25,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 long _allocationSize;
private long _referenceCount;
private bool _disposed;
@@ -53,6 +54,14 @@ namespace Apache.Arrow.Memory
return new IntPtr(handle.Pointer);
}
+ public unsafe IntPtr Acquire(SharedMemoryHandle sharedHandle)
+ {
+ MemoryHandle handle = sharedHandle.Memory.Pin();
+ IntPtr pointer = Reference(handle);
+ _sharedHandles.Add(sharedHandle);
+ return pointer;
+ }
+
public void IncRef()
{
Interlocked.Increment(ref _referenceCount);
@@ -88,6 +97,12 @@ namespace Apache.Arrow.Memory
_handles[i] = default;
}
+ for (int i = 0; i < _sharedHandles.Count; i++)
+ {
+ _sharedHandles[i]?.Dispose();
+ _sharedHandles[i] = default;
+ }
+
GC.RemoveMemoryPressure(_allocationSize);
GC.SuppressFinalize(this);
_disposed = true;
diff --git a/src/Apache.Arrow/Memory/IOwnableAllocation.cs
b/src/Apache.Arrow/Memory/IOwnableAllocation.cs
deleted file mode 100644
index a5e7565..0000000
--- a/src/Apache.Arrow/Memory/IOwnableAllocation.cs
+++ /dev/null
@@ -1,24 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one or more
-// contributor license agreements. See the NOTICE file distributed with
-// this work for additional information regarding copyright ownership.
-// The ASF licenses this file to You under the Apache License, Version 2.0
-// (the "License"); you may not use this file except in compliance with
-// the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing, software
-// distributed under the License is distributed on an "AS IS" BASIS,
-// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-// See the License for the specific language governing permissions and
-// limitations under the License.
-
-using System;
-
-namespace Apache.Arrow.Memory
-{
- internal interface IOwnableAllocation
- {
- bool TryAcquire(out IntPtr ptr, out int offset, out int length);
- }
-}
diff --git a/src/Apache.Arrow/Memory/NativeMemoryManager.cs
b/src/Apache.Arrow/Memory/NativeMemoryManager.cs
index 5fb51be..8bcb705 100644
--- a/src/Apache.Arrow/Memory/NativeMemoryManager.cs
+++ b/src/Apache.Arrow/Memory/NativeMemoryManager.cs
@@ -20,7 +20,7 @@ using System.Threading;
namespace Apache.Arrow.Memory
{
- public class NativeMemoryManager : MemoryManager<byte>, IOwnableAllocation
+ public class NativeMemoryManager : MemoryManager<byte>
{
private IntPtr _ptr;
private int _pinCount;
@@ -91,27 +91,6 @@ namespace Apache.Arrow.Memory
}
}
- bool IOwnableAllocation.TryAcquire(out IntPtr ptr, out int offset, out
int length)
- {
- // TODO: implement refcounted buffers?
-
- if (object.ReferenceEquals(_owner,
NativeMemoryAllocator.ExclusiveOwner))
- {
- ptr = Interlocked.Exchange(ref _ptr, IntPtr.Zero);
- if (ptr != IntPtr.Zero)
- {
- offset = _offset;
- length = _length;
- return true;
- }
- }
-
- ptr = IntPtr.Zero;
- offset = 0;
- length = 0;
- return false;
- }
-
[MethodImpl(MethodImplOptions.AggressiveInlining)]
private unsafe void* CalculatePointer(int index) =>
(_ptr + _offset + index).ToPointer();
diff --git a/src/Apache.Arrow/Memory/SharedMemoryHandle.cs
b/src/Apache.Arrow/Memory/SharedMemoryHandle.cs
new file mode 100644
index 0000000..be2e14d
--- /dev/null
+++ b/src/Apache.Arrow/Memory/SharedMemoryHandle.cs
@@ -0,0 +1,54 @@
+// Licensed to the Apache Software Foundation (ASF) under one or more
+// contributor license agreements. See the NOTICE file distributed with
+// this work for additional information regarding copyright ownership.
+// The ASF licenses this file to You under the Apache License, Version 2.0
+// (the "License"); you may not use this file except in compliance with
+// the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+using System;
+using System.Buffers;
+
+namespace Apache.Arrow.Memory
+{
+ internal sealed class SharedMemoryHandle : IMemoryOwner<byte>
+ {
+ private SharedMemoryOwner _owner;
+
+ public SharedMemoryHandle(SharedMemoryOwner owner)
+ {
+ _owner = owner ?? throw new ArgumentNullException(nameof(owner));
+ }
+
+ public Memory<byte> Memory => _owner.Memory;
+
+ public SharedMemoryHandle Retain()
+ {
+ return _owner.Retain();
+ }
+
+ public void Dispose()
+ {
+ Release();
+ GC.SuppressFinalize(this);
+ }
+
+ ~SharedMemoryHandle()
+ {
+ Release();
+ }
+
+ private void Release()
+ {
+ _owner?.Release();
+ _owner = null;
+ }
+ }
+}
diff --git a/src/Apache.Arrow/Memory/SharedMemoryOwner.cs
b/src/Apache.Arrow/Memory/SharedMemoryOwner.cs
new file mode 100644
index 0000000..522d5e2
--- /dev/null
+++ b/src/Apache.Arrow/Memory/SharedMemoryOwner.cs
@@ -0,0 +1,62 @@
+// Licensed to the Apache Software Foundation (ASF) under one or more
+// contributor license agreements. See the NOTICE file distributed with
+// this work for additional information regarding copyright ownership.
+// The ASF licenses this file to You under the Apache License, Version 2.0
+// (the "License"); you may not use this file except in compliance with
+// the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+using System;
+using System.Buffers;
+using System.Threading;
+
+namespace Apache.Arrow.Memory
+{
+ internal sealed class SharedMemoryOwner
+ {
+ private readonly IMemoryOwner<byte> _inner;
+ private readonly Memory<byte> _memory;
+ private int _refCount;
+
+ public SharedMemoryOwner(IMemoryOwner<byte> inner)
+ {
+ _inner = inner ?? throw new ArgumentNullException(nameof(inner));
+ _memory = inner.Memory;
+ _refCount = 1;
+ }
+
+ public Memory<byte> Memory => _memory;
+
+ public SharedMemoryHandle Retain()
+ {
+ while (true)
+ {
+ int current = Volatile.Read(ref _refCount);
+ if (current <= 0)
+ {
+ throw new
ObjectDisposedException(nameof(SharedMemoryOwner));
+ }
+
+ if (Interlocked.CompareExchange(ref _refCount, current + 1,
current) == current)
+ {
+ return new SharedMemoryHandle(this);
+ }
+ }
+ }
+
+ public void Release()
+ {
+ if (Interlocked.Decrement(ref _refCount) == 0)
+ {
+ _inner.Dispose();
+ }
+ }
+ }
+}
diff --git a/src/Apache.Arrow/RecordBatch.cs b/src/Apache.Arrow/RecordBatch.cs
index 4067ba9..a9d2fd6 100644
--- a/src/Apache.Arrow/RecordBatch.cs
+++ b/src/Apache.Arrow/RecordBatch.cs
@@ -111,6 +111,17 @@ namespace Apache.Arrow
return new RecordBatch(Schema, _arrays.Select(a =>
ArrowArrayFactory.Slice(a, offset, length)), length);
}
+ public RecordBatch SliceShared(int offset, int length)
+ {
+ if (offset > Length)
+ {
+ throw new ArgumentException($"Offset {offset} cannot be
greater than Length {Length} for RecordBatch.SliceShared");
+ }
+
+ length = Math.Min(Length - offset, length);
+ return new RecordBatch(Schema, _arrays.Select(a =>
ArrowArrayFactory.SliceShared(a, offset, length)), length);
+ }
+
public void Accept(IArrowArrayVisitor visitor)
{
switch (visitor)
diff --git a/test/Apache.Arrow.Tests/ArrayDataReferenceCountingTests.cs
b/test/Apache.Arrow.Tests/ArrayDataReferenceCountingTests.cs
new file mode 100644
index 0000000..6b3a046
--- /dev/null
+++ b/test/Apache.Arrow.Tests/ArrayDataReferenceCountingTests.cs
@@ -0,0 +1,703 @@
+// Licensed to the Apache Software Foundation (ASF) under one or more
+// contributor license agreements. See the NOTICE file distributed with
+// this work for additional information regarding copyright ownership.
+// The ASF licenses this file to You under the Apache License, Version 2.0
+// (the "License"); you may not use this file except in compliance with
+// the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+using System;
+using System.Linq;
+using Apache.Arrow.C;
+using Apache.Arrow.Memory;
+using Apache.Arrow.Types;
+using Xunit;
+
+namespace Apache.Arrow.Tests
+{
+ public class ArrayDataReferenceCountingTests
+ {
+ private static Int32Array BuildInt32Array(params int[] values)
+ {
+ var builder = new Int32Array.Builder();
+ builder.AppendRange(values);
+ return builder.Build();
+ }
+
+ private static Int32Array BuildInt32Array(MemoryAllocator allocator,
params int[] values)
+ {
+ var builder = new Int32Array.Builder();
+ builder.AppendRange(values);
+ return builder.Build(allocator);
+ }
+
+ [Fact]
+ public void AcquireAndDispose_SingleOwner()
+ {
+ var array = BuildInt32Array(1, 2, 3);
+ var data = array.Data;
+
+ // After build, data is usable
+ Assert.Equal(3, data.Length);
+
+ // Dispose releases the data
+ array.Dispose();
+ }
+
+
+ [Fact]
+ public void SliceShared_KeepsParentAlive()
+ {
+ var array = BuildInt32Array(0, 1, 2, 3, 4, 5, 6, 7, 8, 9);
+ var sliced = array.Data.SliceShared(2, 5);
+
+ // Dispose the original — the slice keeps the parent alive
+ array.Dispose();
+
+ // Sliced data should still be usable
+ Assert.Equal(5, sliced.Length);
+ Assert.Equal(2, sliced.Offset);
+ var span = sliced.Buffers[1].Span;
+ Assert.True(span.Length > 0);
+
+ // Disposing the slice releases the parent
+ sliced.Dispose();
+ }
+
+ [Fact]
+ public void SliceShared_OfSliceShared_PointsToRoot()
+ {
+ var array = BuildInt32Array(0, 1, 2, 3, 4, 5, 6, 7, 8, 9);
+ var slice1 = array.Data.SliceShared(2, 6);
+ var slice2 = slice1.SliceShared(1, 3);
+
+ // Dispose original and first slice
+ array.Dispose();
+ slice1.Dispose();
+
+ // Second slice keeps the root alive
+ Assert.Equal(3, slice2.Length);
+ Assert.Equal(3, slice2.Offset); // 2 + 1
+
+ var span = slice2.Buffers[1].Span;
+ Assert.True(span.Length > 0);
+
+ slice2.Dispose();
+ }
+
+ [Fact]
+ public void SliceShared_DisposeSliceFirst_ThenOriginal()
+ {
+ var array = BuildInt32Array(0, 1, 2, 3, 4, 5, 6, 7, 8, 9);
+ var sliced = array.Data.SliceShared(3, 4);
+
+ // Disposing the slice first is fine — it just decrements the
parent ref
+ sliced.Dispose();
+
+ // Original is still valid
+ Assert.Equal(10, array.Length);
+ Assert.Equal(0, array.GetValue(0));
+
+ array.Dispose();
+ }
+
+ [Fact]
+ public void ShareColumnsBetweenRecordBatches()
+ {
+ // Build a record batch with two columns
+ var col1 = BuildInt32Array(1, 2, 3);
+ var col2 = BuildInt32Array(4, 5, 6);
+
+ var schema = new Schema.Builder()
+ .Field(new Field("a", Int32Type.Default, false))
+ .Field(new Field("b", Int32Type.Default, false))
+ .Build();
+
+ var batch1 = new RecordBatch(schema, new IArrowArray[] { col1,
col2 }, 3);
+
+ // Share column "a" into a new batch with a different column "c"
+ var sharedA = batch1.Column(0).Data.SliceShared(0, batch1.Length);
+ var col3 = BuildInt32Array(7, 8, 9);
+
+ var schema2 = new Schema.Builder()
+ .Field(new Field("a", Int32Type.Default, false))
+ .Field(new Field("c", Int32Type.Default, false))
+ .Build();
+
+ var batch2 = new RecordBatch(schema2, new IArrowArray[]
+ {
+ ArrowArrayFactory.BuildArray(sharedA),
+ col3,
+ }, 3);
+
+ // Dispose original batch — shared column should stay alive in
batch2
+ batch1.Dispose();
+
+ // Verify batch2's column "a" is still readable
+ var aArray = (Int32Array)batch2.Column(0);
+ Assert.Equal(3, aArray.Length);
+ Assert.Equal(1, aArray.GetValue(0));
+ Assert.Equal(2, aArray.GetValue(1));
+ Assert.Equal(3, aArray.GetValue(2));
+
+ batch2.Dispose();
+ }
+
+ [Fact]
+ public unsafe void ExportArray_OriginalRemainsValid()
+ {
+ var array = BuildInt32Array(10, 20, 30);
+
+ CArrowArray* cArray = CArrowArray.Create();
+ CArrowArrayExporter.ExportArray(array, cArray);
+
+ // Original array should still be valid after export
+ Assert.Equal(3, array.Length);
+ Assert.Equal(10, array.GetValue(0));
+ Assert.Equal(20, array.GetValue(1));
+ Assert.Equal(30, array.GetValue(2));
+
+ // Import the exported copy and verify it
+ using (var imported =
(Int32Array)CArrowArrayImporter.ImportArray(cArray, array.Data.DataType))
+ {
+ Assert.Equal(3, imported.Length);
+ Assert.Equal(10, imported.GetValue(0));
+ Assert.Equal(20, imported.GetValue(1));
+ Assert.Equal(30, imported.GetValue(2));
+ }
+
+ // Original should still be usable after import is disposed
+ Assert.Equal(10, array.GetValue(0));
+
+ array.Dispose();
+ CArrowArray.Free(cArray);
+ }
+
+ [Fact]
+ public unsafe void ExportRecordBatch_OriginalRemainsValid()
+ {
+ var col1 = BuildInt32Array(1, 2, 3);
+ var col2 = BuildInt32Array(4, 5, 6);
+
+ var schema = new Schema.Builder()
+ .Field(new Field("a", Int32Type.Default, false))
+ .Field(new Field("b", Int32Type.Default, false))
+ .Build();
+
+ var batch = new RecordBatch(schema, new IArrowArray[] { col1, col2
}, 3);
+
+ CArrowArray* cArray = CArrowArray.Create();
+ CArrowArrayExporter.ExportRecordBatch(batch, cArray);
+
+ // Original batch should still be valid
+ Assert.Equal(3, batch.Length);
+ var a = (Int32Array)batch.Column(0);
+ Assert.Equal(1, a.GetValue(0));
+ Assert.Equal(2, a.GetValue(1));
+ Assert.Equal(3, a.GetValue(2));
+
+ // Import and verify the exported copy
+ using (var imported =
CArrowArrayImporter.ImportRecordBatch(cArray, schema))
+ {
+ Assert.Equal(3, imported.Length);
+ var importedA = (Int32Array)imported.Column(0);
+ Assert.Equal(1, importedA.GetValue(0));
+ }
+
+ // Original still usable
+ Assert.Equal(1, ((Int32Array)batch.Column(0)).GetValue(0));
+
+ batch.Dispose();
+ CArrowArray.Free(cArray);
+ }
+
+ [Fact]
+ public unsafe void ExportSlicedArray_OriginalRemainsValid()
+ {
+ var array = BuildInt32Array(0, 1, 2, 3, 4, 5, 6, 7, 8, 9);
+ IArrowArray sliced = array.Slice(2, 6);
+
+ CArrowArray* cArray = CArrowArray.Create();
+ CArrowArrayExporter.ExportArray(sliced, cArray);
+
+ // Original array should still be valid
+ Assert.Equal(10, array.Length);
+ Assert.Equal(0, array.GetValue(0));
+
+ // Import the sliced export
+ using (var imported =
(Int32Array)CArrowArrayImporter.ImportArray(cArray, array.Data.DataType))
+ {
+ Assert.Equal(6, imported.Length);
+ Assert.Equal(2, imported.GetValue(0));
+ }
+
+ sliced.Dispose();
+ array.Dispose();
+ CArrowArray.Free(cArray);
+ }
+
+ [Fact]
+ public unsafe void ExportArray_DisposeOriginalBeforeImportRelease()
+ {
+ // Verify that disposing the original C# array before the C
consumer
+ // releases the export does not cause issues.
+ var array = BuildInt32Array(10, 20, 30);
+
+ CArrowArray* cArray = CArrowArray.Create();
+ CArrowArrayExporter.ExportArray(array, cArray);
+
+ // Dispose the original first
+ array.Dispose();
+
+ // The export should still be valid — the ref count keeps the data
alive
+ using (var imported =
(Int32Array)CArrowArrayImporter.ImportArray(cArray, Int32Type.Default))
+ {
+ Assert.Equal(3, imported.Length);
+ Assert.Equal(10, imported.GetValue(0));
+ Assert.Equal(20, imported.GetValue(1));
+ Assert.Equal(30, imported.GetValue(2));
+ }
+
+ CArrowArray.Free(cArray);
+ }
+
+ // ---------------------------------------------------------------
+ // Array.SliceShared tests
+ // ---------------------------------------------------------------
+
+ [Fact]
+ public void Array_SliceShared_KeepsParentAlive()
+ {
+ var array = BuildInt32Array(0, 1, 2, 3, 4, 5, 6, 7, 8, 9);
+ var sliced = (Int32Array)array.SliceShared(2, 5);
+
+ array.Dispose();
+
+ Assert.Equal(5, sliced.Length);
+ Assert.Equal(2, sliced.GetValue(0));
+ Assert.Equal(6, sliced.GetValue(4));
+
+ sliced.Dispose();
+ }
+
+ [Fact]
+ public void Array_SliceShared_DisposeSliceFirst()
+ {
+ var array = BuildInt32Array(0, 1, 2, 3, 4, 5, 6, 7, 8, 9);
+ var sliced = array.SliceShared(3, 4);
+
+ sliced.Dispose();
+
+ Assert.Equal(10, array.Length);
+ Assert.Equal(0, array.GetValue(0));
+
+ array.Dispose();
+ }
+
+ // ---------------------------------------------------------------
+ // RecordBatch.SliceShared tests
+ // ---------------------------------------------------------------
+
+ [Fact]
+ public void RecordBatch_SliceShared_KeepsParentAlive()
+ {
+ var col1 = BuildInt32Array(0, 1, 2, 3, 4);
+ var col2 = BuildInt32Array(10, 11, 12, 13, 14);
+ var schema = new Schema.Builder()
+ .Field(new Field("a", Int32Type.Default, false))
+ .Field(new Field("b", Int32Type.Default, false))
+ .Build();
+
+ var batch = new RecordBatch(schema, new IArrowArray[] { col1, col2
}, 5);
+ var sliced = batch.SliceShared(1, 3);
+
+ batch.Dispose();
+
+ Assert.Equal(3, sliced.Length);
+ var a = (Int32Array)sliced.Column(0);
+ var b = (Int32Array)sliced.Column(1);
+ Assert.Equal(1, a.GetValue(0));
+ Assert.Equal(3, a.GetValue(2));
+ Assert.Equal(11, b.GetValue(0));
+ Assert.Equal(13, b.GetValue(2));
+
+ sliced.Dispose();
+ }
+
+ [Fact]
+ public void RecordBatch_SliceShared_DisposeSliceFirst()
+ {
+ var col1 = BuildInt32Array(0, 1, 2, 3, 4);
+ var schema = new Schema.Builder()
+ .Field(new Field("a", Int32Type.Default, false))
+ .Build();
+
+ var batch = new RecordBatch(schema, new IArrowArray[] { col1 }, 5);
+ var sliced = batch.SliceShared(1, 2);
+
+ sliced.Dispose();
+
+ Assert.Equal(5, batch.Length);
+ Assert.Equal(0, ((Int32Array)batch.Column(0)).GetValue(0));
+
+ batch.Dispose();
+ }
+
+ // ---------------------------------------------------------------
+ // ChunkedArray.SliceShared tests
+ // ---------------------------------------------------------------
+
+ [Fact]
+ public void ChunkedArray_SliceShared_KeepsParentAlive()
+ {
+ var chunk1 = BuildInt32Array(0, 1, 2);
+ var chunk2 = BuildInt32Array(3, 4, 5);
+ var chunked = new ChunkedArray(new IArrowArray[] { chunk1, chunk2
});
+ var sliced = chunked.SliceShared(1, 4);
+
+ chunk1.Dispose();
+ chunk2.Dispose();
+
+ Assert.Equal(4, sliced.Length);
+ var arr0 = (Int32Array)sliced.ArrowArray(0);
+ Assert.Equal(1, arr0.GetValue(0));
+
+ sliced.Dispose();
+ }
+
+ [Fact]
+ public void ChunkedArray_SliceShared_SingleArgOverload()
+ {
+ var chunk1 = BuildInt32Array(0, 1, 2, 3, 4);
+ var chunked = new ChunkedArray(new IArrowArray[] { chunk1 });
+ var sliced = chunked.SliceShared(2);
+
+ chunk1.Dispose();
+
+ Assert.Equal(3, sliced.Length);
+ var arr = (Int32Array)sliced.ArrowArray(0);
+ Assert.Equal(2, arr.GetValue(0));
+ Assert.Equal(4, arr.GetValue(2));
+
+ sliced.Dispose();
+ }
+
+ // ---------------------------------------------------------------
+ // Column.SliceShared tests
+ // ---------------------------------------------------------------
+
+ [Fact]
+ public void Column_SliceShared_KeepsParentAlive()
+ {
+ var array = BuildInt32Array(0, 1, 2, 3, 4);
+ var field = new Field("x", Int32Type.Default, false);
+ var column = new Column(field, new IArrowArray[] { array });
+ var sliced = column.SliceShared(1, 3);
+
+ array.Dispose();
+
+ Assert.Equal(3, sliced.Length);
+ var arr = (Int32Array)sliced.Data.ArrowArray(0);
+ Assert.Equal(1, arr.GetValue(0));
+ Assert.Equal(3, arr.GetValue(2));
+
+ sliced.Dispose();
+ }
+
+ [Fact]
+ public void Column_SliceShared_SingleArgOverload()
+ {
+ var array = BuildInt32Array(10, 20, 30, 40);
+ var field = new Field("x", Int32Type.Default, false);
+ var column = new Column(field, new IArrowArray[] { array });
+ var sliced = column.SliceShared(2);
+
+ array.Dispose();
+
+ Assert.Equal(2, sliced.Length);
+ var arr = (Int32Array)sliced.Data.ArrowArray(0);
+ Assert.Equal(30, arr.GetValue(0));
+ Assert.Equal(40, arr.GetValue(1));
+
+ sliced.Dispose();
+ }
+
+ // ---------------------------------------------------------------
+ // Tracking-allocator tests: verify no leaks and no double-frees.
+ // TestMemoryAllocator tracks outstanding allocations via Rented,
+ // and throws ObjectDisposedException on double-free.
+ // ---------------------------------------------------------------
+
+ [Fact]
+ public void Tracked_SingleOwner_FreesAll()
+ {
+ var allocator = new TestMemoryAllocator();
+ var array = BuildInt32Array(allocator, 1, 2, 3);
+ Assert.True(allocator.Rented > 0);
+
+ array.Dispose();
+ Assert.Equal(0, allocator.Rented);
+ }
+
+ [Fact]
+ public void Tracked_Acquire_FreesOnlyWhenLastRefDisposed()
+ {
+ var allocator = new TestMemoryAllocator();
+ var array = BuildInt32Array(allocator, 1, 2, 3);
+ var shared = array.Data.SliceShared(0, array.Length);
+
+ array.Dispose();
+ Assert.True(allocator.Rented > 0); // Still held by shared ref
+
+ shared.Dispose();
+ Assert.Equal(0, allocator.Rented);
+ }
+
+ [Fact]
+ public void Tracked_MultipleAcquires_FreesOnlyWhenAllDisposed()
+ {
+ var allocator = new TestMemoryAllocator();
+ var array = BuildInt32Array(allocator, 1, 2, 3);
+ var ref1 = array.Data.SliceShared(0, array.Length);
+ var ref2 = array.Data.SliceShared(0, array.Length);
+
+ array.Dispose();
+ Assert.True(allocator.Rented > 0);
+
+ ref1.Dispose();
+ Assert.True(allocator.Rented > 0);
+
+ ref2.Dispose();
+ Assert.Equal(0, allocator.Rented);
+ }
+
+ [Fact]
+ public void Tracked_SliceShared_FreesWhenSliceDisposed()
+ {
+ var allocator = new TestMemoryAllocator();
+ var array = BuildInt32Array(allocator, 0, 1, 2, 3, 4, 5, 6, 7, 8,
9);
+ var sliced = array.Data.SliceShared(2, 5);
+
+ array.Dispose();
+ Assert.True(allocator.Rented > 0); // Slice keeps parent alive
+
+ sliced.Dispose();
+ Assert.Equal(0, allocator.Rented);
+ }
+
+ [Fact]
+ public void Tracked_SliceShared_DisposeSliceFirst()
+ {
+ var allocator = new TestMemoryAllocator();
+ var array = BuildInt32Array(allocator, 0, 1, 2, 3, 4, 5, 6, 7, 8,
9);
+ var sliced = array.Data.SliceShared(3, 4);
+
+ sliced.Dispose();
+ Assert.True(allocator.Rented > 0); // Original still holds buffers
+
+ array.Dispose();
+ Assert.Equal(0, allocator.Rented);
+ }
+
+ [Fact]
+ public void Tracked_ChainedSliceShared_FreesWhenLastDisposed()
+ {
+ var allocator = new TestMemoryAllocator();
+ var array = BuildInt32Array(allocator, 0, 1, 2, 3, 4, 5, 6, 7, 8,
9);
+ var slice1 = array.Data.SliceShared(2, 6);
+ var slice2 = slice1.SliceShared(1, 3);
+
+ array.Dispose();
+ Assert.True(allocator.Rented > 0);
+
+ slice1.Dispose();
+ Assert.True(allocator.Rented > 0);
+
+ slice2.Dispose();
+ Assert.Equal(0, allocator.Rented);
+ }
+
+ [Fact]
+ public void Tracked_ShareColumnsBetweenBatches_FreesAll()
+ {
+ var allocator = new TestMemoryAllocator();
+ var col1 = BuildInt32Array(allocator, 1, 2, 3);
+ var col2 = BuildInt32Array(allocator, 4, 5, 6);
+
+ var schema = new Schema.Builder()
+ .Field(new Field("a", Int32Type.Default, false))
+ .Field(new Field("b", Int32Type.Default, false))
+ .Build();
+
+ var batch1 = new RecordBatch(schema, new IArrowArray[] { col1,
col2 }, 3);
+
+ var sharedA = batch1.Column(0).Data.SliceShared(0, batch1.Length);
+ var col3 = BuildInt32Array(allocator, 7, 8, 9);
+
+ var schema2 = new Schema.Builder()
+ .Field(new Field("a", Int32Type.Default, false))
+ .Field(new Field("c", Int32Type.Default, false))
+ .Build();
+
+ var batch2 = new RecordBatch(schema2, new IArrowArray[]
+ {
+ ArrowArrayFactory.BuildArray(sharedA),
+ col3,
+ }, 3);
+
+ batch1.Dispose();
+ Assert.True(allocator.Rented > 0); // shared column + col3 still
alive
+
+ batch2.Dispose();
+ Assert.Equal(0, allocator.Rented);
+ }
+
+ [Fact]
+ public unsafe void Tracked_ExportArray_FreesAfterBothSidesDispose()
+ {
+ var allocator = new TestMemoryAllocator();
+ var array = BuildInt32Array(allocator, 10, 20, 30);
+
+ CArrowArray* cArray = CArrowArray.Create();
+ CArrowArrayExporter.ExportArray(array, cArray);
+
+ // Dispose C# side first
+ array.Dispose();
+ Assert.True(allocator.Rented > 0); // Export keeps data alive
+
+ // Release the C side via import + dispose
+ using (var imported = CArrowArrayImporter.ImportArray(cArray,
Int32Type.Default))
+ {
+ Assert.Equal(10, ((Int32Array)imported).GetValue(0));
+ }
+
+ Assert.Equal(0, allocator.Rented);
+ CArrowArray.Free(cArray);
+ }
+
+ [Fact]
+ public unsafe void Tracked_ExportArray_ReleaseExportFirst()
+ {
+ var allocator = new TestMemoryAllocator();
+ var array = BuildInt32Array(allocator, 10, 20, 30);
+
+ CArrowArray* cArray = CArrowArray.Create();
+ CArrowArrayExporter.ExportArray(array, cArray);
+
+ // Release the C side first
+ using (var imported = CArrowArrayImporter.ImportArray(cArray,
Int32Type.Default))
+ {
+ Assert.Equal(10, ((Int32Array)imported).GetValue(0));
+ }
+ Assert.True(allocator.Rented > 0); // C# side still holds data
+
+ // Now dispose the original
+ array.Dispose();
+ Assert.Equal(0, allocator.Rented);
+ CArrowArray.Free(cArray);
+ }
+
+ [Fact]
+ public unsafe void Tracked_ExportRecordBatch_FreesAll()
+ {
+ var allocator = new TestMemoryAllocator();
+ var col1 = BuildInt32Array(allocator, 1, 2, 3);
+ var col2 = BuildInt32Array(allocator, 4, 5, 6);
+
+ var schema = new Schema.Builder()
+ .Field(new Field("a", Int32Type.Default, false))
+ .Field(new Field("b", Int32Type.Default, false))
+ .Build();
+
+ var batch = new RecordBatch(schema, new IArrowArray[] { col1, col2
}, 3);
+
+ CArrowArray* cArray = CArrowArray.Create();
+ CArrowArrayExporter.ExportRecordBatch(batch, cArray);
+
+ batch.Dispose();
+ Assert.True(allocator.Rented > 0);
+
+ using (var imported =
CArrowArrayImporter.ImportRecordBatch(cArray, schema))
+ {
+ Assert.Equal(3, imported.Length);
+ }
+
+ Assert.Equal(0, allocator.Rented);
+ CArrowArray.Free(cArray);
+ }
+ [Fact]
+ public void Tracked_Array_SliceShared_FreesWhenSliceDisposed()
+ {
+ var allocator = new TestMemoryAllocator();
+ var array = BuildInt32Array(allocator, 0, 1, 2, 3, 4, 5, 6, 7, 8,
9);
+ var sliced = array.SliceShared(2, 5);
+
+ array.Dispose();
+ Assert.True(allocator.Rented > 0);
+
+ sliced.Dispose();
+ Assert.Equal(0, allocator.Rented);
+ }
+
+ [Fact]
+ public void Tracked_RecordBatch_SliceShared_FreesWhenSliceDisposed()
+ {
+ var allocator = new TestMemoryAllocator();
+ var col1 = BuildInt32Array(allocator, 0, 1, 2, 3, 4);
+ var col2 = BuildInt32Array(allocator, 10, 11, 12, 13, 14);
+ var schema = new Schema.Builder()
+ .Field(new Field("a", Int32Type.Default, false))
+ .Field(new Field("b", Int32Type.Default, false))
+ .Build();
+
+ var batch = new RecordBatch(schema, new IArrowArray[] { col1, col2
}, 5);
+ var sliced = batch.SliceShared(1, 3);
+
+ batch.Dispose();
+ Assert.True(allocator.Rented > 0);
+
+ sliced.Dispose();
+ Assert.Equal(0, allocator.Rented);
+ }
+
+ [Fact]
+ public void Tracked_ChunkedArray_SliceShared_FreesWhenSliceDisposed()
+ {
+ var allocator = new TestMemoryAllocator();
+ var chunk1 = BuildInt32Array(allocator, 0, 1, 2);
+ var chunk2 = BuildInt32Array(allocator, 3, 4, 5);
+ var chunked = new ChunkedArray(new IArrowArray[] { chunk1, chunk2
});
+ var sliced = chunked.SliceShared(1, 4);
+
+ chunk1.Dispose();
+ chunk2.Dispose();
+ Assert.True(allocator.Rented > 0);
+
+ sliced.Dispose();
+ Assert.Equal(0, allocator.Rented);
+ }
+
+ [Fact]
+ public void Tracked_Column_SliceShared_FreesWhenSliceDisposed()
+ {
+ var allocator = new TestMemoryAllocator();
+ var array = BuildInt32Array(allocator, 0, 1, 2, 3, 4);
+ var field = new Field("x", Int32Type.Default, false);
+ var column = new Column(field, new IArrowArray[] { array });
+ var sliced = column.SliceShared(1, 3);
+
+ array.Dispose();
+ Assert.True(allocator.Rented > 0);
+
+ sliced.Dispose();
+ Assert.Equal(0, allocator.Rented);
+ }
+ }
+}
diff --git a/test/Apache.Arrow.Tests/ArrowBufferTests.cs
b/test/Apache.Arrow.Tests/ArrowBufferTests.cs
index 15ab3d3..28f22c2 100644
--- a/test/Apache.Arrow.Tests/ArrowBufferTests.cs
+++ b/test/Apache.Arrow.Tests/ArrowBufferTests.cs
@@ -14,7 +14,7 @@
// limitations under the License.
using System;
-using System.Runtime.InteropServices;
+using System.Threading;
using Apache.Arrow.Tests.Fixtures;
using Xunit;
@@ -110,5 +110,322 @@ namespace Apache.Arrow.Tests
span[2] = 10;
Assert.Equal(10, buffer.Span.CastTo<int>()[2]);
}
+
+ public class Retain
+ {
+ [Fact]
+ public void RetainedBufferSharesMemory()
+ {
+ ArrowBuffer original = new ArrowBuffer.Builder<int>(3)
+ .Append(1).Append(2).Append(3).Build();
+
+ ArrowBuffer retained = original.Retain();
+
+ Assert.True(original.Span.SequenceEqual(retained.Span));
+ Assert.Equal(original.Length, retained.Length);
+
+ // Verify they share the same underlying memory by checking
pointer identity
+ unsafe
+ {
+ fixed (byte* pOriginal = original.Span)
+ fixed (byte* pRetained = retained.Span)
+ {
+ Assert.True(pOriginal == pRetained);
+ }
+ }
+
+ original.Dispose();
+ retained.Dispose();
+ }
+
+ [Fact]
+ public void RetainedBufferSurvivesOriginalDispose()
+ {
+ ArrowBuffer retained;
+ var expected = new int[] { 10, 20, 30 };
+
+ using (ArrowBuffer original = new ArrowBuffer.Builder<int>(3)
+ .Append(10).Append(20).Append(30).Build())
+ {
+ retained = original.Retain();
+ }
+
+ // Original is disposed, but retained should still be valid
+ var span = retained.Span.CastTo<int>();
+ Assert.Equal(10, span[0]);
+ Assert.Equal(20, span[1]);
+ Assert.Equal(30, span[2]);
+
+ retained.Dispose();
+ }
+
+ [Fact]
+ public void MultipleRetainsAllShareMemory()
+ {
+ ArrowBuffer original = new ArrowBuffer.Builder<byte>(4)
+
.Append(0xAA).Append(0xBB).Append(0xCC).Append(0xDD).Build();
+
+ ArrowBuffer r1 = original.Retain();
+ ArrowBuffer r2 = original.Retain();
+ ArrowBuffer r3 = r1.Retain();
+
+ // All share the same data
+ Assert.True(original.Span.SequenceEqual(r1.Span));
+ Assert.True(original.Span.SequenceEqual(r2.Span));
+ Assert.True(original.Span.SequenceEqual(r3.Span));
+
+ // Dispose in arbitrary order; last one standing should still
work
+ original.Dispose();
+ r2.Dispose();
+
+ Assert.Equal(0xAA, r1.Span[0]);
+ Assert.Equal(0xDD, r3.Span[3]);
+
+ r1.Dispose();
+
+ Assert.Equal(0xDD, r3.Span[3]);
+
+ r3.Dispose();
+ }
+
+ [Fact]
+ public void RetainOnManagedMemorySharesMemory()
+ {
+ byte[] data = { 1, 2, 3, 4 };
+ ArrowBuffer original = new ArrowBuffer(data);
+
+ ArrowBuffer retained = original.Retain();
+
+ Assert.True(original.Span.SequenceEqual(retained.Span));
+
+ // Managed memory buffers share via ReadOnlyMemory, so pointer
identity holds
+ unsafe
+ {
+ fixed (byte* pOriginal = original.Span)
+ fixed (byte* pRetained = retained.Span)
+ {
+ Assert.True(pOriginal == pRetained);
+ }
+ }
+
+ original.Dispose();
+ retained.Dispose();
+ }
+
+ [Fact]
+ public void RetainOnEmptyBuffer()
+ {
+ ArrowBuffer empty = ArrowBuffer.Empty;
+
+ ArrowBuffer retained = empty.Retain();
+
+ Assert.True(retained.IsEmpty);
+ Assert.Equal(0, retained.Length);
+
+ empty.Dispose();
+ retained.Dispose();
+ }
+
+ [Fact]
+ public void ConcurrentRetainAndDispose()
+ {
+ ArrowBuffer original = new ArrowBuffer.Builder<long>(100)
+ .AppendRange(new long[100])
+ .Build();
+
+ const int threadCount = 8;
+ const int iterations = 1000;
+ int errors = 0;
+
+ var threads = new Thread[threadCount];
+ for (int t = 0; t < threadCount; t++)
+ {
+ threads[t] = new Thread(() =>
+ {
+ for (int i = 0; i < iterations; i++)
+ {
+ try
+ {
+ ArrowBuffer retained = original.Retain();
+ _ = retained.Length;
+ retained.Dispose();
+ }
+ catch
+ {
+ Interlocked.Increment(ref errors);
+ }
+ }
+ });
+ }
+
+ foreach (var thread in threads) thread.Start();
+ foreach (var thread in threads) thread.Join();
+
+ Assert.Equal(0, errors);
+ original.Dispose();
+ }
+ }
+
+ public class SliceSharedTests
+ {
+ [Fact]
+ public void SliceSharedSurvivesOriginalDispose()
+ {
+ Int32Array original = new Int32Array.Builder()
+ .AppendRange(new[] { 0, 1, 2, 3, 4, 5, 6, 7, 8, 9 })
+ .Build();
+
+ ArrayData sliced = original.Data.SliceShared(2, 5);
+
+ // Dispose the original array
+ original.Dispose();
+
+ // Sliced data should still be usable
+ var array = new Int32Array(sliced);
+ Assert.Equal(5, array.Length);
+ Assert.Equal(2, array.GetValue(0));
+ Assert.Equal(6, array.GetValue(4));
+
+ sliced.Dispose();
+ }
+
+ [Fact]
+ public void DoubleSliceSharedSurvivesDispose()
+ {
+ Int32Array original = new Int32Array.Builder()
+ .AppendRange(new[] { 0, 1, 2, 3, 4, 5 })
+ .Build();
+
+ ArrayData slice1 = original.Data.SliceShared(1, 4);
+ ArrayData slice2 = slice1.SliceShared(1, 2);
+
+ original.Dispose();
+ slice1.Dispose();
+
+ // Double-sliced data should still work
+ var array = new Int32Array(slice2);
+ Assert.Equal(2, array.Length);
+ Assert.Equal(2, array.GetValue(0));
+ Assert.Equal(3, array.GetValue(1));
+
+ slice2.Dispose();
+ }
+
+ [Fact]
+ public void ArraySliceSharedSurvivesOriginalDispose()
+ {
+ Int32Array original = new Int32Array.Builder()
+ .AppendRange(new[] { 10, 20, 30, 40, 50 })
+ .Build();
+
+ Array sliced = original.SliceShared(1, 3);
+
+ original.Dispose();
+
+ var typedSlice = (Int32Array)sliced;
+ Assert.Equal(3, typedSlice.Length);
+ Assert.Equal(20, typedSlice.GetValue(0));
+ Assert.Equal(40, typedSlice.GetValue(2));
+
+ sliced.Dispose();
+ }
+
+ [Fact]
+ public void RecordBatchSliceSharedSurvivesOriginalDispose()
+ {
+ var batch = new RecordBatch.Builder()
+ .Append("values", false, col => col.Int32(b =>
b.AppendRange(new[] { 0, 1, 2, 3, 4 })))
+ .Build();
+
+ RecordBatch sliced = batch.SliceShared(1, 3);
+
+ batch.Dispose();
+
+ Assert.Equal(3, sliced.Length);
+ var column = (Int32Array)sliced.Column(0);
+ Assert.Equal(1, column.GetValue(0));
+ Assert.Equal(3, column.GetValue(2));
+
+ sliced.Dispose();
+ }
+
+ [Fact]
+ public void NestedListSliceSharedSurvivesOriginalDispose()
+ {
+ var listBuilder = new
ListArray.Builder(Apache.Arrow.Types.Int32Type.Default);
+ var valueBuilder =
(Int32Array.Builder)listBuilder.ValueBuilder;
+
+ // List 0: [10, 20]
+ listBuilder.Append();
+ valueBuilder.Append(10);
+ valueBuilder.Append(20);
+ // List 1: [30]
+ listBuilder.Append();
+ valueBuilder.Append(30);
+ // List 2: [40, 50, 60]
+ listBuilder.Append();
+ valueBuilder.Append(40);
+ valueBuilder.Append(50);
+ valueBuilder.Append(60);
+ // List 3: [70]
+ listBuilder.Append();
+ valueBuilder.Append(70);
+
+ ListArray original = listBuilder.Build();
+ ArrayData sliced = original.Data.SliceShared(1, 2);
+
+ original.Dispose();
+
+ // Sliced data (lists 1 and 2) should still be usable,
including child data
+ var slicedArray = new ListArray(sliced);
+ Assert.Equal(2, slicedArray.Length);
+
+ var list1 = (Int32Array)slicedArray.GetSlicedValues(0);
+ Assert.Equal(1, list1.Length);
+ Assert.Equal(30, list1.GetValue(0));
+
+ var list2 = (Int32Array)slicedArray.GetSlicedValues(1);
+ Assert.Equal(3, list2.Length);
+ Assert.Equal(40, list2.GetValue(0));
+ Assert.Equal(60, list2.GetValue(2));
+
+ sliced.Dispose();
+ }
+
+ [Fact]
+ public void DictionaryArrayRetainSurvivesOriginalDispose()
+ {
+ var dictionaryBuilder = new StringArray.Builder();
+ var indicesBuilder = new Int32Array.Builder();
+
+ dictionaryBuilder.Append("alpha");
+ dictionaryBuilder.Append("beta");
+ dictionaryBuilder.Append("gamma");
+ var dictionary = dictionaryBuilder.Build();
+
+ indicesBuilder.Append(0);
+ indicesBuilder.Append(2);
+ indicesBuilder.Append(1);
+ indicesBuilder.Append(0);
+ var indices = indicesBuilder.Build();
+
+ var dictType = new
Apache.Arrow.Types.DictionaryType(Apache.Arrow.Types.Int32Type.Default,
Apache.Arrow.Types.StringType.Default, false);
+ var originalData = new ArrayData(dictType, indices.Length, 0,
0, indices.Data.Buffers, null, dictionary.Data);
+ var original = new DictionaryArray(originalData);
+
+ ArrayData retained = original.Data.Retain();
+
+ original.Dispose();
+
+ // The retained copy should still have a valid dictionary
+ var retainedArray = new DictionaryArray(retained);
+ Assert.Equal(4, retainedArray.Length);
+ var retainedDict = (StringArray)retainedArray.Dictionary;
+ Assert.Equal("alpha", retainedDict.GetString(0));
+ Assert.Equal("gamma", retainedDict.GetString(2));
+
+ retained.Dispose();
+ }
+ }
}
}
diff --git a/test/Apache.Arrow.Tests/CDataInterfaceDataTests.cs
b/test/Apache.Arrow.Tests/CDataInterfaceDataTests.cs
index 3081160..2823df0 100644
--- a/test/Apache.Arrow.Tests/CDataInterfaceDataTests.cs
+++ b/test/Apache.Arrow.Tests/CDataInterfaceDataTests.cs
@@ -111,6 +111,36 @@ namespace Apache.Arrow.Tests
CArrowArray.Free(cArray);
}
+ [Fact]
+ public unsafe void ExportArrayPreservesSource()
+ {
+ Int32Array array = new Int32Array.Builder()
+ .AppendRange(new[] { 10, 20, 30, 40, 50 })
+ .Build();
+
+ CArrowArray* cArray = CArrowArray.Create();
+ CArrowArrayExporter.ExportArray(array, cArray);
+
+ // Source array should still be valid after export
+ Assert.Equal(5, array.Length);
+ Assert.Equal(10, array.GetValue(0));
+ Assert.Equal(50, array.GetValue(4));
+
+ // Import and verify the exported copy is also valid
+ using (var imported =
(Int32Array)CArrowArrayImporter.ImportArray(cArray, array.Data.DataType))
+ {
+ Assert.Equal(5, imported.Length);
+ Assert.Equal(10, imported.GetValue(0));
+ Assert.Equal(50, imported.GetValue(4));
+ }
+
+ // Source should still be valid after the imported copy is disposed
+ Assert.Equal(30, array.GetValue(2));
+
+ CArrowArray.Free(cArray);
+ array.Dispose();
+ }
+
[Fact]
public unsafe void ExportRecordBatch_LargerThan2GB_Succeeds()
{
diff --git a/test/Apache.Arrow.Tests/CDataInterfacePythonTests.cs
b/test/Apache.Arrow.Tests/CDataInterfacePythonTests.cs
index 6a10255..427ce44 100644
--- a/test/Apache.Arrow.Tests/CDataInterfacePythonTests.cs
+++ b/test/Apache.Arrow.Tests/CDataInterfacePythonTests.cs
@@ -740,58 +740,55 @@ namespace Apache.Arrow.Tests
[SkippableFact]
public unsafe void ExportManagedMemoryArray()
{
- using (new EnableManagedExport())
+ var expectedValues = Enumerable.Range(0, 100).Select(i => i % 10
== 3 ? null : (long?)i).ToArray();
+ var gcRefs = new List<WeakReference>();
+
+ void TestExport()
{
- var expectedValues = Enumerable.Range(0, 100).Select(i => i %
10 == 3 ? null : (long?)i).ToArray();
- var gcRefs = new List<WeakReference>();
+ var array = CreateManagedMemoryArray(expectedValues, gcRefs);
- void TestExport()
+ dynamic pyArray;
+ using (Py.GIL())
{
- var array = CreateManagedMemoryArray(expectedValues,
gcRefs);
-
- dynamic pyArray;
- using (Py.GIL())
- {
- dynamic pa = Py.Import("pyarrow");
- pyArray = pa.array(expectedValues);
- }
-
- CArrowArray* cArray = CArrowArray.Create();
- CArrowArrayExporter.ExportArray(array, cArray);
+ dynamic pa = Py.Import("pyarrow");
+ pyArray = pa.array(expectedValues);
+ }
- CArrowSchema* cSchema = CArrowSchema.Create();
- CArrowSchemaExporter.ExportType(array.Data.DataType,
cSchema);
+ CArrowArray* cArray = CArrowArray.Create();
+ CArrowArrayExporter.ExportArray(array, cArray);
- GcCollect();
- foreach (var weakRef in gcRefs)
- {
- Assert.True(weakRef.IsAlive);
- }
+ CArrowSchema* cSchema = CArrowSchema.Create();
+ CArrowSchemaExporter.ExportType(array.Data.DataType, cSchema);
- long arrayPtr = ((IntPtr)cArray).ToInt64();
- long schemaPtr = ((IntPtr)cSchema).ToInt64();
+ GcCollect();
+ foreach (var weakRef in gcRefs)
+ {
+ Assert.True(weakRef.IsAlive);
+ }
- using (Py.GIL())
- {
- dynamic pa = Py.Import("pyarrow");
- dynamic exportedPyArray =
pa.Array._import_from_c(arrayPtr, schemaPtr);
- Assert.True(exportedPyArray == pyArray);
+ long arrayPtr = ((IntPtr)cArray).ToInt64();
+ long schemaPtr = ((IntPtr)cSchema).ToInt64();
- // Required for the Python object to be garbage
collected:
- exportedPyArray.Dispose();
- }
+ using (Py.GIL())
+ {
+ dynamic pa = Py.Import("pyarrow");
+ dynamic exportedPyArray =
pa.Array._import_from_c(arrayPtr, schemaPtr);
+ Assert.True(exportedPyArray == pyArray);
- CArrowArray.Free(cArray);
- CArrowSchema.Free(cSchema);
+ // Required for the Python object to be garbage collected:
+ exportedPyArray.Dispose();
}
- TestExport();
+ CArrowArray.Free(cArray);
+ CArrowSchema.Free(cSchema);
+ }
- GcCollect();
- foreach (var weakRef in gcRefs)
- {
- Assert.False(weakRef.IsAlive);
- }
+ TestExport();
+
+ GcCollect();
+ foreach (var weakRef in gcRefs)
+ {
+ Assert.False(weakRef.IsAlive);
}
}
@@ -1009,7 +1006,6 @@ namespace Apache.Arrow.Tests
var originalBatch = GetTestRecordBatch();
dynamic pyBatch = GetPythonRecordBatch();
- using (new EnableManagedExport())
using (var stream = new MemoryStream())
{
var writer = new ArrowStreamWriter(stream,
originalBatch.Schema);
@@ -1061,7 +1057,6 @@ namespace Apache.Arrow.Tests
var originalBatch = GetTestRecordBatch();
- using (new EnableManagedExport())
using (var stream = new MemoryStream())
{
var writer = new ArrowStreamWriter(stream,
originalBatch.Schema);
@@ -1391,21 +1386,5 @@ namespace Apache.Arrow.Tests
_index = -1;
}
}
-
- sealed class EnableManagedExport : IDisposable
- {
- readonly bool _previousValue;
-
- public EnableManagedExport()
- {
- _previousValue = CArrowArrayExporter.EnableManagedMemoryExport;
- CArrowArrayExporter.EnableManagedMemoryExport = true;
- }
-
- public void Dispose()
- {
- CArrowArrayExporter.EnableManagedMemoryExport = _previousValue;
- }
- }
}
}
