diff --git a/src/Apache.Arrow/Arrays/ArrayDataConcatenator.cs b/src/Apache.Arrow/Arrays/ArrayDataConcatenator.cs index 831fbfa1..e78edfe6 100644 --- a/src/Apache.Arrow/Arrays/ArrayDataConcatenator.cs +++ b/src/Apache.Arrow/Arrays/ArrayDataConcatenator.cs @@ -72,6 +72,8 @@ private static ArrayData ConcatenateExtension(ExtensionType type, IReadOnlyList< private class ArrayDataConcatenationVisitor : IArrowTypeVisitor, IArrowTypeVisitor, + IArrowTypeVisitor, + IArrowTypeVisitor, IArrowTypeVisitor, IArrowTypeVisitor, IArrowTypeVisitor, @@ -125,6 +127,223 @@ public void Visit(FixedWidthType type) Result = new ArrayData(resolvedType, _totalLength, _totalNullCount, 0, new ArrowBuffer[] { validityBuffer, valueBuffer }); } + public void Visit(NullType type) + { + foreach (ArrayData arrayData in _arrayDataList) + { + arrayData.EnsureDataType(type.TypeId); + } + + // A null array has no buffers; every slot is null. + Result = new ArrayData(type, _totalLength, _totalLength, 0, System.Array.Empty()); + } + + public void Visit(DictionaryType type) + { + CheckData(type, 2); + var indexType = (IntegerType)type.IndexType; + + // Inputs with no rows don't contribute any dictionary entries. + var contributing = new List(_arrayDataList.Count); + bool allOrdered = true; + foreach (ArrayData arrayData in _arrayDataList) + { + var otherType = (DictionaryType)arrayData.DataType; + allOrdered &= otherType.Ordered; + if (otherType.IndexType.TypeId != indexType.TypeId) + { + throw new ArgumentException( + $"Cannot concatenate dictionary arrays with different index types: {indexType.Name} vs {otherType.IndexType.Name}"); + } + if (otherType.ValueType.TypeId != type.ValueType.TypeId) + { + throw new ArgumentException( + $"Cannot concatenate dictionary arrays with different value types: {type.ValueType.Name} vs {otherType.ValueType.Name}"); + } + if (arrayData.Dictionary == null) + { + throw new ArgumentException("Cannot concatenate a dictionary array that has no dictionary."); + } + + if (arrayData.Length > 0) + { + contributing.Add(arrayData); + } + } + + ArrayData firstDictionary = (contributing.Count > 0 ? contributing[0] : _arrayDataList[0]).Dictionary; + bool sharedDictionary = true; + foreach (ArrayData arrayData in contributing) + { + if (!SharesContents(firstDictionary, arrayData.Dictionary)) + { + sharedDictionary = false; + break; + } + } + + if (sharedDictionary) + { + // The result reuses a dictionary every input describes, so it is ordered + // only if every input says that dictionary is ordered. + DictionaryType sharedType = type.Ordered == allOrdered + ? type + : new DictionaryType(type.IndexType, type.ValueType, allOrdered); + ArrowBuffer sharedValidityBuffer = ConcatenateValidityBuffer(); + ArrowBuffer indexBuffer = ConcatenateFixedWidthTypeValueBuffer(1, type); + Result = new ArrayData(sharedType, _totalLength, _totalNullCount, 0, new ArrowBuffer[] { sharedValidityBuffer, indexBuffer }, null, firstDictionary.Retain()); + return; + } + + // Append the dictionaries end to end and shift each input's indices past the + // entries of the dictionaries before it. Nothing shows that the appended entries + // are in order, so the result is never ordered. + var dictionaries = new List(contributing.Count); + long dictionaryLength = 0; + foreach (ArrayData arrayData in contributing) + { + dictionaries.Add(arrayData.Dictionary); + dictionaryLength += arrayData.Dictionary.Length; + } + + if (dictionaryLength - 1 > MaxIndexValue(indexType)) + { + throw new OverflowException( + $"Cannot concatenate dictionary arrays: the combined dictionary has {dictionaryLength} entries, which is more than index type {indexType.Name} can address."); + } + + // Concatenate the dictionaries first: it can throw, and nothing else has been allocated yet. + ArrayData combinedDictionary = Concatenate(dictionaries, _allocator); + ArrowBuffer validityBuffer = ArrowBuffer.Empty; + ArrowBuffer shiftedIndexBuffer; + try + { + validityBuffer = ConcatenateValidityBuffer(); + shiftedIndexBuffer = ConcatenateShiftedIndexBuffer(indexType); + } + catch + { + validityBuffer.Dispose(); + combinedDictionary.Dispose(); + throw; + } + + var combinedType = type.Ordered + ? new DictionaryType(type.IndexType, type.ValueType, false) + : type; + Result = new ArrayData(combinedType, _totalLength, _totalNullCount, 0, new ArrowBuffer[] { validityBuffer, shiftedIndexBuffer }, null, combinedDictionary); + } + + private static bool SharesContents(ArrayData left, ArrayData right) + { + if (ReferenceEquals(left, right)) + { + return true; + } + + // Distinct ArrayData objects can still view the same memory, e.g. after Retain or SliceShared. + if (left == null || right == null || + left.DataType.TypeId != right.DataType.TypeId || + left.Length != right.Length || + left.Offset != right.Offset || + left.Buffers.Length != right.Buffers.Length || + (left.Children?.Length ?? 0) != (right.Children?.Length ?? 0)) + { + return false; + } + + for (int i = 0; i < left.Buffers.Length; i++) + { + if (left.Buffers[i].Span != right.Buffers[i].Span) + { + return false; + } + } + + for (int i = 0; i < (left.Children?.Length ?? 0); i++) + { + if (!SharesContents(left.Children[i], right.Children[i])) + { + return false; + } + } + + return left.Dictionary == null && right.Dictionary == null || + SharesContents(left.Dictionary, right.Dictionary); + } + + private static long MaxIndexValue(IntegerType indexType) => indexType.TypeId switch + { + ArrowTypeId.Int8 => sbyte.MaxValue, + ArrowTypeId.UInt8 => byte.MaxValue, + ArrowTypeId.Int16 => short.MaxValue, + ArrowTypeId.UInt16 => ushort.MaxValue, + ArrowTypeId.Int32 => int.MaxValue, + ArrowTypeId.UInt32 => uint.MaxValue, + _ => long.MaxValue, + }; + + private ArrowBuffer ConcatenateShiftedIndexBuffer(IntegerType indexType) + { + int byteWidth = indexType.BitWidth / 8; + var builder = new ArrowBuffer.Builder(_totalLength * byteWidth); + builder.Resize(_totalLength * byteWidth); + Span destination = builder.Span; + int position = 0; + long shift = 0; + + foreach (ArrayData arrayData in _arrayDataList) + { + if (arrayData.Length == 0) + { + continue; + } + + ReadOnlySpan validity = arrayData.Buffers[0].Span; + bool allNull = validity.Length == 0 && arrayData.GetNullCount() == arrayData.Length; + ReadOnlySpan source = arrayData.Buffers[1].Span.Slice(arrayData.Offset * byteWidth, arrayData.Length * byteWidth); + + for (int i = 0; i < arrayData.Length; i++) + { + bool isValid = validity.Length == 0 ? !allNull : BitUtility.GetBit(validity, arrayData.Offset + i); + // A null slot's index is unspecified and may not be in range, so write 0 instead of shifting it. + long index = isValid ? ReadIndex(source, indexType.TypeId, i) + shift : 0; + WriteIndex(destination, indexType.TypeId, position++, index); + } + + shift += arrayData.Dictionary.Length; + } + + return builder.Build(_allocator); + } + + private static long ReadIndex(ReadOnlySpan source, ArrowTypeId typeId, int i) => typeId switch + { + ArrowTypeId.Int8 => source.CastTo()[i], + ArrowTypeId.UInt8 => source[i], + ArrowTypeId.Int16 => source.CastTo()[i], + ArrowTypeId.UInt16 => source.CastTo()[i], + ArrowTypeId.Int32 => source.CastTo()[i], + ArrowTypeId.UInt32 => source.CastTo()[i], + ArrowTypeId.Int64 => source.CastTo()[i], + _ => checked((long)source.CastTo()[i]), + }; + + private static void WriteIndex(Span destination, ArrowTypeId typeId, int i, long value) + { + switch (typeId) + { + case ArrowTypeId.Int8: destination.CastTo()[i] = (sbyte)value; break; + case ArrowTypeId.UInt8: destination[i] = (byte)value; break; + case ArrowTypeId.Int16: destination.CastTo()[i] = (short)value; break; + case ArrowTypeId.UInt16: destination.CastTo()[i] = (ushort)value; break; + case ArrowTypeId.Int32: destination.CastTo()[i] = (int)value; break; + case ArrowTypeId.UInt32: destination.CastTo()[i] = (uint)value; break; + case ArrowTypeId.Int64: destination.CastTo()[i] = value; break; + default: destination.CastTo()[i] = (ulong)value; break; + } + } + public void Visit(BinaryType type) => ConcatenateVariableBinaryArrayData(type); public void Visit(BinaryViewType type) => ConcatenateBinaryViewArrayData(type); diff --git a/test/Apache.Arrow.Tests/ArrowArrayConcatenatorTests.cs b/test/Apache.Arrow.Tests/ArrowArrayConcatenatorTests.cs index 4e781a78..f2b8094b 100644 --- a/test/Apache.Arrow.Tests/ArrowArrayConcatenatorTests.cs +++ b/test/Apache.Arrow.Tests/ArrowArrayConcatenatorTests.cs @@ -1328,6 +1328,202 @@ public void TestFixedSizeListSizeMismatchThrows() Assert.Contains("list size", ex.Message, StringComparison.OrdinalIgnoreCase); } + [Fact] + public void TestNullArrays() + { + var result = ArrowArrayConcatenator.Concatenate(new IArrowArray[] { new NullArray(3), new NullArray(0), new NullArray(2) }); + + var nullArray = Assert.IsType(result); + Assert.Equal(5, nullArray.Length); + Assert.Equal(5, nullArray.NullCount); + } + + [Fact] + public void TestStructWithNullField() + { + var type = new StructType(new[] { new Field("n", NullType.Default, true) }); + StructArray Build(int length) => new StructArray(type, length, new IArrowArray[] { new NullArray(length) }, ArrowBuffer.Empty, 0); + + var result = (StructArray)ArrowArrayConcatenator.Concatenate(new IArrowArray[] { Build(3), Build(2) }); + + Assert.Equal(5, result.Length); + var field = Assert.IsType(result.Fields[0]); + Assert.Equal(5, field.Length); + } + + [Fact] + public void TestDictionaryArraysWithSharedDictionary() + { + var dictionary = new StringArray.Builder().AppendRange(new[] { "a", "b" }).Build(); + var first = BuildDictionaryArray(dictionary, 0, 1); + var second = BuildDictionaryArray(dictionary, 1, 0); + + var result = (DictionaryArray)ArrowArrayConcatenator.Concatenate(new IArrowArray[] { first, second }); + first.Dispose(); + second.Dispose(); + + Assert.Equal(2, result.Dictionary.Length); + Assert.Equal(new[] { "a", "b", "b", "a" }, DictionaryStrings(result)); + result.Dispose(); + } + + [Fact] + public void TestDictionaryArraysWithDifferentDictionaries() + { + var first = BuildDictionaryArray(new StringArray.Builder().AppendRange(new[] { "a", "b" }).Build(), 0, 1, null); + var second = BuildDictionaryArray(new StringArray.Builder().AppendRange(new[] { "x", "y", "z" }).Build(), 2, null, 0); + var third = BuildDictionaryArray(new StringArray.Builder().AppendRange(new[] { "q" }).Build(), 0); + + var result = (DictionaryArray)ArrowArrayConcatenator.Concatenate(new IArrowArray[] { first, second, third }); + first.Dispose(); + second.Dispose(); + third.Dispose(); + + Assert.Equal(6, result.Dictionary.Length); + Assert.Equal(2, result.NullCount); + Assert.Equal(new[] { "a", "b", null, "z", null, "x", "q" }, DictionaryStrings(result)); + result.Dispose(); + } + + [Fact] + public void TestSlicedDictionaryArrays() + { + var first = BuildDictionaryArray(new StringArray.Builder().AppendRange(new[] { "a", "b", "c" }).Build(), 0, 1, 2, 0); + var second = BuildDictionaryArray(new StringArray.Builder().AppendRange(new[] { "x", "y" }).Build(), 1, 0, 1); + + var result = (DictionaryArray)ArrowArrayConcatenator.Concatenate(new IArrowArray[] + { + ArrowArrayFactory.Slice(first, 1, 2), + ArrowArrayFactory.Slice(second, 1, 2), + }); + + Assert.Equal(new[] { "b", "c", "x", "y" }, DictionaryStrings(result)); + } + + [Fact] + public void TestDictionaryArraysSkipEmptyInputDictionaries() + { + var dictionary = new StringArray.Builder().AppendRange(new[] { "a", "b" }).Build(); + var empty = BuildDictionaryArray(new StringArray.Builder().AppendRange(new[] { "x" }).Build()); + + var result = (DictionaryArray)ArrowArrayConcatenator.Concatenate(new IArrowArray[] + { + BuildDictionaryArray(dictionary, 1), empty, BuildDictionaryArray(dictionary, 0), + }); + + Assert.Equal(2, result.Dictionary.Length); + Assert.Equal(new[] { "b", "a" }, DictionaryStrings(result)); + } + + [Fact] + public void TestDictionaryArraysIndexOverflowThrows() + { + var type = new DictionaryType(Int8Type.Default, Int32Type.Default, false); + DictionaryArray Build(int dictionaryLength) => new DictionaryArray( + type, + new Int8Array.Builder().Append(0).Build(), + new Int32Array.Builder().AppendRange(Enumerable.Range(0, dictionaryLength)).Build()); + + Assert.Throws(() => + ArrowArrayConcatenator.Concatenate(new IArrowArray[] { Build(100), Build(29) })); + + var result = (DictionaryArray)ArrowArrayConcatenator.Concatenate(new IArrowArray[] { Build(100), Build(28) }); + Assert.Equal(128, result.Dictionary.Length); + Assert.Equal((sbyte)100, ((Int8Array)result.Indices).GetValue(1)); + } + + [Fact] + public void TestDictionaryArraysSharingMemoryKeepDictionary() + { + var dictionary = new StringArray.Builder().AppendRange(new[] { "a", "b" }).Build(); + var first = BuildDictionaryArray(dictionary, 0, 1); + // A distinct ArrayData over the same memory, as after Retain. + var second = BuildDictionaryArray(new StringArray(dictionary.Data.Retain()), 1, 0); + Assert.NotSame(first.Data.Dictionary, second.Data.Dictionary); + + var result = (DictionaryArray)ArrowArrayConcatenator.Concatenate(new IArrowArray[] { first, second }); + first.Dispose(); + second.Dispose(); + + Assert.Equal(2, result.Dictionary.Length); + Assert.Equal(new[] { "a", "b", "b", "a" }, DictionaryStrings(result)); + result.Dispose(); + } + + [Fact] + public void TestCombinedDictionariesAreNotOrdered() + { + var first = BuildDictionaryArray(new StringArray.Builder().AppendRange(new[] { "a", "b" }).Build(), ordered: true, 0, 1); + var second = BuildDictionaryArray(new StringArray.Builder().AppendRange(new[] { "c", "d" }).Build(), ordered: true, 1, 0); + + var result = (DictionaryArray)ArrowArrayConcatenator.Concatenate(new IArrowArray[] { first, second }); + + Assert.False(((DictionaryType)result.Data.DataType).Ordered); + Assert.Equal(new[] { "a", "b", "d", "c" }, DictionaryStrings(result)); + } + + [Theory] + [InlineData(true, true, true)] + [InlineData(true, false, false)] + [InlineData(false, true, false)] + [InlineData(false, false, false)] + public void TestSharedDictionaryOrderedOnlyIfAllInputsOrdered(bool firstOrdered, bool secondOrdered, bool expected) + { + var dictionary = new StringArray.Builder().AppendRange(new[] { "a", "b" }).Build(); + var first = BuildDictionaryArray(dictionary, firstOrdered, 0, 1); + var second = BuildDictionaryArray(dictionary, secondOrdered, 1, 0); + + var result = (DictionaryArray)ArrowArrayConcatenator.Concatenate(new IArrowArray[] { first, second }); + + Assert.Equal(2, result.Dictionary.Length); + Assert.Equal(expected, ((DictionaryType)result.Data.DataType).Ordered); + } + + [Fact] + public void TestArrayDataConcatenatorKeepsDictionary() + { + var dictionary = new StringArray.Builder().AppendRange(new[] { "a", "b" }).Build(); + + var data = ArrayDataConcatenator.Concatenate(new[] + { + BuildDictionaryArray(dictionary, 0, 1).Data, + BuildDictionaryArray(dictionary, 1, 0).Data, + }); + + Assert.Equal(4, data.Length); + Assert.NotNull(data.Dictionary); + } + + private static DictionaryArray BuildDictionaryArray(StringArray dictionary, params int?[] indices) => + BuildDictionaryArray(dictionary, false, indices); + + private static DictionaryArray BuildDictionaryArray(StringArray dictionary, bool ordered, params int?[] indices) + { + var builder = new Int32Array.Builder(); + foreach (int? index in indices) + { + builder.Append(index); + } + + return new DictionaryArray( + new DictionaryType(Int32Type.Default, StringType.Default, ordered), + builder.Build(), + dictionary); + } + + private static string[] DictionaryStrings(DictionaryArray array) + { + var indices = (Int32Array)array.Indices; + var dictionary = (StringArray)array.Dictionary; + var values = new string[array.Length]; + for (int i = 0; i < array.Length; i++) + { + int? index = indices.GetValue(i); + values[i] = index.HasValue ? dictionary.GetString(index.Value) : null; + } + return values; + } + private static FixedSizeListArray BuildFixedSizeListArray(int listSize, int[] flatValues) { var type = new FixedSizeListType(Int32Type.Default, listSize);