pre filter for ivf (#23117)

pre filter for ivf

Approved by: @heni02, @ouyuanning, @daviszhen, @zhangxu19830126, @gouhongshen, @LeftHandCold, @fengttt, @aunjgr, @XuPeng-SH

Author: iamlinjunhong

pkg/common/bloomfilter/bloomfilter.go

@@ -15,8 +15,13 @@
 package bloomfilter
 
 import (
+	"bytes"
+
+	"github.com/matrixorigin/matrixone/pkg/common/bitmap"
 	"github.com/matrixorigin/matrixone/pkg/common/hashmap"
+	"github.com/matrixorigin/matrixone/pkg/common/moerr"
 	"github.com/matrixorigin/matrixone/pkg/container/hashtable"
+	"github.com/matrixorigin/matrixone/pkg/container/types"
 	"github.com/matrixorigin/matrixone/pkg/container/vector"
 )
 
@@ -114,6 +119,81 @@ func (bf *BloomFilter) TestAndAdd(v *vector.Vector, callBack func(bool, int)) {
 
 }
 
+// Marshal encodes BloomFilter into byte sequence for transmission via runtime filter message within the same CN.
+// Encoding format:
+//
+//	[seedCount:uint32][seeds...:uint64][bitmapLen:uint32][bitmapBytes...]
+func (bf *BloomFilter) Marshal() ([]byte, error) {
+	var buf bytes.Buffer
+
+	seedCount := uint32(len(bf.hashSeed))
+	buf.Write(types.EncodeUint32(&seedCount))
+	for i := 0; i < int(seedCount); i++ {
+		buf.Write(types.EncodeUint64(&bf.hashSeed[i]))
+	}
+
+	bmBytes := bf.bitmap.Marshal()
+	bmLen := uint32(len(bmBytes))
+	buf.Write(types.EncodeUint32(&bmLen))
+	buf.Write(bmBytes)
+
+	return buf.Bytes(), nil
+}
+
+// Unmarshal restores BloomFilter from byte sequence.
+// Initializes internal structures (keys / states / vals / addVals) based on encoded seedCount and bitmap.
+func (bf *BloomFilter) Unmarshal(data []byte) error {
+	if len(data) < 4 {
+		return moerr.NewInternalErrorNoCtx("invalid bloomfilter data")
+	}
+
+	seedCount := int(types.DecodeUint32(data[:4]))
+	data = data[4:]
+
+	if seedCount <= 0 {
+		return moerr.NewInternalErrorNoCtx("invalid bloomfilter seed count")
+	}
+
+	hashSeed := make([]uint64, seedCount)
+	for i := 0; i < seedCount; i++ {
+		if len(data) < 8 {
+			return moerr.NewInternalErrorNoCtx("invalid bloomfilter data (seed truncated)")
+		}
+		hashSeed[i] = types.DecodeUint64(data[:8])
+		data = data[8:]
+	}
+
+	if len(data) < 4 {
+		return moerr.NewInternalErrorNoCtx("invalid bloomfilter data (no bitmap length)")
+	}
+	bmLen := int(types.DecodeUint32(data[:4]))
+	data = data[4:]
+	if bmLen < 0 || len(data) < bmLen {
+		return moerr.NewInternalErrorNoCtx("invalid bloomfilter data (bitmap truncated)")
+	}
+
+	var bm bitmap.Bitmap
+	bm.Unmarshal(data[:bmLen])
+
+	// Reinitialize internal auxiliary structures for subsequent Test/TestAndAdd
+	vals := make([][]uint64, hashmap.UnitLimit)
+	keys := make([][]byte, hashmap.UnitLimit)
+	states := make([][3]uint64, hashmap.UnitLimit)
+	for j := 0; j < hashmap.UnitLimit; j++ {
+		vals[j] = make([]uint64, seedCount*3)
+	}
+
+	bf.bitmap = bm
+	bf.hashSeed = hashSeed
+	bf.keys = keys
+	bf.states = states
+	bf.vals = vals
+	bf.addVals = make([]uint64, hashmap.UnitLimit*3*seedCount)
+	bf.valLength = len(hashSeed) * 3
+
+	return nil
+}
+
 // for an incoming vector, compute the hash value of each of its elements, and manipulate it with func tf.fn
 func (bf *BloomFilter) handle(v *vector.Vector, callBack func(int, int)) {
 	length := v.Length()

pkg/common/bloomfilter/bloomfilter_test.go

@@ -111,3 +111,157 @@ func BenchmarkBloomFiltrerTestAndAdd(b *testing.B) {
 		}
 	}
 }
