diff --git a/README.md b/README.md index f949c1b3..b914e5c3 100644 --- a/README.md +++ b/README.md @@ -259,7 +259,8 @@ colV.Reset() ## Features * OpenTelemetry support * No reflection or `interface{}` -* Generics (go1.18) for `Array[T]`, `LowCardinaliy[T]`, `Map[K, V]`, `Nullable[T]` +* Generics (go1.18) for `Array[T]`, `LowCardinality[T]`, `Map[K, V]`, `Nullable[T]` +* Nested type support with helper methods for flattened array columns * [Reading or writing](#dumps) ClickHouse dumps in `Native` format * **Column**-oriented design that operates directly with **blocks** of data * [Dramatically more efficient](https://github.com/ClickHouse/ch-bench) @@ -300,6 +301,7 @@ colV.Reset() * Bool * Tuple(T1, T2, ..., Tn) * Nullable(T) +* Nested(T1, T2, ...) * Point * Nothing, Interval @@ -350,6 +352,59 @@ q := ch.Query{ arr.Row(0) // ["foo", "bar", "baz"] ``` +### Nested + +Helper for `Nested(name1 Type1, name2 Type2, ...)`. In ClickHouse, Nested columns are stored as parallel arrays with dot notation (e.g., `col.field`). + +```go +// For table: +// CREATE TABLE t ( +// id UInt64, +// events Nested(name String, value Float64) +// ) ENGINE = Memory + +// Writing data +nested := proto.NewNested( + proto.NestedColumn{Name: "name", Data: new(proto.ColStr).Array()}, + proto.NestedColumn{Name: "value", Data: new(proto.ColFloat64).Array()}, +) +// Append a row with 2 events +nested.Append(map[string]any{ + "name": []string{"click", "view"}, + "value": []float64{1.5, 2.5}, +}) + +var id proto.ColUInt64 +id.Append(1) + +input := proto.Input{{Name: "id", Data: id}} +input = append(input, nested.InputColumns("events")...) + +if err := conn.Do(ctx, ch.Query{ + Body: "INSERT INTO t VALUES", + Input: input, +}); err != nil { + panic(err) +} + +// Reading data +var resultID proto.ColUInt64 +resultNested := proto.NewNested( + proto.NestedColumn{Name: "name", Data: new(proto.ColStr).Array()}, + proto.NestedColumn{Name: "value", Data: new(proto.ColFloat64).Array()}, +) + +results := proto.Results{{Name: "id", Data: &resultID}} +results = append(results, resultNested.ResultColumns("events")...) + +if err := conn.Do(ctx, ch.Query{ + Body: "SELECT id, `events.name`, `events.value` FROM t", + Result: results, +}); err != nil { + panic(err) +} +``` + ## Dumps ### Reading @@ -455,7 +510,7 @@ func TestLocalNativeDump(t *testing.T) { - [ ] AggregateFunction - [x] Nothing - [x] Interval - - [ ] Nested + - [x] Nested - [ ] [Geo types](https://clickhouse.com/docs/en/sql-reference/data-types/geo/) - [x] Point - [ ] Ring diff --git a/cht/local_test.go b/cht/local_test.go index d94cb632..116ad214 100644 --- a/cht/local_test.go +++ b/cht/local_test.go @@ -72,15 +72,16 @@ func TestLocalNativeDump(t *testing.T) { Rows int `json:"rows"` Data []struct { Title string `json:"title"` - Data int `json:"data,string"` + Data int `json:"data"` } }{} - require.NoError(t, json.Unmarshal(out.Bytes(), &v), "json") + val := out.Bytes() + require.NoError(t, json.Unmarshal(val, &v), "json") assert.Equal(t, 2, v.Rows) if assert.Len(t, v.Data, 2) { for i, r := range []struct { Title string `json:"title"` - Data int `json:"data,string"` + Data int `json:"data"` }{ {"Foo", 1}, {"Bar", 2}, diff --git a/nested_test.go b/nested_test.go new file mode 100644 index 00000000..8ca27f13 --- /dev/null +++ b/nested_test.go @@ -0,0 +1,773 @@ +package ch + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "os" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/ClickHouse/ch-go/proto" +) + +// Environment variables for ClickHouse connection: +// +// CH_HOST - ClickHouse host (default: localhost) +// CH_PORT - ClickHouse port (default: 9000) +// CH_USER - Username (default: default) +// CH_PASSWORD - Password (default: empty) +// CH_DATABASE - Database (default: default) + +func envOr(key, defaultVal string) string { + if v := os.Getenv(key); v != "" { + return v + } + return defaultVal +} + +// skipIfNoClickHouse skips the test if CH_HOST is not set. +func skipIfNoClickHouse(tb testing.TB) { + tb.Helper() + if os.Getenv("CH_HOST") == "" { + tb.Skip("CH_HOST not set - skipping (need ClickHouse connection)") + } +} + +// getTestClient returns a connected ClickHouse client using environment variables. +// Skips the test if CH_HOST is not set. +func getTestClient(tb testing.TB) *Client { + tb.Helper() + skipIfNoClickHouse(tb) + + ctx := context.Background() + addr := fmt.Sprintf("%s:%s", envOr("CH_HOST", "localhost"), envOr("CH_PORT", "9000")) + + settings := []Setting{ + {Key: "allow_experimental_object_type", Value: "1", Important: true}, + {Key: "output_format_native_write_json_as_string", Value: "1", Important: true}, + } + + client, err := Dial(ctx, Options{ + Address: addr, + User: envOr("CH_USER", "default"), + Password: envOr("CH_PASSWORD", ""), + Database: envOr("CH_DATABASE", "default"), + Settings: settings, + }) + require.NoError(tb, err, "failed to connect to ClickHouse at %s", addr) + + tb.Cleanup(func() { _ = client.Close() }) + return client +} + +func TestNestedBasic(t *testing.T) { + ctx := context.Background() + conn := getTestClient(t) + + // Create table with Nested column + require.NoError(t, conn.Do(ctx, Query{ + Body: `CREATE TABLE IF NOT EXISTS test_nested ( + id UInt64, + events Nested( + event_id UInt32, + event_name String + ) + ) ENGINE = Memory`, + })) + t.Cleanup(func() { + _ = conn.Do(ctx, Query{Body: "DROP TABLE IF EXISTS test_nested"}) + }) + + // Truncate for clean test + _ = conn.Do(ctx, Query{Body: "TRUNCATE TABLE test_nested"}) + + // Prepare data for INSERT + // Nested columns are sent as separate Array columns with dot notation + eventIds := new(proto.ColUInt32).Array() + eventNames := new(proto.ColStr).Array() + + // Row 1: events = [(1, "click"), (2, "view")] + eventIds.Append([]uint32{1, 2}) + eventNames.Append([]string{"click", "view"}) + + // Row 2: events = [(3, "purchase")] + eventIds.Append([]uint32{3}) + eventNames.Append([]string{"purchase"}) + + // Row 3: events = [] (empty) + eventIds.Append([]uint32{}) + eventNames.Append([]string{}) + + // INSERT using flattened columns + require.NoError(t, conn.Do(ctx, Query{ + Body: "INSERT INTO test_nested VALUES", + Input: proto.Input{ + {Name: "id", Data: proto.ColUInt64{1, 2, 3}}, + {Name: "events.event_id", Data: eventIds}, + {Name: "events.event_name", Data: eventNames}, + }, + })) + + // SELECT and verify + var ( + resultId proto.ColUInt64 + resultEventIds = new(proto.ColUInt32).Array() + resultEventNames = new(proto.ColStr).Array() + ) + + require.NoError(t, conn.Do(ctx, Query{ + Body: "SELECT id, `events.event_id`, `events.event_name` FROM test_nested ORDER BY id", + Result: proto.Results{ + {Name: "id", Data: &resultId}, + {Name: "events.event_id", Data: resultEventIds}, + {Name: "events.event_name", Data: resultEventNames}, + }, + })) + + // Verify results + require.Equal(t, 3, resultId.Rows()) + require.Equal(t, 3, resultEventIds.Rows()) + require.Equal(t, 3, resultEventNames.Rows()) + + // Row 1 + assert.Equal(t, uint64(1), resultId.Row(0)) + assert.Equal(t, []uint32{1, 2}, resultEventIds.Row(0)) + assert.Equal(t, []string{"click", "view"}, resultEventNames.Row(0)) + + // Row 2 + assert.Equal(t, uint64(2), resultId.Row(1)) + assert.Equal(t, []uint32{3}, resultEventIds.Row(1)) + assert.Equal(t, []string{"purchase"}, resultEventNames.Row(1)) + + // Row 3 (empty) - empty arrays may be returned as nil + assert.Equal(t, uint64(3), resultId.Row(2)) + assert.Empty(t, resultEventIds.Row(2)) + assert.Empty(t, resultEventNames.Row(2)) +} + +func TestNestedWithHelper(t *testing.T) { + ctx := context.Background() + conn := getTestClient(t) + + // Create table with Nested column + require.NoError(t, conn.Do(ctx, Query{ + Body: `CREATE TABLE IF NOT EXISTS test_nested_helper ( + id UInt64, + tags Nested( + name String, + value Float64 + ) + ) ENGINE = Memory`, + })) + t.Cleanup(func() { + _ = conn.Do(ctx, Query{Body: "DROP TABLE IF EXISTS test_nested_helper"}) + }) + + _ = conn.Do(ctx, Query{Body: "TRUNCATE TABLE test_nested_helper"}) + + // Use ColNested helper + nested := proto.NewNested( + proto.NestedColumn{Name: "name", Data: new(proto.ColStr).Array()}, + proto.NestedColumn{Name: "value", Data: new(proto.ColFloat64).Array()}, + ) + + // Append rows using the helper + require.NoError(t, nested.Append(map[string]any{ + "name": []string{"tag1", "tag2"}, + "value": []float64{1.5, 2.5}, + })) + require.NoError(t, nested.Append(map[string]any{ + "name": []string{"single"}, + "value": []float64{3.5}, + })) + + // INSERT using InputColumns helper + idCol := proto.ColUInt64{100, 200} + input := proto.Input{{Name: "id", Data: idCol}} + input = append(input, nested.InputColumns("tags")...) + + require.NoError(t, conn.Do(ctx, Query{ + Body: "INSERT INTO test_nested_helper VALUES", + Input: input, + })) + + // SELECT using ResultColumns helper + resultNested := proto.NewNested( + proto.NestedColumn{Name: "name", Data: new(proto.ColStr).Array()}, + proto.NestedColumn{Name: "value", Data: new(proto.ColFloat64).Array()}, + ) + + var resultId proto.ColUInt64 + results := proto.Results{{Name: "id", Data: &resultId}} + results = append(results, resultNested.ResultColumns("tags")...) + + require.NoError(t, conn.Do(ctx, Query{ + Body: "SELECT id, `tags.name`, `tags.value` FROM test_nested_helper ORDER BY id", + Result: results, + })) + + // Verify + require.Equal(t, 2, resultId.Rows()) + assert.Equal(t, uint64(100), resultId.Row(0)) + assert.Equal(t, uint64(200), resultId.Row(1)) + + // Access data through the nested helper + nameCol := resultNested.Column("name") + require.NotNil(t, nameCol) + nameArr := nameCol.Data.(*proto.ColArr[string]) + assert.Equal(t, []string{"tag1", "tag2"}, nameArr.Row(0)) + assert.Equal(t, []string{"single"}, nameArr.Row(1)) +} + +func TestNestedMultipleColumns(t *testing.T) { + ctx := context.Background() + conn := getTestClient(t) + + // Create table with multiple Nested columns + require.NoError(t, conn.Do(ctx, Query{ + Body: `CREATE TABLE IF NOT EXISTS test_nested_multi ( + id UInt64, + users Nested( + user_id UInt32, + name String + ), + events Nested( + event_type String, + timestamp Int64 + ) + ) ENGINE = Memory`, + })) + t.Cleanup(func() { + _ = conn.Do(ctx, Query{Body: "DROP TABLE IF EXISTS test_nested_multi"}) + }) + + _ = conn.Do(ctx, Query{Body: "TRUNCATE TABLE test_nested_multi"}) + + // Prepare users nested + usersIds := new(proto.ColUInt32).Array() + usersNames := new(proto.ColStr).Array() + usersIds.Append([]uint32{1, 2}) + usersNames.Append([]string{"alice", "bob"}) + + // Prepare events nested + eventTypes := new(proto.ColStr).Array() + eventTimestamps := new(proto.ColInt64).Array() + eventTypes.Append([]string{"login", "click", "logout"}) + eventTimestamps.Append([]int64{1000, 2000, 3000}) + + require.NoError(t, conn.Do(ctx, Query{ + Body: "INSERT INTO test_nested_multi VALUES", + Input: proto.Input{ + {Name: "id", Data: proto.ColUInt64{1}}, + {Name: "users.user_id", Data: usersIds}, + {Name: "users.name", Data: usersNames}, + {Name: "events.event_type", Data: eventTypes}, + {Name: "events.timestamp", Data: eventTimestamps}, + }, + })) + + // SELECT and verify + var ( + resultId proto.ColUInt64 + resultUsersIds = new(proto.ColUInt32).Array() + resultUsersNames = new(proto.ColStr).Array() + resultEventTypes = new(proto.ColStr).Array() + resultEventTimestamps = new(proto.ColInt64).Array() + ) + + require.NoError(t, conn.Do(ctx, Query{ + Body: "SELECT * FROM test_nested_multi", + Result: proto.Results{ + {Name: "id", Data: &resultId}, + {Name: "users.user_id", Data: resultUsersIds}, + {Name: "users.name", Data: resultUsersNames}, + {Name: "events.event_type", Data: resultEventTypes}, + {Name: "events.timestamp", Data: resultEventTimestamps}, + }, + })) + + require.Equal(t, 1, resultId.Rows()) + assert.Equal(t, []uint32{1, 2}, resultUsersIds.Row(0)) + assert.Equal(t, []string{"alice", "bob"}, resultUsersNames.Row(0)) + assert.Equal(t, []string{"login", "click", "logout"}, resultEventTypes.Row(0)) + assert.Equal(t, []int64{1000, 2000, 3000}, resultEventTimestamps.Row(0)) +} + +func TestNestedTypeInference(t *testing.T) { + // Test that ColNested.Infer works correctly + nested := &proto.ColNested{} + require.NoError(t, nested.Infer("Nested(id UInt64, name String, value Float64)")) + + assert.Equal(t, "Nested(id UInt64, name String, value Float64)", nested.Type().String()) + + cols := nested.Columns() + require.Len(t, cols, 3) + assert.Equal(t, "id", cols[0].Name) + assert.Equal(t, "name", cols[1].Name) + assert.Equal(t, "value", cols[2].Name) + + // Verify internal types are Array(T) + assert.Equal(t, "Array(UInt64)", cols[0].Data.Type().String()) + assert.Equal(t, "Array(String)", cols[1].Data.Type().String()) + assert.Equal(t, "Array(Float64)", cols[2].Data.Type().String()) +} + +func TestNestedLargeData(t *testing.T) { + ctx := context.Background() + conn := getTestClient(t) + + require.NoError(t, conn.Do(ctx, Query{ + Body: `CREATE TABLE IF NOT EXISTS test_nested_large ( + id UInt64, + data Nested( + idx UInt32, + val String + ) + ) ENGINE = Memory`, + })) + t.Cleanup(func() { + _ = conn.Do(ctx, Query{Body: "DROP TABLE IF EXISTS test_nested_large"}) + }) + + _ = conn.Do(ctx, Query{Body: "TRUNCATE TABLE test_nested_large"}) + + // Generate large data + const numRows = 1000 + const elementsPerRow = 100 + + idCol := make(proto.ColUInt64, numRows) + idxCol := new(proto.ColUInt32).Array() + valCol := new(proto.ColStr).Array() + + for i := range numRows { + idCol[i] = uint64(i) + + idxs := make([]uint32, elementsPerRow) + vals := make([]string, elementsPerRow) + for j := range elementsPerRow { + idxs[j] = uint32(j) + vals[j] = "value" + } + idxCol.Append(idxs) + valCol.Append(vals) + } + + require.NoError(t, conn.Do(ctx, Query{ + Body: "INSERT INTO test_nested_large VALUES", + Input: proto.Input{ + {Name: "id", Data: idCol}, + {Name: "data.idx", Data: idxCol}, + {Name: "data.val", Data: valCol}, + }, + })) + + // Verify count + var count proto.ColUInt64 + require.NoError(t, conn.Do(ctx, Query{ + Body: "SELECT count() FROM test_nested_large", + Result: proto.Results{{Name: "count()", Data: &count}}, + })) + assert.Equal(t, uint64(numRows), count.Row(0)) + + // Verify array lengths + var ( + resultIdxCol = new(proto.ColUInt32).Array() + ) + require.NoError(t, conn.Do(ctx, Query{ + Body: "SELECT `data.idx` FROM test_nested_large LIMIT 1", + Result: proto.Results{{Name: "data.idx", Data: resultIdxCol}}, + })) + assert.Equal(t, elementsPerRow, len(resultIdxCol.Row(0))) +} + +// ============================================================================= +// BENCHMARKS +// ============================================================================= +// These benchmarks test encoding/decoding performance at the protocol level, +// following the same pattern as other benchmarks in proto/*_test.go. +// They do NOT require a ClickHouse server. + +// BenchmarkColNested_EncodeColumn benchmarks encoding nested column data. +// Tests encoding of multiple parallel Array columns (how Nested is stored). +func BenchmarkColNested_EncodeColumn(b *testing.B) { + const rows = 1_000 + const elementsPerRow = 10 + + // Prepare data: Nested(id UInt64, name String, value Float64) + nestedIds := new(proto.ColUInt64).Array() + nestedNames := new(proto.ColStr).Array() + nestedValues := new(proto.ColFloat64).Array() + + for i := range rows { + rowIds := make([]uint64, elementsPerRow) + rowNames := make([]string, elementsPerRow) + rowValues := make([]float64, elementsPerRow) + + for j := range elementsPerRow { + rowIds[j] = uint64(i*elementsPerRow + j) + rowNames[j] = fmt.Sprintf("item_%d_%d", i, j) + rowValues[j] = float64(j) * 1.5 + } + + nestedIds.Append(rowIds) + nestedNames.Append(rowNames) + nestedValues.Append(rowValues) + } + + var buf proto.Buffer + nestedIds.EncodeColumn(&buf) + nestedNames.EncodeColumn(&buf) + nestedValues.EncodeColumn(&buf) + + b.SetBytes(int64(len(buf.Buf))) + b.ResetTimer() + b.ReportAllocs() + + for i := 0; i < b.N; i++ { + buf.Reset() + nestedIds.EncodeColumn(&buf) + nestedNames.EncodeColumn(&buf) + nestedValues.EncodeColumn(&buf) + } +} + +// BenchmarkColNested_DecodeColumn benchmarks decoding nested column data. +func BenchmarkColNested_DecodeColumn(b *testing.B) { + const rows = 1_000 + const elementsPerRow = 10 + + // Prepare and encode data + nestedIds := new(proto.ColUInt64).Array() + nestedNames := new(proto.ColStr).Array() + nestedValues := new(proto.ColFloat64).Array() + + for i := 0; i < rows; i++ { + rowIds := make([]uint64, elementsPerRow) + rowNames := make([]string, elementsPerRow) + rowValues := make([]float64, elementsPerRow) + + for j := 0; j < elementsPerRow; j++ { + rowIds[j] = uint64(i*elementsPerRow + j) + rowNames[j] = fmt.Sprintf("item_%d_%d", i, j) + rowValues[j] = float64(j) * 1.5 + } + + nestedIds.Append(rowIds) + nestedNames.Append(rowNames) + nestedValues.Append(rowValues) + } + + // Encode each column separately (as they would be sent over the wire) + var bufIds, bufNames, bufValues proto.Buffer + nestedIds.EncodeColumn(&bufIds) + nestedNames.EncodeColumn(&bufNames) + nestedValues.EncodeColumn(&bufValues) + + totalBytes := len(bufIds.Buf) + len(bufNames.Buf) + len(bufValues.Buf) + b.SetBytes(int64(totalBytes)) + + b.ResetTimer() + b.ReportAllocs() + + for i := 0; i < b.N; i++ { + decIds := new(proto.ColUInt64).Array() + decNames := new(proto.ColStr).Array() + decValues := new(proto.ColFloat64).Array() + + rIds := proto.NewReader(bytes.NewReader(bufIds.Buf)) + rNames := proto.NewReader(bytes.NewReader(bufNames.Buf)) + rValues := proto.NewReader(bytes.NewReader(bufValues.Buf)) + + if err := decIds.DecodeColumn(rIds, rows); err != nil { + b.Fatal(err) + } + if err := decNames.DecodeColumn(rNames, rows); err != nil { + b.Fatal(err) + } + if err := decValues.DecodeColumn(rValues, rows); err != nil { + b.Fatal(err) + } + } +} + +// BenchmarkColNested_Wide_EncodeColumn benchmarks encoding wide nested data (10 columns). +func BenchmarkColNested_Wide_EncodeColumn(b *testing.B) { + const rows = 1_000 + + // 5 UInt64 columns + 5 String columns + arrCols := make([]*proto.ColArr[uint64], 5) + strCols := make([]*proto.ColArr[string], 5) + + for i := range arrCols { + arrCols[i] = new(proto.ColUInt64).Array() + strCols[i] = new(proto.ColStr).Array() + } + + for i := range rows { + for j := range arrCols { + arrCols[j].Append([]uint64{uint64(i), uint64(i + 1)}) + strCols[j].Append([]string{"a", "b"}) + } + } + + var buf proto.Buffer + for _, col := range arrCols { + col.EncodeColumn(&buf) + } + for _, col := range strCols { + col.EncodeColumn(&buf) + } + + b.SetBytes(int64(len(buf.Buf))) + b.ResetTimer() + b.ReportAllocs() + + for i := 0; i < b.N; i++ { + buf.Reset() + for _, col := range arrCols { + col.EncodeColumn(&buf) + } + for _, col := range strCols { + col.EncodeColumn(&buf) + } + } +} + +// BenchmarkColNested_Wide_DecodeColumn benchmarks decoding wide nested data (10 columns). +func BenchmarkColNested_Wide_DecodeColumn(b *testing.B) { + const rows = 1_000 + + // Prepare data + arrCols := make([]*proto.ColArr[uint64], 5) + strCols := make([]*proto.ColArr[string], 5) + + for i := range arrCols { + arrCols[i] = new(proto.ColUInt64).Array() + strCols[i] = new(proto.ColStr).Array() + } + + for i := 0; i < rows; i++ { + for j := range arrCols { + arrCols[j].Append([]uint64{uint64(i), uint64(i + 1)}) + strCols[j].Append([]string{"a", "b"}) + } + } + + // Encode each column + arrBufs := make([]proto.Buffer, 5) + strBufs := make([]proto.Buffer, 5) + var totalBytes int + + for i, col := range arrCols { + col.EncodeColumn(&arrBufs[i]) + totalBytes += len(arrBufs[i].Buf) + } + for i, col := range strCols { + col.EncodeColumn(&strBufs[i]) + totalBytes += len(strBufs[i].Buf) + } + + b.SetBytes(int64(totalBytes)) + b.ResetTimer() + b.ReportAllocs() + + for i := 0; i < b.N; i++ { + for j := range arrCols { + decArr := new(proto.ColUInt64).Array() + decStr := new(proto.ColStr).Array() + rArr := proto.NewReader(bytes.NewReader(arrBufs[j].Buf)) + rStr := proto.NewReader(bytes.NewReader(strBufs[j].Buf)) + + if err := decArr.DecodeColumn(rArr, rows); err != nil { + b.Fatal(err) + } + if err := decStr.DecodeColumn(rStr, rows); err != nil { + b.Fatal(err) + } + } + } +} + +// BenchmarkColNested_Deep_EncodeColumn benchmarks encoding nested data with large arrays (100 elements per row). +func BenchmarkColNested_Deep_EncodeColumn(b *testing.B) { + const rows = 1_000 + const elementsPerRow = 100 + + idxCol := new(proto.ColUInt64).Array() + valCol := new(proto.ColStr).Array() + + for i := 0; i < rows; i++ { + idxs := make([]uint64, elementsPerRow) + vals := make([]string, elementsPerRow) + for j := 0; j < elementsPerRow; j++ { + idxs[j] = uint64(j) + vals[j] = fmt.Sprintf("value_%d", j) + } + idxCol.Append(idxs) + valCol.Append(vals) + } + + var buf proto.Buffer + idxCol.EncodeColumn(&buf) + valCol.EncodeColumn(&buf) + + b.SetBytes(int64(len(buf.Buf))) + b.ResetTimer() + b.ReportAllocs() + + for i := 0; i < b.N; i++ { + buf.Reset() + idxCol.EncodeColumn(&buf) + valCol.EncodeColumn(&buf) + } +} + +// BenchmarkColNested_Deep_DecodeColumn benchmarks decoding nested data with large arrays (100 elements per row). +func BenchmarkColNested_Deep_DecodeColumn(b *testing.B) { + const rows = 1_000 + const elementsPerRow = 100 + + // Prepare data + idxCol := new(proto.ColUInt64).Array() + valCol := new(proto.ColStr).Array() + + for i := 0; i < rows; i++ { + idxs := make([]uint64, elementsPerRow) + vals := make([]string, elementsPerRow) + for j := 0; j < elementsPerRow; j++ { + idxs[j] = uint64(j) + vals[j] = fmt.Sprintf("value_%d", j) + } + idxCol.Append(idxs) + valCol.Append(vals) + } + + // Encode + var bufIdx, bufVal proto.Buffer + idxCol.EncodeColumn(&bufIdx) + valCol.EncodeColumn(&bufVal) + + totalBytes := len(bufIdx.Buf) + len(bufVal.Buf) + b.SetBytes(int64(totalBytes)) + + b.ResetTimer() + b.ReportAllocs() + + for i := 0; i < b.N; i++ { + decIdx := new(proto.ColUInt64).Array() + decVal := new(proto.ColStr).Array() + + rIdx := proto.NewReader(bytes.NewReader(bufIdx.Buf)) + rVal := proto.NewReader(bytes.NewReader(bufVal.Buf)) + + if err := decIdx.DecodeColumn(rIdx, rows); err != nil { + b.Fatal(err) + } + if err := decVal.DecodeColumn(rVal, rows); err != nil { + b.Fatal(err) + } + } +} + +// BenchmarkColStr_JSONMarshal benchmarks JSON marshaling for comparison with native encoding. +// This shows the overhead of using JSON strings vs native Nested types. +func BenchmarkColStr_JSONMarshal(b *testing.B) { + const rows = 1_000 + const elementsPerRow = 10 + + type Element struct { + ID uint64 `json:"id"` + Name string `json:"name"` + Value float64 `json:"value"` + } + + type JSONData struct { + Items []Element `json:"items"` + } + + b.ResetTimer() + b.ReportAllocs() + + var totalBytes int64 + for i := 0; i < b.N; i++ { + var strCol proto.ColStr + for row := 0; row < rows; row++ { + elements := make([]Element, elementsPerRow) + for j := 0; j < elementsPerRow; j++ { + elements[j] = Element{ + ID: uint64(row*elementsPerRow + j), + Name: fmt.Sprintf("item_%d_%d", row, j), + Value: float64(j) * 1.5, + } + } + data := JSONData{Items: elements} + jsonBytes, _ := json.Marshal(data) + strCol.Append(string(jsonBytes)) + } + + var buf proto.Buffer + strCol.EncodeColumn(&buf) + totalBytes = int64(len(buf.Buf)) + } + b.SetBytes(totalBytes) +} + +// BenchmarkColStr_JSONUnmarshal benchmarks JSON unmarshaling for comparison. +func BenchmarkColStr_JSONUnmarshal(b *testing.B) { + const rows = 1_000 + const elementsPerRow = 10 + + type Element struct { + ID uint64 `json:"id"` + Name string `json:"name"` + Value float64 `json:"value"` + } + + type JSONData struct { + Items []Element `json:"items"` + } + + // Prepare encoded data + var strCol proto.ColStr + for row := 0; row < rows; row++ { + elements := make([]Element, elementsPerRow) + for j := 0; j < elementsPerRow; j++ { + elements[j] = Element{ + ID: uint64(row*elementsPerRow + j), + Name: fmt.Sprintf("item_%d_%d", row, j), + Value: float64(j) * 1.5, + } + } + data := JSONData{Items: elements} + jsonBytes, _ := json.Marshal(data) + strCol.Append(string(jsonBytes)) + } + + var buf proto.Buffer + strCol.EncodeColumn(&buf) + + b.SetBytes(int64(len(buf.Buf))) + b.ResetTimer() + b.ReportAllocs() + + for i := 0; i < b.N; i++ { + r := proto.NewReader(bytes.NewReader(buf.Buf)) + + var dec proto.ColStr + if err := dec.DecodeColumn(r, rows); err != nil { + b.Fatal(err) + } + + // Unmarshal each row (the real cost of using JSON strings) + for j := 0; j < dec.Rows(); j++ { + var data JSONData + if err := json.Unmarshal([]byte(dec.Row(j)), &data); err != nil { + b.Fatal(err) + } + } + } +} diff --git a/proto/col_auto.go b/proto/col_auto.go index 658c5a1f..f30330c2 100644 --- a/proto/col_auto.go +++ b/proto/col_auto.go @@ -94,6 +94,14 @@ func (c *ColAuto) Infer(t ColumnType) error { return nil } } + case ColumnTypeNested: + v := new(ColNested) + if err := v.Infer(t); err != nil { + return errors.Wrap(err, "nested") + } + c.Data = v + c.DataType = t + return nil case ColumnTypeDateTime: v := new(ColDateTime) if err := v.Infer(t); err != nil { diff --git a/proto/col_auto_test.go b/proto/col_auto_test.go index 1165d56f..405c2da1 100644 --- a/proto/col_auto_test.go +++ b/proto/col_auto_test.go @@ -59,6 +59,8 @@ func TestColAuto_Infer(t *testing.T) { "Decimal256(4)", "Array(Nullable(Int8))", "Nullable(DateTime64(3))", + "Nested(id UInt64, name String)", + "Nested(value Float64)", } { r := AutoResult("foo") require.NoError(t, r.Data.(Inferable).Infer(columnType)) diff --git a/proto/col_nested.go b/proto/col_nested.go new file mode 100644 index 00000000..8d75eaca --- /dev/null +++ b/proto/col_nested.go @@ -0,0 +1,337 @@ +package proto + +import ( + "reflect" + + "github.com/go-faster/errors" +) + +// ColNested represents Nested(name1 Type1, name2 Type2, ...). +// +// In ClickHouse, Nested types are stored as multiple parallel Array columns. +// For example, Nested(a UInt32, b String) is stored as: +// - col.a Array(UInt32) +// - col.b Array(String) +// +// All arrays in a row must have the same length (this is enforced by ClickHouse). +// +// Since Nested is not a wire-level type, ColNested provides helper methods +// to work with the flattened array columns: +// - InputColumns(prefix) returns []InputColumn for INSERT operations +// - ResultColumns(prefix) returns []ResultColumn for SELECT operations +type ColNested struct { + columns []NestedColumn +} + +// NestedColumn represents a single column within a Nested type. +type NestedColumn struct { + Name string // Field name (e.g., "id") + Data Column // Array column (e.g., *ColArr[uint64]) +} + +// NewNested creates a new ColNested with the given columns. +// +// Each column's Data should be an Array type. Example: +// +// nested := NewNested( +// NestedColumn{Name: "id", Data: new(ColUInt64).Array()}, +// NestedColumn{Name: "name", Data: new(ColStr).Array()}, +// ) +func NewNested(columns ...NestedColumn) *ColNested { + return &ColNested{columns: columns} +} + +// Columns returns the nested columns. +func (c *ColNested) Columns() []NestedColumn { + return c.columns +} + +// Column returns the column with the given name, or nil if not found. +func (c *ColNested) Column(name string) *NestedColumn { + for i := range c.columns { + if c.columns[i].Name == name { + return &c.columns[i] + } + } + return nil +} + +// Type returns the Nested type string, e.g., "Nested(id UInt64, name String)". +func (c *ColNested) Type() ColumnType { + if len(c.columns) == 0 { + return "Nested()" + } + var types []ColumnType + for _, col := range c.columns { + // Extract element type from Array(T) -> T + elemType := col.Data.Type().Elem() + types = append(types, ColumnType(col.Name+" "+elemType.String())) + } + return ColumnTypeNested.Sub(types...) +} + +// Rows returns the number of rows (from the first column). +func (c *ColNested) Rows() int { + if len(c.columns) == 0 { + return 0 + } + return c.columns[0].Data.Rows() +} + +// Reset clears all columns. +func (c *ColNested) Reset() { + for _, col := range c.columns { + col.Data.Reset() + } +} + +// InputColumns returns []InputColumn for INSERT operations. +// The prefix is added to each column name with a dot separator. +// +// Example: +// +// nested.InputColumns("n") +// // Returns: [{Name: "n.id", Data: ...}, {Name: "n.name", Data: ...}] +func (c *ColNested) InputColumns(prefix string) []InputColumn { + result := make([]InputColumn, len(c.columns)) + for i, col := range c.columns { + result[i] = InputColumn{ + Name: prefix + "." + col.Name, + Data: col.Data, + } + } + return result +} + +// ResultColumns returns []ResultColumn for SELECT operations. +// The prefix is added to each column name with a dot separator. +// +// Example: +// +// nested.ResultColumns("n") +// // Returns: [{Name: "n.id", Data: ...}, {Name: "n.name", Data: ...}] +func (c *ColNested) ResultColumns(prefix string) []ResultColumn { + result := make([]ResultColumn, len(c.columns)) + for i, col := range c.columns { + result[i] = ResultColumn{ + Name: prefix + "." + col.Name, + Data: col.Data, + } + } + return result +} + +// Append appends a row to the nested column. +// The values map should contain all column names with their array values. +// All arrays must have the same length. +// +// Example: +// +// nested.Append(map[string]any{ +// "id": []uint64{1, 2, 3}, +// "name": []string{"a", "b", "c"}, +// }) +func (c *ColNested) Append(values map[string]any) error { + if len(values) != len(c.columns) { + return errors.Errorf("expected %d columns, got %d", len(c.columns), len(values)) + } + + // First pass: validate all columns exist and arrays have equal length + expectedLen := -1 + for _, col := range c.columns { + val, ok := values[col.Name] + if !ok { + return errors.Errorf("missing value for column %q", col.Name) + } + + arrLen := reflectArrayLen(val) + if arrLen < 0 { + return errors.Errorf("value for column %q is not a slice", col.Name) + } + + if expectedLen == -1 { + expectedLen = arrLen + } else if arrLen != expectedLen { + return errors.Errorf("column %q has length %d, expected %d (all arrays must have equal length)", + col.Name, arrLen, expectedLen) + } + } + + // Second pass: append to each column + for _, col := range c.columns { + val := values[col.Name] + if err := appendToArray(col.Data, val); err != nil { + return errors.Wrapf(err, "append to column %q", col.Name) + } + } + + return nil +} + +// AppendTyped is a type-safe way to append a row using a struct. +// The struct fields should match the column names (case-sensitive or using tags). +// This is more efficient than Append as it avoids reflection in hot path. +// +// For now, use Append() which uses reflection. Type-safe variants can be +// added for common use cases. + +// Infer implements Inferable interface. +// It parses the Nested type and initializes the columns. +func (c *ColNested) Infer(t ColumnType) error { + if t.Base() != ColumnTypeNested { + return errors.Errorf("expected Nested type, got %q", t.Base()) + } + + fields, err := ParseNestedFields(string(t.Elem())) + if err != nil { + return errors.Wrap(err, "parse nested fields") + } + + // If columns already exist, validate and update them + if len(c.columns) > 0 { + if len(c.columns) != len(fields) { + return errors.Errorf("column count mismatch: have %d, type has %d", + len(c.columns), len(fields)) + } + for i, field := range fields { + if c.columns[i].Name != field.Name { + return errors.Errorf("column name mismatch at %d: have %q, type has %q", + i, c.columns[i].Name, field.Name) + } + // Infer the inner type (as Array) + if infer, ok := c.columns[i].Data.(Inferable); ok { + arrayType := ColumnTypeArray.Sub(field.Type) + if err := infer.Infer(arrayType); err != nil { + return errors.Wrapf(err, "infer column %q", field.Name) + } + } + } + return nil + } + + // Create new columns using ColAuto for automatic type inference + c.columns = make([]NestedColumn, len(fields)) + for i, field := range fields { + inner := new(ColAuto) + arrayType := ColumnTypeArray.Sub(field.Type) + if err := inner.Infer(arrayType); err != nil { + return errors.Wrapf(err, "infer column %q", field.Name) + } + c.columns[i] = NestedColumn{ + Name: field.Name, + Data: inner.Data, + } + } + + return nil +} + +// Prepare implements Preparable interface. +func (c *ColNested) Prepare() error { + for _, col := range c.columns { + if p, ok := col.Data.(Preparable); ok { + if err := p.Prepare(); err != nil { + return errors.Wrapf(err, "prepare column %q", col.Name) + } + } + } + return nil +} + +// DecodeState implements StateDecoder interface. +func (c *ColNested) DecodeState(r *Reader) error { + for _, col := range c.columns { + if s, ok := col.Data.(StateDecoder); ok { + if err := s.DecodeState(r); err != nil { + return errors.Wrapf(err, "decode state for column %q", col.Name) + } + } + } + return nil +} + +// EncodeState implements StateEncoder interface. +func (c *ColNested) EncodeState(b *Buffer) { + for _, col := range c.columns { + if s, ok := col.Data.(StateEncoder); ok { + s.EncodeState(b) + } + } +} + +// EncodeColumn is not supported for Nested types. +// Nested columns must be encoded as separate Array columns using InputColumns(). +// This method exists only to satisfy the Column interface for type inference. +func (c *ColNested) EncodeColumn(b *Buffer) { + // Nested types are sent as multiple Array columns, not a single column. + // Use InputColumns(prefix) to get the flattened columns for INSERT. + panic("ColNested.EncodeColumn: Nested types must be encoded as separate Array columns using InputColumns()") +} + +// WriteColumn is not supported for Nested types. +// Nested columns must be written as separate Array columns using InputColumns(). +func (c *ColNested) WriteColumn(w *Writer) { + // Nested types are sent as multiple Array columns, not a single column. + // Use InputColumns(prefix) to get the flattened columns for INSERT. + panic("ColNested.WriteColumn: Nested types must be written as separate Array columns using InputColumns()") +} + +// DecodeColumn is not supported for Nested types. +// Nested columns must be decoded as separate Array columns using ResultColumns(). +// This method exists only to satisfy the Column interface for type inference. +func (c *ColNested) DecodeColumn(r *Reader, rows int) error { + // Nested types are received as multiple Array columns, not a single column. + // Use ResultColumns(prefix) to get the expected columns for SELECT. + return errors.New("ColNested.DecodeColumn: Nested types must be decoded as separate Array columns using ResultColumns()") +} + +// Note: ColNested does NOT encode/decode directly because ClickHouse sends +// Nested as multiple separate Array columns with dot-notation names. +// Use InputColumns() and ResultColumns() to get the flattened columns. +// +// For INSERT: Use InputColumns(prefix) to get []InputColumn +// For SELECT: Use ResultColumns(prefix) to get []ResultColumn + +// reflectArrayLen returns the length of a slice/array, or -1 if not a slice. +func reflectArrayLen(v any) int { + rv := reflect.ValueOf(v) + if rv.Kind() != reflect.Slice && rv.Kind() != reflect.Array { + return -1 + } + return rv.Len() +} + +// appendToArray appends values to an array column using reflection. +func appendToArray(col Column, values any) error { + rv := reflect.ValueOf(values) + if rv.Kind() != reflect.Slice && rv.Kind() != reflect.Array { + return errors.New("values must be a slice") + } + + // Get the underlying array column and use type assertion to call Append + colVal := reflect.ValueOf(col) + + // Try to find Append method + appendMethod := colVal.MethodByName("Append") + if !appendMethod.IsValid() { + return errors.Errorf("column type %T does not have Append method", col) + } + + // Call Append with the values + results := appendMethod.Call([]reflect.Value{rv}) + if len(results) > 0 && !results[0].IsNil() { + if err, ok := results[0].Interface().(error); ok { + return err + } + } + + return nil +} + +// Compile-time assertions for ColNested. +var ( + _ Inferable = (*ColNested)(nil) + _ Preparable = (*ColNested)(nil) + _ StateEncoder = (*ColNested)(nil) + _ StateDecoder = (*ColNested)(nil) +) diff --git a/proto/col_nested_test.go b/proto/col_nested_test.go new file mode 100644 index 00000000..9ff87ec9 --- /dev/null +++ b/proto/col_nested_test.go @@ -0,0 +1,413 @@ +package proto + +import ( + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestColNestedBasic(t *testing.T) { + t.Run("empty", func(t *testing.T) { + nested := NewNested() + assert.Equal(t, 0, nested.Rows()) + assert.Equal(t, "Nested()", nested.Type().String()) + }) + + t.Run("single column", func(t *testing.T) { + idCol := new(ColUInt64).Array() + nested := NewNested( + NestedColumn{Name: "id", Data: idCol}, + ) + + assert.Equal(t, "Nested(id UInt64)", nested.Type().String()) + assert.Equal(t, 0, nested.Rows()) + + // Append a row + idCol.Append([]uint64{1, 2, 3}) + assert.Equal(t, 1, nested.Rows()) + }) + + t.Run("multiple columns", func(t *testing.T) { + idCol := new(ColUInt64).Array() + nameCol := new(ColStr).Array() + nested := NewNested( + NestedColumn{Name: "id", Data: idCol}, + NestedColumn{Name: "name", Data: nameCol}, + ) + + assert.Equal(t, "Nested(id UInt64, name String)", nested.Type().String()) + }) +} + +func TestColNestedType(t *testing.T) { + tests := []struct { + name string + columns []NestedColumn + expected string + }{ + { + name: "empty", + columns: nil, + expected: "Nested()", + }, + { + name: "single uint64", + columns: []NestedColumn{ + {Name: "id", Data: new(ColUInt64).Array()}, + }, + expected: "Nested(id UInt64)", + }, + { + name: "uint64 and string", + columns: []NestedColumn{ + {Name: "id", Data: new(ColUInt64).Array()}, + {Name: "name", Data: new(ColStr).Array()}, + }, + expected: "Nested(id UInt64, name String)", + }, + { + name: "various types", + columns: []NestedColumn{ + {Name: "id", Data: new(ColUInt64).Array()}, + {Name: "name", Data: new(ColStr).Array()}, + {Name: "value", Data: new(ColFloat64).Array()}, + {Name: "flag", Data: new(ColBool).Array()}, + }, + expected: "Nested(id UInt64, name String, value Float64, flag Bool)", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + nested := NewNested(tt.columns...) + assert.Equal(t, tt.expected, nested.Type().String()) + }) + } +} + +func TestColNestedInputColumns(t *testing.T) { + idCol := new(ColUInt64).Array() + nameCol := new(ColStr).Array() + nested := NewNested( + NestedColumn{Name: "id", Data: idCol}, + NestedColumn{Name: "name", Data: nameCol}, + ) + + inputs := nested.InputColumns("n") + + require.Len(t, inputs, 2) + assert.Equal(t, "n.id", inputs[0].Name) + assert.Equal(t, "n.name", inputs[1].Name) + assert.Same(t, idCol, inputs[0].Data) + assert.Same(t, nameCol, inputs[1].Data) +} + +func TestColNestedResultColumns(t *testing.T) { + idCol := new(ColUInt64).Array() + nameCol := new(ColStr).Array() + nested := NewNested( + NestedColumn{Name: "id", Data: idCol}, + NestedColumn{Name: "name", Data: nameCol}, + ) + + results := nested.ResultColumns("data") + + require.Len(t, results, 2) + assert.Equal(t, "data.id", results[0].Name) + assert.Equal(t, "data.name", results[1].Name) + assert.Same(t, idCol, results[0].Data) + assert.Same(t, nameCol, results[1].Data) +} + +func TestColNestedAppend(t *testing.T) { + t.Run("valid append", func(t *testing.T) { + nested := NewNested( + NestedColumn{Name: "id", Data: new(ColUInt64).Array()}, + NestedColumn{Name: "name", Data: new(ColStr).Array()}, + ) + + err := nested.Append(map[string]any{ + "id": []uint64{1, 2, 3}, + "name": []string{"a", "b", "c"}, + }) + require.NoError(t, err) + assert.Equal(t, 1, nested.Rows()) + + // Append another row + err = nested.Append(map[string]any{ + "id": []uint64{4, 5}, + "name": []string{"d", "e"}, + }) + require.NoError(t, err) + assert.Equal(t, 2, nested.Rows()) + }) + + t.Run("empty arrays", func(t *testing.T) { + nested := NewNested( + NestedColumn{Name: "id", Data: new(ColUInt64).Array()}, + NestedColumn{Name: "name", Data: new(ColStr).Array()}, + ) + + err := nested.Append(map[string]any{ + "id": []uint64{}, + "name": []string{}, + }) + require.NoError(t, err) + assert.Equal(t, 1, nested.Rows()) + }) + + t.Run("mismatched lengths", func(t *testing.T) { + nested := NewNested( + NestedColumn{Name: "id", Data: new(ColUInt64).Array()}, + NestedColumn{Name: "name", Data: new(ColStr).Array()}, + ) + + err := nested.Append(map[string]any{ + "id": []uint64{1, 2, 3}, + "name": []string{"a", "b"}, // different length + }) + require.Error(t, err) + assert.Contains(t, err.Error(), "all arrays must have equal length") + }) + + t.Run("missing column", func(t *testing.T) { + nested := NewNested( + NestedColumn{Name: "id", Data: new(ColUInt64).Array()}, + NestedColumn{Name: "name", Data: new(ColStr).Array()}, + ) + + err := nested.Append(map[string]any{ + "id": []uint64{1, 2, 3}, + // missing "name" + }) + require.Error(t, err) + assert.Contains(t, err.Error(), "expected 2 columns") + }) + + t.Run("wrong column name", func(t *testing.T) { + nested := NewNested( + NestedColumn{Name: "id", Data: new(ColUInt64).Array()}, + NestedColumn{Name: "name", Data: new(ColStr).Array()}, + ) + + err := nested.Append(map[string]any{ + "id": []uint64{1, 2, 3}, + "wrong": []string{"a", "b", "c"}, // wrong column name + }) + require.Error(t, err) + assert.Contains(t, err.Error(), "missing value for column") + }) + + t.Run("not a slice", func(t *testing.T) { + nested := NewNested( + NestedColumn{Name: "id", Data: new(ColUInt64).Array()}, + ) + + err := nested.Append(map[string]any{ + "id": "not a slice", + }) + require.Error(t, err) + assert.Contains(t, err.Error(), "not a slice") + }) +} + +func TestColNestedReset(t *testing.T) { + idCol := new(ColUInt64).Array() + nameCol := new(ColStr).Array() + nested := NewNested( + NestedColumn{Name: "id", Data: idCol}, + NestedColumn{Name: "name", Data: nameCol}, + ) + + // Add some data + idCol.Append([]uint64{1, 2, 3}) + nameCol.Append([]string{"a", "b", "c"}) + assert.Equal(t, 1, nested.Rows()) + + // Reset + nested.Reset() + assert.Equal(t, 0, nested.Rows()) + assert.Equal(t, 0, idCol.Rows()) + assert.Equal(t, 0, nameCol.Rows()) +} + +func TestColNestedColumn(t *testing.T) { + idCol := new(ColUInt64).Array() + nameCol := new(ColStr).Array() + nested := NewNested( + NestedColumn{Name: "id", Data: idCol}, + NestedColumn{Name: "name", Data: nameCol}, + ) + + t.Run("found", func(t *testing.T) { + col := nested.Column("id") + require.NotNil(t, col) + assert.Equal(t, "id", col.Name) + assert.Same(t, idCol, col.Data) + }) + + t.Run("not found", func(t *testing.T) { + col := nested.Column("nonexistent") + assert.Nil(t, col) + }) +} + +func TestColNestedInfer(t *testing.T) { + t.Run("simple types", func(t *testing.T) { + nested := &ColNested{} + err := nested.Infer("Nested(id UInt64, name String)") + require.NoError(t, err) + + require.Len(t, nested.columns, 2) + assert.Equal(t, "id", nested.columns[0].Name) + assert.Equal(t, "name", nested.columns[1].Name) + + // Verify types + assert.Equal(t, "Array(UInt64)", nested.columns[0].Data.Type().String()) + assert.Equal(t, "Array(String)", nested.columns[1].Data.Type().String()) + }) + + t.Run("complex types", func(t *testing.T) { + nested := &ColNested{} + err := nested.Infer("Nested(value Float64, count UInt32, flag Bool)") + require.NoError(t, err) + + require.Len(t, nested.columns, 3) + assert.Equal(t, "value", nested.columns[0].Name) + assert.Equal(t, "count", nested.columns[1].Name) + assert.Equal(t, "flag", nested.columns[2].Name) + + // Each nested field becomes Array(T) + assert.Equal(t, "Array(Float64)", nested.columns[0].Data.Type().String()) + assert.Equal(t, "Array(UInt32)", nested.columns[1].Data.Type().String()) + assert.Equal(t, "Array(Bool)", nested.columns[2].Data.Type().String()) + }) + + t.Run("existing columns - matching", func(t *testing.T) { + nested := NewNested( + NestedColumn{Name: "id", Data: new(ColUInt64).Array()}, + NestedColumn{Name: "name", Data: new(ColStr).Array()}, + ) + + err := nested.Infer("Nested(id UInt64, name String)") + require.NoError(t, err) + }) + + t.Run("existing columns - count mismatch", func(t *testing.T) { + nested := NewNested( + NestedColumn{Name: "id", Data: new(ColUInt64).Array()}, + ) + + err := nested.Infer("Nested(id UInt64, name String)") + require.Error(t, err) + assert.Contains(t, err.Error(), "column count mismatch") + }) + + t.Run("existing columns - name mismatch", func(t *testing.T) { + nested := NewNested( + NestedColumn{Name: "id", Data: new(ColUInt64).Array()}, + NestedColumn{Name: "wrong", Data: new(ColStr).Array()}, + ) + + err := nested.Infer("Nested(id UInt64, name String)") + require.Error(t, err) + assert.Contains(t, err.Error(), "column name mismatch") + }) + + t.Run("not nested type", func(t *testing.T) { + nested := &ColNested{} + err := nested.Infer("Array(String)") + require.Error(t, err) + assert.Contains(t, err.Error(), "expected Nested type") + }) + + t.Run("empty nested", func(t *testing.T) { + nested := &ColNested{} + err := nested.Infer("Nested()") + require.Error(t, err) + }) +} + +func TestColNestedColumns(t *testing.T) { + idCol := new(ColUInt64).Array() + nameCol := new(ColStr).Array() + nested := NewNested( + NestedColumn{Name: "id", Data: idCol}, + NestedColumn{Name: "name", Data: nameCol}, + ) + + cols := nested.Columns() + require.Len(t, cols, 2) + assert.Equal(t, "id", cols[0].Name) + assert.Equal(t, "name", cols[1].Name) +} + +func TestColNestedPrepare(t *testing.T) { + // Create nested with columns that implement Preparable + nested := NewNested( + NestedColumn{Name: "id", Data: new(ColUInt64).Array()}, + ) + + // Prepare should not error for basic types + err := nested.Prepare() + require.NoError(t, err) +} + +func TestColNestedIntegration(t *testing.T) { + // Test a full workflow: create, append, get flattened columns + t.Run("full workflow", func(t *testing.T) { + nested := NewNested( + NestedColumn{Name: "user_id", Data: new(ColUInt64).Array()}, + NestedColumn{Name: "event_name", Data: new(ColStr).Array()}, + NestedColumn{Name: "timestamp", Data: new(ColInt64).Array()}, + ) + + // Append multiple rows + require.NoError(t, nested.Append(map[string]any{ + "user_id": []uint64{1, 2}, + "event_name": []string{"click", "view"}, + "timestamp": []int64{1000, 2000}, + })) + + require.NoError(t, nested.Append(map[string]any{ + "user_id": []uint64{3}, + "event_name": []string{"purchase"}, + "timestamp": []int64{3000}, + })) + + assert.Equal(t, 2, nested.Rows()) + + // Get input columns for INSERT + inputs := nested.InputColumns("events") + require.Len(t, inputs, 3) + assert.Equal(t, "events.user_id", inputs[0].Name) + assert.Equal(t, "events.event_name", inputs[1].Name) + assert.Equal(t, "events.timestamp", inputs[2].Name) + + // Verify data in columns + userIdCol := inputs[0].Data.(*ColArr[uint64]) + assert.Equal(t, 2, userIdCol.Rows()) + assert.Equal(t, []uint64{1, 2}, userIdCol.Row(0)) + assert.Equal(t, []uint64{3}, userIdCol.Row(1)) + }) +} + +func TestColNestedWithInferredTypes(t *testing.T) { + // Test creating nested via Infer and then using it + t.Run("infer then use", func(t *testing.T) { + nested := &ColNested{} + require.NoError(t, nested.Infer("Nested(id UInt64, name String)")) + + // Now the columns are set up, we can use InputColumns/ResultColumns + inputs := nested.InputColumns("n") + require.Len(t, inputs, 2) + assert.Equal(t, "n.id", inputs[0].Name) + assert.Equal(t, "n.name", inputs[1].Name) + + results := nested.ResultColumns("n") + require.Len(t, results, 2) + assert.Equal(t, "n.id", results[0].Name) + assert.Equal(t, "n.name", results[1].Name) + }) +} diff --git a/proto/column.go b/proto/column.go index 9e00cc72..1eed2ab0 100644 --- a/proto/column.go +++ b/proto/column.go @@ -153,7 +153,7 @@ func (c ColumnType) normalizeCommas() ColumnType { // Should we check for escaped commas in enums here? const sep = "," var elems []string - for _, e := range strings.Split(string(c), sep) { + for e := range strings.SplitSeq(string(c), sep) { elems = append(elems, strings.TrimSpace(e)) } return ColumnType(strings.Join(elems, sep)) @@ -255,6 +255,7 @@ const ( ColumnTypeNothing ColumnType = "Nothing" ColumnTypeJSON ColumnType = "JSON" ColumnTypeQBit ColumnType = "QBit" + ColumnTypeNested ColumnType = "Nested" ) // colWrap wraps Column with type t. @@ -268,7 +269,7 @@ func (c colWrap) Type() ColumnType { return c.t } // Wrap Column with type parameters. // // So if c type is T, result type will be T(arg0, arg1, ...). -func Wrap(c Column, args ...interface{}) Column { +func Wrap(c Column, args ...any) Column { var params []string for _, a := range args { params = append(params, fmt.Sprint(a)) diff --git a/proto/nested_parser.go b/proto/nested_parser.go new file mode 100644 index 00000000..0526acae --- /dev/null +++ b/proto/nested_parser.go @@ -0,0 +1,128 @@ +package proto + +import ( + "strings" + "unicode" + + "github.com/go-faster/errors" +) + +// NestedField represents a single field within a Nested type definition. +type NestedField struct { + Name string // Field name (e.g., "id") + Type ColumnType // Field type (e.g., "UInt64" or "Array(String)") +} + +// ParseNestedFields parses the element string inside Nested(...). +// For example, "a UInt32, b String, c Array(Int64)" returns three NestedField entries. +// +// The parser handles: +// - Nested parentheses: "a Array(Array(Int8)), b String" +// - Named fields with space separator: "field_name Type" +// - Comma-separated fields at top level only +// - Whitespace variations +func ParseNestedFields(elem string) ([]NestedField, error) { + elem = strings.TrimSpace(elem) + if elem == "" { + return nil, errors.New("empty nested type definition") + } + + var fields []NestedField + depth := 0 // Track parenthesis depth + start := 0 // Start of current field + + for i := 0; i < len(elem); i++ { + switch elem[i] { + case '(': + depth++ + case ')': + depth-- + if depth < 0 { + return nil, errors.Errorf("unbalanced parentheses at position %d", i) + } + case ',': + if depth == 0 { + // Found top-level separator + field, err := parseNameTypePair(strings.TrimSpace(elem[start:i])) + if err != nil { + return nil, err + } + fields = append(fields, field) + start = i + 1 + } + } + } + + if depth != 0 { + return nil, errors.New("unbalanced parentheses in nested type definition") + } + + // Don't forget the last field + if start < len(elem) { + field, err := parseNameTypePair(strings.TrimSpace(elem[start:])) + if err != nil { + return nil, err + } + fields = append(fields, field) + } + + if len(fields) == 0 { + return nil, errors.New("no fields found in nested type definition") + } + + return fields, nil +} + +// parseNameTypePair parses "name Type" or "name Type(params)" into NestedField. +// Examples: +// - "id UInt64" -> {Name: "id", Type: "UInt64"} +// - "values Array(String)" -> {Name: "values", Type: "Array(String)"} +// - "data Nullable(Int32)" -> {Name: "data", Type: "Nullable(Int32)"} +func parseNameTypePair(s string) (NestedField, error) { + s = strings.TrimSpace(s) + if s == "" { + return NestedField{}, errors.New("empty field definition") + } + + // Find the first whitespace that separates name from type + idx := strings.IndexFunc(s, unicode.IsSpace) + if idx == -1 { + return NestedField{}, errors.Errorf("invalid nested field definition %q: expected 'name Type' format", s) + } + + name := s[:idx] + typ := strings.TrimSpace(s[idx+1:]) + + if name == "" { + return NestedField{}, errors.Errorf("empty field name in %q", s) + } + if typ == "" { + return NestedField{}, errors.Errorf("empty field type in %q", s) + } + + // Validate name doesn't contain invalid characters + if strings.ContainsAny(name, "(),") { + return NestedField{}, errors.Errorf("invalid field name %q: contains invalid characters", name) + } + + return NestedField{ + Name: name, + Type: ColumnType(typ), + }, nil +} + +// ParseNestedType parses a full Nested type string like "Nested(a UInt32, b String)". +// Returns the parsed fields. +func ParseNestedType(t ColumnType) ([]NestedField, error) { + base := t.Base() + if base != ColumnTypeNested { + return nil, errors.Errorf("expected Nested type, got %q", base) + } + + elem := t.Elem() + if elem == "" { + return nil, errors.New("empty Nested type definition") + } + + return ParseNestedFields(string(elem)) +} diff --git a/proto/nested_parser_test.go b/proto/nested_parser_test.go new file mode 100644 index 00000000..ee53f83e --- /dev/null +++ b/proto/nested_parser_test.go @@ -0,0 +1,435 @@ +package proto + +import ( + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestParseNestedFields(t *testing.T) { + tests := []struct { + name string + input string + expected []NestedField + wantErr bool + errMsg string + }{ + // Basic cases + { + name: "single field", + input: "a UInt32", + expected: []NestedField{ + {Name: "a", Type: "UInt32"}, + }, + }, + { + name: "two fields", + input: "a UInt32, b String", + expected: []NestedField{ + {Name: "a", Type: "UInt32"}, + {Name: "b", Type: "String"}, + }, + }, + { + name: "three fields", + input: "id UInt64, name String, value Float64", + expected: []NestedField{ + {Name: "id", Type: "UInt64"}, + {Name: "name", Type: "String"}, + {Name: "value", Type: "Float64"}, + }, + }, + + // Nested types (parentheses) + { + name: "array type", + input: "values Array(String)", + expected: []NestedField{ + {Name: "values", Type: "Array(String)"}, + }, + }, + { + name: "nullable type", + input: "data Nullable(Int32)", + expected: []NestedField{ + {Name: "data", Type: "Nullable(Int32)"}, + }, + }, + { + name: "array and simple", + input: "items Array(Int64), count UInt32", + expected: []NestedField{ + {Name: "items", Type: "Array(Int64)"}, + {Name: "count", Type: "UInt32"}, + }, + }, + { + name: "multiple complex types", + input: "arr Array(String), nullable Nullable(Float64), simple Int8", + expected: []NestedField{ + {Name: "arr", Type: "Array(String)"}, + {Name: "nullable", Type: "Nullable(Float64)"}, + {Name: "simple", Type: "Int8"}, + }, + }, + + // Deeply nested types + { + name: "array of arrays", + input: "matrix Array(Array(Int8))", + expected: []NestedField{ + {Name: "matrix", Type: "Array(Array(Int8))"}, + }, + }, + { + name: "array of arrays with other field", + input: "matrix Array(Array(Int8)), label String", + expected: []NestedField{ + {Name: "matrix", Type: "Array(Array(Int8))"}, + {Name: "label", Type: "String"}, + }, + }, + { + name: "map type", + input: "metadata Map(String, Int32)", + expected: []NestedField{ + {Name: "metadata", Type: "Map(String, Int32)"}, + }, + }, + { + name: "map with array value", + input: "data Map(String, Array(Int64))", + expected: []NestedField{ + {Name: "data", Type: "Map(String, Array(Int64))"}, + }, + }, + { + name: "low cardinality", + input: "category LowCardinality(String)", + expected: []NestedField{ + {Name: "category", Type: "LowCardinality(String)"}, + }, + }, + { + name: "tuple type", + input: "point Tuple(Float64, Float64)", + expected: []NestedField{ + {Name: "point", Type: "Tuple(Float64, Float64)"}, + }, + }, + { + name: "named tuple", + input: "coords Tuple(x Float64, y Float64)", + expected: []NestedField{ + {Name: "coords", Type: "Tuple(x Float64, y Float64)"}, + }, + }, + { + name: "nullable array", + input: "values Nullable(Array(Int32))", + expected: []NestedField{ + {Name: "values", Type: "Nullable(Array(Int32))"}, + }, + }, + { + name: "array of nullable", + input: "values Array(Nullable(Int32))", + expected: []NestedField{ + {Name: "values", Type: "Array(Nullable(Int32))"}, + }, + }, + + // Whitespace variations + { + name: "extra spaces around commas", + input: "a UInt32 , b String , c Int64", + expected: []NestedField{ + {Name: "a", Type: "UInt32"}, + {Name: "b", Type: "String"}, + {Name: "c", Type: "Int64"}, + }, + }, + { + name: "leading/trailing whitespace", + input: " a UInt32, b String ", + expected: []NestedField{ + {Name: "a", Type: "UInt32"}, + {Name: "b", Type: "String"}, + }, + }, + { + name: "tabs and newlines", + input: "a\tUInt32,\nb\tString", + expected: []NestedField{ + {Name: "a", Type: "UInt32"}, + {Name: "b", Type: "String"}, + }, + }, + + // Underscore names + { + name: "underscore in name", + input: "user_id UInt64, created_at DateTime", + expected: []NestedField{ + {Name: "user_id", Type: "UInt64"}, + {Name: "created_at", Type: "DateTime"}, + }, + }, + + // DateTime with parameters + { + name: "datetime64 with precision", + input: "ts DateTime64(3)", + expected: []NestedField{ + {Name: "ts", Type: "DateTime64(3)"}, + }, + }, + { + name: "datetime64 with precision and timezone", + input: "ts DateTime64(3, 'UTC')", + expected: []NestedField{ + {Name: "ts", Type: "DateTime64(3, 'UTC')"}, + }, + }, + + // Fixed string + { + name: "fixed string", + input: "code FixedString(3)", + expected: []NestedField{ + {Name: "code", Type: "FixedString(3)"}, + }, + }, + + // Decimal + { + name: "decimal with precision and scale", + input: "price Decimal(10, 2)", + expected: []NestedField{ + {Name: "price", Type: "Decimal(10, 2)"}, + }, + }, + + // Enum + { + name: "enum8", + input: "status Enum8('active' = 1, 'inactive' = 2)", + expected: []NestedField{ + {Name: "status", Type: "Enum8('active' = 1, 'inactive' = 2)"}, + }, + }, + + // Complex realistic example + { + name: "realistic nested definition", + input: "id UInt64, name String, tags Array(String), metadata Map(String, String), created_at DateTime64(3)", + expected: []NestedField{ + {Name: "id", Type: "UInt64"}, + {Name: "name", Type: "String"}, + {Name: "tags", Type: "Array(String)"}, + {Name: "metadata", Type: "Map(String, String)"}, + {Name: "created_at", Type: "DateTime64(3)"}, + }, + }, + + // Error cases + { + name: "empty input", + input: "", + wantErr: true, + errMsg: "empty nested type definition", + }, + { + name: "whitespace only", + input: " ", + wantErr: true, + errMsg: "empty nested type definition", + }, + { + name: "missing type", + input: "fieldname", + wantErr: true, + errMsg: "expected 'name Type' format", + }, + { + name: "unbalanced open paren", + input: "a Array(String", + wantErr: true, + errMsg: "unbalanced parentheses", + }, + { + name: "unbalanced close paren", + input: "a String)", + wantErr: true, + errMsg: "unbalanced parentheses", + }, + { + name: "empty field between commas", + input: "a UInt32, , b String", + wantErr: true, + errMsg: "empty field definition", + }, + { + name: "trailing comma is tolerated", + input: "a UInt32, b String,", + expected: []NestedField{ + {Name: "a", Type: "UInt32"}, + {Name: "b", Type: "String"}, + }, + }, + { + name: "invalid name with parens", + input: "a() UInt32", + wantErr: true, + errMsg: "invalid field name", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + fields, err := ParseNestedFields(tt.input) + + if tt.wantErr { + require.Error(t, err) + if tt.errMsg != "" { + assert.Contains(t, err.Error(), tt.errMsg) + } + return + } + + require.NoError(t, err) + require.Equal(t, len(tt.expected), len(fields), "field count mismatch") + + for i, expected := range tt.expected { + assert.Equal(t, expected.Name, fields[i].Name, "field %d name", i) + assert.Equal(t, expected.Type, fields[i].Type, "field %d type", i) + } + }) + } +} + +func TestParseNestedType(t *testing.T) { + tests := []struct { + name string + input ColumnType + expected []NestedField + wantErr bool + errMsg string + }{ + { + name: "simple nested type", + input: "Nested(a UInt32, b String)", + expected: []NestedField{ + {Name: "a", Type: "UInt32"}, + {Name: "b", Type: "String"}, + }, + }, + { + name: "nested with array", + input: "Nested(id UInt64, tags Array(String))", + expected: []NestedField{ + {Name: "id", Type: "UInt64"}, + {Name: "tags", Type: "Array(String)"}, + }, + }, + { + name: "complex nested type", + input: "Nested(x Array(Array(Int8)), y Map(String, Int32), z Nullable(Float64))", + expected: []NestedField{ + {Name: "x", Type: "Array(Array(Int8))"}, + {Name: "y", Type: "Map(String, Int32)"}, + {Name: "z", Type: "Nullable(Float64)"}, + }, + }, + { + name: "not a nested type", + input: "Array(String)", + wantErr: true, + errMsg: "expected Nested type", + }, + { + name: "empty nested", + input: "Nested()", + wantErr: true, + errMsg: "empty Nested type definition", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + fields, err := ParseNestedType(tt.input) + + if tt.wantErr { + require.Error(t, err) + if tt.errMsg != "" { + assert.Contains(t, err.Error(), tt.errMsg) + } + return + } + + require.NoError(t, err) + require.Equal(t, len(tt.expected), len(fields), "field count mismatch") + + for i, expected := range tt.expected { + assert.Equal(t, expected.Name, fields[i].Name, "field %d name", i) + assert.Equal(t, expected.Type, fields[i].Type, "field %d type", i) + } + }) + } +} + +func TestParseNameTypePair(t *testing.T) { + tests := []struct { + name string + input string + expected NestedField + wantErr bool + }{ + { + name: "simple", + input: "id UInt64", + expected: NestedField{Name: "id", Type: "UInt64"}, + }, + { + name: "with params", + input: "arr Array(String)", + expected: NestedField{Name: "arr", Type: "Array(String)"}, + }, + { + name: "underscore name", + input: "user_id UInt64", + expected: NestedField{Name: "user_id", Type: "UInt64"}, + }, + { + name: "multiple spaces", + input: "field Type", + expected: NestedField{Name: "field", Type: "Type"}, + }, + { + name: "empty", + input: "", + wantErr: true, + }, + { + name: "no type", + input: "fieldonly", + wantErr: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + field, err := parseNameTypePair(tt.input) + + if tt.wantErr { + require.Error(t, err) + return + } + + require.NoError(t, err) + assert.Equal(t, tt.expected.Name, field.Name) + assert.Equal(t, tt.expected.Type, field.Type) + }) + } +}