opt: inverted-index accelerate filters of the form j->'a' = '{"b": "c"}'

Previously, the optimizer did not plan inverted index scans for queries with
the JSON fetch value operator `->` when an object or array was on the right
side of the equality operator `=`, only when it was a boolean, string, number,
or null.

This change allows the inverted index to be used in these types of queries. It
supports objects with multiple key/value pairs, nested objects, arrays, arrays
nested within objects, and objects nested within arrays.

Fixes: #59605

Release note: None

pkg/sql/logictest/testdata/logic_test/inverted_index

@@ -752,7 +752,22 @@ INSERT INTO f VALUES
   (10, '{"a": {"b": {"c": 1}}}'),
   (11, '{"a": {"b": {"c": 1, "d": 2}}}}'),
   (12, '{"a": {"b": {"d": 2}}}}'),
-  (13, '{"a": {"b": {"c": [1, 2]}}}')
+  (13, '{"a": {"b": {"c": [1, 2]}}}'),
+  (14, '{"a": {"b": {"c": [1, 2, 3]}}}'),
+  (15, '{"a": []}'),
+  (16, '{"a": {}}}'),
+  (17, '{"a": {"b": "c"}}'),
+  (18, '{"a": {"b": ["c", "d", "e"]}}'),
+  (19, '{"a": ["b", "c", "d", "e"]}'),
+  (20, '{"a": ["b", "e", "c", "d"]}'),
+  (21, '{"z": {"a": "b", "c": "d"}}'),
+  (22, '{"z": {"a": "b", "c": "d", "e": "f"}}'),
+  (23, '{"a": "b", "x": ["c", "d", "e"]}}'),
+  (24, '{"a": "b", "c": [{"d": 1}, {"e": 2}]}}'),
+  (25, '{"a": {"b": "c", "d": "e"}}'),
+  (26, '{"a": {"b": "c"}, "d": "e"}'),
+  (27, '[1, 2, {"b": "c"}]'),
+  (28, '[{"a": {"b": "c"}}, "d", "e"]')
 
 query T
 SELECT j FROM f@i WHERE j->'a' = '1' ORDER BY k
@@ -789,6 +804,82 @@ SELECT j FROM f@i WHERE j->'a'->'b'->'c' = '1' ORDER BY k
 {"a": {"b": {"c": 1}}}
 {"a": {"b": {"c": 1, "d": 2}}}
 
+query T
+SELECT j FROM f@i WHERE j->'a' = '[]' ORDER BY k
+----
+{"a": []}
+
+query T
+SELECT j FROM f@i WHERE j->'a' = '{}' ORDER BY k
+----
+{"a": {}}
+
+query T
+SELECT j FROM f@i WHERE j->'a' = '["b"]' ORDER BY k
+----
+
+query T
+SELECT j FROM f@i WHERE j->'a' = '"b"' ORDER BY k
+----
+{"a": "b", "x": ["c", "d", "e"]}
+{"a": "b", "c": [{"d": 1}, {"e": 2}]}
+
+query T
+SELECT j FROM f@i WHERE j->'a' = '{"b": "c"}' ORDER BY k
+----
+{"a": {"b": "c"}}
+{"a": {"b": "c"}, "d": "e"}
+
+query T
+SELECT j FROM f@i WHERE j->'a'->'b'->'c' = '[1, 2]' ORDER BY k
+----
+{"a": {"b": {"c": [1, 2]}}}
+
+query T
+SELECT j FROM f@i WHERE j->'z' = '{"a": "b", "c": "d"}' ORDER BY k
+----
+{"z": {"a": "b", "c": "d"}}
+
+query T
+SELECT j FROM f@i WHERE j->'a' = '["b", "c", "d", "e"]' ORDER BY k
+----
+{"a": ["b", "c", "d", "e"]}
+
+query T
+SELECT j FROM f@i WHERE j->'a' = '["b", "c", "d", "e"]' OR  j->'a' = '["b", "e", "c", "d"]' ORDER BY k
+----
+{"a": ["b", "c", "d", "e"]}
+{"a": ["b", "e", "c", "d"]}
+
+query T
+SELECT j FROM f@i WHERE j->'a' = '{"b": ["c", "d", "e"]}' ORDER BY k
+----
+{"a": {"b": ["c", "d", "e"]}}
+
+query T
+SELECT j FROM f@i WHERE j->'a'->'b' = '["c", "d", "e"]' ORDER BY k
+----
+{"a": {"b": ["c", "d", "e"]}}
+
+query T
+SELECT j FROM f@i WHERE j->'z'->'c' = '"d"' ORDER BY k
+----
+{"z": {"a": "b", "c": "d"}}
+{"z": {"a": "b", "c": "d", "e": "f"}}
+
+query T
+SELECT j FROM f@i WHERE j->'z' = '{"c": "d"}' ORDER BY k
+----
+
+query T
+SELECT j FROM f@i WHERE j->'a' = '"b"' AND j->'c' = '[{"d": 1}]' ORDER BY k
+----
+
+query T
+SELECT j FROM f@i WHERE j->'a' = '"b"' AND j->'c' = '[{"d": 1}, {"e": 2}]' ORDER BY k
+----
+{"a": "b", "c": [{"d": 1}, {"e": 2}]}
+
 subtest arrays
 
 statement ok

