diff --git a/src/Apache.Arrow/Arrays/ArrayDataConcatenator.cs b/src/Apache.Arrow/Arrays/ArrayDataConcatenator.cs index 3d843312..831fbfa1 100644 --- a/src/Apache.Arrow/Arrays/ArrayDataConcatenator.cs +++ b/src/Apache.Arrow/Arrays/ArrayDataConcatenator.cs @@ -32,17 +32,43 @@ public static ArrayData Concatenate(IReadOnlyList 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 arrayDataList, MemoryAllocator allocator) + { + var storageDataList = new List(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, IArrowTypeVisitor, @@ -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 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); } } diff --git a/src/Apache.Arrow/Arrays/ArrowArrayConcatenator.cs b/src/Apache.Arrow/Arrays/ArrowArrayConcatenator.cs index 01ca2438..64a9fcd7 100644 --- a/src/Apache.Arrow/Arrays/ArrowArrayConcatenator.cs +++ b/src/Apache.Arrow/Arrays/ArrowArrayConcatenator.cs @@ -29,7 +29,8 @@ public static IArrowArray Concatenate(IReadOnlyList 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(arrowArrayList.Count); diff --git a/test/Apache.Arrow.Tests/ArrowArrayConcatenatorTests.cs b/test/Apache.Arrow.Tests/ArrowArrayConcatenatorTests.cs index f5ed951b..4e781a78 100644 --- a/test/Apache.Arrow.Tests/ArrowArrayConcatenatorTests.cs +++ b/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(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().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(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(() => 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, IArrowArray>> GenerateTestData(bool slicedArrays = false) { var targetTypes = new List() {