perf: root BatchEqual in the existence bitmap and share plane work across values

Follow-up addressing review findings on the vectorized implementation.

Correctness: the per-value path built results from bitplane clones and
never intersected eBM, so columns present in a plane but absent from
the existence bitmap leaked into results. The checked-in testdata/age
fixture contains exactly this state (plane 4 holds 33.3M bits against
28.0M in eBM): BatchEqual(0, {16}) returned 454,521 columns, all of
them outside eBM, where the pre-optimization implementation returned
none. GetValue, CompareValue, Sum, and the original BatchEqual all
treat eBM as authoritative; the new implementation is rooted in eBM by
construction, and a regression test pins the invariant.

Scaling: the maxVectorizedBatchSize=128 crossover was inverted on large
datasets. Measured on the age fixture, the vectorized path at M=128
cost 3.6s / 2.1GB while the scalar fallback it guarded cost 0.59s: the
guard fired on list size when the real variable is data size. Both
thresholds and the scalar fallback path are deleted.

BatchEqual now descends the bitplanes from the most significant bit as
a binary trie over the sorted, deduplicated query values: each shared
value prefix is intersected once, subtrees whose every bit pattern is
queried collapse to a single operation, and empty intermediates prune
the recursion. Work is bounded by the size of the values' bit trie and
the data actually present rather than len(values) x BitCount().

testdata/age fixture (28M columns, 8 planes), 16 threads, -benchmem,
previous commit -> this commit:

  M=2 {55,57}       ~56ms / 33.5MB / 10.3k allocs -> ~50ms / 26.7MB / 8.2k
  M=128 contiguous   3.6s / 2.1GB                 -> 1.2ms / 9.3MB
  M=128 scattered    3.3s / 2.2GB                 -> 440ms / 283MB
  M=200 (fallback)   590ms / 155MB                -> 3.3ms / 9.5MB

The M=2 delta is within run-to-run noise; allocation counts are exact.
New benchmarks cover both sides of the deleted thresholds and both
contiguous and scattered value shapes.

Author: perfloop-agent

BitSliceIndexing/bsi.go

