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

close #60517

Author: ti-chi-bot

pkg/planner/core/physical_plan_test.go

@@ -517,3 +517,36 @@ func TestPhysicalTableScanExtractCorrelatedCols(t *testing.T) {
 	require.Equal(t, 1, len(correlated))
 	require.Equal(t, "test.t2.company_no", correlated[0].StringWithCtx(errors.RedactLogDisable))
 }
+
+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)
+}

pkg/planner/core/resolve_indices.go

@@ -505,12 +505,17 @@ func (p *PhysicalSelection) ResolveIndices() (err error) {
 
 // ResolveIndicesItself resolve indices for PhysicalPlan 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
 }

pkg/planner/property/physical_property.go

@@ -93,6 +93,19 @@ type MPPPartitionColumn struct {
 	CollateID int32
 }
 
+// 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
+}
+
 func (partitionCol *MPPPartitionColumn) hashCode() []byte {
 	hashcode := partitionCol.Col.HashCode()
 	if partitionCol.CollateID < 0 {