Merge pull request #1 from mrproliu/merge-callback

Author: hanahmily

config.go

@@ -18,6 +18,9 @@ import (
 	"io"
 	"log"
 
+	"github.com/RoaringBitmap/roaring"
+	segment "github.com/blugelabs/bluge_segment_api"
+
 	"github.com/blugelabs/bluge/analysis"
 	"github.com/blugelabs/bluge/analysis/analyzer"
 	"github.com/blugelabs/bluge/index"
@@ -78,6 +81,12 @@ func (config Config) WithSearchStartFunc(f func(size uint64) error) Config {
 	return config
 }
 
+func (config Config) WithPrepareMergeCallback(f func(src []*roaring.Bitmap, segments []segment.Segment, id uint64) (
+	dest []*roaring.Bitmap, err error)) Config {
+	config.indexConfig.PrepareMergeFunc = f
+	return config
+}
+
 func DefaultConfig(path string) Config {
 	indexConfig := index.DefaultConfig(path)
 	return defaultConfig(indexConfig)

index/config.go

@@ -21,6 +21,8 @@ import (
 
 	"github.com/blugelabs/bluge/index/mergeplan"
 	"github.com/blugelabs/ice"
+
+	"github.com/RoaringBitmap/roaring"
 )
 
 type Config struct {
@@ -80,6 +82,9 @@ type Config struct {
 	virtualFields map[string][]segment.Field
 
 	CacheMaxBytes int
+
+	// PerFieldSimilarity allows specifying a custom callback before merge operation
+	PrepareMergeFunc func(src []*roaring.Bitmap, segments []segment.Segment, id uint64) (dest []*roaring.Bitmap, err error)
 }
 
 func (config Config) WithSegmentType(typ string) Config {

index/merge.go

@@ -206,6 +206,7 @@ func (s *Writer) executeMergeTask(merges chan *segmentMerge, task *mergeplan.Mer
 			// decrement the ref counts on skipping introduction.
 			// FIXME stale file that won't get cleaned up
 			_ = seg.Close()
+			_ = s.directory.Remove(ItemKindSegment, newSegmentID)
 		}
 	}
 
@@ -363,9 +364,18 @@ func (s *Writer) mergeSegmentBases(merges chan *segmentMerge, snapshot *Snapshot
 
 func (s *Writer) merge(segments []segment.Segment, drops []*roaring.Bitmap, id uint64) (
 	[][]uint64, error) {
+	var err error
+	if s.config.PrepareMergeFunc != nil {
+		// call the prepare merge function
+		drops, err = s.config.PrepareMergeFunc(drops, segments, id)
+		if err != nil {
+			return nil, fmt.Errorf("prepare merge failed: %w", err)
+		}
+	}
+
 	merger := s.segPlugin.Merge(segments, drops, s.config.MergeBufferSize)
 
-	err := s.directory.Persist(ItemKindSegment, id, merger, s.closeCh)
+	err = s.directory.Persist(ItemKindSegment, id, merger, s.closeCh)
 	if err != nil {
 		return nil, err
 	}