planner: refactor Join and Limit's ResolveIndices (pingcap#45831)

close pingcap#45758, close pingcap#45805

Author: winoros

planner/core/issuetest/BUILD.bazel

@@ -5,6 +5,6 @@ go_test(
     timeout = "short",
     srcs = ["planner_issue_test.go"],
     flaky = True,
-    shard_count = 5,
+    shard_count = 6,
     deps = ["//testkit"],
 )

planner/core/issuetest/planner_issue_test.go

@@ -107,3 +107,13 @@ func TestIssue45036(t *testing.T) {
 		"    └─TableReader_9 10000.00 root partition:all data:TableRangeScan_8",
 		"      └─TableRangeScan_8 10000.00 cop[tikv] table:s range:[1,100000], keep order:false, stats:pseudo"))
 }
+
+func TestIssue45758(t *testing.T) {
+	store := testkit.CreateMockStore(t)
+	tk := testkit.NewTestKit(t, store)
+	tk.MustExec("use test")
+	tk.MustExec("CREATE TABLE tb1 (cid INT, code INT, class VARCHAR(10))")
+	tk.MustExec("CREATE TABLE tb2 (cid INT, code INT, class VARCHAR(10))")
+	// result ok
+	tk.MustExec("UPDATE tb1, (SELECT code AS cid, code, MAX(class) AS class FROM tb2 GROUP BY code) tb3 SET tb1.cid = tb3.cid, tb1.code = tb3.code, tb1.class = tb3.class")
+}

planner/core/resolve_indices.go

