Table.Get SelectorRanges implementation

batch.go

@@ -107,7 +107,7 @@ func (b *_batch) DeleteRange(start []byte, end []byte, opt WriteOptions, _ ...Ba
 }
 
 func (b *_batch) Iter(opt *IterOptions, _ ...Batch) Iterator {
-	return newIteratorFromBatch(b.Batch, opt)
+	return newIterator(b.Batch, opt)
 }
 
 func (b *_batch) Apply(batch Batch, opt WriteOptions) error {

bond.go

@@ -205,7 +205,7 @@ func (db *_db) Iter(opt *IterOptions, batch ...Batch) Iterator {
 	if batch != nil && len(batch) > 0 && batch[0] != nil {
 		return batch[0].Iter(opt)
 	} else {
-		return newIteratorFromDB(db, opt)
+		return newIterator(db.pebble, opt)
 	}
 }
 

index.go

@@ -179,55 +179,111 @@ func (idx *Index[T]) Name() string {
 // The iterator will return the index keys in the order specified by the index.
 // The iterator will return data if created for primary index.
 func (idx *Index[T]) Iter(table Table[T], selector Selector[T], optBatch ...Batch) Iterator {
-	var (
-		lowerBound []byte
-		upperBound []byte
-	)
+	var iterConstructor Iterationer = table.DB()
+	if len(optBatch) > 0 {
+		iterConstructor = optBatch[0]
+	}
 
 	keyBufferPool := table.DB().getKeyBufferPool()
 
 	switch selector.Type() {
 	case SelectorTypePoint:
 		sel := selector.(SelectorPoint[T])
 
-		lowerBound = encodeIndexKey(table, sel.Point(), idx, keyBufferPool.Get()[:0])
-		upperBound = keySuccessor(lowerBound[0:_KeyPrefixSplitIndex(lowerBound)], keyBufferPool.Get()[:0])
+		lowerBound := encodeIndexKey(table, sel.Point(), idx, keyBufferPool.Get()[:0])
+		upperBound := keySuccessor(lowerBound[0:_KeyPrefixSplitIndex(lowerBound)], keyBufferPool.Get()[:0])
+
+		releaseBuffers := func() {
+			keyBufferPool.Put(lowerBound[:0])
+			keyBufferPool.Put(upperBound[:0])
+		}
+
+		return iterConstructor.Iter(&IterOptions{
+			IterOptions: pebble.IterOptions{
+				LowerBound: lowerBound,
+				UpperBound: upperBound,
+			},
+			releaseBufferOnClose: releaseBuffers,
+		})
+	case SelectorTypePoints:
+		sel := selector.(SelectorPoints[T])
+
+		var pebbleOpts []*IterOptions
+		for _, point := range sel.Points() {
+			lowerBound := encodeIndexKey(table, point, idx, keyBufferPool.Get()[:0])
+			upperBound := keySuccessor(lowerBound[0:_KeyPrefixSplitIndex(lowerBound)], keyBufferPool.Get()[:0])
+
+			releaseBuffers := func() {
+				keyBufferPool.Put(lowerBound[:0])
+				keyBufferPool.Put(upperBound[:0])
+			}
+
+			pebbleOpts = append(pebbleOpts, &IterOptions{
+				IterOptions: pebble.IterOptions{
+					LowerBound: lowerBound,
+					UpperBound: upperBound,
+				},
+				releaseBufferOnClose: releaseBuffers,
+			})
+		}
+
+		return newIteratorMulti(iterConstructor, pebbleOpts)
 	case SelectorTypeRange:
 		sel := selector.(SelectorRange[T])
 		low, up := sel.Range()
 
-		lowerBound = encodeIndexKey(table, low, idx, keyBufferPool.Get()[:0])
-		upperBound = encodeIndexKey(table, up, idx, keyBufferPool.Get()[:0])
+		lowerBound := encodeIndexKey(table, low, idx, keyBufferPool.Get()[:0])
+		upperBound := encodeIndexKey(table, up, idx, keyBufferPool.Get()[:0])
 		if bytes.Equal(lowerBound, upperBound) {
 			upperBound = keySuccessor(lowerBound[0:_KeyPrefixSplitIndex(lowerBound)], upperBound[:0])
 		} else {
 			upperBound = keySuccessor(upperBound, upperBound[:0])
 		}
-	default:
-		panic("invalid selector type")
-	}
 
-	releaseBuffers := func() {
-		keyBufferPool.Put(lowerBound[:0])
-		keyBufferPool.Put(upperBound[:0])
-	}
+		releaseBuffers := func() {
+			keyBufferPool.Put(lowerBound[:0])
+			keyBufferPool.Put(upperBound[:0])
+		}
 
-	if len(optBatch) > 0 {
-		return optBatch[0].Iter(&IterOptions{
-			IterOptions: pebble.IterOptions{
-				LowerBound: lowerBound,
-				UpperBound: upperBound,
-			},
-			releaseBufferOnClose: releaseBuffers,
-		})
-	} else {
-		return table.DB().Iter(&IterOptions{
+		return iterConstructor.Iter(&IterOptions{
 			IterOptions: pebble.IterOptions{
 				LowerBound: lowerBound,
 				UpperBound: upperBound,
 			},
 			releaseBufferOnClose: releaseBuffers,
 		})
+	case SelectorTypeRanges:
+		sel := selector.(SelectorRanges[T])
+
+		var pebbleOpts []*IterOptions
+		for _, r := range sel.Ranges() {
+			low, up := r[0], r[1]
+
+			lowerBound := encodeIndexKey(table, low, idx, keyBufferPool.Get()[:0])
+			upperBound := encodeIndexKey(table, up, idx, keyBufferPool.Get()[:0])
+			if bytes.Equal(lowerBound, upperBound) {
+				upperBound = keySuccessor(lowerBound[0:_KeyPrefixSplitIndex(lowerBound)], upperBound[:0])
+			} else {
+				upperBound = keySuccessor(upperBound, upperBound[:0])
+			}
+
+			releaseBuffers := func() {
+				keyBufferPool.Put(lowerBound[:0])
+				keyBufferPool.Put(upperBound[:0])
+			}
+
+			pebbleOpts = append(pebbleOpts, &IterOptions{
+				IterOptions: pebble.IterOptions{
+					LowerBound: lowerBound,
+					UpperBound: upperBound,
+				},
+				releaseBufferOnClose: releaseBuffers,
+			})
+		}
+
+		return newIteratorMulti(iterConstructor, pebbleOpts)
+	default:
+		panic("invalid selector type")
 	}
 }
 

iter.go

@@ -26,18 +26,18 @@ type Iterator interface {
 	Close() error
 }
 
+type _iterConstructor interface {
+	NewIter(opts *pebble.IterOptions) *pebble.Iterator
+}
+
 type _iterator struct {
 	Iterator
 
 	opts *IterOptions
 }
 
-func newIteratorFromDB(db *_db, opts *IterOptions) *_iterator {
-	return &_iterator{Iterator: db.pebble.NewIter(&opts.IterOptions), opts: opts}
-}
-
-func newIteratorFromBatch(batch *pebble.Batch, opts *IterOptions) *_iterator {
-	return &_iterator{Iterator: batch.NewIter(&opts.IterOptions), opts: opts}
+func newIterator(itc _iterConstructor, opts *IterOptions) *_iterator {
+	return &_iterator{Iterator: itc.NewIter(&opts.IterOptions), opts: opts}
 }
 
 func (it *_iterator) Close() error {
@@ -53,3 +53,115 @@ func (it *_iterator) Close() error {
 	}
 	return nil
 }
+
+type _iteratorMulti struct {
+	iteratorOptions      []*IterOptions
+	iteratorOptionsIndex int
+
+	iteratorConstuctor Iterationer
+	iterator           Iterator
+}
+
+func newIteratorMulti(itc Iterationer, opts []*IterOptions) *_iteratorMulti {
+	return &_iteratorMulti{
+		iteratorOptions:      opts,
+		iteratorOptionsIndex: 0,
+		iteratorConstuctor:   itc,
+		iterator:             itc.Iter(opts[0]),
+	}
+}
+
+func (it *_iteratorMulti) First() bool {
+	if it.iteratorOptionsIndex != 0 {
+		_ = it.iterator.Close()
+
+		it.iteratorOptionsIndex = 0
+		it.iterator = it.iteratorConstuctor.Iter(it.iteratorOptions[it.iteratorOptionsIndex])
+	}
+	return it.iterator.First()
+}
+
+func (it *_iteratorMulti) Last() bool {
+	if it.iteratorOptionsIndex != len(it.iteratorOptions)-1 {
+		_ = it.iterator.Close()
+
+		it.iteratorOptionsIndex = len(it.iteratorOptions) - 1
+		it.iterator = it.iteratorConstuctor.Iter(it.iteratorOptions[it.iteratorOptionsIndex])
+	}
+	return it.iterator.Last()
+}
+
+func (it *_iteratorMulti) Prev() bool {
+	if !it.iterator.Prev() {
+		if it.iteratorOptionsIndex == 0 {
+			return false
+		}
+
+		_ = it.iterator.Close()
+
+		it.iteratorOptionsIndex--
+		it.iterator = it.iteratorConstuctor.Iter(it.iteratorOptions[it.iteratorOptionsIndex])
+		return it.iterator.Last()
+	}
+	return true
+}
+
+func (it *_iteratorMulti) Next() bool {
+	if !it.iterator.Next() {
+		if it.iteratorOptionsIndex == len(it.iteratorOptions)-1 {
+			return false
+		}
+
+		_ = it.iterator.Close()
+
+		it.iteratorOptionsIndex++
+		it.iterator = it.iteratorConstuctor.Iter(it.iteratorOptions[it.iteratorOptionsIndex])
+		return it.iterator.First()
+	}
+	return true
+}
+
+func (it *_iteratorMulti) Valid() bool {
+	return it.iterator.Valid()
+}
+
+func (it *_iteratorMulti) Error() error {
+	return it.iterator.Error()
+}
+
+func (it *_iteratorMulti) SeekGE(key []byte) bool {
+	//TODO implement me
+	panic("implement me")
+}
+
+func (it *_iteratorMulti) SeekPrefixGE(key []byte) bool {
+	//TODO implement me
+	panic("implement me")
+}
+
+func (it *_iteratorMulti) SeekLT(key []byte) bool {
+	//TODO implement me
+	panic("implement me")
+}
+
+func (it *_iteratorMulti) Key() []byte {
+	return it.iterator.Key()
+}
+
+func (it *_iteratorMulti) Value() []byte {
+	return it.iterator.Value()
+}
+
+func (it *_iteratorMulti) Close() error {
+	defer func() {
+		for _, opts := range it.iteratorOptions {
+			if opts.releaseBufferOnClose != nil {
+				opts.releaseBufferOnClose()
+			}
+		}
+	}()
+
+	return it.iterator.Close()
+}
+
+var _ Iterator = (*_iteratorMulti)(nil)

table.go

@@ -894,7 +894,7 @@ func (t *_table[T]) Get(ctx context.Context, sel Selector[T], optBatch ...Batch)
 		}
 
 		return trs, nil
-	case SelectorTypeRange:
+	case SelectorTypeRange, SelectorTypeRanges:
 		var trs []T
 
 		err := t.ScanIndex(ctx, t.primaryIndex, sel, &trs, optBatch...)

table_test.go

@@ -2,6 +2,7 @@ package bond
 
 import (
 	"context"
+	"math"
 	"testing"
 
 	"github.com/cockroachdb/pebble"
@@ -318,6 +319,70 @@ func TestBondTable_Get_Points(t *testing.T) {
 	assert.Nil(t, tokenBalances[0])
 }
 
+func TestBondTable_Get_Ranges(t *testing.T) {
+	db := setupDatabase()
+	defer tearDownDatabase(db)
+
+	const (
+		TokenBalanceTableID TableID = 0xC0
+	)
+
+	tokenBalanceTable := NewTable[*TokenBalance](TableOptions[*TokenBalance]{
+		DB:        db,
+		TableID:   TokenBalanceTableID,
+		TableName: "token_balance",
+		TablePrimaryKeyFunc: func(builder KeyBuilder, tb *TokenBalance) []byte {
+			return builder.AddUint64Field(tb.ID).Bytes()
+		},
+		Serializer: &serializers.JsonSerializer{},
+	})
+	require.NotNil(t, tokenBalanceTable)
+
+	// token balances to insert
+	insertTokenBalances := []*TokenBalance{
+		{
+			ID:              1,
+			AccountID:       1,
+			ContractAddress: "0xtestContract",
+			AccountAddress:  "0xtestAccount",
+			Balance:         5,
+		},
+		{
+			ID:              2,
+			AccountID:       1,
+			ContractAddress: "0xtestContract",
+			AccountAddress:  "0xtestAccount",
+			Balance:         5,
+		},
+		{
+			ID:              3,
+			AccountID:       1,
+			ContractAddress: "0xtestContract",
+			AccountAddress:  "0xtestAccount",
+			Balance:         7,
+		},
+	}
+
+	err := tokenBalanceTable.Insert(context.Background(), insertTokenBalances)
+	require.NoError(t, err)
+
+	expectedTokenBalances := []*TokenBalance{
+		insertTokenBalances[0],
+		insertTokenBalances[2],
+	}
+
+	// get token balances with range
+	tokenBalances, err := tokenBalanceTable.Get(context.Background(), NewSelectorRanges([]*TokenBalance{{ID: 0}, {ID: 1}}, []*TokenBalance{{ID: 3}, {ID: math.MaxUint64}}))
+	require.NoError(t, err)
+	require.Equal(t, len(expectedTokenBalances), len(tokenBalances))
+	assert.Equal(t, expectedTokenBalances, tokenBalances)
+
+	// get token balance with non-existing id range
+	tokenBalances, err = tokenBalanceTable.Get(context.Background(), NewSelectorRanges([]*TokenBalance{{ID: 5}, {ID: 5}}))
+	require.NoError(t, err)
+	assert.Equal(t, 0, len(tokenBalances))
+}
+
 func TestBondTable_Insert(t *testing.T) {
 	db := setupDatabase()
 	defer tearDownDatabase(db)