Test pr 997

converter.go

@@ -7,6 +7,7 @@
 	"database/sql"
 	"database/sql/driver"
 	"encoding/hex"
+	"encoding/json"
 	"fmt"
 	"math"
 	"math/big"
@@ -91,10 +92,10 @@
 }
 
 // snowflakeTypeToGo translates Snowflake data type to Go data type.
-func snowflakeTypeToGo(dbtype snowflakeType, scale int64) reflect.Type {
-	switch dbtype {
+func snowflakeTypeToGo(rowType execResponseRowType) reflect.Type {
+	switch getSnowflakeType(rowType.Type) {
 	case fixedType:
-		if scale == 0 {
+		if rowType.Scale == 0 {
 			return reflect.TypeOf(int64(0))
 		}
 		return reflect.TypeOf(float64(0))
@@ -108,8 +109,22 @@
 		return reflect.TypeOf([]byte{})
 	case booleanType:
 		return reflect.TypeOf(true)
+	case vectorType:
+		if len(rowType.Fields) != 1 {
+			logger.Errorf("invalid result metadata fields for vector: length=%d", len(rowType.Fields))
+			return reflect.TypeOf("")
+		}
+		switch getSnowflakeType(rowType.Fields[0].Type) {
+		case fixedType:
+			return reflect.TypeOf([]int32{})
+		case realType:
+			return reflect.TypeOf([]float32{})
+		default:
+			logger.Errorf("invalid element type for vector: %s", rowType.Fields[0].Type)
+			return reflect.TypeOf("")
+		}
 	}
-	logger.Errorf("unsupported dbtype is specified. %v", dbtype)
+	logger.Errorf("unsupported dbtype is specified. %v", rowType.Type)
 	return reflect.TypeOf("")
 }
 
@@ -329,6 +344,27 @@
 		}
 		*dest = b
 		return nil
+	case "vector":
+		if len(srcColumnMeta.Fields) != 1 {
+			return fmt.Errorf("invalid result metadata fields for vector: length=%d", len(srcColumnMeta.Fields))
+		}
+		switch getSnowflakeType(srcColumnMeta.Fields[0].Type) {
+		case fixedType:
+			values := make([]int32, 0, srcColumnMeta.VectorDimension)
+			if err := json.Unmarshal([]byte(*srcValue), &values); err != nil {
+				return err
+			}
+			*dest = values
+		case realType:
+			values := make([]float32, 0, srcColumnMeta.VectorDimension)
+			if err := json.Unmarshal([]byte(*srcValue), &values); err != nil {
+				return err
+			}
+			*dest = values
+		default:
+			return fmt.Errorf("invalid element type for vector: %s", srcColumnMeta.Fields[0].Type)
+		}
+		return nil
 	}
 	*dest = *srcValue
 	return nil
@@ -610,6 +646,33 @@
 			}
 		}
 		return err
+	case vectorType:
+		vectorData := srcValue.(*array.FixedSizeList)
+		datatype := vectorData.DataType().(*arrow.FixedSizeListType)
+		dim := int(datatype.Len())
+		switch datatype.Elem().ID() {
+		case arrow.INT32:
+			values := vectorData.ListValues().(*array.Int32).Int32Values()
+			for i := 0; i < vectorData.Len(); i++ {
+				if vectorData.IsNull(i) {
+					destcol[i] = []int32(nil)
+				} else {
+					destcol[i] = values[i*dim : (i+1)*dim]
+				}
+			}
+		case arrow.FLOAT32:
+			values := vectorData.ListValues().(*array.Float32).Float32Values()
+			for i := 0; i < vectorData.Len(); i++ {
+				if vectorData.IsNull(i) {
+					destcol[i] = []float32(nil)
+				} else {
+					destcol[i] = values[i*dim : (i+1)*dim]
+				}
+			}
+		default:
+			return fmt.Errorf("unsupported element type %q for a vector", datatype.Elem().String())
+		}
+		return err
 	}
 
 	return fmt.Errorf("unsupported data type")

converter_test.go

