ARROW-7387: [C#] Support ListType Serialization

takashi hashida · HashidaTKS · eerhardt · commit 7f905a61dec9 · 2020-01-08T12:44:04.000-06:00
Closes #6030 from HashidaTKS/support_ListType_serialization and squashes the following commits: f2add6e <takashi hashida> Respond to changes of master branch bed012b <takashi hashida> Merge remote-tracking branch 'remotes/upstream/master' into support_ListType_serialization 35ff97b <takashi hashida> Format codes Remove extra brace f026d5a <takashi hashida> Respond to PR feedback cd6ad25 <takashi hashida> Don't use Linq Avoid unnecessary allocation Change RecordBatchManipulator to RecordBatchEnumerator 9b6dd27 <takashi hashida> Fix names 46ef481 <takashi hashida> Support ListType serialization / deserialization Lead-authored-by: takashi hashida <t-hashida@amiya.co.jp> Co-authored-by: takashi hashida <ellipsoid1989@gmail.com> Signed-off-by: Eric Erhardt <eric.erhardt@microsoft.com>
diff --git a/csharp/src/Apache.Arrow/Arrays/ListArray.cs b/csharp/src/Apache.Arrow/Arrays/ListArray.cs
@@ -30,16 +30,21 @@ public ListArray(IArrowType dataType, int length,
             ArrowBuffer valueOffsetsBuffer, IArrowArray values,
             ArrowBuffer nullBitmapBuffer, int nullCount = 0, int offset = 0)
             : this(new ArrayData(dataType, length, nullCount, offset,
-                new[] { nullBitmapBuffer, valueOffsetsBuffer }, new[] { values.Data }))
+                new[] { nullBitmapBuffer, valueOffsetsBuffer }, new[] { values.Data }),
+                values)
         {
-            Values = values;
         }
 
         public ListArray(ArrayData data)
-            : base(data)
+            : this(data, ArrowArrayFactory.BuildArray(data.Children[0]))
+        {
+        }
+
+        private ListArray(ArrayData data, IArrowArray values) : base(data)
         {
             data.EnsureBufferCount(2);
             data.EnsureDataType(ArrowTypeId.List);
+            Values = values;
         }
 
         public override void Accept(IArrowArrayVisitor visitor) => Accept(this, visitor);
diff --git a/csharp/src/Apache.Arrow/Ipc/ArrowReaderImplementation.cs b/csharp/src/Apache.Arrow/Ipc/ArrowReaderImplementation.cs
@@ -21,6 +21,7 @@
 using System.IO;
 using System.Threading;
 using System.Threading.Tasks;
+using Apache.Arrow.Types;
 
 namespace Apache.Arrow.Ipc
 {
@@ -114,35 +115,42 @@ private List<IArrowArray> BuildArrays(
             Flatbuf.RecordBatch recordBatchMessage)
         {
             var arrays = new List<IArrowArray>(recordBatchMessage.NodesLength);
-            int bufferIndex = 0;
 
-            for (var n = 0; n < recordBatchMessage.NodesLength; n++)
+            if (recordBatchMessage.NodesLength == 0)
             {
-                Field field = schema.GetFieldByIndex(n);
-                Flatbuf.FieldNode fieldNode = recordBatchMessage.Nodes(n).GetValueOrDefault();
+                return arrays;
+            }
+
+            var recordBatchEnumerator = new RecordBatchEnumerator(in recordBatchMessage);
+
+            do
+            {
+                var field = schema.GetFieldByIndex(recordBatchEnumerator.CurrentNodeIndex);
+                var fieldNode = recordBatchEnumerator.CurrentNode;
 
-                ArrayData arrayData = field.DataType.IsFixedPrimitive() ?
-                    LoadPrimitiveField(field, fieldNode, recordBatchMessage, messageBuffer, ref bufferIndex) :
-                    LoadVariableField(field, fieldNode, recordBatchMessage, messageBuffer, ref bufferIndex);
+                var arrayData = field.DataType.IsFixedPrimitive()
+                    ? LoadPrimitiveField(ref recordBatchEnumerator, field, in fieldNode, messageBuffer)
+                    : LoadVariableField(ref recordBatchEnumerator, field, in fieldNode, messageBuffer);
 
                 arrays.Add(ArrowArrayFactory.BuildArray(arrayData));
-            }
+            } while (recordBatchEnumerator.MoveNextNode());
 
             return arrays;
         }
 
