Skip to content

Commit 2fdd1b8

Browse files
fallintoplacefallintoplace
authored
perf(arrow/array): batch-pack boolean values (#1237)
## Summary - Pack BooleanBuilder values into the bitmap one byte at a time. - Handle unaligned prefixes and trailing values without changing neighboring bits. - Reuse the 8-value pack helper for value and validity bitmaps. - Add exhaustive bitmap coverage and a BooleanBuilder benchmark. ## Benchmark Apple M1 Pro. 65,536 values. The builder is pre-reserved. Median of 5 runs. | Benchmark | upstream main | this PR | change | | --- | ---: | ---: | ---: | | BooleanBuilder / all false | 145 us | 18.6 us | 7.8x | | BooleanBuilder / alternating | 145 us | 18.6 us | 7.8x | | Boolean / all valid | 160 us | 34.4 us | 4.6x | | Boolean / 50% null | 161 us | 35.7 us | 4.5x | The first two rows use nil validity. The last two use the existing explicit-validity benchmark. ```text go test ./arrow/array -run '^$' -bench '^BenchmarkBooleanBuilderAppendValues$' -benchmem -benchtime=100ms -count=5 go test ./arrow/array -run '^$' -bench '^BenchmarkAppendValuesWithValidity/boolean/' -benchmem -benchtime=100ms -count=5 ``` ## Tests - `go test ./arrow/array -count=1` - `go test -race ./arrow/array -count=1` - `go test ./arrow/bitutil ./arrow/compute/... -count=1` - `go vet ./arrow/array ./arrow/bitutil`
1 parent 6e4e241 commit 2fdd1b8

4 files changed

Lines changed: 166 additions & 15 deletions

File tree

‎arrow/array/booleanbuilder.go‎

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -135,8 +135,12 @@ func (b *BooleanBuilder) AppendValues(v []bool, valid []bool) {
135135
}
136136

137137
b.Reserve(len(v))
138-
for i, vv := range v {
139-
bitutil.SetBitTo(b.rawData, b.length+i, vv)
138+
if len(v) < 8 {
139+
for i, vv := range v {
140+
bitutil.SetBitTo(b.rawData, b.length+i, vv)
141+
}
142+
} else {
143+
packBoolsToBitmap(b.rawData, b.length, v)
140144
}
141145
b.unsafeAppendBoolsToBitmap(valid, len(v))
142146
}

‎arrow/array/booleanbuilder_benchmark_test.go‎

Lines changed: 74 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,74 @@
1+
// Licensed to the Apache Software Foundation (ASF) under one
2+
// or more contributor license agreements. See the NOTICE file
3+
// distributed with this work for additional information
4+
// regarding copyright ownership. The ASF licenses this file
5+
// to you under the Apache License, Version 2.0 (the
6+
// "License"); you may not use this file except in compliance
7+
// with the License. You may obtain a copy of the License at
8+
//
9+
// http://www.apache.org/licenses/LICENSE-2.0
10+
//
11+
// Unless required by applicable law or agreed to in writing, software
12+
// distributed under the License is distributed on an "AS IS" BASIS,
13+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14+
// See the License for the specific language governing permissions and
15+
// limitations under the License.
16+
17+
package array
18+
19+
import (
20+
"fmt"
21+
"testing"
22+
23+
"github.com/apache/arrow-go/v18/arrow/memory"
24+
)
25+
26+
func BenchmarkBooleanBuilderAppendValues(b *testing.B) {
27+
const length = 65536
28+
patterns := []struct {
29+
name string
30+
values []bool
31+
}{
32+
{"all-false", makeBooleanBenchmarkValues(length, func(int) bool { return false })},
33+
{"all-true", makeBooleanBenchmarkValues(length, func(int) bool { return true })},
34+
{"alternating", makeBooleanBenchmarkValues(length, func(i int) bool { return i%2 == 0 })},
35+
{"one-in-three", makeBooleanBenchmarkValues(length, func(i int) bool { return i%3 == 0 })},
36+
}
37+
38+
for _, pattern := range patterns {
39+
b.Run(pattern.name, func(b *testing.B) {
40+
benchmarkAppendValues(b, func() (func(), func()) {
41+
bldr := NewBooleanBuilder(memory.DefaultAllocator)
42+
bldr.Reserve(length)
43+
return func() {
44+
bldr.AppendValues(pattern.values, nil)
45+
}, bldr.Release
46+
})
47+
})
48+
}
49+
}
50+
51+
func BenchmarkBooleanBuilderAppendValuesSmall(b *testing.B) {
52+
for _, length := range []int{1, 2, 3, 7, 8} {
53+
b.Run(fmt.Sprintf("len=%d", length), func(b *testing.B) {
54+
values := makeBooleanBenchmarkValues(length, func(i int) bool { return i%2 == 0 })
55+
bldr := NewBooleanBuilder(memory.DefaultAllocator)
56+
bldr.Reserve(len(values) * b.N)
57+
b.ReportAllocs()
58+
b.ResetTimer()
59+
for i := 0; i < b.N; i++ {
60+
bldr.AppendValues(values, nil)
61+
}
62+
b.StopTimer()
63+
bldr.Release()
64+
})
65+
}
66+
}
67+
68+
func makeBooleanBenchmarkValues(length int, value func(int) bool) []bool {
69+
values := make([]bool, length)
70+
for i := range values {
71+
values[i] = value(i)
72+
}
73+
return values
74+
}

‎arrow/array/builder.go‎

Lines changed: 52 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -238,7 +238,7 @@ func (b *builder) unsafeAppendBoolsToBitmap(valid []bool, length int) {
238238
}
239239

240240
for len(valid) >= 8 {
241-
bitSet := packValidityByte(valid)
241+
bitSet := packBoolsByte(valid)
242242
nullBitmap[byteOffset] = bitSet
243243
b.nulls += 8 - bits.OnesCount8(bitSet)
244244
valid = valid[8:]
@@ -260,36 +260,77 @@ func (b *builder) unsafeAppendBoolsToBitmap(valid []bool, length int) {
260260
b.length += validLength
261261
}
262262

263-
func packValidityByte(valid []bool) byte {
264-
valid = valid[:8]
263+
func packBoolsByte(values []bool) byte {
264+
values = values[:8]
265265
var packed byte
266-
if valid[0] {
266+
if values[0] {
267267
packed |= 1 << 0
268268
}
269-
if valid[1] {
269+
if values[1] {
270270
packed |= 1 << 1
271271
}
272-
if valid[2] {
272+
if values[2] {
273273
packed |= 1 << 2
274274
}
275-
if valid[3] {
275+
if values[3] {
276276
packed |= 1 << 3
277277
}
278-
if valid[4] {
278+
if values[4] {
279279
packed |= 1 << 4
280280
}
281-
if valid[5] {
281+
if values[5] {
282282
packed |= 1 << 5
283283
}
284-
if valid[6] {
284+
if values[6] {
285285
packed |= 1 << 6
286286
}
287-
if valid[7] {
287+
if values[7] {
288288
packed |= 1 << 7
289289
}
290290
return packed
291291
}
292292

293+
func packBoolsToBitmap(dst []byte, offset int, values []bool) {
294+
if len(values) == 0 {
295+
return
296+
}
297+
298+
byteOffset := offset / 8
299+
bitOffset := offset % 8
300+
if bitOffset != 0 {
301+
bitSet := dst[byteOffset]
302+
prefixLength := min(8-bitOffset, len(values))
303+
for i, v := range values[:prefixLength] {
304+
if v {
305+
bitSet |= bitutil.BitMask[bitOffset+i]
306+
} else {
307+
bitSet &= bitutil.FlippedBitMask[bitOffset+i]
308+
}
309+
}
310+
dst[byteOffset] = bitSet
311+
values = values[prefixLength:]
312+
byteOffset++
313+
}
314+
315+
for len(values) >= 8 {
316+
dst[byteOffset] = packBoolsByte(values)
317+
values = values[8:]
318+
byteOffset++
319+
}
320+
321+
if len(values) != 0 {
322+
bitSet := dst[byteOffset]
323+
for i, v := range values {
324+
if v {
325+
bitSet |= bitutil.BitMask[i]
326+
} else {
327+
bitSet &= bitutil.FlippedBitMask[i]
328+
}
329+
}
330+
dst[byteOffset] = bitSet
331+
}
332+
}
333+
293334
// unsafeSetValid sets the next length bits to valid in the validity bitmap.
294335
func (b *builder) unsafeSetValid(length int) {
295336
padToByte := min(8-(b.length%8), length)

‎arrow/array/builder_test.go‎

Lines changed: 34 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -103,13 +103,45 @@ func TestBuilder_UnsafeAppendBoolsToBitmap(t *testing.T) {
103103
}
104104
}
105105

106-
func TestPackValidityByte(t *testing.T) {
106+
func TestPackBoolsByte(t *testing.T) {
107107
for want := 0; want < 1<<8; want++ {
108108
valid := make([]bool, 8)
109109
for i := range valid {
110110
valid[i] = want&(1<<i) != 0
111111
}
112-
assert.Equal(t, byte(want), packValidityByte(valid), "want=%08b", want)
112+
assert.Equal(t, byte(want), packBoolsByte(valid), "want=%08b", want)
113+
}
114+
}
115+
116+
func TestPackBoolsToBitmap(t *testing.T) {
117+
patterns := []struct {
118+
name string
119+
value func(int) bool
120+
}{
121+
{"all false", func(int) bool { return false }},
122+
{"all true", func(int) bool { return true }},
123+
{"alternating", func(i int) bool { return i%2 == 0 }},
124+
{"one in three", func(i int) bool { return i%3 == 0 }},
125+
}
126+
127+
for _, pattern := range patterns {
128+
for offset := 0; offset < 8; offset++ {
129+
for length := 0; length <= 33; length++ {
130+
got := make([]byte, 8)
131+
for i := range got {
132+
got[i] = byte(0x5a + i*31)
133+
}
134+
want := append([]byte(nil), got...)
135+
values := make([]bool, length)
136+
for i := range values {
137+
values[i] = pattern.value(i)
138+
bitutil.SetBitTo(want, offset+i, values[i])
139+
}
140+
141+
packBoolsToBitmap(got, offset, values)
142+
assert.Equal(t, want, got, "%s, offset=%d, length=%d", pattern.name, offset, length)
143+
}
144+
}
113145
}
114146
}
115147

0 commit comments

Comments
 (0)