diff --git a/arrow/array/booleanbuilder.go b/arrow/array/booleanbuilder.go index 88a1e00c..f90ac0b4 100644 --- a/arrow/array/booleanbuilder.go +++ b/arrow/array/booleanbuilder.go @@ -135,8 +135,12 @@ func (b *BooleanBuilder) AppendValues(v []bool, valid []bool) { } b.Reserve(len(v)) - for i, vv := range v { - bitutil.SetBitTo(b.rawData, b.length+i, vv) + if len(v) < 8 { + for i, vv := range v { + bitutil.SetBitTo(b.rawData, b.length+i, vv) + } + } else { + packBoolsToBitmap(b.rawData, b.length, v) } b.unsafeAppendBoolsToBitmap(valid, len(v)) } diff --git a/arrow/array/booleanbuilder_benchmark_test.go b/arrow/array/booleanbuilder_benchmark_test.go new file mode 100644 index 00000000..0273c503 --- /dev/null +++ b/arrow/array/booleanbuilder_benchmark_test.go @@ -0,0 +1,74 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package array + +import ( + "fmt" + "testing" + + "github.com/apache/arrow-go/v18/arrow/memory" +) + +func BenchmarkBooleanBuilderAppendValues(b *testing.B) { + const length = 65536 + patterns := []struct { + name string + values []bool + }{ + {"all-false", makeBooleanBenchmarkValues(length, func(int) bool { return false })}, + {"all-true", makeBooleanBenchmarkValues(length, func(int) bool { return true })}, + {"alternating", makeBooleanBenchmarkValues(length, func(i int) bool { return i%2 == 0 })}, + {"one-in-three", makeBooleanBenchmarkValues(length, func(i int) bool { return i%3 == 0 })}, + } + + for _, pattern := range patterns { + b.Run(pattern.name, func(b *testing.B) { + benchmarkAppendValues(b, func() (func(), func()) { + bldr := NewBooleanBuilder(memory.DefaultAllocator) + bldr.Reserve(length) + return func() { + bldr.AppendValues(pattern.values, nil) + }, bldr.Release + }) + }) + } +} + +func BenchmarkBooleanBuilderAppendValuesSmall(b *testing.B) { + for _, length := range []int{1, 2, 3, 7, 8} { + b.Run(fmt.Sprintf("len=%d", length), func(b *testing.B) { + values := makeBooleanBenchmarkValues(length, func(i int) bool { return i%2 == 0 }) + bldr := NewBooleanBuilder(memory.DefaultAllocator) + bldr.Reserve(len(values) * b.N) + b.ReportAllocs() + b.ResetTimer() + for i := 0; i < b.N; i++ { + bldr.AppendValues(values, nil) + } + b.StopTimer() + bldr.Release() + }) + } +} + +func makeBooleanBenchmarkValues(length int, value func(int) bool) []bool { + values := make([]bool, length) + for i := range values { + values[i] = value(i) + } + return values +} diff --git a/arrow/array/builder.go b/arrow/array/builder.go index 11bffb01..2b5aa2cf 100644 --- a/arrow/array/builder.go +++ b/arrow/array/builder.go @@ -238,7 +238,7 @@ func (b *builder) unsafeAppendBoolsToBitmap(valid []bool, length int) { } for len(valid) >= 8 { - bitSet := packValidityByte(valid) + bitSet := packBoolsByte(valid) nullBitmap[byteOffset] = bitSet b.nulls += 8 - bits.OnesCount8(bitSet) valid = valid[8:] @@ -260,36 +260,77 @@ func (b *builder) unsafeAppendBoolsToBitmap(valid []bool, length int) { b.length += validLength } -func packValidityByte(valid []bool) byte { - valid = valid[:8] +func packBoolsByte(values []bool) byte { + values = values[:8] var packed byte - if valid[0] { + if values[0] { packed |= 1 << 0 } - if valid[1] { + if values[1] { packed |= 1 << 1 } - if valid[2] { + if values[2] { packed |= 1 << 2 } - if valid[3] { + if values[3] { packed |= 1 << 3 } - if valid[4] { + if values[4] { packed |= 1 << 4 } - if valid[5] { + if values[5] { packed |= 1 << 5 } - if valid[6] { + if values[6] { packed |= 1 << 6 } - if valid[7] { + if values[7] { packed |= 1 << 7 } return packed } +func packBoolsToBitmap(dst []byte, offset int, values []bool) { + if len(values) == 0 { + return + } + + byteOffset := offset / 8 + bitOffset := offset % 8 + if bitOffset != 0 { + bitSet := dst[byteOffset] + prefixLength := min(8-bitOffset, len(values)) + for i, v := range values[:prefixLength] { + if v { + bitSet |= bitutil.BitMask[bitOffset+i] + } else { + bitSet &= bitutil.FlippedBitMask[bitOffset+i] + } + } + dst[byteOffset] = bitSet + values = values[prefixLength:] + byteOffset++ + } + + for len(values) >= 8 { + dst[byteOffset] = packBoolsByte(values) + values = values[8:] + byteOffset++ + } + + if len(values) != 0 { + bitSet := dst[byteOffset] + for i, v := range values { + if v { + bitSet |= bitutil.BitMask[i] + } else { + bitSet &= bitutil.FlippedBitMask[i] + } + } + dst[byteOffset] = bitSet + } +} + // unsafeSetValid sets the next length bits to valid in the validity bitmap. func (b *builder) unsafeSetValid(length int) { padToByte := min(8-(b.length%8), length) diff --git a/arrow/array/builder_test.go b/arrow/array/builder_test.go index 166a6a4c..2424c139 100644 --- a/arrow/array/builder_test.go +++ b/arrow/array/builder_test.go @@ -103,13 +103,45 @@ func TestBuilder_UnsafeAppendBoolsToBitmap(t *testing.T) { } } -func TestPackValidityByte(t *testing.T) { +func TestPackBoolsByte(t *testing.T) { for want := 0; want < 1<<8; want++ { valid := make([]bool, 8) for i := range valid { valid[i] = want&(1<