+
         private ArrayData LoadPrimitiveField(
+            ref RecordBatchEnumerator recordBatchEnumerator,
             Field field,
-            Flatbuf.FieldNode fieldNode,
-            Flatbuf.RecordBatch recordBatch,
-            ByteBuffer bodyData,
-            ref int bufferIndex)
+            in Flatbuf.FieldNode fieldNode,
+            ByteBuffer bodyData)
         {
-            var nullBitmapBuffer = recordBatch.Buffers(bufferIndex++).GetValueOrDefault();
-            var valueBuffer = recordBatch.Buffers(bufferIndex++).GetValueOrDefault();
 
-            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;
@@ -158,24 +166,25 @@ private ArrayData LoadPrimitiveField(
             }
 
             var arrowBuff = new[] { nullArrowBuffer, valueArrowBuffer };
+            var children = GetChildren(ref recordBatchEnumerator, field, bodyData);
 
-            return new ArrayData(field.DataType, fieldLength, fieldNullCount, 0, arrowBuff);
+            return new ArrayData(field.DataType, fieldLength, fieldNullCount, 0, arrowBuff, children);
         }
 
+
         private ArrayData LoadVariableField(
+            ref RecordBatchEnumerator recordBatchEnumerator,
             Field field,
-            Flatbuf.FieldNode fieldNode,
-            Flatbuf.RecordBatch recordBatch,
-            ByteBuffer bodyData,
-            ref int bufferIndex)
+            in Flatbuf.FieldNode fieldNode,
+            ByteBuffer bodyData)
         {
-            var nullBitmapBuffer = recordBatch.Buffers(bufferIndex++).GetValueOrDefault();
-            var offsetBuffer = recordBatch.Buffers(bufferIndex++).GetValueOrDefault();
-            var valueBuffer = recordBatch.Buffers(bufferIndex++).GetValueOrDefault();
 
-            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;
@@ -191,8 +200,33 @@ private ArrayData LoadVariableField(
             }
 
             var arrowBuff = new[] { nullArrowBuffer, offsetArrowBuffer, valueArrowBuffer };
+            var children = GetChildren(ref recordBatchEnumerator, field, bodyData);
 
-            return new ArrayData(field.DataType, fieldLength, fieldNullCount, 0, arrowBuff);
+            return new ArrayData(field.DataType, fieldLength, fieldNullCount, 0, arrowBuff, children);
+        }
+
+        private ArrayData[] GetChildren(
+            ref RecordBatchEnumerator recordBatchEnumerator,
+            Field field,
+            ByteBuffer bodyData)
+        {
+            if (!(field.DataType is NestedType type)) return null;
+
+            var childrenCount = type.Children.Count;
+            var children = new ArrayData[childrenCount];
+            for (var index = 0; index < childrenCount; index++)
+            {
+                Flatbuf.FieldNode childFieldNode = recordBatchEnumerator.CurrentNode;
+                recordBatchEnumerator.MoveNextNode();
+
+                var childField = type.Children[index];
+                var child = childField.DataType.IsFixedPrimitive()
+                    ? LoadPrimitiveField(ref recordBatchEnumerator, childField, in childFieldNode, bodyData)
+                    : LoadVariableField(ref recordBatchEnumerator, childField, in childFieldNode, bodyData);
+
+                children[index] = child;
+            }
+            return children;
         }
 
         private ArrowBuffer BuildArrowBuffer(ByteBuffer bodyData, Flatbuf.Buffer buffer)
@@ -209,4 +243,32 @@ private ArrowBuffer BuildArrowBuffer(ByteBuffer bodyData, Flatbuf.Buffer buffer)
             return new ArrowBuffer(data);
         }
     }
+
+    internal struct RecordBatchEnumerator
+    {
+        private Flatbuf.RecordBatch RecordBatch { get; }
+        internal int CurrentBufferIndex { get; private set; }
+        internal int CurrentNodeIndex { get; private set; }
+
+        internal Flatbuf.Buffer CurrentBuffer => RecordBatch.Buffers(CurrentBufferIndex).GetValueOrDefault();
+
+        internal Flatbuf.FieldNode CurrentNode => RecordBatch.Nodes(CurrentNodeIndex).GetValueOrDefault();
+
+        internal bool MoveNextBuffer()
+        {
+            return ++CurrentBufferIndex < RecordBatch.BuffersLength;
+        }
+
+        internal bool MoveNextNode()
+        {
+            return ++CurrentNodeIndex < RecordBatch.NodesLength;
+        }
+
+        internal RecordBatchEnumerator(in Flatbuf.RecordBatch recordBatch)
+        {
+            RecordBatch = recordBatch;
+            CurrentBufferIndex = 0;
+            CurrentNodeIndex = 0;
+        }
+    }
 }
diff --git a/csharp/src/Apache.Arrow/Ipc/ArrowStreamWriter.cs b/csharp/src/Apache.Arrow/Ipc/ArrowStreamWriter.cs
@@ -20,6 +20,7 @@
 using System.IO;
 using System.Threading;
 using System.Threading.Tasks;
+using Apache.Arrow.Types;
 using FlatBuffers;
 
 namespace Apache.Arrow.Ipc
