Merge #78520

78520: sql: fix handling of INT2VECTOR and OIDVECTOR in some cases r=yuzefovich a=yuzefovich

Previously, we incorrectly handled arrays with special oids
(`INT2VECTOR` and `OIDVECTOR`) in a couple of scenarios:
- when casting to those special types
- when the array datum was spilled to disk.

This now fixed.

Release note (bug fix): Previously, CockroachDB could lose INT2VECTOR
and OIDVECTOR type of some arrays, and this is now fixed.

Co-authored-by: Yahor Yuzefovich <[email protected]>

pkg/sql/logictest/testdata/logic_test/array

@@ -226,7 +226,12 @@ SELECT ARRAY[1,2,3]::TEXT[]
 query T
 SELECT ARRAY[1,2,3]::INT2VECTOR
 ----
-{1,2,3}
+1 2 3
+
+query T
+SELECT ARRAY[1,2,3]::OIDVECTOR
+----
+1 2 3
 
 # array subscript access
 
@@ -2223,3 +2228,19 @@ query T
 SELECT array_positions(NULL::letter[], 'b'::letter)
 ----
 NULL
+
+# Regression test for losing the OID of the array during spilling to disk.
+
+# Lower the distsql_workmem so that the sort in the query below has to spill to
+# disk.
+statement ok
+SET distsql_workmem = '2B'
+
+# The output must be printed without the curly braces.
+query T
+SELECT indkey FROM pg_index WHERE indrelid = (SELECT oid FROM pg_class WHERE relname = 'k') ORDER BY 1
+----
+1
+
+statement ok
+RESET distsql_workmem

pkg/sql/rowenc/keyside/array.go

@@ -53,6 +53,9 @@ func decodeArrayKey(
 	}
 
 	result := tree.NewDArray(t.ArrayContents())
+	if err = result.MaybeSetCustomOid(t); err != nil {
+		return nil, nil, err
+	}
 
 	for {
 		if len(buf) == 0 {

pkg/sql/rowenc/valueside/array.go

@@ -63,15 +63,19 @@ func encodeArray(d *tree.DArray, scratch []byte) ([]byte, error) {
 }
 
 // decodeArray decodes the value encoding for an array.
-func decodeArray(a *tree.DatumAlloc, elementType *types.T, b []byte) (tree.Datum, []byte, error) {
+func decodeArray(a *tree.DatumAlloc, arrayType *types.T, b []byte) (tree.Datum, []byte, error) {
 	header, b, err := decodeArrayHeader(b)
 	if err != nil {
 		return nil, b, err
 	}
+	elementType := arrayType.ArrayContents()
 	result := tree.DArray{
 		Array:    make(tree.Datums, header.length),
 		ParamTyp: elementType,
 	}
+	if err = result.MaybeSetCustomOid(arrayType); err != nil {
+		return nil, b, err
+	}
 	var val tree.Datum
 	for i := uint64(0); i < header.length; i++ {
 		if header.isNull(i) {

pkg/sql/rowenc/valueside/array_test.go

@@ -123,7 +123,7 @@ func TestArrayEncoding(t *testing.T) {
 		})
 
 		t.Run("decode "+test.name, func(t *testing.T) {
-			d, _, err := decodeArray(&tree.DatumAlloc{}, test.datum.ParamTyp, test.encoding)
+			d, _, err := decodeArray(&tree.DatumAlloc{}, types.MakeArray(test.datum.ParamTyp), test.encoding)
 			hasNulls := d.(*tree.DArray).HasNulls
 			if test.datum.HasNulls != hasNulls {
 				t.Fatalf("expected %v to have HasNulls=%t, got %t", test.encoding, test.datum.HasNulls, hasNulls)

pkg/sql/rowenc/valueside/decode.go

@@ -196,7 +196,7 @@ func DecodeUntaggedDatum(
 		if err != nil {
 			return nil, nil, err
 		}
-		return decodeArray(a, t.ArrayContents(), b)
+		return decodeArray(a, t, b)
 	case types.TupleFamily:
 		return decodeTuple(a, t, buf)
 	case types.EnumFamily:

pkg/sql/rowenc/valueside/legacy.go

@@ -334,7 +334,7 @@ func UnmarshalLegacy(a *tree.DatumAlloc, typ *types.T, value roachpb.Value) (tre
 		if err != nil {
 			return nil, err
 		}
-		datum, _, err := decodeArray(a, typ.ArrayContents(), v)
+		datum, _, err := decodeArray(a, typ, v)
 		// TODO(yuzefovich): do we want to create a new object via tree.DatumAlloc?
 		return datum, err
 	case types.JsonFamily:

pkg/sql/sem/tree/cast.go

@@ -2448,6 +2448,9 @@ func performCastWithoutPrecisionTruncation(
 			return res, err
 		case *DArray:
 			dcast := NewDArray(t.ArrayContents())
+			if err := dcast.MaybeSetCustomOid(t); err != nil {
+				return nil, err
+			}
 			for _, e := range v.Array {
 				ecast := DNull
 				if e != DNull {

pkg/sql/sem/tree/datum.go

@@ -4255,7 +4255,24 @@ func MustBeDArray(e Expr) *DArray {
 	return i
 }
 
-// ResolvedType implements the TypedExpr interface.
+// MaybeSetCustomOid checks whether t has a special oid that we want to set into
+// d. Must be kept in sync with DArray.ResolvedType. Returns an error if t is
+// not an array type.
+func (d *DArray) MaybeSetCustomOid(t *types.T) error {
+	if t.Family() != types.ArrayFamily {
+		return errors.AssertionFailedf("expected array type, got %s", t.SQLString())
+	}
+	switch t.Oid() {
+	case oid.T_int2vector:
+		d.customOid = oid.T_int2vector
+	case oid.T_oidvector:
+		d.customOid = oid.T_oidvector
+	}
+	return nil
+}
+
+// ResolvedType implements the TypedExpr interface. Must be kept in sync with
+// DArray.MaybeSetCustomOid.
 func (d *DArray) ResolvedType() *types.T {
 	switch d.customOid {
 	case oid.T_int2vector:
@@ -5447,17 +5464,17 @@ func NewDName(d string) Datum {
 	return NewDNameFromDString(NewDString(d))
 }
 
-// NewDIntVectorFromDArray is a helper routine to create a *DIntVector
-// (implemented as a *DOidWrapper) initialized from an existing *DArray.
+// NewDIntVectorFromDArray is a helper routine to create a new *DArray,
+// initialized from an existing *DArray, with the special oid for IntVector.
 func NewDIntVectorFromDArray(d *DArray) Datum {
 	ret := new(DArray)
 	*ret = *d
 	ret.customOid = oid.T_int2vector
 	return ret
 }
 
-// NewDOidVectorFromDArray is a helper routine to create a *DOidVector
-// (implemented as a *DOidWrapper) initialized from an existing *DArray.
+// NewDOidVectorFromDArray is a helper routine to create a new *DArray,
+// initialized from an existing *DArray, with the special oid for OidVector.
 func NewDOidVectorFromDArray(d *DArray) Datum {
 	ret := new(DArray)
 	*ret = *d