diff --git a/arrow/array/binary_test.go b/arrow/array/binary_test.go index 15cf86928..0e27b9bb4 100644 --- a/arrow/array/binary_test.go +++ b/arrow/array/binary_test.go @@ -18,6 +18,7 @@ package array import ( "reflect" + "strconv" "testing" "github.com/apache/arrow-go/v18/arrow" @@ -26,6 +27,12 @@ import ( "github.com/stretchr/testify/assert" ) +type binaryViewNoAllocAllocator struct{} + +func (binaryViewNoAllocAllocator) Allocate(int) []byte { panic("unexpected allocation") } +func (binaryViewNoAllocAllocator) Reallocate(int, []byte) []byte { panic("unexpected allocation") } +func (binaryViewNoAllocAllocator) Free([]byte) {} + func TestBinary(t *testing.T) { mem := memory.NewCheckedAllocator(memory.NewGoAllocator()) defer mem.AssertSize(t, 0) @@ -701,6 +708,30 @@ func TestBinaryStringRoundTrip(t *testing.T) { assert.True(t, Equal(arr, arr1)) } +func TestBinaryViewBuilderRejectsOversizedValues(t *testing.T) { + const expected = "invalid: BinaryView or StringView elements cannot reference strings larger than 2GB" + oversizedLen := int64(viewValueSizeLimit) + 1 + + t.Run("size guard", func(t *testing.T) { + assert.PanicsWithError(t, expected, func() { + checkBinaryViewValueSize(oversizedLen) + }) + }) + + t.Run("ReserveData", func(t *testing.T) { + if strconv.IntSize == 32 { + t.Skip("oversized int requires 64-bit int") + } + + b := NewBinaryViewBuilder(binaryViewNoAllocAllocator{}) + defer b.Release() + + assert.PanicsWithError(t, expected, func() { + b.ReserveData(int(oversizedLen)) + }) + }) +} + func TestBinaryViewStringRoundTrip(t *testing.T) { mem := memory.NewCheckedAllocator(memory.DefaultAllocator) defer mem.AssertSize(t, 0) diff --git a/arrow/array/binarybuilder.go b/arrow/array/binarybuilder.go index 5b881a842..698e7c53d 100644 --- a/arrow/array/binarybuilder.go +++ b/arrow/array/binarybuilder.go @@ -467,11 +467,15 @@ func (b *BinaryViewBuilder) Resize(n int) { b.rawData = arrow.ViewHeaderTraits.CastFromBytes(b.data.Bytes()) } -func (b *BinaryViewBuilder) ReserveData(length int) { - if int32(length) > viewValueSizeLimit { +func checkBinaryViewValueSize(length int64) { + if length > int64(viewValueSizeLimit) { panic(fmt.Errorf("%w: BinaryView or StringView elements cannot reference strings larger than 2GB", arrow.ErrInvalid)) } +} + +func (b *BinaryViewBuilder) ReserveData(length int) { + checkBinaryViewValueSize(int64(length)) b.blockBuilder.Reserve(int(length)) } @@ -480,9 +484,7 @@ func (b *BinaryViewBuilder) Reserve(n int) { } func (b *BinaryViewBuilder) Append(v []byte) { - if int32(len(v)) > viewValueSizeLimit { - panic(fmt.Errorf("%w: BinaryView or StringView elements cannot reference strings larger than 2GB", arrow.ErrInvalid)) - } + checkBinaryViewValueSize(int64(len(v))) if !arrow.IsViewInline(len(v)) { b.ReserveData(len(v))