Skip to content

Commit

Permalink
Merge pull request #406 from tdakkota/feat/go1.23-seqs
Browse files Browse the repository at this point in the history
  • Loading branch information
ernado authored Aug 13, 2024
2 parents 0f7e570 + 823c31a commit b6b0921
Show file tree
Hide file tree
Showing 4 changed files with 132 additions and 0 deletions.
22 changes: 22 additions & 0 deletions proto/col_arr_go123.go
Original file line number Diff line number Diff line change
@@ -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
}
}
}
}
37 changes: 37 additions & 0 deletions proto/col_arr_go123_test.go
Original file line number Diff line number Diff line change
@@ -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)
}
25 changes: 25 additions & 0 deletions proto/col_map_go123.go
Original file line number Diff line number Diff line change
@@ -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
}
}
}
}
48 changes: 48 additions & 0 deletions proto/col_map_go123_test.go
Original file line number Diff line number Diff line change
@@ -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)
}

0 comments on commit b6b0921

Please sign in to comment.