Optimize populated run-optimized BSI query results

BitSliceIndexing/bsi.go

@@ -97,9 +97,6 @@ func (b *BSI) SetValue(columnID uint64, value int64) {
 	if b.MaxValue == 0 && b.MinValue == 0 {
 		for i := bits.Len64(uint64(value)) - b.BitCount(); i > 0; i-- {
 			b.bA = append(b.bA, roaring.NewBitmap())
-			if b.runOptimized {
-				b.bA[i].RunOptimize()
-			}
 		}
 	}
 
@@ -121,9 +118,6 @@ func (b *BSI) SetMany(foundSet *roaring.Bitmap, value int64) {
 	if b.MaxValue == 0 && b.MinValue == 0 {
 		for i := bits.Len64(uint64(value)) - b.BitCount(); i > 0; i-- {
 			b.bA = append(b.bA, roaring.NewBitmap())
-			if b.runOptimized {
-				b.bA[i].RunOptimize()
-			}
 		}
 	}
 
@@ -191,7 +185,12 @@ func parallelExecutor(parallelism int, t *task, e action,
 		ba = append(ba, bm)
 	}
 
-	return roaring.ParOr(0, ba...)
+	results := roaring.ParOr(0, ba...)
+	// Optimize the aggregate returned to the caller after it has been populated.
+	if t.bsi.runOptimized && !results.IsEmpty() {
+		results.RunOptimize()
+	}
+	return results
 
 }
 
@@ -241,6 +240,10 @@ func parallelExecutorBSIResults(parallelism int, input *BSI, e bsiAction, foundS
 	} else {
 		results.ParOr(0, ba...)
 	}
+	// Optimize the aggregate returned to the caller after it has been populated.
+	if input.runOptimized && !results.eBM.IsEmpty() {
+		results.RunOptimize()
+	}
 	return results
 
 }
