dynamically change discriminator precision

Author: SpencerTorres

lib/column/dynamic.go

@@ -20,6 +20,7 @@ package column
 import (
 	"database/sql/driver"
 	"fmt"
+	"math"
 	"reflect"
 	"strings"
 	"time"
@@ -265,14 +266,27 @@ func (c *Dynamic) encodeHeader(buffer *proto.Buffer) error {
 	return nil
 }
 
+func discriminatorWriter(totalTypes int, buffer *proto.Buffer) func(int) {
+	switch {
+	case totalTypes <= math.MaxUint8:
+		return func(d int) { buffer.PutUInt8(uint8(d)) }
+	case totalTypes <= math.MaxUint16:
+		return func(d int) { buffer.PutUInt16(uint16(d)) }
+	case totalTypes <= math.MaxUint32:
+		return func(d int) { buffer.PutUInt32(uint32(d)) }
+	default:
+		return func(d int) { buffer.PutUInt64(uint64(d)) }
+	}
+}
+
 func (c *Dynamic) encodeData(buffer *proto.Buffer) {
-	// TODO: dynamically switch types based on size requirements
+	writeDiscriminator := discriminatorWriter(c.totalTypes, buffer)
 	for _, typeIndex := range c.discriminators {
 		if typeIndex == DynamicNullDiscriminator {
 			typeIndex = c.totalTypes
 		}
 
-		buffer.PutByte(uint8(typeIndex))
+		writeDiscriminator(typeIndex)
 	}
 
 	for _, col := range c.columns {
@@ -344,18 +358,42 @@ func (c *Dynamic) decodeHeader(reader *proto.Reader) error {
 	return nil
 }
 
+func discriminatorReader(totalTypes int, reader *proto.Reader) func() (int, error) {
+	switch {
+	case totalTypes <= math.MaxUint8:
+		return func() (int, error) {
+			v, err := reader.UInt8()
+			return int(v), err
+		}
+	case totalTypes <= math.MaxUint16:
+		return func() (int, error) {
+			v, err := reader.UInt16()
+			return int(v), err
+		}
+	case totalTypes <= math.MaxUint32:
+		return func() (int, error) {
+			v, err := reader.UInt32()
+			return int(v), err
+		}
+	default:
+		return func() (int, error) {
+			v, err := reader.UInt64()
+			return int(v), err
+		}
+	}
+}
+
 func (c *Dynamic) decodeData(reader *proto.Reader, rows int) error {
 	c.discriminators = make([]int, rows)
 	c.offsets = make([]int, rows)
 	rowCountByType := make([]int, len(c.columns))
 
+	readDiscriminator := discriminatorReader(c.totalTypes, reader)
 	for i := 0; i < rows; i++ {
-		// TODO: dynamically switch types based on size requirements
-		discByte, err := reader.ReadByte()
+		disc, err := readDiscriminator()
 		if err != nil {
 			return fmt.Errorf("failed to read discriminator at index %d: %w", i, err)
 		}
-		disc := int(discByte)
 
 		c.discriminators[i] = disc
 		if disc != c.totalTypes {

tests/dynamic_test.go

@@ -19,6 +19,7 @@ package tests
 
 import (
 	"context"
+	"fmt"
 	"testing"
 	"time"
 
@@ -202,6 +203,46 @@ func TestDynamicMaxTypes(t *testing.T) {
 	})
 }
 
+// Discriminator precision must grow dynamically depending on the number of types within the Dynamic.
+// This test confirms that we can go beyond UInt8/255 types.
+func TestDynamicExceededTypes(t *testing.T) {
+	conn := setupDynamicTest(t, clickhouse.Native)
+	ctx := context.Background()
+
+	const ddl = `
+		CREATE TABLE IF NOT EXISTS test_dynamic_exceeded_types (
+			  c Dynamic
+		) 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_dynamic_exceeded_types"))
+	}()
+
+	testTypeCount := func(typeCount int) func(t *testing.T) {
+		return func(t *testing.T) {
+			batch, err := conn.PrepareBatch(ctx, "INSERT INTO test_dynamic_exceeded_types (c)")
+			require.NoError(t, err)
+
+			for i := 0; i < typeCount; i++ {
+				typeName := fmt.Sprintf("Tuple(\"%d\" Int64)", i)
+				require.NoError(t, batch.Append(clickhouse.NewDynamicWithType([]int64{int64(i)}, typeName)))
+			}
+			require.NoError(t, batch.Send())
+
+			rows, err := conn.Query(ctx, "SELECT c FROM test_dynamic_exceeded_types")
+			require.NoError(t, err)
+
+			require.NoError(t, rows.Close())
+			require.NoError(t, rows.Err())
+		}
+	}
+
+	t.Run("less than UInt8", testTypeCount(16))
+	t.Run("UInt8 bounds", testTypeCount(255))
+	t.Run("UInt16 range", testTypeCount(300))
+}
+
 func TestDynamicArray(t *testing.T) {
 	TestProtocols(t, func(t *testing.T, protocol clickhouse.Protocol) {
 		conn := setupDynamicTest(t, protocol)

tests/std/dynamic_test.go

@@ -40,7 +40,7 @@ func setupDynamicTest(t *testing.T) *sql.DB {
 	})
 	require.NoError(t, err)
 
-	if !CheckMinServerVersion(conn, 24, 8, 0) {
+	if !CheckMinServerVersion(conn, 25, 6, 0) {
 		t.Skip(fmt.Errorf("unsupported clickhouse version for Dynamic type"))
 		return nil
 	}