Go - Batch Iterator

batch_iterator.go

package util

type IteratorCallback = func(start int, v []interface{}) error

type BatchIterator struct {
	batchSize int
}

func NewBatchIterator(batchSize uint) BatchIterator {
	return BatchIterator{
		batchSize: int(batchSize),
	}
}

func (r BatchIterator) ReadStrings(cb IteratorCallback, strs ...string) error {
	values := make([]interface{}, 0, len(strs))
	for _, s := range strs {
		values = append(values, s)
	}

	return r.Read(cb, values...)
}

func (r BatchIterator) Read(cb IteratorCallback, values ...interface{}) error {
	keysCount := len(values)
	startKey := 0
	endKey := 0

	for {
		endKey = startKey + r.batchSize
		if endKey >= keysCount {
			return cb(startKey, values[startKey:keysCount])
		}

		if err := cb(startKey, values[startKey:endKey]); err != nil {
			return err
		}
		startKey = endKey
	}
}

batch_iterator_test.go

package util

import (
	"github.com/stretchr/testify/require"
	"testing"
)

func TestRead(t *testing.T) {
	r := require.New(t)
	items := []interface{}{
		0x0, 0x1, 0x2, 0x3, 0x4, 0x5,
		0x6, 0x7, 0x8, 0x9, 0xA, 0xB,
		0xC, 0xD, 0xE, 0xF,
	}
	expect := [][]interface{}{
		{0x0, 0x1, 0x2},
		{0x3, 0x4, 0x5},
		{0x6, 0x7, 0x8},
		{0x9, 0xA, 0xB},
		{0xC, 0xD, 0xE},
		{0xF},
	}
	iter := NewBatchIterator(3)

	i := 0
	err := iter.Read(func(start int, v []interface{}) error {
		r.Equal(expect[i][0], start)
		r.Equal(expect[i], v)
		i++
		return nil
	}, items...)

	r.NoError(err)
}