executor: fill correlated column value in late materialization filter conditions (#49244)

pkg/executor/builder.go

@@ -2953,13 +2953,18 @@ func markChildrenUsedCols(outputCols []*expression.Column, childSchemas ...*expr
 
 func (*executorBuilder) corColInDistPlan(plans []plannercore.PhysicalPlan) bool {
 	for _, p := range plans {
-		x, ok := p.(*plannercore.PhysicalSelection)
-		if !ok {
-			continue
-		}
-		for _, cond := range x.Conditions {
-			if len(expression.ExtractCorColumns(cond)) > 0 {
-				return true
+		switch x := p.(type) {
+		case *plannercore.PhysicalSelection:
+			for _, cond := range x.Conditions {
+				if len(expression.ExtractCorColumns(cond)) > 0 {
+					return true
+				}
+			}
+		case *plannercore.PhysicalTableScan:
+			for _, cond := range x.LateMaterializationFilterCondition {
+				if len(expression.ExtractCorColumns(cond)) > 0 {
+					return true
+				}
 			}
 		}
 	}

pkg/executor/table_reader.go

@@ -112,7 +112,8 @@ type TableReaderExecutor struct {
 	byItems   []*util.ByItems
 	paging    bool
 	storeType kv.StoreType
-	// corColInFilter tells whether there's correlated column in filter.
+	// corColInFilter tells whether there's correlated column in filter (both conditions in PhysicalSelection and LateMaterializationFilterCondition in PhysicalTableScan)
+	// If true, we will need to revise the dagPB (fill correlated column value in filter) each time call Open().
 	corColInFilter bool
 	// corColInAccess tells whether there's correlated column in access conditions.
 	corColInAccess bool
@@ -156,6 +157,7 @@ func (e *TableReaderExecutor) Open(ctx context.Context) error {
 
 	var err error
 	if e.corColInFilter {
+		// If there's correlated column in filter, need to rewrite dagPB
 		if e.storeType == kv.TiFlash {
 			execs, err := builder.ConstructTreeBasedDistExec(e.Ctx(), e.tablePlan)
 			if err != nil {

pkg/planner/core/explain.go

@@ -215,11 +215,11 @@ func (p *PhysicalTableScan) OperatorInfo(normalized bool) string {
 	}
 	if p.SCtx().GetSessionVars().EnableLateMaterialization && len(p.filterCondition) > 0 && p.StoreType == kv.TiFlash {
 		buffer.WriteString("pushed down filter:")
-		if len(p.lateMaterializationFilterCondition) > 0 {
+		if len(p.LateMaterializationFilterCondition) > 0 {
 			if normalized {
-				buffer.Write(expression.SortedExplainNormalizedExpressionList(p.lateMaterializationFilterCondition))
+				buffer.Write(expression.SortedExplainNormalizedExpressionList(p.LateMaterializationFilterCondition))
 			} else {
-				buffer.Write(expression.SortedExplainExpressionList(p.lateMaterializationFilterCondition))
+				buffer.Write(expression.SortedExplainExpressionList(p.LateMaterializationFilterCondition))
 			}
 		} else {
 			buffer.WriteString("empty")

pkg/planner/core/physical_plans.go

@@ -829,10 +829,10 @@ type PhysicalTableScan struct {
 	// AccessCondition is used to calculate range.
 	AccessCondition []expression.Expression
 	filterCondition []expression.Expression
-	// lateMaterializationFilterCondition is used to record the filter conditions
+	// LateMaterializationFilterCondition is used to record the filter conditions
 	// that are pushed down to table scan from selection by late materialization.
 	// TODO: remove this field after we support pushing down selection to coprocessor.
-	lateMaterializationFilterCondition []expression.Expression
+	LateMaterializationFilterCondition []expression.Expression
 
 	Table   *model.TableInfo
 	Columns []*model.ColumnInfo

pkg/planner/core/plan_to_pb.go

@@ -256,10 +256,10 @@ func (p *PhysicalTableScan) ToPB(ctx sessionctx.Context, storeType kv.StoreType)
 	tsExec.KeepOrder = &keepOrder
 	tsExec.IsFastScan = &(ctx.GetSessionVars().TiFlashFastScan)
 
-	if len(p.lateMaterializationFilterCondition) > 0 {
+	if len(p.LateMaterializationFilterCondition) > 0 {
 		sc := ctx.GetSessionVars().StmtCtx
 		client := ctx.GetClient()
-		conditions, err := expression.ExpressionsToPBList(sc, p.lateMaterializationFilterCondition, client)
+		conditions, err := expression.ExpressionsToPBList(sc, p.LateMaterializationFilterCondition, client)
 		if err != nil {
 			return nil, err
 		}
@@ -296,10 +296,10 @@ func (p *PhysicalTableScan) partitionTableScanToPBForFlash(ctx sessionctx.Contex
 		telemetry.CurrentTiflashTableScanWithFastScanCount.Inc()
 	}
 
-	if len(p.lateMaterializationFilterCondition) > 0 {
+	if len(p.LateMaterializationFilterCondition) > 0 {
 		sc := ctx.GetSessionVars().StmtCtx
 		client := ctx.GetClient()
-		conditions, err := expression.ExpressionsToPBList(sc, p.lateMaterializationFilterCondition, client)
+		conditions, err := expression.ExpressionsToPBList(sc, p.LateMaterializationFilterCondition, client)
 		if err != nil {
 			return nil, err
 		}

pkg/planner/core/tiflash_selection_late_materialization.go

@@ -255,7 +255,7 @@ func predicatePushDownToTableScanImpl(sctx sessionctx.Context, physicalSelection
 	// remove the pushed down conditions from selection
 	removeSpecificExprsFromSelection(physicalSelection, selectedConds)
 	// add the pushed down conditions to table scan
-	physicalTableScan.lateMaterializationFilterCondition = selectedConds
+	physicalTableScan.LateMaterializationFilterCondition = selectedConds
 	// Update the row count of table scan after pushing down the conditions.
 	physicalTableScan.StatsInfo().RowCount *= selectedSelectivity
 }