From 4c5f7326ee9db6ac25f431a4cf0b9c9df7ca6ec2 Mon Sep 17 00:00:00 2001 From: Yahor Yuzefovich Date: Fri, 25 Mar 2022 12:04:52 -0700 Subject: [PATCH] sql: fix handling of INT2VECTOR and OIDVECTOR in some cases 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. --- pkg/sql/logictest/testdata/logic_test/array | 23 +++++++++++++++++- pkg/sql/rowenc/keyside/array.go | 3 +++ pkg/sql/rowenc/valueside/array.go | 6 ++++- pkg/sql/rowenc/valueside/array_test.go | 2 +- pkg/sql/rowenc/valueside/decode.go | 2 +- pkg/sql/rowenc/valueside/legacy.go | 2 +- pkg/sql/sem/tree/cast.go | 3 +++ pkg/sql/sem/tree/datum.go | 27 +++++++++++++++++---- 8 files changed, 58 insertions(+), 10 deletions(-) diff --git a/pkg/sql/logictest/testdata/logic_test/array b/pkg/sql/logictest/testdata/logic_test/array index f32598c3f9bc..097097afccf4 100644 --- a/pkg/sql/logictest/testdata/logic_test/array +++ b/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 diff --git a/pkg/sql/rowenc/keyside/array.go b/pkg/sql/rowenc/keyside/array.go index f68b88b79c5f..2c66a524d89b 100644 --- a/pkg/sql/rowenc/keyside/array.go +++ b/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 { diff --git a/pkg/sql/rowenc/valueside/array.go b/pkg/sql/rowenc/valueside/array.go index 55d361762bed..4ea8c4efc4c6 100644 --- a/pkg/sql/rowenc/valueside/array.go +++ b/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) { diff --git a/pkg/sql/rowenc/valueside/array_test.go b/pkg/sql/rowenc/valueside/array_test.go index 7704c961ed6a..47f44c0afe68 100644 --- a/pkg/sql/rowenc/valueside/array_test.go +++ b/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) diff --git a/pkg/sql/rowenc/valueside/decode.go b/pkg/sql/rowenc/valueside/decode.go index 3b6ff9b6da3a..1d93cc92015d 100644 --- a/pkg/sql/rowenc/valueside/decode.go +++ b/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: diff --git a/pkg/sql/rowenc/valueside/legacy.go b/pkg/sql/rowenc/valueside/legacy.go index 1fa6f04d8d60..49ffc653a472 100644 --- a/pkg/sql/rowenc/valueside/legacy.go +++ b/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: diff --git a/pkg/sql/sem/tree/cast.go b/pkg/sql/sem/tree/cast.go index d66ff8b47ef7..5455564f3786 100644 --- a/pkg/sql/sem/tree/cast.go +++ b/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 { diff --git a/pkg/sql/sem/tree/datum.go b/pkg/sql/sem/tree/datum.go index 947422168f29..a9b2817b0a09 100644 --- a/pkg/sql/sem/tree/datum.go +++ b/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,8 +5464,8 @@ 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 @@ -5456,8 +5473,8 @@ func NewDIntVectorFromDArray(d *DArray) Datum { 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