lightning: duplicate resolution support to use DELETE to cleanup clustered table

br/pkg/lightning/backend/kv/BUILD.bazel

@@ -48,14 +48,15 @@ go_test(
     name = "kv_test",
     timeout = "short",
     srcs = [
+        "kv2sql_test.go",
         "session_internal_test.go",
         "session_test.go",
         "sql2kv_test.go",
     ],
     embed = [":kv"],
     flaky = True,
     race = "on",
-    shard_count = 16,
+    shard_count = 18,
     deps = [
         "//br/pkg/lightning/backend/encode",
         "//br/pkg/lightning/common",
@@ -70,6 +71,7 @@ go_test(
         "//parser/mysql",
         "//planner/core",
         "//sessionctx",
+        "//sessionctx/variable",
         "//table",
         "//table/tables",
         "//tablecodec",

br/pkg/lightning/backend/kv/kv2sql.go

@@ -88,10 +88,16 @@ func (t *TableKVDecoder) IterRawIndexKeys(h kv.Handle, rawRow []byte, fn func([]
 	}
 
 	indices := t.tbl.Indices()
+	isCommonHandle := t.tbl.Meta().IsCommonHandle
 
 	var buffer []types.Datum
 	var indexBuffer []byte
 	for _, index := range indices {
+		// skip clustered PK
+		if index.Meta().Primary && isCommonHandle {
+			continue
+		}
+
 		indexValues, err := index.FetchValues(row, buffer)
 		if err != nil {
 			return err

br/pkg/lightning/backend/kv/kv2sql_test.go

@@ -0,0 +1,111 @@
+// Copyright 2023 PingCAP, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package kv_test
+
+import (
+	"testing"
+
+	"github.com/pingcap/tidb/br/pkg/lightning/backend/encode"
+	"github.com/pingcap/tidb/br/pkg/lightning/backend/kv"
+	"github.com/pingcap/tidb/br/pkg/lightning/log"
+	"github.com/pingcap/tidb/ddl"
+	"github.com/pingcap/tidb/parser"
+	"github.com/pingcap/tidb/parser/ast"
+	"github.com/pingcap/tidb/parser/model"
+	"github.com/pingcap/tidb/parser/mysql"
+	"github.com/pingcap/tidb/sessionctx/variable"
+	"github.com/pingcap/tidb/table/tables"
+	"github.com/pingcap/tidb/types"
+	"github.com/pingcap/tidb/util/mock"
+	"github.com/stretchr/testify/require"
+)
+
+func TestIterRawIndexKeysClusteredPK(t *testing.T) {
+	p := parser.New()
+	node, _, err := p.ParseSQL("create table t (a varchar(10) primary key, b int, index idx(b));")
+	require.NoError(t, err)
+	mockSctx := mock.NewContext()
+	mockSctx.GetSessionVars().EnableClusteredIndex = variable.ClusteredIndexDefModeOn
+	info, err := ddl.MockTableInfo(mockSctx, node[0].(*ast.CreateTableStmt), 1)
+	require.NoError(t, err)
+	info.State = model.StatePublic
+	require.True(t, info.IsCommonHandle)
+	tbl, err := tables.TableFromMeta(kv.NewPanickingAllocators(0), info)
+	require.NoError(t, err)
+
+	sessionOpts := &encode.SessionOptions{
+		SQLMode:   mysql.ModeStrictAllTables,
+		Timestamp: 1234567890,
+	}
+	decoder, err := kv.NewTableKVDecoder(tbl, "`test`.`c1`", sessionOpts, log.L())
+	require.NoError(t, err)
+
+	sctx := kv.NewSession(sessionOpts, log.L())
+	handle, err := tbl.AddRecord(sctx, []types.Datum{types.NewIntDatum(1), types.NewIntDatum(2)})
+	require.NoError(t, err)
+	paris := sctx.TakeKvPairs()
+	require.Len(t, paris.Pairs, 2)
+	_, rowValue := paris.Pairs[0].Key, paris.Pairs[0].Val
+	idxKey := paris.Pairs[1].Key
+
+	deleteKeys := make([][]byte, 0, 2)
+	err = decoder.IterRawIndexKeys(handle, rowValue, func(bs []byte) error {
+		deleteKeys = append(deleteKeys, bs)
+		return nil
+	})
+	require.NoError(t, err)
+
+	expected := [][]byte{idxKey}
+	require.Equal(t, expected, deleteKeys)
+}
+
+func TestIterRawIndexKeysIntPK(t *testing.T) {
+	p := parser.New()
+	node, _, err := p.ParseSQL("create table t (a int primary key, b int, index idx(b));")
+	require.NoError(t, err)
+	mockSctx := mock.NewContext()
+	mockSctx.GetSessionVars().EnableClusteredIndex = variable.ClusteredIndexDefModeOn
+	info, err := ddl.MockTableInfo(mockSctx, node[0].(*ast.CreateTableStmt), 1)
+	require.NoError(t, err)
+	info.State = model.StatePublic
+	require.True(t, info.PKIsHandle)
+	tbl, err := tables.TableFromMeta(kv.NewPanickingAllocators(0), info)
+	require.NoError(t, err)
+
+	sessionOpts := &encode.SessionOptions{
+		SQLMode:   mysql.ModeStrictAllTables,
+		Timestamp: 1234567890,
+	}
+	decoder, err := kv.NewTableKVDecoder(tbl, "`test`.`c1`", sessionOpts, log.L())
+	require.NoError(t, err)
+
+	sctx := kv.NewSession(sessionOpts, log.L())
+	handle, err := tbl.AddRecord(sctx, []types.Datum{types.NewIntDatum(1), types.NewIntDatum(2)})
+	require.NoError(t, err)
+	paris := sctx.TakeKvPairs()
+	require.Len(t, paris.Pairs, 2)
+	_, rowValue := paris.Pairs[0].Key, paris.Pairs[0].Val
+	idxKey := paris.Pairs[1].Key
+
+	deleteKeys := make([][]byte, 0, 2)
+	err = decoder.IterRawIndexKeys(handle, rowValue, func(bs []byte) error {
+		deleteKeys = append(deleteKeys, bs)
+		return nil
+	})
+	require.NoError(t, err)
+
+	expected := [][]byte{idxKey}
+	require.Equal(t, expected, deleteKeys)
+}

br/tests/lightning_duplicate_resolution/run.sh

@@ -18,48 +18,56 @@ set -eux
 
 check_cluster_version 5 2 0 'duplicate detection' || exit 0
 
-run_lightning
+for i in {1..2} ; do
+    run_lightning
 
-# Ensure all tables are consistent.
-run_sql 'admin check table dup_resolve.a'
-run_sql 'admin check table dup_resolve.b'
-run_sql 'admin check table dup_resolve.c'
+    # Ensure all tables are consistent.
+    run_sql 'admin check table dup_resolve.a'
+    run_sql 'admin check table dup_resolve.b'
+    run_sql 'admin check table dup_resolve.c'
 
-## Table "a" has a clustered integer key and generated column.
+    ## Table "a" has a clustered integer key and generated column.
 
-# only one row remains (2, 1, 4.csv). all others are duplicates 🤷
-run_sql 'select count(*) from dup_resolve.a'
-check_contains 'count(*): 1'
+    # only one row remains (2, 1, 4.csv). all others are duplicates 🤷
+    run_sql 'select count(*) from dup_resolve.a'
+    check_contains 'count(*): 1'
 
-run_sql 'select * from dup_resolve.a'
-check_contains 'a: 2'
-check_contains 'b: 1'
-check_contains 'c: 4.csv'
-check_contains 'd: 3'
+    run_sql 'select * from dup_resolve.a'
+    check_contains 'a: 2'
+    check_contains 'b: 1'
+    check_contains 'c: 4.csv'
+    check_contains 'd: 3'
 
-## Table "b" has a nonclustered integer key and nullable unique column.
-run_sql 'select count(*) from dup_resolve.b'
-check_contains 'count(*): 3'
+    ## Table "b" has a nonclustered integer key and nullable unique column.
+    run_sql 'select count(*) from dup_resolve.b'
+    check_contains 'count(*): 3'
 
-run_sql 'select * from dup_resolve.b where a = 2'
-check_contains 'b: 1'
-check_contains 'c: 4.csv'
-run_sql 'select * from dup_resolve.b where a = 7'
-check_contains 'b: NULL'
-check_contains 'c: 8.csv#2'
-run_sql 'select * from dup_resolve.b where a = 8'
-check_contains 'b: NULL'
-check_contains 'c: 8.csv#3'
+    run_sql 'select * from dup_resolve.b where a = 2'
+    check_contains 'b: 1'
+    check_contains 'c: 4.csv'
+    run_sql 'select * from dup_resolve.b where a = 7'
+    check_contains 'b: NULL'
+    check_contains 'c: 8.csv#2'
+    run_sql 'select * from dup_resolve.b where a = 8'
+    check_contains 'b: NULL'
+    check_contains 'c: 8.csv#3'
 
-## Table "c" has a clustered non-integer key.
-run_sql 'select count(*) from dup_resolve.c'
-check_contains 'count(*): 4'
+    ## Table "c" has a clustered non-integer key.
+    run_sql 'select count(*) from dup_resolve.c'
+    check_contains 'count(*): 4'
 
-run_sql 'select c from dup_resolve.c where a = 2 and b = 1'
-check_contains 'c: 1.csv#1'
-run_sql 'select c from dup_resolve.c where a = 7 and b = 0'
-check_contains 'c: 1.csv#4'
-run_sql 'select c from dup_resolve.c where a = 7 and b = 1'
-check_contains 'c: 1.csv#7'
-run_sql 'select c from dup_resolve.c where a = 9 and b = 0'
-check_contains 'c: 1.csv#8'
+    run_sql 'select c from dup_resolve.c where a = 2 and b = 1'
+    check_contains 'c: 1.csv#1'
+    run_sql 'select c from dup_resolve.c where a = 7 and b = 0'
+    check_contains 'c: 1.csv#4'
+    run_sql 'select c from dup_resolve.c where a = 7 and b = 1'
+    check_contains 'c: 1.csv#7'
+    run_sql 'select c from dup_resolve.c where a = 9 and b = 0'
+    check_contains 'c: 1.csv#8'
+
+    # test can use DELETE to cleanup data
+    run_sql 'delete from dup_resolve.a'
+    run_sql 'truncate table dup_resolve.b' # TODO: non-clustered index table should also support DELETE
+    run_sql 'delete from dup_resolve.c'
+    run_sql 'drop database lightning_task_info'
+done