planner: fix flaky test TestPhysicalTableScanExtractCorrelatedCols (#51355)

close pingcap/tidb#51289

pkg/planner/core/physical_plan_test.go

@@ -21,6 +21,7 @@ import (
 	"testing"
 
 	"github.com/pingcap/tidb/pkg/domain"
+	"github.com/pingcap/tidb/pkg/expression"
 	"github.com/pingcap/tidb/pkg/infoschema"
 	"github.com/pingcap/tidb/pkg/kv"
 	"github.com/pingcap/tidb/pkg/parser"
@@ -459,31 +460,51 @@ func TestPhysicalTableScanExtractCorrelatedCols(t *testing.T) {
 	p, ok := info.Plan.(core.Plan)
 	require.True(t, ok)
 
-	var findTableScan func(p core.Plan) *core.PhysicalTableScan
-	findTableScan = func(p core.Plan) *core.PhysicalTableScan {
+	var findSelection func(p core.Plan) *core.PhysicalSelection
+	findSelection = func(p core.Plan) *core.PhysicalSelection {
 		if p == nil {
 			return nil
 		}
 		switch v := p.(type) {
-		case *core.PhysicalTableScan:
-			if v.Table.Name.L == "t1" {
-				return v
+		case *core.PhysicalSelection:
+			if len(v.Children()) == 1 {
+				if ts, ok := v.Children()[0].(*core.PhysicalTableScan); ok && ts.Table.Name.L == "t1" {
+					return v
+				}
 			}
 			return nil
 		case *core.PhysicalTableReader:
-			return findTableScan(v.TablePlans[0])
+			for _, child := range v.TablePlans {
+				if sel := findSelection(child); sel != nil {
+					return sel
+				}
+			}
+			return nil
 		default:
 			physicayPlan := p.(core.PhysicalPlan)
 			for _, child := range physicayPlan.Children() {
-				if ts := findTableScan(child); ts != nil {
-					return ts
+				if sel := findSelection(child); sel != nil {
+					return sel
 				}
 			}
 			return nil
 		}
 	}
-	ts := findTableScan(p)
+	sel := findSelection(p)
+	require.NotNil(t, sel)
+	ts := sel.Children()[0].(*core.PhysicalTableScan)
 	require.NotNil(t, ts)
+	// manually push down the condition `client_no = c.company_no`
+	var selected expression.Expression
+	for _, cond := range sel.Conditions {
+		if sf, ok := cond.(*expression.ScalarFunction); ok && sf.Function.PbCode() == tipb.ScalarFuncSig_EQString {
+			selected = cond
+			break
+		}
+	}
+	if selected != nil {
+		core.PushedDown(sel, ts, []expression.Expression{selected}, 0.1)
+	}
 
 	pb, err := ts.ToPB(tk.Session().GetPlanCtx(), kv.TiFlash)
 	require.NoError(t, err)

pkg/planner/core/tiflash_selection_late_materialization.go

@@ -251,10 +251,16 @@ func predicatePushDownToTableScanImpl(sctx PlanContext, physicalSelection *Physi
 		return
 	}
 	logutil.BgLogger().Debug("planner: push down conditions to table scan", zap.String("table", physicalTableScan.Table.Name.L), zap.String("conditions", string(expression.SortedExplainExpressionList(sctx, selectedConds))))
-	// remove the pushed down conditions from selection
-	removeSpecificExprsFromSelection(physicalSelection, selectedConds)
-	// add the pushed down conditions to table scan
-	physicalTableScan.LateMaterializationFilterCondition = selectedConds
-	// Update the row count of table scan after pushing down the conditions.
-	physicalTableScan.StatsInfo().RowCount *= selectedSelectivity
+	PushedDown(physicalSelection, physicalTableScan, selectedConds, selectedSelectivity)
+}
+
+// PushedDown is used to push down the selected conditions from PhysicalSelection to PhysicalTableScan.
+// Used in unit test, so it is exported.
+func PushedDown(sel *PhysicalSelection, ts *PhysicalTableScan, selectedConds []expression.Expression, selectedSelectivity float64) {
+	// remove the pushed down conditions from selection
+	removeSpecificExprsFromSelection(sel, selectedConds)
+	// add the pushed down conditions to table scan
+	ts.LateMaterializationFilterCondition = selectedConds
+	// Update the row count of table scan after pushing down the conditions.
+	ts.StatsInfo().RowCount *= selectedSelectivity
 }