GH-36795: [C#] Implement support for dense and sparse unions

csharp/src/Apache.Arrow/Arrays/Array.cs

@@ -62,16 +62,7 @@ internal static void Accept<T>(T array, IArrowArrayVisitor visitor)
 
         public Array Slice(int offset, int length)
         {
-            if (offset > Length)
-            {
-                throw new ArgumentException($"Offset {offset} cannot be greater than Length {Length} for Array.Slice");
-            }
-
-            length = Math.Min(Data.Length - offset, length);
-            offset += Data.Offset;
-
-            ArrayData newData = Data.Slice(offset, length);
-            return ArrowArrayFactory.BuildArray(newData) as Array;
+            return ArrowArrayFactory.Slice(this, offset, length) as Array;
         }
 
         public void Dispose()
@@ -88,4 +79,4 @@ protected virtual void Dispose(bool disposing)
             }
         }
     }
-}
+}

csharp/src/Apache.Arrow/Arrays/ArrayDataConcatenator.cs

@@ -49,7 +49,8 @@ private class ArrayDataConcatenationVisitor :
             IArrowTypeVisitor<StringType>,
             IArrowTypeVisitor<ListType>,
             IArrowTypeVisitor<FixedSizeListType>,
-            IArrowTypeVisitor<StructType>
+            IArrowTypeVisitor<StructType>,
+            IArrowTypeVisitor<UnionType>
         {
             public ArrayData Result { get; private set; }
             private readonly IReadOnlyList<ArrayData> _arrayDataList;
@@ -123,6 +124,33 @@ public void Visit(StructType type)
                 Result = new ArrayData(type, _arrayDataList[0].Length, _arrayDataList[0].NullCount, 0, _arrayDataList[0].Buffers, children);
             }
 
+            public void Visit(UnionType type)
+            {
+                int bufferCount = type.Mode switch
+                {
+                    UnionMode.Sparse => 1,
+                    UnionMode.Dense => 2,
+                    _ => throw new InvalidOperationException("TODO"),
+                };
+
+                CheckData(type, bufferCount);
+                List<ArrayData> children = new List<ArrayData>(type.Fields.Count);
+
+                for (int i = 0; i < type.Fields.Count; i++)
+                {
+                    children.Add(Concatenate(SelectChildren(i), _allocator));
+                }
+
+                ArrowBuffer[] buffers = new ArrowBuffer[bufferCount];
+                buffers[0] = ConcatenateUnionTypeBuffer();
+                if (bufferCount > 1)
+                {
+                    buffers[1] = ConcatenateUnionOffsetBuffer();
+                }
+
+                Result = new ArrayData(type, _totalLength, _totalNullCount, 0, buffers, children);
+            }
+
             public void Visit(IArrowType type)
             {
                 throw new NotImplementedException($"Concatenation for {type.Name} is not supported yet.");
@@ -231,6 +259,38 @@ private ArrowBuffer ConcatenateOffsetBuffer()
                 return builder.Build(_allocator);
             }
 
+            private ArrowBuffer ConcatenateUnionTypeBuffer()
+            {
+                var builder = new ArrowBuffer.Builder<byte>(_totalLength);
+
+                foreach (ArrayData arrayData in _arrayDataList)
+                {
+                    builder.Append(arrayData.Buffers[0]);
+                }
+
+                return builder.Build(_allocator);
+            }
+
+            private ArrowBuffer ConcatenateUnionOffsetBuffer()
+            {
+                var builder = new ArrowBuffer.Builder<int>(_totalLength);
+                int baseOffset = 0;
+
+                foreach (ArrayData arrayData in _arrayDataList)
+                {
+                    ReadOnlySpan<int> span = arrayData.Buffers[1].Span.CastTo<int>();
+                    foreach (int offset in span)
+                    {
+                        builder.Append(baseOffset + offset);
+                    }
+
+                    // The next offset must start from the current last offset.
+                    baseOffset += span[arrayData.Length];
+                }
+
+                return builder.Build(_allocator);
+            }
+
             private List<ArrayData> SelectChildren(int index)
             {
                 var children = new List<ArrayData>(_arrayDataList.Count);

csharp/src/Apache.Arrow/Arrays/ArrayDataTypeComparer.cs

@@ -27,7 +27,8 @@ internal sealed class ArrayDataTypeComparer :
         IArrowTypeVisitor<FixedSizeBinaryType>,
         IArrowTypeVisitor<ListType>,
         IArrowTypeVisitor<FixedSizeListType>,
-        IArrowTypeVisitor<StructType>
+        IArrowTypeVisitor<StructType>,
+        IArrowTypeVisitor<UnionType>
     {
         private readonly IArrowType _expectedType;
         private bool _dataTypeMatch;
@@ -122,6 +123,15 @@ public void Visit(StructType actualType)
             }
         }
 
