Optimize repeated row iterator

pkg/iter/batch_async.go

@@ -0,0 +1,119 @@
+package iter
+
+type AsyncBatchIterator[T, N any] struct {
+	idx      int
+	batch    []N
+	buffered []N
+
+	close    chan struct{}
+	done     chan struct{}
+	c        chan batch[N]
+	delegate Iterator[T]
+
+	clone   func(T) N
+	release func([]N)
+}
+
+type batch[T any] struct {
+	buffered []T
+	done     chan struct{}
+}
+
+const minBatchSize = 64
+
+func NewAsyncBatchIterator[T, N any](
+	iterator Iterator[T],
+	size int,
+	clone func(T) N,
+	release func([]N),
+) *AsyncBatchIterator[T, N] {
+	if size == 0 {
+		size = minBatchSize
+	}
+	x := &AsyncBatchIterator[T, N]{
+		idx:      -1,
+		batch:    make([]N, 0, size),
+		buffered: make([]N, 0, size),
+		close:    make(chan struct{}),
+		done:     make(chan struct{}),
+		c:        make(chan batch[N]),
+		clone:    clone,
+		release:  release,
+		delegate: iterator,
+	}
+	go x.iterate()
+	return x
+}
+
+func (x *AsyncBatchIterator[T, N]) Next() bool {
+	if x.idx < 0 || x.idx >= len(x.batch)-1 {
+		if !x.loadBatch() {
+			return false
+		}
+	}
+	x.idx++
+	return true
+}
+
+func (x *AsyncBatchIterator[T, N]) At() N { return x.batch[x.idx] }
+
+func (x *AsyncBatchIterator[T, N]) iterate() {
+	defer func() {
+		close(x.c)
+		close(x.done)
+	}()
+	for x.fillBuffer() {
+		b := batch[N]{
+			buffered: x.buffered,
+			done:     make(chan struct{}),
+		}
+		select {
+		case x.c <- b:
+			// Wait for the next loadBatch call.
+			<-b.done
+		case <-x.close:
+			return
+		}
+	}
+}
+
+func (x *AsyncBatchIterator[T, N]) loadBatch() bool {
+	var b batch[N]
+	select {
+	case b = <-x.c:
+	case <-x.done:
+	}
+	if len(b.buffered) == 0 {
+		return false
+	}
+	// Swap buffers and signal "iterate" goroutine
+	// that x.buffered can be used: it will
+	// immediately start filling the buffer.
+	x.buffered, x.batch = x.batch, b.buffered
+	x.idx = -1
+	close(b.done)
+	return true
+}
+
+func (x *AsyncBatchIterator[T, N]) fillBuffer() bool {
+	x.buffered = x.buffered[:cap(x.buffered)]
+	x.release(x.buffered)
+	for i := range x.buffered {
+		if !x.delegate.Next() {
+			x.buffered = x.buffered[:i]
+			break
+		}
+		x.buffered[i] = x.clone(x.delegate.At())
+	}
+	return len(x.buffered) > 0
+}
+
+func (x *AsyncBatchIterator[T, N]) Close() error {
+	close(x.close)
+	<-x.done
+	return x.delegate.Close()
+}
+
+func (x *AsyncBatchIterator[T, N]) Err() error {
+	return x.delegate.Err()
+}

pkg/iter/batch_async_test.go

@@ -0,0 +1,116 @@
+package iter
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func Test_AsyncBatchIter(t *testing.T) {
+	type testCase struct {
+		description string
+		seqSize     int
+		bufSize     int
+	}
+	testCases := []testCase{
+		{
+			description: "empty iterator",
+			seqSize:     0,
+			bufSize:     1,
+		},
+		{
+			description: "empty iterator, zero buffer",
+			seqSize:     0,
+			bufSize:     0,
+		},
+		{
+			description: "zero buffer",
+			seqSize:     10,
+			bufSize:     0,
+		},
+		{
+			description: "iterator < buffer",
+			seqSize:     5,
+			bufSize:     10,
+		},
+		{
+			description: "iterator == buffer",
+			seqSize:     10,
+			bufSize:     10,
+		},
+		{
+			description: "iterator > buffer",
+			seqSize:     25,
+			bufSize:     10,
+		},
+	}
+	for _, tc := range testCases {
+		tc := tc
+		t.Run(tc.description, func(t *testing.T) {
+			x := NewAsyncBatchIterator[int, int](
+				newSeqIterator(tc.seqSize),
+				tc.bufSize,
+				func(i int) int { return i },
+				func([]int) {},
+			)
+			var p, c int
+			for x.Next() {
+				i := x.At()
+				require.Equal(t, 1, i-p)
+				p = i
+				c++
+			}
+			require.Equal(t, tc.seqSize, c)
+			require.NoError(t, x.Err())
+			require.NoError(t, x.Close())
+		})
+	}
+}
+
+type seqIterator struct{ n, c int }
+
+func newSeqIterator(n int) *seqIterator {
+	return &seqIterator{n: n}
+}
+
+func (x *seqIterator) Next() bool {
+	if x.c < x.n {
+		x.c++
+		return true
+	}
+	return false
+}
+
+func (x *seqIterator) At() int { return x.c }
+
+func (x *seqIterator) Close() error { return nil }
+func (x *seqIterator) Err() error   { return nil }
+
+// Benchmark_AsyncBatchIterator-10    	   91417	     13353 ns/op	   17017 B/op	      10 allocs/op
+func Benchmark_AsyncBatchIterator(b *testing.B) {
+	b.ReportAllocs()
+	var n int
+	for i := 0; i < b.N; i++ {
+		x := NewAsyncBatchIterator[int, int](
+			newSeqIterator(1<<20),
+			1<<10,
+			func(i int) int { return i },
+			func([]int) {},
+		)
+		for x.Next() {
+			n += x.At()
+		}
+	}
+}
+
+// Benchmark_BufferedIterator-10    	      12	  99730976 ns/op	   10047 B/op	       8 allocs/op
+func Benchmark_BufferedIterator(b *testing.B) {
+	b.ReportAllocs()
+	var n int
+	for i := 0; i < b.N; i++ {
+		x := NewBufferedIterator[int](newSeqIterator(1<<20), 1<<10)
+		for x.Next() {
+			n += x.At()
+		}
+	}
+}

pkg/iter/iter.go

@@ -152,6 +152,9 @@ func (i *sliceIterator[A]) Close() error {
 }
 
 func Slice[T any](it Iterator[T]) ([]T, error) {
+	if s, ok := it.(*sliceIterator[T]); ok {
+		return s.list, nil
+	}
 	var result []T
 	defer it.Close()
 	for it.Next() {