perf(arrow/array): batch-pack boolean values

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))
 }

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
+}

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)

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<<i) != 0
 		}
-		assert.Equal(t, byte(want), packValidityByte(valid), "want=%08b", want)
+		assert.Equal(t, byte(want), packBoolsByte(valid), "want=%08b", want)
+	}
+}
+
+func TestPackBoolsToBitmap(t *testing.T) {
+	patterns := []struct {
+		name  string
+		value func(int) bool
+	}{
+		{"all false", func(int) bool { return false }},
+		{"all true", func(int) bool { return true }},
+		{"alternating", func(i int) bool { return i%2 == 0 }},
+		{"one in three", func(i int) bool { return i%3 == 0 }},
+	}
+
+	for _, pattern := range patterns {
+		for offset := 0; offset < 8; offset++ {
+			for length := 0; length <= 33; length++ {
+				got := make([]byte, 8)
+				for i := range got {
+					got[i] = byte(0x5a + i*31)
+				}
+				want := append([]byte(nil), got...)
+				values := make([]bool, length)
+				for i := range values {
+					values[i] = pattern.value(i)
+					bitutil.SetBitTo(want, offset+i, values[i])
+				}
+
+				packBoolsToBitmap(got, offset, values)
+				assert.Equal(t, want, got, "%s, offset=%d, length=%d", pattern.name, offset, length)
+			}
+		}
 	}
 }