@@ -108,7 +109,7 @@ private void CreateBuffers(BooleanArray array)
             }
 
             private void CreateBuffers<T>(PrimitiveArray<T> array)
-                where T: struct
+                where T : struct
             {
                 _buffers.Add(CreateBuffer(array.NullBitmapBuffer));
                 _buffers.Add(CreateBuffer(array.ValueBuffer));
@@ -175,6 +176,42 @@ public ArrowStreamWriter(Stream baseStream, Schema schema, bool leaveOpen, IpcOp
             _options = options ?? IpcOptions.Default;
         }
 
+
+        private void CreateSelfAndChildrenFieldNodes(ArrayData data)
+        {
+            if (data.DataType is NestedType)
+            {
+                // 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);
+        }
+
+        private int CountAllNodes()
+        {
+            var count = 0;
+            foreach (var arrowArray in Schema.Fields.Values)
+            {
+                CountSelfAndChildrenNodes(arrowArray.DataType, ref count);
+            }
+            return count;
+        }
+
+        private void CountSelfAndChildrenNodes(IArrowType type, ref int count)
+        {
+            if (type is NestedType nestedType)
+            {
+                foreach (var childField in nestedType.Children)
+                {
+                    CountSelfAndChildrenNodes(childField.DataType, ref count);
+                }
+            }
+            count++;
+        }
+
         private protected async Task WriteRecordBatchInternalAsync(RecordBatch recordBatch,
             CancellationToken cancellationToken = default)
         {
@@ -192,13 +229,12 @@ private protected async Task WriteRecordBatchInternalAsync(RecordBatch recordBat
 
             var fieldCount = Schema.Fields.Count;
 
-            Flatbuf.RecordBatch.StartNodesVector(Builder, fieldCount);
+            Flatbuf.RecordBatch.StartNodesVector(Builder, CountAllNodes());
 
             // flatbuffer struct vectors have to be created in reverse order
             for (var i = fieldCount - 1; i >= 0; i--)
             {
-                var fieldArray = recordBatch.Column(i);
-                Flatbuf.FieldNode.CreateFieldNode(Builder, fieldArray.Length, fieldArray.NullCount);
+                CreateSelfAndChildrenFieldNodes(recordBatch.Column(i).Data);
             }
 
             var fieldNodesVectorOffset = Builder.EndVector();
@@ -285,7 +321,7 @@ public virtual Task WriteRecordBatchAsync(RecordBatch recordBatch, CancellationT
         {
             return WriteRecordBatchInternalAsync(recordBatch, cancellationToken);
         }
-        
+
         public async Task WriteEndAsync(CancellationToken cancellationToken = default)
         {
             if (!HasWrittenEnd)
@@ -307,19 +343,18 @@ private ValueTask WriteBufferAsync(ArrowBuffer arrowBuffer, CancellationToken ca
             // Build fields
 
             var fieldOffsets = new Offset<Flatbuf.Field>[schema.Fields.Count];
-            var fieldChildren = new List<Offset<Flatbuf.Field>>();
 
             for (var i = 0; i < fieldOffsets.Length; i++)
             {
                 var field = schema.GetFieldByIndex(i);
                 var fieldNameOffset = Builder.CreateString(field.Name);
                 var fieldType = _fieldTypeBuilder.BuildFieldType(field);
 
-                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);
@@ -332,6 +367,30 @@ private ValueTask WriteBufferAsync(ArrowBuffer arrowBuffer, CancellationToken ca
                 Builder, endianness, fieldsVectorOffset);
         }
 
+        private protected Offset<Flatbuf.Field>[] GetChildrenFieldOffsets(Field field)
+        {
+            if (!(field.DataType is NestedType type))
+            {
+                return System.Array.Empty<Offset<Flatbuf.Field>>();
+            }
+
+            var childrenCount = type.Children.Count;
+            var children = new Offset<Flatbuf.Field>[childrenCount];
+
+            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)
         {
             Builder.Clear();
@@ -357,7 +416,7 @@ await WriteMessageAsync(Flatbuf.MessageHeader.Schema, schemaOffset, 0, cancellat
         private async ValueTask<long> WriteMessageAsync<T>(
             Flatbuf.MessageHeader headerType, Offset<T> headerOffset, int bodyLength,
             CancellationToken cancellationToken)
-            where T: struct
+            where T : struct
         {
             var messageOffset = Flatbuf.Message.CreateMessage(
                 Builder, CurrentMetadataVersion, headerType, headerOffset.Value,
diff --git a/csharp/src/Apache.Arrow/Ipc/ArrowTypeFlatbufferBuilder.cs b/csharp/src/Apache.Arrow/Ipc/ArrowTypeFlatbufferBuilder.cs
@@ -98,7 +98,10 @@ public void Visit(BinaryType type)
 
             public void Visit(ListType type)
             {
-                throw new NotImplementedException();
+                Flatbuf.List.StartList(Builder);
+                Result = FieldType.Build(
+                    Flatbuf.Type.List, 
+                    Flatbuf.List.EndList(Builder));
             }
 
             public void Visit(UnionType type)
diff --git a/csharp/src/Apache.Arrow/Ipc/MessageSerializer.cs b/csharp/src/Apache.Arrow/Ipc/MessageSerializer.cs
@@ -125,6 +125,12 @@ private static Types.IArrowType GetFieldArrowType(Flatbuf.Field field)
                     return new Types.StringType();
                 case Flatbuf.Type.Binary:
                     return Types.BinaryType.Default;
+                case Flatbuf.Type.List:
+                    if (field.ChildrenLength != 1)
+                    {
+                        throw new InvalidDataException($"List type must have only one child.");
+                    }
+                    return new Types.ListType(GetFieldArrowType(field.Children(0).GetValueOrDefault()));
                 default:
                     throw new InvalidDataException($"Arrow primitive '{field.TypeType}' is unsupported.");
             }
diff --git a/csharp/src/Apache.Arrow/Types/ListType.cs b/csharp/src/Apache.Arrow/Types/ListType.cs
@@ -17,21 +17,19 @@
 
 namespace Apache.Arrow.Types
 {
-    public sealed class ListType : ArrowType
+    public sealed class ListType : NestedType
     {
         public override ArrowTypeId TypeId => ArrowTypeId.List;
         public override string Name => "list";
 
-        public Field ValueField { get; }
-        public IArrowType ValueDataType { get; }
+        public Field ValueField => Children[0];
+
+        public IArrowType ValueDataType => Children[0].DataType;
 
         public ListType(Field valueField)
-        {
-            ValueField = valueField ?? throw new ArgumentNullException(nameof(valueField));
-            ValueDataType = ValueField.DataType;
-        }
+           : base(valueField) { }
 
-        public ListType(IArrowType valueDataType) 
+        public ListType(IArrowType valueDataType)
             : this(new Field("item", valueDataType, true)) { }
 
         public override void Accept(IArrowTypeVisitor visitor) => Accept(this, visitor);
diff --git a/csharp/src/Apache.Arrow/Types/NestedType.cs b/csharp/src/Apache.Arrow/Types/NestedType.cs
diff --git a/csharp/test/Apache.Arrow.Tests/ArrowReaderVerifier.cs b/csharp/test/Apache.Arrow.Tests/ArrowReaderVerifier.cs
diff --git a/csharp/test/Apache.Arrow.Tests/TestData.cs b/csharp/test/Apache.Arrow.Tests/TestData.cs

Original file line number	Diff line number	Diff line change
`@@ -30,16 +30,21 @@ public ListArray(IArrowType dataType, int length,`
`30`	`30`	`ArrowBuffer valueOffsetsBuffer, IArrowArray values,`
`31`	`31`	`ArrowBuffer nullBitmapBuffer, int nullCount = 0, int offset = 0)`
`32`	`32`	`: this(new ArrayData(dataType, length, nullCount, offset,`
`33`		`- new[] { nullBitmapBuffer, valueOffsetsBuffer }, new[] { values.Data }))`
	`33`	`+ new[] { nullBitmapBuffer, valueOffsetsBuffer }, new[] { values.Data }),`
	`34`	`+ values)`
`34`	`35`	`{`
`35`		`- Values = values;`
`36`	`36`	`}`
`37`	`37`
`38`	`38`	`public ListArray(ArrayData data)`
`39`		`- : base(data)`
	`39`	`+ : this(data, ArrowArrayFactory.BuildArray(data.Children[0]))`
	`40`	`+ {`
	`41`	`+ }`
	`42`	`+`
	`43`	`+ private ListArray(ArrayData data, IArrowArray values) : base(data)`
`40`	`44`	`{`
`41`	`45`	`data.EnsureBufferCount(2);`
`42`	`46`	`data.EnsureDataType(ArrowTypeId.List);`
	`47`	`+ Values = values;`
`43`	`48`	`}`
`44`	`49`
`45`	`50`	`public override void Accept(IArrowArrayVisitor visitor) => Accept(this, visitor);`
Original file line number	Diff line number	Diff line change
`@@ -98,7 +98,10 @@ public void Visit(BinaryType type)`
`98`	`98`
`99`	`99`	`public void Visit(ListType type)`
`100`	`100`	`{`
`101`		`- throw new NotImplementedException();`
	`101`	`+ Flatbuf.List.StartList(Builder);`
	`102`	`+ Result = FieldType.Build(`
	`103`	`+ Flatbuf.Type.List,`
	`104`	`+ Flatbuf.List.EndList(Builder));`
`102`	`105`	`}`
`103`	`106`
`104`	`107`	`public void Visit(UnionType type)`