From e7c25024337b1ba9b380c78549499e81f5018049 Mon Sep 17 00:00:00 2001
From: Ti Chi Robot <ti-community-prow-bot@tidb.io>
Date: Tue, 11 Feb 2025 14:07:49 +0800
Subject: [PATCH] planner: fix flaky test
 TestPhysicalTableScanExtractCorrelatedCols (#51355) (#58952)

close pingcap/tidb#51289
---
 planner/core/physical_plan_test.go            | 39 ++++++++++++++-----
 .../tiflash_selection_late_materialization.go | 12 ++++--
 2 files changed, 39 insertions(+), 12 deletions(-)

diff --git a/planner/core/physical_plan_test.go b/planner/core/physical_plan_test.go
index d0eedee151604..792b3392173a0 100644
--- a/planner/core/physical_plan_test.go
+++ b/planner/core/physical_plan_test.go
@@ -21,6 +21,7 @@ import (
 	"testing"
 
 	"github.com/pingcap/tidb/domain"
+	"github.com/pingcap/tidb/expression"
 	"github.com/pingcap/tidb/infoschema"
 	"github.com/pingcap/tidb/kv"
 	"github.com/pingcap/tidb/parser"
@@ -546,31 +547,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(), kv.TiFlash)
 	require.NoError(t, err)
diff --git a/planner/core/tiflash_selection_late_materialization.go b/planner/core/tiflash_selection_late_materialization.go
index 8f4d4a0c806ac..d3ec2a2b0e056 100644
--- a/planner/core/tiflash_selection_late_materialization.go
+++ b/planner/core/tiflash_selection_late_materialization.go
@@ -250,10 +250,16 @@ func predicatePushDownToTableScanImpl(sctx sessionctx.Context, physicalSelection
 		return
 	}
 	logutil.BgLogger().Debug("planner: push down conditions to table scan", zap.String("table", physicalTableScan.Table.Name.L), zap.String("conditions", string(expression.SortedExplainExpressionList(selectedConds))))
+	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(physicalSelection, selectedConds)
+	removeSpecificExprsFromSelection(sel, selectedConds)
 	// add the pushed down conditions to table scan
-	physicalTableScan.LateMaterializationFilterCondition = selectedConds
+	ts.LateMaterializationFilterCondition = selectedConds
 	// Update the row count of table scan after pushing down the conditions.
-	physicalTableScan.stats.RowCount *= selectedSelectivity
+	ts.stats.RowCount *= selectedSelectivity
 }