Use ArrayPool for Reader and Writer Buffers

csharp/LogicalBatchReader/NestedReader.cs

@@ -1,4 +1,5 @@
 using System;
+using System.Buffers;
 
 namespace ParquetSharp.LogicalBatchReader
 {
@@ -11,28 +12,36 @@ internal sealed class NestedReader<TItem> : ILogicalBatchReader<Nested<TItem>>
         public NestedReader(ILogicalBatchReader<TItem> innerReader, int bufferLength)
         {
             _innerReader = innerReader;
-            _buffer = new TItem[bufferLength];
+            _bufferLength = bufferLength;
         }
 
         public int ReadBatch(Span<Nested<TItem>> destination)
         {
             // Read batches of values from the underlying reader and convert them to nested values
             var totalRead = 0;
-            while (totalRead < destination.Length)
+            var buffer = ArrayPool<TItem>.Shared.Rent(_bufferLength);
+            try
             {
-                var readSize = Math.Min(destination.Length - totalRead, _buffer.Length);
-                var valuesRead = _innerReader.ReadBatch(_buffer.AsSpan(0, readSize));
-                for (var i = 0; i < valuesRead; ++i)
+                while (totalRead < destination.Length)
                 {
-                    destination[totalRead + i] = new Nested<TItem>(_buffer[i]);
-                }
+                    var readSize = Math.Min(destination.Length - totalRead, buffer.Length);
+                    var valuesRead = _innerReader.ReadBatch(buffer.AsSpan(0, readSize));
+                    for (var i = 0; i < valuesRead; ++i)
+                    {
+                        destination[totalRead + i] = new Nested<TItem>(buffer[i]);
+                    }
 
-                totalRead += valuesRead;
-                if (valuesRead < readSize)
-                {
-                    break;
+                    totalRead += valuesRead;
+                    if (valuesRead < readSize)
+                    {
+                        break;
+                    }
                 }
             }
+            finally
+            {
+                ArrayPool<TItem>.Shared.Return(buffer);
+            }
 
             return totalRead;
         }
@@ -48,6 +57,6 @@ public long Skip(long numRowsToSkip)
         }
 
         private readonly ILogicalBatchReader<TItem> _innerReader;
-        private readonly TItem[] _buffer;
+        private readonly int _bufferLength;
     }
 }

csharp/LogicalBatchWriter/NestedWriter.cs

@@ -1,4 +1,5 @@
 using System;
+using System.Buffers;
 
 namespace ParquetSharp.LogicalBatchWriter
 {
@@ -14,28 +15,36 @@ public NestedWriter(
         {
             _firstInnerWriter = firstInnerWriter;
             _innerWriter = innerWriter;
-            _buffer = new TItem[bufferLength];
+            _bufferLength = bufferLength;
         }
 
         public void WriteBatch(ReadOnlySpan<Nested<TItem>> values)
         {
             var offset = 0;
             var writer = _firstInnerWriter;
-            while (offset < values.Length)
+            var buffer = ArrayPool<TItem>.Shared.Rent(_bufferLength);
+            try
             {
-                var batchSize = Math.Min(values.Length - offset, _buffer.Length);
-                for (var i = 0; i < batchSize; ++i)
+                while (offset < values.Length)
                 {
-                    _buffer[i] = values[offset + i].Value;
+                    var batchSize = Math.Min(values.Length - offset, buffer.Length);
+                    for (var i = 0; i < batchSize; ++i)
+                    {
+                        buffer[i] = values[offset + i].Value;
+                    }
+                    writer.WriteBatch(buffer.AsSpan(0, batchSize));
+                    offset += batchSize;
+                    writer = _innerWriter;
                 }
-                writer.WriteBatch(_buffer.AsSpan(0, batchSize));
-                offset += batchSize;
-                writer = _innerWriter;
+            }
+            finally
+            {
+                ArrayPool<TItem>.Shared.Return(buffer);
             }
         }
 
         private readonly ILogicalBatchWriter<TItem> _firstInnerWriter;
         private readonly ILogicalBatchWriter<TItem> _innerWriter;
-        private readonly TItem[] _buffer;
+        private readonly int _bufferLength;
     }
 }

csharp/LogicalColumnReader.cs

@@ -1,8 +1,7 @@
 using System;
+using System.Buffers;
 using System.Collections;
 using System.Collections.Generic;
-using System.Linq;
-using ParquetSharp.Schema;
 using ParquetSharp.LogicalBatchReader;
 
 namespace ParquetSharp
@@ -159,17 +158,23 @@ public override TReturn Apply<TReturn>(ILogicalColumnReaderVisitor<TReturn> visi
 
         public IEnumerator<TElement> GetEnumerator()
         {
-            var buffer = new TElement[BufferLength];
-
-            while (HasNext)
+            var buffer = ArrayPool<TElement>.Shared.Rent(BufferLength);
+            try
             {
-                var read = ReadBatch(buffer);
-
-                for (int i = 0; i != read; ++i)
+                while (HasNext)
                 {
-                    yield return buffer[i];
+                    var read = ReadBatch(buffer);
+
+                    for (int i = 0; i != read; ++i)
+                    {
+                        yield return buffer[i];
+                    }
                 }
             }
+            finally
+            {
+                ArrayPool<TElement>.Shared.Return(buffer);
+            }
         }
 
         IEnumerator IEnumerable.GetEnumerator()