@@ -139,35 +139,41 @@ func TestGoTypeToSnowflake(t *testing.T) {
 }
 
 type tcSnowflakeTypeToGo struct {
-	in    snowflakeType
-	scale int64
-	out   reflect.Type
+	in  execResponseRowType
+	out reflect.Type
 }
 
 func TestSnowflakeTypeToGo(t *testing.T) {
 	testcases := []tcSnowflakeTypeToGo{
-		{in: fixedType, scale: 0, out: reflect.TypeOf(int64(0))},
-		{in: fixedType, scale: 2, out: reflect.TypeOf(float64(0))},
-		{in: realType, scale: 0, out: reflect.TypeOf(float64(0))},
-		{in: textType, scale: 0, out: reflect.TypeOf("")},
-		{in: dateType, scale: 0, out: reflect.TypeOf(time.Now())},
-		{in: timeType, scale: 0, out: reflect.TypeOf(time.Now())},
-		{in: timestampLtzType, scale: 0, out: reflect.TypeOf(time.Now())},
-		{in: timestampNtzType, scale: 0, out: reflect.TypeOf(time.Now())},
-		{in: timestampTzType, scale: 0, out: reflect.TypeOf(time.Now())},
-		{in: objectType, scale: 0, out: reflect.TypeOf("")},
-		{in: variantType, scale: 0, out: reflect.TypeOf("")},
-		{in: arrayType, scale: 0, out: reflect.TypeOf("")},
-		{in: binaryType, scale: 0, out: reflect.TypeOf([]byte{})},
-		{in: booleanType, scale: 0, out: reflect.TypeOf(true)},
-		{in: sliceType, scale: 0, out: reflect.TypeOf("")},
+		{in: execResponseRowType{Type: "FIXED"}, out: reflect.TypeOf(int64(0))},
+		{in: execResponseRowType{Type: "FIXED", Scale: 2}, out: reflect.TypeOf(float64(0))},
+		{in: execResponseRowType{Type: "REAL"}, out: reflect.TypeOf(float64(0))},
+		{in: execResponseRowType{Type: "TEXT"}, out: reflect.TypeOf("")},
+		{in: execResponseRowType{Type: "TIME"}, out: reflect.TypeOf(time.Now())},
+		{in: execResponseRowType{Type: "TIME"}, out: reflect.TypeOf(time.Now())},
+		{in: execResponseRowType{Type: "TIMESTAMP_LTZ"}, out: reflect.TypeOf(time.Now())},
+		{in: execResponseRowType{Type: "TIMESTAMP_NTZ"}, out: reflect.TypeOf(time.Now())},
+		{in: execResponseRowType{Type: "TIMESTAMP_TZ"}, out: reflect.TypeOf(time.Now())},
+		{in: execResponseRowType{Type: "OBJECT"}, out: reflect.TypeOf("")},
+		{in: execResponseRowType{Type: "VARIANT"}, out: reflect.TypeOf("")},
+		{in: execResponseRowType{Type: "ARRAY"}, out: reflect.TypeOf("")},
+		{in: execResponseRowType{Type: "BINARY"}, out: reflect.TypeOf([]byte{})},
+		{in: execResponseRowType{Type: "BOOLEAN"}, out: reflect.TypeOf(true)},
+		{in: execResponseRowType{Type: "SLICE"}, out: reflect.TypeOf("")},
+		{
+			in:  execResponseRowType{Type: "VECTOR", Fields: []execResponseRowFieldType{{Type: "FIXED"}}},
+			out: reflect.TypeOf([]int32{}),
+		},
+		{
+			in:  execResponseRowType{Type: "VECTOR", Fields: []execResponseRowFieldType{{Type: "REAL"}}},
+			out: reflect.TypeOf([]float32{}),
+		},
 	}
 	for _, test := range testcases {
 		t.Run(fmt.Sprintf("%v_%v", test.in, test.out), func(t *testing.T) {
-			a := snowflakeTypeToGo(test.in, test.scale)
+			a := snowflakeTypeToGo(test.in)
 			if a != test.out {
-				t.Errorf("failed. in: %v, scale: %v, expected: %v, got: %v",
-					test.in, test.scale, test.out, a)
+				t.Errorf("failed. in: %v, expected: %v, got: %v", test.in, test.out, a)
 			}
 		})
 	}
@@ -849,6 +855,64 @@ func TestArrowToValue(t *testing.T) {
 				return -1
 			},
 		},
