planner: support ScalarSubQuery to display them in EXPLAIN (#45252)

close #22076

executor/executor.go

@@ -2083,6 +2083,7 @@ func ResetContextOfStmt(ctx sessionctx.Context, s ast.StmtNode) (err error) {
 	if explainStmt, ok := s.(*ast.ExplainStmt); ok {
 		sc.InExplainStmt = true
 		sc.ExplainFormat = explainStmt.Format
+		sc.InExplainAnalyzeStmt = explainStmt.Analyze
 		sc.IgnoreExplainIDSuffix = strings.ToLower(explainStmt.Format) == types.ExplainFormatBrief
 		sc.InVerboseExplain = strings.ToLower(explainStmt.Format) == types.ExplainFormatVerbose
 		s = explainStmt.Stmt
@@ -2091,6 +2092,7 @@ func ResetContextOfStmt(ctx sessionctx.Context, s ast.StmtNode) (err error) {
 	}
 	if explainForStmt, ok := s.(*ast.ExplainForStmt); ok {
 		sc.InExplainStmt = true
+		sc.InExplainAnalyzeStmt = true
 		sc.InVerboseExplain = strings.ToLower(explainForStmt.Format) == types.ExplainFormatVerbose
 	}
 

expression/expression.go

@@ -49,6 +49,7 @@ const (
 	columnFlag         byte = 1
 	scalarFunctionFlag byte = 3
 	parameterFlag      byte = 4
+	ScalarSubQFlag     byte = 5
 )
 
 // EvalAstExpr evaluates ast expression directly.

planner/core/BUILD.bazel

@@ -70,6 +70,7 @@ go_library(
         "rule_topn_push_down.go",
         "runtime_filter.go",
         "runtime_filter_generator.go",
+        "scalar_subq_expression.go",
         "show_predicate_extractor.go",
         "stats.go",
         "stringer.go",

planner/core/casetest/scalarsubquery/BUILD.bazel

@@ -0,0 +1,20 @@
+load("@io_bazel_rules_go//go:def.bzl", "go_test")
+
+go_test(
+    name = "scalarsubquery_test",
+    timeout = "short",
+    srcs = [
+        "cases_test.go",
+        "main_test.go",
+    ],
+    data = glob(["testdata/**"]),
+    flaky = True,
+    deps = [
+        "//testkit",
+        "//testkit/testdata",
+        "//testkit/testmain",
+        "//testkit/testsetup",
+        "@com_github_stretchr_testify//require",
+        "@org_uber_go_goleak//:goleak",
+    ],
+)

planner/core/casetest/scalarsubquery/cases_test.go

@@ -0,0 +1,92 @@
+// Copyright 2023 PingCAP, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package scalarsubquery
+
+import (
+	"fmt"
+	"testing"
+
+	"github.com/pingcap/tidb/testkit"
+	"github.com/pingcap/tidb/testkit/testdata"
+	"github.com/stretchr/testify/require"
+)
+
+func TestExplainNonEvaledSubquery(t *testing.T) {
+	var (
+		input []struct {
+			SQL              string
+			IsExplainAnalyze bool
+			HasErr           bool
+		}
+		output []struct {
+			SQL   string
+			Plan  []string
+			Error string
+		}
+	)
+	planSuiteData := GetPlanSuiteData()
+	planSuiteData.LoadTestCases(t, &input, &output)
+	store := testkit.CreateMockStore(t)
+	tk := testkit.NewTestKit(t, store)
+
+	tk.MustExec("use test")
+	tk.MustExec("create table t1(a int, b int, c int)")
+	tk.MustExec("create table t2(a int, b int, c int)")
+	tk.MustExec("create table t3(a varchar(5), b varchar(5), c varchar(5))")
+	tk.MustExec("set @@tidb_opt_enable_non_eval_scalar_subquery=true")
+
+	cutExecutionInfoFromExplainAnalyzeOutput := func(rows [][]interface{}) [][]interface{} {
+		// The columns are id, estRows, actRows, task type, access object, execution info, operator info, memory, disk
+		// We need to cut the unstable output of execution info, memory and disk.
+		for i := range rows {
+			rows[i] = rows[i][:6] // cut the final memory and disk.
+			rows[i] = append(rows[i][:5], rows[i][6:]...)
+		}
+		return rows
+	}
+
+	for i, ts := range input {
+		testdata.OnRecord(func() {
+			output[i].SQL = ts.SQL
+			if ts.HasErr {
+				err := tk.ExecToErr(ts.SQL)
+				require.NotNil(t, err, fmt.Sprintf("Failed at case #%d", i))
+				output[i].Error = err.Error()
+				output[i].Plan = nil
+			} else {
+				rows := tk.MustQuery(ts.SQL).Rows()
+				if ts.IsExplainAnalyze {
+					rows = cutExecutionInfoFromExplainAnalyzeOutput(rows)
+				}
+				output[i].Plan = testdata.ConvertRowsToStrings(rows)
+				output[i].Error = ""
+			}
+		})
+		if ts.HasErr {
+			err := tk.ExecToErr(ts.SQL)
+			require.NotNil(t, err, fmt.Sprintf("Failed at case #%d", i))
+		} else {
+			rows := tk.MustQuery(ts.SQL).Rows()
+			if ts.IsExplainAnalyze {
+				rows = cutExecutionInfoFromExplainAnalyzeOutput(rows)
+			}
+			require.Equal(t,
+				testdata.ConvertRowsToStrings(testkit.Rows(output[i].Plan...)),
+				testdata.ConvertRowsToStrings(rows),
+				fmt.Sprintf("Failed at case #%d, SQL: %v", i, ts.SQL),
+			)
+		}
+	}
+}

planner/core/casetest/scalarsubquery/main_test.go

@@ -0,0 +1,52 @@
+// Copyright 2023 PingCAP, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package scalarsubquery
+
+import (
+	"flag"
+	"testing"
+
+	"github.com/pingcap/tidb/testkit/testdata"
+	"github.com/pingcap/tidb/testkit/testmain"
+	"github.com/pingcap/tidb/testkit/testsetup"
+	"go.uber.org/goleak"
+)
+
+var testDataMap = make(testdata.BookKeeper)
+
+func TestMain(m *testing.M) {
+	testsetup.SetupForCommonTest()
+	flag.Parse()
+	testDataMap.LoadTestSuiteData("testdata", "plan_suite")
+	opts := []goleak.Option{
+		goleak.IgnoreTopFunction("github.com/golang/glog.(*fileSink).flushDaemon"),
+		goleak.IgnoreTopFunction("github.com/lestrrat-go/httprc.runFetchWorker"),
+		goleak.IgnoreTopFunction("go.etcd.io/etcd/client/pkg/v3/logutil.(*MergeLogger).outputLoop"),
+		goleak.IgnoreTopFunction("gopkg.in/natefinch/lumberjack%2ev2.(*Logger).millRun"),
+		goleak.IgnoreTopFunction("github.com/tikv/client-go/v2/txnkv/transaction.keepAlive"),
+		goleak.IgnoreTopFunction("go.opencensus.io/stats/view.(*worker).start"),
+	}
+
+	callback := func(i int) int {
+		testDataMap.GenerateOutputIfNeeded()
+		return i
+	}
+
+	goleak.VerifyTestMain(testmain.WrapTestingM(m, callback), opts...)
+}
+
+func GetPlanSuiteData() testdata.TestData {
+	return testDataMap["plan_suite"]
+}

planner/core/casetest/scalarsubquery/testdata/plan_suite_in.json

@@ -0,0 +1,61 @@
+[
+  {
+    "name": "TestExplainNonEvaledSubquery",
+    "cases": [
+      // Test normal non-correlated scalar sub query.
+      {
+        "SQL": "explain format = 'brief' select * from t1 where a = (select a from t2 limit 1)",
+        "IsExplainAnalyze": false,
+        "HasErr": false
+      },
+      {
+        "SQL": "explain analyze format = 'brief' select * from t1 where a = (select a from t2 limit 1)",
+        "IsExplainAnalyze": true,
+        "HasErr": false
+      },
+      // Test EXISTS non-correlated scalar sub query.
+      {
+        "SQL": "explain format = 'brief' select * from t1 where exists(select 1 from t2 where a = 1)",
+        "IsExplainAnalyze": false,
+        "HasErr": false
+      },
+      {
+        "SQL": "explain analyze format = 'brief' select * from t1 where exists(select 1 from t2 where a = 1)",
+        "IsExplainAnalyze": true,
+        "HasErr": false
+      },
+      {
+        "SQL": "explain format = 'brief' select * from t1 where not exists(select 1 from t2 where a = 1)",
+        "IsExplainAnalyze": false,
+        "HasErr": false
+      },
+      {
+        "SQL": "explain analyze format = 'brief' select * from t1 where not exists(select 1 from t2 where a = 1)",
+        "IsExplainAnalyze": true,
+        "HasErr": false
+      },
+      // Test with constant propagation.
+      {
+        "SQL": "explain format = 'brief' select * from t1 where exists(select 1 from t2 where a = (select a from t3 limit 1) and b = a);",
+        "IsExplainAnalyze": false,
+        "HasErr": false
+      },
+      {
+        "SQL": "explain analyze format = 'brief' select * from t1 where exists(select 1 from t2 where a = (select a from t3 limit 1) and b = a);",
+        "IsExplainAnalyze": true,
+        "HasErr": false
+      },
+      // Test multiple returns.
+      {
+        "SQL": "explain format = 'brief' select * from t1 where (a, b) = (select a, b from t2 limit 1)",
+        "IsExplainAnalyze": false,
+        "HasErr": false
+      },
+      {
+        "SQL": "explain analyze format = 'brief' select * from t1 where (a, b) = (select a, b from t2 limit 1)",
+        "IsExplainAnalyze": true,
+        "HasErr": false
+      }
+    ]
+  }
+]

planner/core/casetest/scalarsubquery/testdata/plan_suite_out.json

@@ -0,0 +1,119 @@
+[
+  {
+    "Name": "TestExplainNonEvaledSubquery",
+    "Cases": [
+      {
+        "SQL": "explain format = 'brief' select * from t1 where a = (select a from t2 limit 1)",
+        "Plan": [
+          "Selection 8000.00 root  eq(test.t1.a, ScalarQueryCol#9)",
+          "└─TableReader 10000.00 root  data:TableFullScan",
+          "  └─TableFullScan 10000.00 cop[tikv] table:t1 keep order:false, stats:pseudo",
+          "ScalarSubQuery N/A root  Output: ScalarQueryCol#9",
+          "└─MaxOneRow 1.00 root  ",
+          "  └─Limit 1.00 root  offset:0, count:1",
+          "    └─TableReader 1.00 root  data:Limit",
+          "      └─Limit 1.00 cop[tikv]  offset:0, count:1",
+          "        └─TableFullScan 1.00 cop[tikv] table:t2 keep order:false, stats:pseudo"
+        ],
+        "Error": ""
+      },
+      {
+        "SQL": "explain analyze format = 'brief' select * from t1 where a = (select a from t2 limit 1)",
+        "Plan": [
+          "TableDual 0.00 0 root "
+        ],
+        "Error": ""
+      },
+      {
+        "SQL": "explain format = 'brief' select * from t1 where exists(select 1 from t2 where a = 1)",
+        "Plan": [
+          "Selection 8000.00 root  ScalarQueryCol#10",
+          "└─TableReader 10000.00 root  data:TableFullScan",
+          "  └─TableFullScan 10000.00 cop[tikv] table:t1 keep order:false, stats:pseudo",
+          "ScalarSubQuery N/A root  Output: ScalarQueryCol#10",
+          "└─TableReader 10.00 root  data:Selection",
+          "  └─Selection 10.00 cop[tikv]  eq(test.t2.a, 1)",
+          "    └─TableFullScan 10000.00 cop[tikv] table:t2 keep order:false, stats:pseudo"
+        ],
+        "Error": ""
+      },
+      {
+        "SQL": "explain analyze format = 'brief' select * from t1 where exists(select 1 from t2 where a = 1)",
+        "Plan": [
+          "TableDual 0.00 0 root "
+        ],
+        "Error": ""
+      },
+      {
+        "SQL": "explain format = 'brief' select * from t1 where not exists(select 1 from t2 where a = 1)",
+        "Plan": [
+          "Selection 8000.00 root  not(ScalarQueryCol#10)",
+          "└─TableReader 10000.00 root  data:TableFullScan",
+          "  └─TableFullScan 10000.00 cop[tikv] table:t1 keep order:false, stats:pseudo",
+          "ScalarSubQuery N/A root  Output: ScalarQueryCol#10",
+          "└─TableReader 10.00 root  data:Selection",
+          "  └─Selection 10.00 cop[tikv]  eq(test.t2.a, 1)",
+          "    └─TableFullScan 10000.00 cop[tikv] table:t2 keep order:false, stats:pseudo"
+        ],
+        "Error": ""
+      },
+      {
+        "SQL": "explain analyze format = 'brief' select * from t1 where not exists(select 1 from t2 where a = 1)",
+        "Plan": [
+          "TableReader 10000.00 0 root ",
+          "└─TableFullScan 10000.00 0 cop[tikv] table:t1"
+        ],
+        "Error": ""
+      },
+      {
+        "SQL": "explain format = 'brief' select * from t1 where exists(select 1 from t2 where a = (select a from t3 limit 1) and b = a);",
+        "Plan": [
+          "Selection 8000.00 root  ScalarQueryCol#15",
+          "└─TableReader 10000.00 root  data:TableFullScan",
+          "  └─TableFullScan 10000.00 cop[tikv] table:t1 keep order:false, stats:pseudo",
+          "ScalarSubQuery N/A root  Output: ScalarQueryCol#13",
+          "└─MaxOneRow 1.00 root  ",
+          "  └─Limit 1.00 root  offset:0, count:1",
+          "    └─TableReader 1.00 root  data:Limit",
+          "      └─Limit 1.00 cop[tikv]  offset:0, count:1",
+          "        └─TableFullScan 1.00 cop[tikv] table:t3 keep order:false, stats:pseudo",
+          "ScalarSubQuery N/A root  Output: ScalarQueryCol#15",
+          "└─Selection 6400.00 root  eq(cast(test.t2.a, double BINARY), cast(ScalarQueryCol#13, double BINARY)), eq(cast(test.t2.b, double BINARY), cast(ScalarQueryCol#13, double BINARY))",
+          "  └─TableReader 8000.00 root  data:Selection",
+          "    └─Selection 8000.00 cop[tikv]  eq(test.t2.b, test.t2.a)",
+          "      └─TableFullScan 10000.00 cop[tikv] table:t2 keep order:false, stats:pseudo"
+        ],
+        "Error": ""
+      },
+      {
+        "SQL": "explain analyze format = 'brief' select * from t1 where exists(select 1 from t2 where a = (select a from t3 limit 1) and b = a);",
+        "Plan": [
+          "TableDual 0.00 0 root "
+        ],
+        "Error": ""
+      },
+      {
+        "SQL": "explain format = 'brief' select * from t1 where (a, b) = (select a, b from t2 limit 1)",
+        "Plan": [
+          "Selection 8000.00 root  eq(test.t1.a, ScalarQueryCol#9), eq(test.t1.b, ScalarQueryCol#10)",
+          "└─TableReader 10000.00 root  data:TableFullScan",
+          "  └─TableFullScan 10000.00 cop[tikv] table:t1 keep order:false, stats:pseudo",
+          "ScalarSubQuery N/A root  Output: ScalarQueryCol#9,ScalarQueryCol#10",
+          "└─MaxOneRow 1.00 root  ",
+          "  └─Limit 1.00 root  offset:0, count:1",
+          "    └─TableReader 1.00 root  data:Limit",
+          "      └─Limit 1.00 cop[tikv]  offset:0, count:1",
+          "        └─TableFullScan 1.00 cop[tikv] table:t2 keep order:false, stats:pseudo"
+        ],
+        "Error": ""
+      },
+      {
+        "SQL": "explain analyze format = 'brief' select * from t1 where (a, b) = (select a, b from t2 limit 1)",
+        "Plan": [
+          "TableDual 0.00 0 root "
+        ],
+        "Error": ""
+      }
+    ]
+  }
+]