sql/jsonpath: validate JSON column in jsonb_path_exists for inverted index acceleration

Informs: #154729

Previously, the jsonb_path_exists function would attempt inverted index
acceleration even when the JSON parameter referenced a column that was not
the source column of the inverted index, or when it was a constant JSON
value. This resulted in incorrect query plans and potentially wrong results.

This commit adds proper validation to ensure that inverted index acceleration
is only used when the JSON parameter is a variable expression that references
the actual source column of the inverted index. For all other cases, the
query will fall back to a full table scan with post-filtering.

Release note: None

Author: ZhouXing19

pkg/sql/logictest/testdata/logic_test/jsonb_path_exists_index_acceleration

@@ -4,7 +4,8 @@
 statement ok
 CREATE TABLE json_tab (
   a INT PRIMARY KEY,
-  b JSONB
+  b JSONB,
+  c JSONB
 )
 
 statement ok
@@ -25,7 +26,7 @@ INSERT INTO json_tab VALUES
 
 # To confirm which rows contain the specified path.
 query IT
-SELECT * FROM json_tab@primary WHERE jsonb_path_exists(b, '$.a.b') ORDER BY a;
+SELECT a, b FROM json_tab@primary WHERE jsonb_path_exists(b, '$.a.b') ORDER BY a, b;
 ----
 3  {"a": {"b": "c", "d": "e"}}
 4  {"a": {"b": [1, 2, 3, 4]}}
@@ -35,7 +36,7 @@ SELECT * FROM json_tab@primary WHERE jsonb_path_exists(b, '$.a.b') ORDER BY a;
 9  {"a": {"b": "c"}}
 
 query IT
-SELECT * FROM json_tab@foo_inv WHERE jsonb_path_exists(b, '$.a.b') ORDER BY a;
+SELECT a, b FROM json_tab@foo_inv WHERE jsonb_path_exists(b, '$.a.b') ORDER BY a, b;
 ----
 3  {"a": {"b": "c", "d": "e"}}
 4  {"a": {"b": [1, 2, 3, 4]}}
@@ -505,3 +506,42 @@ SELECT a FROM json_tab@foo_inv WHERE jsonb_path_exists(b, '$.a[*][*] ? (@.b == 1
 16
 17
 18
+
+# If json parameter is not part of the inverted index, the filter should
+# not be index accelerated.
+statement ok
+INSERT INTO json_tab VALUES (19, '{"a": {"b": "c"}}', '{"a": {"b": "c"}}')
+
+query I
+SELECT a FROM json_tab@primary WHERE jsonb_path_exists(c, '$.a.b') ORDER BY a;
+----
+19
+
+statement error index "foo_inv" is inverted and cannot be used for this query
+SELECT a FROM json_tab@foo_inv WHERE jsonb_path_exists(c, '$.a.b') ORDER BY a;
+
+query I
+SELECT a FROM json_tab@primary WHERE jsonb_path_exists('{"a": {"b": "c"}}', '$.a.b') ORDER BY a;
+----
+1
+2
+3
+4
+5
+6
+7
+8
+9
+10
+11
+12
+13
+14
+15
+16
+17
+18
+19
+
+statement error index "foo_inv" is inverted and cannot be used for this query
+SELECT a FROM json_tab@foo_inv WHERE jsonb_path_exists('{"a": {"b": "c"}}', '$.a.b') ORDER BY a;

pkg/sql/opt/invertedidx/json_array.go

@@ -417,15 +417,19 @@ func (j *jsonOrArrayFilterPlanner) extractInvertedFilterConditionFromLeaf(
 	case *memo.FunctionExpr:
 		if t.Properties.Category == builtinconstants.CategoryJsonpath && t.Name == "jsonb_path_exists" {
 			if len(t.Args) > 1 {
-				if ce, ok := t.Args[1].(*memo.ConstExpr); ok {
-					if dJsonPath, ok := ce.Value.(*tree.DJsonpath); ok {
-						if dJsonPath.Strict {
-							return inverted.NonInvertedColExpression{}, expr, nil
+				// The first parameter has to be a column reference.
+				if isIndexColumn(j.tabID, j.index, t.Args[0], j.computedColumns) {
+					if ce, ok := t.Args[1].(*memo.ConstExpr); ok {
+						if dJsonPath, ok := ce.Value.(*tree.DJsonpath); ok {
+							if dJsonPath.Strict {
+								return inverted.NonInvertedColExpression{}, expr, nil
+							}
+							jp := dJsonPath.Path
+							invertedExpr = j.extractJSONPathCondition(ctx, evalCtx, jp)
 						}
-						jp := dJsonPath.Path
-						invertedExpr = j.extractJSONPathCondition(ctx, evalCtx, jp)
 					}
 				}
+
 			}
 		}
 	}