implement database iterator functionality (#326)

add low-level custom iterator for rosedb

Author: XiXi-2024

examples/iterate/main.go

@@ -8,7 +8,6 @@ import (
 )
 
 // this file shows how to use the iterate operations of rosedb
-// you can use Ascend, Descend(and some other similar methods) to iterate all keys and values in order.
 func main() {
 	// specify the options
 	options := rosedb.DefaultOptions
@@ -28,12 +27,19 @@ func main() {
 		_ = db.Close()
 	}()
 
+	// prepare sample data
 	_ = db.Put([]byte("key13"), []byte("value13"))
 	_ = db.Put([]byte("key11"), []byte("value11"))
 	_ = db.Put([]byte("key35"), []byte("value35"))
 	_ = db.Put([]byte("key27"), []byte("value27"))
 	_ = db.Put([]byte("key41"), []byte("value41"))
 
+	dbIteratorExample(db)
+	customIteratorExample(db)
+}
+
+// dbIteratorExample demonstrates the built-in database iterator methods
+func dbIteratorExample(db *rosedb.DB) {
 	// iterate all keys in order
 	db.AscendKeys(nil, true, func(k []byte) (bool, error) {
 		fmt.Println("key = ", string(k))
@@ -57,10 +63,43 @@ func main() {
 		fmt.Printf("key = %s, value = %s\n", string(k), string(v))
 		return true, nil
 	})
-
 	// you can also use some other similar methods to iterate the data.
 	// db.AscendRange()
 	// db.AscendGreaterOrEqual()
 	// db.DescendRange()
 	// db.DescendLessOrEqual()
 }
+
+// customIteratorExample demonstrates how to use the low-level iterator API
+func customIteratorExample(db *rosedb.DB) {
+	// 1: Using iterator with ContinueOnError = true
+	iterOpts := rosedb.DefaultIteratorOptions
+	iterOpts.ContinueOnError = true
+	iter1 := db.NewIterator(iterOpts)
+	defer iter1.Close()
+
+	for iter1.Rewind(); iter1.Valid(); iter1.Next() {
+		item := iter1.Item()
+		if item != nil {
+			fmt.Printf("key = %s, value = %s\n", string(item.Key), string(item.Value))
+		}
+	}
+	if err := iter1.Err(); err != nil {
+		fmt.Printf("Iterator encountered errors but continued: %v\n", err)
+	}
+
+	// 2: Using iterator with ContinueOnError = false
+	iterOpts.ContinueOnError = false
+	iter2 := db.NewIterator(iterOpts)
+	defer iter2.Close()
+
+	for iter2.Rewind(); iter2.Valid(); iter2.Next() {
+		item := iter2.Item()
+		if item != nil {
+			fmt.Printf("key = %s, value = %s\n", string(item.Key), string(item.Value))
+		}
+	}
+	if err := iter2.Err(); err != nil {
+		fmt.Printf("Iterator stopped due to error: %v\n", err)
+	}
+}

index/btree.go

@@ -147,3 +147,126 @@ func (mt *MemoryBTree) DescendLessOrEqual(key []byte, handleFn func(key []byte,
 		return cont
 	})
 }
+
+func (mt *MemoryBTree) Iterator(reverse bool) IndexIterator {
+	if mt.tree == nil {
+		return nil
+	}
+	mt.lock.RLock()
+	defer mt.lock.RUnlock()
+
+	return newMemoryBTreeIterator(mt.tree, reverse)
+}
+
+// memoryBTreeIterator represents a B-tree index iterator implementation
+type memoryBTreeIterator struct {
+	tree    *btree.BTree // underlying B-tree implementation
+	reverse bool         // indicates whether to traverse in descending order
+	current *item        // current element being traversed
+	valid   bool         // indicates if the iterator is valid
+}
+
+func newMemoryBTreeIterator(tree *btree.BTree, reverse bool) *memoryBTreeIterator {
+	var current *item
+	var valid bool
+	if tree.Len() > 0 {
+		if reverse {
+			current = tree.Max().(*item)
+		} else {
+			current = tree.Min().(*item)
+		}
+		valid = true
+	}
+	return &memoryBTreeIterator{
+		tree:    tree.Clone(),
+		reverse: reverse,
+		current: current,
+		valid:   valid,
+	}
+}
+
+func (it *memoryBTreeIterator) Rewind() {
+	if it.tree == nil || it.tree.Len() == 0 {
+		return
+	}
+	if it.reverse {
+		it.current = it.tree.Max().(*item)
+	} else {
+		it.current = it.tree.Min().(*item)
+	}
+}
+
+func (it *memoryBTreeIterator) Seek(key []byte) {
+	if it.tree == nil || !it.valid {
+		return
+	}
+	seekItem := &item{key: key}
+	it.valid = false
+	if it.reverse {
+		it.tree.DescendLessOrEqual(seekItem, func(i btree.Item) bool {
+			it.current = i.(*item)
+			it.valid = true
+			return false
+		})
+	} else {
+		it.tree.AscendGreaterOrEqual(seekItem, func(i btree.Item) bool {
+			it.current = i.(*item)
+			it.valid = true
+			return false
+		})
+	}
+}
+
+func (it *memoryBTreeIterator) Next() {
+	if it.tree == nil || !it.valid {
+		return
+	}
+	it.valid = false
+	if it.reverse {
+		it.tree.DescendLessOrEqual(it.current, func(i btree.Item) bool {
+			if !i.(*item).Less(it.current) {
+				return true
+			}
+			it.current = i.(*item)
+			it.valid = true
+			return false
+		})
+	} else {
+		it.tree.AscendGreaterOrEqual(it.current, func(i btree.Item) bool {
+			if !it.current.Less(i.(*item)) {
+				return true
+			}
+			it.current = i.(*item)
+			it.valid = true
+			return false
+		})
+	}
+	if !it.valid {
+		it.current = nil
+	}
+}
+
+func (it *memoryBTreeIterator) Valid() bool {
+	return it.valid
+}
+
+func (it *memoryBTreeIterator) Key() []byte {
+	if !it.valid {
+		return nil
+	}
+	return it.current.key
+}
+
+func (it *memoryBTreeIterator) Value() *wal.ChunkPosition {
+	if !it.valid {
+		return nil
+	}
+	return it.current.pos
+}
+
+func (it *memoryBTreeIterator) Close() {
+	it.tree.Clear(true)
+	it.tree = nil
+	it.current = nil
+	it.valid = false
+}

index/btree_test.go

@@ -3,6 +3,7 @@ package index
 import (
 	"bytes"
 	"fmt"
+	"github.com/stretchr/testify/assert"
 	"testing"
 
 	"github.com/rosedblabs/wal"
@@ -193,3 +194,93 @@ func TestMemoryBTree_AscendGreaterOrEqual_DescendLessOrEqual(t *testing.T) {
 		return true, nil
 	})
 }
