(arrow-go) branch main updated: perf(internal/bitutils): use four-word fast path for unaligned bitmaps (#1398)

[ze...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

zeroshade pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-go.git


The following commit(s) were added to refs/heads/main by this push:
     new 3c410074 perf(internal/bitutils): use four-word fast path for unaligned bitmaps (#1398)
3c410074 is described below

commit 3c4100746627e225eabe773a5cedfd2360cb4e5a
Author: Minh Vu <vu...@gmail.com>
AuthorDate: Thu Oct 22 19:52:51 2026 +0200

    perf(internal/bitutils): use four-word fast path for unaligned bitmaps (#1398)
    
    ## What
    Use the four-word fast path for unaligned bitmap counts.
    
    ## Why
    The existing threshold is higher than the five backing words needed by
    the shifted reads.
    
    ## Implementation
    - Lower the threshold to the minimum safe range.
    - Add boundary coverage and a package benchmark.
    
    ## Benchmark
    Go 1.26.3, Apple M1 Pro (`darwin/arm64`), median of five runs; 0 B/op
    and 0 allocs/op.
    
    | Metric | Before | After | Change |
    | --- | ---: | ---: | ---: |
    | 1,024 bits, offset 1 | 281.1 ns/op | 95.64 ns/op | -66.0% |
    | 4,096 bits, offset 1 | 415.8 ns/op | 215.3 ns/op | -48.2% |
    
    ---------
    
    Co-authored-by: Matt Topol <zo...@gmail.com>
---
 internal/bitutils/bit_block_counter.go             |  3 +-
 .../bitutils/bit_block_counter_benchmark_test.go   | 54 +++++++++++++++
 internal/bitutils/bit_block_counter_guard_test.go  | 78 ++++++++++++++++++++++
 internal/bitutils/bit_block_counter_test.go        | 32 +++++++++
 4 files changed, 166 insertions(+), 1 deletion(-)

diff --git a/internal/bitutils/bit_block_counter.go b/internal/bitutils/bit_block_counter.go
index 43b2fbb2..53d25c1b 100644
--- a/internal/bitutils/bit_block_counter.go
+++ b/internal/bitutils/bit_block_counter.go
@@ -111,7 +111,8 @@ func (b *BitBlockCounter) NextFourWords() BitBlockCount {
 	} else {
 		// When the offset is > 0, we need there to be a word beyond the last
 		// aligned word in the bitmap for the bit shifting logic.
-		if b.bitsRemaining < 5*fourWordsBits-int64(b.bitOffset) {
+		// Reading four shifted 64-bit words needs five backing words.
+		if b.bitsRemaining < fourWordsBits+wordBits-int64(b.bitOffset) {
 			return b.getBlockSlow(fourWordsBits)
 		}
 
diff --git a/internal/bitutils/bit_block_counter_benchmark_test.go b/internal/bitutils/bit_block_counter_benchmark_test.go
new file mode 100644
index 00000000..28344071
--- /dev/null
+++ b/internal/bitutils/bit_block_counter_benchmark_test.go
@@ -0,0 +1,54 @@
+// 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 bitutils_test
+
+import (
+	"fmt"
+	"math/rand/v2"
+	"testing"
+
+	"github.com/apache/arrow-go/v18/arrow/bitutil"
+	"github.com/apache/arrow-go/v18/internal/bitutils"
+)
+
+var fourWordsBenchmarkSink int64
+
+func BenchmarkFourWordsUnaligned(b *testing.B) {
+	for _, length := range []int64{512, 1024, 4096} {
+		for _, offset := range []int64{1, 3, 7} {
+			bitmap := make([]byte, bitutil.BytesForBits(offset+length))
+			rng := &rand.ChaCha8{}
+			_, _ = rng.Read(bitmap)
+
+			b.Run(fmt.Sprintf("%dbits/offset%d", length, offset), func(b *testing.B) {
+				b.ReportAllocs()
+				b.SetBytes(int64(bitutil.BytesForBits(length)))
+
+				var totalPopcnt int64
+				for b.Loop() {
+					counter := bitutils.NewBitBlockCounter(bitmap, offset, length)
+					for pos := int64(0); pos < length; {
+						block := counter.NextFourWords()
+						pos += int64(block.Len)
+						totalPopcnt += int64(block.Popcnt)
+					}
+				}
+				fourWordsBenchmarkSink = totalPopcnt
+			})
+		}
+	}
+}
diff --git a/internal/bitutils/bit_block_counter_guard_test.go b/internal/bitutils/bit_block_counter_guard_test.go
new file mode 100644
index 00000000..67ea4a6a
--- /dev/null
+++ b/internal/bitutils/bit_block_counter_guard_test.go
@@ -0,0 +1,78 @@
+// 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.
+
+//go:build linux || darwin
+
+package bitutils_test
+
+import (
+	"fmt"
+	"math/rand/v2"
+	"syscall"
+	"testing"
+
+	"github.com/apache/arrow-go/v18/arrow/bitutil"
+	"github.com/apache/arrow-go/v18/internal/bitutils"
+)
+
+func TestFourWordsUnalignedGuardPage(t *testing.T) {
+	pageSize := syscall.Getpagesize()
+	for offset := int64(1); offset < 8; offset++ {
+		for _, length := range []int64{
+			255, 256, 257,
+			320 - offset - 1, 320 - offset, 320 - offset + 1,
+			512, 1024, 1280 - offset - 1, 1280 - offset,
+		} {
+			t.Run(fmt.Sprintf("offset%d/length%d", offset, length), func(t *testing.T) {
+				mapped, err := syscall.Mmap(-1, 0, 2*pageSize, syscall.PROT_READ|syscall.PROT_WRITE, syscall.MAP_PRIVATE|syscall.MAP_ANON)
+				if err != nil {
+					t.Fatalf("mmap: %v", err)
+				}
+				t.Cleanup(func() {
+					if err := syscall.Munmap(mapped); err != nil {
+						t.Errorf("munmap: %v", err)
+					}
+				})
+				if err := syscall.Mprotect(mapped[pageSize:], syscall.PROT_NONE); err != nil {
+					t.Fatalf("mprotect: %v", err)
+				}
+
+				// An unsafe read beyond the last required byte reaches the protected page.
+				numBytes := int(bitutil.BytesForBits(offset + length))
+				bitmap := mapped[pageSize-numBytes : pageSize]
+				rng := &rand.ChaCha8{}
+				_, _ = rng.Read(bitmap)
+
+				counter := bitutils.NewBitBlockCounter(bitmap, offset, length)
+				for pos := int64(0); pos < length; {
+					block := counter.NextFourWords()
+					wantLen := min(int64(256), length-pos)
+					if int64(block.Len) != wantLen {
+						t.Fatalf("offset=%d length=%d pos=%d: block length=%d, want %d", offset, length, pos, block.Len, wantLen)
+					}
+					wantPopcnt := bitutil.CountSetBits(bitmap, int(offset+pos), int(block.Len))
+					if int(block.Popcnt) != wantPopcnt {
+						t.Fatalf("offset=%d length=%d pos=%d: popcount=%d, want %d", offset, length, pos, block.Popcnt, wantPopcnt)
+					}
+					pos += int64(block.Len)
+				}
+				if got := counter.NextFourWords().Len; got != 0 {
+					t.Fatalf("offset=%d length=%d: final block length=%d, want 0", offset, length, got)
+				}
+			})
+		}
+	}
+}
diff --git a/internal/bitutils/bit_block_counter_test.go b/internal/bitutils/bit_block_counter_test.go
index 8ebbbe88..5b9f364e 100644
--- a/internal/bitutils/bit_block_counter_test.go
+++ b/internal/bitutils/bit_block_counter_test.go
@@ -199,3 +199,35 @@ func TestFourWordsRandomData(t *testing.T) {
 		checkWithOffset(offsetI)
 	}
 }
+
+func TestFourWordsUnalignedBoundaries(t *testing.T) {
+	for offset := int64(1); offset < 8; offset++ {
+		for _, length := range []int64{
+			255, 256, 257,
+			320 - offset - 1, 320 - offset, 320 - offset + 1,
+			512, 1024, 1280 - offset - 1, 1280 - offset,
+		} {
+			// Use an exact-length bitmap for functional boundary coverage.
+			// The guard-page test separately detects unsafe overreads.
+			bitmap := make([]byte, bitutil.BytesForBits(offset+length))
+			rng := &rand.ChaCha8{}
+			_, _ = rng.Read(bitmap)
+
+			counter := bitutils.NewBitBlockCounter(bitmap, offset, length)
+			for pos := int64(0); pos < length; {
+				block := counter.NextFourWords()
+				if block.Len <= 0 {
+					t.Fatalf("no progress at offset=%d length=%d pos=%d", offset, length, pos)
+				}
+				assert.EqualValues(t, min(int64(256), length-pos), block.Len,
+					"offset=%d length=%d pos=%d", offset, length, pos)
+				assert.EqualValues(t,
+					bitutil.CountSetBits(bitmap, int(offset+pos), int(block.Len)),
+					block.Popcnt, "offset=%d length=%d pos=%d", offset, length, pos)
+				pos += int64(block.Len)
+			}
+
+			assert.Zero(t, counter.NextFourWords().Len)
+		}
+	}
+}