Merge pull request #359 from ClickHouse/feat/proto-map-kv

feat(proto.map): add RowKV

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

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))
+	})
+}