diff --git a/arrow/array/builder.go b/arrow/array/builder.go index ed4f08e72..f90bb87c8 100644 --- a/arrow/array/builder.go +++ b/arrow/array/builder.go @@ -70,6 +70,12 @@ type Builder interface { // AppendValueFromString adds a new value from a string. Inverse of array.ValueStr(i int) string AppendValueFromString(string) error + // Truncate reduces the length of the array being built to n elements, + // discarding any elements beyond the first n. It does not reallocate or + // shrink the underlying buffers. If n is greater than or equal to the + // current length, Truncate is a no-op. + Truncate(n int) + // Reserve ensures there is enough space for appending n elements // by checking the capacity and calling Resize if necessary. Reserve(n int) @@ -248,6 +254,22 @@ func (b *builder) unsafeSetValid(length int) { b.length = newLength } +// Truncate reduces the length of the array being built to n elements and +// recomputes the null count from the validity bitmap. The underlying buffers +// are left untouched. +func (b *builder) Truncate(n int) { + if n < 0 { + panic("arrow/array: cannot truncate to a negative length") + } + if n >= b.length { + return + } + b.length = n + if b.nullBitmap != nil { + b.nulls = n - bitutil.CountSetBits(b.nullBitmap.Bytes(), 0, n) + } +} + func (b *builder) UnsafeAppendBoolToBitmap(isValid bool) { if isValid { bitutil.SetBit(b.nullBitmap.Bytes(), b.length) diff --git a/arrow/array/builder_test.go b/arrow/array/builder_test.go index 045317dc8..cbeb24693 100644 --- a/arrow/array/builder_test.go +++ b/arrow/array/builder_test.go @@ -17,8 +17,10 @@ package array import ( + "strconv" "testing" + "github.com/apache/arrow-go/v18/arrow" "github.com/apache/arrow-go/v18/arrow/internal/testing/tools" "github.com/apache/arrow-go/v18/arrow/memory" "github.com/stretchr/testify/assert" @@ -126,3 +128,408 @@ func TestBuilder_SetNull(t *testing.T) { } } } + +func countNulls(arr arrow.Array) int { + n := 0 + for i := 0; i < arr.Len(); i++ { + if arr.IsNull(i) { + n++ + } + } + return n +} + +func TestBuilderTruncate(t *testing.T) { + scalarValid := []bool{true, false, true, true, false, true, true, false} + + cases := []struct { + name string + // build returns a builder that already has `full` elements appended. + build func(mem memory.Allocator) Builder + full int + }{ + { + name: "boolean", + build: func(mem memory.Allocator) Builder { + b := NewBooleanBuilder(mem) + b.AppendValues([]bool{true, false, true, false, true, false, true, false}, scalarValid) + return b + }, + full: 8, + }, + { + name: "int8", + build: func(mem memory.Allocator) Builder { + b := NewInt8Builder(mem) + b.AppendValues([]int8{1, 2, 3, 4, 5, 6, 7, 8}, scalarValid) + return b + }, + full: 8, + }, + { + name: "int16", + build: func(mem memory.Allocator) Builder { + b := NewInt16Builder(mem) + b.AppendValues([]int16{1, 2, 3, 4, 5, 6, 7, 8}, scalarValid) + return b + }, + full: 8, + }, + { + name: "int32", + build: func(mem memory.Allocator) Builder { + b := NewInt32Builder(mem) + b.AppendValues([]int32{1, 2, 3, 4, 5, 6, 7, 8}, scalarValid) + return b + }, + full: 8, + }, + { + name: "int64", + build: func(mem memory.Allocator) Builder { + b := NewInt64Builder(mem) + b.AppendValues([]int64{1, 2, 3, 4, 5, 6, 7, 8}, scalarValid) + return b + }, + full: 8, + }, + { + name: "uint8", + build: func(mem memory.Allocator) Builder { + b := NewUint8Builder(mem) + b.AppendValues([]uint8{1, 2, 3, 4, 5, 6, 7, 8}, scalarValid) + return b + }, + full: 8, + }, + { + name: "uint16", + build: func(mem memory.Allocator) Builder { + b := NewUint16Builder(mem) + b.AppendValues([]uint16{1, 2, 3, 4, 5, 6, 7, 8}, scalarValid) + return b + }, + full: 8, + }, + { + name: "uint32", + build: func(mem memory.Allocator) Builder { + b := NewUint32Builder(mem) + b.AppendValues([]uint32{1, 2, 3, 4, 5, 6, 7, 8}, scalarValid) + return b + }, + full: 8, + }, + { + name: "uint64", + build: func(mem memory.Allocator) Builder { + b := NewUint64Builder(mem) + b.AppendValues([]uint64{1, 2, 3, 4, 5, 6, 7, 8}, scalarValid) + return b + }, + full: 8, + }, + { + name: "float32", + build: func(mem memory.Allocator) Builder { + b := NewFloat32Builder(mem) + b.AppendValues([]float32{1, 2, 3, 4, 5, 6, 7, 8}, scalarValid) + return b + }, + full: 8, + }, + { + name: "float64", + build: func(mem memory.Allocator) Builder { + b := NewFloat64Builder(mem) + b.AppendValues([]float64{1, 2, 3, 4, 5, 6, 7, 8}, scalarValid) + return b + }, + full: 8, + }, + { + name: "string", + build: func(mem memory.Allocator) Builder { + b := NewStringBuilder(mem) + b.AppendValues([]string{"a", "bb", "ccc", "d", "ee", "fff", "g", "hh"}, scalarValid) + return b + }, + full: 8, + }, + { + name: "binary", + build: func(mem memory.Allocator) Builder { + b := NewBinaryBuilder(mem, arrow.BinaryTypes.Binary) + b.AppendValues([][]byte{[]byte("a"), []byte("bb"), []byte("ccc"), []byte("d"), []byte("ee"), []byte("fff"), []byte("g"), []byte("hh")}, scalarValid) + return b + }, + full: 8, + }, + { + name: "fixed_size_binary", + build: func(mem memory.Allocator) Builder { + b := NewFixedSizeBinaryBuilder(mem, &arrow.FixedSizeBinaryType{ByteWidth: 3}) + for i, v := range scalarValid { + if v { + b.Append([]byte{byte(i), byte(i + 1), byte(i + 2)}) + } else { + b.AppendNull() + } + } + return b + }, + full: 8, + }, + { + name: "decimal128", + build: func(mem memory.Allocator) Builder { + b := NewDecimal128Builder(mem, &arrow.Decimal128Type{Precision: 10, Scale: 1}) + for i, v := range scalarValid { + if v { + assert.NoError(t, b.AppendValueFromString(strconv.Itoa(i)+".0")) + } else { + b.AppendNull() + } + } + return b + }, + full: 8, + }, + { + name: "decimal256", + build: func(mem memory.Allocator) Builder { + b := NewDecimal256Builder(mem, &arrow.Decimal256Type{Precision: 10, Scale: 1}) + for i, v := range scalarValid { + if v { + assert.NoError(t, b.AppendValueFromString(strconv.Itoa(i)+".0")) + } else { + b.AppendNull() + } + } + return b + }, + full: 8, + }, + { + name: "date32", + build: func(mem memory.Allocator) Builder { + b := NewDate32Builder(mem) + b.AppendValues([]arrow.Date32{1, 2, 3, 4, 5, 6, 7, 8}, scalarValid) + return b + }, + full: 8, + }, + { + name: "timestamp", + build: func(mem memory.Allocator) Builder { + b := NewTimestampBuilder(mem, &arrow.TimestampType{Unit: arrow.Second}) + b.AppendValues([]arrow.Timestamp{1, 2, 3, 4, 5, 6, 7, 8}, scalarValid) + return b + }, + full: 8, + }, + { + name: "list", + build: func(mem memory.Allocator) Builder { + b := NewListBuilder(mem, arrow.PrimitiveTypes.Int32) + vb := b.ValueBuilder().(*Int32Builder) + lengths := []int{3, 0, 2, 1, 4, 2} + valid := []bool{true, false, true, true, true, false} + v := int32(0) + for i, l := range lengths { + b.Append(valid[i]) + for j := 0; j < l; j++ { + vb.Append(v) + v++ + } + } + return b + }, + full: 6, + }, + { + name: "large_list", + build: func(mem memory.Allocator) Builder { + b := NewLargeListBuilder(mem, arrow.PrimitiveTypes.Int32) + vb := b.ValueBuilder().(*Int32Builder) + lengths := []int{3, 0, 2, 1, 4, 2} + valid := []bool{true, false, true, true, true, false} + v := int32(0) + for i, l := range lengths { + b.Append(valid[i]) + for j := 0; j < l; j++ { + vb.Append(v) + v++ + } + } + return b + }, + full: 6, + }, + { + name: "fixed_size_list", + build: func(mem memory.Allocator) Builder { + b := NewFixedSizeListBuilder(mem, 2, arrow.PrimitiveTypes.Int32) + vb := b.ValueBuilder().(*Int32Builder) + valid := []bool{true, false, true, true, false, true} + v := int32(0) + for _, ok := range valid { + b.Append(ok) + vb.Append(v) + vb.Append(v + 1) + v += 2 + } + return b + }, + full: 6, + }, + { + name: "map", + build: func(mem memory.Allocator) Builder { + b := NewMapBuilder(mem, arrow.PrimitiveTypes.Int32, arrow.PrimitiveTypes.Int32, false) + kb := b.KeyBuilder().(*Int32Builder) + ib := b.ItemBuilder().(*Int32Builder) + lengths := []int{2, 0, 1, 3, 2, 1} + valid := []bool{true, false, true, true, true, false} + k := int32(0) + for i, l := range lengths { + b.Append(valid[i]) + for j := 0; j < l; j++ { + kb.Append(k) + ib.Append(k * 10) + k++ + } + } + return b + }, + full: 6, + }, + { + name: "struct", + build: func(mem memory.Allocator) Builder { + dt := arrow.StructOf( + arrow.Field{Name: "a", Type: arrow.PrimitiveTypes.Int32, Nullable: true}, + arrow.Field{Name: "b", Type: arrow.BinaryTypes.String, Nullable: true}, + ) + b := NewStructBuilder(mem, dt) + ab := b.FieldBuilder(0).(*Int32Builder) + bb := b.FieldBuilder(1).(*StringBuilder) + valid := []bool{true, false, true, true, false, true, true, false} + for i, ok := range valid { + b.Append(ok) + ab.Append(int32(i)) + bb.Append(strconv.Itoa(i)) + } + return b + }, + full: 8, + }, + { + name: "dictionary", + build: func(mem memory.Allocator) Builder { + dt := &arrow.DictionaryType{IndexType: arrow.PrimitiveTypes.Int8, ValueType: arrow.BinaryTypes.String} + b := NewDictionaryBuilder(mem, dt).(*BinaryDictionaryBuilder) + values := []string{"x", "y", "x", "z", "y", "x", "w", "z"} + for i, ok := range scalarValid { + if ok { + assert.NoError(t, b.AppendString(values[i])) + } else { + b.AppendNull() + } + } + return b + }, + full: 8, + }, + { + name: "sparse_union", + build: func(mem memory.Allocator) Builder { + i8 := NewInt8Builder(mem) + str := NewStringBuilder(mem) + b := NewSparseUnionBuilderWithBuilders(mem, + arrow.SparseUnionOf([]arrow.Field{ + {Name: "i8", Type: arrow.PrimitiveTypes.Int8, Nullable: true}, + {Name: "str", Type: arrow.BinaryTypes.String, Nullable: true}, + }, []arrow.UnionTypeCode{0, 1}), + []Builder{i8, str}) + types := []arrow.UnionTypeCode{0, 1, 0, 1, 0, 1} + for i, tc := range types { + b.Append(tc) + if tc == 0 { + i8.Append(int8(i)) + str.AppendEmptyValue() + } else { + str.Append(strconv.Itoa(i)) + i8.AppendEmptyValue() + } + } + i8.Release() + str.Release() + return b + }, + full: 6, + }, + { + name: "dense_union", + build: func(mem memory.Allocator) Builder { + i8 := NewInt8Builder(mem) + str := NewStringBuilder(mem) + b := NewDenseUnionBuilderWithBuilders(mem, + arrow.DenseUnionOf([]arrow.Field{ + {Name: "i8", Type: arrow.PrimitiveTypes.Int8, Nullable: true}, + {Name: "str", Type: arrow.BinaryTypes.String, Nullable: true}, + }, []arrow.UnionTypeCode{0, 1}), + []Builder{i8, str}) + types := []arrow.UnionTypeCode{0, 1, 0, 1, 0, 1} + for i, tc := range types { + b.Append(tc) + if tc == 0 { + i8.Append(int8(i)) + } else { + str.Append(strconv.Itoa(i)) + } + } + i8.Release() + str.Release() + return b + }, + full: 6, + }, + { + name: "run_end_encoded", + build: func(mem memory.Allocator) Builder { + b := NewRunEndEncodedBuilder(mem, arrow.PrimitiveTypes.Int32, arrow.BinaryTypes.String) + vb := b.ValueBuilder().(*StringBuilder) + runs := []uint64{2, 3, 1, 4, 2, 3} + for i, r := range runs { + b.Append(r) + vb.Append(strconv.Itoa(i)) + } + return b + }, + full: 15, + }, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + for _, n := range []int{tc.full - 1, tc.full / 2, 0} { + t.Run(strconv.Itoa(n), func(t *testing.T) { + mem := memory.NewCheckedAllocator(memory.NewGoAllocator()) + defer mem.AssertSize(t, 0) + + b := tc.build(mem) + defer b.Release() + + b.Truncate(n) + + arr := b.NewArray() + defer arr.Release() + + assert.Equal(t, n, arr.Len(), "length after truncate") + assert.Equal(t, countNulls(arr), arr.NullN(), "null count consistency") + }) + } + }) + } +} diff --git a/arrow/array/dictionary.go b/arrow/array/dictionary.go index 38a43f77c..80e8c8aa9 100644 --- a/arrow/array/dictionary.go +++ b/arrow/array/dictionary.go @@ -688,6 +688,21 @@ func (b *dictionaryBuilder) Resize(n int) { b.length = b.idxBuilder.Len() } +// Truncate reduces the length of the dictionary builder to n elements by +// truncating the underlying index builder. The dictionary (memo table) is left +// untouched. No buffers are reallocated. +func (b *dictionaryBuilder) Truncate(n int) { + if n < 0 { + panic("arrow/array: cannot truncate to a negative length") + } + if n >= b.length { + return + } + b.idxBuilder.Truncate(n) + b.length = n + b.nulls = b.idxBuilder.NullN() +} + func (b *dictionaryBuilder) ResetFull() { b.reset() b.idxBuilder.NewArray().Release() diff --git a/arrow/array/encoded.go b/arrow/array/encoded.go index 8d628ffc2..e4c289217 100644 --- a/arrow/array/encoded.go +++ b/arrow/array/encoded.go @@ -370,6 +370,50 @@ func (b *RunEndEncodedBuilder) finishRun() { func (b *RunEndEncodedBuilder) ValueBuilder() Builder { return b.values } +func (b *RunEndEncodedBuilder) runEndValue(i int) int { + switch bldr := b.runEnds.(type) { + case *Int16Builder: + return int(bldr.Value(i)) + case *Int32Builder: + return int(bldr.Value(i)) + case *Int64Builder: + return int(bldr.Value(i)) + } + return 0 +} + +// Truncate reduces the logical length of the run-end encoded builder to n +// elements. It keeps only the runs needed to cover the first n elements and +// leaves the final run pending so newData clamps its end to n. No buffers are +// reallocated. +func (b *RunEndEncodedBuilder) Truncate(n int) { + if n < 0 { + panic("arrow/array: cannot truncate to a negative length") + } + if n >= b.length { + return + } + // materialize the in-progress run so every run is represented in runEnds + b.finishRun() + keep, prev := 0, 0 + for i := 0; i < b.runEnds.Len(); i++ { + if prev >= n { + break + } + prev = b.runEndValue(i) + keep = i + 1 + } + // Keep `keep` values, but drop the last run end so it is treated as an + // in-progress run; newData's finishRun will append the clamped end (n). + b.values.Truncate(keep) + if keep > 0 { + b.runEnds.Truncate(keep - 1) + } else { + b.runEnds.Truncate(0) + } + b.length = n +} + func (b *RunEndEncodedBuilder) Append(n uint64) { b.finishRun() b.addLength(n) diff --git a/arrow/array/fixed_size_list.go b/arrow/array/fixed_size_list.go index ca9210ab4..22e7a5fdd 100644 --- a/arrow/array/fixed_size_list.go +++ b/arrow/array/fixed_size_list.go @@ -273,6 +273,17 @@ func (b *FixedSizeListBuilder) init(capacity int) { b.builder.init(capacity) } +// Truncate reduces the length of the fixed-size-list builder to n elements, +// truncating the values builder to n * listSize elements. No buffers are +// reallocated. +func (b *FixedSizeListBuilder) Truncate(n int) { + if n >= b.length { + return + } + b.builder.Truncate(n) + b.values.Truncate(n * int(b.n)) +} + // Reserve ensures there is enough space for appending n elements // by checking the capacity and calling Resize if necessary. func (b *FixedSizeListBuilder) Reserve(n int) { diff --git a/arrow/array/list.go b/arrow/array/list.go index 0379c7869..adf995619 100644 --- a/arrow/array/list.go +++ b/arrow/array/list.go @@ -523,6 +523,30 @@ func (b *baseListBuilder) resizeHelper(n int) { } } +func (b *baseListBuilder) offsetAt(i int) int { + switch o := b.offsets.(type) { + case *Int32Builder: + return int(o.Value(i)) + case *Int64Builder: + return int(o.Value(i)) + } + return 0 +} + +// Truncate reduces the length of the list builder to n elements. It keeps the +// first n+1 offsets (so the end boundary of the last kept list is preserved) +// and truncates the values builder to that boundary. No buffers are +// reallocated. +func (b *baseListBuilder) Truncate(n int) { + if n >= b.length { + return + } + valueLen := b.offsetAt(n) + b.builder.Truncate(n) + b.offsets.Truncate(n + 1) + b.values.Truncate(valueLen) +} + func (b *baseListBuilder) ValueBuilder() Builder { return b.values } diff --git a/arrow/array/map.go b/arrow/array/map.go index 1d13ea887..ae01f53a5 100644 --- a/arrow/array/map.go +++ b/arrow/array/map.go @@ -259,6 +259,13 @@ func (b *MapBuilder) Reserve(n int) { b.listBuilder.Reserve(n) } // b.Cap(), additional memory will be allocated. If n is smaller, the allocated memory may be reduced. func (b *MapBuilder) Resize(n int) { b.listBuilder.Resize(n) } +// Truncate reduces the length of the map builder to n elements by delegating to +// the underlying list builder. No buffers are reallocated. +func (b *MapBuilder) Truncate(n int) { + b.adjustStructBuilderLen() + b.listBuilder.Truncate(n) +} + // AppendValues is for bulk appending a group of elements with offsets provided // and validity booleans provided. func (b *MapBuilder) AppendValues(offsets []int32, valid []bool) { diff --git a/arrow/array/struct.go b/arrow/array/struct.go index 47e30071d..1a86b455f 100644 --- a/arrow/array/struct.go +++ b/arrow/array/struct.go @@ -420,6 +420,18 @@ func (b *StructBuilder) resizeHelper(n int) { } } +// Truncate reduces the length of the struct builder to n elements, truncating +// each of its field builders to the same length. No buffers are reallocated. +func (b *StructBuilder) Truncate(n int) { + if n >= b.length { + return + } + b.builder.Truncate(n) + for _, f := range b.fields { + f.Truncate(n) + } +} + func (b *StructBuilder) NumField() int { return len(b.fields) } func (b *StructBuilder) FieldBuilder(i int) Builder { return b.fields[i] } diff --git a/arrow/array/union.go b/arrow/array/union.go index 58524a768..723a4710c 100644 --- a/arrow/array/union.go +++ b/arrow/array/union.go @@ -945,6 +945,22 @@ func (b *SparseUnionBuilder) Resize(n int) { b.typesBuilder.resize(n) } +// Truncate reduces the length of the sparse union builder to n elements, +// shrinking the types buffer and truncating every child builder to the same +// length. No buffers are reallocated. +func (b *SparseUnionBuilder) Truncate(n int) { + if n < 0 { + panic("arrow/array: cannot truncate to a negative length") + } + if n >= b.Len() { + return + } + b.typesBuilder.SetLength(n) + for _, c := range b.children { + c.Truncate(n) + } +} + // AppendNull will append a null to the first child and an empty value // (implementation-defined) to the rest of the children. func (b *SparseUnionBuilder) AppendNull() { @@ -1186,6 +1202,21 @@ func (b *DenseUnionBuilder) Resize(n int) { b.offsetsBuilder.resize(n * arrow.Int32SizeBytes) } +// Truncate reduces the length of the dense union builder to n elements, +// shrinking the types and offsets buffers. The child builders are left +// untouched; their extra elements simply become unreferenced. No buffers are +// reallocated. +func (b *DenseUnionBuilder) Truncate(n int) { + if n < 0 { + panic("arrow/array: cannot truncate to a negative length") + } + if n >= b.Len() { + return + } + b.typesBuilder.SetLength(n) + b.offsetsBuilder.SetLength(n * arrow.Int32SizeBytes) +} + // AppendNull will only append a null value arbitrarily to the first child // and use that offset for this element of the array. func (b *DenseUnionBuilder) AppendNull() {