infoschema: optimize performance of FindTableByPartitionID for infosc…

…hema v2 (#52479)

close #52239

pkg/infoschema/infoschema_v2.go

@@ -97,6 +97,17 @@ type Data struct {
 
 	// For information_schema/metrics_schema/performance_schema etc
 	specials map[string]*schemaTables
+
+	// pid2tid is used by FindTableInfoByPartitionID, it stores {partitionID, schemaVersion} => table ID
+	// Need full data in memory!
+	pid2tid *btree.BTreeG[partitionItem]
+}
+
+type partitionItem struct {
+	partitionID   int64
+	schemaVersion int64
+	tableID       int64
+	tomb          bool
 }
 
 func (isd *Data) getVersionByTS(ts uint64) (int64, bool) {
@@ -148,6 +159,7 @@ func NewData() *Data {
 		// TODO: limit by size instead of by table count.
 		tableCache: sieve.New[tableCacheKey, table.Table](1000),
 		specials:   make(map[string]*schemaTables),
+		pid2tid:    btree.NewBTreeG[partitionItem](comparePartitionItem),
 	}
 	return ret
 }
@@ -156,6 +168,11 @@ func (isd *Data) add(item tableItem, tbl table.Table) {
 	isd.byID.Set(item)
 	isd.byName.Set(item)
 	isd.tableCache.Set(tableCacheKey{item.tableID, item.schemaVersion}, tbl)
+	if pi := tbl.Meta().GetPartitionInfo(); pi != nil {
+		for _, def := range pi.Definitions {
+			isd.pid2tid.Set(partitionItem{def.ID, item.schemaVersion, tbl.Meta().ID, false})
+		}
+	}
 }
 
 func (isd *Data) addSpecialDB(di *model.DBInfo, tables *schemaTables) {
@@ -207,6 +224,16 @@ func compareByName(a, b tableItem) bool {
 	return a.schemaVersion < b.schemaVersion
 }
 
+func comparePartitionItem(a, b partitionItem) bool {
+	if a.partitionID < b.partitionID {
+		return true
+	}
+	if a.partitionID > b.partitionID {
+		return false
+	}
+	return a.schemaVersion < b.schemaVersion
+}
+
 func compareSchemaItem(a, b schemaItem) bool {
 	if a.Name() < b.Name() {
 		return true
@@ -464,23 +491,52 @@ func (is *infoschemaV2) SchemaExists(schema model.CIStr) bool {
 }
 
 func (is *infoschemaV2) FindTableByPartitionID(partitionID int64) (table.Table, *model.DBInfo, *model.PartitionDefinition) {
-	// TODO: This is quite inefficient! we need some better way or avoid this API.
-	dbInfos := is.AllSchemas()
-	for _, dbInfo := range dbInfos {
-		tbls := is.SchemaTables(dbInfo.Name)
-		for _, tbl := range tbls {
-			pi := tbl.Meta().GetPartitionInfo()
-			if pi == nil {
-				continue
+	var ok bool
+	var pi partitionItem
+	is.pid2tid.Descend(partitionItem{partitionID: partitionID, schemaVersion: math.MaxInt64},
+		func(item partitionItem) bool {
+			if item.partitionID != partitionID {
+				return false
 			}
-			for _, p := range pi.Definitions {
-				if p.ID == partitionID {
-					return tbl, dbInfo, &p
-				}
+			if item.schemaVersion > is.infoSchema.schemaMetaVersion {
+				// Skip the record.
+				return true
 			}
+			if item.schemaVersion <= is.infoSchema.schemaMetaVersion {
+				ok = !item.tomb
+				pi = item
+				return false
+			}
+			return true
+		})
+	if !ok {
+		return nil, nil, nil
+	}
+
+	tbl, ok := is.TableByID(pi.tableID)
+	if !ok {
+		// something wrong?
+		return nil, nil, nil
+	}
+
+	dbID := tbl.Meta().DBID
+	dbInfo, ok := is.SchemaByID(dbID)
+	if !ok {
+		// something wrong?
+		return nil, nil, nil
+	}
+
+	partInfo := tbl.Meta().GetPartitionInfo()
+	var def *model.PartitionDefinition
+	for i := 0; i < len(partInfo.Definitions); i++ {
+		pdef := &partInfo.Definitions[i]
+		if pdef.ID == partitionID {
+			def = pdef
+			break
 		}
 	}
-	return nil, nil, nil
+
+	return tbl, dbInfo, def
 }
 
 func (is *infoschemaV2) TableExists(schema, table model.CIStr) bool {
@@ -777,6 +833,12 @@ func (b *Builder) applyDropTableV2(diff *model.SchemaDiff, dbInfo *model.DBInfo,
 	// The old DBInfo still holds a reference to old table info, we need to remove it.
 	b.infoSchema.deleteReferredForeignKeys(dbInfo.Name, table.Meta())
 
+	if pi := table.Meta().GetPartitionInfo(); pi != nil {
+		for _, def := range pi.Definitions {
+			b.infoData.pid2tid.Set(partitionItem{def.ID, diff.Version, table.Meta().ID, true})
+		}
+	}
+
 	b.infoData.remove(tableItem{
 		dbName:        dbInfo.Name.L,
 		dbID:          dbInfo.ID,

pkg/infoschema/test/infoschemav2test/BUILD.bazel

@@ -12,6 +12,8 @@ go_test(
         "//pkg/domain",
         "//pkg/infoschema",
         "//pkg/parser/auth",
+        "//pkg/parser/model",
+        "//pkg/table",
         "//pkg/testkit",
         "//pkg/testkit/testsetup",
         "@com_github_stretchr_testify//require",

pkg/infoschema/test/infoschemav2test/v2_test.go

@@ -20,6 +20,8 @@ import (
 	"github.com/pingcap/tidb/pkg/domain"
 	"github.com/pingcap/tidb/pkg/infoschema"
 	"github.com/pingcap/tidb/pkg/parser/auth"
+	"github.com/pingcap/tidb/pkg/parser/model"
+	"github.com/pingcap/tidb/pkg/table"
 	"github.com/pingcap/tidb/pkg/testkit"
 	"github.com/stretchr/testify/require"
 )
@@ -64,3 +66,83 @@ func TestSpecialSchemas(t *testing.T) {
 
 	tk.MustExec("set @@global.tidb_schema_cache_size = default;")
 }
+
+func checkPIDNotExist(t *testing.T, dom *domain.Domain, pid int64) {
+	is := dom.InfoSchema()
+	ptbl, dbInfo, pdef := is.FindTableByPartitionID(pid)
+	require.Nil(t, ptbl)
+	require.Nil(t, dbInfo)
+	require.Nil(t, pdef)
+}
+
+func getPIDForP3(t *testing.T, dom *domain.Domain) (int64, table.Table) {
+	is := dom.InfoSchema()
+	tbl, err := is.TableByName(model.NewCIStr("test"), model.NewCIStr("pt"))
+	require.NoError(t, err)
+	pi := tbl.Meta().GetPartitionInfo()
+	pid := pi.GetPartitionIDByName("p3")
+	ptbl, _, _ := is.FindTableByPartitionID(pid)
+	require.Equal(t, ptbl.Meta().ID, tbl.Meta().ID)
+	return pid, tbl
+}
+
+func TestFindTableByPartitionID(t *testing.T) {
+	store, dom := testkit.CreateMockStoreAndDomain(t)
+	tk := testkit.NewTestKit(t, store)
+	tk.MustExec("use test")
+	tk.MustExec(`create table pt (id int) partition by range (id) (
+partition p0 values less than (10),
+partition p1 values less than (20),
+partition p2 values less than (30),
+partition p3 values less than (40))`)
+
+	pid, tbl := getPIDForP3(t, dom)
+	is := dom.InfoSchema()
+	tbl1, dbInfo, pdef := is.FindTableByPartitionID(pid)
+	require.Equal(t, tbl1.Meta().ID, tbl.Meta().ID)
+	require.Equal(t, dbInfo.Name.L, "test")
+	require.Equal(t, pdef.ID, pid)
+
+	// Test FindTableByPartitionID after dropping a unrelated partition.
+	tk.MustExec("alter table pt drop partition p2")
+	is = dom.InfoSchema()
+	tbl2, dbInfo, pdef := is.FindTableByPartitionID(pid)
+	require.Equal(t, tbl2.Meta().ID, tbl.Meta().ID)
+	require.Equal(t, dbInfo.Name.L, "test")
+	require.Equal(t, pdef.ID, pid)
+
+	// Test FindTableByPartitionID after dropping that partition.
+	tk.MustExec("alter table pt drop partition p3")
+	checkPIDNotExist(t, dom, pid)
+
+	// Test FindTableByPartitionID after adding back the partition.
+	tk.MustExec("alter table pt add partition (partition p3 values less than (35))")
+	checkPIDNotExist(t, dom, pid)
+	pid, _ = getPIDForP3(t, dom)
+
+	// Test FindTableByPartitionID after truncate partition.
+	tk.MustExec("alter table pt truncate partition p3")
+	checkPIDNotExist(t, dom, pid)
+	pid, _ = getPIDForP3(t, dom)
+
+	// Test FindTableByPartitionID after reorganize partition.
+	tk.MustExec(`alter table pt reorganize partition p1,p3 INTO (
+PARTITION p3 VALUES LESS THAN (1970),
+PARTITION p5 VALUES LESS THAN (1980))`)
+	checkPIDNotExist(t, dom, pid)
+	_, _ = getPIDForP3(t, dom)
+
+	// Test FindTableByPartitionID after exchange partition.
+	tk.MustExec("create table nt (id int)")
+	is = dom.InfoSchema()
+	ntbl, err := is.TableByName(model.NewCIStr("test"), model.NewCIStr("nt"))
+	require.NoError(t, err)
+
+	tk.MustExec("alter table pt exchange partition p3 with table nt")
+	is = dom.InfoSchema()
+	ptbl, err := is.TableByName(model.NewCIStr("test"), model.NewCIStr("pt"))
+	require.NoError(t, err)
+	pi := ptbl.Meta().GetPartitionInfo()
+	pid = pi.GetPartitionIDByName("p3")
+	require.Equal(t, pid, ntbl.Meta().ID)
+}