feat: Improve the implementation of TimestampWithOffsetArray.Builder

src/Apache.Arrow/Arrays/TimestampWithOffsetArray.cs

@@ -16,6 +16,7 @@
 using System;
 using System.Collections;
 using System.Collections.Generic;
+using Apache.Arrow.Memory;
 using Apache.Arrow.Types;
 
 namespace Apache.Arrow
@@ -196,39 +197,57 @@ public TimestampWithOffsetArray(TimestampWithOffsetType type, IArrowArray storag
         /// <summary>
         /// Builder for <see cref="TimestampWithOffsetArray"/>.
         /// </summary>
-        public class Builder
+        public class Builder : IArrowArrayBuilder<DateTimeOffset, TimestampWithOffsetArray, Builder>
         {
+            private static readonly DateTimeOffset s_nullPlaceholder =
+                new DateTimeOffset(1970, 1, 1, 0, 0, 0, TimeSpan.Zero);
+
             private readonly TimestampArray.Builder _timestampBuilder;
             private readonly Int16Array.Builder _offsetBuilder;
             private readonly ArrowBuffer.BitmapBuilder _validityBuilder;
             private readonly TimestampWithOffsetType _type;
-            private int _length;
-            private int _nullCount;
 
             public Builder(TimeUnit unit = TimeUnit.Microsecond)
+                : this(new TimestampWithOffsetType(unit))
+            {
+            }
+
+            public Builder(TimestampWithOffsetType type)
             {
-                _type = new TimestampWithOffsetType(unit);
-                _timestampBuilder = new TimestampArray.Builder(unit, "UTC");
+                _type = type ?? throw new ArgumentNullException(nameof(type));
+                _timestampBuilder = new TimestampArray.Builder(type.Unit, "UTC");
                 _offsetBuilder = new Int16Array.Builder();
                 _validityBuilder = new ArrowBuffer.BitmapBuilder();
             }
 
+            public int Length => _validityBuilder.Length;
+
             public Builder Append(DateTimeOffset value)
             {
                 _timestampBuilder.Append(value.ToUniversalTime());
                 _offsetBuilder.Append(checked((short)value.Offset.TotalMinutes));
                 _validityBuilder.Append(true);
-                _length++;
+                return this;
+            }
+
+            public Builder Append(DateTimeOffset? value) =>
+                value.HasValue ? Append(value.Value) : AppendNull();
+
+            public Builder Append(ReadOnlySpan<DateTimeOffset> values)
+            {
+                Reserve(values.Length);
+                foreach (DateTimeOffset value in values)
+                {
+                    Append(value);
+                }
                 return this;
             }
 
             public Builder AppendNull()
             {
-                _timestampBuilder.Append(default(DateTimeOffset));
+                _timestampBuilder.Append(s_nullPlaceholder);
                 _offsetBuilder.Append(0);
                 _validityBuilder.Append(false);
-                _length++;
-                _nullCount++;
                 return this;
             }
 
@@ -259,17 +278,72 @@ public Builder AppendRange(IEnumerable<DateTimeOffset?> values)
                 return this;
             }
 
-            public TimestampWithOffsetArray Build()
+            public Builder Reserve(int capacity)
+            {
+                _timestampBuilder.Reserve(capacity);
+                _offsetBuilder.Reserve(capacity);
+                _validityBuilder.Reserve(capacity);
+                return this;
+            }
+
+            public Builder Resize(int length)
+            {
+                if (length < 0)
+                    throw new ArgumentOutOfRangeException(nameof(length));
+
+                if (length < Length)
+                {
+                    _timestampBuilder.Resize(length);
+                    _offsetBuilder.Resize(length);
+                    _validityBuilder.Resize(length);
+                }
+                else
+                {
+                    Reserve(length - Length);
+                    while (Length < length)
+                    {
+                        AppendNull();
+                    }
+                }
+                return this;
+            }
+
+            public Builder Set(int index, DateTimeOffset value)
+            {
+                _timestampBuilder.Set(index, value.ToUniversalTime());
+                _offsetBuilder.Set(index, checked((short)value.Offset.TotalMinutes));
+                _validityBuilder.Set(index, true);
+                return this;
+            }
+
+            public Builder Swap(int i, int j)
+            {
+                _timestampBuilder.Swap(i, j);
+                _offsetBuilder.Swap(i, j);
+                _validityBuilder.Swap(i, j);
+                return this;
+            }
+
+            public Builder Clear()
+            {
+                _timestampBuilder.Clear();
+                _offsetBuilder.Clear();
+                _validityBuilder.Clear();
+                return this;
+            }
+
+            public TimestampWithOffsetArray Build(MemoryAllocator allocator = default)
             {
-                TimestampArray timestamps = _timestampBuilder.Build();
-                Int16Array offsets = _offsetBuilder.Build();
-                ArrowBuffer validityBuffer = _nullCount > 0 ? _validityBuilder.Build() : ArrowBuffer.Empty;
+                TimestampArray timestamps = _timestampBuilder.Build(allocator);
+                Int16Array offsets = _offsetBuilder.Build(allocator);
+                int nullCount = _validityBuilder.UnsetBitCount;
+                ArrowBuffer validityBuffer = nullCount > 0 ? _validityBuilder.Build(allocator) : ArrowBuffer.Empty;
 
                 var structType = (StructType)_type.StorageType;
                 var structArray = new StructArray(
-                    structType, _length,
+                    structType, Length,
                     new IArrowArray[] { timestamps, offsets },
-                    validityBuffer, _nullCount);
+                    validityBuffer, nullCount);
 
                 return new TimestampWithOffsetArray(_type, structArray);
             }

