fix: rebuild rank cache for set imports (ImportRoaringShard)

the ranked cache must be fully rebuilt as the BitmapRewriter doesn't
have an easy way to track which rows had how many bits changed which
we would need to update the cache.

We also needed to add a Clear method to topn cache to completely
remove old values before the rebuild as otherwise they'd sit there and
pollute the cache after it was rebuilt.

This also includes fixing a strange idiosyncrasy where the _exists
field was a set field, but didn't have its type explicitly set. This
was causing it to have a ranked cache even though that option was
turned off. Hoping this doesn't have any weird follow-on effects... or
if it does the tests catch them.

cache.go

@@ -42,21 +42,26 @@ type cache interface {
 
 	// SetStats defines the stats client used in the cache.
 	SetStats(s stats.StatsClient)
+
+	// Clear removes everything from the cache. If possible it should leave allocated structures in place to be reused.
+	Clear()
 }
 
 // lruCache represents a least recently used Cache implementation.
 type lruCache struct {
-	cache  *lru.Cache
-	counts map[uint64]uint64
-	stats  stats.StatsClient
+	cache      *lru.Cache
+	counts     map[uint64]uint64
+	stats      stats.StatsClient
+	maxEntries uint32
 }
 
 // newLRUCache returns a new instance of LRUCache.
 func newLRUCache(maxEntries uint32) *lruCache {
 	c := &lruCache{
-		cache:  lru.New(int(maxEntries)),
-		counts: make(map[uint64]uint64),
-		stats:  stats.NopStatsClient,
+		cache:      lru.New(int(maxEntries)),
+		counts:     make(map[uint64]uint64),
+		stats:      stats.NopStatsClient,
+		maxEntries: maxEntries,
 	}
 	c.cache.OnEvicted = c.onEvicted
 	return c
@@ -118,6 +123,13 @@ func (c *lruCache) SetStats(s stats.StatsClient) {
 	c.stats = s
 }
 
+func (c *lruCache) Clear() {
+	for k := range c.counts {
+		delete(c.counts, k)
+	}
+	c.cache = lru.New(int(c.maxEntries))
+}
+
 func (c *lruCache) onEvicted(key lru.Key, _ interface{}) { delete(c.counts, key.(uint64)) }
 
 // Ensure LRUCache implements Cache.
@@ -125,6 +137,7 @@ var _ cache = &lruCache{}
 
 // rankCache represents a cache with sorted entries.
 type rankCache struct {
+	// TODO why does this have a lock and lruCache doesn't?
 	mu           sync.Mutex
 	entries      map[uint64]uint64
 	rankings     bitmapPairs // cached, ordered list
@@ -157,6 +170,21 @@ func NewRankCache(maxEntries uint32) *rankCache {
 	}
 }
 
+func (c *rankCache) Clear() {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	for k := range c.entries {
+		delete(c.entries, k)
+	}
+	c.rankings = c.rankings[:0]
+	c.rankingsRead = false
+	c.dirty = false
+
+	c.updateN = 0
+	c.updateTime = time.Time{}
+	c.thresholdValue = 0
+}
+
 // Add adds a count to the cache.
 func (c *rankCache) Add(id uint64, n uint64) {
 	c.mu.Lock()
@@ -594,6 +622,7 @@ func (c nopCache) Invalidate()                {}
 func (c nopCache) Len() int                   { return 0 }
 func (c nopCache) Recalculate()               {}
 func (c nopCache) SetStats(stats.StatsClient) {}
+func (c nopCache) Clear()                     {}
 
 func (c nopCache) Top() []bitmapPair {
 	return []bitmapPair{}

client/batch_test.go

@@ -10,6 +10,7 @@ import (
 	"testing"
 	"time"
 
+	featurebase "github.com/molecula/featurebase/v3"
 	"github.com/molecula/featurebase/v3/test"
 
 	"github.com/pkg/errors"
@@ -27,7 +28,6 @@ func TestAgainstCluster(t *testing.T) {
 	c := test.MustRunCluster(t, 1)
 	defer c.Close()
 	client := NewTestClient(t, c)
-
 	t.Run("string-slice-combos", func(t *testing.T) { testStringSliceCombos(t, c, client) })
 	t.Run("import-batch-ints", func(t *testing.T) { testImportBatchInts(t, c, client) })
 	t.Run("import-batch-sorting", func(t *testing.T) { testImportBatchSorting(t, c, client) })
@@ -39,6 +39,8 @@ func TestAgainstCluster(t *testing.T) {
 	t.Run("batches-strings-ids", func(t *testing.T) { testBatchesStringIDs(t, c, client) })
 	t.Run("test-batch-staleness", func(t *testing.T) { testBatchStaleness(t, c, client) })
 	t.Run("test-import-batch-multiple-ints", func(t *testing.T) { testImportBatchMultipleInts(t, c, client) })
+	t.Run("test-import-batch-sets-clears", func(t *testing.T) { testImportBatchSetsAndClears(t, c, client) })
+	t.Run("test-topn-cache-regression", func(t *testing.T) { testTopNCacheRegression(t, c, client) })
 }
 
 func testStringSliceCombos(t *testing.T, c *test.Cluster, client *Client) {
@@ -1387,3 +1389,134 @@ func testImportBatchMultipleInts(t *testing.T, c *test.Cluster, client *Client)
 	}
 
 }
+
+func testImportBatchSetsAndClears(t *testing.T, c *test.Cluster, client *Client) {
+	schema := NewSchema()
+	idx := schema.Index("test-import-batch-set-and-clear")
+	field := idx.Field("aset", OptFieldTypeSet(featurebase.DefaultCacheType, featurebase.DefaultCacheSize))
+	err := client.SyncSchema(schema)
+	if err != nil {
+		t.Fatalf("syncing schema: %v", err)
+	}
+
+	b, err := NewBatch(client, 6, idx, []*Field{field}, OptUseShardTransactionalEndpoint(true))
+	if err != nil {
+		t.Fatalf("getting batch: %v", err)
+	}
+
+	r := Row{
+		Values: make([]interface{}, 1),
+		Clears: make(map[int]interface{}),
+	}
+
+	vals := []uint64{1, 2, 3, 1, 5, 6}
+	clears := []interface{}{nil, uint64(1), uint64(3), nil, uint64(2), uint64(4)}
+	for i := uint64(0); i < 6; i++ {
+		r.ID = i%3 + 1
+		r.Values[0] = vals[i]
+		if clears[i] != nil {
+			r.Clears[0] = clears[i]
+		}
+		err := b.Add(r)
+		if err != nil && err != ErrBatchNowFull {
+			t.Fatalf("adding to batch: %v", err)
+		}
+	}
+	err = b.Import()
+	if err != nil {
+		t.Fatalf("importing: %v", err)
+	}
+
+	if resp, err := client.Query(field.TopN(6)); err != nil {
+		t.Fatalf("querying topn: %v", err)
+	} else if res := resp.Result().CountItems(); len(res) != 3 {
+		t.Fatalf("unexpected topn: %+v", res)
+	}
+
+	exp := [][]uint64{
+		{},
+		{1},
+		{},
+		{},
+		{},
+		{2},
+		{3},
+	}
+	for row := 0; row < 7; row++ {
+		resp, err := client.Query(field.Row(row))
+		if err != nil {
+			t.Fatalf("querying: %v", err)
+		}
+		res := resp.Results()[0].Row().Columns
+		if !reflect.DeepEqual(exp[row], res) && !(len(exp[row]) == 0 && len(res) == 0) {
+			t.Errorf("row: %d, exp: %v, got %v", row, exp[row], res)
+		}
+	}
+
+}
+
+// testTopNCacheRegression recreates an issue we saw in an IDK test
+// where if a value is completely removed (all bits unset from a row),
+// it didn't get removed from the cache beacuse a full recalculation
+// had no way to clear the cache, it would just reset existing
+// values. We added Clear on the cache interface to fix this.
+func testTopNCacheRegression(t *testing.T, c *test.Cluster, client *Client) {
+	schema := NewSchema()
+	idx := schema.Index("test-topn-cache-regression")
+	field := idx.Field("aset", OptFieldTypeSet(featurebase.DefaultCacheType, featurebase.DefaultCacheSize))
+	err := client.SyncSchema(schema)
+	if err != nil {
+		t.Fatalf("syncing schema: %v", err)
+	}
+
+	b, err := NewBatch(client, 3, idx, []*Field{field}, OptUseShardTransactionalEndpoint(true))
+	if err != nil {
+		t.Fatalf("getting batch: %v", err)
+	}
+
+	records := []struct {
+		ID    uint64
+		Set   interface{}
+		Clear interface{}
+	}{
+		{0, 1, nil},
+		{featurebase.ShardWidth, 1, nil},
+		{featurebase.ShardWidth * 2, nil, 1},
+		{featurebase.ShardWidth * 2, nil, 1},
+		{0, nil, 1},
+		{featurebase.ShardWidth, nil, 1},
+		{featurebase.ShardWidth, 1, nil},
+		{featurebase.ShardWidth, nil, nil},
+	}
+
+	for _, rec := range records {
+		if rec.Set != nil {
+			rec.Set = uint64(rec.Set.(int))
+		}
+		row := Row{
+			ID:     rec.ID,
+			Values: []interface{}{rec.Set},
+		}
+		if rec.Clear != nil {
+			row.Clears = map[int]interface{}{0: uint64(rec.Clear.(int))}
+		}
+
+		err := b.Add(row)
+		if err == ErrBatchNowFull {
+			if err := b.Import(); err != nil {
+				t.Fatalf("importing: %v", err)
+			}
+		}
+	}
+	if err := b.Import(); err != nil {
+		t.Fatalf("importing: %v", err)
+	}
+
+	if resp, err := client.Query(field.TopN(6)); err != nil {
+		t.Fatalf("querying topn: %v", err)
+	} else if res := resp.Result().CountItems(); len(res) != 1 {
+		t.Fatalf("unexpected topn: %+v", res)
+	} else if res[0].ID != 1 || res[0].Count != 1 {
+		t.Fatalf("unexpected topn result: %v", res)
+	}
+}

fragment.go

@@ -2401,7 +2401,19 @@ func (f *fragment) ImportRoaringClearAndSet(ctx context.Context, tx Tx, clear, s
 	}
 
 	err = tx.ApplyRewriter(f.index(), f.field(), f.view(), f.shard, 0, rewriter)
-	return errors.Wrap(err, "applying rewriter")
+	if err != nil {
+		errors.Wrap(err, "applying rewriter")
+	}
+	if f.CacheType != CacheTypeNone {
+		// TODO this may be quite a bit slower than the way
+		// importRoaring does it as it tracks the number of bits
+		// changed per row. We could do that, but I think it'd require
+		// significant changes to the Rewriter API.
+		f.mu.Lock()
+		defer f.mu.Unlock()
+		return f.rebuildRankCache(ctx, tx)
+	}
+	return nil
 }
 
 // ImportRoaringBSI interprets "clear" as a single row specifying
@@ -2539,17 +2551,13 @@ func (f *fragment) FlushCache() error {
 	defer f.mu.Unlock()
 	return f.flushCache()
 }
-func (f *fragment) RebuildRankCache(ctx context.Context) error {
+
+func (f *fragment) rebuildRankCache(ctx context.Context, tx Tx) error {
 	if f.CacheType != CacheTypeRanked {
-		return nil //only rebuild ranked caches
+		return nil // only rebuild ranked caches
 	}
-	f.mu.Lock()
-	defer f.mu.Unlock()
-	tx, err := f.holder.BeginTx(false, f.idx, f.shard)
-	if err != nil {
-		return err
-	}
-	defer tx.Rollback()
+
+	f.cache.Clear()
 	rows, err := f.unprotectedRows(ctx, tx, uint64(0))
 	if err != nil {
 		return err
@@ -2565,6 +2573,20 @@ func (f *fragment) RebuildRankCache(ctx context.Context) error {
 	return nil
 }
 
+func (f *fragment) RebuildRankCache(ctx context.Context) error {
+	if f.CacheType != CacheTypeRanked {
+		return nil //only rebuild ranked caches
+	}
+	f.mu.Lock()
+	defer f.mu.Unlock()
+	tx, err := f.holder.BeginTx(false, f.idx, f.shard)
+	if err != nil {
+		return err
+	}
+	defer tx.Rollback()
+	return f.rebuildRankCache(ctx, tx)
+}
+
 func (f *fragment) flushCache() error {
 	if f.cache == nil {
 		return nil

index.go

@@ -355,7 +355,7 @@ func (i *Index) openExistenceField() error {
 		Index:     i.name,
 		Field:     existenceFieldName,
 		CreatedAt: 0,
-		Meta:      &FieldOptions{CacheType: CacheTypeNone, CacheSize: 0},
+		Meta:      &FieldOptions{Type: FieldTypeSet, CacheType: CacheTypeNone, CacheSize: 0},
 	}
 
 	// First try opening the existence field from disk. If it doesn't already