pkg/sql/opt/invertedidx/json_array.go

@@ -342,23 +342,18 @@ func (j *jsonOrArrayFilterPlanner) extractJSONOrArrayContainsCondition(
 // expression in the form [col]->[index0]->[index1]->...->[indexN] where col is
 // a variable or expression referencing the inverted column in the inverted
 // index and each index is a constant string. The right expression must be a
-// constant JSON value that is not an object or an array.
+// constant JSON value.
 func (j *jsonOrArrayFilterPlanner) extractJSONFetchValEqCondition(
 	evalCtx *tree.EvalContext, left *memo.FetchValExpr, right opt.ScalarExpr,
 ) inverted.Expression {
-	// The right side of the equals expression should be a constant JSON value
-	// that is not an object or array.
+	// The right side of the equals expression should be a constant JSON value.
 	if !memo.CanExtractConstDatum(right) {
 		return inverted.NonInvertedColExpression{}
 	}
 	val, ok := memo.ExtractConstDatum(right).(*tree.DJSON)
 	if !ok {
 		return inverted.NonInvertedColExpression{}
 	}
-	typ := val.JSON.Type()
-	if typ == json.ObjectJSONType || typ == json.ArrayJSONType {
-		return inverted.NonInvertedColExpression{}
-	}
 
 	// Recursively traverse fetch val expressions and collect keys with which to
 	// build the InvertedExpression. If it is not possible to build an inverted
@@ -431,5 +426,15 @@ func (j *jsonOrArrayFilterPlanner) extractJSONFetchValEqCondition(
 		obj = b.Build()
 	}
 
-	return getInvertedExprForJSONOrArrayIndex(evalCtx, tree.NewDJSON(obj))
+	invertedExpr := getInvertedExprForJSONOrArrayIndex(evalCtx, tree.NewDJSON(obj))
+
+	// When the right side is an array or object, the InvertedExpression
+	// generated is not tight. We must indicate it is non-tight so an additional
+	// filter is added.
+	typ := val.JSON.Type()
+	if typ == json.ArrayJSONType || typ == json.ObjectJSONType {
+		invertedExpr.SetNotTight()
+	}
+
+	return invertedExpr
 }

pkg/sql/opt/invertedidx/json_array_test.go

@@ -407,16 +407,22 @@ func TestTryFilterJsonOrArrayIndex(t *testing.T) {
 			ok:       false,
 		},
 		{
-			// Arrays on the right side of the equality are not yet supported.
-			filters:  "j->'a' = '[1]'",
-			indexOrd: jsonOrd,
-			ok:       false,
+			// Arrays on the right side of the equality are supported.
+			filters:          "j->'a' = '[1]'",
+			indexOrd:         jsonOrd,
+			ok:               true,
+			tight:            false,
+			unique:           true,
+			remainingFilters: "j->'a' = '[1]'",
 		},
 		{
-			// Objects on the right side of the equality are not yet supported.
-			filters:  `j->'a' = '{"b": "c"}'`,
-			indexOrd: jsonOrd,
-			ok:       false,
+			// Objects on the right side of the equality are supported.
+			filters:          `j->'a' = '{"b": "c"}'`,
+			indexOrd:         jsonOrd,
+			ok:               true,
+			tight:            false,
+			unique:           true,
+			remainingFilters: `j->'a' = '{"b": "c"}'`,
 		},
 		{
 			// Wrong index ordinal.