Merge pull request #1591 from ClickHouse/variant_fix_prefix

Fix variant prefix

Author: SpencerTorres

lib/column/variant.go

@@ -263,8 +263,18 @@ func (c *Variant) AppendRow(v any) error {
 	return fmt.Errorf("value \"%v\" cannot be stored in variant column: no compatible types", v)
 }
 
-func (c *Variant) encodeHeader(buffer *proto.Buffer) {
+func (c *Variant) encodeHeader(buffer *proto.Buffer) error {
 	buffer.PutUInt64(SupportedVariantSerializationVersion)
+
+	for _, col := range c.columns {
+		if serialize, ok := col.(CustomSerialization); ok {
+			if err := serialize.WriteStatePrefix(buffer); err != nil {
+				return fmt.Errorf("failed to write prefix for type %s in variant: %w", col.Type(), err)
+			}
+		}
+	}
+
+	return nil
 }
 
 func (c *Variant) encodeData(buffer *proto.Buffer) {
@@ -276,9 +286,7 @@ func (c *Variant) encodeData(buffer *proto.Buffer) {
 }
 
 func (c *Variant) WriteStatePrefix(buffer *proto.Buffer) error {
-	c.encodeHeader(buffer)
-
-	return nil
+	return c.encodeHeader(buffer)
 }
 
 func (c *Variant) Encode(buffer *proto.Buffer) {

tests/dynamic_test.go

@@ -76,6 +76,8 @@ func TestDynamic(t *testing.T) {
 		require.NoError(t, batch.Append(colNil))
 		colSliceUInt8 := []uint8{0xA, 0xB, 0xC}
 		require.NoError(t, batch.Append(clickhouse.NewDynamicWithType(colSliceUInt8, "Array(UInt8)")))
+		colSliceSliceUInt8 := [][]uint8{{0xA, 0xB}, {0xC, 0xD}}
+		require.NoError(t, batch.Append(clickhouse.NewDynamicWithType(colSliceSliceUInt8, "Array(Array(UInt8))")))
 		colSliceMapStringString := []map[string]string{{"key1": "value1", "key2": "value2"}, {"key3": "value3"}}
 		require.NoError(t, batch.Append(clickhouse.NewDynamicWithType(colSliceMapStringString, "Array(Map(String, String))")))
 		colMapStringString := map[string]string{"key1": "value1", "key2": "value2"}
@@ -119,6 +121,11 @@ func TestDynamic(t *testing.T) {
 		require.NoError(t, err)
 		require.Equal(t, colSliceUInt8, row.Any())
 
+		require.True(t, rows.Next())
+		err = rows.Scan(&row)
+		require.NoError(t, err)
+		require.Equal(t, colSliceSliceUInt8, row.Any())
+
 		require.True(t, rows.Next())
 		err = rows.Scan(&row)
 		require.NoError(t, err)

tests/variant_test.go

@@ -156,6 +156,43 @@ func TestVariant(t *testing.T) {
 	})
 }
 
+func TestVariantPrefix(t *testing.T) {
+	TestProtocols(t, func(t *testing.T, protocol clickhouse.Protocol) {
+		conn := setupVariantTest(t, protocol)
+		ctx := context.Background()
+
+		const ddl = `
+			CREATE TABLE IF NOT EXISTS test_variant_prefix (
+				  c Variant(LowCardinality(String))                  
+			) Engine = MergeTree() ORDER BY tuple()
+		`
+		require.NoError(t, conn.Exec(ctx, ddl))
+		defer func() {
+			require.NoError(t, conn.Exec(ctx, "DROP TABLE IF EXISTS test_variant_prefix"))
+		}()
+
+		batch, err := conn.PrepareBatch(ctx, "INSERT INTO test_variant_prefix (c)")
+		require.NoError(t, err)
+
+		val := "a"
+		require.NoError(t, batch.Append(clickhouse.NewVariantWithType(val, "LowCardinality(String)")))
+		require.NoError(t, batch.Send())
+
+		rows, err := conn.Query(ctx, "SELECT c FROM test_variant_prefix")
+		require.NoError(t, err)
+
+		var row clickhouse.Variant
+
+		require.True(t, rows.Next())
+		err = rows.Scan(&row)
+		require.NoError(t, err)
+		require.Equal(t, val, row.Any())
+
+		require.NoError(t, rows.Close())
+		require.NoError(t, rows.Err())
+	})
+}
+
 func TestVariantArray(t *testing.T) {
 	TestProtocols(t, func(t *testing.T, protocol clickhouse.Protocol) {
 		conn := setupVariantTest(t, protocol)