fix: Concatenate extension arrays and keep concatenated results valid after inputs are disposed

src/Apache.Arrow/Arrays/ArrayDataConcatenator.cs

@@ -32,17 +32,43 @@ public static ArrayData Concatenate(IReadOnlyList<ArrayData> arrayDataList, Memo
 
             if (arrayDataList.Count == 1)
             {
-                return arrayDataList[0];
+                // The result must stay valid after the input is disposed, so it takes its own references.
+                return arrayDataList[0].Retain();
+            }
+
+            IArrowType type = arrayDataList[0].DataType;
+            if (type is ExtensionType extensionType)
+            {
+                return ConcatenateExtension(extensionType, arrayDataList, allocator);
             }
 
             var arrowArrayConcatenationVisitor = new ArrayDataConcatenationVisitor(arrayDataList, allocator);
 
-            IArrowType type = arrayDataList[0].DataType;
             type.Accept(arrowArrayConcatenationVisitor);
 
             return arrowArrayConcatenationVisitor.Result;
         }
 
+        private static ArrayData ConcatenateExtension(ExtensionType type, IReadOnlyList<ArrayData> arrayDataList, MemoryAllocator allocator)
+        {
+            var storageDataList = new List<ArrayData>(arrayDataList.Count);
+            foreach (ArrayData arrayData in arrayDataList)
+            {
+                if (!(arrayData.DataType is ExtensionType otherType) ||
+                    otherType.Name != type.Name ||
+                    otherType.ExtensionMetadata != type.ExtensionMetadata)
+                {
+                    throw new ArgumentException(
+                        $"Cannot concatenate extension type {type.Name} with {arrayData.DataType.Name}");
+                }
+
+                storageDataList.Add(new ArrayData(otherType.StorageType, arrayData.Length, arrayData.NullCount, arrayData.Offset, arrayData.Buffers, arrayData.Children, arrayData.Dictionary));
+            }
+
+            ArrayData storage = Concatenate(storageDataList, allocator);
+            return new ArrayData(type, storage.Length, storage.NullCount, storage.Offset, storage.Buffers, storage.Children, storage.Dictionary);
+        }
+
         private class ArrayDataConcatenationVisitor :
             IArrowTypeVisitor<BooleanType>,
             IArrowTypeVisitor<FixedWidthType>,
@@ -443,7 +469,7 @@ public void Visit(RunEndEncodedType type)
                     // All inputs were empty. Reuse the first input's values child sliced to length
                     // 0 so we get a valid ArrayData with the correct buffer/child layout for the
                     // values type, regardless of what that type is.
-                    valuesResult = _arrayDataList[0].Children[1].Slice(0, 0);
+                    valuesResult = _arrayDataList[0].Children[1].SliceShared(0, 0);
                 }
                 else
                 {
@@ -704,9 +730,16 @@ private void ConcatenateBinaryViewArrayData(IArrowType type)
                 int index = 2;
                 foreach (ArrayData arrayData in _arrayDataList)
                 {
+                    // Must skip the same inputs as ConcatenateViewBuffer so buffer indices line up.
+                    if (arrayData.Length == 0)
+                    {
+                        continue;
+                    }
+
+                    // The result shares the inputs' data buffers, so it needs its own references to them.
                     for (int i = 2; i < arrayData.Buffers.Length; i++)
                     {
-                        buffers[index++] = arrayData.Buffers[i];
+                        buffers[index++] = arrayData.Buffers[i].Retain();
                     }
                 }
 
@@ -798,7 +831,7 @@ private ArrayData ConcatenateChildrenOrEmpty(List<ArrayData> children)
                     {
                         // All parent arrays are empty, but the nested array still needs a real
                         // zero-length child to preserve Arrow's structural invariant.
-                        return arrayData.Children[0].Slice(0, 0);
+                        return arrayData.Children[0].SliceShared(0, 0);
                     }
                 }
 