+
+func TestMemoryBTree_Iterator(t *testing.T) {
+	mt := newBTree()
+	// Test iterator for empty tree
+	it1 := mt.Iterator(false)
+	assert.Equal(t, false, it1.Valid())
+
+	// Build test data
+	testData := map[string]*wal.ChunkPosition{
+		"acee": {SegmentId: 1, BlockNumber: 2, ChunkOffset: 3, ChunkSize: 100},
+		"bbcd": {SegmentId: 2, BlockNumber: 3, ChunkOffset: 4, ChunkSize: 200},
+		"code": {SegmentId: 3, BlockNumber: 4, ChunkOffset: 5, ChunkSize: 300},
+		"eede": {SegmentId: 4, BlockNumber: 5, ChunkOffset: 6, ChunkSize: 400},
+	}
+
+	// Insert test data
+	for k, v := range testData {
+		mt.Put([]byte(k), v)
+	}
+
+	// Test ascending iteration
+	iter := mt.Iterator(false)
+	var prevKey string
+	count := 0
+	for iter.Rewind(); iter.Valid(); iter.Next() {
+		currKey := string(iter.Key())
+		pos := iter.Value()
+
+		// Verify key order
+		if prevKey != "" {
+			assert.True(t, currKey > prevKey)
+		}
+
+		// Verify value correctness
+		expectedPos := testData[currKey]
+		assert.Equal(t, expectedPos, pos)
+
+		prevKey = currKey
+		count++
+	}
+	assert.Equal(t, len(testData), count)
+
+	// Test descending iteration
+	iter = mt.Iterator(true)
+	prevKey = ""
+	count = 0
+	for iter.Rewind(); iter.Valid(); iter.Next() {
+		currKey := string(iter.Key())
+		pos := iter.Value()
+
+		// Verify key order
+		if prevKey != "" {
+			assert.True(t, currKey < prevKey)
+		}
+
+		// Verify value correctness
+		expectedPos := testData[currKey]
+		assert.Equal(t, expectedPos, pos)
+
+		prevKey = currKey
+		count++
+	}
+	assert.Equal(t, len(testData), count)
+
+	// Test Seek operation
+	testCases := []struct {
+		seekKey    string
+		expectKey  string
+		shouldFind bool
+	}{
+		{"b", "bbcd", true},   // Should find bbcd
+		{"cc", "code", true},  // Should find code
+		{"d", "eede", true},   // Should find eede
+		{"f", "", false},      // Should not find any element
+		{"aaa", "acee", true}, // Should find acee
+	}
+
+	for _, tc := range testCases {
+		iter = mt.Iterator(false)
+		iter.Seek([]byte(tc.seekKey))
+
+		if tc.shouldFind {
+			assert.True(t, iter.Valid())
+			assert.Equal(t, tc.expectKey, string(iter.Key()))
+			assert.Equal(t, testData[tc.expectKey], iter.Value())
+		} else {
+			assert.False(t, iter.Valid())
+		}
+	}
+}

index/index.go

@@ -42,6 +42,9 @@ type Indexer interface {
 	// DescendLessOrEqual iterates in descending order, starting from key <= given key,
 	// invoking handleFn. Stops if handleFn returns false.
 	DescendLessOrEqual(key []byte, handleFn func(key []byte, position *wal.ChunkPosition) (bool, error))
+
+	// IndexIterator returns an index iterator.
+	Iterator(reverse bool) IndexIterator
 }
 
 type IndexerType = byte
@@ -61,3 +64,27 @@ func NewIndexer() Indexer {
 		panic("unexpected index type")
 	}
 }
+
+// IndexIterator represents a generic index iterator interface.
+type IndexIterator interface {
+	// Rewind resets the iterator to its initial position.
+	Rewind()
+
+	// Seek positions the cursor to the element with the specified key.
+	Seek(key []byte)
+
+	// Next moves the cursor to the next element.
+	Next()
+
+	// Valid checks if the iterator is still valid for reading.
+	Valid() bool
+
+	// Key returns the key of the current element.
+	Key() []byte
+
+	// Value returns the value (chunk position) of the current element.
+	Value() *wal.ChunkPosition
+
+	// Close releases the resources associated with the iterator.
+	Close()
+}