planner: fix bug in PhysicalExchangeSender::ResolveIndicesItself (#60520) (#60551)

close #60517

Author: ti-chi-bot

planner/core/physical_plan_test.go

@@ -25,6 +25,7 @@ import (
 	"github.com/pingcap/kvproto/pkg/metapb"
 	"github.com/pingcap/tidb/domain"
 	"github.com/pingcap/tidb/executor"
+	"github.com/pingcap/tidb/expression"
 	"github.com/pingcap/tidb/infoschema"
 	"github.com/pingcap/tidb/kv"
 	"github.com/pingcap/tidb/parser"
@@ -2615,3 +2616,36 @@ func TestAlwaysTruePredicateWithSubquery(t *testing.T) {
 		tk.MustQuery(ts).Check(testkit.Rows(output[i].Plan...))
 	}
 }
+
+func TestExchangeSenderResolveIndices(t *testing.T) {
+	schemaCols1 := make([]*expression.Column, 0, 4)
+	schemaCols1 = append(schemaCols1, &expression.Column{UniqueID: 1})
+	schemaCols1 = append(schemaCols1, &expression.Column{UniqueID: 2})
+	schemaCols1 = append(schemaCols1, &expression.Column{UniqueID: 3})
+	schemaCols1 = append(schemaCols1, &expression.Column{UniqueID: 4})
+	schema1 := expression.NewSchema(schemaCols1...)
+
+	schemaCols2 := make([]*expression.Column, 0, 2)
+	schemaCols2 = append(schemaCols2, &expression.Column{UniqueID: 3})
+	schemaCols2 = append(schemaCols2, &expression.Column{UniqueID: 4})
+	schema2 := expression.NewSchema(schemaCols2...)
+
+	partitionCol1 := &property.MPPPartitionColumn{Col: &expression.Column{UniqueID: 4}}
+
+	// two exchange sender share the same MPPPartitionColumn
+	exchangeSender1 := &core.PhysicalExchangeSender{
+		HashCols: []*property.MPPPartitionColumn{partitionCol1},
+	}
+	exchangeSender2 := &core.PhysicalExchangeSender{
+		HashCols: []*property.MPPPartitionColumn{partitionCol1},
+	}
+
+	err := exchangeSender1.ResolveIndicesItselfWithSchema(schema1)
+	require.NoError(t, err)
+
+	err = exchangeSender2.ResolveIndicesItselfWithSchema(schema2)
+	require.NoError(t, err)
+
+	// after resolving, the partition col in two different exchange sender should have different index
+	require.NotEqual(t, exchangeSender1.HashCols[0].Col.Index, exchangeSender2.HashCols[0].Col.Index)
+}

planner/core/resolve_indices.go

@@ -490,12 +490,17 @@ func (p *PhysicalSelection) ResolveIndices() (err error) {
 
 // ResolveIndicesItself resolve indices for PhyicalPlan itself
 func (p *PhysicalExchangeSender) ResolveIndicesItself() (err error) {
-	for i, col := range p.HashCols {
-		colExpr, err1 := col.Col.ResolveIndices(p.children[0].Schema())
-		if err1 != nil {
-			return err1
+	return p.ResolveIndicesItselfWithSchema(p.children[0].Schema())
+}
+
+// ResolveIndicesItselfWithSchema is added for test usage
+func (p *PhysicalExchangeSender) ResolveIndicesItselfWithSchema(inputSchema *expression.Schema) (err error) {
+	for i, hashCol := range p.HashCols {
+		newHashCol, err := hashCol.ResolveIndices(inputSchema)
+		if err != nil {
+			return err
 		}
-		p.HashCols[i].Col, _ = colExpr.(*expression.Column)
+		p.HashCols[i] = newHashCol
 	}
 	return
 }

planner/property/physical_property.go

@@ -105,6 +105,19 @@ func (partitionCol *MPPPartitionColumn) hashCode(ctx *stmtctx.StatementContext)
 	return hashcode
 }
 
+// ResolveIndices resolve index for MPPPartitionColumn
+func (partitionCol *MPPPartitionColumn) ResolveIndices(schema *expression.Schema) (*MPPPartitionColumn, error) {
+	newColExpr, err := partitionCol.Col.ResolveIndices(schema)
+	if err != nil {
+		return nil, err
+	}
+	newCol, _ := newColExpr.(*expression.Column)
+	return &MPPPartitionColumn{
+		Col:       newCol,
+		CollateID: partitionCol.CollateID,
+	}, nil
+}
+
 // Equal returns true if partitionCol == other
 func (partitionCol *MPPPartitionColumn) Equal(other *MPPPartitionColumn) bool {
 	if partitionCol.CollateID < 0 {