planner: push necessary predicates without virtual column down through UnionScan (#54985) (#54991)

close #54870

Author: ti-chi-bot

pkg/planner/core/integration_test.go

@@ -2322,6 +2322,68 @@ func TestIssue48257(t *testing.T) {
 	))
 }
 
+func TestIssue54870(t *testing.T) {
+	store := testkit.CreateMockStore(t)
+	tk := testkit.NewTestKit(t, store)
+
+	tk.MustExec("use test")
+	tk.MustExec(`create table t (id int,
+deleted_at datetime(3) NOT NULL DEFAULT '1970-01-01 01:00:01.000',
+is_deleted tinyint(1) GENERATED ALWAYS AS ((deleted_at > _utf8mb4'1970-01-01 01:00:01.000')) VIRTUAL NOT NULL,
+key k(id, is_deleted))`)
+	tk.MustExec(`begin`)
+	tk.MustExec(`insert into t (id, deleted_at) values (1, now())`)
+	tk.MustHavePlan(`select 1 from t where id=1 and is_deleted=true`, "IndexRangeScan")
+}
+
+func TestIssue53951(t *testing.T) {
+	store := testkit.CreateMockStore(t)
+	tk := testkit.NewTestKit(t, store)
+	tk.MustExec("use test")
+	tk.MustExec(`CREATE TABLE gholla_dummy1 (
+  id varchar(10) NOT NULL,
+  mark int,
+  deleted_at datetime(3) NOT NULL DEFAULT '1970-01-01 01:00:01.000',
+  account_id varchar(10) NOT NULL,
+  metastore_id varchar(10) NOT NULL,
+  is_deleted tinyint(1) GENERATED ALWAYS AS ((deleted_at > _utf8mb4'1970-01-01 01:00:01.000')) VIRTUAL NOT NULL,
+  PRIMARY KEY (account_id,metastore_id,id),
+  KEY isDeleted_accountId_metastoreId (is_deleted,account_id,metastore_id)
+) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_bin;`)
+	tk.MustExec(`CREATE TABLE gholla_dummy2 (
+  id varchar(10) NOT NULL,
+  mark int,
+  deleted_at datetime(3) NOT NULL DEFAULT '1970-01-01 01:00:01.000',
+  account_id varchar(10) NOT NULL,
+  metastore_id varchar(10) NOT NULL,
+  is_deleted tinyint(1) GENERATED ALWAYS AS ((deleted_at > _utf8mb4'1970-01-01 01:00:01.000')) VIRTUAL NOT NULL,
+  PRIMARY KEY (account_id,metastore_id,id),
+  KEY isDeleted_accountId_metastoreId (is_deleted,account_id,metastore_id)
+) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_bin; `)
+	tk.MustExec(`INSERT INTO gholla_dummy1 (id,mark,deleted_at,account_id,metastore_id) VALUES ('ABC', 1, '1970-01-01 01:00:01.000', 'ABC', 'ABC');`)
+	tk.MustExec(`INSERT INTO gholla_dummy2 (id,mark,deleted_at,account_id,metastore_id) VALUES ('ABC', 1, '1970-01-01 01:00:01.000', 'ABC', 'ABC');`)
+	tk.MustExec(`start transaction;`)
+	tk.MustExec(`update gholla_dummy2 set deleted_at = NOW(), mark=2 where account_id = 'ABC' and metastore_id = 'ABC' and id = 'ABC';`)
+	tk.MustQuery(`select
+  /*+ INL_JOIN(g1, g2) */
+  g1.account_id,
+  g2.mark
+from
+  gholla_dummy1 g1 FORCE INDEX(isDeleted_accountId_metastoreId)
+STRAIGHT_JOIN
+  gholla_dummy2 g2 FORCE INDEX (PRIMARY)
+ON
+  g1.account_id = g2.account_id AND
+  g1.metastore_id = g2.metastore_id AND
+  g1.id = g2.id
+WHERE
+  g1.account_id = 'ABC' AND
+  g1.metastore_id = 'ABC' AND
+  g1.is_deleted = FALSE AND
+  g2.is_deleted = FALSE;`).Check(testkit.Rows()) // empty result, no error
+	tk.MustExec(`rollback`)
+}
+
 func TestIssue52472(t *testing.T) {
 	store := testkit.CreateMockStore(t)
 	tk := testkit.NewTestKit(t, store)

pkg/planner/core/rule_predicate_push_down.go

@@ -124,10 +124,22 @@ func (p *LogicalSelection) PredicatePushDown(predicates []expression.Expression,
 
 // PredicatePushDown implements LogicalPlan PredicatePushDown interface.
 func (p *LogicalUnionScan) PredicatePushDown(predicates []expression.Expression, opt *util.LogicalOptimizeOp) ([]expression.Expression, LogicalPlan) {
+	var predicatesWithVCol, predicatesWithoutVCol []expression.Expression
+	// predicates with virtual columns can't be pushed down to TiKV/TiFlash so they'll be put into a Projection
+	// below the UnionScan, but the current UnionScan doesn't support placing Projection below it, see #53951.
+	for _, expr := range predicates {
+		if expression.ContainVirtualColumn([]expression.Expression{expr}) {
+			predicatesWithVCol = append(predicatesWithVCol, expr)
+		} else {
+			predicatesWithoutVCol = append(predicatesWithoutVCol, expr)
+		}
+	}
+	predicates = predicatesWithoutVCol
 	retainedPredicates, _ := p.children[0].PredicatePushDown(predicates, opt)
 	p.conditions = make([]expression.Expression, 0, len(predicates))
 	p.conditions = append(p.conditions, predicates...)
 	// The conditions in UnionScan is only used for added rows, so parent Selection should not be removed.
+	retainedPredicates = append(retainedPredicates, predicatesWithVCol...)
 	return retainedPredicates, p
 }
 

tests/integrationtest/r/explain_generate_column_substitute.result

@@ -573,15 +573,15 @@ begin;
 delete from t2 where c_decimal > c_double/2 order by c_int, c_str, c_double, c_decimal limit 1;
 desc format='brief' select t2.c_enum from t2,t1 where t1.c_int - 1 = t2.c_int - 1 order  by t2.c_enum;
 id	estRows	task	access object	operator info
-Sort	12487.50	root		explain_generate_column_substitute.t2.c_enum
-└─HashJoin	12487.50	root		inner join, equal:[eq(minus(explain_generate_column_substitute.t1.c_int, 1), minus(explain_generate_column_substitute.t2.c_int, 1))]
-  ├─IndexReader(Build)	9990.00	root		index:IndexFullScan
-  │ └─IndexFullScan	9990.00	cop[tikv]	table:t1, index:expression_index_2(`c_int` - 1)	keep order:false, stats:pseudo
-  └─Projection(Probe)	10000.00	root		explain_generate_column_substitute.t2.c_enum, minus(explain_generate_column_substitute.t2.c_int, 1), explain_generate_column_substitute.t2._tidb_rowid
-    └─UnionScan	8000.00	root		not(isnull(minus(explain_generate_column_substitute.t2.c_int, 1)))
-      └─Selection	8000.00	root		not(isnull(minus(explain_generate_column_substitute.t2.c_int, 1)))
-        └─TableReader	10000.00	root		data:TableFullScan
-          └─TableFullScan	10000.00	cop[tikv]	table:t2	keep order:false, stats:pseudo
+Sort	10000.00	root		explain_generate_column_substitute.t2.c_enum
+└─HashJoin	10000.00	root		inner join, equal:[eq(minus(explain_generate_column_substitute.t1.c_int, 1), minus(explain_generate_column_substitute.t2.c_int, 1))]
+  ├─Selection(Build)	8000.00	root		not(isnull(minus(explain_generate_column_substitute.t2.c_int, 1)))
+  │ └─Projection	10000.00	root		explain_generate_column_substitute.t2.c_enum, minus(explain_generate_column_substitute.t2.c_int, 1), explain_generate_column_substitute.t2._tidb_rowid
+  │   └─UnionScan	10000.00	root		
+  │     └─TableReader	10000.00	root		data:TableFullScan
+  │       └─TableFullScan	10000.00	cop[tikv]	table:t2	keep order:false, stats:pseudo
+  └─IndexReader(Probe)	9990.00	root		index:IndexFullScan
+    └─IndexFullScan	9990.00	cop[tikv]	table:t1, index:expression_index_2(`c_int` - 1)	keep order:false, stats:pseudo
 select t2.c_enum from t2,t1 where t1.c_int - 1 = t2.c_int - 1 order  by t2.c_enum;
 c_enum
 orange