From 823c31a0d19ed4c65bdd5fe2a87493f772ededac Mon Sep 17 00:00:00 2001
From: tdakkota <tanc13@yandex.ru>
Date: Fri, 7 Jun 2024 22:02:47 +0300
Subject: [PATCH] feat(proto): add go1.23 iterators over `ColMap` and `ColArr`

---
 proto/col_arr_go123.go      | 22 +++++++++++++++++
 proto/col_arr_go123_test.go | 37 ++++++++++++++++++++++++++++
 proto/col_map_go123.go      | 25 +++++++++++++++++++
 proto/col_map_go123_test.go | 48 +++++++++++++++++++++++++++++++++++++
 4 files changed, 132 insertions(+)
 create mode 100644 proto/col_arr_go123.go
 create mode 100644 proto/col_arr_go123_test.go
 create mode 100644 proto/col_map_go123.go
 create mode 100644 proto/col_map_go123_test.go

diff --git a/proto/col_arr_go123.go b/proto/col_arr_go123.go
new file mode 100644
index 00000000..c421b419
--- /dev/null
+++ b/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
+			}
+		}
+	}
+}
diff --git a/proto/col_arr_go123_test.go b/proto/col_arr_go123_test.go
new file mode 100644
index 00000000..bd648b01
--- /dev/null
+++ b/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)
+}
diff --git a/proto/col_map_go123.go b/proto/col_map_go123.go
new file mode 100644
index 00000000..e9d83c58
--- /dev/null
+++ b/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
+			}
+		}
+	}
+}
diff --git a/proto/col_map_go123_test.go b/proto/col_map_go123_test.go
new file mode 100644
index 00000000..cb5a5c9d
--- /dev/null
+++ b/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)
+}