planner: allow queries on virtual columns to use index

cmd/explaintest/r/generated_columns.result

@@ -136,3 +136,42 @@ Union_13	23263.33	root
 └─TableReader_34	3323.33	root	data:Selection_33
   └─Selection_33	3323.33	cop	lt(test.sgc3.a, 7)
     └─TableScan_32	10000.00	cop	table:sgc3, partition:max, range:[-inf,+inf], keep order:false, stats:pseudo
+DROP TABLE IF EXISTS t1;
+CREATE TABLE t1(a INT, b INT AS (a+1) VIRTUAL, c INT AS (b+1) VIRTUAL, d INT AS (c+1) VIRTUAL, KEY(b), INDEX IDX(c, d));
+INSERT INTO t1 (a) VALUES (0);
+EXPLAIN SELECT b FROM t1 WHERE b=1;
+id	count	task	operator info
+IndexReader_6	10.00	root	index:IndexScan_5
+└─IndexScan_5	10.00	cop	table:t1, index:b, range:[1,1], keep order:false, stats:pseudo
+EXPLAIN SELECT b, c, d FROM t1 WHERE b=1;
+id	count	task	operator info
+Projection_11	10.00	root	cast(plus(test.t1.a, 1)), cast(plus(cast(plus(test.t1.a, 1)), 1)), cast(plus(cast(plus(cast(plus(test.t1.a, 1)), 1)), 1))
+└─IndexLookUp_12	10.00	root	
+  ├─IndexScan_9	10.00	cop	table:t1, index:b, range:[1,1], keep order:false, stats:pseudo
+  └─TableScan_10	10.00	cop	table:t1, keep order:false, stats:pseudo
+EXPLAIN SELECT * FROM t1 WHERE b=1;
+id	count	task	operator info
+Projection_11	10.00	root	test.t1.a, cast(plus(test.t1.a, 1)), cast(plus(cast(plus(test.t1.a, 1)), 1)), cast(plus(cast(plus(cast(plus(test.t1.a, 1)), 1)), 1))
+└─IndexLookUp_12	10.00	root	
+  ├─IndexScan_9	10.00	cop	table:t1, index:b, range:[1,1], keep order:false, stats:pseudo
+  └─TableScan_10	10.00	cop	table:t1, keep order:false, stats:pseudo
+EXPLAIN SELECT c FROM t1 WHERE c=2 AND d=3;
+id	count	task	operator info
+Projection_4	0.10	root	test.t1.c
+└─IndexReader_6	0.10	root	index:IndexScan_5
+  └─IndexScan_5	0.10	cop	table:t1, index:c, d, range:[2 3,2 3], keep order:false, stats:pseudo
+DROP TABLE IF EXISTS person;
+CREATE TABLE person (
+id INT NOT NULL AUTO_INCREMENT PRIMARY KEY,
+name VARCHAR(255) NOT NULL,
+address_info JSON,
+city_no INT AS (JSON_EXTRACT(address_info, '$.city_no')) VIRTUAL,
+KEY(city_no));
+INSERT INTO person (name, address_info) VALUES ("John", CAST('{"city_no": 1}' AS JSON));
+EXPLAIN SELECT name FROM person where city_no=1;
+id	count	task	operator info
+Projection_4	10.00	root	test.person.name
+└─Projection_11	10.00	root	test.person.name, cast(json_extract(test.person.address_info, "$.city_no"))
+  └─IndexLookUp_12	10.00	root	
+    ├─IndexScan_9	10.00	cop	table:person, index:city_no, range:[1,1], keep order:false, stats:pseudo
+    └─TableScan_10	10.00	cop	table:person, keep order:false, stats:pseudo

cmd/explaintest/t/generated_columns.test

@@ -90,3 +90,24 @@ PARTITION max VALUES LESS THAN MAXVALUE);
 EXPLAIN SELECT * FROM sgc3 WHERE a <= 1;
 EXPLAIN SELECT * FROM sgc3 WHERE a < 7;
 
