copy: fix vectorized copy for INT2 and INT4

This commit fixes a recently introduced bug where we forgot to implement
the special behavior in the vec handler for INT2 and INT4 types which
would then could lead to a runtime crash if such a type is used in the
schema. Int types are special because the vectorized engine handles them
precisely (i.e. it uses int16, int32, and int64 accordingly) whereas the
row engine always internally defaults to int64.

Release note (bug fix): In alpha and beta 23.1.0 releases CockroachDB
could crash when evaluating COPY command in some cases when the schema
had INT2 and / or INT4 type, and this is now fixed.

pkg/col/coldataext/vec_handler.go

@@ -37,7 +37,14 @@ func MakeVecHandler(vec coldata.Vec) tree.ValueHandler {
 	case types.DecimalFamily:
 		v.decimals = vec.Decimal()
 	case types.IntFamily:
-		v.ints = vec.Int64()
+		switch vec.Type().Width() {
+		case 16:
+			v.int16s = vec.Int16()
+		case 32:
+			v.int32s = vec.Int32()
+		default:
+			v.ints = vec.Int64()
+		}
 	case types.FloatFamily:
 		v.floats = vec.Float64()
 	case types.TimestampTZFamily:
@@ -55,13 +62,12 @@ func MakeVecHandler(vec coldata.Vec) tree.ValueHandler {
 }
 
 type vecHandler struct {
-	nulls    *coldata.Nulls
-	bools    coldata.Bools
-	bytes    *coldata.Bytes
-	decimals coldata.Decimals
-	// TODO(cucaroach): implement small int types
-	//int16s     coldata.Int16s
-	//int32s     coldata.Int32s
+	nulls      *coldata.Nulls
+	bools      coldata.Bools
+	bytes      *coldata.Bytes
+	decimals   coldata.Decimals
+	int16s     coldata.Int16s
+	int32s     coldata.Int32s
 	ints       coldata.Int64s
 	floats     coldata.Float64s
 	timestamps coldata.Times
@@ -134,6 +140,18 @@ func (v *vecHandler) Float(f float64) {
 	v.row++
 }
 
+// Int16 is part of the tree.ValueHandler interface.
+func (v *vecHandler) Int16(i int16) {
+	v.int16s[v.row] = i
+	v.row++
+}
+
+// Int32 is part of the tree.ValueHandler interface.
+func (v *vecHandler) Int32(i int32) {
+	v.int32s[v.row] = i
+	v.row++
+}
+
 // Int is part of the tree.ValueHandler interface.
 func (v *vecHandler) Int(i int64) {
 	v.ints[v.row] = i

pkg/sql/copy/copy_in_test.go

@@ -160,6 +160,8 @@ func TestCopyFromRandom(t *testing.T) {
 			cs TEXT COLLATE en_us_u_ks_level2,
 			o BOOL,
 			i INT,
+			i2 INT2,
+			i4 INT4,
 			f FLOAT,
 			e DECIMAL,
 			t TIME,
@@ -184,7 +186,7 @@ func TestCopyFromRandom(t *testing.T) {
 		t.Fatal(err)
 	}
 
-	stmt, err := txn.Prepare(pq.CopyInSchema("d", "t", "id", "n", "cs", "o", "i", "f", "e", "t", "ttz", "ts", "s", "b", "u", "ip", "tz", "geography", "geometry", "box2d"))
+	stmt, err := txn.Prepare(pq.CopyInSchema("d", "t", "id", "n", "cs", "o", "i", "i2", "i4", "f", "e", "t", "ttz", "ts", "s", "b", "u", "ip", "tz", "geography", "geometry", "box2d"))
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -196,6 +198,8 @@ func TestCopyFromRandom(t *testing.T) {
 		types.MakeCollatedString(types.String, "en_us_u_ks_level2"),
 		types.Bool,
 		types.Int,
+		types.Int2,
+		types.Int4,
 		types.Float,
 		types.Decimal,
 		types.Time,

pkg/sql/sem/tree/parse_string.go

@@ -170,6 +170,8 @@ type ValueHandler interface {
 	// Decimal returns a pointer into the vec for in place construction.
 	Decimal() *apd.Decimal
 	Float(f float64)
+	Int16(i int16)
+	Int32(i int32)
 	Int(i int64)
 	Duration(d duration.Duration)
 	JSON(j json.JSON)
@@ -219,10 +221,25 @@ func ParseAndRequireStringHandler(
 		}
 	case types.IntFamily:
 		var i int64
-		if i, err = strconv.ParseInt(s, 0, 64); err == nil {
-			vh.Int(i)
-		} else {
-			err = MakeParseError(s, types.Int, err)
+		switch t.Width() {
+		case 16:
+			if i, err = strconv.ParseInt(s, 0, 16); err == nil {
+				vh.Int16(int16(i))
+			} else {
+				err = MakeParseError(s, t, err)
+			}
+		case 32:
+			if i, err = strconv.ParseInt(s, 0, 32); err == nil {
+				vh.Int32(int32(i))
+			} else {
+				err = MakeParseError(s, t, err)
+			}
+		default:
+			if i, err = strconv.ParseInt(s, 0, 64); err == nil {
+				vh.Int(i)
+			} else {
+				err = MakeParseError(s, t, err)
+			}
 		}
 	case types.JsonFamily:
 		var j json.JSON

pkg/sql/sem/tree/parse_string_test.go

@@ -151,6 +151,8 @@ func (a *anyHandler) Bool(b bool)                  { a.val = b }
 func (a *anyHandler) Bytes(b []byte)               { a.val = b }
 func (a *anyHandler) Decimal() *apd.Decimal        { return &a.dec }
 func (a *anyHandler) Float(f float64)              { a.val = f }
+func (a *anyHandler) Int16(i int16)                { a.val = i }
+func (a *anyHandler) Int32(i int32)                { a.val = i }
 func (a *anyHandler) Int(i int64)                  { a.val = i }
 func (a *anyHandler) Duration(d duration.Duration) { a.val = d }
 func (a *anyHandler) JSON(j json.JSON)             { a.val = j }