Merge #151983 #154783

151983: sqlsmith: add support for DO block r=ZhouXing19 a=ZhouXing19

informs #106368

A `DO` block can be either a top-level sql statement or be part of the plpgsql function body. In this commit we introduce the DO block for both of these 2 usages. We also introduced `DisableDoBlocks` which disables both usages of a DO block.

Release note: None

154783: sql/jsonpath: validate JSON column in jsonb_path_exists for inverted index acceleration r=ZhouXing19 a=ZhouXing19

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. 

We also don't allow index acceleration for jsonb_path_exists filter if it takes a variable list in its parameters. For example, `jsonb_path_exists(b, '$.a ? (`@.b` == $x)', '{"x": "c"}')` (where `'{"x": "c"}'` is the variable list) is not allowed.

Release note: None

Co-authored-by: ZhouXing19 <[email protected]>

Author: craig[bot]

pkg/ccl/changefeedccl/changefeed_test.go

@@ -1692,6 +1692,7 @@ func TestChangefeedRandomExpressions(t *testing.T) {
 			sqlsmith.DisableWith(),
 			sqlsmith.DisableMutations(),
 			sqlsmith.DisableLimits(),
+			sqlsmith.DisableDoBlocks(),
 			sqlsmith.DisableAggregateFuncs(),
 			sqlsmith.DisableWindowFuncs(),
 			sqlsmith.DisableJoins(),

pkg/cmd/smith/main.go

@@ -66,6 +66,7 @@ var (
 		"DisableCrossJoins":          sqlsmith.DisableCrossJoins(),
 		"DisableDDLs":                sqlsmith.DisableDDLs(),
 		"DisableDecimals":            sqlsmith.DisableDecimals(),
+		"DisableDoBlocks":            sqlsmith.DisableDoBlocks(),
 		"DisableEverything":          sqlsmith.DisableEverything(),
 		"DisableIndexHints":          sqlsmith.DisableIndexHints(),
 		"DisableInsertSelect":        sqlsmith.DisableInsertSelect(),

pkg/internal/sqlsmith/plpgsql.go

@@ -141,6 +141,7 @@ var (
 		{2, makePLpgSQLReturn},
 		{2, makePLpgSQLIf},
 		{2, makePLpgSQLWhile},
+		{2, makeDoBlockWithinScope},
 		{2, makePLpgSQLForLoop},
 		{5, makePLpgSQLNull},
 		{10, makePLpgSQLAssign},
@@ -395,3 +396,20 @@ func (s *plpgsqlBlockScope) addVariable(name string, typ *types.T, constant bool
 		s.constants[name] = struct{}{}
 	}
 }
+
+func (s *Smither) makeDoBlockTreeStmt() (*tree.DoBlock, bool) {
+	scope := makeBlockScope(0, types.Unknown, tree.RoutineVolatile)
+	doBlock, ok := makeDoBlockWithinScope(s, scope)
+	if !ok {
+		return nil, false
+	}
+	return &tree.DoBlock{Code: doBlock.(*ast.DoBlock)}, true
+}
+
+func makeDoBlockWithinScope(s *Smither, scope plpgsqlBlockScope) (ast.Statement, bool) {
+	if s.disableDoBlock {
+		return nil, false
+	}
+	block := s.makePLpgSQLBlock(scope)
+	return &ast.DoBlock{Block: block}, true
+}

pkg/internal/sqlsmith/relational.go

@@ -110,9 +110,11 @@ var (
 		{1, makeImport},
 		{1, makeCreateStats},
 		{1, makeSetSessionCharacteristics},
+		{1, makeDoBlock},
 	}
 	nonMutatingStatements = []statementWeight{
 		{10, makeSelect},
+		{1, makeDoBlock},
 	}
 	allStatements = append(mutatingStatements, nonMutatingStatements...)
 
@@ -928,6 +930,10 @@ func makeCreateFunc(s *Smither) (tree.Statement, bool) {
 	return s.makeCreateFunc()
 }
 
+func makeDoBlock(s *Smither) (tree.Statement, bool) {
+	return s.makeDoBlockTreeStmt()
+}
+
 func makeDropFunc(s *Smither) (tree.Statement, bool) {
 	if s.disableUDFCreation {
 		return nil, false

pkg/internal/sqlsmith/sqlsmith.go

@@ -118,7 +118,11 @@ type Smither struct {
 	// disableUDFCreation indicates whether we're not allowed to create UDFs.
 	// It follows that if we haven't created any UDFs, we have no UDFs to invoke
 	// too.
-	disableUDFCreation     bool
+	disableUDFCreation bool
+	// disableDoBlock indicates whether we're not allowed to create DO blocks,
+	// both as top level statements and inside plpgsql function body
+	// definition.
+	disableDoBlock         bool
 	disableIsolationChange bool
 
 	bulkSrv     *httptest.Server
@@ -612,6 +616,11 @@ var DisableUDFs = simpleOption("disable udfs", func(s *Smither) {
 	s.disableUDFCreation = true
 })
 
+// DisableDoBlocks causes the Smither to disable DO blocks.
+var DisableDoBlocks = simpleOption("disable do block", func(s *Smither) {
+	s.disableDoBlock = true
+})
+
 // DisableIsolationChange causes the Smither to disable stmts that modify the
 // txn isolation level.
 var DisableIsolationChange = simpleOption("disable isolation change", func(s *Smither) {

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,53 @@ 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;
+
+# We don't index accelerate jsonb_path_exists clause with variables.
+query I
+SELECT a FROM json_tab@primary WHERE jsonb_path_exists(b, '$.a ? (@.b == $x)', '{"x": "c"}') ORDER BY a;
+----
+3
+9
+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(b, '$.a ? (@.b == $x)', '{"x": "c"}') ORDER BY a;

pkg/sql/opt/invertedidx/json_array.go

@@ -416,16 +416,20 @@ func (j *jsonOrArrayFilterPlanner) extractInvertedFilterConditionFromLeaf(
 		invertedExpr = j.extractArrayOverlapsCondition(ctx, evalCtx, t.Left, t.Right)
 	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
+			if len(t.Args) == 2 {
+				// 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)
 					}
 				}
+
 			}
 		}
 	}