planner: inherit the index join info and stats bottom-up for index join inner side (#61148)

ref #60106

Author: AilinKid

pkg/planner/core/exhaust_physical_plans.go

@@ -1215,7 +1215,7 @@ func buildDataSource2IndexScanByIndexJoinProp(
 	// here we don't need to construct physical index join here anymore, because we will encapsulate it bottom-up.
 	// chosenPath and lastColManager of indexJoinResult should be returned to the caller (seen by index join to keep
 	// index join aware of indexColLens and compareFilters).
-	completeIndexJoinFeedBackInfo(innerTask.(*CopTask), indexJoinResult, indexJoinResult.chosenRanges.Range(), keyOff2IdxOff)
+	completeIndexJoinFeedBackInfo(innerTask.(*CopTask), indexJoinResult, indexJoinResult.chosenRanges, keyOff2IdxOff)
 	return innerTask
 }
 
@@ -2157,6 +2157,13 @@ func tryToGetIndexJoin(p *logicalop.LogicalJoin, prop *property.PhysicalProperty
 	return filterIndexJoinBySessionVars(p.SCtx(), candidates), false
 }
 
+func enumerationContainIndexJoin(candidates []base.PhysicalPlan) bool {
+	return slices.ContainsFunc(candidates, func(candidate base.PhysicalPlan) bool {
+		_, _, ok := getIndexJoinSideAndMethod(candidate)
+		return ok
+	})
+}
+
 // handleFilterIndexJoinHints is trying to avoid generating index join or index hash join when no-index-join related
 // hint is specified in the query. So we can do it in physic enumeration phase.
 func handleFilterIndexJoinHints(p *logicalop.LogicalJoin, candidates []base.PhysicalPlan) []base.PhysicalPlan {
@@ -2186,6 +2193,10 @@ func recordIndexJoinHintWarnings(lp base.LogicalPlan, prop *property.PhysicalPro
 	if !ok {
 		return nil
 	}
+	if !p.PreferAny(h.PreferRightAsINLJInner, h.PreferRightAsINLHJInner, h.PreferRightAsINLMJInner,
+		h.PreferLeftAsINLJInner, h.PreferLeftAsINLHJInner, h.PreferLeftAsINLMJInner) {
+		return nil // no force index join hints
+	}
 	// Cannot find any valid index join plan with these force hints.
 	// Print warning message if any hints cannot work.
 	// If the required property is not empty, we will enforce it and try the hint again.

pkg/planner/core/task.go

@@ -62,6 +62,9 @@ var HeavyFunctionNameMap = map[string]struct{}{
 }
 
 func attachPlan2Task(p base.PhysicalPlan, t base.Task) base.Task {
+	// since almost all current physical plan will be attached to bottom encapsulated task.
+	// we do the stats inheritance here for all the index join inner task.
+	inheritStatsFromBottomTaskForIndexJoinInner(p, t)
 	switch v := t.(type) {
 	case *CopTask:
 		if v.indexPlanFinished {
@@ -130,8 +133,8 @@ func (p *PhysicalUnionScan) Attach2Task(tasks ...base.Task) base.Task {
 			sel.SetChildren(pj.Children()...)
 			p.SetChildren(sel)
 			p.SetStats(task.Plan().StatsInfo())
-			rt, _ := task.(*RootTask)
-			rt.SetPlan(p)
+			rt := task.(*RootTask)
+			rt.SetPlan(p) // root task plan current is p headed.
 			pj.SetChildren(p)
 			return pj.Attach2Task(task)
 		}
@@ -160,6 +163,8 @@ func (p *PhysicalApply) Attach2Task(tasks ...base.Task) base.Task {
 	p.schema = BuildPhysicalJoinSchema(p.JoinType, p)
 	t := &RootTask{}
 	t.SetPlan(p)
+	// inherit left and right child's warnings.
+	t.warnings.CopyFrom(&lTask.(*RootTask).warnings, &rTask.(*RootTask).warnings)
 	return t
 }
 
@@ -200,6 +205,7 @@ func indexHashJoinAttach2TaskV2(p *PhysicalIndexHashJoin, tasks ...base.Task) ba
 	}
 	t := &RootTask{}
 	t.SetPlan(p)
+	t.warnings.CopyFrom(&outerTask.(*RootTask).warnings, &innerTask.(*RootTask).warnings)
 	return t
 }
 
@@ -232,6 +238,7 @@ func indexJoinAttach2TaskV2(p *PhysicalIndexJoin, tasks ...base.Task) base.Task
 	}
 	t := &RootTask{}
 	t.SetPlan(p)
+	t.warnings.CopyFrom(&outerTask.(*RootTask).warnings, &innerTask.(*RootTask).warnings)
 	return t
 }
 
@@ -264,6 +271,7 @@ func (p *PhysicalHashJoin) Attach2Task(tasks ...base.Task) base.Task {
 	p.SetChildren(lTask.Plan(), rTask.Plan())
 	task := &RootTask{}
 	task.SetPlan(p)
+	task.warnings.CopyFrom(&rTask.(*RootTask).warnings, &lTask.(*RootTask).warnings)
 	return task
 }
 
@@ -504,6 +512,7 @@ func (p *PhysicalHashJoin) attach2TaskForMpp(tasks ...base.Task) base.Task {
 		partTp:   outerTask.partTp,
 		hashCols: outerTask.hashCols,
 	}
+	task.warnings.CopyFrom(&rTask.warnings, &lTask.warnings)
 	// Current TiFlash doesn't support receive Join executors' schema info directly from TiDB.
 	// Instead, it calculates Join executors' output schema using algorithm like BuildPhysicalJoinSchema which
 	// produces full semantic schema.
@@ -594,6 +603,7 @@ func (p *PhysicalHashJoin) attach2TaskForTiFlash(tasks ...base.Task) base.Task {
 		indexPlanFinished: true,
 		tablePlan:         p,
 	}
+	task.warnings.CopyFrom(&rTask.warnings, &lTask.warnings)
 	return task
 }
 
@@ -604,6 +614,7 @@ func (p *PhysicalMergeJoin) Attach2Task(tasks ...base.Task) base.Task {
 	p.SetChildren(lTask.Plan(), rTask.Plan())
 	t := &RootTask{}
 	t.SetPlan(p)
+	t.warnings.CopyFrom(&rTask.(*RootTask).warnings, &lTask.(*RootTask).warnings)
 	return t
 }
 
@@ -1503,6 +1514,48 @@ func (sel *PhysicalSelection) Attach2Task(tasks ...base.Task) base.Task {
 	return attachPlan2Task(sel, t)
 }
 
+func inheritStatsFromBottomElemForIndexJoinInner(p base.PhysicalPlan, indexJoinInfo *IndexJoinInfo, stats *property.StatsInfo) {
+	var isIndexJoin bool
+	switch p.(type) {
+	case *PhysicalIndexJoin, *PhysicalIndexHashJoin, *PhysicalIndexMergeJoin:
+		isIndexJoin = true
+	default:
+	}
+	// indexJoinInfo != nil means the child Task comes from an index join inner side.
+	// !isIndexJoin means the childTask only be passed through to indexJoin as an END.
+	if !isIndexJoin && indexJoinInfo != nil {
+		switch p.(type) {
+		case *PhysicalSelection:
+			// todo: for simplicity, we can just inherit it from child.
+			p.StatsInfo().ScaleByExpectCnt(stats.RowCount)
+		case *PhysicalProjection:
+			// mainly about the rowEst, proj doesn't change that.
+			p.StatsInfo().ScaleByExpectCnt(stats.RowCount)
+		case *PhysicalHashAgg, *PhysicalStreamAgg:
+			// todo: for simplicity, we can just inherit it from child.
+			p.StatsInfo().ScaleByExpectCnt(stats.RowCount)
+		case *PhysicalUnionScan:
+			// todo: for simplicity, we can just inherit it from child.
+			p.StatsInfo().ScaleByExpectCnt(stats.RowCount)
+		default:
+			p.StatsInfo().ScaleByExpectCnt(stats.RowCount)
+		}
+	}
+}
+
+func inheritStatsFromBottomTaskForIndexJoinInner(p base.PhysicalPlan, t base.Task) {
+	var indexJoinInfo *IndexJoinInfo
+	switch v := t.(type) {
+	case *CopTask:
+		indexJoinInfo = v.IndexJoinInfo
+	case *RootTask:
+		indexJoinInfo = v.IndexJoinInfo
+	default:
+		// index join's inner side couldn't be a mppTask, leave it.
+	}
+	inheritStatsFromBottomElemForIndexJoinInner(p, indexJoinInfo, t.Plan().StatsInfo())
+}
+
 // CheckAggCanPushCop checks whether the aggFuncs and groupByItems can
 // be pushed down to coprocessor.
 func CheckAggCanPushCop(sctx base.PlanContext, aggFuncs []*aggregation.AggFuncDesc, groupByItems []expression.Expression, storeType kv.StoreType) bool {
@@ -2179,6 +2232,10 @@ func (p *PhysicalStreamAgg) Attach2Task(tasks ...base.Task) base.Task {
 			if partialAgg != nil {
 				if cop.tablePlan != nil {
 					cop.finishIndexPlan()
+					// the partialAgg attachment didn't follow the attachPlan2Task function, so here we actively call
+					// inheritStatsFromBottomForIndexJoinInner(p, t) to inherit stats from the bottom plan for index
+					// join inner side. note: partialAgg will share stats with finalAgg.
+					inheritStatsFromBottomElemForIndexJoinInner(partialAgg, cop.IndexJoinInfo, cop.tablePlan.StatsInfo())
 					partialAgg.SetChildren(cop.tablePlan)
 					cop.tablePlan = partialAgg
 					// If needExtraProj is true, a projection will be created above the PhysicalIndexLookUpReader to make sure
@@ -2189,10 +2246,15 @@ func (p *PhysicalStreamAgg) Attach2Task(tasks ...base.Task) base.Task {
 					// the partial agg, and the schema will be broken.
 					cop.needExtraProj = false
 				} else {
+					// the partialAgg attachment didn't follow the attachPlan2Task function, so here we actively call
+					// inheritStatsFromBottomForIndexJoinInner(p, t) to inherit stats from the bottom plan for index
+					// join inner side. note: partialAgg will share stats with finalAgg.
+					inheritStatsFromBottomElemForIndexJoinInner(partialAgg, cop.IndexJoinInfo, cop.indexPlan.StatsInfo())
 					partialAgg.SetChildren(cop.indexPlan)
 					cop.indexPlan = partialAgg
 				}
 			}
+			// COP Task -> Root Task, warnings inherited inside.
 			t = cop.ConvertToRootTask(p.SCtx())
 			attachPlan2Task(finalAgg, t)
 		}
@@ -2680,6 +2742,10 @@ func (p *PhysicalHashAgg) Attach2Task(tasks ...base.Task) base.Task {
 			if partialAgg != nil {
 				if cop.tablePlan != nil {
 					cop.finishIndexPlan()
+					// the partialAgg attachment didn't follow the attachPlan2Task function, so here we actively call
+					// inheritStatsFromBottomForIndexJoinInner(p, t) to inherit stats from the bottom plan for index
+					// join inner side. note: partialAgg will share stats with finalAgg.
+					inheritStatsFromBottomElemForIndexJoinInner(partialAgg, cop.IndexJoinInfo, cop.tablePlan.StatsInfo())
 					partialAgg.SetChildren(cop.tablePlan)
 					cop.tablePlan = partialAgg
 					// If needExtraProj is true, a projection will be created above the PhysicalIndexLookUpReader to make sure
@@ -2690,6 +2756,10 @@ func (p *PhysicalHashAgg) Attach2Task(tasks ...base.Task) base.Task {
 					// the partial agg, and the schema will be broken.
 					cop.needExtraProj = false
 				} else {
+					// the partialAgg attachment didn't follow the attachPlan2Task function, so here we actively call
+					// inheritStatsFromBottomForIndexJoinInner(p, t) to inherit stats from the bottom plan for index
+					// join inner side. note: partialAgg will share stats with finalAgg.
+					inheritStatsFromBottomElemForIndexJoinInner(partialAgg, cop.IndexJoinInfo, cop.indexPlan.StatsInfo())
 					partialAgg.SetChildren(cop.indexPlan)
 					cop.indexPlan = partialAgg
 				}
@@ -2744,17 +2814,20 @@ func (p *PhysicalCTEStorage) Attach2Task(tasks ...base.Task) base.Task {
 	t := tasks[0].Copy()
 	if mpp, ok := t.(*MppTask); ok {
 		p.SetChildren(t.Plan())
-		return &MppTask{
+		nt := &MppTask{
 			p:           p,
 			partTp:      mpp.partTp,
 			hashCols:    mpp.hashCols,
 			tblColHists: mpp.tblColHists,
 		}
+		nt.warnings.CopyFrom(&mpp.warnings)
+		return nt
 	}
 	t.ConvertToRootTask(p.SCtx())
 	p.SetChildren(t.Plan())
 	ta := &RootTask{}
 	ta.SetPlan(p)
+	ta.warnings.CopyFrom(&t.(*RootTask).warnings)
 	return ta
 }
 
@@ -2782,6 +2855,21 @@ func (p *PhysicalSequence) Attach2Task(tasks ...base.Task) base.Task {
 		hashCols:    lastTask.hashCols,
 		tblColHists: lastTask.tblColHists,
 	}
+	tmpWarnings := make([]*simpleWarnings, 0, len(tasks))
+	for _, t := range tasks {
+		if mpp, ok := t.(*MppTask); ok {
+			tmpWarnings = append(tmpWarnings, &mpp.warnings)
+			continue
+		}
+		if root, ok := t.(*RootTask); ok {
+			tmpWarnings = append(tmpWarnings, &root.warnings)
+			continue
+		}
+		if cop, ok := t.(*CopTask); ok {
+			tmpWarnings = append(tmpWarnings, &cop.warnings)
+		}
+	}
+	mppTask.warnings.CopyFrom(tmpWarnings...)
 	return mppTask
 }
 
@@ -2880,11 +2968,13 @@ func (t *MppTask) enforceExchangerImpl(prop *property.PhysicalProperty) *MppTask
 	sender.SetChildren(t.p)
 	receiver := PhysicalExchangeReceiver{}.Init(ctx, t.p.StatsInfo())
 	receiver.SetChildren(sender)
-	return &MppTask{
+	nt := &MppTask{
 		p:        receiver,
 		partTp:   prop.MPPPartitionTp,
 		hashCols: prop.MPPPartitionCols,
 	}
+	nt.warnings.CopyFrom(&t.warnings)
+	return nt
 }
 
 // IndexJoinInfo is generated by index join's inner ds, which will build their own index choice based

pkg/planner/core/task_base.go

@@ -37,8 +37,6 @@ var (
 	_ base.Task = &CopTask{}
 )
 
-var _ context.WarnGetterAppender = &simpleWarnings{}
-
 type simpleWarnings struct {
 	warnings []*context.SQLWarn
 }
@@ -189,6 +187,11 @@ func (t *RootTask) MemoryUsage() (sum int64) {
 	return sum
 }
 
+// AppendWarning appends a warning
+func (t *RootTask) AppendWarning(err error) {
+	t.warnings.AppendWarning(err)
+}
+
 // ************************************* RootTask End ******************************************
 
 // ************************************* MPPTask Start ******************************************
@@ -263,6 +266,11 @@ func (t *MppTask) MemoryUsage() (sum int64) {
 	return
 }
 
+// AppendWarning appends a warning
+func (t *MppTask) AppendWarning(err error) {
+	t.warnings.AppendWarning(err)
+}
+
 // ConvertToRootTaskImpl implements Task interface.
 func (t *MppTask) ConvertToRootTaskImpl(ctx base.PlanContext) (rt *RootTask) {
 	defer func() {
@@ -367,6 +375,11 @@ type CopTask struct {
 	warnings simpleWarnings
 }
 
+// AppendWarning appends a warning
+func (t *CopTask) AppendWarning(err error) {
+	t.warnings.AppendWarning(err)
+}
+
 // Invalid implements Task interface.
 func (t *CopTask) Invalid() bool {
 	return t.tablePlan == nil && t.indexPlan == nil && len(t.idxMergePartPlans) == 0

pkg/util/context/warn.go

@@ -85,13 +85,6 @@ type WarnAppender interface {
 	AppendNote(msg error)
 }
 
-// WarnGetterAppender provides a function to add a warning and get all warnings.
-type WarnGetterAppender interface {
-	WarnAppender
-	// GetWarnings gets all warnings. The slice is not copied, so it should not be modified.
-	GetWarnings() []SQLWarn
-}
-
 // WarnHandler provides a handler to append and get warnings.
 type WarnHandler interface {
 	WarnAppender