Merge pull request #104 from go-faster/feat/col-map-of

feat(proto): implement Map(K, V)

go.mod

@@ -21,6 +21,7 @@ require (
 	go.uber.org/atomic v1.9.0
 	go.uber.org/multierr v1.8.0
 	go.uber.org/zap v1.21.0
+	golang.org/x/exp v0.0.0-20220518171630-0b5c67f07fdf
 	golang.org/x/sync v0.0.0-20210220032951-036812b2e83c
 )
 

go.sum

@@ -67,6 +67,8 @@ go.uber.org/zap v1.21.0 h1:WefMeulhovoZ2sYXz7st6K0sLj7bBhpiFaud4r4zST8=
 go.uber.org/zap v1.21.0/go.mod h1:wjWOCqI0f2ZZrJF/UufIOkiC8ii6tm1iqIsLo76RfJw=
 golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w=
 golang.org/x/crypto v0.0.0-20191011191535-87dc89f01550/go.mod h1:yigFU9vqHzYiE8UmvKecakEJjdnWj3jj499lnFckfCI=
+golang.org/x/exp v0.0.0-20220518171630-0b5c67f07fdf h1:oXVg4h2qJDd9htKxb5SCpFBHLipW6hXmL3qpUixS2jw=
+golang.org/x/exp v0.0.0-20220518171630-0b5c67f07fdf/go.mod h1:yh0Ynu2b5ZUe3MQfp2nM0ecK7wsgouWTDN0FNeJuIys=
 golang.org/x/lint v0.0.0-20190930215403-16217165b5de/go.mod h1:6SW0HCj/g11FgYtHlgUYUwCkIfeOF89ocIRzGO/8vkc=
 golang.org/x/mod v0.4.2/go.mod h1:s0Qsj1ACt9ePp/hMypM3fl4fZqREWJwdYDEqhRiZZUA=
 golang.org/x/mod v0.6.0-dev.0.20220106191415-9b9b3d81d5e3 h1:kQgndtyPBW/JIYERgdxfwMYh3AVStj88WQTlNDi2a+o=

proto/_golden/col_map_of_str_str.hex

@@ -0,0 +1,5 @@
+00000000  02 00 00 00 00 00 00 00  05 00 00 00 00 00 00 00  |................|
+00000010  03 62 61 7a 03 66 6f 6f  07 64 69 73 6c 69 6b 65  |.baz.foo.dislike|
+00000020  04 6c 69 6b 65 06 72 65  73 75 6c 74 05 68 65 6c  |.like.result.hel|
+00000030  6c 6f 03 62 61 72 03 32  30 30 03 31 30 30 08 31  |lo.bar.200.100.1|
+00000040  30 30 30 20 2d 20 37                              |000 - 7|

proto/col_map_of.go

@@ -1,33 +1,34 @@
 package proto
 
