Implement shotgun4Intersect

Author: alldroll

setutil.go

@@ -582,6 +582,104 @@ mainwhile:
 	return pos
 }
 
+// shotgun4Intersect performs intersection between small and large arrays described in
+// https://lemire.me/blog/2019/01/16/faster-intersections-between-sorted-arrays-with-shotgun/
+func shotgun4Intersect(small, large, buf []uint16) int {
+	if len(small) == 0 {
+		return 0
+	}
+
+	nS, nL := len(small), len(large)
+	buf = buf[:cap(buf)]
+	idxS, idxL := 0, 0
+	pos := 0
+
+	for (idxS+4 <= nS) && idxL < nL {
+		t1, t2, t3, t4 := small[idxS], small[idxS+1], small[idxS+2], small[idxS+3]
+		idx1, idx2, idx3, idx4 := idxL, idxL, idxL, idxL
+		n := nL - idxL
+
+		for n > 1 {
+			m := n >> 1
+
+			if large[idx1+m] < t1 {
+				idx1 += m
+			}
+
+			if large[idx2+m] < t2 {
+				idx2 += m
+			}
+
+			if large[idx3+m] < t3 {
+				idx3 += m
+			}
+
+			if large[idx4+m] < t4 {
+				idx4 += m
+			}
+
+			n -= m
+		}
+
+		if large[idx1] < t1 {
+			idx1++
+		}
+
+		if large[idx2] < t2 {
+			idx2++
+		}
+
+		if large[idx3] < t3 {
+			idx3++
+		}
+
+		if large[idx4] < t4 {
+			idx4++
+		}
+
+		if idx1 < nL && large[idx1] == t1 {
+			buf[pos] = t1
+			pos++
+		}
+
+		if idx2 < nL && large[idx2] == t2 {
+			buf[pos] = t2
+			pos++
+		}
+
+		if idx3 < nL && large[idx3] == t3 {
+			buf[pos] = t3
+			pos++
+		}
+
+		if idx4 < nL && large[idx4] == t4 {
+			buf[pos] = t4
+			pos++
+		}
+
+		idxS += 4
+		idxL = idx4
+	}
+
+	for idxS < nS && idxL < nL {
+		s := small[idxS]
+		idxL = advanceUntil(large, idxL, nL, s)
+
+		if idxL == nL {
+			break
+		}
+
+		if large[idxL] == s {
+			buf[pos] = s
+			pos++
+		}
+
+		idxS++
+	}
+
+	return pos
+}
+
 func binarySearch(array []uint16, ikey uint16) int {
 	low := 0
 	high := len(array) - 1

setutil_test.go

@@ -4,6 +4,7 @@ package roaring
 
 import (
 	"github.com/stretchr/testify/assert"
+	"math/rand"
 	"testing"
 )
 
@@ -92,16 +93,31 @@ func TestSetUtilIntersection(t *testing.T) {
 	assert.Equal(t, expectedresult, result)
 }
 
-func TestSetUtilIntersection2(t *testing.T) {
-	data1 := []uint16{0, 2, 4, 6, 8, 10, 12, 14, 16, 18}
-	data2 := []uint16{0, 3, 6, 9, 12, 15, 18}
-	result := make([]uint16, 0, len(data1)+len(data2))
-	expectedresult := []uint16{0, 6, 12, 18}
-	nl := intersection2by2(data1, data2, result)
-	result = result[:nl]
-	result = result[:len(expectedresult)]
+func TestSetUtilIntersectionCases(t *testing.T) {
+	cases := []struct {
+		name string
+		algo func(a, b, buf []uint16) int
+	}{
+		{
+			name: "onesidedgallopingintersect2by2",
+			algo: onesidedgallopingintersect2by2,
+		},
+		{
+			name: "shotgun4Intersect",
+			algo: shotgun4Intersect,
+		},
+	}
 
-	assert.Equal(t, expectedresult, result)
+	data1 := []uint16{0, 3, 6, 9, 12, 15, 18}
+	data2 := []uint16{0, 2, 4, 6, 8, 10, 12, 14, 16, 18}
+	expected := []uint16{0, 6, 12, 18}
+
+	for _, c := range cases {
+		result := make([]uint16, 0, len(data1)+len(data2))
+		n := c.algo(data1, data2, result)
+
+		assert.Equalf(t, expected, result[:n], "failed algorithm: %s", c.name)
+	}
 }
 
 func TestSetUtilBinarySearch(t *testing.T) {
@@ -119,3 +135,48 @@ func TestSetUtilBinarySearch(t *testing.T) {
 		}
 	}
 }
+
+func BenchmarkIntersectAlgorithms(b *testing.B) {
+	sz1 := 1000
+	s1 := make([]uint16, sz1)
+
+	sz2 := MaxUint16
+	s2 := make([]uint16, sz2)
+
+	for i := 0; i < sz2; i++ {
+		s2[i] = uint16(i)
+	}
+
+	r := rand.New(rand.NewSource(0))
+	k := 0
+
+	for i := 0; i < sz1 && k < sz2; i++ {
+		n := r.Intn(100)
+		k += n
+
+		// prevent adding duplicates
+		if n == 0 && i > 0 {
+			k++
+		}
+
+		s1[i] = uint16(s2[k])
+	}
+
+	buf := make([]uint16, sz1+sz2)
+
+	b.Run("onesidedgallopingintersect2by2", func(b *testing.B) {
+		b.ResetTimer()
+
+		for i := 0; i < b.N; i++ {
+			onesidedgallopingintersect2by2(s1, s2, buf)
+		}
+	})
+
+	b.Run("shotgun4", func(b *testing.B) {
+		b.ResetTimer()
+
+		for i := 0; i < b.N; i++ {
+			shotgun4Intersect(s1, s2, buf)
+		}
+	})
+}