Don't use Linq

Avoid unnecessary allocation
Change RecordBatchManipulator to RecordBatchEnumerator

Author: takashi hashida

csharp/src/Apache.Arrow/Ipc/ArrowReaderImplementation.cs

@@ -19,7 +19,6 @@
 using System.Collections.Generic;
 using System.Diagnostics;
 using System.IO;
-using System.Linq;
 using System.Threading;
 using System.Threading.Tasks;
 using Apache.Arrow.Types;
@@ -115,38 +114,43 @@ private List<IArrowArray> BuildArrays(
             ByteBuffer messageBuffer,
             Flatbuf.RecordBatch recordBatchMessage)
         {
-            return CreateInner().ToList();
+            var arrays = new List<IArrowArray>(recordBatchMessage.NodesLength);
 
-            IEnumerable<IArrowArray> CreateInner()
+            if (recordBatchMessage.NodesLength == 0)
             {
-                var recordBatchManipulator = new RecordBatchManipulator(in recordBatchMessage);
+                return arrays;
+            }
+
+            var recordBatchEnumerator = new RecordBatchEnumerator(in recordBatchMessage);
 
-                while (!recordBatchManipulator.IsAllNodeRead)
-                {
-                    var field = schema.GetFieldByIndex(recordBatchManipulator.CurrentNodeIndex);
-                    Flatbuf.FieldNode fieldNode = recordBatchManipulator.ShiftNode();
+            do
+            {
+                var field = schema.GetFieldByIndex(recordBatchEnumerator.CurrentNodeIndex);
+                var fieldNode = recordBatchEnumerator.CurrentNode;
 
-                    var arrayData = field.DataType.IsFixedPrimitive() ?
-                        LoadPrimitiveField(recordBatchManipulator, field, in fieldNode, messageBuffer) :
-                        LoadVariableField(recordBatchManipulator, field, in fieldNode, messageBuffer);
+                var arrayData = field.DataType.IsFixedPrimitive()
+                    ? LoadPrimitiveField(recordBatchEnumerator, field, in fieldNode, messageBuffer)
+                    : LoadVariableField(recordBatchEnumerator, field, in fieldNode, messageBuffer);
 
-                    yield return ArrowArrayFactory.BuildArray(arrayData);
-                }
-            }
+                arrays.Add(ArrowArrayFactory.BuildArray(arrayData));
+            } while (recordBatchEnumerator.MoveNextNode());
+
+            return arrays;
         }
 
 
         private ArrayData LoadPrimitiveField(
-            RecordBatchManipulator recordBatchManipulator,
+            RecordBatchEnumerator recordBatchEnumerator,
             Field field,
             in Flatbuf.FieldNode fieldNode,
             ByteBuffer bodyData)
         {
-            var nullBitmapBuffer = recordBatchManipulator.ShiftBuffer();
-            var valueBuffer = recordBatchManipulator.ShiftBuffer();
 
-            ArrowBuffer nullArrowBuffer = BuildArrowBuffer(bodyData, nullBitmapBuffer);
-            ArrowBuffer valueArrowBuffer = BuildArrowBuffer(bodyData, valueBuffer);
+            ArrowBuffer nullArrowBuffer = BuildArrowBuffer(bodyData, recordBatchEnumerator.CurrentBuffer);
+            recordBatchEnumerator.MoveNextBuffer();
+            ArrowBuffer valueArrowBuffer = BuildArrowBuffer(bodyData, recordBatchEnumerator.CurrentBuffer);
+            recordBatchEnumerator.MoveNextBuffer();
+
 
             var fieldLength = (int)fieldNode.Length;
             var fieldNullCount = (int)fieldNode.NullCount;
@@ -162,25 +166,25 @@ private ArrayData LoadPrimitiveField(
             }
 
             var arrowBuff = new[] { nullArrowBuffer, valueArrowBuffer };
-            var offspring = GetOffspring(recordBatchManipulator, field, bodyData);
+            var children = GetChildren(recordBatchEnumerator, field, bodyData);
 
-            return new ArrayData(field.DataType, fieldLength, fieldNullCount, 0, arrowBuff, offspring.ToArray());
+            return new ArrayData(field.DataType, fieldLength, fieldNullCount, 0, arrowBuff, children);
         }
 
 
         private ArrayData LoadVariableField(
-            RecordBatchManipulator recordBatchManipulator,
+            RecordBatchEnumerator recordBatchEnumerator,
             Field field,
             in Flatbuf.FieldNode fieldNode,
             ByteBuffer bodyData)
         {
-            var nullBitmapBuffer = recordBatchManipulator.ShiftBuffer();
-            var offsetBuffer = recordBatchManipulator.ShiftBuffer();
-            var valueBuffer = recordBatchManipulator.ShiftBuffer();
 
-            ArrowBuffer nullArrowBuffer = BuildArrowBuffer(bodyData, nullBitmapBuffer);
-            ArrowBuffer offsetArrowBuffer = BuildArrowBuffer(bodyData, offsetBuffer);
-            ArrowBuffer valueArrowBuffer = BuildArrowBuffer(bodyData, valueBuffer);
+            ArrowBuffer nullArrowBuffer = BuildArrowBuffer(bodyData, recordBatchEnumerator.CurrentBuffer);
+            recordBatchEnumerator.MoveNextBuffer();
+            ArrowBuffer offsetArrowBuffer = BuildArrowBuffer(bodyData, recordBatchEnumerator.CurrentBuffer);
+            recordBatchEnumerator.MoveNextBuffer();
+            ArrowBuffer valueArrowBuffer = BuildArrowBuffer(bodyData, recordBatchEnumerator.CurrentBuffer);
+            recordBatchEnumerator.MoveNextBuffer();
 
             var fieldLength = (int)fieldNode.Length;
             var fieldNullCount = (int)fieldNode.NullCount;
@@ -196,24 +200,33 @@ private ArrayData LoadVariableField(
             }
 
             var arrowBuff = new[] { nullArrowBuffer, offsetArrowBuffer, valueArrowBuffer };
-            var offspring = GetOffspring(recordBatchManipulator, field, bodyData);
+            var children = GetChildren(recordBatchEnumerator, field, bodyData);
 
-            return new ArrayData(field.DataType, fieldLength, fieldNullCount, 0, arrowBuff, offspring.ToArray());
+            return new ArrayData(field.DataType, fieldLength, fieldNullCount, 0, arrowBuff, children);
         }
 
