Skip to content

Commit dac837d

Browse files
authored
Merge pull request #532 from perfloop/perfloop-pr-open-j88wgmc05t
bsi: vectorize and simplify batch equal with direct bitwise set operations
2 parents bd4128b + 7fe6c1a commit dac837d

2 files changed

Lines changed: 321 additions & 23 deletions

File tree

BitSliceIndexing/bsi.go

Lines changed: 87 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@ import (
44
"fmt"
55
"math/bits"
66
"runtime"
7+
"sort"
78
"sync"
89
"sync/atomic"
910

@@ -271,7 +272,6 @@ type task struct {
271272
op Operation
272273
valueOrStart int64
273274
end int64
274-
values map[int64]struct{}
275275
bits *roaring.Bitmap
276276
}
277277

@@ -743,36 +743,100 @@ func (b *BSI) MarshalBinary() ([][]byte, error) {
743743
return data, nil
744744
}
745745

746-
// BatchEqual returns a bitmap containing the column IDs where the values are contained within the list of values provided.
746+
// BatchEqual returns a bitmap containing the column IDs where the values are contained
747+
// within the list of values provided. The parallelism parameter is accepted for API
748+
// compatibility; work is shared across values, so the query runs on the calling goroutine.
747749
func (b *BSI) BatchEqual(parallelism int, values []int64) *roaring.Bitmap {
748-
749-
valMap := make(map[int64]struct{}, len(values))
750-
for i := 0; i < len(values); i++ {
751-
valMap[values[i]] = struct{}{}
750+
if b.eBM.IsEmpty() || len(values) == 0 {
751+
return roaring.NewBitmap()
752752
}
753-
comp := &task{bsi: b, values: valMap}
754-
return parallelExecutor(parallelism, comp, batchEqual, b.eBM)
755-
}
756-
757-
func batchEqual(e *task, batch []uint32, resultsChan chan *roaring.Bitmap,
758-
wg *sync.WaitGroup) {
759753

760-
defer wg.Done()
754+
bitCount := b.BitCount()
761755

762-
results := roaring.NewBitmap()
763-
if e.bsi.runOptimized {
764-
results.RunOptimize()
756+
// Deduplicate, and drop values that cannot be represented in bitCount
757+
// planes: GetValue can never observe such a value, so it matches no column.
758+
seen := make(map[uint64]struct{}, len(values))
759+
vals := make([]uint64, 0, len(values))
760+
for _, v := range values {
761+
u := uint64(v)
762+
if bitCount < 64 && (v < 0 || u >= uint64(1)<<uint(bitCount)) {
763+
continue
764+
}
765+
if _, ok := seen[u]; ok {
766+
continue
767+
}
768+
seen[u] = struct{}{}
769+
vals = append(vals, u)
765770
}
771+
if len(vals) == 0 {
772+
return roaring.NewBitmap()
773+
}
774+
sort.Slice(vals, func(i, j int) bool { return vals[i] < vals[j] })
766775

767-
for i := 0; i < len(batch); i++ {
768-
cID := batch[i]
769-
if value, ok := e.bsi.GetValue(uint64(cID)); ok {
770-
if _, yes := e.values[value]; yes {
771-
results.Add(cID)
772-
}
776+
result := b.matchTrie(vals, bitCount-1, b.eBM, false)
777+
if b.runOptimized {
778+
result.RunOptimize()
779+
}
780+
return result
781+
}
782+
783+
// matchTrie returns the columns in prefix whose bits on planes [0, p] equal the low
784+
// p+1 bits of any value in vals. vals must be unique, sorted, and share identical bits
785+
// above plane p; prefix holds the columns matching those shared upper bits, so results
786+
// stay rooted in the existence bitmap the recursion starts from. owned reports whether
787+
// prefix is a private intermediate the callee may consume in place. Each shared value
788+
// prefix is intersected exactly once and empty intermediates prune the recursion, so
789+
// total work is bounded by the size of the values' bit trie and the data actually
790+
// present, not len(vals) × BitCount().
791+
func (b *BSI) matchTrie(vals []uint64, p int, prefix *roaring.Bitmap, owned bool) *roaring.Bitmap {
792+
if prefix.IsEmpty() {
793+
if owned {
794+
return prefix
795+
}
796+
return roaring.NewBitmap()
797+
}
798+
// Either all planes are consumed (exactly one value remains and prefix matched
799+
// every one of its bits), or vals covers every bit pattern of the remaining
800+
// planes, which collapses dense value ranges to a single operation. In both
801+
// cases every column in prefix matches.
802+
if p < 0 || (p < 63 && uint64(len(vals)) == uint64(1)<<uint(p+1)) {
803+
if owned {
804+
return prefix
805+
}
806+
return prefix.Clone()
807+
}
808+
// vals is sorted and shares all bits above p, so bit p partitions it in place.
809+
cut := sort.Search(len(vals), func(i int) bool { return vals[i]&(uint64(1)<<uint(p)) != 0 })
810+
lo, hi := vals[:cut], vals[cut:]
811+
switch {
812+
case len(hi) == 0:
813+
return b.matchTrie(lo, p-1, planeChild(prefix, b.bA[p], false, owned), true)
814+
case len(lo) == 0:
815+
return b.matchTrie(hi, p-1, planeChild(prefix, b.bA[p], true, owned), true)
816+
default:
817+
// Materialize the set branch before the clear branch may consume prefix.
818+
hiBM := roaring.And(prefix, b.bA[p])
819+
result := b.matchTrie(lo, p-1, planeChild(prefix, b.bA[p], false, owned), true)
820+
result.Or(b.matchTrie(hi, p-1, hiBM, true))
821+
return result
822+
}
823+
}
824+
825+
// planeChild narrows prefix by one bitplane: prefix ∧ plane when set, prefix ∧ ¬plane
826+
// otherwise. When the caller owns prefix it is consumed in place to avoid a copy.
827+
func planeChild(prefix, plane *roaring.Bitmap, set, owned bool) *roaring.Bitmap {
828+
if owned {
829+
if set {
830+
prefix.And(plane)
831+
} else {
832+
prefix.AndNot(plane)
773833
}
834+
return prefix
774835
}
775-
resultsChan <- results
836+
if set {
837+
return roaring.And(prefix, plane)
838+
}
839+
return roaring.AndNot(prefix, plane)
776840
}
777841

778842
// ClearBits cleared the bits that exist in the target if they are also in the found set.
Lines changed: 234 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,234 @@
1+
package roaring
2+
3+
import (
4+
"math/rand"
5+
"testing"
6+
7+
"github.com/RoaringBitmap/roaring/v2"
8+
"github.com/stretchr/testify/assert"
9+
)
10+
11+
func BenchmarkBatchEqual(b *testing.B) {
12+
bsi := setupLargeBSI(b)
13+
if bsi == nil {
14+
b.Skip("skipping, large BSI setup failed")
15+
return
16+
}
17+
b.ResetTimer()
18+
for i := 0; i < b.N; i++ {
19+
res := bsi.BatchEqual(0, []int64{55, 57})
20+
_ = res
21+
}
22+
}
23+
24+
func TestBatchEqualEdgeCases(t *testing.T) {
25+
// 1. Empty or Nil inputs
26+
bsi := NewDefaultBSI()
27+
res := bsi.BatchEqual(0, nil)
28+
assert.True(t, res.IsEmpty())
29+
30+
res = bsi.BatchEqual(0, []int64{})
31+
assert.True(t, res.IsEmpty())
32+
33+
// 2. Set some values
34+
bsi.SetValue(10, 42)
35+
bsi.SetValue(20, 100)
36+
bsi.SetValue(30, 42)
37+
bsi.SetValue(40, -5)
38+
39+
// Test matching positive values
40+
res = bsi.BatchEqual(0, []int64{42})
41+
assert.Equal(t, uint64(2), res.GetCardinality())
42+
assert.True(t, res.Contains(10))
43+
assert.True(t, res.Contains(30))
44+
45+
// Test matching multiple values including non-existent and duplicates
46+
res = bsi.BatchEqual(0, []int64{42, 100, 42, 999})
47+
assert.Equal(t, uint64(3), res.GetCardinality())
48+
assert.True(t, res.Contains(10))
49+
assert.True(t, res.Contains(20))
50+
assert.True(t, res.Contains(30))
51+
52+
// Test negative value
53+
res = bsi.BatchEqual(0, []int64{-5})
54+
assert.Equal(t, uint64(1), res.GetCardinality())
55+
assert.True(t, res.Contains(40))
56+
57+
// Test 1<<62 edge case explicitly
58+
bsi62 := NewBSI(1<<62, 0)
59+
bsi62.SetValue(10, 5)
60+
res = bsi62.BatchEqual(0, []int64{5})
61+
assert.Equal(t, uint64(1), res.GetCardinality())
62+
assert.True(t, res.Contains(10))
63+
}
64+
65+
func TestBatchEqualSub64Bit(t *testing.T) {
66+
// NewBSI(100,0) (BitCount()==7) queried with {-5} and with {200} (>= 2^7)
67+
// must both return an empty bitmap and must equal a GetValue-derived ground truth.
68+
bsi := NewBSI(100, 0)
69+
assert.Equal(t, 7, bsi.BitCount())
70+
71+
// Set some values inside range
72+
bsi.SetValue(10, 42)
73+
bsi.SetValue(20, 99)
74+
75+
// Ground truth function
76+
getGroundTruth := func(query []int64) *roaring.Bitmap {
77+
expected := roaring.NewBitmap()
78+
valMap := make(map[int64]bool)
79+
for _, q := range query {
80+
valMap[q] = true
81+
}
82+
iter := bsi.GetExistenceBitmap().Iterator()
83+
for iter.HasNext() {
84+
col := iter.Next()
85+
val, ok := bsi.GetValue(uint64(col))
86+
if ok && valMap[val] {
87+
expected.Add(col)
88+
}
89+
}
90+
return expected
91+
}
92+
93+
for _, q := range []int64{-5, 200} {
94+
res := bsi.BatchEqual(0, []int64{q})
95+
assert.True(t, res.IsEmpty())
96+
97+
expected := getGroundTruth([]int64{q})
98+
assert.True(t, expected.IsEmpty())
99+
assert.True(t, res.Equals(expected))
100+
}
101+
}
102+
103+
func TestBatchEqualResultIsolation(t *testing.T) {
104+
bsi := NewDefaultBSI()
105+
bsi.SetValue(10, 42)
106+
bsi.SetValue(20, 100)
107+
108+
// Get batch equal result
109+
res := bsi.BatchEqual(0, []int64{42})
110+
assert.True(t, res.Contains(10))
111+
112+
// Mutate the returned bitmap
113+
res.Add(999)
114+
res.Remove(10)
115+
116+
// Assert that the source BSI's internal state (existence bitmap and bit planes) is completely unaffected
117+
assert.False(t, bsi.GetExistenceBitmap().Contains(999))
118+
assert.True(t, bsi.GetExistenceBitmap().Contains(10))
119+
120+
val, ok := bsi.GetValue(10)
121+
assert.True(t, ok)
122+
assert.Equal(t, int64(42), val)
123+
124+
val, ok = bsi.GetValue(999)
125+
assert.False(t, ok)
126+
}
127+
128+
func TestBatchEqualConsistentWithGetValue(t *testing.T) {
129+
rg := rand.New(rand.NewSource(42))
130+
for run := 0; run < 15; run++ {
131+
// Create a randomized BSI
132+
bsi := NewDefaultBSI()
133+
numCols := rg.Intn(1000) + 10
134+
for col := 0; col < numCols; col++ {
135+
if rg.Float64() < 0.8 {
136+
val := rg.Int63n(500) - 250 // Mix of positive, zero, and negative values
137+
bsi.SetValue(uint64(col), val)
138+
}
139+
}
140+
141+
// Generate query values (small, medium, and large list sizes to test the hybrid threshold)
142+
querySizes := []int{rg.Intn(10) + 1, rg.Intn(50) + 50, rg.Intn(200) + 100}
143+
for _, querySize := range querySizes {
144+
query := make([]int64, querySize)
145+
for i := range query {
146+
query[i] = rg.Int63n(600) - 300
147+
}
148+
149+
// Ground truth
150+
expected := roaring.NewBitmap()
151+
valMap := make(map[int64]bool)
152+
for _, q := range query {
153+
valMap[q] = true
154+
}
155+
iter := bsi.GetExistenceBitmap().Iterator()
156+
for iter.HasNext() {
157+
col := iter.Next()
158+
val, ok := bsi.GetValue(uint64(col))
159+
if ok && valMap[val] {
160+
expected.Add(col)
161+
}
162+
}
163+
164+
// Test different parallelism settings
165+
for _, parallelism := range []int{0, 1, 2, 4} {
166+
actual := bsi.BatchEqual(parallelism, query)
167+
if !actual.Equals(expected) {
168+
t.Fatalf("Mismatch in run %d querySize %d parallelism %d. Query: %v. Expected: %v, Got: %v", run, querySize, parallelism, query, expected.ToArray(), actual.ToArray())
169+
}
170+
}
171+
}
172+
}
173+
}
174+
175+
// TestBatchEqualExistenceAuthority pins BatchEqual results to the existence
176+
// bitmap. UnmarshalBinary accepts plane data that is not a subset of eBM (the
177+
// checked-in testdata/age fixture is such data), and every read path treats
178+
// eBM as authoritative, so columns present in a plane but absent from eBM must
179+
// never appear in results.
180+
func TestBatchEqualExistenceAuthority(t *testing.T) {
181+
// Synthetic state: column 2 has bits in plane 0 but is absent from eBM.
182+
ebm := roaring.BitmapOf(1)
183+
plane := roaring.BitmapOf(1, 2)
184+
ebmData, err := ebm.MarshalBinary()
185+
if err != nil {
186+
t.Fatal(err)
187+
}
188+
planeData, err := plane.MarshalBinary()
189+
if err != nil {
190+
t.Fatal(err)
191+
}
192+
bsi := NewDefaultBSI()
193+
if err := bsi.UnmarshalBinary([][]byte{ebmData, planeData}); err != nil {
194+
t.Fatal(err)
195+
}
196+
res := bsi.BatchEqual(0, []int64{1})
197+
assert.True(t, res.Contains(1))
198+
assert.False(t, res.Contains(2), "column 2 is not in eBM and must not match")
199+
200+
// The age fixture ships with plane cardinalities above the eBM cardinality;
201+
// results must still be a subset of eBM.
202+
large := setupLargeBSI(t)
203+
if large == nil {
204+
t.Skip("skipping, large BSI setup failed")
205+
}
206+
for _, vals := range [][]int64{{16}, {55, 57}, {0, 1, 2, 3}} {
207+
res := large.BatchEqual(0, vals)
208+
outside := roaring.AndNot(res, large.GetExistenceBitmap())
209+
assert.True(t, outside.IsEmpty(), "BatchEqual(%v) returned %d columns outside eBM", vals, outside.GetCardinality())
210+
}
211+
}
212+
213+
// Benchmarks across query-list shapes: work sharing behaves differently for
214+
// small lists, dense contiguous ranges (which collapse to a few plane
215+
// operations), and scattered values (no range collapse).
216+
func BenchmarkBatchEqualM128(b *testing.B) { benchmarkBatchEqualM(b, 128, 1) }
217+
func BenchmarkBatchEqualM128Scattered(b *testing.B) { benchmarkBatchEqualM(b, 128, 2) }
218+
func BenchmarkBatchEqualM200(b *testing.B) { benchmarkBatchEqualM(b, 200, 1) }
219+
220+
func benchmarkBatchEqualM(b *testing.B, m int, stride int64) {
221+
bsi := setupLargeBSI(b)
222+
if bsi == nil {
223+
b.Skip("skipping, large BSI setup failed")
224+
}
225+
vals := make([]int64, m)
226+
for i := range vals {
227+
vals[i] = int64(i)*stride + stride - 1
228+
}
229+
b.ResetTimer()
230+
for i := 0; i < b.N; i++ {
231+
res := bsi.BatchEqual(0, vals)
232+
_ = res
233+
}
234+
}

0 commit comments

Comments
 (0)