+
+func TestMarshal(t *testing.T) {
+	mp := mpool.MustNewZero()
+	vec := testutil.NewVector(testCount, types.New(types.T_int64, 0, 0), mp, false, nil)
+	defer vec.Free(mp)
+
+	bf := New(testCount, testRate)
+	bf.Add(vec)
+
+	data, err := bf.Marshal()
+	require.NoError(t, err)
+	require.NotNil(t, data)
+	require.Greater(t, len(data), 0)
+
+	// Verify basic structure: should have at least seedCount (4 bytes) + some seeds + bitmapLen (4 bytes) + bitmap data
+	require.GreaterOrEqual(t, len(data), 8, "marshaled data should have at least seedCount and bitmapLen")
+}
+
+func TestUnmarshal(t *testing.T) {
+	mp := mpool.MustNewZero()
+	vec := testutil.NewVector(testCount, types.New(types.T_int64, 0, 0), mp, false, nil)
+	defer vec.Free(mp)
+
+	// Test normal case: marshal and unmarshal
+	bf1 := New(testCount, testRate)
+	bf1.Add(vec)
+
+	data, err := bf1.Marshal()
+	require.NoError(t, err)
+
+	bf2 := BloomFilter{}
+	err = bf2.Unmarshal(data)
+	require.NoError(t, err)
+
+	// Verify that unmarshaled filter has same functionality
+	allFound := true
+	bf2.Test(vec, func(exists bool, _ int) {
+		allFound = allFound && exists
+	})
+	require.True(t, allFound, "unmarshaled filter should find all added elements")
+
+	// Test error cases
+	tests := []struct {
+		name    string
+		data    []byte
+		wantErr bool
+	}{
+		{
+			name:    "empty data",
+			data:    []byte{},
+			wantErr: true,
+		},
+		{
+			name:    "data too short (< 4 bytes)",
+			data:    []byte{1, 2, 3},
+			wantErr: true,
+		},
+		{
+			name:    "invalid seed count (0)",
+			data:    []byte{0, 0, 0, 0}, // seedCount = 0
+			wantErr: true,
+		},
+		{
+			name:    "seed data truncated",
+			data:    []byte{1, 0, 0, 0, 1, 2, 3, 4, 5, 6, 7}, // seedCount=1, but only 7 bytes for seed (need 8)
+			wantErr: true,
+		},
+		{
+			name:    "no bitmap length",
+			data:    []byte{1, 0, 0, 0, 1, 2, 3, 4, 5, 6, 7, 8}, // seedCount=1, one seed, but no bitmapLen
+			wantErr: true,
+		},
+		{
+			name:    "bitmap data truncated",
+			data:    []byte{1, 0, 0, 0, 1, 2, 3, 4, 5, 6, 7, 8, 10, 0, 0, 0, 1, 2, 3}, // seedCount=1, one seed, bitmapLen=10, but only 3 bytes
+			wantErr: true,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			bf := BloomFilter{}
+			err := bf.Unmarshal(tt.data)
+			if tt.wantErr {
+				require.Error(t, err)
+			} else {
+				require.NoError(t, err)
+			}
+		})
+	}
+}
+
+func TestMarshalUnmarshal(t *testing.T) {
+	mp := mpool.MustNewZero()
+	vecs := make([]*vector.Vector, vecCount)
+	for i := 0; i < vecCount; i++ {
+		vecs[i] = testutil.NewVector(testCount/vecCount, types.New(types.T_int64, 0, 0), mp, false, nil)
+	}
+	defer func() {
+		for i := 0; i < vecCount; i++ {
+			vecs[i].Free(mp)
+		}
+	}()
+
+	// Create and populate original filter
+	bf1 := New(testCount, testRate)
+	for j := 0; j < vecCount; j++ {
+		bf1.Add(vecs[j])
+	}
+
+	// Marshal
+	data, err := bf1.Marshal()
+	require.NoError(t, err)
+	require.NotNil(t, data)
+
+	// Unmarshal into new filter
+	bf2 := BloomFilter{}
+	err = bf2.Unmarshal(data)
+	require.NoError(t, err)
+
+	// Verify both filters behave the same
+	testVec := testutil.NewVector(testCount/vecCount, types.New(types.T_int64, 0, 0), mp, false, nil)
+	defer testVec.Free(mp)
+
+	// Test original filter
+	allFound1 := true
+	bf1.Test(testVec, func(exists bool, _ int) {
+		allFound1 = allFound1 && exists
+	})
+
+	// Test unmarshaled filter
+	allFound2 := true
+	bf2.Test(testVec, func(exists bool, _ int) {
+		allFound2 = allFound2 && exists
+	})
+
+	require.Equal(t, allFound1, allFound2, "original and unmarshaled filters should behave the same")
+
+	// Test with new data that wasn't added
+	newVec := testutil.NewVector(testCount*2, types.New(types.T_int64, 0, 0), mp, false, nil)
+	defer newVec.Free(mp)
+
+	allFoundNew1 := true
+	bf1.Test(newVec, func(exists bool, _ int) {
+		allFoundNew1 = allFoundNew1 && exists
+	})
+
+	allFoundNew2 := true
+	bf2.Test(newVec, func(exists bool, _ int) {
+		allFoundNew2 = allFoundNew2 && exists
+	})
+
+	require.Equal(t, allFoundNew1, allFoundNew2, "original and unmarshaled filters should behave the same for new data")
+}

