Merge pull request #406 from tdakkota/feat/go1.23-seqs

proto/col_arr_go123.go

@@ -0,0 +1,22 @@
+//go:build go1.23
+
+package proto
+
+import "iter"
+
+// RowRange returns a [iter.Seq] iterator over i-th row.
+func (c ColArr[T]) RowRange(i int) iter.Seq[T] {
+	var start int
+	end := int(c.Offsets[i])
+	if i > 0 {
+		start = int(c.Offsets[i-1])
+	}
+
+	return func(yield func(T) bool) {
+		for idx := start; idx < end; idx++ {
+			if !yield(c.Data.Row(idx)) {
+				return
+			}
+		}
+	}
+}

proto/col_arr_go123_test.go

@@ -0,0 +1,37 @@
+//go:build go1.23
+
+package proto
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestColArrRange(t *testing.T) {
+	var (
+		enc      = new(ColStr).Array()
+		expected = [][]string{
+			{"foo", "bar", "foo", "foo", "baz"},
+			{"foo", "baz"},
+		}
+	)
+	enc.AppendArr(expected)
+
+	var buf Buffer
+	enc.EncodeColumn(&buf)
+
+	var (
+		dec = new(ColStr).Array()
+		got [][]string
+	)
+	require.NoError(t, dec.DecodeColumn(buf.Reader(), enc.Rows()))
+	for rowIdx := range dec.Rows() {
+		var row []string
+		for e := range dec.RowRange(rowIdx) {
+			row = append(row, e)
+		}
+		got = append(got, row)
+	}
+	require.Equal(t, expected, got)
+}

proto/col_map_go123.go

@@ -0,0 +1,25 @@
+//go:build go1.23
+
+package proto
+
+import "iter"
+
+// RowRange returns a [iter.Seq2] iterator over i-th row.
+func (c ColMap[K, V]) RowRange(i int) iter.Seq2[K, V] {
+	var start int
+	end := int(c.Offsets[i])
+	if i > 0 {
+		start = int(c.Offsets[i-1])
+	}
+
+	return func(yield func(K, V) bool) {
+		for idx := start; idx < end; idx++ {
+			if !yield(
+				c.Keys.Row(idx),
+				c.Values.Row(idx),
+			) {
+				return
+			}
+		}
+	}
+}

proto/col_map_go123_test.go

@@ -0,0 +1,48 @@
+//go:build go1.23
+
+package proto
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestColMapRange(t *testing.T) {
+	var (
+		enc = &ColMap[string, string]{
+			Keys:   &ColStr{},
+			Values: &ColStr{},
+		}
+		expected = []map[string]string{
+			{
+				"a": "b",
+				"c": "d",
+			},
+			{
+				"e": "f",
+			},
+		}
+	)
+	enc.AppendArr(expected)
+
+	var buf Buffer
+	enc.EncodeColumn(&buf)
+
+	var (
+		dec = &ColMap[string, string]{
+			Keys:   &ColStr{},
+			Values: &ColStr{},
+		}
+		got []map[string]string
+	)
+	require.NoError(t, dec.DecodeColumn(buf.Reader(), enc.Rows()))
+	for rowIdx := range dec.Rows() {
+		row := map[string]string{}
+		for k, v := range dec.RowRange(rowIdx) {
+			row[k] = v
+		}
+		got = append(got, row)
+	}
+	require.Equal(t, expected, got)
+}