@@ -140,12 +140,25 @@ func (p *PhysicalHashJoin) ResolveIndicesItself() (err error) {
 	shallowColSlice := make([]*expression.Column, p.schema.Len())
 	copy(shallowColSlice, p.schema.Columns)
 	p.schema = expression.NewSchema(shallowColSlice...)
-	for i := 0; i < colsNeedResolving; i++ {
-		newCol, err := p.schema.Columns[i].ResolveIndices(mergedSchema)
-		if err != nil {
-			return err
+	foundCnt := 0
+	// The two column sets are all ordered. And the colsNeedResolving is the subset of the mergedSchema.
+	// So we can just move forward j if there's no matching is found.
+	// We don't use the normal ResolvIndices here since there might be duplicate columns in the schema.
+	//   e.g. The schema of child_0 is [col0, col0, col1]
+	//        ResolveIndices will only resolve all col0 reference of the current plan to the first col0.
+	for i, j := 0, 0; i < colsNeedResolving && j < len(mergedSchema.Columns); {
+		if !p.schema.Columns[i].Equal(nil, mergedSchema.Columns[j]) {
+			j++
+			continue
 		}
-		p.schema.Columns[i] = newCol.(*expression.Column)
+		p.schema.Columns[i] = p.schema.Columns[i].Clone().(*expression.Column)
+		p.schema.Columns[i].Index = j
+		i++
+		j++
+		foundCnt++
+	}
+	if foundCnt < colsNeedResolving {
+		return errors.Errorf("Some columns of %v cannot find the reference from its child(ren)", p.ExplainID().String())
 	}
 
 	return
@@ -213,12 +226,25 @@ func (p *PhysicalMergeJoin) ResolveIndices() (err error) {
 	shallowColSlice := make([]*expression.Column, p.schema.Len())
 	copy(shallowColSlice, p.schema.Columns)
 	p.schema = expression.NewSchema(shallowColSlice...)
-	for i := 0; i < colsNeedResolving; i++ {
-		newCol, err := p.schema.Columns[i].ResolveIndices(mergedSchema)
-		if err != nil {
-			return err
+	foundCnt := 0
+	// The two column sets are all ordered. And the colsNeedResolving is the subset of the mergedSchema.
+	// So we can just move forward j if there's no matching is found.
+	// We don't use the normal ResolvIndices here since there might be duplicate columns in the schema.
+	//   e.g. The schema of child_0 is [col0, col0, col1]
+	//        ResolveIndices will only resolve all col0 reference of the current plan to the first col0.
+	for i, j := 0, 0; i < colsNeedResolving && j < len(mergedSchema.Columns); {
+		if !p.schema.Columns[i].Equal(nil, mergedSchema.Columns[j]) {
+			j++
+			continue
 		}
-		p.schema.Columns[i] = newCol.(*expression.Column)
+		p.schema.Columns[i] = p.schema.Columns[i].Clone().(*expression.Column)
+		p.schema.Columns[i].Index = j
+		i++
+		j++
+		foundCnt++
+	}
+	if foundCnt < colsNeedResolving {
+		return errors.Errorf("Some columns of %v cannot find the reference from its child(ren)", p.ExplainID().String())
 	}
 	return
 }
@@ -296,12 +322,25 @@ func (p *PhysicalIndexJoin) ResolveIndices() (err error) {
 	shallowColSlice := make([]*expression.Column, p.schema.Len())
 	copy(shallowColSlice, p.schema.Columns)
 	p.schema = expression.NewSchema(shallowColSlice...)
-	for i := 0; i < colsNeedResolving; i++ {
-		newCol, err := p.schema.Columns[i].ResolveIndices(mergedSchema)
-		if err != nil {
-			return err
+	foundCnt := 0
+	// The two column sets are all ordered. And the colsNeedResolving is the subset of the mergedSchema.
+	// So we can just move forward j if there's no matching is found.
+	// We don't use the normal ResolvIndices here since there might be duplicate columns in the schema.
+	//   e.g. The schema of child_0 is [col0, col0, col1]
+	//        ResolveIndices will only resolve all col0 reference of the current plan to the first col0.
+	for i, j := 0, 0; i < colsNeedResolving && j < len(mergedSchema.Columns); {
+		if !p.schema.Columns[i].Equal(nil, mergedSchema.Columns[j]) {
+			j++
+			continue
 		}
-		p.schema.Columns[i] = newCol.(*expression.Column)
+		p.schema.Columns[i] = p.schema.Columns[i].Clone().(*expression.Column)
+		p.schema.Columns[i].Index = j
+		i++
+		j++
+		foundCnt++
+	}
+	if foundCnt < colsNeedResolving {
+		return errors.Errorf("Some columns of %v cannot find the reference from its child(ren)", p.ExplainID().String())
 	}
 
 	return
@@ -670,12 +709,25 @@ func (p *PhysicalLimit) ResolveIndices() (err error) {
 	shallowColSlice := make([]*expression.Column, p.schema.Len())
 	copy(shallowColSlice, p.schema.Columns)
 	p.schema = expression.NewSchema(shallowColSlice...)
-	for i, col := range p.schema.Columns {
-		newCol, err := col.ResolveIndices(p.children[0].Schema())
-		if err != nil {
-			return err
+	foundCnt := 0
+	// The two column sets are all ordered. And the colsNeedResolving is the subset of the mergedSchema.
+	// So we can just move forward j if there's no matching is found.
+	// We don't use the normal ResolvIndices here since there might be duplicate columns in the schema.
+	//   e.g. The schema of child_0 is [col0, col0, col1]
+	//        ResolveIndices will only resolve all col0 reference of the current plan to the first col0.
+	for i, j := 0, 0; i < p.schema.Len() && j < p.children[0].Schema().Len(); {
+		if !p.schema.Columns[i].Equal(nil, p.children[0].Schema().Columns[j]) {
+			j++
+			continue
 		}
-		p.schema.Columns[i] = newCol.(*expression.Column)
+		p.schema.Columns[i] = p.schema.Columns[i].Clone().(*expression.Column)
+		p.schema.Columns[i].Index = j
+		i++
+		j++
+		foundCnt++
+	}
+	if foundCnt < p.schema.Len() {
+		return errors.Errorf("Some columns of %v cannot find the reference from its child(ren)", p.ExplainID().String())
 	}
 	return
 }