+-- Virtual generated columns as indices
+
+DROP TABLE IF EXISTS t1;
+CREATE TABLE t1(a INT, b INT AS (a+1) VIRTUAL, c INT AS (b+1) VIRTUAL, d INT AS (c+1) VIRTUAL, KEY(b), INDEX IDX(c, d));
+INSERT INTO t1 (a) VALUES (0);
+
+EXPLAIN SELECT b FROM t1 WHERE b=1;
+EXPLAIN SELECT b, c, d FROM t1 WHERE b=1;
+EXPLAIN SELECT * FROM t1 WHERE b=1;
+EXPLAIN SELECT c FROM t1 WHERE c=2 AND d=3;
+
+DROP TABLE IF EXISTS person;
+CREATE TABLE person (
+id INT NOT NULL AUTO_INCREMENT PRIMARY KEY,
+name VARCHAR(255) NOT NULL,
+address_info JSON,
+city_no INT AS (JSON_EXTRACT(address_info, '$.city_no')) VIRTUAL,
+KEY(city_no));
+
+INSERT INTO person (name, address_info) VALUES ("John", CAST('{"city_no": 1}' AS JSON));
+EXPLAIN SELECT name FROM person where city_no=1;

executor/builder.go

@@ -1946,7 +1946,12 @@ type dataReaderBuilder struct {
 
 func (builder *dataReaderBuilder) buildExecutorForIndexJoin(ctx context.Context, lookUpContents []*indexJoinLookUpContent,
 	IndexRanges []*ranger.Range, keyOff2IdxOff []int, cwc *plannercore.ColWithCmpFuncManager) (Executor, error) {
-	switch v := builder.Plan.(type) {
+	return builder.buildReaderForIndexJoin(ctx, builder.Plan, lookUpContents, IndexRanges, keyOff2IdxOff, cwc)
+}
+
+func (builder *dataReaderBuilder) buildReaderForIndexJoin(ctx context.Context, p plannercore.Plan,
+	lookUpContents []*indexJoinLookUpContent, IndexRanges []*ranger.Range, keyOff2IdxOff []int, cwc *plannercore.ColWithCmpFuncManager) (Executor, error) {
+	switch v := p.(type) {
 	case *plannercore.PhysicalTableReader:
 		return builder.buildTableReaderForIndexJoin(ctx, v, lookUpContents)
 	case *plannercore.PhysicalIndexReader:
@@ -1955,10 +1960,31 @@ func (builder *dataReaderBuilder) buildExecutorForIndexJoin(ctx context.Context,
 		return builder.buildIndexLookUpReaderForIndexJoin(ctx, v, lookUpContents, IndexRanges, keyOff2IdxOff, cwc)
 	case *plannercore.PhysicalUnionScan:
 		return builder.buildUnionScanForIndexJoin(ctx, v, lookUpContents, IndexRanges, keyOff2IdxOff, cwc)
+	case *plannercore.PhysicalProjection:
+		return builder.buildProjectionForIndexJoin(ctx, v, lookUpContents, IndexRanges, keyOff2IdxOff, cwc)
 	}
+
 	return nil, errors.New("Wrong plan type for dataReaderBuilder")
 }
 
+func (builder *dataReaderBuilder) buildProjectionForIndexJoin(ctx context.Context, v *plannercore.PhysicalProjection,
+	lookUpContents []*indexJoinLookUpContent, IndexRanges []*ranger.Range, keyOff2IdxOff []int, cwc *plannercore.ColWithCmpFuncManager) (Executor, error) {
+	childExec, err := builder.buildReaderForIndexJoin(ctx, v.Children()[0], lookUpContents, IndexRanges, keyOff2IdxOff, cwc)
+	if err != nil {
+		return nil, err
+	}
+	e := &ProjectionExec{
+		baseExecutor:     newBaseExecutor(builder.ctx, v.Schema(), v.ExplainID(), childExec),
+		numWorkers:       0, // always run in un-parallel mode to avoid too many number of goroutines.
+		evaluatorSuit:    expression.NewEvaluatorSuite(v.Exprs, v.AvoidColumnEvaluator),
+		calculateNoDelay: v.CalculateNoDelay,
+	}
+	if err := e.open(ctx); err != nil {
+		return nil, nil
+	}
+	return e, nil
+}
+
 func (builder *dataReaderBuilder) buildUnionScanForIndexJoin(ctx context.Context, v *plannercore.PhysicalUnionScan,
 	values []*indexJoinLookUpContent, indexRanges []*ranger.Range, keyOff2IdxOff []int, cwc *plannercore.ColWithCmpFuncManager) (Executor, error) {
 	childBuilder := &dataReaderBuilder{Plan: v.Children()[0], executorBuilder: builder.executorBuilder}

executor/projection.go

@@ -78,6 +78,10 @@ func (e *ProjectionExec) Open(ctx context.Context) error {
 		return err
 	}
 
+	return e.open(ctx)
+}
+
+func (e *ProjectionExec) open(ctx context.Context) error {
 	e.prepared = false
 	e.parentReqRows = int64(e.maxChunkSize)
 

planner/core/cbo_test.go

@@ -959,32 +959,79 @@ func (s *testAnalyzeSuite) TestIssue9805(c *C) {
 
 	// Expected output is like:
 	//
-	// +--------------------------------+----------+------+----------------------------------------------------------------------------------+----------------------------------+
-	// | id                             | count    | task | operator info                                                                    | execution info                   |
-	// +--------------------------------+----------+------+----------------------------------------------------------------------------------+----------------------------------+
-	// | Projection_9                   | 10.00    | root | test.t1.id, test.t2.a                                                            | time:203.355µs, loops:1, rows:0  |
-	// | └─IndexJoin_13                 | 10.00    | root | inner join, inner:IndexLookUp_12, outer key:test.t1.a, inner key:test.t2.d       | time:199.633µs, loops:1, rows:0  |
-	// |   ├─Projection_16              | 8.00     | root | test.t1.id, test.t1.a, test.t1.b, cast(mod(test.t1.a, 30))                       | time:164.587µs, loops:1, rows:0  |
-	// |   │ └─Selection_17             | 8.00     | root | eq(cast(mod(test.t1.a, 30)), 4)                                                  | time:157.768µs, loops:1, rows:0  |
-	// |   │   └─TableReader_20         | 10.00    | root | data:Selection_19                                                                | time:154.61µs, loops:1, rows:0   |
-	// |   │     └─Selection_19         | 10.00    | cop  | eq(test.t1.b, "t2")                                                              | time:28.824µs, loops:1, rows:0   |
-	// |   │       └─TableScan_18       | 10000.00 | cop  | table:t1, range:[-inf,+inf], keep order:false, stats:pseudo                      | time:27.654µs, loops:1, rows:0   |
-	// |   └─IndexLookUp_12             | 10.00    | root |                                                                                  | time:0ns, loops:0, rows:0        |
-	// |     ├─IndexScan_10             | 10.00    | cop  | table:t2, index:d, range: decided by [test.t1.a], keep order:false, stats:pseudo | time:0ns, loops:0, rows:0        |
-	// |     └─TableScan_11             | 10.00    | cop  | table:t2, keep order:false, stats:pseudo                                         | time:0ns, loops:0, rows:0        |
-	// +--------------------------------+----------+------+----------------------------------------------------------------------------------+----------------------------------+
-	// 10 rows in set (0.00 sec)
+	//+----------------------------+-------+------+-------------------------------------------------------------------------------------------------+----------------------------------+
+	//| id                         | count | task | operator info                                                                                   | execution info                   |
+	//+----------------------------+-------+------+-------------------------------------------------------------------------------------------------+----------------------------------+
+	//| Projection_7               | 0.12  | root | test.t1.id, test.t2.a                                                                           | time:3.844593ms, loops:1, rows:0 |
+	//| └─IndexJoin_11             | 0.12  | root | inner join, inner:IndexLookUp_10, outer key:test.t1.a, inner key:test.t2.d                      | time:3.830714ms, loops:1, rows:0 |
+	//|   ├─Projection_21          | 0.10  | root | test.t1.id, test.t1.a, test.t1.b, cast(mod(test.t1.a, 30))                                      | time:3.735174ms, loops:1, rows:0 |
+	//|   │ └─IndexLookUp_22       | 0.10  | root |                                                                                                 | time:3.569946ms, loops:1, rows:0 |
+	//|   │   ├─IndexScan_19       | 0.10  | cop  | table:t1, index:d, b, c, range:[4 "t2",4 "t2"], keep order:false, stats:pseudo                  | time:50.542µs, loops:1, rows:0   |
+	//|   │   └─TableScan_20       | 0.10  | cop  | table:t1, keep order:false, stats:pseudo                                                        | time:0s, loops:0, rows:0         |
+	//|   └─IndexLookUp_10         | 10.00 | root |                                                                                                 | time:0ns, loops:0, rows:0        |
+	//|     ├─IndexScan_8          | 10.00 | cop  | table:t2, index:d, range: decided by [eq(test.t2.d, test.t1.a)], keep order:false, stats:pseudo | time:0ns, loops:0, rows:0        |
+	//|     └─TableScan_9          | 10.00 | cop  | table:t2, keep order:false, stats:pseudo                                                        | time:0ns, loops:0, rows:0        |
+	//+----------------------------+-------+------+-------------------------------------------------------------------------------------------------+----------------------------------+
+	//9 rows in set (0.01 sec)
 	//
-	c.Assert(rs.Rows(), HasLen, 10)
-	hasIndexLookUp12 := false
+	c.Assert(rs.Rows(), HasLen, 9)
+	hasIndexLookUp10 := false
+	hasIndexLookUp22 := false
 	for _, row := range rs.Rows() {
 		c.Assert(row, HasLen, 6)
-		if strings.HasSuffix(row[0].(string), "IndexLookUp_12") {
-			hasIndexLookUp12 = true
+		if strings.Contains(row[0].(string), "IndexLookUp_10") {
+			hasIndexLookUp10 = true
 			c.Assert(row[4], Equals, "time:0ns, loops:0, rows:0")
 		}
+		if strings.Contains(row[0].(string), "IndexLookUp_22") {
+			hasIndexLookUp22 = true
+		}
 	}
-	c.Assert(hasIndexLookUp12, IsTrue)
+	c.Assert(hasIndexLookUp10, IsTrue)
+	c.Assert(hasIndexLookUp22, IsTrue)
+}
+
+func (s *testAnalyzeSuite) TestVirtualGeneratedColumn(c *C) {
+	defer testleak.AfterTest(c)()
+	store, dom, err := newStoreWithBootstrap()
+	c.Assert(err, IsNil)
+	tk := testkit.NewTestKit(c, store)
+	defer func() {
+		dom.Close()
+		store.Close()
+	}()
+	tk.MustExec("use test")
+	tk.MustExec("drop table if exists t1")
+	tk.MustExec("create table t1(a int, b int as (a+1) virtual, c int as (b+1) virtual, d int as (c+1) virtual, key(b), index idx(c, d))")
+	tk.MustExec("insert into t1 (a) values (0)")
+
+	tk.MustIndexRead("select b from t1 where b=1").Check(testkit.Rows("1"))
+	tk.MustIndexLookup("select b, c, d from t1 where b=1").Check(testkit.Rows("1 2 3"))
+	tk.MustIndexLookup("select * from t1 where b=1").Check(testkit.Rows("0 1 2 3"))
+	tk.MustIndexRead("select c from t1 where c=2 and d=3").Check(testkit.Rows("2"))
+
+	tk.MustExec("insert into t1 (a) values (1)")
+	tk.MustQuery("select /*+ TIDB_INLJ(o, i) */ i.b, o.a from t1 o, t1 i where i.b = o.a").Check(testkit.Rows("1 1"))
+	tk.MustQuery("explain select /*+ TIDB_INLJ(o, i) */ * from t1 o, t1 i where i.b = o.a").Check(testkit.Rows(
+		"IndexJoin_12 12487.50 root inner join, inner:Projection_10, outer key:test.o.a, inner key:test.i.b",
+		"├─Projection_15 9990.00 root test.o.a, cast(plus(test.o.a, 1)), cast(plus(cast(plus(test.o.a, 1)), 1)), cast(plus(cast(plus(cast(plus(test.o.a, 1)), 1)), 1))",
+		"│ └─TableReader_16 9990.00 root data:Selection_14",
+		"│   └─Selection_14 9990.00 cop not(isnull(test.o.a))",
+		"│     └─TableScan_13 10000.00 cop table:o, range:[-inf,+inf], keep order:false, stats:pseudo",
+		"└─Projection_10 9.99 root test.i.a, cast(plus(test.i.a, 1)), cast(plus(cast(plus(test.i.a, 1)), 1)), cast(plus(cast(plus(cast(plus(test.i.a, 1)), 1)), 1))",
+		"  └─IndexLookUp_11 9.99 root ",
+		"    ├─Selection_9 9.99 cop not(isnull(test.i.b))",
+		"    │ └─IndexScan_7 10.00 cop table:i, index:b, range: decided by [eq(test.i.b, test.o.a)], keep order:false, stats:pseudo",
+		"    └─TableScan_8 9.99 cop table:t1, keep order:false, stats:pseudo"))
+
+	tk.MustExec(`CREATE TABLE person (
+    id INT NOT NULL AUTO_INCREMENT PRIMARY KEY,
+    name VARCHAR(255) NOT NULL,
+    address_info JSON,
+    city_no INT AS (JSON_EXTRACT(address_info, '$.city_no')) VIRTUAL,
+    KEY(city_no))`)
+	tk.MustExec(`INSERT INTO person (name, address_info) VALUES ("John", CAST('{"city_no": 1}' AS JSON))`)
+	tk.MustIndexLookup(`SELECT name FROM person where city_no=1`).Check(testkit.Rows("John"))
 }
 
 func (s *testAnalyzeSuite) TestLimitCrossEstimation(c *C) {

planner/core/exhaust_physical_plans.go

@@ -435,7 +435,11 @@ func (p *LogicalJoin) getIndexJoinByOuterIdx(prop *property.PhysicalProperty, ou
 			keyOff2IdxOff[i] = 0
 		}
 		if pkMatched {
-			innerPlan := p.constructInnerTableScan(ds, pkCol, outerJoinKeys, us)
+			innerPlan, err := p.constructInnerTableScan(ds, pkCol, outerJoinKeys, us)
+			if err != nil {
+				logutil.BgLogger().Error("construct inner table scan error", zap.Error(err))
+				return nil
+			}
 			// Since the primary key means one value corresponding to exact one row, this will always be a no worse one
 			// comparing to other index.
 			return p.constructIndexJoin(prop, outerIdx, innerPlan, nil, keyOff2IdxOff, nil, nil)
@@ -464,7 +468,11 @@ func (p *LogicalJoin) getIndexJoinByOuterIdx(prop *property.PhysicalProperty, ou
 		}
 		idxCols, lens := expression.IndexInfo2Cols(ds.schema.Columns, helper.chosenIndexInfo)
 		rangeInfo := helper.buildRangeDecidedByInformation(idxCols, outerJoinKeys)
-		innerPlan := p.constructInnerIndexScan(ds, helper.chosenIndexInfo, helper.chosenRemained, outerJoinKeys, us, rangeInfo)
+		innerPlan, err := p.constructInnerIndexScan(ds, helper.chosenIndexInfo, helper.chosenRemained, outerJoinKeys, us, rangeInfo)
+		if err != nil {
+			logutil.BgLogger().Error("construct inner index scan error", zap.Error(err))
+			return nil
+		}
 		return p.constructIndexJoin(prop, outerIdx, innerPlan, helper.chosenRanges, keyOff2IdxOff, lens, helper.lastColManager)
 	}
 	return nil
@@ -514,7 +522,7 @@ func (ijHelper *indexJoinBuildHelper) buildRangeDecidedByInformation(idxCols []*
 }
 
 // constructInnerTableScan is specially used to construct the inner plan for PhysicalIndexJoin.
-func (p *LogicalJoin) constructInnerTableScan(ds *DataSource, pk *expression.Column, outerJoinKeys []*expression.Column, us *LogicalUnionScan) PhysicalPlan {
+func (p *LogicalJoin) constructInnerTableScan(ds *DataSource, pk *expression.Column, outerJoinKeys []*expression.Column, us *LogicalUnionScan) (PhysicalPlan, error) {
 	ranges := ranger.FullIntRange(mysql.HasUnsignedFlag(pk.RetType.Flag))
 	ts := PhysicalTableScan{
 		Table:           ds.tableInfo,
@@ -533,15 +541,20 @@ func (p *LogicalJoin) constructInnerTableScan(ds *DataSource, pk *expression.Col
 		ts.stats.StatsVersion = statistics.PseudoVersion
 	}
 
-	copTask := &copTask{
+	cop := &copTask{
 		tablePlan:         ts,
 		indexPlanFinished: true,
 	}
 	selStats := ts.stats.Scale(selectionFactor)
-	ts.addPushedDownSelection(copTask, selStats)
-	t := finishCopTask(ds.ctx, copTask)
+	t, err := ds.pushDownSelAndResolveVirtualCols(cop, nil, selStats)
+	if err != nil {
+		return nil, err
+	}
+	if cop, ok := t.(*copTask); ok {
+		t = finishCopTask(ds.ctx, cop)
+	}
 	reader := t.plan()
-	return p.constructInnerUnionScan(us, reader)
+	return p.constructInnerUnionScan(us, reader), nil
 }
 
 func (p *LogicalJoin) constructInnerUnionScan(us *LogicalUnionScan, reader PhysicalPlan) PhysicalPlan {
@@ -557,7 +570,7 @@ func (p *LogicalJoin) constructInnerUnionScan(us *LogicalUnionScan, reader Physi
 
 // constructInnerIndexScan is specially used to construct the inner plan for PhysicalIndexJoin.
 func (p *LogicalJoin) constructInnerIndexScan(ds *DataSource, idx *model.IndexInfo, filterConds []expression.Expression,
-	outerJoinKeys []*expression.Column, us *LogicalUnionScan, rangeInfo string) PhysicalPlan {
+	outerJoinKeys []*expression.Column, us *LogicalUnionScan, rangeInfo string) (PhysicalPlan, error) {
 	is := PhysicalIndexScan{
 		Table:            ds.tableInfo,
 		TableAsName:      ds.TableAsName,
@@ -611,10 +624,15 @@ func (p *LogicalJoin) constructInnerIndexScan(ds *DataSource, idx *model.IndexIn
 	}
 	selectivity := ds.stats.RowCount / ds.tableStats.RowCount
 	finalStats := ds.stats.ScaleByExpectCnt(selectivity * rowCount)
-	is.addPushedDownSelection(cop, ds, path, finalStats)
-	t := finishCopTask(ds.ctx, cop)
+	t, err := ds.pushDownSelAndResolveVirtualCols(cop, path, finalStats)
+	if err != nil {
+		return nil, err
+	}
+	if cop, ok := t.(*copTask); ok {
+		t = finishCopTask(ds.ctx, cop)
+	}
 	reader := t.plan()
-	return p.constructInnerUnionScan(us, reader)
+	return p.constructInnerUnionScan(us, reader), nil
 }
 
 var symmetricOp = map[string]string{