+        public void Visit(UnionType actualType)
+        {
+            if (_expectedType is UnionType expectedType
+                && CompareNested(expectedType, actualType))
+            {
+                _dataTypeMatch = true;
+            }
+        }
+
         private static bool CompareNested(NestedType expectedType, NestedType actualType)
         {
             if (expectedType.Fields.Count != actualType.Fields.Count)

csharp/src/Apache.Arrow/Arrays/ArrowArrayFactory.cs

@@ -62,7 +62,7 @@ public static IArrowArray BuildArray(ArrayData data)
                 case ArrowTypeId.Struct:
                     return new StructArray(data);
                 case ArrowTypeId.Union:
-                    return new UnionArray(data);
+                    return UnionArray.Create(data);
                 case ArrowTypeId.Date64:
                     return new Date64Array(data);
                 case ArrowTypeId.Date32:
@@ -91,5 +91,19 @@ public static IArrowArray BuildArray(ArrayData data)
                     throw new NotSupportedException($"An ArrowArray cannot be built for type {data.DataType.TypeId}.");
             }
         }
+
+        public static IArrowArray Slice(IArrowArray array, int offset, int length)
+        {
+            if (offset > array.Length)
+            {
+                throw new ArgumentException($"Offset {offset} cannot be greater than Length {array.Length} for Array.Slice");
+            }
+
+            length = Math.Min(array.Data.Length - offset, length);
+            offset += array.Data.Offset;
+
+            ArrayData newData = array.Data.Slice(offset, length);
+            return BuildArray(newData);
+        }
     }
 }

csharp/src/Apache.Arrow/Arrays/DenseUnionArray.cs

@@ -0,0 +1,52 @@
+// 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 Apache.Arrow.Types;
+using System;
+using System.Collections.Generic;
+using System.Linq;
+
+namespace Apache.Arrow
+{
+    public class DenseUnionArray : UnionArray
+    {
+        public ArrowBuffer ValueOffsetBuffer => Data.Buffers[1];
+
+        public ReadOnlySpan<int> ValueOffsets => ValueOffsetBuffer.Span.CastTo<int>();
+
+        public DenseUnionArray(
+            IArrowType dataType,
+            int length,
+            IEnumerable<IArrowArray> children,
+            ArrowBuffer typeIds,
+            ArrowBuffer valuesOffsetBuffer,
+            int nullCount = 0,
+            int offset = 0)
+            : base(new ArrayData(
+                dataType, length, nullCount, offset, new[] { typeIds, valuesOffsetBuffer },
+                children.Select(child => child.Data)))
+        {
+            _fields = children.ToArray();
+            ValidateMode(UnionMode.Dense, Type.Mode);
+        }
+
+        public DenseUnionArray(ArrayData data) 
+            : base(data)
+        {
+            ValidateMode(UnionMode.Dense, Type.Mode);
+            data.EnsureBufferCount(2);
+        }
+    }
+}

csharp/src/Apache.Arrow/Arrays/PrimitiveArrayBuilder.cs

@@ -137,6 +137,9 @@ public TBuilder Append(T value)
             return Instance;
         }
 
