util: refine chunk.SwapColumn to rebuild the column reference

util/chunk/chunk.go

@@ -139,7 +139,27 @@ func (c *Chunk) MakeRef(srcColIdx, dstColIdx int) {
 
 // SwapColumn swaps column "c.columns[colIdx]" with column "other.columns[otherIdx]".
 func (c *Chunk) SwapColumn(colIdx int, other *Chunk, otherIdx int) {
+	// If there exists columns refer to the column to be swapped, we need to
+	// re-build the reference.
+	refColsIdx := make([]int, 0, len(c.columns))
+	for i := colIdx; i < len(c.columns); i++ {
+		if c.columns[i] == c.columns[colIdx] {
+			refColsIdx = append(refColsIdx, i)
+		}
+	}
+	refColsIdx4Other := make([]int, 0, len(other.columns))
+	for i := otherIdx; i < len(other.columns); i++ {
+		if other.columns[i] == other.columns[otherIdx] {
+			refColsIdx4Other = append(refColsIdx4Other, i)
+		}
+	}
 	c.columns[colIdx], other.columns[otherIdx] = other.columns[otherIdx], c.columns[colIdx]
+	for _, i := range refColsIdx {
+		c.MakeRef(colIdx, i)
+	}
+	for _, i := range refColsIdx4Other {
+		other.MakeRef(otherIdx, i)
+	}
 }
 
 // SwapColumns swaps columns with another Chunk.

util/chunk/chunk_test.go

@@ -471,6 +471,26 @@ func (s *testChunkSuite) TestChunkMemoryUsage(c *check.C) {
 	c.Assert(memUsage, check.Equals, int64(expectedUsage))
 }
 
+func (s *testChunkSuite) TestSwapColumn(c *check.C) {
+	fieldTypes := make([]*types.FieldType, 0, 2)
+	fieldTypes = append(fieldTypes, &types.FieldType{Tp: mysql.TypeFloat})
+	fieldTypes = append(fieldTypes, &types.FieldType{Tp: mysql.TypeFloat})
+
+	chk1 := NewChunkWithCapacity(fieldTypes, 1)
+	chk1.AppendInt64(0, 1)
+	chk1.MakeRef(0, 1)
+	c.Assert(chk1.columns[0] == chk1.columns[1], check.IsTrue)
+
+	chk2 := NewChunkWithCapacity(fieldTypes, 1)
+	chk2.AppendFloat64(0, 1)
+	chk2.AppendFloat64(1, 2)
+
+	chk1.SwapColumn(0, chk2, 0)
+	chk2.Reset()
+	c.Assert(chk1.columns[0] == chk1.columns[1], check.IsTrue)
+	c.Assert(chk1.columns[0] == chk2.columns[0], check.IsFalse)
+}
+
 func BenchmarkAppendInt(b *testing.B) {
 	b.ReportAllocs()
 	chk := newChunk(8)