Repository: apache/arrow
Pull request #46190: GH-46189 [C#] Use pooled buffers in ArrowStreamWriter
Author: login: m-v-w
Labels: name: Component: C#
Created: 2025-04-21T21:07:02Z | Merged: 2025-05-27T13:33:47Z

### Rationale for this change
`ArrowStreamWriter` should not make large heap allocations for temporary buffers. Pooled buffers should be used. 
All allocations made by `MemoryAllocator` should be disposed before they are finalized.

### What changes are included in this PR?

- `ArrowStreamWriter` makes sure all temporary buffers are disposed.
- `ArrowStreamWriter` uses `MemoryAllocator` for decompression buffers instead of `MemoryStream`.
- Added `TryCompress` method to `ICompressionCodec` and implemented it.

### Are these changes tested?

Yes. 
- `MemoryOwnerDisposal` in `Apache.Arrow.Compression.Tests.ArrowStreamWriterTests`
- `MemoryOwnerDisposalSlicedArray` in `Apache.Arrow.Tests.ArrowStreamWriterTests`

### Are there any user-facing changes?

Yes. `ICompressionCodec` requires the implementation of a new method 'TryCompress'

**This PR includes breaking changes to public APIs.**

`ICompressionCodec` requires the implementation of a new method 'TryCompress'

* GitHub Issue: #46189


Diff:
diff --git a/csharp/src/Apache.Arrow.Compression/Lz4CompressionCodec.cs b/csharp/src/Apache.Arrow.Compression/Lz4CompressionCodec.cs
index df19c16a3021..ec48d8990902 100644
--- a/csharp/src/Apache.Arrow.Compression/Lz4CompressionCodec.cs
+++ b/csharp/src/Apache.Arrow.Compression/Lz4CompressionCodec.cs
@@ -33,7 +33,7 @@ public Lz4CompressionCodec(int? compressionLevel = null)
                 {
                     _settings = new LZ4EncoderSettings
                     {
-                        CompressionLevel = (LZ4Level) compressionLevel,
+                        CompressionLevel = (LZ4Level) compressionLevel
                     };
                 }
                 else
diff --git a/csharp/src/Apache.Arrow.Compression/ZstdCompressionCodec.cs b/csharp/src/Apache.Arrow.Compression/ZstdCompressionCodec.cs
index cc340a7cd1b9..d89480ebfac4 100644
--- a/csharp/src/Apache.Arrow.Compression/ZstdCompressionCodec.cs
+++ b/csharp/src/Apache.Arrow.Compression/ZstdCompressionCodec.cs
@@ -20,7 +20,7 @@
 
 namespace Apache.Arrow.Compression
 {
-    internal sealed class ZstdCompressionCodec : ICompressionCodec
+    internal sealed class ZstdCompressionCodec : ITryCompressionCodec
     {
         private readonly Decompressor _decompressor;
         private readonly Compressor _compressor;
@@ -52,6 +52,11 @@ public void Compress(ReadOnlyMemory<byte> source, Stream destination)
             compressor.Write(source.Span);
         }
 
+        public bool TryCompress(ReadOnlyMemory<byte> source, Memory<byte> destination, out int bytesWritten)
+        {
+            return _compressor.TryWrap(source.Span, destination.Span, out bytesWritten);
+        }
+
         public void Dispose()
         {
             _decompressor.Dispose();
diff --git a/csharp/src/Apache.Arrow/Ipc/ArrowStreamWriter.cs b/csharp/src/Apache.Arrow/Ipc/ArrowStreamWriter.cs
index ff0b64c09eeb..f31709382240 100644
--- a/csharp/src/Apache.Arrow/Ipc/ArrowStreamWriter.cs
+++ b/csharp/src/Apache.Arrow/Ipc/ArrowStreamWriter.cs
@@ -87,36 +87,38 @@ public FieldNode(int length, int nullCount)
                 }
             }
 
-            public readonly struct Buffer
+            public readonly struct Buffer : IDisposable
             {
                 public readonly ReadOnlyMemory<byte> DataBuffer;
                 public readonly int Offset;
+                private readonly IDisposable _owner;
 
-                public Buffer(ReadOnlyMemory<byte> buffer, int offset)
+                public Buffer(ReadOnlyMemory<byte> buffer, int offset, IDisposable owner)
                 {
                     DataBuffer = buffer;
                     Offset = offset;
+                    _owner = owner;
                 }
+                public void Dispose() => _owner?.Dispose();
             }
 
             private readonly List<FieldNode> _fieldNodes;
             private readonly List<Buffer> _buffers;
             private readonly ICompressionCodec _compressionCodec;
             private readonly MemoryAllocator _allocator;
-            private readonly MemoryStream _compressionStream;
-
+            private readonly MemoryStream _fallbackCompressionStream;
             public IReadOnlyList<FieldNode> FieldNodes => _fieldNodes;
             public IReadOnlyList<Buffer> Buffers => _buffers;
 