pkg/defines/type.go

@@ -241,6 +241,10 @@ type VarScopeKey struct{}
 // Determine if it is a stored procedure
 type InSp struct{}
 
+// IvfBloomFilter carries BloomFilter bytes for ivf entries scan in internal SQL executor.
+// This key is set on context when invoking internal SQL from ivf_search.
+type IvfBloomFilter struct{}
+
 // PkCheckByTN whether TN does primary key uniqueness check against transaction's workspace or not.
 type PkCheckByTN struct{}
 

pkg/objectio/types.go

@@ -20,9 +20,9 @@ import (
 	"github.com/matrixorigin/matrixone/pkg/common/mpool"
 	"github.com/matrixorigin/matrixone/pkg/container/batch"
 	"github.com/matrixorigin/matrixone/pkg/container/types"
-	"github.com/matrixorigin/matrixone/pkg/container/vector"
 	"github.com/matrixorigin/matrixone/pkg/fileservice"
 	"github.com/matrixorigin/matrixone/pkg/vectorindex/metric"
+	"github.com/matrixorigin/matrixone/pkg/vm/engine/tae/containers"
 	"github.com/matrixorigin/matrixone/pkg/vm/engine/tae/index"
 )
 
@@ -42,13 +42,14 @@ type ColumnMetaFetcher interface {
 	MustGetColumn(seqnum uint16) ColumnMeta
 }
 
-type ReadFilterSearchFuncType func(*vector.Vector) []int64
+type ReadFilterSearchFuncType func(containers.Vectors) []int64
 
 type BlockReadFilter struct {
 	HasFakePK          bool
 	Valid              bool
 	SortedSearchFunc   ReadFilterSearchFuncType
 	UnSortedSearchFunc ReadFilterSearchFuncType
+	Cleanup            func() // Cleanup function to release resources (e.g., reusableTempVec)
 }
 
 func (f BlockReadFilter) DecideSearchFunc(isSortedBlk bool) ReadFilterSearchFuncType {