+        public TBuilder Append(T? value) =>
+            (value == null) ? AppendNull() : Append(value.Value);
+
         public TBuilder Append(ReadOnlySpan<T> span)
         {
             int len = ValueBuffer.Length;

csharp/src/Apache.Arrow/Arrays/SparseUnionArray.cs

@@ -0,0 +1,46 @@
+// 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 Apache.Arrow.Types;
+using System.Collections.Generic;
+using System.Linq;
+
+namespace Apache.Arrow
+{
+    public class SparseUnionArray : UnionArray
+    {
+        public SparseUnionArray(
+            IArrowType dataType,
+            int length,
+            IEnumerable<IArrowArray> children,
+            ArrowBuffer typeIds,
+            int nullCount = 0,
+            int offset = 0)
+            : base(new ArrayData(
+                dataType, length, nullCount, offset, new[] { typeIds },
+                children.Select(child => child.Data)))
+        {
+            _fields = children.ToArray();
+            ValidateMode(UnionMode.Sparse, Type.Mode);
+        }
+
+        public SparseUnionArray(ArrayData data) 
+            : base(data)
+        {
+            ValidateMode(UnionMode.Sparse, Type.Mode);
+            data.EnsureBufferCount(1);
+        }
+    }
+}

csharp/src/Apache.Arrow/Arrays/UnionArray.cs

@@ -15,37 +15,88 @@
 
 using Apache.Arrow.Types;
 using System;
+using System.Collections.Generic;
+using System.Threading;
 
 namespace Apache.Arrow
 {
-    public class UnionArray: Array
+    public abstract class UnionArray : IArrowArray
     {
-        public UnionType Type => Data.DataType as UnionType;
+        protected IReadOnlyList<IArrowArray> _fields;
 
-        public UnionMode Mode => Type.Mode;
+        public IReadOnlyList<IArrowArray> Fields =>
+            LazyInitializer.EnsureInitialized(ref _fields, () => InitializeFields());
+
+        public ArrayData Data { get; }
 
-        public ArrowBuffer TypeBuffer => Data.Buffers[1];
+        public UnionType Type => (UnionType)Data.DataType;
 
-        public ArrowBuffer ValueOffsetBuffer => Data.Buffers[2];
+        public UnionMode Mode => Type.Mode;
+
+        public ArrowBuffer TypeBuffer => Data.Buffers[0];
 
         public ReadOnlySpan<byte> TypeIds => TypeBuffer.Span;
 
-        public ReadOnlySpan<int> ValueOffsets => ValueOffsetBuffer.Span.CastTo<int>().Slice(0, Length + 1);
+        public int Length => Data.Length;
+
+        public int Offset => Data.Offset;
 
-        public UnionArray(ArrayData data) 
-            : base(data)
+        public int NullCount => Data.NullCount;
+
+        public bool IsValid(int index) => NullCount == 0 || Fields[TypeIds[index]].IsValid(index);
+
+        public bool IsNull(int index) => !IsValid(index);
+
+        protected UnionArray(ArrayData data) 
         {
+            Data = data;
             data.EnsureDataType(ArrowTypeId.Union);
-            data.EnsureBufferCount(3);
         }
 
-        public IArrowArray GetChild(int index)
+        public static UnionArray Create(ArrayData data)
         {
-            // TODO: Implement
-            throw new NotImplementedException();
+            return ((UnionType)data.DataType).Mode switch
+            {
+                UnionMode.Dense => new DenseUnionArray(data),
+                UnionMode.Sparse => new SparseUnionArray(data),
+                _ => throw new InvalidOperationException("unknown union mode in array creation")
+            };
         }
 
-        public override void Accept(IArrowArrayVisitor visitor) => Accept(this, visitor);
+        public void Accept(IArrowArrayVisitor visitor) => Array.Accept(this, visitor);
 
+        public void Dispose()
+        {
+            Dispose(true);
+            GC.SuppressFinalize(this);
+        }
+
+        protected virtual void Dispose(bool disposing)
+        {
+            if (disposing)
+            {
+                Data.Dispose();
+            }
+        }
+
+        protected static void ValidateMode(UnionMode expected, UnionMode actual)
+        {
+            if (expected != actual)
+            {
+                throw new ArgumentException(
+                    $"Specified union mode <{actual}> does not match expected mode <{expected}>",
+                    "Mode");
+            }
+        }
+
+        private IReadOnlyList<IArrowArray> InitializeFields()
+        {
+            IArrowArray[] result = new IArrowArray[Data.Children.Length];
+            for (int i = 0; i < Data.Children.Length; i++)
+            {
+                result[i] = ArrowArrayFactory.BuildArray(Data.Children[i]);
+            }
+            return result;
+        }
     }
 }