-            public List<long> VariadicCounts { get; private set; } 
+            public List<long> VariadicCounts { get; private set; }
             public int TotalLength { get; private set; }
 
             public ArrowRecordBatchFlatBufferBuilder(
-                ICompressionCodec compressionCodec, MemoryAllocator allocator, MemoryStream compressionStream)
+                ICompressionCodec compressionCodec, MemoryAllocator allocator, MemoryStream fallbackCompressionStream)
             {
                 _compressionCodec = compressionCodec;
-                _compressionStream = compressionStream;
                 _allocator = allocator;
+                _fallbackCompressionStream = fallbackCompressionStream;
                 _fieldNodes = new List<FieldNode>();
                 _buffers = new List<Buffer>();
                 TotalLength = 0;
@@ -155,20 +157,20 @@ public void VisitArray(IArrowArray array)
             private void VisitPrimitiveArray<T>(PrimitiveArray<T> array)
                 where T : struct, IEquatable<T>
             {
-                _buffers.Add(CreateBitmapBuffer(array.NullBitmapBuffer, array.Offset, array.Length));
-                _buffers.Add(CreateSlicedBuffer<T>(array.ValueBuffer, array.Offset, array.Length));
+                _buffers.Add(CreateBitmapBuffer(array.NullBitmapBuffer, array.Offset, array.Length, false));
+                _buffers.Add(CreateSlicedBuffer<T>(array.ValueBuffer, array.Offset, array.Length, false));
             }
 
             public void Visit(BooleanArray array)
             {
-                _buffers.Add(CreateBitmapBuffer(array.NullBitmapBuffer, array.Offset, array.Length));
-                _buffers.Add(CreateBitmapBuffer(array.ValueBuffer, array.Offset, array.Length));
+                _buffers.Add(CreateBitmapBuffer(array.NullBitmapBuffer, array.Offset, array.Length, false));
+                _buffers.Add(CreateBitmapBuffer(array.ValueBuffer, array.Offset, array.Length, false));
             }
 
             public void Visit(ListArray array)
             {
-                _buffers.Add(CreateBitmapBuffer(array.NullBitmapBuffer, array.Offset, array.Length));
-                _buffers.Add(CreateBuffer(GetZeroBasedValueOffsets(array.ValueOffsetsBuffer, array.Offset, array.Length)));
+                _buffers.Add(CreateBitmapBuffer(array.NullBitmapBuffer, array.Offset, array.Length, false));
+                _buffers.Add(CreateBuffer(GetZeroBasedValueOffsets(array.ValueOffsetsBuffer, array.Offset, array.Length), true));
 
                 int valuesOffset = 0;
                 int valuesLength = 0;
@@ -191,9 +193,9 @@ public void Visit(ListViewArray array)
             {
                 var (valueOffsetsBuffer, minOffset, maxEnd) = GetZeroBasedListViewOffsets(array);
 
-                _buffers.Add(CreateBitmapBuffer(array.NullBitmapBuffer, array.Offset, array.Length));
-                _buffers.Add(CreateBuffer(valueOffsetsBuffer));
-                _buffers.Add(CreateSlicedBuffer<int>(array.SizesBuffer, array.Offset, array.Length));
+                _buffers.Add(CreateBitmapBuffer(array.NullBitmapBuffer, array.Offset, array.Length, false));
+                _buffers.Add(CreateBuffer(valueOffsetsBuffer, true));
+                _buffers.Add(CreateSlicedBuffer<int>(array.SizesBuffer, array.Offset, array.Length, false));
 
                 IArrowArray values = array.Values;
                 if (minOffset != 0 || values.Length != maxEnd)
@@ -206,8 +208,8 @@ public void Visit(ListViewArray array)
 
             public void Visit(LargeListArray array)
             {
-                _buffers.Add(CreateBitmapBuffer(array.NullBitmapBuffer, array.Offset, array.Length));
-                _buffers.Add(CreateBuffer(GetZeroBasedLongValueOffsets(array.ValueOffsetsBuffer, array.Offset, array.Length)));
+                _buffers.Add(CreateBitmapBuffer(array.NullBitmapBuffer, array.Offset, array.Length, false));
+                _buffers.Add(CreateBuffer(GetZeroBasedLongValueOffsets(array.ValueOffsetsBuffer, array.Offset, array.Length), true));
 
                 int valuesOffset = 0;
                 int valuesLength = 0;
@@ -228,7 +230,7 @@ public void Visit(LargeListArray array)
 
             public void Visit(FixedSizeListArray array)
             {
-                _buffers.Add(CreateBitmapBuffer(array.NullBitmapBuffer, array.Offset, array.Length));
+                _buffers.Add(CreateBitmapBuffer(array.NullBitmapBuffer, array.Offset, array.Length, false));
 
                 var listSize = ((FixedSizeListType)array.Data.DataType).ListSize;
                 var valuesSlice =
@@ -245,8 +247,8 @@ public void Visit(FixedSizeListArray array)
 
             public void Visit(BinaryArray array)
             {
-                _buffers.Add(CreateBitmapBuffer(array.NullBitmapBuffer, array.Offset, array.Length));
-                _buffers.Add(CreateBuffer(GetZeroBasedValueOffsets(array.ValueOffsetsBuffer, array.Offset, array.Length)));
+                _buffers.Add(CreateBitmapBuffer(array.NullBitmapBuffer, array.Offset, array.Length, false));
+                _buffers.Add(CreateBuffer(GetZeroBasedValueOffsets(array.ValueOffsetsBuffer, array.Offset, array.Length), true));
 
                 int valuesOffset = 0;
                 int valuesLength = 0;
@@ -256,16 +258,16 @@ public void Visit(BinaryArray array)
                     valuesLength = array.ValueOffsets[array.Length] - valuesOffset;
                 }
 
-                _buffers.Add(CreateSlicedBuffer<byte>(array.ValueBuffer, valuesOffset, valuesLength));
+                _buffers.Add(CreateSlicedBuffer<byte>(array.ValueBuffer, valuesOffset, valuesLength, false));
             }
 
             public void Visit(BinaryViewArray array)
             {
-                _buffers.Add(CreateBitmapBuffer(array.NullBitmapBuffer, array.Offset, array.Length));
-                _buffers.Add(CreateSlicedBuffer<Scalars.BinaryView>(array.ViewsBuffer, array.Offset, array.Length));
+                _buffers.Add(CreateBitmapBuffer(array.NullBitmapBuffer, array.Offset, array.Length, false));
+                _buffers.Add(CreateSlicedBuffer<Scalars.BinaryView>(array.ViewsBuffer, array.Offset, array.Length, false));
                 for (int i = 0; i < array.DataBufferCount; i++)
                 {
-                    _buffers.Add(CreateBuffer(array.DataBuffer(i)));
+                    _buffers.Add(CreateBuffer(array.DataBuffer(i), false));
                 }
                 VariadicCounts = VariadicCounts ?? new List<long>();
                 VariadicCounts.Add(array.DataBufferCount);
@@ -273,8 +275,8 @@ public void Visit(BinaryViewArray array)
 
             public void Visit(LargeBinaryArray array)
             {
-                _buffers.Add(CreateBitmapBuffer(array.NullBitmapBuffer, array.Offset, array.Length));
-                _buffers.Add(CreateBuffer(GetZeroBasedLongValueOffsets(array.ValueOffsetsBuffer, array.Offset, array.Length)));
+                _buffers.Add(CreateBitmapBuffer(array.NullBitmapBuffer, array.Offset, array.Length, false));
+                _buffers.Add(CreateBuffer(GetZeroBasedLongValueOffsets(array.ValueOffsetsBuffer, array.Offset, array.Length), true));
 
                 int valuesOffset = 0;
                 int valuesLength = 0;
@@ -284,14 +286,14 @@ public void Visit(LargeBinaryArray array)
                     valuesLength = checked((int)array.ValueOffsets[array.Length]) - valuesOffset;
                 }
 
-                _buffers.Add(CreateSlicedBuffer<byte>(array.ValueBuffer, valuesOffset, valuesLength));
+                _buffers.Add(CreateSlicedBuffer<byte>(array.ValueBuffer, valuesOffset, valuesLength, false));
             }
 
             public void Visit(FixedSizeBinaryArray array)
             {
                 var itemSize = ((FixedSizeBinaryType)array.Data.DataType).ByteWidth;
-                _buffers.Add(CreateBitmapBuffer(array.NullBitmapBuffer, array.Offset, array.Length));
-                _buffers.Add(CreateSlicedBuffer(array.ValueBuffer, itemSize, array.Offset, array.Length));
+                _buffers.Add(CreateBitmapBuffer(array.NullBitmapBuffer, array.Offset, array.Length, false));
+                _buffers.Add(CreateSlicedBuffer(array.ValueBuffer, itemSize, array.Offset, array.Length, false));
             }
 
             public void Visit(Decimal32Array array) => Visit(array as FixedSizeBinaryArray);
@@ -304,7 +306,7 @@ public void Visit(FixedSizeBinaryArray array)
 
             public void Visit(StructArray array)
             {
-                _buffers.Add(CreateBitmapBuffer(array.NullBitmapBuffer, array.Offset, array.Length));
+                _buffers.Add(CreateBitmapBuffer(array.NullBitmapBuffer, array.Offset, array.Length, false));
 
                 for (int i = 0; i < array.Fields.Count; i++)
                 {
@@ -315,12 +317,12 @@ public void Visit(StructArray array)
 
             public void Visit(UnionArray array)
             {
-                _buffers.Add(CreateSlicedBuffer<byte>(array.TypeBuffer, array.Offset, array.Length));
+                _buffers.Add(CreateSlicedBuffer<byte>(array.TypeBuffer, array.Offset, array.Length, false));
 
                 ArrowBuffer? offsets = (array as DenseUnionArray)?.ValueOffsetBuffer;
                 if (offsets != null)
                 {
-                    _buffers.Add(CreateSlicedBuffer<int>(offsets.Value, array.Offset, array.Length));
+                    _buffers.Add(CreateSlicedBuffer<int>(offsets.Value, array.Offset, array.Length, false));
                 }
 
                 for (int i = 0; i < array.Fields.Count; i++)
@@ -351,7 +353,7 @@ private ArrowBuffer GetZeroBasedValueOffsets(ArrowBuffer valueOffsetsBuffer, int
                 {
                     // Array has been sliced, so we need to shift and adjust the offsets
                     var originalOffsets = valueOffsetsBuffer.Span.CastTo<int>().Slice(arrayOffset, arrayLength + 1);
-                    var firstOffset = arrayLength > 0 ? originalOffsets[0] : 0;
+                    int firstOffset = arrayLength > 0 ? originalOffsets[0] : 0;
 
                     var newValueOffsetsBuffer = _allocator.Allocate(requiredBytes);
                     var newValueOffsets = newValueOffsetsBuffer.Memory.Span.CastTo<int>();
@@ -371,8 +373,8 @@ private ArrowBuffer GetZeroBasedValueOffsets(ArrowBuffer valueOffsetsBuffer, int
                 }
                 else
                 {
-                    // Use the full buffer
-                    return valueOffsetsBuffer;
+                    // Use the full buffer, but create a copy that is externally owned
+                    return new ArrowBuffer(valueOffsetsBuffer.Memory);
                 }
             }
 
@@ -393,7 +395,6 @@ private ArrowBuffer GetZeroBasedLongValueOffsets(ArrowBuffer valueOffsetsBuffer,
                     {
                         newValueOffsets[i] = originalOffsets[i] - firstOffset;
                     }
-
                     return new ArrowBuffer(newValueOffsetsBuffer);
                 }
                 else if (valueOffsetsBuffer.Length > requiredBytes)
@@ -404,8 +405,8 @@ private ArrowBuffer GetZeroBasedLongValueOffsets(ArrowBuffer valueOffsetsBuffer,
                 }
                 else
                 {
-                    // Use the full buffer
-                    return valueOffsetsBuffer;
+                    // Use the full buffer, create new buffer that is externally owned
+                    return new ArrowBuffer(valueOffsetsBuffer.Memory);
                 }
             }
 
@@ -439,12 +440,17 @@ private ArrowBuffer GetZeroBasedLongValueOffsets(ArrowBuffer valueOffsetsBuffer,
                 if (minOffset == 0)
                 {
                     // No need to adjust the offsets, but we may need to slice the offsets buffer.
-                    ArrowBuffer buffer = array.ValueOffsetsBuffer;
-                    if (array.Offset != 0 || buffer.Length > requiredBytes)
+                    ArrowBuffer buffer;
+                    if (array.Offset != 0 || array.ValueOffsetsBuffer.Length > requiredBytes)
                     {
                         var byteOffset = sizeof(int) * array.Offset;
-                        var sliceLength = Math.Min(requiredBytes, buffer.Length - byteOffset);
-                        buffer = new ArrowBuffer(buffer.Memory.Slice(byteOffset, sliceLength));
+                        var sliceLength = Math.Min(requiredBytes, array.ValueOffsetsBuffer.Length - byteOffset);
+                        buffer = new ArrowBuffer(array.ValueOffsetsBuffer.Memory.Slice(byteOffset, sliceLength));
+                    }
+                    else
+                    {
+                        // create a copy that's externally owned
+                        buffer = new ArrowBuffer(array.ValueOffsetsBuffer.Memory);
                     }
 
                     return (buffer, minOffset, maxEnd);
@@ -461,11 +467,11 @@ private ArrowBuffer GetZeroBasedLongValueOffsets(ArrowBuffer valueOffsetsBuffer,
                 return (new ArrowBuffer(newOffsetsBuffer), minOffset, maxEnd);
             }
 
-            private Buffer CreateBitmapBuffer(ArrowBuffer buffer, int offset, int length)
+            private Buffer CreateBitmapBuffer(ArrowBuffer buffer, int offset, int length, bool locallyOwned)
             {
                 if (buffer.IsEmpty)
                 {
-                    return CreateBuffer(buffer.Memory);
+                    return CreateBuffer(buffer, locallyOwned);
                 }
 
                 var paddedLength = CalculatePaddedBufferLength(BitUtility.ByteCount(length));
@@ -474,7 +480,7 @@ private Buffer CreateBitmapBuffer(ArrowBuffer buffer, int offset, int length)
                     var byteOffset = offset / 8;
                     var sliceLength = Math.Min(paddedLength, buffer.Length - byteOffset);
 
-                    return CreateBuffer(buffer.Memory.Slice(byteOffset, sliceLength));
+                    return CreateBuffer(buffer.Memory.Slice(byteOffset, sliceLength), locallyOwned ? buffer : null);
                 }
                 else
                 {
@@ -487,17 +493,17 @@ private Buffer CreateBitmapBuffer(ArrowBuffer buffer, int offset, int length)
                         BitUtility.SetBit(outputSpan, i, BitUtility.GetBit(inputSpan, offset + i));
                     }
 
-                    return CreateBuffer(memoryOwner.Memory);
+                    return CreateBuffer(memoryOwner);
                 }
             }
 
-            private Buffer CreateSlicedBuffer<T>(ArrowBuffer buffer, int offset, int length)
+            private Buffer CreateSlicedBuffer<T>(ArrowBuffer buffer, int offset, int length, bool owned)
                 where T : struct
             {
-                return CreateSlicedBuffer(buffer, Unsafe.SizeOf<T>(), offset, length);
+                return CreateSlicedBuffer(buffer, Unsafe.SizeOf<T>(), offset, length, owned);
             }
 
-            private Buffer CreateSlicedBuffer(ArrowBuffer buffer, int itemSize, int offset, int length)
+            private Buffer CreateSlicedBuffer(ArrowBuffer buffer, int itemSize, int offset, int length, bool owned)
             {
                 var byteLength = length * itemSize;
                 var paddedLength = CalculatePaddedBufferLength(byteLength);
@@ -505,23 +511,35 @@ private Buffer CreateSlicedBuffer(ArrowBuffer buffer, int itemSize, int offset,
                 {
                     var byteOffset = offset * itemSize;
                     var sliceLength = Math.Min(paddedLength, buffer.Length - byteOffset);
-                    return CreateBuffer(buffer.Memory.Slice(byteOffset, sliceLength));
+                    return CreateBuffer(buffer.Memory.Slice(byteOffset, sliceLength),
+                        owned ? buffer : null);
                 }
 
-                return CreateBuffer(buffer.Memory);
+                return CreateBuffer(buffer, owned);
             }
 
-            private Buffer CreateBuffer(ArrowBuffer buffer)
+            private Buffer CreateBuffer(ArrowBuffer buffer, bool locallyOwned)
             {
-                return CreateBuffer(buffer.Memory);
+                if (locallyOwned)
+                {
+                    // this buffer is locally owned, we can dispose it
+                    return CreateBuffer(buffer.Memory, buffer);
+                }
+                return CreateBuffer(buffer.Memory, null);
             }
 
-            private Buffer CreateBuffer(ReadOnlyMemory<byte> buffer)
+            private Buffer CreateBuffer(IMemoryOwner<byte>  bufferOwner)
+            {
+                return CreateBuffer(bufferOwner.Memory, bufferOwner);
+            }
+
+            private Buffer CreateBuffer(ReadOnlyMemory<byte> buffer, IDisposable localBufferOwner)
             {
                 int offset = TotalLength;
                 const int UncompressedLengthSize = 8;
 
                 ReadOnlyMemory<byte> bufferToWrite;
+                IMemoryOwner<byte> bufferOwner=null;
                 if (_compressionCodec == null)
                 {
                     bufferToWrite = buffer;
@@ -529,40 +547,60 @@ private Buffer CreateBuffer(ReadOnlyMemory<byte> buffer)
                 else if (buffer.Length == 0)
                 {
                     // Write zero length and skip compression
-                    var uncompressedLengthBytes = _allocator.Allocate(UncompressedLengthSize);
-                    BinaryPrimitives.WriteInt64LittleEndian(uncompressedLengthBytes.Memory.Span, 0);
-                    bufferToWrite = uncompressedLengthBytes.Memory;
+                    bufferOwner = _allocator.Allocate(UncompressedLengthSize);
+                    BinaryPrimitives.WriteInt64LittleEndian(bufferOwner.Memory.Span, 0);
+                    bufferToWrite = bufferOwner.Memory.Slice(0, UncompressedLengthSize);
+                    // the local source buffer owner can be disposed, its memory is no longer needed
+                    localBufferOwner?.Dispose();
                 }
                 else
                 {
                     // See format/Message.fbs, and the BUFFER BodyCompressionMethod for documentation on how
                     // compressed buffers are stored.
-                    _compressionStream.Seek(0, SeekOrigin.Begin);
-                    _compressionStream.SetLength(0);
-                    _compressionCodec.Compress(buffer, _compressionStream);
-                    if (_compressionStream.Length < buffer.Length)
+                    int newBufferLength = UncompressedLengthSize + buffer.Length;
+                    bufferOwner = _allocator.Allocate(newBufferLength);
+
+                    if(TryCompress(buffer, bufferOwner.Memory.Slice(UncompressedLengthSize, buffer.Length), out int bytesWritten))
                     {
-                        var newBuffer = _allocator.Allocate((int) _compressionStream.Length + UncompressedLengthSize);
-                        BinaryPrimitives.WriteInt64LittleEndian(newBuffer.Memory.Span, buffer.Length);
-                        _compressionStream.Seek(0, SeekOrigin.Begin);
-                        _compressionStream.ReadFullBuffer(newBuffer.Memory.Slice(UncompressedLengthSize));
-                        bufferToWrite = newBuffer.Memory;
+                        // Write the uncompressed length to the start of the buffer
+                        BinaryPrimitives.WriteInt64LittleEndian(bufferOwner.Memory.Span, buffer.Length);
+                        bufferToWrite = bufferOwner.Memory.Slice(0, bytesWritten+UncompressedLengthSize);
                     }
                     else
                     {
+                        // TryCompress failed, because the buffer is too small.
                         // If the compressed buffer is larger than the uncompressed buffer, use the uncompressed
                         // buffer instead, and indicate this by setting the uncompressed length to -1
-                        var newBuffer = _allocator.Allocate(buffer.Length + UncompressedLengthSize);
-                        BinaryPrimitives.WriteInt64LittleEndian(newBuffer.Memory.Span, -1);
-                        buffer.CopyTo(newBuffer.Memory.Slice(UncompressedLengthSize));
-                        bufferToWrite = newBuffer.Memory;
+                        BinaryPrimitives.WriteInt64LittleEndian(bufferOwner.Memory.Span, -1);
+                        buffer.CopyTo(bufferOwner.Memory.Slice(UncompressedLengthSize));
+                        bufferToWrite = bufferOwner.Memory.Slice(0, newBufferLength);
                     }
+                    // the local source buffer owner can be disposed, its memory was copied
+                    localBufferOwner?.Dispose();
                 }
 
                 int paddedLength = checked((int)BitUtility.RoundUpToMultipleOf8(bufferToWrite.Length));
                 TotalLength += paddedLength;
+                // if the buffer is owned locally or allocated in the method, we can dispose it after writing the buffer to the stream
+                return new Buffer(bufferToWrite, offset, (IDisposable)bufferOwner ?? localBufferOwner);
+            }
 
-                return new Buffer(bufferToWrite, offset);
+            private bool TryCompress(ReadOnlyMemory<byte> source, Memory<byte> destination, out int bytesWritten)
+            {
+                if(_compressionCodec is ITryCompressionCodec tryCompressionCodec)
+                    return tryCompressionCodec.TryCompress(source, destination, out bytesWritten);
+                // Fallback to using a memory stream for compression
+                _fallbackCompressionStream.Seek(0, SeekOrigin.Begin);
+                _fallbackCompressionStream.SetLength(0);
+                _compressionCodec.Compress(source, _fallbackCompressionStream);
+                if (_fallbackCompressionStream.Length >= destination.Length)
+                {
+                    bytesWritten = 0;
+                    return false;
+                }
+                _fallbackCompressionStream.Seek(0, SeekOrigin.Begin);
+                bytesWritten = _fallbackCompressionStream.ReadFullBuffer(destination);
+                return true;
             }
 
             public void Visit(IArrowArray array)
@@ -590,8 +628,9 @@ public void Visit(IArrowArray array)
         private readonly bool _leaveOpen;
         private readonly IpcOptions _options;
         private readonly MemoryAllocator _allocator;
-        // Reuse a single memory stream for writing compressed data to, to reduce memory allocations
-        private readonly MemoryStream _compressionStream = new MemoryStream();
+        // Fallback for compression codec, not implementing TryCompress.
+        // Reuse a single memory stream for writing compressed
+        private readonly MemoryStream _fallbackCompressionStream = new MemoryStream();
 
         private protected const Flatbuf.MetadataVersion CurrentMetadataVersion = Flatbuf.MetadataVersion.V5;
 
@@ -742,19 +781,17 @@ private long WriteBufferData(IReadOnlyList<ArrowRecordBatchFlatBufferBuilder.Buf
 
             for (int i = 0; i < buffers.Count; i++)
             {
-                ReadOnlyMemory<byte> buffer = buffers[i].DataBuffer;
-                if (buffer.IsEmpty)
+                using var buffer = buffers[i]; // Dispose the buffer after writing
+                ReadOnlyMemory<byte> dataBuffer = buffer.DataBuffer;
+                if (dataBuffer.IsEmpty)
                     continue;
-
-                BaseStream.Write(buffer);
-
-                int paddedLength = checked((int)BitUtility.RoundUpToMultipleOf8(buffer.Length));
-                int padding = paddedLength - buffer.Length;
+                BaseStream.Write(dataBuffer);
+                int paddedLength = checked((int)BitUtility.RoundUpToMultipleOf8(dataBuffer.Length));
+                int padding = paddedLength - dataBuffer.Length;
                 if (padding > 0)
                 {
                     WritePadding(padding);
                 }
-
                 bodyLength += paddedLength;
             }
 
@@ -773,14 +810,15 @@ private async ValueTask<long> WriteBufferDataAsync(IReadOnlyList<ArrowRecordBatc
 
             for (int i = 0; i < buffers.Count; i++)
             {
-                ReadOnlyMemory<byte> buffer = buffers[i].DataBuffer;
-                if (buffer.IsEmpty)
+                using var buffer = buffers[i]; // Dispose the buffer after writing
+                ReadOnlyMemory<byte> dataBuffer = buffer.DataBuffer;
+                if (dataBuffer.IsEmpty)
                     continue;
 
-                await BaseStream.WriteAsync(buffer, cancellationToken).ConfigureAwait(false);
+                await BaseStream.WriteAsync(dataBuffer, cancellationToken).ConfigureAwait(false);
 
-                int paddedLength = checked((int)BitUtility.RoundUpToMultipleOf8(buffer.Length));
-                int padding = paddedLength - buffer.Length;
+                int paddedLength = checked((int)BitUtility.RoundUpToMultipleOf8(dataBuffer.Length));
+                int padding = paddedLength - dataBuffer.Length;
                 if (padding > 0)
                 {
                     await WritePaddingAsync(padding).ConfigureAwait(false);
@@ -813,7 +851,7 @@ private Tuple<ArrowRecordBatchFlatBufferBuilder, VectorOffset, VectorOffset> Pre
                 ? _options.CompressionCodecFactory.CreateCodec(_options.CompressionCodec.Value, _options.CompressionLevel)
                 : null;
 
-            var recordBatchBuilder = new ArrowRecordBatchFlatBufferBuilder(compressionCodec, _allocator, _compressionStream);
+            var recordBatchBuilder = new ArrowRecordBatchFlatBufferBuilder(compressionCodec, _allocator, _fallbackCompressionStream);
 
             // Visit all arrays recursively
             for (int i = 0; i < fields.Count; i++)
@@ -1313,7 +1351,7 @@ public virtual void Dispose()
             {
                 BaseStream.Dispose();
             }
-            _compressionStream.Dispose();
+            _fallbackCompressionStream.Dispose();
         }
     }
 
diff --git a/csharp/src/Apache.Arrow/Ipc/ITryCompressionCodec.cs b/csharp/src/Apache.Arrow/Ipc/ITryCompressionCodec.cs
new file mode 100644
index 000000000000..123d9eb83721
--- /dev/null
+++ b/csharp/src/Apache.Arrow/Ipc/ITryCompressionCodec.cs
@@ -0,0 +1,32 @@
+// 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.Ipc
+{
+    public interface ITryCompressionCodec : ICompressionCodec
+    {
+        /// <summary>
+        /// Try to write compressed data to a fixed length memory span
+        /// </summary>
+        /// <param name="source">The data to compress</param>
+        /// <param name="destination">Memory to write compressed data to</param>
+        /// <param name="bytesWritten">The number of bytes written to the destination</param>
+        /// <returns>true if compression was successful, false if the destination buffer is too small</returns>
+        bool TryCompress(ReadOnlyMemory<byte> source, Memory<byte> destination, out int bytesWritten);
+
+    }
+}
diff --git a/csharp/test/Apache.Arrow.Compression.Tests/ArrowStreamReaderTests.cs b/csharp/test/Apache.Arrow.Compression.Tests/ArrowStreamReaderTests.cs
index b9a886c88c5a..6f6a26e3f893 100644
--- a/csharp/test/Apache.Arrow.Compression.Tests/ArrowStreamReaderTests.cs
+++ b/csharp/test/Apache.Arrow.Compression.Tests/ArrowStreamReaderTests.cs
@@ -16,6 +16,7 @@
 using Apache.Arrow.Ipc;
 using System;
 using System.Reflection;
+using Apache.Arrow.Tests;
 using Xunit;
 
 namespace Apache.Arrow.Compression.Tests
@@ -65,6 +66,26 @@ public void ErrorReadingCompressedStreamWithoutCodecFactory()
             Assert.Contains("no ICompressionCodecFactory has been configured", exception.Message);
         }
 
+        [Theory]
+        [InlineData("ipc_lz4_compression.arrow_stream")]
+        [InlineData("ipc_zstd_compression.arrow_stream")]
+        public void MemoryPoolDisposedOnReadCompressedIpcStream(string fileName)
+        {
+            var assembly = Assembly.GetExecutingAssembly();
+            using var stream = assembly.GetManifestResourceStream($"Apache.Arrow.Compression.Tests.Resources.{fileName}");
+            Assert.NotNull(stream);
+            var codecFactory = new CompressionCodecFactory();
+            var allocator = new TestMemoryAllocator();
+            using var reader = new ArrowStreamReader(stream, allocator, codecFactory, false);
+            using (var recordBatch = reader.ReadNextRecordBatch())
+            {
+                VerifyCompressedIpcFileBatch(recordBatch);
+            }
+            Assert.True(allocator.Statistics.Allocations > 0);
+            Assert.Equal(0, allocator.Rented);
+
+        }
+
         private static void VerifyCompressedIpcFileBatch(RecordBatch batch)
         {
             var intArray = (Int32Array) batch.Column("integers");
diff --git a/csharp/test/Apache.Arrow.Compression.Tests/ArrowStreamWriterTests.cs b/csharp/test/Apache.Arrow.Compression.Tests/ArrowStreamWriterTests.cs
index 3b09dc26a343..1f0f61c1d004 100644
--- a/csharp/test/Apache.Arrow.Compression.Tests/ArrowStreamWriterTests.cs
+++ b/csharp/test/Apache.Arrow.Compression.Tests/ArrowStreamWriterTests.cs
@@ -18,6 +18,7 @@
 using System.IO;
 using System.Threading.Tasks;
 using Apache.Arrow.Ipc;
+using Apache.Arrow.Memory;
 using Apache.Arrow.Tests;
 using K4os.Compression.LZ4;
 using Xunit;
@@ -111,6 +112,20 @@ public void ThrowsForInvalidCompressionLevel(CompressionCodecType codec)
                 writer.WriteEnd();
             });
         }
+        [Theory]
+        [InlineData(CompressionCodecType.Zstd)]
+        [InlineData(CompressionCodecType.Lz4Frame)]
+        public async Task MemoryOwnerDisposal(CompressionCodecType codec)
+        {
+            var allocator = new TestMemoryAllocator();
+            var originalBatch = TestData.CreateSampleRecordBatch(length: 100);
+            var options = new IpcOptions() { CompressionCodecFactory = new CompressionCodecFactory(), CompressionCodec = codec };
+            await TestRoundTripRecordBatchesAsync(new List<RecordBatch> () {originalBatch}, options, options.CompressionCodecFactory,
+                allocator);
+            Assert.True(allocator.Statistics.Allocations > 0);
+            // make sure all memory allocated by the writer was disposed
+            Assert.Equal(0,allocator.Rented);
+        }
 
         private static void TestRoundTripRecordBatches(
             IReadOnlyList<RecordBatch> originalBatches, IpcOptions options, ICompressionCodecFactory codecFactory)
@@ -147,11 +162,11 @@ private static void TestRoundTripRecordBatches(
         }
 
         private static async Task TestRoundTripRecordBatchesAsync(
-            IReadOnlyList<RecordBatch> originalBatches, IpcOptions options, ICompressionCodecFactory codecFactory)
+            IReadOnlyList<RecordBatch> originalBatches, IpcOptions options, ICompressionCodecFactory codecFactory, MemoryAllocator writerAllocator = null)
         {
             using var stream = new MemoryStream();
 
-            using (var writer = new ArrowStreamWriter(stream, originalBatches[0].Schema, leaveOpen: true, options))
+            using (var writer = new ArrowStreamWriter(stream, originalBatches[0].Schema, leaveOpen: true, options, writerAllocator))
             {
                 foreach (var originalBatch in originalBatches)
                 {
diff --git a/csharp/test/Apache.Arrow.Tests/ArrowStreamWriterTests.cs b/csharp/test/Apache.Arrow.Tests/ArrowStreamWriterTests.cs
index d4e06a91faef..570fb03ed2f7 100644
--- a/csharp/test/Apache.Arrow.Tests/ArrowStreamWriterTests.cs
+++ b/csharp/test/Apache.Arrow.Tests/ArrowStreamWriterTests.cs
@@ -22,6 +22,7 @@
 using System.Net.Sockets;
 using System.Threading.Tasks;
 using Apache.Arrow.Ipc;
+using Apache.Arrow.Memory;
 using Apache.Arrow.Types;
 using Xunit;
 
@@ -239,7 +240,7 @@ private static void TestRoundTripRecordBatches(List<RecordBatch> originalBatches
         {
             using (MemoryStream stream = new MemoryStream())
             {
-                using (var writer = new ArrowStreamWriter(stream, originalBatches[0].Schema, leaveOpen: true, options))
+                using (var writer = new ArrowStreamWriter(stream, originalBatches[0].Schema, leaveOpen: true, options, new TestMemoryAllocator()))
                 {
                     foreach (RecordBatch originalBatch in originalBatches)
                     {
@@ -261,11 +262,11 @@ private static void TestRoundTripRecordBatches(List<RecordBatch> originalBatches
             }
         }
 
-        private static async Task TestRoundTripRecordBatchesAsync(List<RecordBatch> originalBatches, IpcOptions options = null, bool strictCompare = true)
+        private static async Task TestRoundTripRecordBatchesAsync(List<RecordBatch> originalBatches, IpcOptions options = null, bool strictCompare = true, MemoryAllocator memoryAllocator = null)
         {
             using (MemoryStream stream = new MemoryStream())
             {
-                using (var writer = new ArrowStreamWriter(stream, originalBatches[0].Schema, leaveOpen: true, options))
+                using (var writer = new ArrowStreamWriter(stream, originalBatches[0].Schema, leaveOpen: true, options, memoryAllocator ??  new TestMemoryAllocator()))
                 {
                     foreach (RecordBatch originalBatch in originalBatches)
                     {
@@ -715,5 +716,24 @@ public async Task WritesEmptyFileAsync()
             Assert.Null(readBatch);
             SchemaComparer.Compare(originalBatch.Schema, reader.Schema);
         }
+
+
+        [Theory]
+        [InlineData(0, 45)]
+        [InlineData(3, 45)]
+        [InlineData(16, 45)]
+        public async Task MemoryOwnerDisposalSlicedArray(int sliceOffset, int sliceLength)
+        {
+            var originalBatch = TestData.CreateSampleRecordBatch(length: 100);
+            var slicedArrays = originalBatch.Arrays
+                .Select(array => ArrowArrayFactory.Slice(array, sliceOffset, sliceLength))
+                .ToList();
+            var slicedBatch = new RecordBatch(originalBatch.Schema, slicedArrays, sliceLength);
+            var allocator = new TestMemoryAllocator();
+            await TestRoundTripRecordBatchesAsync(new List<RecordBatch> () {slicedBatch}, null, false, allocator);
+            if(sliceOffset % 8 != 0)
+                Assert.True(allocator.Statistics.Allocations > 0);
+            Assert.Equal(0,allocator.Rented);
+        }
     }
 }
diff --git a/csharp/test/Apache.Arrow.Tests/TestMemoryAllocator.cs b/csharp/test/Apache.Arrow.Tests/TestMemoryAllocator.cs
index e0e36af17296..0243c3952911 100644
--- a/csharp/test/Apache.Arrow.Tests/TestMemoryAllocator.cs
+++ b/csharp/test/Apache.Arrow.Tests/TestMemoryAllocator.cs
@@ -13,17 +13,56 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+using System;
 using Apache.Arrow.Memory;
 using System.Buffers;
+using System.Threading;
 
 namespace Apache.Arrow.Tests
 {
     public class TestMemoryAllocator : MemoryAllocator
     {
+        private int _rented = 0;
+        public int Rented => _rented;
+
         protected override IMemoryOwner<byte> AllocateInternal(int length, out int bytesAllocated)
         {
-            bytesAllocated = length;
-            return MemoryPool<byte>.Shared.Rent(length);
+            var mem = MemoryPool<byte>.Shared.Rent(length);
+            bytesAllocated = mem.Memory.Length;
+            Interlocked.Increment(ref _rented);
+            return new TestMemoryOwner(mem, this);
+        }
+
+        private class TestMemoryOwner : IMemoryOwner<byte>
+        {
+            private readonly IMemoryOwner<byte> _inner;
+            private readonly TestMemoryAllocator _allocator;
+            private bool _disposed;
+
+            public TestMemoryOwner(IMemoryOwner<byte> inner, TestMemoryAllocator allocator)
+            {
+                _inner = inner;
+                _allocator = allocator;
+            }
+
+            public Memory<byte> Memory
+            {
+                get
+                {
+                    if (_disposed)
+                        throw new ObjectDisposedException(nameof(TestMemoryOwner));
+                    return _inner.Memory;
+                }
+            }
+
+            public void Dispose()
+            {
+                if (_disposed)
+                    return;
+                _disposed = true;
+                Interlocked.Decrement(ref _allocator._rented);
+                _inner?.Dispose();
+            }
         }
     }
 }


Reviews:
user: login: adamreeve
body: Hi @m-v-w, thanks for the PR. I've left some comments and questions.

It's not clear to me that the change to how compression works when writing is actually an improvement, it looks like this would be worse when LZ4 is used. Do you have benchmarks showing the reduction in memory use and allocations from disposing the local buffers, and then the effect of the compression change separately for both LZ4 and Zstd?
state: CHANGES_REQUESTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-22T02:38:43Z
user: login: CurtHagenlocher
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-22T03:55:07Z
user: login: CurtHagenlocher
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-22T17:12:01Z
user: login: m-v-w
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-22T19:44:09Z
user: login: adamreeve
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-22T22:44:39Z
user: login: adamreeve
body: This looks good to me now thanks, I've just suggested some minor changes to comments
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-23T01:39:47Z
user: login: adamreeve
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-23T01:40:32Z
user: login: adamreeve
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-05-23T01:42:21Z

Review comments:
diff_hunk: @@ -44,5 +44,21 @@ void Compress(ReadOnlyMemory<byte> source, Stream destination)
 #else
         ;
 #endif
+
+        /// <summary>
+        /// try to write compressed data to span
+        /// </summary>
+        /// <param name="source">The data to compress</param>
+        /// <param name="destination">Span to write compressed data to</param>
+        /// <param name="bytesWritten">The number of bytes written to the destination</param>
+        /// <returns>true if compressed was successful, false if the destination buffer is too small</returns>
+        bool TryCompress(ReadOnlyMemory<byte> source, Memory<byte> destination, out int bytesWritten)
+#if NET6_0_OR_GREATER
+        {
+            throw new NotImplementedException("This codec does not support compression");
path: csharp/src/Apache.Arrow/Ipc/ICompressionCodec.cs
user: login: adamreeve
body: Could you add a default implementation that uses the `Compress` method? Then this is only breaking for .NET < 6.0.
created_at: 2025-05-22T01:52:44Z
author_association: CONTRIBUTOR
diff_hunk: @@ -15,9 +15,14 @@
 
 using System;
 using System.IO;
+using System.Runtime.CompilerServices;
path: csharp/src/Apache.Arrow.Compression/Lz4CompressionCodec.cs
user: login: adamreeve
body: I don't think any of these new using statements are actually needed
created_at: 2025-05-22T01:54:26Z
author_association: CONTRIBUTOR
diff_hunk: @@ -13,17 +13,54 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+using System;
 using Apache.Arrow.Memory;
 using System.Buffers;
+using System.Diagnostics;
path: csharp/test/Apache.Arrow.Tests/TestMemoryAllocator.cs
user: login: adamreeve
body: `System.Diagnostics` and `Xunit` usings aren't needed.
created_at: 2025-05-22T01:56:52Z
author_association: CONTRIBUTOR
diff_hunk: @@ -590,10 +607,8 @@ public void Visit(IArrowArray array)
         private readonly bool _leaveOpen;
         private readonly IpcOptions _options;
         private readonly MemoryAllocator _allocator;
-        // Reuse a single memory stream for writing compressed data to, to reduce memory allocations
-        private readonly MemoryStream _compressionStream = new MemoryStream();
 
-        private protected const Flatbuf.MetadataVersion CurrentMetadataVersion = Flatbuf.MetadataVersion.V5;
+        internal const Flatbuf.MetadataVersion CurrentMetadataVersion = Flatbuf.MetadataVersion.V5;
path: csharp/src/Apache.Arrow/Ipc/ArrowStreamWriter.cs
user: login: adamreeve
body: It doesn't look like this needed to be changed to internal?
created_at: 2025-05-22T02:10:16Z
author_association: CONTRIBUTOR
diff_hunk: @@ -56,6 +61,21 @@ public void Compress(ReadOnlyMemory<byte> source, Stream destination)
             encoder.WriteManyBytes(source.Span);
         }
 
+        public bool TryCompress(ReadOnlyMemory<byte> source, Memory<byte> destination, out int bytesWritten)
+        {
+            using var memoryStream = new MemoryStream(destination.Length);
path: csharp/src/Apache.Arrow.Compression/Lz4CompressionCodec.cs
user: login: adamreeve
body: Is this not a regression in behaviour by creating a new `MemoryStream` here for each `TryCompress` call rather than reusing the stream from `ArrowRecordBatchFlatBufferBuilder._compressionStream`?
created_at: 2025-05-22T02:25:18Z
author_association: CONTRIBUTOR
diff_hunk: @@ -13,17 +13,54 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+using System;
 using Apache.Arrow.Memory;
 using System.Buffers;
+using System.Diagnostics;
+using System.Threading;
+using Xunit;
 
 namespace Apache.Arrow.Tests
 {
     public class TestMemoryAllocator : MemoryAllocator
     {
+        private int _rented = 0;
path: csharp/test/Apache.Arrow.Tests/TestMemoryAllocator.cs
user: login: adamreeve
body: Could the formatting in this class please be tidied up by adding blank lines around methods.
created_at: 2025-05-22T02:28:09Z
author_association: CONTRIBUTOR
diff_hunk: @@ -44,5 +44,21 @@ void Compress(ReadOnlyMemory<byte> source, Stream destination)
 #else
         ;
 #endif
+
+        /// <summary>
+        /// try to write compressed data to span
+        /// </summary>
+        /// <param name="source">The data to compress</param>
+        /// <param name="destination">Span to write compressed data to</param>
+        /// <param name="bytesWritten">The number of bytes written to the destination</param>
+        /// <returns>true if compressed was successful, false if the destination buffer is too small</returns>
+        bool TryCompress(ReadOnlyMemory<byte> source, Memory<byte> destination, out int bytesWritten)
+#if NET6_0_OR_GREATER
+        {
+            throw new NotImplementedException("This codec does not support compression");
path: csharp/src/Apache.Arrow/Ipc/ICompressionCodec.cs
user: login: CurtHagenlocher
body: The problem is that default interface implementations aren't supported on downlevel .NET. I think the only way to make this not be a breaking change is to have a new interface and check dynamically for the capability at runtime (at least for Desktop framework).
created_at: 2025-05-22T03:55:07Z
author_association: CONTRIBUTOR
diff_hunk: @@ -44,5 +44,21 @@ void Compress(ReadOnlyMemory<byte> source, Stream destination)
 #else
         ;
 #endif
+
+        /// <summary>
+        /// try to write compressed data to span
+        /// </summary>
+        /// <param name="source">The data to compress</param>
+        /// <param name="destination">Span to write compressed data to</param>
+        /// <param name="bytesWritten">The number of bytes written to the destination</param>
+        /// <returns>true if compressed was successful, false if the destination buffer is too small</returns>
+        bool TryCompress(ReadOnlyMemory<byte> source, Memory<byte> destination, out int bytesWritten)
+#if NET6_0_OR_GREATER
+        {
+            throw new NotImplementedException("This codec does not support compression");
path: csharp/src/Apache.Arrow/Ipc/ICompressionCodec.cs
user: login: CurtHagenlocher
body: I just realized that I misunderstood what @adamreeve wrote, as you acknowledge that this is breaking for downlevel .NET. I suspect that custom implementations of `ICompressionCodec` are pretty rare, and am fighting to square this with my dislike for breaking changes of any kind.

One approach to avoiding the break is to define a new interface `ICompressionCodec2` which derives from `ICompressionCodec` and then adds the member. The consuming code can then wrap any `ICompressionCodec` it gets with a helper that implements `ICompressionCodec2`.

And yes, I know the name `ICompressionCodec2` is awful.
created_at: 2025-05-22T17:12:01Z
author_association: CONTRIBUTOR
diff_hunk: @@ -44,5 +44,21 @@ void Compress(ReadOnlyMemory<byte> source, Stream destination)
 #else
         ;
 #endif
+
+        /// <summary>
+        /// try to write compressed data to span
+        /// </summary>
+        /// <param name="source">The data to compress</param>
+        /// <param name="destination">Span to write compressed data to</param>
+        /// <param name="bytesWritten">The number of bytes written to the destination</param>
+        /// <returns>true if compressed was successful, false if the destination buffer is too small</returns>
+        bool TryCompress(ReadOnlyMemory<byte> source, Memory<byte> destination, out int bytesWritten)
+#if NET6_0_OR_GREATER
+        {
+            throw new NotImplementedException("This codec does not support compression");
path: csharp/src/Apache.Arrow/Ipc/ICompressionCodec.cs
user: login: m-v-w
body: Totally agree, using a new interface would have been better. Here some name suggestions: 
`ITryCompressionCodec`, `ICompressionCodecExt`
created_at: 2025-05-22T19:44:09Z
author_association: CONTRIBUTOR
diff_hunk: @@ -44,5 +44,21 @@ void Compress(ReadOnlyMemory<byte> source, Stream destination)
 #else
         ;
 #endif
+
+        /// <summary>
+        /// try to write compressed data to span
+        /// </summary>
+        /// <param name="source">The data to compress</param>
+        /// <param name="destination">Span to write compressed data to</param>
+        /// <param name="bytesWritten">The number of bytes written to the destination</param>
+        /// <returns>true if compressed was successful, false if the destination buffer is too small</returns>
+        bool TryCompress(ReadOnlyMemory<byte> source, Memory<byte> destination, out int bytesWritten)
+#if NET6_0_OR_GREATER
+        {
+            throw new NotImplementedException("This codec does not support compression");
path: csharp/src/Apache.Arrow/Ipc/ICompressionCodec.cs
user: login: adamreeve
body: :+1: thanks, the new approach looks much better to me, and also addresses my concerns about the LZ4 implementation.
created_at: 2025-05-22T22:44:39Z
author_association: CONTRIBUTOR
diff_hunk: @@ -487,82 +493,114 @@ private Buffer CreateBitmapBuffer(ArrowBuffer buffer, int offset, int length)
                         BitUtility.SetBit(outputSpan, i, BitUtility.GetBit(inputSpan, offset + i));
                     }
 
-                    return CreateBuffer(memoryOwner.Memory);
+                    return CreateBuffer(memoryOwner);
                 }
             }
 
-            private Buffer CreateSlicedBuffer<T>(ArrowBuffer buffer, int offset, int length)
+            private Buffer CreateSlicedBuffer<T>(ArrowBuffer buffer, int offset, int length, bool owned)
                 where T : struct
             {
-                return CreateSlicedBuffer(buffer, Unsafe.SizeOf<T>(), offset, length);
+                return CreateSlicedBuffer(buffer, Unsafe.SizeOf<T>(), offset, length, owned);
             }
 
-            private Buffer CreateSlicedBuffer(ArrowBuffer buffer, int itemSize, int offset, int length)
+            private Buffer CreateSlicedBuffer(ArrowBuffer buffer, int itemSize, int offset, int length, bool owned)
             {
                 var byteLength = length * itemSize;
                 var paddedLength = CalculatePaddedBufferLength(byteLength);
                 if (offset != 0 || paddedLength < buffer.Length)
                 {
                     var byteOffset = offset * itemSize;
                     var sliceLength = Math.Min(paddedLength, buffer.Length - byteOffset);
-                    return CreateBuffer(buffer.Memory.Slice(byteOffset, sliceLength));
+                    return CreateBuffer(buffer.Memory.Slice(byteOffset, sliceLength),
+                        owned ? buffer : null);
                 }
 
-                return CreateBuffer(buffer.Memory);
+                return CreateBuffer(buffer, owned);
             }
 
-            private Buffer CreateBuffer(ArrowBuffer buffer)
+            private Buffer CreateBuffer(ArrowBuffer buffer, bool locallyOwned)
             {
-                return CreateBuffer(buffer.Memory);
+                if (locallyOwned)
+                {
+                    // this buffer is locally owned, we can dispose it
+                    return CreateBuffer(buffer.Memory, buffer);
+                }
+                return CreateBuffer(buffer.Memory, null);
             }
 
-            private Buffer CreateBuffer(ReadOnlyMemory<byte> buffer)
+            private Buffer CreateBuffer(IMemoryOwner<byte>  bufferOwner)
+            {
+                return CreateBuffer(bufferOwner.Memory, bufferOwner);
+            }
+
+            private Buffer CreateBuffer(ReadOnlyMemory<byte> buffer, IDisposable localBufferOwner)
             {
                 int offset = TotalLength;
                 const int UncompressedLengthSize = 8;
 
                 ReadOnlyMemory<byte> bufferToWrite;
+                IMemoryOwner<byte> bufferOwner=null;
                 if (_compressionCodec == null)
                 {
                     bufferToWrite = buffer;
                 }
                 else if (buffer.Length == 0)
                 {
                     // Write zero length and skip compression
-                    var uncompressedLengthBytes = _allocator.Allocate(UncompressedLengthSize);
-                    BinaryPrimitives.WriteInt64LittleEndian(uncompressedLengthBytes.Memory.Span, 0);
-                    bufferToWrite = uncompressedLengthBytes.Memory;
+                    bufferOwner = _allocator.Allocate(UncompressedLengthSize);
+                    BinaryPrimitives.WriteInt64LittleEndian(bufferOwner.Memory.Span, 0);
+                    bufferToWrite = bufferOwner.Memory.Slice(0, UncompressedLengthSize);
+                    // the local source buffer owner can be disposed, it's memory is no longer needed
path: csharp/src/Apache.Arrow/Ipc/ArrowStreamWriter.cs
user: login: adamreeve
body: ```suggestion
                    // the local source buffer owner can be disposed, its memory is no longer needed
```
created_at: 2025-05-23T01:32:50Z
author_association: CONTRIBUTOR
diff_hunk: @@ -487,82 +493,114 @@ private Buffer CreateBitmapBuffer(ArrowBuffer buffer, int offset, int length)
                         BitUtility.SetBit(outputSpan, i, BitUtility.GetBit(inputSpan, offset + i));
                     }
 
-                    return CreateBuffer(memoryOwner.Memory);
+                    return CreateBuffer(memoryOwner);
                 }
             }
 
-            private Buffer CreateSlicedBuffer<T>(ArrowBuffer buffer, int offset, int length)
+            private Buffer CreateSlicedBuffer<T>(ArrowBuffer buffer, int offset, int length, bool owned)
                 where T : struct
             {
-                return CreateSlicedBuffer(buffer, Unsafe.SizeOf<T>(), offset, length);
+                return CreateSlicedBuffer(buffer, Unsafe.SizeOf<T>(), offset, length, owned);
             }
 
-            private Buffer CreateSlicedBuffer(ArrowBuffer buffer, int itemSize, int offset, int length)
+            private Buffer CreateSlicedBuffer(ArrowBuffer buffer, int itemSize, int offset, int length, bool owned)
             {
                 var byteLength = length * itemSize;
                 var paddedLength = CalculatePaddedBufferLength(byteLength);
                 if (offset != 0 || paddedLength < buffer.Length)
                 {
                     var byteOffset = offset * itemSize;
                     var sliceLength = Math.Min(paddedLength, buffer.Length - byteOffset);
-                    return CreateBuffer(buffer.Memory.Slice(byteOffset, sliceLength));
+                    return CreateBuffer(buffer.Memory.Slice(byteOffset, sliceLength),
+                        owned ? buffer : null);
                 }
 
-                return CreateBuffer(buffer.Memory);
+                return CreateBuffer(buffer, owned);
             }
 
-            private Buffer CreateBuffer(ArrowBuffer buffer)
+            private Buffer CreateBuffer(ArrowBuffer buffer, bool locallyOwned)
             {
-                return CreateBuffer(buffer.Memory);
+                if (locallyOwned)
+                {
+                    // this buffer is locally owned, we can dispose it
+                    return CreateBuffer(buffer.Memory, buffer);
+                }
+                return CreateBuffer(buffer.Memory, null);
             }
 
-            private Buffer CreateBuffer(ReadOnlyMemory<byte> buffer)
+            private Buffer CreateBuffer(IMemoryOwner<byte>  bufferOwner)
+            {
+                return CreateBuffer(bufferOwner.Memory, bufferOwner);
+            }
+
+            private Buffer CreateBuffer(ReadOnlyMemory<byte> buffer, IDisposable localBufferOwner)
             {
                 int offset = TotalLength;
                 const int UncompressedLengthSize = 8;
 
                 ReadOnlyMemory<byte> bufferToWrite;
+                IMemoryOwner<byte> bufferOwner=null;
                 if (_compressionCodec == null)
                 {
                     bufferToWrite = buffer;
                 }
                 else if (buffer.Length == 0)
                 {
                     // Write zero length and skip compression
-                    var uncompressedLengthBytes = _allocator.Allocate(UncompressedLengthSize);
-                    BinaryPrimitives.WriteInt64LittleEndian(uncompressedLengthBytes.Memory.Span, 0);
-                    bufferToWrite = uncompressedLengthBytes.Memory;
+                    bufferOwner = _allocator.Allocate(UncompressedLengthSize);
+                    BinaryPrimitives.WriteInt64LittleEndian(bufferOwner.Memory.Span, 0);
+                    bufferToWrite = bufferOwner.Memory.Slice(0, UncompressedLengthSize);
+                    // the local source buffer owner can be disposed, it's memory is no longer needed
+                    localBufferOwner?.Dispose();
                 }
                 else
                 {
                     // See format/Message.fbs, and the BUFFER BodyCompressionMethod for documentation on how
                     // compressed buffers are stored.
-                    _compressionStream.Seek(0, SeekOrigin.Begin);
-                    _compressionStream.SetLength(0);
-                    _compressionCodec.Compress(buffer, _compressionStream);
-                    if (_compressionStream.Length < buffer.Length)
+                    int newBufferLength = UncompressedLengthSize + buffer.Length;
+                    bufferOwner = _allocator.Allocate(newBufferLength);
+
+                    if(TryCompress(buffer, bufferOwner.Memory.Slice(UncompressedLengthSize, buffer.Length), out int bytesWritten))
                     {
-                        var newBuffer = _allocator.Allocate((int) _compressionStream.Length + UncompressedLengthSize);
-                        BinaryPrimitives.WriteInt64LittleEndian(newBuffer.Memory.Span, buffer.Length);
-                        _compressionStream.Seek(0, SeekOrigin.Begin);
-                        _compressionStream.ReadFullBuffer(newBuffer.Memory.Slice(UncompressedLengthSize));
-                        bufferToWrite = newBuffer.Memory;
+                        // Write the uncompressed length to the start of the buffer
+                        BinaryPrimitives.WriteInt64LittleEndian(bufferOwner.Memory.Span, buffer.Length);
+                        bufferToWrite = bufferOwner.Memory.Slice(0, bytesWritten+UncompressedLengthSize);
                     }
                     else
                     {
+                        // TryCompress failed, because the buffer is too small.
                         // If the compressed buffer is larger than the uncompressed buffer, use the uncompressed
                         // buffer instead, and indicate this by setting the uncompressed length to -1
-                        var newBuffer = _allocator.Allocate(buffer.Length + UncompressedLengthSize);
-                        BinaryPrimitives.WriteInt64LittleEndian(newBuffer.Memory.Span, -1);
-                        buffer.CopyTo(newBuffer.Memory.Slice(UncompressedLengthSize));
-                        bufferToWrite = newBuffer.Memory;
+                        BinaryPrimitives.WriteInt64LittleEndian(bufferOwner.Memory.Span, -1);
+                        buffer.CopyTo(bufferOwner.Memory.Slice(UncompressedLengthSize));
+                        bufferToWrite = bufferOwner.Memory.Slice(0, newBufferLength);
                     }
+                    // the local source buffer owner can be disposed, it's memory was copied
path: csharp/src/Apache.Arrow/Ipc/ArrowStreamWriter.cs
user: login: adamreeve
body: ```suggestion
                    // the local source buffer owner can be disposed, its memory was copied
```
created_at: 2025-05-23T01:33:06Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,32 @@
+// 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.Ipc
+{
+    public interface ITryCompressionCodec : ICompressionCodec
+    {
+        /// <summary>
+        /// try to write compressed data to span
path: csharp/src/Apache.Arrow/Ipc/ITryCompressionCodec.cs
user: login: adamreeve
body: ```suggestion
        /// Try to write compressed data to a fixed length memory span
```
created_at: 2025-05-23T01:36:38Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,32 @@
+// 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.Ipc
+{
+    public interface ITryCompressionCodec : ICompressionCodec
+    {
+        /// <summary>
+        /// try to write compressed data to span
+        /// </summary>
+        /// <param name="source">The data to compress</param>
+        /// <param name="destination">Span to write compressed data to</param>
path: csharp/src/Apache.Arrow/Ipc/ITryCompressionCodec.cs
user: login: adamreeve
body: ```suggestion
        /// <param name="destination">Memory to write compressed data to</param>
```
created_at: 2025-05-23T01:38:35Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,32 @@
+// 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.Ipc
+{
+    public interface ITryCompressionCodec : ICompressionCodec
+    {
+        /// <summary>
+        /// try to write compressed data to span
+        /// </summary>
+        /// <param name="source">The data to compress</param>
+        /// <param name="destination">Span to write compressed data to</param>
+        /// <param name="bytesWritten">The number of bytes written to the destination</param>
+        /// <returns>true if compressed was successful, false if the destination buffer is too small</returns>
path: csharp/src/Apache.Arrow/Ipc/ITryCompressionCodec.cs
user: login: adamreeve
body: ```suggestion
        /// <returns>true if compression was successful, false if the destination buffer is too small</returns>
```
created_at: 2025-05-23T01:40:32Z
author_association: CONTRIBUTOR

PR comments:
user: login: github-actions[bot]
created_at: 2025-04-21T21:07:30Z
body: :warning: GitHub issue #46189 **has been automatically assigned in GitHub** to PR creator.
author_association: NONE
user: login: m-v-w
created_at: 2025-05-22T19:40:30Z
body: Hi @adamreeve , @CurtHagenlocher ,

thanks for considering my PR. 

The issue I encountered with `ArrowStreamWriter` is it allocates large buffers on compression and array slicing which will only be released on finalization by the GC. This caused GC pressure in my app.
Not sure what a good benchmark would look like. It depends on the `MemoryAllocator` implementation. Making sure all allocations are disposed as soon as possible would allow the user to use a pooled `MemoryAllocator` implementation for instance.
My implementation of the `TryCompress` method for LZ4 is bad, it wasn't a priority since I don't use it. But I'll come up with a better solution. 
I think having the possibility to implement a compression codec that's allocation free is useful. 

I'll address the raised issues and make the required changes.



author_association: CONTRIBUTOR
user: login: conbench-apache-arrow[bot]
created_at: 2025-05-27T21:22:42Z
body: After merging your PR, Conbench analyzed the 4 benchmarking runs that have been run so far on merge-commit 1cabcb32326ccb5e227e8a0574d6a81bcbde6f12.

There were no benchmark performance regressions. 🎉

The [full Conbench report](https://github.com/apache/arrow/runs/42996218596) has more details. It also includes information about 3 possible false positives for unstable benchmarks that are known to sometimes produce them.
author_association: NONE