@@ -4,6 +4,7 @@ import (
 	"fmt"
 	"math/bits"
 	"runtime"
+	"sort"
 	"sync"
 	"sync/atomic"
 
@@ -271,7 +272,6 @@ type task struct {
 	op           Operation
 	valueOrStart int64
 	end          int64
-	values       map[int64]struct{}
 	bits         *roaring.Bitmap
 }
 
@@ -743,134 +743,100 @@ func (b *BSI) MarshalBinary() ([][]byte, error) {
 	return data, nil
 }
 
-// maxVectorizedBatchSize is the threshold above which BatchEqual falls back
-// to the scalar parallel executor. This prevents a memory allocation/CPU scaling
-// cliff when querying a very large list of unique values (large M) on a small/sparse
-// database, which would otherwise allocate and process M intermediate bitmaps.
-const maxVectorizedBatchSize = 128
-
-// maxInPlaceCombineSize is the threshold below which BatchEqual combines the
-// computed match bitmaps sequentially in-place using Bitmap.Or rather than roaring.ParOr.
-// This avoids the goroutine spawning, chunking spec, and channel allocation overhead
-// of ParOr when combining a small number of bitmaps, keeping allocs/op low and
-// avoiding a performance regression on typical small queries.
-const maxInPlaceCombineSize = 16
-
-// BatchEqual returns a bitmap containing the column IDs where the values are contained within the list of values provided.
+// BatchEqual returns a bitmap containing the column IDs where the values are contained
+// within the list of values provided. The parallelism parameter is accepted for API
+// compatibility; work is shared across values, so the query runs on the calling goroutine.
 func (b *BSI) BatchEqual(parallelism int, values []int64) *roaring.Bitmap {
 	if b.eBM.IsEmpty() || len(values) == 0 {
 		return roaring.NewBitmap()
 	}
 
-	if parallelism < 0 {
-		parallelism = 0
-	}
+	bitCount := b.BitCount()
 
-	// Deduplicate values
-	valMap := make(map[int64]struct{}, len(values))
+	// Deduplicate, and drop values that cannot be represented in bitCount
+	// planes: GetValue can never observe such a value, so it matches no column.
+	seen := make(map[uint64]struct{}, len(values))
+	vals := make([]uint64, 0, len(values))
 	for _, v := range values {
-		valMap[v] = struct{}{}
-	}
-
-	// Filter and collect valid unique values
-	bitCount := b.BitCount()
-	validValues := make([]int64, 0, len(valMap))
-	for v := range valMap {
-		if bitCount >= 64 {
-			validValues = append(validValues, v)
-		} else {
-			if v >= 0 && uint64(v) < (uint64(1)<<uint64(bitCount)) {
-				validValues = append(validValues, v)
-			}
+		u := uint64(v)
+		if bitCount < 64 && (v < 0 || u >= uint64(1)<<uint(bitCount)) {
+			continue
+		}
+		if _, ok := seen[u]; ok {
+			continue
 		}
+		seen[u] = struct{}{}
+		vals = append(vals, u)
 	}
-
-	if len(validValues) == 0 {
+	if len(vals) == 0 {
 		return roaring.NewBitmap()
 	}
+	sort.Slice(vals, func(i, j int) bool { return vals[i] < vals[j] })
 
-	// Fallback to scalar parallel executor if query list is very large
-	if len(validValues) > maxVectorizedBatchSize {
-		comp := &task{bsi: b, values: valMap}
-		return parallelExecutor(parallelism, comp, batchEqualFallback, b.eBM)
-	}
-
-	bitmaps := make([]*roaring.Bitmap, len(validValues))
-	for i, v := range validValues {
-		bitmaps[i] = b.getEqualBitmap(v)
-	}
-
-	var finalResult *roaring.Bitmap
-	if len(bitmaps) == 1 {
-		finalResult = bitmaps[0]
-	} else if len(bitmaps) <= maxInPlaceCombineSize {
-		finalResult = bitmaps[0]
-		for i := 1; i < len(bitmaps); i++ {
-			finalResult.Or(bitmaps[i])
-		}
-	} else {
-		finalResult = roaring.ParOr(parallelism, bitmaps...)
-	}
-
+	result := b.matchTrie(vals, bitCount-1, b.eBM, false)
 	if b.runOptimized {
-		finalResult.RunOptimize()
+		result.RunOptimize()
 	}
-	return finalResult
+	return result
 }
 
-// getEqualBitmap returns a bitmap of column IDs whose value is exactly v
-func (b *BSI) getEqualBitmap(v int64) *roaring.Bitmap {
-	bitCount := b.BitCount()
-	// Find first set bit to start with a sparse bitmap instead of eBM
-	firstSetBit := -1
-	for i := 0; i < bitCount; i++ {
-		if (uint64(v) & (1 << uint64(i))) != 0 {
-			firstSetBit = i
-			break
+// matchTrie returns the columns in prefix whose bits on planes [0, p] equal the low
+// p+1 bits of any value in vals. vals must be unique, sorted, and share identical bits
+// above plane p; prefix holds the columns matching those shared upper bits, so results
+// stay rooted in the existence bitmap the recursion starts from. owned reports whether
+// prefix is a private intermediate the callee may consume in place. Each shared value
+// prefix is intersected exactly once and empty intermediates prune the recursion, so
+// total work is bounded by the size of the values' bit trie and the data actually
+// present, not len(vals) × BitCount().
+func (b *BSI) matchTrie(vals []uint64, p int, prefix *roaring.Bitmap, owned bool) *roaring.Bitmap {
+	if prefix.IsEmpty() {
+		if owned {
+			return prefix
 		}
+		return roaring.NewBitmap()
 	}
-
-	var vBM *roaring.Bitmap
-	if firstSetBit != -1 {
-		vBM = b.bA[firstSetBit].Clone()
-		for i := 0; i < bitCount; i++ {
-			if i == firstSetBit {
-				continue
-			}
-			if (uint64(v) & (1 << uint64(i))) != 0 {
-				vBM.And(b.bA[i])
-			} else {
-				vBM.AndNot(b.bA[i])
-			}
-		}
-	} else {
-		vBM = b.eBM.Clone()
-		for i := 0; i < bitCount; i++ {
-			vBM.AndNot(b.bA[i])
+	// Either all planes are consumed (exactly one value remains and prefix matched
+	// every one of its bits), or vals covers every bit pattern of the remaining
+	// planes, which collapses dense value ranges to a single operation. In both
+	// cases every column in prefix matches.
+	if p < 0 || (p < 63 && uint64(len(vals)) == uint64(1)<<uint(p+1)) {
+		if owned {
+			return prefix
+		}
+		return prefix.Clone()
+	}
+	// vals is sorted and shares all bits above p, so bit p partitions it in place.
+	cut := sort.Search(len(vals), func(i int) bool { return vals[i]&(uint64(1)<<uint(p)) != 0 })
+	lo, hi := vals[:cut], vals[cut:]
+	switch {
+	case len(hi) == 0:
+		return b.matchTrie(lo, p-1, planeChild(prefix, b.bA[p], false, owned), true)
+	case len(lo) == 0:
+		return b.matchTrie(hi, p-1, planeChild(prefix, b.bA[p], true, owned), true)
+	default:
+		// Materialize the set branch before the clear branch may consume prefix.
+		hiBM := roaring.And(prefix, b.bA[p])
+		result := b.matchTrie(lo, p-1, planeChild(prefix, b.bA[p], false, owned), true)
+		result.Or(b.matchTrie(hi, p-1, hiBM, true))
+		return result
+	}
+}
+
+// planeChild narrows prefix by one bitplane: prefix ∧ plane when set, prefix ∧ ¬plane
+// otherwise. When the caller owns prefix it is consumed in place to avoid a copy.
+func planeChild(prefix, plane *roaring.Bitmap, set, owned bool) *roaring.Bitmap {
+	if owned {
+		if set {
+			prefix.And(plane)
+		} else {
+			prefix.AndNot(plane)
 		}
+		return prefix
 	}
-	return vBM
-}
-
-func batchEqualFallback(e *task, batch []uint32, resultsChan chan *roaring.Bitmap,
-	wg *sync.WaitGroup) {
-
-	defer wg.Done()
-
-	results := roaring.NewBitmap()
-	if e.bsi.runOptimized {
-		results.RunOptimize()
+	if set {
+		return roaring.And(prefix, plane)
 	}
-
-	for i := 0; i < len(batch); i++ {
-		cID := batch[i]
-		if value, ok := e.bsi.GetValue(uint64(cID)); ok {
-			if _, yes := e.values[value]; yes {
-				results.Add(cID)
-			}
-		}
-	}
-	resultsChan <- results
+	return roaring.AndNot(prefix, plane)
 }
 
 // ClearBits cleared the bits that exist in the target if they are also in the found set.

BitSliceIndexing/bsi_benchmark_test.go

@@ -171,3 +171,64 @@ func TestBatchEqualConsistentWithGetValue(t *testing.T) {
 		}
 	}
 }
+
+// TestBatchEqualExistenceAuthority pins BatchEqual results to the existence
+// bitmap. UnmarshalBinary accepts plane data that is not a subset of eBM (the
+// checked-in testdata/age fixture is such data), and every read path treats
+// eBM as authoritative, so columns present in a plane but absent from eBM must
+// never appear in results.
+func TestBatchEqualExistenceAuthority(t *testing.T) {
+	// Synthetic state: column 2 has bits in plane 0 but is absent from eBM.
+	ebm := roaring.BitmapOf(1)
+	plane := roaring.BitmapOf(1, 2)
+	ebmData, err := ebm.MarshalBinary()
+	if err != nil {
+		t.Fatal(err)
+	}
+	planeData, err := plane.MarshalBinary()
+	if err != nil {
+		t.Fatal(err)
+	}
+	bsi := NewDefaultBSI()
+	if err := bsi.UnmarshalBinary([][]byte{ebmData, planeData}); err != nil {
+		t.Fatal(err)
+	}
+	res := bsi.BatchEqual(0, []int64{1})
+	assert.True(t, res.Contains(1))
+	assert.False(t, res.Contains(2), "column 2 is not in eBM and must not match")
+
+	// The age fixture ships with plane cardinalities above the eBM cardinality;
+	// results must still be a subset of eBM.
+	large := setupLargeBSI(t)
+	if large == nil {
+		t.Skip("skipping, large BSI setup failed")
+	}
+	for _, vals := range [][]int64{{16}, {55, 57}, {0, 1, 2, 3}} {
+		res := large.BatchEqual(0, vals)
+		outside := roaring.AndNot(res, large.GetExistenceBitmap())
+		assert.True(t, outside.IsEmpty(), "BatchEqual(%v) returned %d columns outside eBM", vals, outside.GetCardinality())
+	}
+}
+
+// Benchmarks across query-list shapes: work sharing behaves differently for
+// small lists, dense contiguous ranges (which collapse to a few plane
+// operations), and scattered values (no range collapse).
+func BenchmarkBatchEqualM128(b *testing.B)          { benchmarkBatchEqualM(b, 128, 1) }
+func BenchmarkBatchEqualM128Scattered(b *testing.B) { benchmarkBatchEqualM(b, 128, 2) }
+func BenchmarkBatchEqualM200(b *testing.B)          { benchmarkBatchEqualM(b, 200, 1) }
+
+func benchmarkBatchEqualM(b *testing.B, m int, stride int64) {
+	bsi := setupLargeBSI(b)
+	if bsi == nil {
+		b.Skip("skipping, large BSI setup failed")
+	}
+	vals := make([]int64, m)
+	for i := range vals {
+		vals[i] = int64(i)*stride + stride - 1
+	}
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		res := bsi.BatchEqual(0, vals)
+		_ = res
+	}
+}