+		{
+			logical: "vector",
+			values:  [][]int32{nil, {1, 2, 3}},
+			builder: array.NewFixedSizeListBuilder(pool, 3, &arrow.Int32Type{}),
+			append: func(b array.Builder, vs interface{}) {
+				for _, v := range vs.([][]int32) {
+					lb := b.(*array.FixedSizeListBuilder)
+					vb := lb.ValueBuilder().(*array.Int32Builder)
+					if len(v) == 0 {
+						lb.AppendNull()
+						vb.AppendValues([]int32{-1, -1, -1}, nil)
+						continue
+					}
+
+					lb.Append(true)
+					for _, e := range v {
+						vb.Append(e)
+					}
+				}
+			},
+			compare: func(src interface{}, dst []snowflakeValue) int {
+				for i, v := range src.([][]int32) {
+					if !reflect.DeepEqual(v, dst[i]) {
+						return i
+					}
+				}
+				return -1
+			},
+		},
+		{
+			logical: "vector",
+			values:  [][]float32{nil, {1.1, 2.2, 3}},
+			builder: array.NewFixedSizeListBuilder(pool, 3, &arrow.Float32Type{}),
+			append: func(b array.Builder, vs interface{}) {
+				for _, v := range vs.([][]float32) {
+					lb := b.(*array.FixedSizeListBuilder)
+					vb := lb.ValueBuilder().(*array.Float32Builder)
+					if len(v) == 0 {
+						lb.AppendNull()
+						vb.AppendValues([]float32{-1, -1, -1}, nil)
+						continue
+					}
+
+					lb.Append(true)
+					for _, e := range v {
+						vb.Append(e)
+					}
+				}
+			},
+			compare: func(src interface{}, dst []snowflakeValue) int {
+				for i, v := range src.([][]float32) {
+					if !reflect.DeepEqual(v, dst[i]) {
+						return i
+					}
+				}
+				return -1
+			},
+		},
 	} {
 		testName := tc.logical
 		if tc.physical != "" {

datatype.go

@@ -26,6 +26,7 @@
 	binaryType
 	timeType
 	booleanType
+	vectorType
 	// the following are not snowflake types per se but internal types
 	nullType
 	sliceType
@@ -47,6 +48,7 @@
 	"BINARY":        binaryType,
 	"TIME":          timeType,
 	"BOOLEAN":       booleanType,
+	"VECTOR":        vectorType,
 	"NULL":          nullType,
 	"SLICE":         sliceType,
 	"CHANGE_TYPE":   changeType,
@@ -104,6 +106,8 @@
 	DataTypeTime = []byte{timeType.Byte()}
 	// DataTypeBoolean is a BOOLEAN datatype.
 	DataTypeBoolean = []byte{booleanType.Byte()}
+	// DataTypeVector is a VECTOR datatype.
+	DataTypeVector = []byte{vectorType.Byte()}
 )
 
 // dataTypeMode returns the subsequent data type in a string representation.
@@ -131,6 +135,40 @@
 	return tsmode, nil
 }
 
+type vectorElements interface {
+	~int32 | ~float32
+}
+
+// SQLVector is a wrapper type used to support deserializing SQL values into slices
+// in database/sql scans. Cast slice pointers as *SQLVector[T] when passing them to
+// a database/sql Scan method. The slice will be populated with the corresponding
+// column value when the scan completes.
+//
+// Here is an example:
+//
+//	var v []int32
+//	err := rows.Scan((*SQLVector[int32])(&v))
+type SQLVector[T vectorElements] []T
+
+// Vector is syntactic sugar for wrapping slices in SQLVector[t] so that they
+// can be deserialized in database/sql scans.
+//
+// Here is an example:
+//
+//	var v []int32
+//	err := rows.Scan(Vector(&v))
+func Vector[T vectorElements](value *[]T) *SQLVector[T] {
+	return (*SQLVector[T])(value)
+}
+
+func (v *SQLVector[T]) Scan(src any) error {
+	if vec, ok := src.([]T); ok {
+		*v = vec
+		return nil
+	}
+	return fmt.Errorf("cannot convert %T to a vector of type %T", src, *v)
+}
+
 // SnowflakeParameter includes the columns output from SHOW PARAMETER command.
 type SnowflakeParameter struct {
 	Key                       string