Skip to content

Commit 4cc8e21

Browse files
guymolinariguymolinari
committed
Optimize roaring64 BSI BatchEqual
1 parent 44559cd commit 4cc8e21

2 files changed

Lines changed: 317 additions & 7 deletions

File tree

‎roaring64/bsi64.go‎

Lines changed: 119 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@ import (
55
"io"
66
"math/big"
77
"runtime"
8+
"sort"
89
"sync"
910
)
1011

@@ -970,25 +971,136 @@ func (b *BSI) WriteTo(w io.Writer) (n int64, err error) {
970971

971972
// BatchEqual returns a bitmap containing the column IDs where the values are contained within the list of values provided.
972973
func (b *BSI) BatchEqual(parallelism int, values []int64) *Bitmap {
973-
//convert list of int64 values to big.Int(s)
974-
bigValues := make([]*big.Int, len(values))
975-
for i, v := range values {
976-
bigValues[i] = big.NewInt(v)
974+
if b.eBM.IsEmpty() || len(values) == 0 {
975+
return NewBitmap()
977976
}
978-
return b.BatchEqualBig(parallelism, bigValues)
977+
978+
bitCount := b.BitCount()
979+
if bitCount >= 64 {
980+
// Fall back to the arbitrary-precision path when the BSI has more than
981+
// int64's finite bit width. This preserves correctness for big-value BSIs.
982+
bigValues := make([]*big.Int, len(values))
983+
for i, v := range values {
984+
bigValues[i] = big.NewInt(v)
985+
}
986+
return b.BatchEqualBig(parallelism, bigValues)
987+
}
988+
989+
seen := make(map[uint64]struct{}, len(values))
990+
vals := make([]uint64, 0, len(values))
991+
for _, v := range values {
992+
if !bsi64ValueFitsBitCount(v, bitCount) {
993+
continue
994+
}
995+
encoded := encodeBSI64Value(v, bitCount)
996+
if _, ok := seen[encoded]; ok {
997+
continue
998+
}
999+
seen[encoded] = struct{}{}
1000+
vals = append(vals, encoded)
1001+
}
1002+
if len(vals) == 0 {
1003+
return NewBitmap()
1004+
}
1005+
1006+
sort.Slice(vals, func(i, j int) bool { return vals[i] < vals[j] })
1007+
result := b.matchInt64Trie(vals, bitCount, &b.eBM, false)
1008+
if b.runOptimized {
1009+
result.RunOptimize()
1010+
}
1011+
return result
1012+
}
1013+
1014+
func bsi64ValueFitsBitCount(value int64, bitCount int) bool {
1015+
if bitCount >= 63 {
1016+
return true
1017+
}
1018+
min := -(int64(1) << uint(bitCount))
1019+
max := (int64(1) << uint(bitCount)) - 1
1020+
return value >= min && value <= max
1021+
}
1022+
1023+
func encodeBSI64Value(value int64, bitCount int) uint64 {
1024+
if bitCount >= 63 {
1025+
return uint64(value)
1026+
}
1027+
mask := (uint64(1) << uint(bitCount+1)) - 1
1028+
return uint64(value) & mask
1029+
}
1030+
1031+
func (b *BSI) matchInt64Trie(vals []uint64, p int, prefix *Bitmap, owned bool) *Bitmap {
1032+
if prefix.IsEmpty() {
1033+
if owned {
1034+
return prefix
1035+
}
1036+
return NewBitmap()
1037+
}
1038+
if p < 0 || (p < 63 && uint64(len(vals)) == uint64(1)<<uint(p+1)) {
1039+
if owned {
1040+
return prefix
1041+
}
1042+
return prefix.Clone()
1043+
}
1044+
1045+
mask := uint64(1) << uint(p)
1046+
cut := sort.Search(len(vals), func(i int) bool { return vals[i]&mask != 0 })
1047+
lo, hi := vals[:cut], vals[cut:]
1048+
switch {
1049+
case len(hi) == 0:
1050+
return b.matchInt64Trie(lo, p-1, bsi64PlaneChild(prefix, &b.bA[p], false, owned), true)
1051+
case len(lo) == 0:
1052+
return b.matchInt64Trie(hi, p-1, bsi64PlaneChild(prefix, &b.bA[p], true, owned), true)
1053+
default:
1054+
hiBM := And(prefix, &b.bA[p])
1055+
result := b.matchInt64Trie(lo, p-1, bsi64PlaneChild(prefix, &b.bA[p], false, owned), true)
1056+
result.Or(b.matchInt64Trie(hi, p-1, hiBM, true))
1057+
return result
1058+
}
1059+
}
1060+
1061+
func bsi64PlaneChild(prefix, plane *Bitmap, set, owned bool) *Bitmap {
1062+
if owned {
1063+
if set {
1064+
prefix.And(plane)
1065+
} else {
1066+
prefix.AndNot(plane)
1067+
}
1068+
return prefix
1069+
}
1070+
if set {
1071+
return And(prefix, plane)
1072+
}
1073+
return AndNot(prefix, plane)
9791074
}
9801075

9811076
// BatchEqualBig returns a bitmap containing the column IDs where the values are contained within the list of values provided.
9821077
func (b *BSI) BatchEqualBig(parallelism int, values []*big.Int) *Bitmap {
1078+
if b.eBM.IsEmpty() || len(values) == 0 {
1079+
return NewBitmap()
1080+
}
9831081

9841082
valMap := make(map[string]struct{}, len(values))
9851083
for i := 0; i < len(values); i++ {
986-
valMap[string(values[i].Bytes())] = struct{}{}
1084+
if values[i] == nil {
1085+
continue
1086+
}
1087+
valMap[batchEqualBigKey(values[i])] = struct{}{}
1088+
}
1089+
if len(valMap) == 0 {
1090+
return NewBitmap()
9871091
}
9881092
comp := &task{bsi: b, values: valMap}
9891093
return parallelExecutor(parallelism, comp, batchEqual, &b.eBM)
9901094
}
9911095

1096+
func batchEqualBigKey(value *big.Int) string {
1097+
bytes := value.Bytes()
1098+
key := make([]byte, len(bytes)+1)
1099+
key[0] = byte(value.Sign() + 1)
1100+
copy(key[1:], bytes)
1101+
return string(key)
1102+
}
1103+
9921104
func batchEqual(e *task, batch []uint64, resultsChan chan *Bitmap,
9931105
wg *sync.WaitGroup) {
9941106

@@ -1002,7 +1114,7 @@ func batchEqual(e *task, batch []uint64, resultsChan chan *Bitmap,
10021114
for i := 0; i < len(batch); i++ {
10031115
cID := batch[i]
10041116
if value, ok := e.bsi.GetBigValue(cID); ok {
1005-
if _, yes := e.values[string(value.Bytes())]; yes {
1117+
if _, yes := e.values[batchEqualBigKey(value)]; yes {
10061118
results.Add(cID)
10071119
}
10081120
}

‎roaring64/bsi64_batch_equal_test.go‎

Lines changed: 198 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,198 @@
1+
package roaring64
2+
3+
import (
4+
"math/rand"
5+
"testing"
6+
7+
"github.com/stretchr/testify/assert"
8+
)
9+
10+
func expectedBSI64BatchEqual(bsi *BSI, query []int64) *Bitmap {
11+
expected := NewBitmap()
12+
want := make(map[int64]struct{}, len(query))
13+
for _, q := range query {
14+
want[q] = struct{}{}
15+
}
16+
iter := bsi.GetExistenceBitmap().Iterator()
17+
for iter.HasNext() {
18+
col := iter.Next()
19+
val, ok := bsi.GetValue(col)
20+
if ok {
21+
if _, hit := want[val]; hit {
22+
expected.Add(col)
23+
}
24+
}
25+
}
26+
return expected
27+
}
28+
29+
func TestBSI64BatchEqualEdgeCases(t *testing.T) {
30+
bsi := NewDefaultBSI()
31+
res := bsi.BatchEqual(0, nil)
32+
assert.True(t, res.IsEmpty())
33+
34+
res = bsi.BatchEqual(0, []int64{})
35+
assert.True(t, res.IsEmpty())
36+
37+
bsi.SetValue(10, 42)
38+
bsi.SetValue(20, 100)
39+
bsi.SetValue(30, 42)
40+
bsi.SetValue(40, -5)
41+
bsi.SetValue(50, 5)
42+
43+
res = bsi.BatchEqual(0, []int64{42})
44+
assert.Equal(t, uint64(2), res.GetCardinality())
45+
assert.True(t, res.Contains(10))
46+
assert.True(t, res.Contains(30))
47+
48+
res = bsi.BatchEqual(0, []int64{42, 100, 42, 999})
49+
assert.Equal(t, uint64(3), res.GetCardinality())
50+
assert.True(t, res.Contains(10))
51+
assert.True(t, res.Contains(20))
52+
assert.True(t, res.Contains(30))
53+
54+
res = bsi.BatchEqual(0, []int64{-5})
55+
assert.Equal(t, uint64(1), res.GetCardinality())
56+
assert.True(t, res.Contains(40))
57+
assert.False(t, res.Contains(50), "negative and positive values with the same magnitude must not collide")
58+
59+
res = bsi.BatchEqual(0, []int64{5})
60+
assert.Equal(t, uint64(1), res.GetCardinality())
61+
assert.True(t, res.Contains(50))
62+
assert.False(t, res.Contains(40), "positive and negative values with the same magnitude must not collide")
63+
64+
bsi62 := NewBSI(1<<62, 0)
65+
bsi62.SetValue(10, 5)
66+
res = bsi62.BatchEqual(0, []int64{5})
67+
assert.Equal(t, uint64(1), res.GetCardinality())
68+
assert.True(t, res.Contains(10))
69+
}
70+
71+
func TestBSI64BatchEqualSubBitWidthMatchesGetValue(t *testing.T) {
72+
bsi := NewBSI(100, 0)
73+
assert.Equal(t, 7, bsi.BitCount())
74+
75+
bsi.SetValue(10, 42)
76+
bsi.SetValue(20, 99)
77+
78+
for _, query := range [][]int64{{-5}, {200}, {-5, 42, 200}} {
79+
expected := expectedBSI64BatchEqual(bsi, query)
80+
actual := bsi.BatchEqual(0, query)
81+
assert.True(t, actual.Equals(expected), "query %v expected %v got %v", query, expected.ToArray(), actual.ToArray())
82+
}
83+
}
84+
85+
func TestBSI64BatchEqualResultIsolation(t *testing.T) {
86+
bsi := NewDefaultBSI()
87+
bsi.SetValue(10, 42)
88+
bsi.SetValue(20, 100)
89+
90+
res := bsi.BatchEqual(0, []int64{42})
91+
assert.True(t, res.Contains(10))
92+
93+
res.Add(999)
94+
res.Remove(10)
95+
96+
assert.False(t, bsi.GetExistenceBitmap().Contains(999))
97+
assert.True(t, bsi.GetExistenceBitmap().Contains(10))
98+
99+
val, ok := bsi.GetValue(10)
100+
assert.True(t, ok)
101+
assert.Equal(t, int64(42), val)
102+
103+
_, ok = bsi.GetValue(999)
104+
assert.False(t, ok)
105+
}
106+
107+
func TestBSI64BatchEqualConsistentWithGetValue(t *testing.T) {
108+
rg := rand.New(rand.NewSource(42))
109+
for run := 0; run < 15; run++ {
110+
bsi := NewDefaultBSI()
111+
numCols := rg.Intn(1000) + 10
112+
for col := 0; col < numCols; col++ {
113+
if rg.Float64() < 0.8 {
114+
val := rg.Int63n(500) - 250
115+
bsi.SetValue(uint64(col), val)
116+
}
117+
}
118+
119+
querySizes := []int{rg.Intn(10) + 1, rg.Intn(50) + 50, rg.Intn(200) + 100}
120+
for _, querySize := range querySizes {
121+
query := make([]int64, querySize)
122+
for i := range query {
123+
query[i] = rg.Int63n(600) - 300
124+
}
125+
expected := expectedBSI64BatchEqual(bsi, query)
126+
127+
for _, parallelism := range []int{0, 1, 2, 4} {
128+
actual := bsi.BatchEqual(parallelism, query)
129+
if !actual.Equals(expected) {
130+
t.Fatalf("run=%d querySize=%d parallelism=%d query=%v expected=%v actual=%v",
131+
run, querySize, parallelism, query, expected.ToArray(), actual.ToArray())
132+
}
133+
}
134+
}
135+
}
136+
}
137+
138+
func TestBSI64BatchEqualExistenceAuthority(t *testing.T) {
139+
ebm := BitmapOf(1)
140+
plane := BitmapOf(1, 2)
141+
ebmData, err := ebm.MarshalBinary()
142+
if err != nil {
143+
t.Fatal(err)
144+
}
145+
planeData, err := plane.MarshalBinary()
146+
if err != nil {
147+
t.Fatal(err)
148+
}
149+
bsi := NewDefaultBSI()
150+
if err := bsi.UnmarshalBinary([][]byte{ebmData, planeData}); err != nil {
151+
t.Fatal(err)
152+
}
153+
res := bsi.BatchEqual(0, []int64{1})
154+
assert.True(t, res.Contains(1))
155+
assert.False(t, res.Contains(2), "column 2 is not in eBM and must not match")
156+
157+
large := setupLargeBSI(t)
158+
if large == nil {
159+
t.Skip("skipping, large BSI setup failed")
160+
}
161+
for _, vals := range [][]int64{{16}, {55, 57}, {0, 1, 2, 3}} {
162+
res := large.BatchEqual(0, vals)
163+
outside := AndNot(res, large.GetExistenceBitmap())
164+
assert.True(t, outside.IsEmpty(), "BatchEqual(%v) returned %d columns outside eBM", vals, outside.GetCardinality())
165+
}
166+
}
167+
168+
func BenchmarkBSI64BatchEqualLargeAgeFixture(b *testing.B) {
169+
bsi := setupLargeBSI(b)
170+
if bsi == nil {
171+
b.Skip("skipping, large BSI setup failed")
172+
}
173+
b.ResetTimer()
174+
for i := 0; i < b.N; i++ {
175+
res := bsi.BatchEqual(0, []int64{55, 57})
176+
_ = res
177+
}
178+
}
179+
180+
func BenchmarkBSI64BatchEqualM128(b *testing.B) { benchmarkBSI64BatchEqualM(b, 128, 1) }
181+
func BenchmarkBSI64BatchEqualM128Scattered(b *testing.B) { benchmarkBSI64BatchEqualM(b, 128, 2) }
182+
func BenchmarkBSI64BatchEqualM200(b *testing.B) { benchmarkBSI64BatchEqualM(b, 200, 1) }
183+
184+
func benchmarkBSI64BatchEqualM(b *testing.B, m int, stride int64) {
185+
bsi := setupLargeBSI(b)
186+
if bsi == nil {
187+
b.Skip("skipping, large BSI setup failed")
188+
}
189+
vals := make([]int64, m)
190+
for i := range vals {
191+
vals[i] = int64(i)*stride + stride - 1
192+
}
193+
b.ResetTimer()
194+
for i := 0; i < b.N; i++ {
195+
res := bsi.BatchEqual(0, vals)
196+
_ = res
197+
}
198+
}

0 commit comments

Comments
 (0)