diff --git a/src/Apache.Arrow/Memory/NativeBuffer.cs b/src/Apache.Arrow/Memory/NativeBuffer.cs index f5b8f619..b042999b 100644 --- a/src/Apache.Arrow/Memory/NativeBuffer.cs +++ b/src/Apache.Arrow/Memory/NativeBuffer.cs @@ -15,6 +15,7 @@ using System; using System.Buffers; +using System.Diagnostics; using System.Runtime.CompilerServices; using System.Runtime.InteropServices; using System.Threading; @@ -90,15 +91,11 @@ public void Grow(int newElementCount, bool zeroFill = true) if (newElementCount <= Length) return; - // Exponential growth (2x) to amortise repeated grows - // TODO: There might be a size that's big enough to work for this case but not too big to overflow. - // We could use that instead of blindly doubling. - int newCount = Math.Max(newElementCount, checked(Length * 2)); int elementSize = Unsafe.SizeOf(); + int newCount = ComputeGrowCount(Length, newElementCount, elementSize); int needed = checked(newCount * elementSize); - var owner = _owner ?? throw new ObjectDisposedException(nameof(NativeBuffer)); - owner.Reallocate(needed); + _owner.Reallocate(needed); if (zeroFill) { @@ -109,6 +106,21 @@ public void Grow(int newElementCount, bool zeroFill = true) Length = newCount; } + /// + /// The element count to grow to: double the current length to amortise repeated grows, but never + /// past the largest buffer that can be addressed, and never below what the caller asked for. + /// + internal static int ComputeGrowCount(int length, int newElementCount, int elementSize) + { + // Always Unsafe.SizeOf() for an unmanaged TItem, so never below one; the parameter + // exists so the boundary can be tested without allocating a buffer of that size. + Debug.Assert(elementSize > 0); + + int maxCount = int.MaxValue / elementSize; + long doubled = (long)length * 2; + return (int)Math.Max(newElementCount, Math.Min(doubled, maxCount)); + } + public void Dispose() { IDisposable disposable = _owner; diff --git a/test/Apache.Arrow.Tests/NativeBufferTests.cs b/test/Apache.Arrow.Tests/NativeBufferTests.cs index 84d050cf..87ea4405 100644 --- a/test/Apache.Arrow.Tests/NativeBufferTests.cs +++ b/test/Apache.Arrow.Tests/NativeBufferTests.cs @@ -84,6 +84,69 @@ public void GrowWithSmallerOrEqualCountIsNoOp() Assert.Equal(42, buf.Span[0]); } + // Growth doubles to stay amortised, but must saturate rather than overflow. Doubling used to be + // unconditional and checked, so a buffer past half the maximum threw OverflowException on its + // next grow however little was asked for — a byte buffer could not grow beyond about 1 GiB. + // + // The arithmetic is tested directly: reproducing it through Grow would mean allocating more than + // a gigabyte, which is not something to put in a unit test. + [Theory] + // length, requested, elementSize, expected + [InlineData(0, 1, 1, 1)] // nothing to double yet + [InlineData(3, 10, 4, 10)] // request exceeds the doubling + [InlineData(8, 10, 4, 16)] // doubling exceeds the request + [InlineData(5, 5, 4, 10)] // equal: doubling still wins + public void ComputeGrowCountDoublesWhileItFits( + int length, int requested, int elementSize, int expected) + { + Assert.Equal( + expected, + NativeBuffer.ComputeGrowCount(length, requested, elementSize)); + } + + [Fact] + public void ComputeGrowCountSaturatesInsteadOfOverflowing() + { + // Past half the maximum, doubling would overflow. The result saturates at the largest + // addressable count and still covers the request. + const int elementSize = 1; + int overHalf = (int.MaxValue / 2) + 1000; + + int grown = NativeBuffer.ComputeGrowCount( + overHalf, overHalf + 1, elementSize); + + Assert.Equal(int.MaxValue, grown); + Assert.True(grown >= overHalf + 1); + } + + [Fact] + public void ComputeGrowCountSaturatesPerElementSize() + { + // The ceiling is a byte count, so a wider element saturates at proportionally fewer of them. + const int elementSize = 8; + int maxCount = int.MaxValue / elementSize; + int overHalf = (maxCount / 2) + 1000; + + int grown = NativeBuffer.ComputeGrowCount( + overHalf, overHalf + 1, elementSize); + + Assert.Equal(maxCount, grown); + Assert.True((long)grown * elementSize <= int.MaxValue); + } + + [Fact] + public void ComputeGrowCountNeverReturnsLessThanRequested() + { + // A request larger than the ceiling is not silently truncated; Grow still refuses it when it + // works out the byte size. + const int elementSize = 8; + int beyond = (int.MaxValue / elementSize) + 1; + + Assert.Equal( + beyond, + NativeBuffer.ComputeGrowCount(0, beyond, elementSize)); + } + [Fact] public void BuildTransfersOwnershipToArrowBuffer() {