sql: fix behaviour of contains for arrays+scalars

Closes cockroachdb#23897.

The problem here wasn't actually with our inverted index queries - we
had the implementation of this operator wrong from Postgres. Only
*top-level* arrays contain scalars.

Release note (bug fix): fix behaviour of @> operator with arrays and
scalars.

pkg/sql/logictest/testdata/logic_test/json

@@ -472,6 +472,30 @@ SELECT '[1, 2, 3]'::JSONB @> '[1, 2]'::JSONB
 ----
 true
 
+query B
+SELECT '{"a": [1, 2, 3]}'::JSONB->'a' @> '2'::JSONB
+----
+true
+
+statement ok
+CREATE TABLE x (j JSONB)
+
+statement ok
+INSERT INTO x VALUES ('{"a": [1,2,3]}')
+
+query B
+SELECT true FROM x WHERE j->'a' @> '2'::JSONB
+----
+true
+
+statement ok
+CREATE INVERTED INDEX ON x (j)
+
+query B
+SELECT true FROM x WHERE j->'a' @> '2'::JSONB
+----
+true
+
 query T
 SELECT '{"foo": {"bar": 1}}'::JSONB #- ARRAY['foo', 'bar']
 ----

pkg/util/json/contains.go

@@ -31,6 +31,15 @@ import "sort"
 // also the arrays and objects in separate arrays, so that we can do the fast
 // thing for the subset of the arrays which are scalars.
 func Contains(a, b JSON) (bool, error) {
+	if a.Type() == ArrayJSONType && b.isScalar() {
+		decoded, err := a.tryDecode()
+		if err != nil {
+			return false, err
+		}
+		ary := decoded.(jsonArray)
+		return checkArrayContainsScalar(ary, b)
+	}
+
 	preA, err := a.preprocessForContains()
 	if err != nil {
 		return false, err
@@ -42,6 +51,24 @@ func Contains(a, b JSON) (bool, error) {
 	return preA.contains(preB)
 }
 
+// checkArrayContainsScalar performs a unique case of contains (and is
+// described as such in the Postgres docs) - a top-level array contains a
+// scalar which is an element of it.  This contradicts the general rule of
+// contains that the contained object must have the same "shape" as the
+// containing object.
+func checkArrayContainsScalar(ary jsonArray, s JSON) (bool, error) {
+	for _, j := range ary {
+		cmp, err := j.Compare(s)
+		if err != nil {
+			return false, err
+		}
+		if cmp == 0 {
+			return true, nil
+		}
+	}
+	return false, nil
+}
+
 // containsable is an interface used internally for the implementation of @>.
 type containsable interface {
 	contains(other containsable) (bool, error)
@@ -147,35 +174,6 @@ func (j containsableScalar) contains(other containsable) (bool, error) {
 }
 
 func (j containsableArray) contains(other containsable) (bool, error) {
-	// This is a unique case of contains (and is described as such in the
-	// Postgres docs) - an array contains a scalar which is an element of it.
-	// This contradicts the general rule of contains that the contained object
-	// must have the same "shape" as the containing object.
-	if o, ok := other.(containsableScalar); ok {
-		var err error
-		found := sort.Search(len(j.scalars), func(i int) bool {
-			if err != nil {
-				return false
-			}
-			var c int
-			c, err = j.scalars[i].JSON.Compare(o.JSON)
-			return c >= 0
-		})
-		if err != nil {
-			return false, err
-		}
-
-		if found >= len(j.scalars) {
-			return false, nil
-		}
-
-		c, err := j.scalars[found].JSON.Compare(o.JSON)
-		if err != nil {
-			return false, err
-		}
-		return c == 0, nil
-	}
-
 	if contained, ok := other.(containsableArray); ok {
 		// Since both slices of scalars are sorted via the preprocessing, we can
 		// step through them together via binary search.

pkg/util/json/contains_testers.go

@@ -29,6 +29,27 @@ type containsTester interface {
 	subdocument(isRoot bool, rng *rand.Rand) JSON
 }
 
+func slowContains(a, b JSON) bool {
+	// This is a unique case of contains (and is described as such in the
+	// Postgres docs) - an array contains a scalar which is an element of it.
+	// This contradicts the general rule of contains that the contained object
+	// must have the same "shape" as the containing object.
+	if a.Type() == ArrayJSONType {
+		ary := a.MaybeDecode().(jsonArray)
+		if b.isScalar() {
+			for _, j := range ary {
+				cmp, _ := j.Compare(b)
+				if cmp == 0 {
+					return true
+				}
+			}
+			return false
+		}
+	}
+
+	return a.(containsTester).slowContains(b)
+}
+
 func (j jsonNull) slowContains(other JSON) bool {
 	c, _ := j.Compare(other)
 	return c == 0
@@ -51,15 +72,6 @@ func (j jsonString) slowContains(other JSON) bool {
 }
 
 func (j jsonArray) slowContains(other JSON) bool {
-	if other.isScalar() {
-		for i := 0; i < len(j); i++ {
-			c, _ := j[i].Compare(other)
-			if c == 0 {
-				return true
-			}
-		}
-	}
-
 	other = other.MaybeDecode()
 	if ary, ok := other.(jsonArray); ok {
 		for i := 0; i < len(ary); i++ {

pkg/util/json/json_test.go

@@ -1463,6 +1463,7 @@ func TestJSONContains(t *testing.T) {
 			{`{"a": [3], "c": {}}`, true},
 			{`{"a": [4], "c": {}}`, false},
 			{`{"a": [3], "c": {"foo": "gup"}}`, false},
+			{`{"a": 3}`, false},
 		},
 		`[{"a": 1}, {"b": 2, "c": 3}, [1], true, false, null, "hello"]`: {
 			{`[]`, true},
@@ -1483,6 +1484,13 @@ func TestJSONContains(t *testing.T) {
 			{`["hello", "hello", []]`, true},
 			{`["hello", {"a": 1}, "hello", []]`, true},
 			{`["hello", {"a": 1, "b": 2}, "hello", []]`, false},
+			{`"hello"`, true},
+			{`true`, true},
+			{`false`, true},
+			{`null`, true},
+			{`[1]`, false},
+			{`[[1]]`, true},
+			{`1`, false},
 		},
 		`[{"Ck@P":{"7RZ2\"mZBH":null,"|__v":[1.685483]},"EM%&":{"I{TH":[],"}p@]7sIKC\\$":[]},"f}?#z~":{"#e9>m\"v75'&+":false,"F;,+&r9":{}}},[{}],false,0.498401]`: {
 			{`[false,{}]`, true},
@@ -1506,7 +1514,7 @@ func TestJSONContains(t *testing.T) {
 					t.Fatal(err)
 				}
 
-				checkResult := left.(containsTester).slowContains(other)
+				checkResult := slowContains(left, other)
 				if result != checkResult {
 					t.Fatal("mismatch between actual contains and slowContains")
 				}
@@ -1540,7 +1548,7 @@ func TestPositiveRandomJSONContains(t *testing.T) {
 		if !c {
 			t.Fatal(fmt.Sprintf("%s should contain %s", j, subdoc))
 		}
-		if !j.(containsTester).slowContains(subdoc) {
+		if !slowContains(j, subdoc) {
 			t.Fatal(fmt.Sprintf("%s should slowContains %s", j, subdoc))
 		}
 	}
@@ -1564,9 +1572,8 @@ func TestNegativeRandomJSONContains(t *testing.T) {
 		if err != nil {
 			t.Fatal(err)
 		}
-		slowResult := j1.(containsTester).slowContains(j2)
+		slowResult := slowContains(j1, j2)
 		if realResult != slowResult {
-			fmt.Println("realResult=", realResult)
 			t.Fatal("mismatch for document " + j1.String() + " @> " + j2.String())
 		}
 	}