-import "github.com/go-faster/errors"
+import (
+	"github.com/go-faster/errors"
+	"golang.org/x/exp/constraints"
+	"golang.org/x/exp/maps"
+	"golang.org/x/exp/slices"
+)
 
 // Compile-time assertions for ColMapOf.
 var (
-	_ ColInput = ColMapOf[string, string]{
-		Keys:   &ColStr{},
-		Values: &ColStr{},
-	}
-	_ ColResult = &ColMapOf[string, string]{
-		Keys:   &ColStr{},
-		Values: &ColStr{},
-	}
-	_ Column = &ColMapOf[string, string]{
-		Keys:   &ColStr{},
-		Values: &ColStr{},
+	_ ColInput                    = (*ColMapOf[string, string])(nil)
+	_ ColResult                   = (*ColMapOf[string, string])(nil)
+	_ Column                      = (*ColMapOf[string, string])(nil)
+	_ ColumnOf[map[string]string] = (*ColMapOf[string, string])(nil)
+	_ StateEncoder                = (*ColMapOf[string, string])(nil)
+	_ StateDecoder                = (*ColMapOf[string, string])(nil)
+
+	_ = ColMapOf[int64, string]{
+		Keys:   new(ColInt64),
+		Values: new(ColStr),
 	}
 )
 
-type ColMapOf[K comparable, V any] struct {
+// ColMapOf implements Map(K, V) as ColumnOf[map[K]V].
+type ColMapOf[K constraints.Ordered, V any] struct {
 	Offsets ColUInt64
 	Keys    ColumnOf[K]
 	Values  ColumnOf[V]
 }
 
-func (c ColMapOf[K, V]) Get(k K) (v V, ok bool) {
-	return v, ok
-}
-
 func (c ColMapOf[K, V]) Type() ColumnType {
 	return ColumnTypeMap.Sub(c.Keys.Type(), c.Values.Type())
 }
@@ -36,6 +37,61 @@ func (c ColMapOf[K, V]) Rows() int {
 	return c.Offsets.Rows()
 }
 
+func (c *ColMapOf[K, V]) DecodeState(r *Reader) error {
+	if s, ok := c.Keys.(StateDecoder); ok {
+		if err := s.DecodeState(r); err != nil {
+			return errors.Wrap(err, "keys state")
+		}
+	}
+	if s, ok := c.Values.(StateDecoder); ok {
+		if err := s.DecodeState(r); err != nil {
+			return errors.Wrap(err, "values state")
+		}
+	}
+	return nil
+}
+
+func (c ColMapOf[K, V]) EncodeState(b *Buffer) {
+	if s, ok := c.Keys.(StateEncoder); ok {
+		s.EncodeState(b)
+	}
+	if s, ok := c.Values.(StateEncoder); ok {
+		s.EncodeState(b)
+	}
+}
+
+func (c ColMapOf[K, V]) Row(i int) map[K]V {
+	m := make(map[K]V)
+	var start int
+	end := int(c.Offsets[i])
+	if i > 0 {
+		start = int(c.Offsets[i-1])
+	}
+	for idx := start; idx < end; idx++ {
+		m[c.Keys.Row(idx)] = c.Values.Row(idx)
+	}
+	return m
+}
+
+func (c *ColMapOf[K, V]) Append(v map[K]V) {
+	// Make marshaling deterministic and sort map.
+	keys := maps.Keys(v)
+	slices.Sort(keys)
+
+	for _, k := range keys {
+		c.Keys.Append(k)
+		c.Values.Append(v[k])
+	}
+
+	c.Offsets.Append(uint64(c.Keys.Rows()))
+}
+
+func (c *ColMapOf[K, V]) AppendArr(v []map[K]V) {
+	for _, m := range v {
+		c.Append(m)
+	}
+}
+
 func (c *ColMapOf[K, V]) DecodeColumn(r *Reader, rows int) error {
 	if rows == 0 {
 		return nil

proto/col_map_of_test.go

@@ -1,16 +1,61 @@
 package proto
 
 import (
+	"bytes"
+	"io"
 	"testing"
 
 	"github.com/stretchr/testify/require"
+
+	"github.com/go-faster/ch/internal/gold"
 )
 
 func TestColMapOf(t *testing.T) {
 	v := ColMapOf[string, string]{
-		Keys:   &ColStr{},
-		Values: &ColStr{},
+		Keys: &ColStr{}, Values: &ColStr{},
 	}
-	_, _ = v.Get("foo")
 	require.Equal(t, ColumnType("Map(String, String)"), v.Type())
+	v.Append(map[string]string{
+		"foo": "bar",
+		"baz": "hello",
+	})
+	v.Append(map[string]string{
+		"like":    "100",
+		"dislike": "200",
+		"result":  "1000 - 7",
+	})
+	const rows = 2
+
+	var buf Buffer
+	v.EncodeColumn(&buf)
+
+	t.Run("Golden", func(t *testing.T) {
+		gold.Bytes(t, buf.Buf, "col_map_of_str_str")
+	})
+	t.Run("Ok", func(t *testing.T) {
+		br := bytes.NewReader(buf.Buf)
+		r := NewReader(br)
+		dec := &ColMapOf[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))
+		}
+		dec.Reset()
+		require.Equal(t, 0, dec.Rows())
+	})
+	t.Run("ErrUnexpectedEOF", func(t *testing.T) {
+		r := NewReader(bytes.NewReader(nil))
+		dec := &ColMapOf[string, string]{
+			Keys: &ColStr{}, Values: &ColStr{},
+		}
+		require.ErrorIs(t, dec.DecodeColumn(r, rows), io.ErrUnexpectedEOF)
+	})
+	t.Run("NoShortRead", func(t *testing.T) {
+		dec := &ColMapOf[string, string]{
+			Keys: &ColStr{}, Values: &ColStr{},
+		}
+		requireNoShortRead(t, buf.Buf, colAware(dec, rows))
+	})
 }

query_test.go

@@ -722,6 +722,56 @@ func TestClient_Query(t *testing.T) {
 		require.NoError(t, Conn(t).Do(ctx, selectRand))
 		require.Equal(t, numbers, total)
 	})
+	t.Run("InsertMapOfStringInt64", func(t *testing.T) {
+		t.Parallel()
+		conn := Conn(t)
+
+		// Create table, no data fetch.
+		createTable := Query{
+			Body: "CREATE TABLE test_table (v Map(String, Int64)) ENGINE = TinyLog",
+		}
+		require.NoError(t, conn.Do(ctx, createTable), "create table")
+
+		data := &proto.ColMapOf[string, int64]{
+			Keys:   new(proto.ColStr),
+			Values: new(proto.ColInt64),
+		}
+		data.AppendArr([]map[string]int64{
+			{
+				"foo": 1,
+				"bar": 100,
+			},
+			{
+				"clickhouse_1": -7,
+				"clickhouse_2": 130,
+				"clickhouse_3": 110,
+			},
+		})
+
+		insertQuery := Query{
+			Body: "INSERT INTO test_table VALUES",
+			Input: []proto.InputColumn{
+				{Name: "v", Data: data},
+			},
+		}
+		require.NoError(t, conn.Do(ctx, insertQuery), "insert")
+
+		gotData := &proto.ColMapOf[string, int64]{
+			Keys:   new(proto.ColStr),
+			Values: new(proto.ColInt64),
+		}
+		selectData := Query{
+			Body: "SELECT * FROM test_table",
+			Result: proto.Results{
+				{Name: "v", Data: gotData},
+			},
+		}
+		require.NoError(t, conn.Do(ctx, selectData), "select")
+		require.Equal(t, data.Rows(), gotData.Rows())
+		for i := 0; i < data.Rows(); i++ {
+			require.Equal(t, data.Row(i), gotData.Row(i))
+		}
+	})
 }
 
 func TestClientCompression(t *testing.T) {