diff --git a/proto/col_map.go b/proto/col_map.go index cc462d53..90925fb2 100644 --- a/proto/col_map.go +++ b/proto/col_map.go @@ -80,6 +80,23 @@ func (c ColMap[K, V]) Row(i int) map[K]V { return m } +// RowKV returns a slice of KV[K, V] for a given row. +func (c ColMap[K, V]) RowKV(i int) []KV[K, V] { + var start int + end := int(c.Offsets[i]) + if i > 0 { + start = int(c.Offsets[i-1]) + } + v := make([]KV[K, V], 0, end-start) + for idx := start; idx < end; idx++ { + v = append(v, KV[K, V]{ + Key: c.Keys.Row(idx), + Value: c.Values.Row(idx), + }) + } + return v +} + // KV is a key-value pair. type KV[K comparable, V any] struct { Key K diff --git a/proto/col_map_test.go b/proto/col_map_test.go index d1fa83b7..a0a3317f 100644 --- a/proto/col_map_test.go +++ b/proto/col_map_test.go @@ -117,3 +117,57 @@ func TestColMap(t *testing.T) { requireNoShortRead(t, buf.Buf, colAware(dec, rows)) }) } + +func TestColMap_RowKV(t *testing.T) { + v := ColMap[string, string]{ + Keys: &ColStr{}, Values: &ColStr{}, + } + require.Equal(t, ColumnType("Map(String, String)"), v.Type()) + v.AppendKV([]KV[string, string]{ + {"foo", "bar"}, + {"baz", "hello"}, + }) + v.AppendKV([]KV[string, string]{ + {"like", "100"}, + {"dislike", "200"}, + }) + const rows = 2 + + var buf Buffer + v.EncodeColumn(&buf) + + t.Run("Ok", func(t *testing.T) { + br := bytes.NewReader(buf.Buf) + r := NewReader(br) + dec := &ColMap[string, string]{ + Keys: &ColStr{}, Values: &ColStr{}, + } + require.NoError(t, dec.DecodeColumn(r, rows)) + for i := 0; i < rows; i++ { + require.Equal(t, v.Row(i), v.Row(i)) + } + require.Equal(t, []KV[string, string]{ + {"foo", "bar"}, + {"baz", "hello"}, + }, dec.RowKV(0)) + require.Equal(t, []KV[string, string]{ + {"like", "100"}, + {"dislike", "200"}, + }, dec.RowKV(1)) + dec.Reset() + require.Equal(t, 0, dec.Rows()) + }) + t.Run("EOF", func(t *testing.T) { + r := NewReader(bytes.NewReader(nil)) + dec := &ColMap[string, string]{ + Keys: &ColStr{}, Values: &ColStr{}, + } + require.ErrorIs(t, dec.DecodeColumn(r, rows), io.EOF) + }) + t.Run("NoShortRead", func(t *testing.T) { + dec := &ColMap[string, string]{ + Keys: &ColStr{}, Values: &ColStr{}, + } + requireNoShortRead(t, buf.Buf, colAware(dec, rows)) + }) +}