Merge pull request ClickHouse#282 from afterdusk/read-nested-arrays

Add ability to read multidimensional arrays

Author: kshvakov

clickhouse_test.go

@@ -1255,57 +1255,89 @@ func TestArrayArrayT(t *testing.T) {
 	const (
 		ddl = `
 			CREATE TABLE clickhouse_test_array_array_t (
-				String Array(Array(String))
-				, String2 Array(String)
-				, Int32 Array(Int32)
+				String1 Array(Array(String)),
+				String2 Array(Array(Array(String))),
+				Int32   Array(Array(Int32))
 			) Engine=Memory
 		`
 		dml = `
-			INSERT INTO clickhouse_test_array_array_t (String, String2, Int32) VALUES (?)
+			INSERT INTO clickhouse_test_array_array_t (String1, String2, Int32) VALUES (?)
 		`
 		query = `
 			SELECT
-				String
+				String1,
+				String2,
+				Int32
 			FROM clickhouse_test_array_array_t
 		`
 	)
+
+	items := []struct {
+		String1, String2, Int32 interface{}
+	}{
+		{
+			[][]string{
+				[]string{"A"},
+				[]string{"BC"},
+				[]string{"DEF"},
+			},
+			[][][]string{
+				[][]string{
+					[]string{"X"},
+					[]string{"Y"},
+				},
+				[][]string{
+					[]string{"ZZ"},
+				},
+			},
+			[][]int32{
+				[]int32{1},
+				[]int32{2, 3},
+			},
+		},
+		{
+			[][][]byte{
+				[][]byte{[]byte("AA")},
+				[][]byte{[]byte("BB")},
+				[][]byte{[]byte("C4C")},
+			},
+			[][][][]byte{
+				[][][]byte{
+					[][]byte{[]byte("XX"), []byte("YY")},
+				},
+			},
+			[][]int32{
+				[]int32{4, 5, 6},
+			},
+		},
+	}
+
 	if connect, err := sql.Open("clickhouse", "tcp://127.0.0.1:9000?debug=true"); assert.NoError(t, err) && assert.NoError(t, connect.Ping()) {
 		if _, err := connect.Exec("DROP TABLE IF EXISTS clickhouse_test_array_array_t"); assert.NoError(t, err) {
 			if _, err := connect.Exec(ddl); assert.NoError(t, err) {
 				if tx, err := connect.Begin(); assert.NoError(t, err) {
 					if stmt, err := tx.Prepare(dml); assert.NoError(t, err) {
-						_, err = stmt.Exec([][]string{[]string{"A"}, []string{"B"}, []string{"C"}}, []string{"X", "Y"}, []int32{1, 2, 3})
-						if !assert.NoError(t, err) {
-							return
-						}
-						_, err = stmt.Exec([][]string{[]string{"AA"}, []string{"BB"}, []string{"C4C"}}, []string{"XX", "YY"}, []int32{4, 5, 6})
-						if !assert.NoError(t, err) {
-							return
-						}
-						_, err = stmt.Exec(
-							[][][]byte{
-								[][]byte{[]byte("AA")},
-								[][]byte{[]byte("BB")},
-								[][]byte{[]byte("C4C")},
-							},
-							[][]byte{[]byte("XX"), []byte("YY")},
-							[]int32{4, 5, 6},
-						)
-						if !assert.NoError(t, err) {
-							return
+						for _, item := range items {
+							_, err = stmt.Exec(item.String1, item.String2, item.Int32)
+							if !assert.NoError(t, err) {
+								return
+							}
 						}
-					} else {
-						return
+
 					}
 					if assert.NoError(t, tx.Commit()) {
-						/*	var value []string
-							if err := connect.QueryRow(query).Scan(&value); assert.NoError(t, err) {
-								if !assert.NoError(t, err) {
-									return
-								}
-							}
-							assert.Equal(t, []string{"A", "C"}, value)
-						*/
+						var result struct {
+							String1 [][]string
+							String2 [][][]string
+							Int32   [][]int32
+						}
+
+						row := connect.QueryRow(query)
+						if err := row.Scan(&result.String1, &result.String2, &result.Int32); assert.NoError(t, err) {
+							assert.Equal(t, items[0].String1, result.String1)
+							assert.Equal(t, items[0].String2, result.String2)
+							assert.Equal(t, items[0].Int32, result.Int32)
+						}
 					}
 				}
 			}

lib/column/array.go

@@ -26,32 +26,53 @@ func (array *Array) Write(encoder *binary.Encoder, v interface{}) error {
 
 func (array *Array) ReadArray(decoder *binary.Decoder, rows int) (_ []interface{}, err error) {
 	var (
+		offsets = make([][]uint64, array.depth)
 		values  = make([]interface{}, rows)
-		offsets = make([]uint64, rows)
 	)
-	for i := 0; i < rows; i++ {
-		offset, err := decoder.UInt64()
-		if err != nil {
-			return nil, err
+
+	// Read offsets
+	lastOffset := uint64(rows)
+	for i := 0; i < array.depth; i++ {
+		offset := make([]uint64, lastOffset)
+		for j := uint64(0); j < lastOffset; j++ {
+			if offset[j], err = decoder.UInt64(); err != nil {
+				return nil, err
+			}
 		}
 		offsets[i] = offset
-	}
-	for n, offset := range offsets {
-		ln := offset
-		if n != 0 {
-			ln = ln - offsets[n-1]
+		lastOffset = 0
+		if len(offset) > 0 {
+			lastOffset = offset[len(offset)-1]
 		}
-		if values[n], err = array.read(decoder, int(ln)); err != nil {
+	}
+
+	// Read values
+	for i := 0; i < rows; i++ {
+		if values[i], err = array.read(decoder, offsets, uint64(i), 0); err != nil {
 			return nil, err
 		}
 	}
 	return values, nil
 }
 
-func (array *Array) read(decoder *binary.Decoder, ln int) (interface{}, error) {
-	slice := reflect.MakeSlice(array.valueOf.Type(), 0, ln)
-	for i := 0; i < ln; i++ {
-		value, err := array.column.Read(decoder)
+func (array *Array) read(decoder *binary.Decoder, offsets [][]uint64, index uint64, level int) (interface{}, error) {
+	end := offsets[level][index]
+	start := uint64(0)
+	if index > 0 {
+		start = offsets[level][index-1]
+	}
+
+	slice := reflect.MakeSlice(array.arrayType(level), 0, int(end-start))
+	for i := start; i < end; i++ {
+		var (
+			value interface{}
+			err   error
+		)
+		if level == array.depth-1 {
+			value, err = array.column.Read(decoder)
+		} else {
+			value, err = array.read(decoder, offsets, i, level+1)
+		}
 		if err != nil {
 			return nil, err
 		}
@@ -60,6 +81,14 @@ func (array *Array) read(decoder *binary.Decoder, ln int) (interface{}, error) {
 	return slice.Interface(), nil
 }
 
+func (array *Array) arrayType(level int) reflect.Type {
+	t := array.column.ScanType()
+	for i := 0; i < array.depth-level; i++ {
+		t = reflect.SliceOf(t)
+	}
+	return t
+}
+
 func (array *Array) Depth() int {
 	return array.depth
 }