GH-46189 [C#] Use pooled buffers in ArrowStreamWriter (#46190)

### 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

Lead-authored-by: Michael <[email protected]>
Co-authored-by: Adam Reeve <[email protected]>
Signed-off-by: Curt Hagenlocher <[email protected]>

Author: m-v-w

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

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

csharp/src/Apache.Arrow/Ipc/ArrowStreamWriter.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);
+
+    }
+}

csharp/src/Apache.Arrow/Ipc/ITryCompressionCodec.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");

csharp/test/Apache.Arrow.Compression.Tests/ArrowStreamReaderTests.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)
                 {

csharp/test/Apache.Arrow.Compression.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);
+        }
     }
 }

csharp/test/Apache.Arrow.Tests/ArrowStreamWriterTests.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();
+            }
         }
     }
 }