src/Apache.Arrow/Arrays/ArrowArrayConcatenator.cs

@@ -29,7 +29,8 @@ public static IArrowArray Concatenate(IReadOnlyList<IArrowArray> arrowArrayList,
 
             if (arrowArrayList.Count == 1)
             {
-                return arrowArrayList[0];
+                // The result must stay valid after the input is disposed, so it takes its own references.
+                return ArrowArrayFactory.BuildArray(arrowArrayList[0].Data.Retain());
             }
 
             var arrayDataList = new List<ArrayData>(arrowArrayList.Count);

test/Apache.Arrow.Tests/ArrowArrayConcatenatorTests.cs

@@ -424,6 +424,156 @@ public void TestRunEndEncodedWithEmptyInput()
             Assert.Equal("B", values.GetString(1));
         }
 
+        [Fact]
+        public void TestExtensionArrays()
+        {
+            var guids = new[] { Guid.NewGuid(), Guid.NewGuid(), Guid.NewGuid() };
+            var a = new GuidArray.Builder().Append(guids[0]).AppendNull().Build();
+            var b = new GuidArray.Builder().Append(guids[1]).Append(guids[2]).Build();
+
+            var concatenated = Assert.IsType<GuidArray>(ArrowArrayConcatenator.Concatenate(new IArrowArray[] { a, b }));
+
+            Assert.Equal(4, concatenated.Length);
+            Assert.Equal(1, concatenated.NullCount);
+            Assert.Equal(new Guid?[] { guids[0], null, guids[1], guids[2] }, concatenated.ToArray());
+        }
+
+        [Fact]
+        public void TestListOfExtensionArrays()
+        {
+            var guids = new[] { Guid.NewGuid(), Guid.NewGuid() };
+            ListArray List(Guid guid)
+            {
+                var values = new GuidArray.Builder().Append(guid).Build();
+                var offsets = new ArrowBuffer.Builder<int>().Append(0).Append(1).Build();
+                return new ListArray(new ListType(GuidType.Default), 1, offsets, values, ArrowBuffer.Empty);
+            }
+
+            var concatenated = (ListArray)ArrowArrayConcatenator.Concatenate(new IArrowArray[] { List(guids[0]), List(guids[1]) });
+
+            var values = Assert.IsType<GuidArray>(concatenated.Values);
+            Assert.Equal(new Guid?[] { guids[0], guids[1] }, values.ToArray());
+        }
+
+        [Fact]
+        public void TestMismatchedExtensionTypesThrow()
+        {
+            var a = new GuidArray.Builder().Append(Guid.NewGuid()).Build();
+            var b = new FixedSizeBinaryArray(new ArrayData(new FixedSizeBinaryType(16), a.Length, 0, 0, a.Data.Buffers));
+
+            Assert.Throws<ArgumentException>(() => ArrowArrayConcatenator.Concatenate(new IArrowArray[] { a, b }));
+        }
+
+        [Fact]
+        public void TestStringViewResultOutlivesInputs()
+        {
+            var a = BuildStringView("a string longer than twelve bytes");
+            var b = BuildStringView("another string longer than twelve");
+
+            var concatenated = (StringViewArray)ArrowArrayConcatenator.Concatenate(new IArrowArray[] { a, b });
+            a.Dispose();
+            b.Dispose();
+
+            Assert.Equal("a string longer than twelve bytes", concatenated.GetString(0));
+            Assert.Equal("another string longer than twelve", concatenated.GetString(1));
+            concatenated.Dispose();
+        }
+
+        [Fact]
+        public void TestStringViewResultOutlivesSharedSliceInputs()
+        {
+            var source = BuildStringView("a string longer than twelve bytes", "another string longer than twelve");
+            var a = new StringViewArray(source.Data.SliceShared(0, 1));
+            var b = new StringViewArray(source.Data.SliceShared(1, 1));
+            source.Dispose();
+
+            var concatenated = (StringViewArray)ArrowArrayConcatenator.Concatenate(new IArrowArray[] { a, b });
+            a.Dispose();
+            b.Dispose();
+
+            Assert.Equal("a string longer than twelve bytes", concatenated.GetString(0));
+            Assert.Equal("another string longer than twelve", concatenated.GetString(1));
+            concatenated.Dispose();
+        }
+
+        [Fact]
+        public void TestStringViewWithEmptyInputHavingDataBuffers()
+        {
+            var empty = new StringViewArray(BuildStringView("zzzzzzzzzzzzzzzzzzzzzzzzzzzzzz").Data.Slice(0, 0));
+            var b = BuildStringView("another string longer than twelve");
+            var c = BuildStringView("short", "yet another string longer than twelve");
+
+            var concatenated = (StringViewArray)ArrowArrayConcatenator.Concatenate(new IArrowArray[] { empty, b, empty, c, empty });
+
+            Assert.Equal(3, concatenated.Length);
+            Assert.Equal("another string longer than twelve", concatenated.GetString(0));
+            Assert.Equal("short", concatenated.GetString(1));
+            Assert.Equal("yet another string longer than twelve", concatenated.GetString(2));
+        }
+
+        [Fact]
+        public void TestSingleElementResultOutlivesInput()
+        {
+            var array = BuildStringView("a string longer than twelve bytes");
+
+            var concatenated = (StringViewArray)ArrowArrayConcatenator.Concatenate(new IArrowArray[] { array });
+            array.Dispose();
+
+            Assert.Equal("a string longer than twelve bytes", concatenated.GetString(0));
+            concatenated.Dispose();
+        }
+
+        [Fact]
+        public void TestListWithSingleNonEmptyInputOutlivesInputs()
+        {
+            var builder = new ListArray.Builder(StringType.Default);
+            var valueBuilder = (StringArray.Builder)builder.ValueBuilder;
+            builder.Append();
+            valueBuilder.Append("x");
+            valueBuilder.Append("y");
+            var nonEmpty = builder.Build();
+            var empty = new ListArray.Builder(StringType.Default).Build();
+
+            var concatenated = (ListArray)ArrowArrayConcatenator.Concatenate(new IArrowArray[] { empty, nonEmpty });
+            nonEmpty.Dispose();
+            empty.Dispose();
+
+            var values = (StringArray)concatenated.Values;
+            Assert.Equal("x", values.GetString(0));
+            Assert.Equal("y", values.GetString(1));
+            concatenated.Dispose();
+        }
+
+        [Fact]
+        public void TestRunEndEncodedWithSingleNonEmptyInputOutlivesInputs()
+        {
+            var empty = new RunEndEncodedArray(
+                new Int32Array.Builder().Build(),
+                new StringArray.Builder().Build());
+            var nonEmpty = new RunEndEncodedArray(
+                new Int32Array.Builder().AppendRange(new[] { 2, 4 }).Build(),
+                new StringArray.Builder().AppendRange(new[] { "A", "B" }).Build());
+
+            var concatenated = (RunEndEncodedArray)ArrowArrayConcatenator.Concatenate(new IArrowArray[] { empty, nonEmpty });
+            nonEmpty.Dispose();
+            empty.Dispose();
+
+            var values = (StringArray)concatenated.Values;
+            Assert.Equal("A", values.GetString(0));
+            Assert.Equal("B", values.GetString(1));
+            concatenated.Dispose();
+        }
+
+        private static StringViewArray BuildStringView(params string[] values)
+        {
+            var builder = new StringViewArray.Builder();
+            foreach (string value in values)
+            {
+                builder.Append(value);
+            }
+            return builder.Build();
+        }
+
         private static IEnumerable<Tuple<List<IArrowArray>, IArrowArray>> GenerateTestData(bool slicedArrays = false)
         {
             var targetTypes = new List<IArrowType>() {