-        private IEnumerable<ArrayData> GetOffspring(
-            RecordBatchManipulator recordBatchManipulator,
+        private ArrayData[] GetChildren(
+            RecordBatchEnumerator recordBatchEnumerator,
             Field field,
             ByteBuffer bodyData)
         {
-            if (!(field.DataType is NestedType type)) yield break;
-            foreach (var childField in type.Children)
+            if (!(field.DataType is NestedType type)) return null;
+
+            var childrenCount = type.ChildrenCount;
+            var children = new ArrayData[childrenCount];
+            for (var index = 0; index < childrenCount; index++)
             {
-                Flatbuf.FieldNode childFieldNode = recordBatchManipulator.ShiftNode();
-                yield return childField.DataType.IsFixedPrimitive()
-                    ? LoadPrimitiveField(recordBatchManipulator, childField, in childFieldNode, bodyData)
-                    : LoadVariableField(recordBatchManipulator, childField, in childFieldNode, bodyData);
+                Flatbuf.FieldNode childFieldNode = recordBatchEnumerator.CurrentNode;
+                recordBatchEnumerator.MoveNextNode();
+
+                var childField = type.Children[index];
+                var child = childField.DataType.IsFixedPrimitive()
+                    ? LoadPrimitiveField(recordBatchEnumerator, childField, in childFieldNode, bodyData)
+                    : LoadVariableField(recordBatchEnumerator, childField, in childFieldNode, bodyData);
+
+                children[index] = child;
             }
+            return children;
         }
 
         private ArrowBuffer BuildArrowBuffer(ByteBuffer bodyData, Flatbuf.Buffer buffer)
@@ -231,27 +244,29 @@ private ArrowBuffer BuildArrowBuffer(ByteBuffer bodyData, Flatbuf.Buffer buffer)
         }
     }
 
-    class RecordBatchManipulator
+    internal class RecordBatchEnumerator
     {
-        private int CurrentBufferIndex { get; set; }
         private Flatbuf.RecordBatch RecordBatch { get; }
+        internal int CurrentBufferIndex { get; set; }
         internal int CurrentNodeIndex { get; set; }
-        internal bool IsAllNodeRead => CurrentNodeIndex >= RecordBatch.NodesLength;
 
-        internal RecordBatchManipulator(in Flatbuf.RecordBatch recordBatch)
+        internal Flatbuf.Buffer CurrentBuffer => RecordBatch.Buffers(CurrentBufferIndex).GetValueOrDefault();
+
+        internal Flatbuf.FieldNode CurrentNode => RecordBatch.Nodes(CurrentNodeIndex).GetValueOrDefault();
+
+        internal bool MoveNextBuffer()
         {
-            RecordBatch = recordBatch;
+            return ++CurrentBufferIndex < RecordBatch.BuffersLength;
         }
 
-        internal Flatbuf.Buffer ShiftBuffer()
+        internal bool MoveNextNode()
         {
-            return RecordBatch.Buffers(CurrentBufferIndex++).GetValueOrDefault();
+            return ++CurrentNodeIndex < RecordBatch.NodesLength;
         }
 
-        internal Flatbuf.FieldNode ShiftNode()
+        internal RecordBatchEnumerator(in Flatbuf.RecordBatch recordBatch)
         {
-            return RecordBatch.Nodes(CurrentNodeIndex++).GetValueOrDefault();
+            RecordBatch = recordBatch;
         }
-
     }
 }

