planner: fix projection of generate column for indexLookup task (ping…

…cap#31684) (pingcap#32049) close pingcap#31682
ti-srebot · Apr 25, 2022 · 11dbb91 · 11dbb91
1 parent 9e4a150
commit 11dbb91
Show file tree

Hide file tree

Showing 5 changed files with 38 additions and 15 deletions.
diff --git a/cmd/explaintest/r/explain_generate_column_substitute.result b/cmd/explaintest/r/explain_generate_column_substitute.result
@@ -577,3 +577,15 @@ select reverse(c_str) from t use index(idx);
 reverse(c_str)
 nnamreh lufituaeb
 drop table t;
+drop table if exists t1,t2;
+create table t1 (c_int int, c_str varchar(40) ,primary key (c_int)  , key(c_str(36))  , key((c_int + 1))) partition by hash (c_int) partitions 4 ;
+create table t2  like t1 ;
+insert into t1 values (1, 'sleepy kowalevski');
+insert into t2 values (3, 'unruffled chaplygin');
+select (select t2.c_str from t2 where t2.c_int + 1 = 4 order by t2.c_str) x from t1;
+x
+unruffled chaplygin
+select (select t2.c_str from t2 where t2.c_int = 3 order by t2.c_str) x from t1;
+x
+unruffled chaplygin
+drop table t1,t2;
diff --git a/cmd/explaintest/r/generated_columns.result b/cmd/explaintest/r/generated_columns.result
@@ -170,9 +170,10 @@ INSERT INTO person (name, address_info) VALUES ("John", CAST('{"city_no": 1}' AS
 EXPLAIN format = 'brief' SELECT name FROM person where city_no=1;
 id	estRows	task	access object	operator info
 Projection	10.00	root		test.person.name
-└─IndexLookUp	10.00	root		
-  ├─IndexRangeScan(Build)	10.00	cop[tikv]	table:person, index:city_no(city_no)	range:[1,1], keep order:false, stats:pseudo
-  └─TableRowIDScan(Probe)	10.00	cop[tikv]	table:person	keep order:false, stats:pseudo
+└─Projection	10.00	root		test.person.name, test.person.city_no
+  └─IndexLookUp	10.00	root		
+    ├─IndexRangeScan(Build)	10.00	cop[tikv]	table:person, index:city_no(city_no)	range:[1,1], keep order:false, stats:pseudo
+    └─TableRowIDScan(Probe)	10.00	cop[tikv]	table:person	keep order:false, stats:pseudo
 DROP TABLE IF EXISTS t1;
 CREATE TABLE t1 (a INT,
 b INT GENERATED ALWAYS AS (-a) VIRTUAL,

diff --git a/cmd/explaintest/r/vitess_hash.result b/cmd/explaintest/r/vitess_hash.result
@@ -19,10 +19,11 @@ explain format = 'brief' select id from t where (vitess_hash(customer_id) >> 56)
 id	estRows	task	access object	operator info
 Projection	0.50	root		test.t.id
 └─Sort	0.50	root		test.t.id
-  └─IndexLookUp	0.50	root		
-    ├─IndexRangeScan(Build)	20.00	cop[tikv]	table:t, index:t_vitess_shard(vitess_hash(`customer_id`) >> 56)	range:[224,224], [225,225], keep order:false, stats:pseudo
-    └─Selection(Probe)	0.50	cop[tikv]		ge(test.t.id, 2), le(test.t.id, 5)
-      └─TableRowIDScan	20.00	cop[tikv]	table:t	keep order:false, stats:pseudo
+  └─Projection	0.50	root		test.t.id, rightshift(vitess_hash(test.t.customer_id), 56)
+    └─IndexLookUp	0.50	root		
+      ├─IndexRangeScan(Build)	20.00	cop[tikv]	table:t, index:t_vitess_shard(vitess_hash(`customer_id`) >> 56)	range:[224,224], [225,225], keep order:false, stats:pseudo
+      └─Selection(Probe)	0.50	cop[tikv]		ge(test.t.id, 2), le(test.t.id, 5)
+        └─TableRowIDScan	20.00	cop[tikv]	table:t	keep order:false, stats:pseudo
 explain format = 'brief' select hex(vitess_hash(1123)) from t;
 id	estRows	task	access object	operator info
 Projection	10000.00	root		31B565D41BDF8CA->Column#7

diff --git a/cmd/explaintest/t/explain_generate_column_substitute.test b/cmd/explaintest/t/explain_generate_column_substitute.test
@@ -266,3 +266,13 @@ replace into t (c_int, c_str) values (9, "beautiful hermann");
 select reverse(c_str) from t use index(idx);
 
 drop table t;
+
+drop table if exists t1,t2;
+create table t1 (c_int int, c_str varchar(40) ,primary key (c_int)  , key(c_str(36))  , key((c_int + 1))) partition by hash (c_int) partitions 4 ;
+create table t2  like t1 ;
+insert into t1 values (1, 'sleepy kowalevski');
+insert into t2 values (3, 'unruffled chaplygin');
+select (select t2.c_str from t2 where t2.c_int + 1 = 4 order by t2.c_str) x from t1;
+select (select t2.c_str from t2 where t2.c_int = 3 order by t2.c_str) x from t1;
+
+drop table t1,t2;
diff --git a/planner/core/task.go b/planner/core/task.go
@@ -1018,8 +1018,6 @@ func (t *copTask) convertToRootTaskImpl(ctx sessionctx.Context) *rootTask {
 	// the number of regions involved, we simply use DistSQLScanConcurrency.
 	copIterWorkers := float64(t.plan().SCtx().GetSessionVars().DistSQLScanConcurrency())
 	t.finishIndexPlan()
-	needExtraProj := false
-	var prevSchema *expression.Schema
 	// Network cost of transferring rows of table scan to TiDB.
 	if t.tablePlan != nil {
 		t.cst += t.count() * sessVars.GetNetworkFactor(nil) * t.tblColHists.GetAvgRowSize(ctx, t.tablePlan.Schema().Columns, false, false)
@@ -1035,11 +1033,12 @@ func (t *copTask) convertToRootTaskImpl(ctx sessionctx.Context) *rootTask {
 		}
 		ts := tp.(*PhysicalTableScan)
 		prevColumnLen := len(ts.Columns)
-		prevSchema = ts.schema.Clone()
+		prevSchema := ts.schema.Clone()
 		ts.Columns = ExpandVirtualColumn(ts.Columns, ts.schema, ts.Table.Columns)
-		if len(ts.Columns) > prevColumnLen {
+		if !t.needExtraProj && len(ts.Columns) > prevColumnLen {
 			// Add an projection to make sure not to output extract columns.
-			needExtraProj = true
+			t.needExtraProj = true
+			t.originSchema = prevSchema
 		}
 	}
 	t.cst /= copIterWorkers
@@ -1106,9 +1105,9 @@ func (t *copTask) convertToRootTaskImpl(ctx sessionctx.Context) *rootTask {
 			aggPushedDown = true
 		}
 
-		if needExtraProj && !aggPushedDown {
-			proj := PhysicalProjection{Exprs: expression.Column2Exprs(prevSchema.Columns)}.Init(ts.ctx, ts.stats, ts.SelectBlockOffset(), nil)
-			proj.SetSchema(prevSchema)
+		if t.needExtraProj && !aggPushedDown {
+			proj := PhysicalProjection{Exprs: expression.Column2Exprs(t.originSchema.Columns)}.Init(ts.ctx, ts.stats, ts.SelectBlockOffset(), nil)
+			proj.SetSchema(t.originSchema)
 			proj.SetChildren(p)
 			proj.cost = t.cost()
 			newTask.p = proj