test/Apache.Arrow.Tests/TimestampWithOffsetArrayTests.cs

@@ -113,6 +113,67 @@ public void BuilderAppendRangeNullable()
             Assert.NotNull(array.GetValue(2));
         }
 
+        [Fact]
+        public void BuilderSupportsStandardOperations()
+        {
+            var first = new DateTimeOffset(2024, 1, 1, 1, 0, 0, TimeSpan.FromHours(1));
+            var second = new DateTimeOffset(2024, 2, 1, 2, 0, 0, TimeSpan.FromHours(2));
+            var replacement = new DateTimeOffset(2024, 3, 1, 3, 0, 0, TimeSpan.FromHours(3));
+            var builder = new TimestampWithOffsetArray.Builder();
+
+            Assert.Same(builder, builder.Reserve(4));
+            Assert.Same(builder, builder.Append(new[] { first, second }.AsSpan()));
+            Assert.Equal(2, builder.Length);
+            Assert.Same(builder, builder.Swap(0, 1));
+            Assert.Same(builder, builder.Set(1, replacement));
+
+            var array = builder.Build();
+            Assert.Equal(second, array.GetValue(0));
+            Assert.Equal(replacement, array.GetValue(1));
+
+            Assert.Same(builder, builder.Resize(4));
+            Assert.Equal(4, builder.Length);
+            array = builder.Build();
+            Assert.Equal(2, array.NullCount);
+            Assert.Null(array.GetValue(2));
+            Assert.Null(array.GetValue(3));
+
+            Assert.Same(builder, builder.Resize(1));
+            array = builder.Build();
+            Assert.Single(array);
+            Assert.Equal(second, array.GetValue(0));
+
+            Assert.Same(builder, builder.Clear());
+            Assert.Equal(0, builder.Length);
+            Assert.Empty(builder.Build());
+        }
+
+        [Fact]
+        public void BuilderImplementsArrayBuilderInterface()
+        {
+            IArrowArrayBuilder<DateTimeOffset, TimestampWithOffsetArray, TimestampWithOffsetArray.Builder> builder =
+                new TimestampWithOffsetArray.Builder();
+
+            TimestampWithOffsetArray array = builder
+                .Append(new DateTimeOffset(2024, 1, 1, 0, 0, 0, TimeSpan.Zero))
+                .Build();
+
+            Assert.Single(array);
+        }
+
+        [Fact]
+        public void NanosecondBuilderCanAppendNullAndResize()
+        {
+            var array = new TimestampWithOffsetArray.Builder(TimeUnit.Nanosecond)
+                .AppendNull()
+                .Resize(2)
+                .Build();
+
+            Assert.Equal(2, array.Length);
+            Assert.Equal(2, array.NullCount);
+            Assert.All(array, value => Assert.Null(value));
+        }
+
         [Fact]
         public void EmptyArray()
         {