importinto: fix incorrect index kv group count when table have common…

… handle (#51615) ref #50752
pingcap · Mar 8, 2024 · b55c855 · b55c855
1 parent 13304a2
commit b55c855
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 19 deletions.
diff --git a/pkg/disttask/importinto/encode_and_sort_operator.go b/pkg/disttask/importinto/encode_and_sort_operator.go
@@ -247,8 +247,10 @@ func getNumOfIndexGenKV(tblInfo *model.TableInfo) int {
 		if idxInfo.State != model.StatePublic {
 			continue
 		}
-		if idxInfo.Primary && !tblInfo.HasClusteredIndex() {
-			nonClusteredPK = true
+		if idxInfo.Primary {
+			if !tblInfo.HasClusteredIndex() {
+				nonClusteredPK = true
+			}
 			continue
 		}
 		count++

diff --git a/pkg/disttask/importinto/encode_and_sort_operator_test.go b/pkg/disttask/importinto/encode_and_sort_operator_test.go
@@ -185,6 +185,12 @@ func TestGetWriterMemorySizeLimit(t *testing.T) {
 			dataKVMemSizePerCon:     460175067,
 			perIndexKVMemSizePerCon: 153391689,
 		},
+		{
+			createSQL:               "create table t (a int, b int, c int, primary key(a,b,c) clustered, key(b,c), unique(b), unique(c), key(a,b))",
+			numOfIndexGenKV:         4,
+			dataKVMemSizePerCon:     460175067,
+			perIndexKVMemSizePerCon: 153391689,
+		},
 		{
 			createSQL:               "create table t (a int, b int, c int, primary key(a,b,c) nonclustered, key(b,c), unique(b), unique(c), key(a,b))",
 			numOfIndexGenKV:         5,
@@ -195,25 +201,27 @@ func TestGetWriterMemorySizeLimit(t *testing.T) {
 
 	for _, c := range cases {
 		p := parser.New()
-		node, err := p.ParseOneStmt(c.createSQL, "", "")
-		require.NoError(t, err)
-		sctx := utilmock.NewContext()
-		info, err := ddl.MockTableInfo(sctx, node.(*ast.CreateTableStmt), 1)
-		require.NoError(t, err)
-		info.State = model.StatePublic
+		t.Run(c.createSQL, func(t *testing.T) {
+			node, err := p.ParseOneStmt(c.createSQL, "", "")
+			require.NoError(t, err)
+			sctx := utilmock.NewContext()
+			info, err := ddl.MockTableInfo(sctx, node.(*ast.CreateTableStmt), 1)
+			require.NoError(t, err)
+			info.State = model.StatePublic
 
-		require.Equal(t, c.numOfIndexGenKV, getNumOfIndexGenKV(info), c.createSQL)
-		dataKVMemSizePerCon, perIndexKVMemSizePerCon := getWriterMemorySizeLimit(&proto.StepResource{
-			Mem: proto.NewAllocatable(2 * units.GiB),
-		}, &importer.Plan{
-			DesiredTableInfo: info,
-			ThreadCnt:        1,
+			require.Equal(t, c.numOfIndexGenKV, getNumOfIndexGenKV(info), c.createSQL)
+			dataKVMemSizePerCon, perIndexKVMemSizePerCon := getWriterMemorySizeLimit(&proto.StepResource{
+				Mem: proto.NewAllocatable(2 * units.GiB),
+			}, &importer.Plan{
+				DesiredTableInfo: info,
+				ThreadCnt:        1,
+			})
+			require.Equal(t, c.dataKVMemSizePerCon, dataKVMemSizePerCon, c.createSQL)
+			if c.numOfIndexGenKV > 0 {
+				require.Equal(t, c.perIndexKVMemSizePerCon, perIndexKVMemSizePerCon, c.createSQL)
+			}
+			require.LessOrEqual(t, c.dataKVMemSizePerCon+c.perIndexKVMemSizePerCon*uint64(c.numOfIndexGenKV), uint64(units.GiB))
 		})
-		require.Equal(t, c.dataKVMemSizePerCon, dataKVMemSizePerCon, c.createSQL)
-		if c.numOfIndexGenKV > 0 {
-			require.Equal(t, c.perIndexKVMemSizePerCon, perIndexKVMemSizePerCon, c.createSQL)
-		}
-		require.LessOrEqual(t, c.dataKVMemSizePerCon+c.perIndexKVMemSizePerCon*uint64(c.numOfIndexGenKV), uint64(units.GiB))
 	}
 }