csharp/src/Apache.Arrow/Ipc/ArrowStreamWriter.cs

@@ -18,7 +18,6 @@
 using System.Buffers.Binary;
 using System.Collections.Generic;
 using System.IO;
-using System.Linq;
 using System.Threading;
 using System.Threading.Tasks;
 using Apache.Arrow.Types;
@@ -177,47 +176,40 @@ public ArrowStreamWriter(Stream baseStream, Schema schema, bool leaveOpen, IpcOp
             _options = options ?? IpcOptions.Default;
         }
 
-        private VectorOffset CreateFieldVector(IEnumerable<IArrowArray> fieldArrayList)
-        {
-            var allArrowArrayList = GetAll(fieldArrayList).ToList();
-            
-            Flatbuf.RecordBatch.StartNodesVector(Builder, allArrowArrayList.Count);
 
-            foreach (var array in allArrowArrayList)
+        private void CreateSelfAndChildrenFieldNodes(ArrayData data)
+        {
+            if (data.DataType is NestedType)
             {
-                Flatbuf.FieldNode.CreateFieldNode(Builder, array.Length, array.NullCount);
+                // flatbuffer struct vectors have to be created in reverse order
+                for (var i = data.Children.Length - 1; i >= 0; i--)
+                {
+                    CreateSelfAndChildrenFieldNodes(data.Children[i]);
+                }
             }
+            Flatbuf.FieldNode.CreateFieldNode(Builder, data.Length, data.NullCount);
+        }
 
-            return Builder.EndVector();
-
-
-            //Inner methods
-
-            IEnumerable<IArrowArray> GetAll(IEnumerable<IArrowArray> targetArrayList)
+        private int CountAllNodes()
+        {
+            var count = 0;
+            foreach (var arrowArray in Schema.Fields.Values)
             {
-                foreach (var arrowArray in targetArrayList)
-                {
-                    foreach (var arr in GetSelfAndOffspring(arrowArray))
-                    {
-                        yield return arr;
-                    }
-                }
+                CountSelfAndChildrenNodes(arrowArray.DataType, ref count);
             }
+            return count;
+        }
 
-            IEnumerable<IArrowArray> GetSelfAndOffspring(IArrowArray targetArray)
+        private void CountSelfAndChildrenNodes(IArrowType type, ref int count)
+        {
+            if (type is NestedType nestedType)
             {
-                if (targetArray.Data.DataType is NestedType)
+                foreach (var childField in nestedType.Children)
                 {
-                    foreach (var child in targetArray.Data.Children)
-                    {
-                        foreach (var offspring in GetSelfAndOffspring(ArrowArrayFactory.BuildArray(child)))
-                        {
-                            yield return offspring;
-                        }
-                    }
+                    CountSelfAndChildrenNodes(childField.DataType, ref count);
                 }
-                yield return targetArray;
             }
+            count++;
         }
 
         private protected async Task WriteRecordBatchInternalAsync(RecordBatch recordBatch,
@@ -234,13 +226,22 @@ private protected async Task WriteRecordBatchInternalAsync(RecordBatch recordBat
             Builder.Clear();
 
             // Serialize field nodes
+
+            var fieldCount = Schema.Fields.Count;
+
+            Flatbuf.RecordBatch.StartNodesVector(Builder, CountAllNodes());
+
             // flatbuffer struct vectors have to be created in reverse order
-            var fieldNodesVectorOffset = CreateFieldVector(recordBatch.Arrays.Reverse());
+            for (var i = fieldCount - 1; i >= 0; i--)
+            {
+                CreateSelfAndChildrenFieldNodes(recordBatch.Column(i).Data);
+            }
+
+            var fieldNodesVectorOffset = Builder.EndVector();
 
             // Serialize buffers
-            var recordBatchBuilder = new ArrowRecordBatchFlatBufferBuilder();
 
-            var fieldCount = Schema.Fields.Count;
+            var recordBatchBuilder = new ArrowRecordBatchFlatBufferBuilder();
             for (var i = 0; i < fieldCount; i++)
             {
                 var fieldArray = recordBatch.Column(i);
@@ -349,12 +350,11 @@ private ValueTask WriteBufferAsync(ArrowBuffer arrowBuffer, CancellationToken ca
                 var fieldNameOffset = Builder.CreateString(field.Name);
                 var fieldType = _fieldTypeBuilder.BuildFieldType(field);
 
-                var fieldChildren = GetChildrenFieldOffset(field).ToArray();
-                var fieldChildrenOffsets = Builder.CreateVectorOfTables(fieldChildren.ToArray());
+                var fieldChildrenVectorOffset = Builder.CreateVectorOfTables(GetChildrenFieldOffsets(field));
 
                 fieldOffsets[i] = Flatbuf.Field.CreateField(Builder,
                     fieldNameOffset, field.IsNullable, fieldType.Type, fieldType.Offset,
-                    default, fieldChildrenOffsets, default);
+                    default, fieldChildrenVectorOffset, default);
             }
 
             var fieldsVectorOffset = Flatbuf.Schema.CreateFieldsVector(Builder, fieldOffsets);
@@ -367,19 +367,28 @@ private ValueTask WriteBufferAsync(ArrowBuffer arrowBuffer, CancellationToken ca
                 Builder, endianness, fieldsVectorOffset);
         }
 
-        private protected IEnumerable<Offset<Flatbuf.Field>> GetChildrenFieldOffset(Field field)
+        private protected Offset<Flatbuf.Field>[] GetChildrenFieldOffsets(Field field)
         {
-            if (!(field.DataType is NestedType type)) yield break;
-            foreach (var child in type.Children)
+            if (!(field.DataType is NestedType type))
             {
-                var fieldNameOffset = Builder.CreateString(child.Name);
-                var fieldType = _fieldTypeBuilder.BuildFieldType(child);
-                var fieldChildrenOffsets = Builder.CreateVectorOfTables(GetChildrenFieldOffset(child).ToArray());
+                return new Offset<Flatbuf.Field>[0];
+            }
+
+            var childrenCount = type.ChildrenCount;
+            var children = new Offset<Flatbuf.Field>[childrenCount];
 
-                yield return Flatbuf.Field.CreateField(Builder,
-                    fieldNameOffset, child.IsNullable, fieldType.Type, fieldType.Offset,
-                    default, fieldChildrenOffsets, default);
+            for (var i = 0; i < childrenCount; i++)
+            {
+                var childField = type.Children[i];
+                var childFieldNameOffset = Builder.CreateString(childField.Name);
+                var childFieldType = _fieldTypeBuilder.BuildFieldType(childField);
+                var childFieldChildrenVectorOffset = Builder.CreateVectorOfTables(GetChildrenFieldOffsets(childField));
+
+                children[i] = Flatbuf.Field.CreateField(Builder,
+                    childFieldNameOffset, childField.IsNullable, childFieldType.Type, childFieldType.Offset,
+                    default, childFieldChildrenVectorOffset, default);
             }
+            return children;
         }
 
         private async ValueTask<Offset<Flatbuf.Schema>> WriteSchemaAsync(Schema schema, CancellationToken cancellationToken)

csharp/test/Apache.Arrow.Tests/TestData.cs

@@ -16,7 +16,6 @@
 using Apache.Arrow.Types;
 using System;
 using System.Collections.Generic;
-using System.Linq;
 
 namespace Apache.Arrow.Tests
 {
@@ -86,7 +85,7 @@ private static IArrowArray CreateArray(Field field, int length)
             field.DataType.Accept(creator);
 
             ArrayData data = new ArrayData(field.DataType, length, 0, 0,
-                new[] { ArrowBuffer.Empty, creator.Buffer }, creator.Children?.Select(_ => _.Data));
+                new[] { ArrowBuffer.Empty, creator.Buffer }, creator.Children);
 
             return ArrowArrayFactory.BuildArray(data);
         }
@@ -110,7 +109,7 @@ private class ArrayBufferCreator :
             private readonly int _length;
             public ArrowBuffer Buffer { get; private set; }
 
-            public List<IArrowArray> Children { get; private set; }
+            public ArrayData[] Children { get; private set; }
 
             public ArrayBufferCreator(int length)
             {
@@ -174,7 +173,7 @@ public void Visit(UInt16Type type)
             public void Visit(ListType type)
             {
                 //Buffer is valueOffsetsBuffer
-                Children = new List<IArrowArray> { CreateArray(type.ValueField, _length) };
+                Children = new[] { CreateArray(type.ValueField, _length).Data };
                 ArrowBuffer.Builder<int> builder = new ArrowBuffer.Builder<int>(_length);
                 for (int i = 0; i < _length; i++)
                     builder.Append(i);