Repository: pingcap/tidb
File: pkg/planner/core/operator/physicalop/physical_indexlookup.go (Go)

Code before:

func resetPlanIDRecursively(ctx base.PlanContext, p base.PhysicalPlan) {
	p.SetID(int(ctx.GetSessionVars().PlanID.Add(1)))
	for _, child := range p.Children() {
		resetPlanIDRecursively(ctx, child)
	}
}

func buildPushDownIndexLookUpPlan(
	ctx base.PlanContext, indexPlan base.PhysicalPlan, tablePlan base.PhysicalPlan, isCommonHandle bool,
) (indexLookUpPlan base.PhysicalPlan, err error) {
	tablePlan, err = tablePlan.Clone(ctx)
	if err != nil {
		return nil, err
	}
	resetPlanIDRecursively(ctx, tablePlan)

	var indexHandleOffsets []uint32
	if !isCommonHandle {
		// - If common handle, we don't need to set the indexHandleOffsets to build the common handle key
		// which can be read from the index value directly.
		// - If int handle, it is the last column in the index schema.
		//   - If the last column is ExtraHandleID, or a non-negative column ID, handle is the last column.
		//   - Otherwise, we need to find the last column whose ID is not ExtraHandleID and is negative.
		//     For example, partition table's IndexScan needs to append ExtraPhysTblID
		//     to the last if it is under `UnionScanExec`
		offset := indexPlan.Schema().Len() - 1
		for offset >= 0 {
			col := indexPlan.Schema().Columns[offset]
			if col.ID >= 0 || col.ID == model.ExtraHandleID {
				break
			}
			offset--
			intest.Assert(offset >= 0, "cannot find handle column in index schema")
		}

		if offset < 0 {
			return nil, errors.New("cannot find handle column in index schema")
		}
		indexHandleOffsets = []uint32{uint32(offset)}
	}

	tableScanPlan, parentOfTableScan := detachRootTableScanPlan(tablePlan)
	indexLookUpPlan = PhysicalLocalIndexLookUp{
		IndexHandleOffsets: indexHandleOffsets,
	}.Init(ctx, indexPlan, tableScanPlan, tablePlan.QueryBlockOffset())

	if parentOfTableScan != nil {
		parentOfTableScan.SetChildren(indexLookUpPlan)
		indexLookUpPlan = tablePlan
	}

Review comment (bug):
The comment is incorrect. The code searches backwards to find a column whose ID is ExtraHandleID or non-negative (>= 0), not a column "whose ID is not ExtraHandleID and is negative". The comment should be revised to accurately describe that we're looking for the handle column by skipping over extra columns with special negative IDs (like ExtraPhysTblID = -3) until we find either a regular column (ID >= 0) or the ExtraHandleID (-1).
```suggestion
		//   - Otherwise, we search backwards, skipping extra columns with special negative IDs
		//     (for example ExtraPhysTblID = -3), until we find either a regular column (ID >= 0)
		//     or the ExtraHandleID (-1), which is treated as the handle column.
```

Code after:

func resetPlanIDRecursively(ctx base.PlanContext, p base.PhysicalPlan) {
	p.SetID(int(ctx.GetSessionVars().PlanID.Add(1)))
	for _, child := range p.Children() {
		resetPlanIDRecursively(ctx, child)
	}
}

func buildPushDownIndexLookUpPlan(
	ctx base.PlanContext, indexPlan base.PhysicalPlan, tablePlan base.PhysicalPlan, isCommonHandle bool,
) (indexLookUpPlan base.PhysicalPlan, err error) {
	tablePlan, err = tablePlan.Clone(ctx)
	if err != nil {
		return nil, err
	}
	resetPlanIDRecursively(ctx, tablePlan)

	var indexHandleOffsets []uint32
	if !isCommonHandle {
		// - If common handle, we don't need to set the indexHandleOffsets to build the common handle key
		// which can be read from the index value directly.
		// - If int handle, it is the last column in the index schema.
		//   - If the last column is ExtraHandleID, or a non-negative column ID, handle is the last column.
		//   - Otherwise, we need to find the last column whose ID is not ExtraHandleID and is negative.
		//     For example, when a partition table needs to append ExtraPhysTblID
		//     to the end for the upper UnionScanExec.
		offset := indexPlan.Schema().Len() - 1
		for offset >= 0 {
			col := indexPlan.Schema().Columns[offset]
			if col.ID >= 0 || col.ID == model.ExtraHandleID {
				break
			}
			offset--
			intest.Assert(offset >= 0, "cannot find handle column in index schema")
		}

		if offset < 0 {
			return nil, errors.New("cannot find handle column in index schema")
		}
		indexHandleOffsets = []uint32{uint32(offset)}
	}

	tableScanPlan, parentOfTableScan := detachRootTableScanPlan(tablePlan)
	indexLookUpPlan = PhysicalLocalIndexLookUp{
		IndexHandleOffsets: indexHandleOffsets,
	}.Init(ctx, indexPlan, tableScanPlan, tablePlan.QueryBlockOffset())

	if parentOfTableScan != nil {
		parentOfTableScan.SetChildren(indexLookUpPlan)
		indexLookUpPlan = tablePlan
	}