Skip to content

Commit fe80908

Browse files
rogerykrogeryk
authored
Support [n]byte/[]byte type Scan/Append to FixedString column (#1205)
1 parent 34af932 commit fe80908

File tree

2 files changed

+97
-5
lines changed

2 files changed

+97
-5
lines changed

‎lib/column/fixed_string.go

+70-1
Original file line numberDiff line numberDiff line change
@@ -22,9 +22,10 @@ import (
2222
"database/sql/driver"
2323
"encoding"
2424
"fmt"
25-
"github.com/ClickHouse/ch-go/proto"
2625
"reflect"
2726

27+
"github.com/ClickHouse/ch-go/proto"
28+
2829
"github.com/ClickHouse/clickhouse-go/v2/lib/binary"
2930
)
3031

@@ -77,7 +78,27 @@ func (col *FixedString) ScanRow(dest any, row int) error {
7778
**d = col.row(row)
7879
case encoding.BinaryUnmarshaler:
7980
return d.UnmarshalBinary(col.rowBytes(row))
81+
case *[]byte:
82+
*d = col.rowBytes(row)
8083
default:
84+
// handle for *[n]byte
85+
if t := reflect.TypeOf(dest); t.Kind() == reflect.Pointer &&
86+
t.Elem().Kind() == reflect.Array &&
87+
t.Elem().Elem() == reflect.TypeOf(byte(0)) {
88+
size := t.Elem().Len()
89+
if size != col.col.Size {
90+
return &ColumnConverterError{
91+
Op: "ScanRow",
92+
To: fmt.Sprintf("%T", dest),
93+
From: "FixedString",
94+
Hint: fmt.Sprintf("invalid size %d, expect %d", size, col.col.Size),
95+
}
96+
}
97+
rv := reflect.ValueOf(dest).Elem()
98+
reflect.Copy(rv, reflect.ValueOf(col.row(row)))
99+
return nil
100+
}
101+
81102
if scan, ok := dest.(sql.Scanner); ok {
82103
return scan.Scan(col.row(row))
83104
}
@@ -125,7 +146,39 @@ func (col *FixedString) Append(v any) (nulls []uint8, err error) {
125146
}
126147
col.col.Append(data)
127148
nulls = make([]uint8, len(data)/col.col.Size)
149+
case [][]byte:
150+
nulls = make([]uint8, len(v))
151+
for i, v := range v {
152+
if v == nil {
153+
nulls[i] = 1
154+
}
155+
n := len(v)
156+
if n == 0 {
157+
col.col.Append(make([]byte, col.col.Size))
158+
} else if n >= col.col.Size {
159+
col.col.Append(v[0:col.col.Size])
160+
} else {
161+
data := make([]byte, col.col.Size)
162+
copy(data, v)
163+
col.col.Append(data)
164+
}
165+
}
128166
default:
167+
// handle for [][n]byte
168+
if t := reflect.TypeOf(v); t.Kind() == reflect.Slice &&
169+
t.Elem().Kind() == reflect.Array &&
170+
t.Elem().Elem() == reflect.TypeOf(byte(0)) {
171+
rv := reflect.ValueOf(v)
172+
nulls = make([]uint8, rv.Len())
173+
for i := 0; i < rv.Len(); i++ {
174+
e := rv.Index(i)
175+
data := make([]byte, e.Len())
176+
reflect.Copy(reflect.ValueOf(data), e)
177+
col.col.Append(data)
178+
}
179+
return
180+
}
181+
129182
if s, ok := v.(driver.Valuer); ok {
130183
val, err := s.Value()
131184
if err != nil {
@@ -150,6 +203,8 @@ func (col *FixedString) Append(v any) (nulls []uint8, err error) {
150203
func (col *FixedString) AppendRow(v any) (err error) {
151204
data := make([]byte, col.col.Size)
152205
switch v := v.(type) {
206+
case []byte:
207+
copy(data, v)
153208
case string:
154209
if v != "" {
155210
data = binary.Str2Bytes(v, col.col.Size)
@@ -166,6 +221,20 @@ func (col *FixedString) AppendRow(v any) (err error) {
166221
return err
167222
}
168223
default:
224+
if t := reflect.TypeOf(v); t.Kind() == reflect.Array && t.Elem() == reflect.TypeOf(byte(0)) {
225+
if t.Len() != col.col.Size {
226+
return &ColumnConverterError{
227+
Op: "AppendRow",
228+
To: "FixedString",
229+
From: fmt.Sprintf("%T", v),
230+
Hint: fmt.Sprintf("invalid size %d, expect %d", t.Len(), col.col.Size),
231+
}
232+
}
233+
reflect.Copy(reflect.ValueOf(data), reflect.ValueOf(v))
234+
col.col.Append(data)
235+
return nil
236+
}
237+
169238
if s, ok := v.(driver.Valuer); ok {
170239
val, err := s.Value()
171240
if err != nil {

‎tests/fixed_string_test.go

+27-4
Original file line numberDiff line numberDiff line change
@@ -22,9 +22,10 @@ import (
2222
"crypto/rand"
2323
"database/sql/driver"
2424
"fmt"
25-
"github.com/stretchr/testify/require"
2625
"testing"
2726

27+
"github.com/stretchr/testify/require"
28+
2829
"github.com/ClickHouse/clickhouse-go/v2"
2930
"github.com/stretchr/testify/assert"
3031
)
@@ -56,6 +57,8 @@ func TestFixedString(t *testing.T) {
5657
, Col4 Array(FixedString(10))
5758
, Col5 Array(Nullable(FixedString(10)))
5859
, Col6 FixedString(12)
60+
, Col7 FixedString(10)
61+
, Col8 FixedString(10)
5962
) Engine MergeTree() ORDER BY tuple()
6063
`
6164
defer func() {
@@ -71,10 +74,12 @@ func TestFixedString(t *testing.T) {
7174
col4Data = []string{"ClickHouse", "ClickHouse", "ClickHouse"}
7275
col5Data = []*string{&col1Data, nil, &col1Data}
7376
col6Data = "clickhouse"
77+
col7Data = []byte("clickhouse")
78+
col8Data = [10]byte{99, 108, 105, 99, 107, 104, 111, 117, 115, 101}
7479
)
7580
_, err = rand.Read(col2Data.data[:])
7681
require.NoError(t, err)
77-
require.NoError(t, batch.Append(col1Data, col2Data, col3Data, col4Data, col5Data, col6Data))
82+
require.NoError(t, batch.Append(col1Data, col2Data, col3Data, col4Data, col5Data, col6Data, col7Data, col8Data))
7883
require.Equal(t, 1, batch.Rows())
7984
require.NoError(t, batch.Send())
8085
var (
@@ -84,14 +89,18 @@ func TestFixedString(t *testing.T) {
8489
col4 []string
8590
col5 []*string
8691
col6 string
92+
col7 []byte
93+
col8 [10]byte
8794
)
88-
require.NoError(t, conn.QueryRow(ctx, "SELECT * FROM test_fixed_string").Scan(&col1, &col2, &col3, &col4, &col5, &col6))
95+
require.NoError(t, conn.QueryRow(ctx, "SELECT * FROM test_fixed_string").Scan(&col1, &col2, &col3, &col4, &col5, &col6, &col7, &col8))
8996
assert.Equal(t, col1Data, col1)
9097
assert.Equal(t, col2Data.data, col2.data)
9198
assert.Equal(t, col3Data, col3)
9299
assert.Equal(t, col4Data, col4)
93100
assert.Equal(t, col5Data, col5)
94101
assert.Equal(t, col6Data+string([]byte{0, 0}), col6)
102+
assert.Equal(t, col7Data, col7)
103+
assert.Equal(t, col8Data, col8)
95104
rows, err := conn.Query(ctx, "SELECT CAST('RU' AS FixedString(2)) FROM system.numbers_mt LIMIT 10")
96105
require.NoError(t, err)
97106
var count int
@@ -206,6 +215,8 @@ func TestColumnarFixedString(t *testing.T) {
206215
, Col3 Nullable(FixedString(10))
207216
, Col4 Array(FixedString(10))
208217
, Col5 Array(Nullable(FixedString(10)))
218+
, Col6 FixedString(10)
219+
, Col7 FixedString(10)
209220
) Engine MergeTree() ORDER BY tuple()
210221
`
211222
defer func() {
@@ -220,6 +231,8 @@ func TestColumnarFixedString(t *testing.T) {
220231
col3Data = &col1Data
221232
col4Data = []string{"ClickHouse", "ClickHouse", "ClickHouse"}
222233
col5Data = []*string{&col1Data, nil, &col1Data}
234+
col6Data = []byte("clickhouse")
235+
col7Data = [10]byte{99, 108, 105, 99, 107, 104, 111, 117, 115, 101}
223236
)
224237
require.NoError(t, batch.Column(0).Append([]string{
225238
col1Data, col1Data, col1Data, col1Data, col1Data,
@@ -236,6 +249,12 @@ func TestColumnarFixedString(t *testing.T) {
236249
require.NoError(t, batch.Column(4).Append([][]*string{
237250
col5Data, col5Data, col5Data, col5Data, col5Data,
238251
}))
252+
require.NoError(t, batch.Column(5).Append([][]byte{
253+
col6Data, col6Data, col6Data, col6Data, col6Data,
254+
}))
255+
require.NoError(t, batch.Column(6).Append([][10]byte{
256+
col7Data, col7Data, col7Data, col7Data, col7Data,
257+
}))
239258
require.Equal(t, 5, batch.Rows())
240259
require.NoError(t, batch.Send())
241260
var (
@@ -244,13 +263,17 @@ func TestColumnarFixedString(t *testing.T) {
244263
col3 *string
245264
col4 []string
246265
col5 []*string
266+
col6 []byte
267+
col7 [10]byte
247268
)
248-
require.NoError(t, conn.QueryRow(ctx, "SELECT * FROM test_fixed_string LIMIT 1").Scan(&col1, &col2, &col3, &col4, &col5))
269+
require.NoError(t, conn.QueryRow(ctx, "SELECT * FROM test_fixed_string LIMIT 1").Scan(&col1, &col2, &col3, &col4, &col5, &col6, &col7))
249270
assert.Equal(t, col1Data, col1)
250271
assert.Equal(t, col2Data, col2)
251272
assert.Equal(t, col3Data, col3)
252273
assert.Equal(t, col4Data, col4)
253274
assert.Equal(t, col5Data, col5)
275+
assert.Equal(t, col6Data, col6)
276+
assert.Equal(t, col7Data, col7)
254277
}
255278

256279
func BenchmarkFixedString(b *testing.B) {

0 commit comments

Comments
 (0)