Merge pull request #86 from go-faster/fix/proto-array-of-low-cord

fix(proto): set low cardinality index properly during encoding

proto/_golden/col_arr_of_low_cord_str.hex

@@ -0,0 +1,4 @@
+00000000  05 00 00 00 00 00 00 00  07 00 00 00 00 00 00 00  |................|
+00000010  00 06 00 00 00 00 00 00  03 00 00 00 00 00 00 00  |................|
+00000020  03 66 6f 6f 03 62 61 72  03 62 61 7a 07 00 00 00  |.foo.bar.baz....|
+00000030  00 00 00 00 00 01 00 00  02 00 02                 |...........|

proto/_golden/col_low_cardinality_of_str.hex

@@ -1,3 +1,3 @@
 00000000  00 06 00 00 00 00 00 00  03 00 00 00 00 00 00 00  |................|
 00000010  03 66 6f 6f 03 62 61 72  03 62 61 7a 05 00 00 00  |.foo.bar.baz....|
-00000020  00 00 00 00 00 00 00 00  00                       |.........|
+00000020  00 00 00 00 00 01 00 00  02                       |.........|

proto/col_arr_of.go

@@ -59,6 +59,16 @@ func (c *ColArrOf[T]) EncodeState(b *Buffer) {
 	}
 }
 
+// Prepare ensures Preparable column propagation.
+func (c *ColArrOf[T]) Prepare() error {
+	if v, ok := c.Data.(Preparable); ok {
+		if err := v.Prepare(); err != nil {
+			return errors.Wrap(err, "prepare data")
+		}
+	}
+	return nil
+}
+
 // RowAppend appends i-th row to target and returns it.
 func (c ColArrOf[T]) RowAppend(i int, target []T) []T {
 	var start int

proto/col_arr_of_test.go

@@ -62,7 +62,7 @@ func TestColArrOfStr(t *testing.T) {
 	var buf Buffer
 	col.EncodeColumn(&buf)
 	t.Run("Golden", func(t *testing.T) {
-		gold.Bytes(t, buf.Buf, "col_col_arr_of_str")
+		gold.Bytes(t, buf.Buf, "col_arr_of_str")
 	})
 	t.Run("Ok", func(t *testing.T) {
 		br := bytes.NewReader(buf.Buf)
@@ -85,3 +85,37 @@ func TestColArrOfStr(t *testing.T) {
 		requireNoShortRead(t, buf.Buf, colAware(dec, col.Rows()))
 	})
 }
+
+func TestArrOfLowCordStr(t *testing.T) {
+	col := ArrayOf[string](new(ColStr).LowCardinality())
+	col.Append([]string{"foo", "bar", "foo", "foo", "baz"})
+	col.Append([]string{"foo", "baz"})
+
+	require.NoError(t, col.Prepare())
+
+	var buf Buffer
+	col.EncodeColumn(&buf)
+	t.Run("Golden", func(t *testing.T) {
+		gold.Bytes(t, buf.Buf, "col_arr_of_low_cord_str")
+	})
+	t.Run("Ok", func(t *testing.T) {
+		br := bytes.NewReader(buf.Buf)
+		r := NewReader(br)
+		dec := ArrayOf[string](new(ColStr).LowCardinality())
+
+		require.NoError(t, dec.DecodeColumn(r, col.Rows()))
+		require.Equal(t, col.Rows(), dec.Rows())
+		require.Equal(t, ColumnType("Array(LowCardinality(String))"), dec.Type())
+		require.Equal(t, []string{"foo", "bar", "foo", "foo", "baz"}, dec.Row(0))
+		require.Equal(t, []string{"foo", "baz"}, dec.Row(1))
+	})
+	t.Run("ErrUnexpectedEOF", func(t *testing.T) {
+		r := NewReader(bytes.NewReader(nil))
+		dec := ArrayOf[string](new(ColStr).LowCardinality())
+		require.ErrorIs(t, dec.DecodeColumn(r, col.Rows()), io.ErrUnexpectedEOF)
+	})
+	t.Run("NoShortRead", func(t *testing.T) {
+		dec := ArrayOf[string](new(ColStr).LowCardinality())
+		requireNoShortRead(t, buf.Buf, colAware(dec, col.Rows()))
+	})
+}

proto/col_low_cardinality_of.go

@@ -248,6 +248,7 @@ func (c *ColLowCardinalityOf[T]) Prepare() error {
 		if !ok {
 			c.index.Append(v)
 			c.kv[v] = last
+			idx = last
 			last++
 		}
 		c.keys[i] = idx

proto/col_low_cardinality_of_test.go

@@ -5,6 +5,7 @@ import (
 	"io"
 	"testing"
 
+	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 
 	"github.com/go-faster/ch/internal/gold"
@@ -19,7 +20,8 @@ func TestLowCardinalityOf(t *testing.T) {
 
 func TestLowCardinalityOfStr(t *testing.T) {
 	col := (&ColStr{}).LowCardinality()
-	col.AppendArr([]string{"foo", "bar", "foo", "foo", "baz"})
+	v := []string{"foo", "bar", "foo", "foo", "baz"}
+	col.AppendArr(v)
 
 	require.NoError(t, col.Prepare())
 
@@ -35,7 +37,10 @@ func TestLowCardinalityOfStr(t *testing.T) {
 
 		require.NoError(t, dec.DecodeColumn(r, col.Rows()))
 		require.Equal(t, col.Rows(), dec.Rows())
-		require.Equal(t, ColumnType("LowCardinality(String)"), dec.Type())
+		for i, s := range v {
+			assert.Equal(t, s, col.Row(i))
+		}
+		assert.Equal(t, ColumnType("LowCardinality(String)"), dec.Type())
 	})
 	t.Run("ErrUnexpectedEOF", func(t *testing.T) {
 		r := NewReader(bytes.NewReader(nil))

query_test.go

@@ -478,6 +478,35 @@ func TestClient_Query(t *testing.T) {
 		require.Len(t, data, 1)
 		require.Equal(t, data, gotData)
 	})
+	t.Run("ArrayLowCardinality", func(t *testing.T) {
+		t.Parallel()
+		conn := Conn(t)
+		require.NoError(t, conn.Do(ctx, Query{
+			Body: "CREATE TABLE test_table (v Array(LowCardinality(String))) ENGINE = Memory",
+		}), "create table")
+
+		v := proto.ArrayOf[string](new(proto.ColStr).LowCardinality())
+		v.Append([]string{"foo", "bar"})
+		v.Append([]string{"baz"})
+
+		require.NoError(t, conn.Do(ctx, Query{
+			Body: "INSERT INTO test_table VALUES",
+			Input: []proto.InputColumn{
+				{Name: "v", Data: v},
+			},
+		}), "insert")
+
+		gotData := proto.ArrayOf[string](new(proto.ColStr).LowCardinality())
+		require.NoError(t, conn.Do(ctx, Query{
+			Body: "SELECT * FROM test_table",
+			Result: proto.Results{
+				{Name: "v", Data: gotData},
+			},
+		}), "select")
+
+		assert.Equal(t, []string{"foo", "bar"}, gotData.Row(0))
+		assert.Equal(t, []string{"baz"}, gotData.Row(1))
+	})
 	t.Run("InsertLowCardinalityString", func(t *testing.T) {
 		t.Parallel()
 		conn := Conn(t)