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

### What changes are included in this PR? Support dense and sparse unions in the C# implementation. Adds Archery support for C# unions. ### Are these changes tested? Yes ### Are there any user-facing changes? Unions are now supported in the C# implementation. **This PR includes breaking changes to public APIs.** The public APIs for the UnionArray and UnionType were changed fairly substantially. As these were previously not implemented properly, the impact of the changes ought to be minimal. The ChunkedArray and Column classes were changed to hold IArrowArrays instead of Arrays. To accomodate this, a constructor was added which may introduce ambiguity in calling code. This could be avoided by changing the overloaded constructor to instead be a factory method. This didn't seem worthwhile but could be reconsidered. The metadata version was finally increased to V5. * Closes: #36795 Authored-by: Curt Hagenlocher <[email protected]> Signed-off-by: David Li <[email protected]>
apache · Sep 25, 2023 · e55f912 · e55f912
1 parent 64acef8
commit e55f912
Show file tree

Hide file tree

Showing 32 changed files with 797 additions and 119 deletions.
diff --git a/csharp/src/Apache.Arrow/Arrays/Array.cs b/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)
             }
         }
     }
-}
+}
diff --git a/csharp/src/Apache.Arrow/Arrays/ArrayDataConcatenator.cs b/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);

diff --git a/csharp/src/Apache.Arrow/Arrays/ArrayDataTypeComparer.cs b/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)

diff --git a/csharp/src/Apache.Arrow/Arrays/ArrowArrayFactory.cs b/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);
+        }
     }
 }
diff --git a/csharp/src/Apache.Arrow/Arrays/DenseUnionArray.cs b/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);
+        }
+    }
+}
diff --git a/csharp/src/Apache.Arrow/Arrays/PrimitiveArrayBuilder.cs b/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;

diff --git a/csharp/src/Apache.Arrow/Arrays/SparseUnionArray.cs b/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);
+        }
+    }
+}
diff --git a/csharp/src/Apache.Arrow/Arrays/UnionArray.cs b/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;
+        }
     }
 }