planner, core: implement partial order TopN attach2Task and partial order flow (#65799)

close #65813

Author: elsa0520

pkg/planner/core/find_best_task.go

@@ -753,28 +753,10 @@ type candidatePath struct {
 	accessCondsColMap util.Col2Len // accessCondsColMap maps Column.UniqueID to column length for the columns in AccessConds.
 	indexCondsColMap  util.Col2Len // indexCondsColMap maps Column.UniqueID to column length for the columns in AccessConds and indexFilters.
 	matchPropResult   property.PhysicalPropMatchResult
-
-	indexJoinCols int // how many index columns are used in access conditions in this IndexJoin.
-
 	// partialOrderMatch records the partial order match result for TopN optimization.
-	// When this field is not nil, it means this path can provide partial order using prefix index.
-	partialOrderMatch *PartialOrderMatchResult
-}
-
-// PartialOrderMatchResult records the result of matching partial order for TopN optimization.
-type PartialOrderMatchResult struct {
-	// Matched indicates whether the index can provide partial order
-	Matched bool
-	// PrefixColID is the last and only one prefix column ID of index, only used for executor part
-	// For example:
-	// Query ORDER BY a,b,c
-	// Index: a, b, c(10)
-	// ColIds: a=0, b=1, c=2
-	// PrefixColID: 2, the col id of c
-	// PrefixLen: 10, the col length of c in index
-	PrefixColID int64
-	// PrefixLen is the prefix length in bytes for prefix index,  only used for executor part
-	PrefixLen int
+	// When the matched is true, it means this path can provide partial order using prefix index.
+	partialOrderMatchResult property.PartialOrderMatchResult // Result of matching partial order property
+	indexJoinCols           int                              // how many index columns are used in access conditions in this IndexJoin.
 }
 
 func compareBool(l, r bool) int {
@@ -1170,73 +1152,86 @@ func matchProperty(ds *logicalop.DataSource, path *util.AccessPath, prop *proper
 // the index can **completely match** the **prefix** of the order by column.
 // Case 1: Prefix index INDEX idx(col(N)) matches ORDER BY col
 // For example:
-// query: order by a, b
-// index: prefix(a)
-// index: a, prefix(b)
+//
+//	query: order by a, b
+//	index: a(10)
+//	index: a, b(10)
+//
+// On success, this function will return `PartialOrderMatchResult.Matched`= true, `PartialOrderMatchResult.PrefixCol` and `PartialOrderMatchResult.PrefixLen`.
+// Otherwise it returns false and not initialize `PrefixCol and PrefixLen`.
 // TODO
 // Case 2: Composite index INDEX idx(a, b) matches ORDER BY a, b, c (index provides order for a, b)
 // In fact, there is a Case3 that can also be supported, but it will not be explained in detail here.
 // Please refer to the design document for details.
-func matchPartialOrderProperty(path *util.AccessPath, partialOrderInfo *property.PartialOrderInfo) *PartialOrderMatchResult {
+func matchPartialOrderProperty(path *util.AccessPath, partialOrderInfo *property.PartialOrderInfo) property.PartialOrderMatchResult {
+	emptyResult := property.PartialOrderMatchResult{Matched: false}
+
 	if partialOrderInfo == nil || path.Index == nil || len(path.IdxCols) == 0 {
-		return nil
+		return emptyResult
 	}
 
 	sortItems := partialOrderInfo.SortItems
 	if len(sortItems) == 0 {
-		return nil
+		return emptyResult
 	}
 
 	allSameOrder, _ := partialOrderInfo.AllSameOrder()
 	if !allSameOrder {
-		return nil
+		return emptyResult
 	}
 
 	// Case 1: Prefix index INDEX idx(col(N)) matches ORDER BY col
 	// Check if index columns can match ORDER BY columns (allowing prefix index)
-	// Constraint 1: The number of index columns must be <= the number of ORDER BY columns
-	if len(path.IdxCols) > len(sortItems) {
-		return nil
-	}
-	// Constraint 2: The last column of the index must be a prefix column
-	if path.IdxColLens[len(path.IdxCols)-1] == types.UnspecifiedLength {
+	//
+	// NOTE: We use path.Index.Columns to get the actual index definition columns count,
+	// because path.IdxCols may include additional handle columns (e.g., primary key `id`)
+	// appended for non-unique indexes on tables with PKIsHandle.
+	// For example, for `index idx_name_prefix (name(10))` on a table with `id int primary key`,
+	// path.IdxCols = [name, id] but path.Index.Columns only contains [name].
+	// We should only consider the actual index definition columns for partial order matching.
+	indexColCount := len(path.Index.Columns)
+
+	// Constraint 1: The number of index definition columns must be <= the number of ORDER BY columns
+	if indexColCount > len(sortItems) {
+		return emptyResult
+	}
+	// Constraint 2: The last column of the index definition must be a prefix column
+	lastIdxColLen := path.Index.Columns[indexColCount-1].Length
+	if lastIdxColLen == types.UnspecifiedLength {
 		// The last column is not a prefix column, skip this index
-		return nil
+		return emptyResult
 	}
 	// Extract ORDER BY columns
 	orderByCols := make([]*expression.Column, 0, len(sortItems))
 	for _, item := range sortItems {
 		orderByCols = append(orderByCols, item.Col)
 	}
 
-	var prefixColumnID int64
-	var prefixLen int
-	for i := range len(path.IdxCols) {
+	// Only iterate over the actual index definition columns, not the appended handle columns
+	for i := range indexColCount {
 		// check if the same column
 		if !orderByCols[i].EqualColumn(path.IdxCols[i]) {
-			return nil
+			return emptyResult
 		}
 
 		// meet prefix index column, match termination
 		if path.IdxColLens[i] != types.UnspecifiedLength {
-			// If we meet a prefix column but it's not the last index column, it's not supported.
+			// If we meet a prefix column but it's not the last index definition column, it's not supported.
 			// e.g. prefix(a), b cannot provide partial order for ORDER BY a, b.
-			if i != len(path.IdxCols)-1 {
-				return nil
+			if i != indexColCount-1 {
+				return emptyResult
 			}
 			// Encountered a prefix index column.
 			// This prefix index column can provide partial order, but subsequent columns cannot match.
-			prefixColumnID = path.IdxCols[i].UniqueID
-			prefixLen = path.IdxColLens[i]
-			return &PartialOrderMatchResult{
-				Matched:     true,
-				PrefixColID: prefixColumnID,
-				PrefixLen:   prefixLen,
+			return property.PartialOrderMatchResult{
+				Matched:   true,
+				PrefixCol: path.IdxCols[i],
+				PrefixLen: path.IdxColLens[i],
 			}
 		}
 	}
 
-	return nil
+	return emptyResult
 }
 
 // GroupRangesByCols groups the ranges by the values of the columns specified by groupByColIdxs.
@@ -1462,6 +1457,13 @@ func getTableCandidate(ds *logicalop.DataSource, path *util.AccessPath, prop *pr
 func getIndexCandidate(ds *logicalop.DataSource, path *util.AccessPath, prop *property.PhysicalProperty) *candidatePath {
 	candidate := &candidatePath{path: path}
 	candidate.matchPropResult = matchProperty(ds, path, prop)
+	// Because the skyline pruning already prune the indexes that cannot provide partial order
+	// when prop has PartialOrderInfo physical property,
+	// So here we just need to record the partial order match result(prefixCol, prefixLen).
+	// The partialOrderMatchResult.Matched() will be always true after skyline pruning.
+	if ds.SCtx().GetSessionVars().IsPartialOrderedIndexForTopNEnabled() && prop.PartialOrderInfo != nil {
+		candidate.partialOrderMatchResult = matchPartialOrderProperty(path, prop.PartialOrderInfo)
+	}
 	candidate.accessCondsColMap = util.ExtractCol2Len(ds.SCtx().GetExprCtx().GetEvalCtx(), path.AccessConds, path.IdxCols, path.IdxColLens)
 	candidate.indexCondsColMap = util.ExtractCol2Len(ds.SCtx().GetExprCtx().GetEvalCtx(), append(path.AccessConds, path.IndexFilters...), path.FullIdxCols, path.FullIdxColLens)
 	return candidate
@@ -1528,17 +1530,29 @@ func skylinePruning(ds *logicalop.DataSource, prop *property.PhysicalProperty) [
 		}
 		var currentCandidate *candidatePath
 		if path.IsTablePath() {
+			if prop.PartialOrderInfo != nil {
+				// skyline pruning table path with partial order property is not supported yet.
+				// TODO: support it in the future after we support prefix column as partial order.
+				continue
+			}
 			currentCandidate = getTableCandidate(ds, path, prop)
 		} else {
-			// Check if candidate path match the partial order property
-			// To consider an index for partial order optimization, the following conditions must be met together:
-			// 1. OptPartialOrderedIndexForTopN should be enabled.
-			// 2. Physical property should has PartialOrderInfo
-			// 3. The path should match PartialOrderInfo
-			// TODO: use the PartialOrderMatchResult in the further PR to construct the special TopN and Limit executor
-			matchPartialOrderIndex := ds.SCtx().GetSessionVars().OptPartialOrderedIndexForTopN &&
-				prop.PartialOrderInfo != nil &&
-				matchPartialOrderProperty(path, prop.PartialOrderInfo) != nil
+			// Check if this path can be used for partial order optimization
+			var matchPartialOrderIndex bool
+			if ds.SCtx().GetSessionVars().IsPartialOrderedIndexForTopNEnabled() &&
+				prop.PartialOrderInfo != nil {
+				if !matchPartialOrderProperty(path, prop.PartialOrderInfo).Matched {
+					// skyline pruning all indexes that cannot provide partial order when we are looking for
+					continue
+				}
+				matchPartialOrderIndex = true
+				// If the index can match partial order requirement and user use "use/force index" in hint.
+				// If the index can't match partial order requirement and use use "use/force index" and enable partial order optimization together,
+				// the behavior will degenerate into normal index use behavior without considering partial order optimization.
+				if path.Forced {
+					path.ForcePartialOrder = true
+				}
+			}
 
 			// We will use index to generate physical plan if any of the following conditions is satisfied:
 			// 1. This path's access cond is not nil.
@@ -1551,6 +1565,8 @@ func skylinePruning(ds *logicalop.DataSource, prop *property.PhysicalProperty) [
 				// If none of the above conditions are met, this index will be directly pruned here.
 				continue
 			}
+
+			// After passing the check, generate the candidate
 			currentCandidate = getIndexCandidate(ds, path, prop)
 		}
 		pruned := false
@@ -2245,17 +2261,29 @@ func convertToIndexScan(ds *logicalop.DataSource, prop *property.PhysicalPropert
 		// If it's parent requires double read task, return max cost.
 		return base.InvalidTask, nil
 	}
+	// Check if sort items can be matched. If not, return Invalid task
 	if !prop.IsSortItemEmpty() && !candidate.matchPropResult.Matched() {
 		return base.InvalidTask, nil
 	}
 	// If we need to keep order for the index scan, we should forbid the non-keep-order index scan when we try to generate the path.
-	if prop.IsSortItemEmpty() && candidate.path.ForceKeepOrder {
+	if !prop.NeedKeepOrder() && candidate.path.ForceKeepOrder {
 		return base.InvalidTask, nil
 	}
 	// If we don't need to keep order for the index scan, we should forbid the non-keep-order index scan when we try to generate the path.
-	if !prop.IsSortItemEmpty() && candidate.path.ForceNoKeepOrder {
+	if prop.NeedKeepOrder() && candidate.path.ForceNoKeepOrder {
 		return base.InvalidTask, nil
 	}
+	// If we want to force partial order, then we should remove all others property candidate path such as: full order and no order.
+	if candidate.path.ForcePartialOrder && prop.PartialOrderInfo == nil {
+		return base.InvalidTask, nil
+	}
+	// For partial order property
+	// We **don't need to check** the partial order property is matched in here.
+	// Because if the index scan cannot satisfy partial order, it will be pruned at the SkylinePruning phase
+	// (which is the previous phase before this function).
+	// So, if an index can enter this function and also contains the requirement of a partial order property,
+	// then it must meet the requirements.
+
 	path := candidate.path
 	is := physicalop.GetOriginalPhysicalIndexScan(ds, prop, path, candidate.matchPropResult.Matched(), candidate.path.IsSingleScan)
 	cop := &physicalop.CopTask{
@@ -2264,12 +2292,17 @@ func convertToIndexScan(ds *logicalop.DataSource, prop *property.PhysicalPropert
 		TblCols:     ds.TblCols,
 		ExpectCnt:   uint64(prop.ExpectedCnt),
 	}
+	// Store partial order match result in CopTask for use in attach2Task
+	if candidate.partialOrderMatchResult.Matched {
+		cop.PartialOrderMatchResult = &candidate.partialOrderMatchResult
+	}
 	cop.PhysPlanPartInfo = &physicalop.PhysPlanPartInfo{
 		PruningConds:   ds.AllConds,
 		PartitionNames: ds.PartitionNames,
 		Columns:        ds.TblCols,
 		ColumnNames:    ds.OutputNames(),
 	}
+
 	if !candidate.path.IsSingleScan {
 		// On this way, it's double read case.
 		ts := physicalop.PhysicalTableScan{
@@ -2307,7 +2340,8 @@ func convertToIndexScan(ds *logicalop.DataSource, prop *property.PhysicalPropert
 			}
 		}
 	}
-	if candidate.matchPropResult.Matched() {
+	// handles both normal sort (SortItems) and partial order (PartialOrderInfo)
+	if prop.NeedKeepOrder() {
 		cop.KeepOrder = true
 		if cop.TablePlan != nil && !ds.TableInfo.IsCommonHandle {
 			col, isNew := cop.TablePlan.(*physicalop.PhysicalTableScan).AppendExtraHandleCol(ds)
@@ -2322,8 +2356,9 @@ func convertToIndexScan(ds *logicalop.DataSource, prop *property.PhysicalPropert
 		// Case 3: both
 		if (ds.TableInfo.GetPartitionInfo() != nil && !is.Index.Global) ||
 			candidate.matchPropResult == property.PropMatchedNeedMergeSort {
-			byItems := make([]*util.ByItems, 0, len(prop.SortItems))
-			for _, si := range prop.SortItems {
+			sortItems := prop.GetSortItemsForKeepOrder()
+			byItems := make([]*util.ByItems, 0, len(sortItems))
+			for _, si := range sortItems {
 				byItems = append(byItems, &util.ByItems{
 					Expr: si.Col,
 					Desc: si.Desc,

pkg/planner/core/hint_test.go

@@ -144,34 +144,29 @@ func TestSetVarPartialOrderedIndexForTopN(t *testing.T) {
 		testKit.MustExec(`use test`)
 
 		// Test default value
-		testKit.MustQuery(`select @@tidb_opt_partial_ordered_index_for_topn`).Check(testkit.Rows("0"))
+		testKit.MustQuery(`select @@tidb_opt_partial_ordered_index_for_topn`).Check(testkit.Rows("DISABLE"))
 
 		// Test set_var hint changes the value during query execution
-		testKit.MustExec(`set @@tidb_opt_partial_ordered_index_for_topn = 0`)
-		testKit.MustQuery(`select /*+ set_var(tidb_opt_partial_ordered_index_for_topn=ON) */ @@tidb_opt_partial_ordered_index_for_topn`).Check(testkit.Rows("1"))
+		testKit.MustExec(`set @@tidb_opt_partial_ordered_index_for_topn = DISABLE`)
+		testKit.MustQuery(`select /*+ set_var(tidb_opt_partial_ordered_index_for_topn=COST) */ @@tidb_opt_partial_ordered_index_for_topn`).Check(testkit.Rows("COST"))
 		// Value should be restored after query
-		testKit.MustQuery(`select @@tidb_opt_partial_ordered_index_for_topn`).Check(testkit.Rows("0"))
+		testKit.MustQuery(`select @@tidb_opt_partial_ordered_index_for_topn`).Check(testkit.Rows("DISABLE"))
 
 		// Test set_var hint with OFF
-		testKit.MustExec(`set @@tidb_opt_partial_ordered_index_for_topn = 1`)
-		testKit.MustQuery(`select /*+ set_var(tidb_opt_partial_ordered_index_for_topn=OFF) */ @@tidb_opt_partial_ordered_index_for_topn`).Check(testkit.Rows("0"))
+		testKit.MustExec(`set @@tidb_opt_partial_ordered_index_for_topn = COST`)
+		testKit.MustQuery(`select /*+ set_var(tidb_opt_partial_ordered_index_for_topn=DISABLE) */ @@tidb_opt_partial_ordered_index_for_topn`).Check(testkit.Rows("DISABLE"))
 		// Value should be restored after query
-		testKit.MustQuery(`select @@tidb_opt_partial_ordered_index_for_topn`).Check(testkit.Rows("1"))
-
-		// Test set_var hint with numeric values
-		testKit.MustExec(`set @@tidb_opt_partial_ordered_index_for_topn = 0`)
-		testKit.MustQuery(`select /*+ set_var(tidb_opt_partial_ordered_index_for_topn=1) */ @@tidb_opt_partial_ordered_index_for_topn`).Check(testkit.Rows("1"))
-		testKit.MustQuery(`select @@tidb_opt_partial_ordered_index_for_topn`).Check(testkit.Rows("0"))
+		testKit.MustQuery(`select @@tidb_opt_partial_ordered_index_for_topn`).Check(testkit.Rows("COST"))
 
 		// Test set_var hint with multiple queries
 		testKit.MustExec(`create table t(a int, b varchar(10), index idx_b(b(5)));`)
-		testKit.MustExec(`set @@tidb_opt_partial_ordered_index_for_topn = 0`)
-		testKit.MustExec(`select /*+ set_var(tidb_opt_partial_ordered_index_for_topn=ON) */ * from t order by b limit 10;`)
-		testKit.MustQuery(`select @@tidb_opt_partial_ordered_index_for_topn`).Check(testkit.Rows("0"))
+		testKit.MustExec(`set @@tidb_opt_partial_ordered_index_for_topn = DISABLE`)
+		testKit.MustExec(`select /*+ set_var(tidb_opt_partial_ordered_index_for_topn=COST) */ * from t order by b limit 10;`)
+		testKit.MustQuery(`select @@tidb_opt_partial_ordered_index_for_topn`).Check(testkit.Rows("DISABLE"))
 
 		// Test with EXPLAIN (should not change the value)
-		testKit.MustExec(`set @@tidb_opt_partial_ordered_index_for_topn = 0`)
-		testKit.MustExec(`explain select /*+ set_var(tidb_opt_partial_ordered_index_for_topn=ON) */ * from t order by b limit 10;`)
-		testKit.MustQuery(`select @@tidb_opt_partial_ordered_index_for_topn`).Check(testkit.Rows("0"))
+		testKit.MustExec(`set @@tidb_opt_partial_ordered_index_for_topn = DISABLE`)
+		testKit.MustExec(`explain select /*+ set_var(tidb_opt_partial_ordered_index_for_topn=COST) */ * from t order by b limit 10;`)
+		testKit.MustQuery(`select @@tidb_opt_partial_ordered_index_for_topn`).Check(testkit.Rows("DISABLE"))
 	})
 }

pkg/planner/core/operator/physicalop/physical_index_scan.go

@@ -688,8 +688,9 @@ func GetOriginalPhysicalIndexScan(ds *logicalop.DataSource, prop *property.Physi
 	if usedStats != nil && usedStats.GetUsedInfo(is.PhysicalTableID) != nil {
 		is.UsedStatsInfo = usedStats.GetUsedInfo(is.PhysicalTableID)
 	}
-	if isMatchProp {
-		is.Desc = prop.SortItems[0].Desc
+	// Index scan should maintain order (true for both normal sorting via SortItems and partial order via PartialOrderInfo)
+	if prop.NeedKeepOrder() {
+		is.Desc = prop.GetSortDescForKeepOrder()
 		is.KeepOrder = true
 	}
 	return is

pkg/planner/core/operator/physicalop/physical_limit.go

@@ -37,6 +37,15 @@ type PhysicalLimit struct {
 	PartitionBy []property.SortItem
 	Offset      uint64
 	Count       uint64
+
+	// PrefixCol is the prefix index column for partial order optimization.
+	// Used for both execution (via UniqueID) and explain (via column name).
+	// If prefix index optimization is not used, this field is nil.
+	PrefixCol *expression.Column
+
+	// PrefixLen is the prefix index length (in bytes) for TiKV-side short-circuiting.
+	// If prefix index optimization is not used, this field is 0.
+	PrefixLen int
 }
 
 // ExhaustPhysicalPlans4LogicalLimit will be called by LogicalLimit in logicalOp pkg.
@@ -101,7 +110,10 @@ func (p *PhysicalLimit) MemoryUsage() (sum int64) {
 		return
 	}
 
-	sum = p.PhysicalSchemaProducer.MemoryUsage() + size.SizeOfUint64*2
+	sum = p.PhysicalSchemaProducer.MemoryUsage() +
+		size.SizeOfUint64*2 + // Offset, Count
+		size.SizeOfInt64 + // PrefixColID
+		size.SizeOfInt // PrefixLen
 	return
 }
 
@@ -116,10 +128,23 @@ func (p *PhysicalLimit) ExplainInfo() string {
 	}
 	if redact == perrors.RedactLogDisable {
 		fmt.Fprintf(buffer, "offset:%v, count:%v", p.Offset, p.Count)
+		if p.PrefixCol != nil {
+			prefixColName := p.PrefixCol.ColumnExplainInfo(ectx, false)
+			fmt.Fprintf(buffer, ", prefix_col:%v, prefix_len:%v",
+				prefixColName, p.PrefixLen)
+		}
 	} else if redact == perrors.RedactLogMarker {
 		fmt.Fprintf(buffer, "offset:‹%v›, count:‹%v›", p.Offset, p.Count)
+		if p.PrefixCol != nil {
+			prefixColName := p.PrefixCol.ColumnExplainInfo(ectx, false)
+			fmt.Fprintf(buffer, ", prefix_col:‹%v›, prefix_len:‹%v›",
+				prefixColName, p.PrefixLen)
+		}
 	} else if redact == perrors.RedactLogEnable {
 		fmt.Fprintf(buffer, "offset:?, count:?")
+		if p.PrefixCol != nil {
+			fmt.Fprintf(buffer, ", prefix_col:?, prefix_len:?")
+		}
 	}
 	return buffer.String()
 }