@@ -297,9 +300,6 @@ func compareValue(e *task, batch []uint32, resultsChan chan *roaring.Bitmap, wg
 	defer wg.Done()
 
 	results := roaring.NewBitmap()
-	if e.bsi.runOptimized {
-		results.RunOptimize()
-	}
 	if len(batch) == 0 {
 		resultsChan <- results
 		return
@@ -610,9 +610,6 @@ func transpose(e *task, batch []uint32, resultsChan chan *roaring.Bitmap, wg *sy
 	defer wg.Done()
 
 	results := roaring.NewBitmap()
-	if e.bsi.runOptimized {
-		results.RunOptimize()
-	}
 	for _, cID := range batch {
 		if value, ok := e.bsi.GetValue(uint64(cID)); ok {
 			results.Add(uint32(value))
@@ -1031,9 +1028,6 @@ func transposeWithCounts(input *BSI, batch []uint32, resultsChan chan *BSI, wg *
 	defer wg.Done()
 
 	results := NewDefaultBSI()
-	if input.runOptimized {
-		results.RunOptimize()
-	}
 	for _, cID := range batch {
 		if value, ok := input.GetValue(uint64(cID)); ok {
 			if val, ok2 := results.GetValue(uint64(value)); !ok2 {

BitSliceIndexing/bsi_benchmark_test.go

@@ -21,6 +21,64 @@ func BenchmarkBatchEqual(b *testing.B) {
 	}
 }
 
+const runOptimizedBenchmarkResultCardinality = 50_000
+
+func setupRunOptimizedBenchmark(b testing.TB) *BSI {
+	return setupRunOptimizedQueryBSI(b, runOptimizedBenchmarkResultCardinality, func(int) int64 { return 1 })
+}
+
+func BenchmarkBatchEqualRunOptimizedEndToEnd(b *testing.B) {
+	bsi := setupRunOptimizedBenchmark(b)
+	values := []int64{1}
+	var result *roaring.Bitmap
+
+	b.ResetTimer()
+	for b.Loop() {
+		result = bsi.BatchEqual(0, values)
+	}
+
+	if result.GetCardinality() != runOptimizedBenchmarkResultCardinality {
+		b.Fatalf("BatchEqual cardinality = %d, want %d", result.GetCardinality(), runOptimizedBenchmarkResultCardinality)
+	}
+}
+
+func BenchmarkCompareValueRunOptimizedEndToEnd(b *testing.B) {
+	bsi := setupRunOptimizedBenchmark(b)
+	var result *roaring.Bitmap
+
+	b.ResetTimer()
+	for b.Loop() {
+		result = bsi.CompareValue(4, EQ, 1, 0, nil)
+	}
+
+	if result.GetCardinality() != runOptimizedBenchmarkResultCardinality {
+		b.Fatalf("CompareValue cardinality = %d, want %d", result.GetCardinality(), runOptimizedBenchmarkResultCardinality)
+	}
+}
+
+func BenchmarkCompareValueRunOptimizedResultSerialization(b *testing.B) {
+	bsi := setupRunOptimizedBenchmark(b)
+	result := bsi.CompareValue(4, EQ, 1, 0, nil)
+	if result.GetCardinality() != runOptimizedBenchmarkResultCardinality {
+		b.Fatalf("CompareValue cardinality = %d, want %d", result.GetCardinality(), runOptimizedBenchmarkResultCardinality)
+	}
+
+	var encoded []byte
+	b.ReportAllocs()
+	b.ResetTimer()
+	for b.Loop() {
+		var err error
+		encoded, err = result.MarshalBinary()
+		if err != nil {
+			b.Fatal(err)
+		}
+	}
+
+	if len(encoded) == 0 {
+		b.Fatal("MarshalBinary returned an empty result")
+	}
+}
+
 func TestBatchEqualEdgeCases(t *testing.T) {
 	// 1. Empty or Nil inputs
 	bsi := NewDefaultBSI()

BitSliceIndexing/bsi_test.go

@@ -127,6 +127,19 @@ func setupRandom() *BSI {
 	return bsi
 }
 
+const runOptimizedQueryResultCardinality = 10_000
+
+func setupRunOptimizedQueryBSI(t testing.TB, cardinality int, valueForColumn func(int) int64) *BSI {
+	t.Helper()
+
+	bsi := NewDefaultBSI()
+	for columnID := 0; columnID < cardinality; columnID++ {
+		bsi.SetValue(uint64(columnID), valueForColumn(columnID))
+	}
+	bsi.RunOptimize()
+	return bsi
+}
+
 func TestEQ(t *testing.T) {
 	bsi := setup()
 	eq := bsi.CompareValue(0, EQ, 50, 0, nil)
@@ -365,6 +378,91 @@ func TestTransposeWithCounts(t *testing.T) {
 	assert.Equal(t, int64(2), a)
 }
 
+func TestRunOptimizedBitmapQueryResults(t *testing.T) {
+	expected := roaring.NewBitmap()
+	expected.AddRange(0, runOptimizedQueryResultCardinality)
+
+	constantValues := setupRunOptimizedQueryBSI(t, runOptimizedQueryResultCardinality, func(int) int64 { return 1 })
+	sequentialValues := setupRunOptimizedQueryBSI(t, runOptimizedQueryResultCardinality, func(columnID int) int64 { return int64(columnID) })
+	results := []struct {
+		name   string
+		result *roaring.Bitmap
+	}{
+		{"BatchEqual", constantValues.BatchEqual(4, []int64{1})},
+		{"CompareValue", constantValues.CompareValue(4, EQ, 1, 0, nil)},
+		{"Transpose", sequentialValues.IntersectAndTranspose(4, nil)},
+	}
+
+	for _, test := range results {
+		t.Run(test.name, func(t *testing.T) {
+			assert.True(t, test.result.Equals(expected))
+			assert.True(t, test.result.HasRunCompression())
+		})
+	}
+}
+
+func TestRunOptimizedTransposeWithCountsResult(t *testing.T) {
+	input := setupRunOptimizedQueryBSI(t, runOptimizedQueryResultCardinality, func(columnID int) int64 { return int64(columnID) })
+	result := input.TransposeWithCounts(4, nil)
+
+	assert.Equal(t, uint64(runOptimizedQueryResultCardinality), result.GetCardinality())
+	assert.True(t, result.HasRunCompression())
+	assert.True(t, result.GetExistenceBitmap().HasRunCompression())
+	require.Len(t, result.bA, 1)
+	assert.True(t, result.bA[0].HasRunCompression())
+	for _, columnID := range []uint64{0, 1, runOptimizedQueryResultCardinality - 1} {
+		value, exists := result.GetValue(columnID)
+		assert.True(t, exists)
+		assert.Equal(t, int64(1), value)
+	}
+}
+
+func TestRunOptimizedTransposeWithCountsEmptyResultIsMutable(t *testing.T) {
+	input := NewDefaultBSI()
+	input.SetValue(0, 1)
+	input.RunOptimize()
+
+	result := input.TransposeWithCounts(4, roaring.NewBitmap())
+	assert.Zero(t, result.GetCardinality())
+	assert.False(t, result.HasRunCompression())
+
+	result.SetValue(0, 1)
+	value, exists := result.GetValue(0)
+	assert.True(t, exists)
+	assert.Equal(t, int64(1), value)
+}
+
+func TestRunOptimizedBSISettersCanGrowBitSlices(t *testing.T) {
+	const expandedValue int64 = 1 << 10
+
+	t.Run("SetValue", func(t *testing.T) {
+		bsi := NewDefaultBSI()
+		bsi.SetValue(0, 1)
+		bsi.RunOptimize()
+		bsi.SetValue(1, expandedValue)
+
+		value, exists := bsi.GetValue(1)
+		assert.True(t, exists)
+		assert.Equal(t, expandedValue, value)
+		assert.Equal(t, 11, bsi.BitCount())
+	})
+
+	t.Run("SetMany", func(t *testing.T) {
+		bsi := NewDefaultBSI()
+		bsi.SetValue(0, 1)
+		bsi.RunOptimize()
+		foundSet := roaring.BitmapOf(1, 2)
+		bsi.SetMany(foundSet, expandedValue)
+
+		for _, columnID := range []uint64{1, 2} {
+			value, exists := bsi.GetValue(columnID)
+			assert.True(t, exists)
+			assert.Equal(t, expandedValue, value)
+		}
+		assert.Equal(t, 11, bsi.BitCount())
+	})
+}
+
 func TestRangeAllNegative(t *testing.T) {
 	bsi := setupAllNegative()
 	assert.Equal(t, uint64(100), bsi.GetCardinality())