From 96f9954cf0ef6d49601fdeadbd5672d9e51cd28e Mon Sep 17 00:00:00 2001 From: tangenta Date: Thu, 25 Jul 2024 22:35:11 +0800 Subject: [PATCH 01/13] infoschema: adapt TableByName for query of memory table --- pkg/executor/infoschema_reader.go | 60 ++++++++++++++++--- .../core/memtable_predicate_extractor.go | 47 +++++++++++++++ pkg/util/set/string_set.go | 7 +++ 3 files changed, 106 insertions(+), 8 deletions(-) diff --git a/pkg/executor/infoschema_reader.go b/pkg/executor/infoschema_reader.go index 1d730126d8c12..9ab29ed297a6d 100644 --- a/pkg/executor/infoschema_reader.go +++ b/pkg/executor/infoschema_reader.go @@ -112,7 +112,14 @@ func (e *memtableRetriever) retrieve(ctx context.Context, sctx sessionctx.Contex if !e.initialized { is := sctx.GetInfoSchema().(infoschema.InfoSchema) e.is = is - dbs := is.AllSchemaNames() + + var dbs []model.CIStr + if ex, ok := e.extractor.(plannercore.TableSchemaSelector); ok { + dbs = ex.SelectedSchemaNames() + } + if len(dbs) == 0 { + dbs = is.AllSchemaNames() + } slices.SortFunc(dbs, func(a, b model.CIStr) int { return strings.Compare(a.L, b.L) }) @@ -563,6 +570,46 @@ func (e *memtableRetriever) updateStatsCacheIfNeed() bool { return false } +func getAllTableInfos( + ctx context.Context, + extractor base.MemTablePredicateExtractor, + schema model.CIStr, + is infoschema.InfoSchema, +) ([]*model.TableInfo, error) { + ex := extractor.(plannercore.TableSchemaSelector) + var tables []*model.TableInfo + if ex != nil { + // Find all table infos from where condition. + selectedNames := ex.SelectedTableNames() + tables = make([]*model.TableInfo, 0, len(selectedNames)) + for _, n := range selectedNames { + tbl, err := is.TableByName(ctx, schema, n) + if err != nil { + return nil, errors.Trace(err) + } + tables = append(tables, tbl.Meta()) + } + } + // There is no specified table in where condition. + if len(tables) == 0 { + var err error + tables, err = is.SchemaTableInfos(ctx, schema) + if err != nil { + return nil, errors.Trace(err) + } + if ex != nil { + tmp := tables[:0] + for _, tbl := range tables { + if ex.ContainTableName(tbl.Name.L) { + tmp = append(tmp, tbl) + } + } + tables = tmp + } + } + return tables, nil +} + func (e *memtableRetriever) setDataFromTables(ctx context.Context, sctx sessionctx.Context, schemas []model.CIStr) error { useStatsCache := e.updateStatsCacheIfNeed() checker := privilege.GetPrivilegeManager(sctx) @@ -573,22 +620,19 @@ func (e *memtableRetriever) setDataFromTables(ctx context.Context, sctx sessionc if loc == nil { loc = time.Local } - extractor, ok := e.extractor.(*plannercore.InfoSchemaTablesExtractor) - if ok && extractor.SkipRequest { + tableEx := e.extractor.(*plannercore.InfoSchemaTablesExtractor) + if tableEx != nil && tableEx.SkipRequest { return nil } for _, schema := range schemas { - if ok && extractor.Filter("table_schema", schema.L) { + if tableEx != nil && !tableEx.ContainSchemaName(schema.L) { continue } - tables, err := e.is.SchemaTableInfos(ctx, schema) + tables, err := getAllTableInfos(ctx, e.extractor, schema, e.is) if err != nil { return errors.Trace(err) } for _, table := range tables { - if ok && extractor.Filter("table_name", table.Name.L) { - continue - } collation := table.Collate if collation == "" { collation = mysql.DefaultCollationName diff --git a/pkg/planner/core/memtable_predicate_extractor.go b/pkg/planner/core/memtable_predicate_extractor.go index 7187771bb67c5..81bae1128c6c2 100644 --- a/pkg/planner/core/memtable_predicate_extractor.go +++ b/pkg/planner/core/memtable_predicate_extractor.go @@ -30,6 +30,7 @@ import ( "github.com/pingcap/tidb/pkg/infoschema" "github.com/pingcap/tidb/pkg/kv" "github.com/pingcap/tidb/pkg/parser/ast" + "github.com/pingcap/tidb/pkg/parser/model" "github.com/pingcap/tidb/pkg/parser/mysql" "github.com/pingcap/tidb/pkg/planner/core/base" "github.com/pingcap/tidb/pkg/planner/util" @@ -1896,3 +1897,49 @@ func (e *InfoSchemaTablesExtractor) Filter(colName string, val string) bool { // No need to filter records since no predicate for the column exists. return false } + +var _ TableSchemaSelector = (*InfoSchemaTablesExtractor)(nil) + +type TableSchemaSelector interface { + ContainTableName(table string) bool + ContainSchemaName(schema string) bool + + SelectedTableNames() []model.CIStr + SelectedSchemaNames() []model.CIStr +} + +// ContainTableName returns whether the given table is contained in selection. +func (e *InfoSchemaTablesExtractor) ContainTableName(table string) bool { + return !e.Filter("table_name", table) +} + +// ContainSchemaName returns whether the given schema is contained in selection. +func (e *InfoSchemaTablesExtractor) ContainSchemaName(schema string) bool { + return !e.Filter("table_schema", schema) +} + +// SelectedTableNames gets the table names specified in selection. +func (e *InfoSchemaTablesExtractor) SelectedTableNames() []model.CIStr { + return e.getSchemaObjectNames("table_name") +} + +// SelectedSchemaNames gets the schema names specified in selection. +func (e *InfoSchemaTablesExtractor) SelectedSchemaNames() []model.CIStr { + return e.getSchemaObjectNames("table_schema") +} + +// getSchemaObjectNames gets the schema object names specified in selection of given column name. +func (e *InfoSchemaTablesExtractor) getSchemaObjectNames(colName string) []model.CIStr { + predVals, ok := e.ColPredicates[colName] + if ok && len(predVals) > 0 { + tableNames := make([]model.CIStr, 0, len(predVals)) + predVals.IterateWith(func(n string) { + tableNames = append(tableNames, model.NewCIStr(n)) + }) + slices.SortFunc(tableNames, func(a, b model.CIStr) int { + return strings.Compare(a.L, b.L) + }) + return tableNames + } + return nil +} diff --git a/pkg/util/set/string_set.go b/pkg/util/set/string_set.go index 5a74790971070..e61f46182e390 100644 --- a/pkg/util/set/string_set.go +++ b/pkg/util/set/string_set.go @@ -85,3 +85,10 @@ func (s StringSet) Empty() bool { func (s StringSet) Clear() { maps.Clear(s) } + +// IterateWith iterate items in StringSet and pass it to `fn`. +func (s StringSet) IterateWith(fn func(string)) { + for k := range s { + fn(k) + } +} From 3d7c4a532f5cb9b28504b06ed86a0ec53d81ba6a Mon Sep 17 00:00:00 2001 From: tangenta Date: Thu, 25 Jul 2024 22:47:12 +0800 Subject: [PATCH 02/13] fix linter --- pkg/planner/core/memtable_predicate_extractor.go | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/pkg/planner/core/memtable_predicate_extractor.go b/pkg/planner/core/memtable_predicate_extractor.go index 81bae1128c6c2..451f2b63e949c 100644 --- a/pkg/planner/core/memtable_predicate_extractor.go +++ b/pkg/planner/core/memtable_predicate_extractor.go @@ -1900,6 +1900,8 @@ func (e *InfoSchemaTablesExtractor) Filter(colName string, val string) bool { var _ TableSchemaSelector = (*InfoSchemaTablesExtractor)(nil) +// TableSchemaSelector is used to help determine if a specified table/schema name contained in predicate condition, +// and return all specified table/schema names in predicate condition. type TableSchemaSelector interface { ContainTableName(table string) bool ContainSchemaName(schema string) bool @@ -1908,27 +1910,27 @@ type TableSchemaSelector interface { SelectedSchemaNames() []model.CIStr } -// ContainTableName returns whether the given table is contained in selection. +// ContainTableName returns whether the given table is contained in predicate condition. func (e *InfoSchemaTablesExtractor) ContainTableName(table string) bool { return !e.Filter("table_name", table) } -// ContainSchemaName returns whether the given schema is contained in selection. +// ContainSchemaName returns whether the given schema is contained in predicate condition. func (e *InfoSchemaTablesExtractor) ContainSchemaName(schema string) bool { return !e.Filter("table_schema", schema) } -// SelectedTableNames gets the table names specified in selection. +// SelectedTableNames gets the table names specified in predicate condition. func (e *InfoSchemaTablesExtractor) SelectedTableNames() []model.CIStr { return e.getSchemaObjectNames("table_name") } -// SelectedSchemaNames gets the schema names specified in selection. +// SelectedSchemaNames gets the schema names specified in predicate condition. func (e *InfoSchemaTablesExtractor) SelectedSchemaNames() []model.CIStr { return e.getSchemaObjectNames("table_schema") } -// getSchemaObjectNames gets the schema object names specified in selection of given column name. +// getSchemaObjectNames gets the schema object names specified in predicate condition of given column name. func (e *InfoSchemaTablesExtractor) getSchemaObjectNames(colName string) []model.CIStr { predVals, ok := e.ColPredicates[colName] if ok && len(predVals) > 0 { From 518958cbe2be283452a3eea7e29f399d442d4b90 Mon Sep 17 00:00:00 2001 From: tangenta Date: Fri, 26 Jul 2024 15:42:53 +0800 Subject: [PATCH 03/13] fix table not exists error --- pkg/executor/infoschema_reader.go | 98 +++++++++++++++++++------------ pkg/expression/context/context.go | 2 +- 2 files changed, 60 insertions(+), 40 deletions(-) diff --git a/pkg/executor/infoschema_reader.go b/pkg/executor/infoschema_reader.go index 9ab29ed297a6d..92f7ce964271c 100644 --- a/pkg/executor/infoschema_reader.go +++ b/pkg/executor/infoschema_reader.go @@ -47,6 +47,7 @@ import ( "github.com/pingcap/tidb/pkg/parser/charset" "github.com/pingcap/tidb/pkg/parser/model" "github.com/pingcap/tidb/pkg/parser/mysql" + "github.com/pingcap/tidb/pkg/parser/terror" plannercore "github.com/pingcap/tidb/pkg/planner/core" "github.com/pingcap/tidb/pkg/planner/core/base" "github.com/pingcap/tidb/pkg/privilege" @@ -113,37 +114,42 @@ func (e *memtableRetriever) retrieve(ctx context.Context, sctx sessionctx.Contex is := sctx.GetInfoSchema().(infoschema.InfoSchema) e.is = is - var dbs []model.CIStr - if ex, ok := e.extractor.(plannercore.TableSchemaSelector); ok { - dbs = ex.SelectedSchemaNames() + var getAllSchemas = func() []model.CIStr { + dbs := is.AllSchemaNames() + slices.SortFunc(dbs, func(a, b model.CIStr) int { + return strings.Compare(a.L, b.L) + }) + return dbs } - if len(dbs) == 0 { - dbs = is.AllSchemaNames() - } - slices.SortFunc(dbs, func(a, b model.CIStr) int { - return strings.Compare(a.L, b.L) - }) + var err error switch e.table.Name.O { case infoschema.TableSchemata: + dbs := getAllSchemas() e.setDataFromSchemata(sctx, dbs) case infoschema.TableStatistics: + dbs := getAllSchemas() err = e.setDataForStatistics(ctx, sctx, dbs) case infoschema.TableTables: - err = e.setDataFromTables(ctx, sctx, dbs) + err = e.setDataFromTables(ctx, sctx) case infoschema.TableReferConst: + dbs := getAllSchemas() err = e.setDataFromReferConst(ctx, sctx, dbs) case infoschema.TableSequences: + dbs := getAllSchemas() err = e.setDataFromSequences(ctx, sctx, dbs) case infoschema.TablePartitions: + dbs := getAllSchemas() err = e.setDataFromPartitions(ctx, sctx, dbs) case infoschema.TableClusterInfo: err = e.dataForTiDBClusterInfo(sctx) case infoschema.TableAnalyzeStatus: err = e.setDataForAnalyzeStatus(ctx, sctx) case infoschema.TableTiDBIndexes: + dbs := getAllSchemas() err = e.setDataFromIndexes(ctx, sctx, dbs) case infoschema.TableViews: + dbs := getAllSchemas() err = e.setDataFromViews(ctx, sctx, dbs) case infoschema.TableEngines: e.setDataFromEngines() @@ -152,6 +158,7 @@ func (e *memtableRetriever) retrieve(ctx context.Context, sctx sessionctx.Contex case infoschema.TableCollations: e.setDataFromCollations() case infoschema.TableKeyColumn: + dbs := getAllSchemas() err = e.setDataFromKeyColumnUsage(ctx, sctx, dbs) case infoschema.TableMetricTables: e.setDataForMetricTables() @@ -170,12 +177,14 @@ func (e *memtableRetriever) retrieve(ctx context.Context, sctx sessionctx.Contex case infoschema.TableTiDBHotRegions: err = e.setDataForTiDBHotRegions(ctx, sctx) case infoschema.TableConstraints: + dbs := getAllSchemas() err = e.setDataFromTableConstraints(ctx, sctx, dbs) case infoschema.TableSessionVar: e.rows, err = infoschema.GetDataFromSessionVariables(ctx, sctx) case infoschema.TableTiDBServersInfo: err = e.setDataForServersInfo(sctx) case infoschema.TableTiFlashReplica: + dbs := getAllSchemas() err = e.dataForTableTiFlashReplica(ctx, sctx, dbs) case infoschema.TableTiKVStoreStatus: err = e.dataForTiKVStoreStatus(ctx, sctx) @@ -208,14 +217,18 @@ func (e *memtableRetriever) retrieve(ctx context.Context, sctx sessionctx.Contex case infoschema.TableRunawayWatches: err = e.setDataFromRunawayWatches(sctx) case infoschema.TableCheckConstraints: + dbs := getAllSchemas() err = e.setDataFromCheckConstraints(ctx, sctx, dbs) case infoschema.TableTiDBCheckConstraints: + dbs := getAllSchemas() err = e.setDataFromTiDBCheckConstraints(ctx, sctx, dbs) case infoschema.TableKeywords: err = e.setDataFromKeywords() case infoschema.TableTiDBIndexUsage: + dbs := getAllSchemas() err = e.setDataFromIndexUsage(ctx, sctx, dbs) case infoschema.ClusterTableTiDBIndexUsage: + dbs := getAllSchemas() err = e.setDataForClusterIndexUsage(ctx, sctx, dbs) } if err != nil { @@ -570,47 +583,53 @@ func (e *memtableRetriever) updateStatsCacheIfNeed() bool { return false } -func getAllTableInfos( +func getMatchSchemas( ctx context.Context, extractor base.MemTablePredicateExtractor, - schema model.CIStr, is infoschema.InfoSchema, -) ([]*model.TableInfo, error) { +) []model.CIStr { ex := extractor.(plannercore.TableSchemaSelector) - var tables []*model.TableInfo if ex != nil { - // Find all table infos from where condition. - selectedNames := ex.SelectedTableNames() - tables = make([]*model.TableInfo, 0, len(selectedNames)) - for _, n := range selectedNames { - tbl, err := is.TableByName(ctx, schema, n) - if err != nil { - return nil, errors.Trace(err) - } - tables = append(tables, tbl.Meta()) + if schemas := ex.SelectedSchemaNames(); len(schemas) > 0 { + return schemas } } - // There is no specified table in where condition. - if len(tables) == 0 { - var err error - tables, err = is.SchemaTableInfos(ctx, schema) - if err != nil { - return nil, errors.Trace(err) - } - if ex != nil { - tmp := tables[:0] - for _, tbl := range tables { - if ex.ContainTableName(tbl.Name.L) { - tmp = append(tmp, tbl) + schemas := is.AllSchemaNames() + slices.SortFunc(schemas, func(a, b model.CIStr) int { + return strings.Compare(a.L, b.L) + }) + return schemas +} + +func getMatchTableInfos( + ctx context.Context, + extractor base.MemTablePredicateExtractor, + schema model.CIStr, + is infoschema.InfoSchema, +) ([]*model.TableInfo, error) { + ex := extractor.(plannercore.TableSchemaSelector) + if ex != nil { + if selectedNames := ex.SelectedTableNames(); len(selectedNames) > 0 { + // Find all table infos from where condition. + tables := make([]*model.TableInfo, 0, len(selectedNames)) + for _, n := range selectedNames { + tbl, err := is.TableByName(ctx, schema, n) + if err != nil { + if terror.ErrorEqual(err, infoschema.ErrTableNotExists) { + continue + } + return nil, errors.Trace(err) } + tables = append(tables, tbl.Meta()) } - tables = tmp + return tables, nil } } - return tables, nil + // There is no specified table in where condition. + return is.SchemaTableInfos(ctx, schema) } -func (e *memtableRetriever) setDataFromTables(ctx context.Context, sctx sessionctx.Context, schemas []model.CIStr) error { +func (e *memtableRetriever) setDataFromTables(ctx context.Context, sctx sessionctx.Context) error { useStatsCache := e.updateStatsCacheIfNeed() checker := privilege.GetPrivilegeManager(sctx) @@ -624,11 +643,12 @@ func (e *memtableRetriever) setDataFromTables(ctx context.Context, sctx sessionc if tableEx != nil && tableEx.SkipRequest { return nil } + schemas := getMatchSchemas(ctx, e.extractor, e.is) for _, schema := range schemas { if tableEx != nil && !tableEx.ContainSchemaName(schema.L) { continue } - tables, err := getAllTableInfos(ctx, e.extractor, schema, e.is) + tables, err := getMatchTableInfos(ctx, e.extractor, schema, e.is) if err != nil { return errors.Trace(err) } diff --git a/pkg/expression/context/context.go b/pkg/expression/context/context.go index 8079d4205c40a..8164b6e422d69 100644 --- a/pkg/expression/context/context.go +++ b/pkg/expression/context/context.go @@ -231,6 +231,6 @@ func AssertLocationWithSessionVars(ctxLoc *time.Location, vars *variable.Session stmtLocStr := vars.StmtCtx.TimeZone().String() intest.Assert(ctxLocStr == varsLocStr && ctxLocStr == stmtLocStr, "location mismatch, ctxLoc: %s, varsLoc: %s, stmtLoc: %s", - ctxLoc.String(), ctxLocStr, stmtLocStr, + ctxLoc.String(), varsLocStr, stmtLocStr, ) } From 597acb91d5f4b5e1880f54e2ae9055d123d93fef Mon Sep 17 00:00:00 2001 From: tangenta Date: Fri, 26 Jul 2024 15:47:20 +0800 Subject: [PATCH 04/13] refine code --- pkg/executor/infoschema_reader.go | 14 ++++++++------ pkg/planner/core/memtable_predicate_extractor.go | 13 ------------- 2 files changed, 8 insertions(+), 19 deletions(-) diff --git a/pkg/executor/infoschema_reader.go b/pkg/executor/infoschema_reader.go index 92f7ce964271c..65f05e335991b 100644 --- a/pkg/executor/infoschema_reader.go +++ b/pkg/executor/infoschema_reader.go @@ -584,14 +584,19 @@ func (e *memtableRetriever) updateStatsCacheIfNeed() bool { } func getMatchSchemas( - ctx context.Context, extractor base.MemTablePredicateExtractor, is infoschema.InfoSchema, ) []model.CIStr { ex := extractor.(plannercore.TableSchemaSelector) if ex != nil { if schemas := ex.SelectedSchemaNames(); len(schemas) > 0 { - return schemas + tmp := schemas[:0] + for _, s := range schemas { + if is.SchemaExists(s) { + tmp = append(tmp, s) + } + } + return tmp } } schemas := is.AllSchemaNames() @@ -643,11 +648,8 @@ func (e *memtableRetriever) setDataFromTables(ctx context.Context, sctx sessionc if tableEx != nil && tableEx.SkipRequest { return nil } - schemas := getMatchSchemas(ctx, e.extractor, e.is) + schemas := getMatchSchemas(e.extractor, e.is) for _, schema := range schemas { - if tableEx != nil && !tableEx.ContainSchemaName(schema.L) { - continue - } tables, err := getMatchTableInfos(ctx, e.extractor, schema, e.is) if err != nil { return errors.Trace(err) diff --git a/pkg/planner/core/memtable_predicate_extractor.go b/pkg/planner/core/memtable_predicate_extractor.go index 451f2b63e949c..e4e452dcbd691 100644 --- a/pkg/planner/core/memtable_predicate_extractor.go +++ b/pkg/planner/core/memtable_predicate_extractor.go @@ -1903,23 +1903,10 @@ var _ TableSchemaSelector = (*InfoSchemaTablesExtractor)(nil) // TableSchemaSelector is used to help determine if a specified table/schema name contained in predicate condition, // and return all specified table/schema names in predicate condition. type TableSchemaSelector interface { - ContainTableName(table string) bool - ContainSchemaName(schema string) bool - SelectedTableNames() []model.CIStr SelectedSchemaNames() []model.CIStr } -// ContainTableName returns whether the given table is contained in predicate condition. -func (e *InfoSchemaTablesExtractor) ContainTableName(table string) bool { - return !e.Filter("table_name", table) -} - -// ContainSchemaName returns whether the given schema is contained in predicate condition. -func (e *InfoSchemaTablesExtractor) ContainSchemaName(schema string) bool { - return !e.Filter("table_schema", schema) -} - // SelectedTableNames gets the table names specified in predicate condition. func (e *InfoSchemaTablesExtractor) SelectedTableNames() []model.CIStr { return e.getSchemaObjectNames("table_name") From 1eb039290392f3423dfcad0a40b2a274522653d2 Mon Sep 17 00:00:00 2001 From: tangenta Date: Fri, 26 Jul 2024 16:47:11 +0800 Subject: [PATCH 05/13] fix TestAlterDBPlacement --- pkg/executor/infoschema_reader.go | 4 ++-- .../core/memtable_predicate_extractor.go | 19 ++++++++++++++++++- 2 files changed, 20 insertions(+), 3 deletions(-) diff --git a/pkg/executor/infoschema_reader.go b/pkg/executor/infoschema_reader.go index 65f05e335991b..c820a02afa46d 100644 --- a/pkg/executor/infoschema_reader.go +++ b/pkg/executor/infoschema_reader.go @@ -592,8 +592,8 @@ func getMatchSchemas( if schemas := ex.SelectedSchemaNames(); len(schemas) > 0 { tmp := schemas[:0] for _, s := range schemas { - if is.SchemaExists(s) { - tmp = append(tmp, s) + if n, ok := is.SchemaByName(s); ok { + tmp = append(tmp, n.Name) } } return tmp diff --git a/pkg/planner/core/memtable_predicate_extractor.go b/pkg/planner/core/memtable_predicate_extractor.go index e4e452dcbd691..118a3ac6df142 100644 --- a/pkg/planner/core/memtable_predicate_extractor.go +++ b/pkg/planner/core/memtable_predicate_extractor.go @@ -1826,7 +1826,8 @@ func (e *InfoSchemaTablesExtractor) Extract(ctx base.PlanContext, predicates []expression.Expression, ) (remained []expression.Expression) { var resultSet, resultSet1 set.StringSet - e.colNames = []string{"table_schema", "constraint_schema", "table_name", "constraint_name", "sequence_schema", "sequence_name", "partition_name", "schema_name", "index_name"} + e.colNames = []string{"table_schema", "constraint_schema", "table_name", "constraint_name", + "sequence_schema", "sequence_name", "partition_name", "schema_name", "index_name", "tidb_table_id"} e.ColPredicates = make(map[string]set.StringSet) remained = predicates for _, colName := range e.colNames { @@ -1905,6 +1906,7 @@ var _ TableSchemaSelector = (*InfoSchemaTablesExtractor)(nil) type TableSchemaSelector interface { SelectedTableNames() []model.CIStr SelectedSchemaNames() []model.CIStr + SelectedTableIDs() []int64 } // SelectedTableNames gets the table names specified in predicate condition. @@ -1917,6 +1919,21 @@ func (e *InfoSchemaTablesExtractor) SelectedSchemaNames() []model.CIStr { return e.getSchemaObjectNames("table_schema") } +// SelectedTableIDs get table IDs specified in predicate condition. +func (e *InfoSchemaTablesExtractor) SelectedTableIDs() []int64 { + strs := e.getSchemaObjectNames("tidb_table_id") + tableIDs := make([]int64, 0, len(strs)) + for _, s := range strs { + v, err := strconv.ParseInt(s.L, 10, 64) + if err != nil { + continue + } + tableIDs = append(tableIDs, v) + } + slices.Sort(tableIDs) + return tableIDs +} + // getSchemaObjectNames gets the schema object names specified in predicate condition of given column name. func (e *InfoSchemaTablesExtractor) getSchemaObjectNames(colName string) []model.CIStr { predVals, ok := e.ColPredicates[colName] From c5a97d27e551adebec546ded7239fdeb9300fe93 Mon Sep 17 00:00:00 2001 From: tangenta Date: Mon, 29 Jul 2024 12:08:48 +0800 Subject: [PATCH 06/13] find table info by tidb_table_id --- pkg/executor/infoschema_reader.go | 321 +++++++++++++++++------------- 1 file changed, 183 insertions(+), 138 deletions(-) diff --git a/pkg/executor/infoschema_reader.go b/pkg/executor/infoschema_reader.go index c820a02afa46d..9d68c26188dd1 100644 --- a/pkg/executor/infoschema_reader.go +++ b/pkg/executor/infoschema_reader.go @@ -255,8 +255,13 @@ func (e *memtableRetriever) retrieve(ctx context.Context, sctx sessionctx.Contex return adjustColumns(ret, e.columns, e.table), nil } -func getAutoIncrementID(ctx context.Context, sctx sessionctx.Context, schema model.CIStr, tblInfo *model.TableInfo) (int64, error) { - is := sctx.GetInfoSchema().(infoschema.InfoSchema) +func getAutoIncrementID( + ctx context.Context, + is infoschema.InfoSchema, + sctx sessionctx.Context, + schema model.CIStr, + tblInfo *model.TableInfo, +) (int64, error) { tbl, err := is.TableByName(ctx, schema, tblInfo.Name) if err != nil { return 0, err @@ -590,13 +595,13 @@ func getMatchSchemas( ex := extractor.(plannercore.TableSchemaSelector) if ex != nil { if schemas := ex.SelectedSchemaNames(); len(schemas) > 0 { - tmp := schemas[:0] + ret := schemas[:0] for _, s := range schemas { if n, ok := is.SchemaByName(s); ok { - tmp = append(tmp, n.Name) + ret = append(ret, n.Name) } } - return tmp + return ret } } schemas := is.AllSchemaNames() @@ -614,19 +619,42 @@ func getMatchTableInfos( ) ([]*model.TableInfo, error) { ex := extractor.(plannercore.TableSchemaSelector) if ex != nil { - if selectedNames := ex.SelectedTableNames(); len(selectedNames) > 0 { - // Find all table infos from where condition. - tables := make([]*model.TableInfo, 0, len(selectedNames)) - for _, n := range selectedNames { - tbl, err := is.TableByName(ctx, schema, n) - if err != nil { - if terror.ErrorEqual(err, infoschema.ErrTableNotExists) { - continue - } - return nil, errors.Trace(err) + tables := make([]*model.TableInfo, 0, 8) + // Find all table infos from where condition. + fetchAllTables := true + for _, n := range ex.SelectedTableNames() { + fetchAllTables = false + tbl, err := is.TableByName(ctx, schema, n) + if err != nil { + if terror.ErrorEqual(err, infoschema.ErrTableNotExists) { + continue } - tables = append(tables, tbl.Meta()) + return nil, errors.Trace(err) + } + tables = append(tables, tbl.Meta()) + } + for _, id := range ex.SelectedTableIDs() { + fetchAllTables = false + tbl, ok := is.TableByID(id) + if !ok { + continue } + _, err := is.TableByName(ctx, schema, tbl.Meta().Name) + if err != nil { + continue + } + found := false + for _, t := range tables { + if t.ID == tbl.Meta().ID { + found = true // deduplicate + } + } + if found { + continue + } + tables = append(tables, tbl.Meta()) + } + if !fetchAllTables { return tables, nil } } @@ -634,20 +662,153 @@ func getMatchTableInfos( return is.SchemaTableInfos(ctx, schema) } +func (e *memtableRetriever) setDataFromOneTable( + ctx context.Context, + sctx sessionctx.Context, + loc *time.Location, + checker privilege.Manager, + schema model.CIStr, + table *model.TableInfo, + rows [][]types.Datum, + useStatsCache bool, +) ([][]types.Datum, error) { + collation := table.Collate + if collation == "" { + collation = mysql.DefaultCollationName + } + createTime := types.NewTime(types.FromGoTime(table.GetUpdateTime().In(loc)), mysql.TypeDatetime, types.DefaultFsp) + + createOptions := "" + + if checker != nil && !checker.RequestVerification(sctx.GetSessionVars().ActiveRoles, schema.L, table.Name.L, "", mysql.AllPrivMask) { + return rows, nil + } + pkType := "NONCLUSTERED" + if !table.IsView() { + if table.GetPartitionInfo() != nil { + createOptions = "partitioned" + } else if table.TableCacheStatusType == model.TableCacheStatusEnable { + createOptions = "cached=on" + } + var err error + var autoIncID any + hasAutoIncID, _ := infoschema.HasAutoIncrementColumn(table) + if hasAutoIncID { + autoIncID, err = getAutoIncrementID(ctx, e.is, sctx, schema, table) + if err != nil { + return rows, err + } + } + tableType := "BASE TABLE" + if util.IsSystemView(schema.L) { + tableType = "SYSTEM VIEW" + } + if table.IsSequence() { + tableType = "SEQUENCE" + } + if table.HasClusteredIndex() { + pkType = "CLUSTERED" + } + shardingInfo := infoschema.GetShardingInfo(schema, table) + var policyName any + if table.PlacementPolicyRef != nil { + policyName = table.PlacementPolicyRef.Name.O + } + + var rowCount, avgRowLength, dataLength, indexLength uint64 + if useStatsCache { + if table.GetPartitionInfo() == nil { + err := cache.TableRowStatsCache.UpdateByID(sctx, table.ID) + if err != nil { + return rows, err + } + } else { + // needs to update all partitions for partition table. + for _, pi := range table.GetPartitionInfo().Definitions { + err := cache.TableRowStatsCache.UpdateByID(sctx, pi.ID) + if err != nil { + return rows, err + } + } + } + rowCount, avgRowLength, dataLength, indexLength = cache.TableRowStatsCache.EstimateDataLength(table) + } + + record := types.MakeDatums( + infoschema.CatalogVal, // TABLE_CATALOG + schema.O, // TABLE_SCHEMA + table.Name.O, // TABLE_NAME + tableType, // TABLE_TYPE + "InnoDB", // ENGINE + uint64(10), // VERSION + "Compact", // ROW_FORMAT + rowCount, // TABLE_ROWS + avgRowLength, // AVG_ROW_LENGTH + dataLength, // DATA_LENGTH + uint64(0), // MAX_DATA_LENGTH + indexLength, // INDEX_LENGTH + uint64(0), // DATA_FREE + autoIncID, // AUTO_INCREMENT + createTime, // CREATE_TIME + nil, // UPDATE_TIME + nil, // CHECK_TIME + collation, // TABLE_COLLATION + nil, // CHECKSUM + createOptions, // CREATE_OPTIONS + table.Comment, // TABLE_COMMENT + table.ID, // TIDB_TABLE_ID + shardingInfo, // TIDB_ROW_ID_SHARDING_INFO + pkType, // TIDB_PK_TYPE + policyName, // TIDB_PLACEMENT_POLICY_NAME + ) + rows = append(rows, record) + } else { + record := types.MakeDatums( + infoschema.CatalogVal, // TABLE_CATALOG + schema.O, // TABLE_SCHEMA + table.Name.O, // TABLE_NAME + "VIEW", // TABLE_TYPE + nil, // ENGINE + nil, // VERSION + nil, // ROW_FORMAT + nil, // TABLE_ROWS + nil, // AVG_ROW_LENGTH + nil, // DATA_LENGTH + nil, // MAX_DATA_LENGTH + nil, // INDEX_LENGTH + nil, // DATA_FREE + nil, // AUTO_INCREMENT + createTime, // CREATE_TIME + nil, // UPDATE_TIME + nil, // CHECK_TIME + nil, // TABLE_COLLATION + nil, // CHECKSUM + nil, // CREATE_OPTIONS + "VIEW", // TABLE_COMMENT + table.ID, // TIDB_TABLE_ID + nil, // TIDB_ROW_ID_SHARDING_INFO + pkType, // TIDB_PK_TYPE + nil, // TIDB_PLACEMENT_POLICY_NAME + ) + rows = append(rows, record) + } + return rows, nil +} + func (e *memtableRetriever) setDataFromTables(ctx context.Context, sctx sessionctx.Context) error { useStatsCache := e.updateStatsCacheIfNeed() checker := privilege.GetPrivilegeManager(sctx) var rows [][]types.Datum - createTimeTp := mysql.TypeDatetime loc := sctx.GetSessionVars().TimeZone if loc == nil { loc = time.Local } - tableEx := e.extractor.(*plannercore.InfoSchemaTablesExtractor) - if tableEx != nil && tableEx.SkipRequest { + ex := e.extractor.(*plannercore.InfoSchemaTablesExtractor) + if ex != nil && ex.SkipRequest { return nil } + schemas := getMatchSchemas(e.extractor, e.is) for _, schema := range schemas { tables, err := getMatchTableInfos(ctx, e.extractor, schema, e.is) @@ -655,125 +816,9 @@ func (e *memtableRetriever) setDataFromTables(ctx context.Context, sctx sessionc return errors.Trace(err) } for _, table := range tables { - collation := table.Collate - if collation == "" { - collation = mysql.DefaultCollationName - } - createTime := types.NewTime(types.FromGoTime(table.GetUpdateTime().In(loc)), createTimeTp, types.DefaultFsp) - - createOptions := "" - - if checker != nil && !checker.RequestVerification(sctx.GetSessionVars().ActiveRoles, schema.L, table.Name.L, "", mysql.AllPrivMask) { - continue - } - pkType := "NONCLUSTERED" - if !table.IsView() { - if table.GetPartitionInfo() != nil { - createOptions = "partitioned" - } else if table.TableCacheStatusType == model.TableCacheStatusEnable { - createOptions = "cached=on" - } - var err error - var autoIncID any - hasAutoIncID, _ := infoschema.HasAutoIncrementColumn(table) - if hasAutoIncID { - autoIncID, err = getAutoIncrementID(ctx, sctx, schema, table) - if err != nil { - return err - } - } - tableType := "BASE TABLE" - if util.IsSystemView(schema.L) { - tableType = "SYSTEM VIEW" - } - if table.IsSequence() { - tableType = "SEQUENCE" - } - if table.HasClusteredIndex() { - pkType = "CLUSTERED" - } - shardingInfo := infoschema.GetShardingInfo(schema, table) - var policyName any - if table.PlacementPolicyRef != nil { - policyName = table.PlacementPolicyRef.Name.O - } - - var rowCount, avgRowLength, dataLength, indexLength uint64 - if useStatsCache { - if table.GetPartitionInfo() == nil { - err := cache.TableRowStatsCache.UpdateByID(sctx, table.ID) - if err != nil { - return err - } - } else { - // needs to update all partitions for partition table. - for _, pi := range table.GetPartitionInfo().Definitions { - err := cache.TableRowStatsCache.UpdateByID(sctx, pi.ID) - if err != nil { - return err - } - } - } - rowCount, avgRowLength, dataLength, indexLength = cache.TableRowStatsCache.EstimateDataLength(table) - } - - record := types.MakeDatums( - infoschema.CatalogVal, // TABLE_CATALOG - schema.O, // TABLE_SCHEMA - table.Name.O, // TABLE_NAME - tableType, // TABLE_TYPE - "InnoDB", // ENGINE - uint64(10), // VERSION - "Compact", // ROW_FORMAT - rowCount, // TABLE_ROWS - avgRowLength, // AVG_ROW_LENGTH - dataLength, // DATA_LENGTH - uint64(0), // MAX_DATA_LENGTH - indexLength, // INDEX_LENGTH - uint64(0), // DATA_FREE - autoIncID, // AUTO_INCREMENT - createTime, // CREATE_TIME - nil, // UPDATE_TIME - nil, // CHECK_TIME - collation, // TABLE_COLLATION - nil, // CHECKSUM - createOptions, // CREATE_OPTIONS - table.Comment, // TABLE_COMMENT - table.ID, // TIDB_TABLE_ID - shardingInfo, // TIDB_ROW_ID_SHARDING_INFO - pkType, // TIDB_PK_TYPE - policyName, // TIDB_PLACEMENT_POLICY_NAME - ) - rows = append(rows, record) - } else { - record := types.MakeDatums( - infoschema.CatalogVal, // TABLE_CATALOG - schema.O, // TABLE_SCHEMA - table.Name.O, // TABLE_NAME - "VIEW", // TABLE_TYPE - nil, // ENGINE - nil, // VERSION - nil, // ROW_FORMAT - nil, // TABLE_ROWS - nil, // AVG_ROW_LENGTH - nil, // DATA_LENGTH - nil, // MAX_DATA_LENGTH - nil, // INDEX_LENGTH - nil, // DATA_FREE - nil, // AUTO_INCREMENT - createTime, // CREATE_TIME - nil, // UPDATE_TIME - nil, // CHECK_TIME - nil, // TABLE_COLLATION - nil, // CHECKSUM - nil, // CREATE_OPTIONS - "VIEW", // TABLE_COMMENT - table.ID, // TIDB_TABLE_ID - nil, // TIDB_ROW_ID_SHARDING_INFO - pkType, // TIDB_PK_TYPE - nil, // TIDB_PLACEMENT_POLICY_NAME - ) - rows = append(rows, record) + rows, err = e.setDataFromOneTable(ctx, sctx, loc, checker, schema, table, rows, useStatsCache) + if err != nil { + return errors.Trace(err) } } } From b67366ed4e6defe87eddd18dea0b7cf8d034353f Mon Sep 17 00:00:00 2001 From: tangenta Date: Tue, 30 Jul 2024 10:47:24 +0800 Subject: [PATCH 07/13] fix integration test infoschema/infoschema --- pkg/executor/infoschema_reader.go | 24 +++++++++++-------- .../r/infoschema/infoschema.result | 1 + 2 files changed, 15 insertions(+), 10 deletions(-) diff --git a/pkg/executor/infoschema_reader.go b/pkg/executor/infoschema_reader.go index a5bad146e2909..e1514d1c98039 100644 --- a/pkg/executor/infoschema_reader.go +++ b/pkg/executor/infoschema_reader.go @@ -643,25 +643,29 @@ func getMatchTableInfos( if err != nil { continue } - found := false - for _, t := range tables { - if t.ID == tbl.Meta().ID { - found = true // deduplicate - } - } - if found { - continue - } tables = append(tables, tbl.Meta()) } if !fetchAllTables { - return tables, nil + return deduplicateTableInfos(tables), nil } } // There is no specified table in where condition. return is.SchemaTableInfos(ctx, schema) } +func deduplicateTableInfos(tables []*model.TableInfo) []*model.TableInfo { + idMap := make(map[int64]struct{}, len(tables)) + tmp := tables[:0] + for _, t := range tables { + _, found := idMap[t.ID] + if !found { + idMap[t.ID] = struct{}{} + tmp = append(tmp, t) + } + } + return tmp +} + func (e *memtableRetriever) setDataFromOneTable( ctx context.Context, sctx sessionctx.Context, diff --git a/tests/integrationtest/r/infoschema/infoschema.result b/tests/integrationtest/r/infoschema/infoschema.result index 63e7d48570d3d..e2d2e08c0ab4f 100644 --- a/tests/integrationtest/r/infoschema/infoschema.result +++ b/tests/integrationtest/r/infoschema/infoschema.result @@ -115,6 +115,7 @@ Projection_4 8000.00 root Column#5, Column#10 └─MemTableScan_6 10000.00 root table:TABLES table_schema:["infoschema__infoschema"] select engine, DATA_LENGTH from information_schema.tables where lower(table_name) = 't5' and upper(table_schema) = 'INFOSCHEMA__INFOSCHEMA'; engine DATA_LENGTH +InnoDB 8 explain select engine, DATA_LENGTH from information_schema.tables where (table_name ='t4' or lower(table_name) = 't5') and upper(table_schema) = 'INFOSCHEMA__INFOSCHEMA'; id estRows task access object operator info Projection_4 8000.00 root Column#5, Column#10 From d4b8acb8eac0b809c141318a57fffd1d0a776cf0 Mon Sep 17 00:00:00 2001 From: tangenta Date: Tue, 30 Jul 2024 12:06:09 +0800 Subject: [PATCH 08/13] adapt FindTableByPartitionID for information_schema.partitions --- pkg/executor/infoschema_reader.go | 109 +++++++++++------- .../core/memtable_predicate_extractor.go | 39 +++++-- .../r/infoschema/infoschema.result | 28 +++++ .../t/infoschema/infoschema.test | 7 ++ 4 files changed, 133 insertions(+), 50 deletions(-) diff --git a/pkg/executor/infoschema_reader.go b/pkg/executor/infoschema_reader.go index e1514d1c98039..fe9430efc1142 100644 --- a/pkg/executor/infoschema_reader.go +++ b/pkg/executor/infoschema_reader.go @@ -139,8 +139,7 @@ func (e *memtableRetriever) retrieve(ctx context.Context, sctx sessionctx.Contex dbs := getAllSchemas() err = e.setDataFromSequences(ctx, sctx, dbs) case infoschema.TablePartitions: - dbs := getAllSchemas() - err = e.setDataFromPartitions(ctx, sctx, dbs) + err = e.setDataFromPartitions(ctx, sctx) case infoschema.TableClusterInfo: err = e.dataForTiDBClusterInfo(sctx) case infoschema.TableAnalyzeStatus: @@ -611,46 +610,80 @@ func getMatchSchemas( return schemas } -func getMatchTableInfos( +func getMatchTableInfosForPartitions( ctx context.Context, extractor base.MemTablePredicateExtractor, schema model.CIStr, is infoschema.InfoSchema, ) ([]*model.TableInfo, error) { ex := extractor.(plannercore.TableSchemaSelector) - if ex != nil { - tables := make([]*model.TableInfo, 0, 8) - // Find all table infos from where condition. - fetchAllTables := true - for _, n := range ex.SelectedTableNames() { - fetchAllTables = false - tbl, err := is.TableByName(ctx, schema, n) - if err != nil { - if terror.ErrorEqual(err, infoschema.ErrTableNotExists) { - continue - } - return nil, errors.Trace(err) + if ex == nil || !ex.HasTables() { + // There is no specified table in predicate. + return is.SchemaTableInfos(ctx, schema) + } + tables := make([]*model.TableInfo, 0, 8) + // Find all table infos from predicate. + for _, n := range ex.SelectedTableNames() { + tbl, err := is.TableByName(ctx, schema, n) + if err != nil { + if terror.ErrorEqual(err, infoschema.ErrTableNotExists) { + continue } - tables = append(tables, tbl.Meta()) + return nil, errors.Trace(err) } - for _, id := range ex.SelectedTableIDs() { - fetchAllTables = false - tbl, ok := is.TableByID(id) - if !ok { + tables = append(tables, tbl.Meta()) + } + for _, pid := range ex.SelectedPartitionIDs() { + tbl, db, _ := is.FindTableByPartitionID(pid) + if tbl == nil { + continue + } + if db.Name.L != schema.L { + continue + } + tables = append(tables, tbl.Meta()) + } + return deduplicateTableInfos(tables), nil +} + +func getMatchTableInfos( + ctx context.Context, + extractor base.MemTablePredicateExtractor, + schema model.CIStr, + is infoschema.InfoSchema, +) ([]*model.TableInfo, error) { + ex := extractor.(plannercore.TableSchemaSelector) + if ex == nil || !ex.HasTables() { + // There is no specified table in predicate. + return is.SchemaTableInfos(ctx, schema) + } + tables := make([]*model.TableInfo, 0, 8) + // Find all table infos from predicate. + for _, n := range ex.SelectedTableNames() { + tbl, err := is.TableByName(ctx, schema, n) + if err != nil { + if terror.ErrorEqual(err, infoschema.ErrTableNotExists) { continue } - _, err := is.TableByName(ctx, schema, tbl.Meta().Name) - if err != nil { + return nil, errors.Trace(err) + } + tables = append(tables, tbl.Meta()) + } + for _, id := range ex.SelectedTableIDs() { + tbl, ok := is.TableByID(id) + if !ok { + continue + } + _, err := is.TableByName(ctx, schema, tbl.Meta().Name) + if err != nil { + if terror.ErrorEqual(err, infoschema.ErrTableNotExists) { continue } - tables = append(tables, tbl.Meta()) - } - if !fetchAllTables { - return deduplicateTableInfos(tables), nil + return nil, errors.Trace(err) } + tables = append(tables, tbl.Meta()) } - // There is no specified table in where condition. - return is.SchemaTableInfos(ctx, schema) + return deduplicateTableInfos(tables), nil } func deduplicateTableInfos(tables []*model.TableInfo) []*model.TableInfo { @@ -1227,29 +1260,23 @@ func calcCharOctLength(lenInChar int, cs string) int { return lenInBytes } -func (e *memtableRetriever) setDataFromPartitions(ctx context.Context, sctx sessionctx.Context, schemas []model.CIStr) error { +func (e *memtableRetriever) setDataFromPartitions(ctx context.Context, sctx sessionctx.Context) error { useStatsCache := e.updateStatsCacheIfNeed() checker := privilege.GetPrivilegeManager(sctx) var rows [][]types.Datum createTimeTp := mysql.TypeDatetime - extractor, ok := e.extractor.(*plannercore.InfoSchemaTablesExtractor) - if ok && extractor.SkipRequest { + ex, ok := e.extractor.(*plannercore.InfoSchemaTablesExtractor) + if ok && ex.SkipRequest { return nil } - + schemas := getMatchSchemas(e.extractor, e.is) for _, schema := range schemas { - if ok && extractor.Filter("table_schema", schema.L) { - continue - } - tables, err := e.is.SchemaTableInfos(ctx, schema) + tables, err := getMatchTableInfosForPartitions(ctx, e.extractor, schema, e.is) if err != nil { return errors.Trace(err) } for _, table := range tables { - if ok && extractor.Filter("table_name", table.Name.L) { - continue - } if checker != nil && !checker.RequestVerification(sctx.GetSessionVars().ActiveRoles, schema.L, table.Name.L, "", mysql.SelectPriv) { continue } @@ -1265,7 +1292,7 @@ func (e *memtableRetriever) setDataFromPartitions(ctx context.Context, sctx sess } else { // needs to update needed partitions for partition table. for _, pi := range table.GetPartitionInfo().Definitions { - if ok && extractor.Filter("partition_name", pi.Name.L) { + if ok && ex.Filter("partition_name", pi.Name.L) { continue } err := cache.TableRowStatsCache.UpdateByID(sctx, pi.ID) @@ -1314,7 +1341,7 @@ func (e *memtableRetriever) setDataFromPartitions(ctx context.Context, sctx sess rows = append(rows, record) } else { for i, pi := range table.GetPartitionInfo().Definitions { - if ok && extractor.Filter("partition_name", pi.Name.L) { + if ok && ex.Filter("partition_name", pi.Name.L) { continue } rowCount = cache.TableRowStatsCache.GetTableRows(pi.ID) diff --git a/pkg/planner/core/memtable_predicate_extractor.go b/pkg/planner/core/memtable_predicate_extractor.go index 118a3ac6df142..3f299fd140ab9 100644 --- a/pkg/planner/core/memtable_predicate_extractor.go +++ b/pkg/planner/core/memtable_predicate_extractor.go @@ -1901,29 +1901,50 @@ func (e *InfoSchemaTablesExtractor) Filter(colName string, val string) bool { var _ TableSchemaSelector = (*InfoSchemaTablesExtractor)(nil) -// TableSchemaSelector is used to help determine if a specified table/schema name contained in predicate condition, -// and return all specified table/schema names in predicate condition. +// TableSchemaSelector is used to help determine if a specified table/schema name contained in predicate, +// and return all specified table/schema names in predicate. type TableSchemaSelector interface { - SelectedTableNames() []model.CIStr SelectedSchemaNames() []model.CIStr + + HasTables() bool + SelectedTableNames() []model.CIStr SelectedTableIDs() []int64 + SelectedPartitionIDs() []int64 +} + +// HasTables returns true if there is table names or table IDs specified in predicate. +func (e *InfoSchemaTablesExtractor) HasTables() bool { + _, hasTableName := e.ColPredicates["table_name"] + _, hasTableID := e.ColPredicates["tidb_table_id"] + _, hasPartID := e.ColPredicates["tidb_partition_id"] + return hasTableName || hasTableID || hasPartID } -// SelectedTableNames gets the table names specified in predicate condition. +// SelectedTableNames gets the table names specified in predicate. func (e *InfoSchemaTablesExtractor) SelectedTableNames() []model.CIStr { return e.getSchemaObjectNames("table_name") } -// SelectedSchemaNames gets the schema names specified in predicate condition. +// SelectedSchemaNames gets the schema names specified in predicate. func (e *InfoSchemaTablesExtractor) SelectedSchemaNames() []model.CIStr { return e.getSchemaObjectNames("table_schema") } -// SelectedTableIDs get table IDs specified in predicate condition. +// SelectedTableIDs get table IDs specified in predicate. func (e *InfoSchemaTablesExtractor) SelectedTableIDs() []int64 { strs := e.getSchemaObjectNames("tidb_table_id") - tableIDs := make([]int64, 0, len(strs)) - for _, s := range strs { + return e.parseIDs(strs) +} + +// SelectedPartitionIDs get partitions IDs specified in predicate. +func (e *InfoSchemaTablesExtractor) SelectedPartitionIDs() []int64 { + strs := e.getSchemaObjectNames("tidb_partition_id") + return e.parseIDs(strs) +} + +func (e *InfoSchemaTablesExtractor) parseIDs(ids []model.CIStr) []int64 { + tableIDs := make([]int64, 0, len(ids)) + for _, s := range ids { v, err := strconv.ParseInt(s.L, 10, 64) if err != nil { continue @@ -1934,7 +1955,7 @@ func (e *InfoSchemaTablesExtractor) SelectedTableIDs() []int64 { return tableIDs } -// getSchemaObjectNames gets the schema object names specified in predicate condition of given column name. +// getSchemaObjectNames gets the schema object names specified in predicate of given column name. func (e *InfoSchemaTablesExtractor) getSchemaObjectNames(colName string) []model.CIStr { predVals, ok := e.ColPredicates[colName] if ok && len(predVals) > 0 { diff --git a/tests/integrationtest/r/infoschema/infoschema.result b/tests/integrationtest/r/infoschema/infoschema.result index e2d2e08c0ab4f..d34456770dddb 100644 --- a/tests/integrationtest/r/infoschema/infoschema.result +++ b/tests/integrationtest/r/infoschema/infoschema.result @@ -131,3 +131,31 @@ MemTableScan_5 10000.00 root table:TABLES table_name:["T4","t4"], table_schema:[ select engine, DATA_LENGTH from information_schema.tables where table_name ='t4' and upper(table_name) ='T4' and table_schema = 'infoschema__infoschema'; engine DATA_LENGTH InnoDB 8 +create table pt1(a int primary key, b int) partition by hash(a) partitions 4; +create table pt2(a int primary key, b int) partition by hash(a) partitions 4; +select TABLE_NAME, PARTITION_NAME from information_schema.partitions where table_schema = 'infoschema__infoschema'; +TABLE_NAME PARTITION_NAME +t4 NULL +t5 NULL +pt1 p0 +pt1 p1 +pt1 p2 +pt1 p3 +pt2 p0 +pt2 p1 +pt2 p2 +pt2 p3 +select TABLE_NAME, PARTITION_NAME from information_schema.partitions where table_name = 'pt1'; +TABLE_NAME PARTITION_NAME +pt1 p0 +pt1 p1 +pt1 p2 +pt1 p3 +select TABLE_NAME, PARTITION_NAME from information_schema.partitions where table_name = 'pt2' and table_schema = 'infoschema__infoschema'; +TABLE_NAME PARTITION_NAME +pt2 p0 +pt2 p1 +pt2 p2 +pt2 p3 +select TABLE_NAME, PARTITION_NAME from information_schema.partitions where table_name = 'pt0'; +TABLE_NAME PARTITION_NAME diff --git a/tests/integrationtest/t/infoschema/infoschema.test b/tests/integrationtest/t/infoschema/infoschema.test index faad9b5573054..1c67a74fcf503 100644 --- a/tests/integrationtest/t/infoschema/infoschema.test +++ b/tests/integrationtest/t/infoschema/infoschema.test @@ -58,3 +58,10 @@ select engine, DATA_LENGTH from information_schema.tables where (table_name ='t4 explain select engine, DATA_LENGTH from information_schema.tables where table_name ='t4' and upper(table_name) ='T4' and table_schema = 'infoschema__infoschema'; select engine, DATA_LENGTH from information_schema.tables where table_name ='t4' and upper(table_name) ='T4' and table_schema = 'infoschema__infoschema'; +# TestPartitionsColumn +create table pt1(a int primary key, b int) partition by hash(a) partitions 4; +create table pt2(a int primary key, b int) partition by hash(a) partitions 4; +select TABLE_NAME, PARTITION_NAME from information_schema.partitions where table_schema = 'infoschema__infoschema'; +select TABLE_NAME, PARTITION_NAME from information_schema.partitions where table_name = 'pt1'; +select TABLE_NAME, PARTITION_NAME from information_schema.partitions where table_name = 'pt2' and table_schema = 'infoschema__infoschema'; +select TABLE_NAME, PARTITION_NAME from information_schema.partitions where table_name = 'pt0'; From c011cd4cd67226c24c03e369a151b3dfed4fa4d4 Mon Sep 17 00:00:00 2001 From: tangenta Date: Tue, 30 Jul 2024 12:16:38 +0800 Subject: [PATCH 09/13] fix linter --- pkg/planner/core/memtable_predicate_extractor.go | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/pkg/planner/core/memtable_predicate_extractor.go b/pkg/planner/core/memtable_predicate_extractor.go index 3f299fd140ab9..b6c41a109dfca 100644 --- a/pkg/planner/core/memtable_predicate_extractor.go +++ b/pkg/planner/core/memtable_predicate_extractor.go @@ -1933,16 +1933,16 @@ func (e *InfoSchemaTablesExtractor) SelectedSchemaNames() []model.CIStr { // SelectedTableIDs get table IDs specified in predicate. func (e *InfoSchemaTablesExtractor) SelectedTableIDs() []int64 { strs := e.getSchemaObjectNames("tidb_table_id") - return e.parseIDs(strs) + return parseIDs(strs) } // SelectedPartitionIDs get partitions IDs specified in predicate. func (e *InfoSchemaTablesExtractor) SelectedPartitionIDs() []int64 { strs := e.getSchemaObjectNames("tidb_partition_id") - return e.parseIDs(strs) + return parseIDs(strs) } -func (e *InfoSchemaTablesExtractor) parseIDs(ids []model.CIStr) []int64 { +func parseIDs(ids []model.CIStr) []int64 { tableIDs := make([]int64, 0, len(ids)) for _, s := range ids { v, err := strconv.ParseInt(s.L, 10, 64) From a7719ab6c7b15d04a79dd7fb639c150688d8a34f Mon Sep 17 00:00:00 2001 From: tangenta Date: Tue, 30 Jul 2024 17:43:51 +0800 Subject: [PATCH 10/13] fix integration test infoschema/infoschema --- tests/integrationtest/r/infoschema/infoschema.result | 6 +++--- tests/integrationtest/t/infoschema/infoschema.test | 4 ++-- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/tests/integrationtest/r/infoschema/infoschema.result b/tests/integrationtest/r/infoschema/infoschema.result index d34456770dddb..29624f0be5500 100644 --- a/tests/integrationtest/r/infoschema/infoschema.result +++ b/tests/integrationtest/r/infoschema/infoschema.result @@ -135,7 +135,6 @@ create table pt1(a int primary key, b int) partition by hash(a) partitions 4; create table pt2(a int primary key, b int) partition by hash(a) partitions 4; select TABLE_NAME, PARTITION_NAME from information_schema.partitions where table_schema = 'infoschema__infoschema'; TABLE_NAME PARTITION_NAME -t4 NULL t5 NULL pt1 p0 pt1 p1 @@ -145,7 +144,8 @@ pt2 p0 pt2 p1 pt2 p2 pt2 p3 -select TABLE_NAME, PARTITION_NAME from information_schema.partitions where table_name = 'pt1'; +t4 NULL +select TABLE_NAME, PARTITION_NAME from information_schema.partitions where table_name = 'pt1' and table_schema = 'infoschema__infoschema'; TABLE_NAME PARTITION_NAME pt1 p0 pt1 p1 @@ -157,5 +157,5 @@ pt2 p0 pt2 p1 pt2 p2 pt2 p3 -select TABLE_NAME, PARTITION_NAME from information_schema.partitions where table_name = 'pt0'; +select TABLE_NAME, PARTITION_NAME from information_schema.partitions where table_name = 'pt0' and table_schema = 'infoschema__infoschema'; TABLE_NAME PARTITION_NAME diff --git a/tests/integrationtest/t/infoschema/infoschema.test b/tests/integrationtest/t/infoschema/infoschema.test index 1c67a74fcf503..ba34626d2360c 100644 --- a/tests/integrationtest/t/infoschema/infoschema.test +++ b/tests/integrationtest/t/infoschema/infoschema.test @@ -62,6 +62,6 @@ select engine, DATA_LENGTH from information_schema.tables where table_name ='t4' create table pt1(a int primary key, b int) partition by hash(a) partitions 4; create table pt2(a int primary key, b int) partition by hash(a) partitions 4; select TABLE_NAME, PARTITION_NAME from information_schema.partitions where table_schema = 'infoschema__infoschema'; -select TABLE_NAME, PARTITION_NAME from information_schema.partitions where table_name = 'pt1'; +select TABLE_NAME, PARTITION_NAME from information_schema.partitions where table_name = 'pt1' and table_schema = 'infoschema__infoschema'; select TABLE_NAME, PARTITION_NAME from information_schema.partitions where table_name = 'pt2' and table_schema = 'infoschema__infoschema'; -select TABLE_NAME, PARTITION_NAME from information_schema.partitions where table_name = 'pt0'; +select TABLE_NAME, PARTITION_NAME from information_schema.partitions where table_name = 'pt0' and table_schema = 'infoschema__infoschema'; From 5bd5a55778abd56766074eb8e59987aa8e512664 Mon Sep 17 00:00:00 2001 From: tangenta Date: Tue, 30 Jul 2024 18:41:59 +0800 Subject: [PATCH 11/13] fix integration test infoschema/infoschema --- tests/integrationtest/r/infoschema/infoschema.result | 2 +- tests/integrationtest/t/infoschema/infoschema.test | 4 ++++ 2 files changed, 5 insertions(+), 1 deletion(-) diff --git a/tests/integrationtest/r/infoschema/infoschema.result b/tests/integrationtest/r/infoschema/infoschema.result index 29624f0be5500..2edffeff5e3d9 100644 --- a/tests/integrationtest/r/infoschema/infoschema.result +++ b/tests/integrationtest/r/infoschema/infoschema.result @@ -135,7 +135,6 @@ create table pt1(a int primary key, b int) partition by hash(a) partitions 4; create table pt2(a int primary key, b int) partition by hash(a) partitions 4; select TABLE_NAME, PARTITION_NAME from information_schema.partitions where table_schema = 'infoschema__infoschema'; TABLE_NAME PARTITION_NAME -t5 NULL pt1 p0 pt1 p1 pt1 p2 @@ -145,6 +144,7 @@ pt2 p1 pt2 p2 pt2 p3 t4 NULL +t5 NULL select TABLE_NAME, PARTITION_NAME from information_schema.partitions where table_name = 'pt1' and table_schema = 'infoschema__infoschema'; TABLE_NAME PARTITION_NAME pt1 p0 diff --git a/tests/integrationtest/t/infoschema/infoschema.test b/tests/integrationtest/t/infoschema/infoschema.test index ba34626d2360c..f4e0f993e817f 100644 --- a/tests/integrationtest/t/infoschema/infoschema.test +++ b/tests/integrationtest/t/infoschema/infoschema.test @@ -61,7 +61,11 @@ select engine, DATA_LENGTH from information_schema.tables where table_name ='t4' # TestPartitionsColumn create table pt1(a int primary key, b int) partition by hash(a) partitions 4; create table pt2(a int primary key, b int) partition by hash(a) partitions 4; +-- sorted_result select TABLE_NAME, PARTITION_NAME from information_schema.partitions where table_schema = 'infoschema__infoschema'; +-- sorted_result select TABLE_NAME, PARTITION_NAME from information_schema.partitions where table_name = 'pt1' and table_schema = 'infoschema__infoschema'; +-- sorted_result select TABLE_NAME, PARTITION_NAME from information_schema.partitions where table_name = 'pt2' and table_schema = 'infoschema__infoschema'; +-- sorted_result select TABLE_NAME, PARTITION_NAME from information_schema.partitions where table_name = 'pt0' and table_schema = 'infoschema__infoschema'; From f17548305c469c0fc5b7bfa9f2e9d5985ae5c2c9 Mon Sep 17 00:00:00 2001 From: tangenta Date: Wed, 31 Jul 2024 14:20:40 +0800 Subject: [PATCH 12/13] address comment --- pkg/executor/infoschema_reader.go | 46 ++++++++----------- .../core/memtable_predicate_extractor.go | 5 +- 2 files changed, 21 insertions(+), 30 deletions(-) diff --git a/pkg/executor/infoschema_reader.go b/pkg/executor/infoschema_reader.go index fe9430efc1142..e66009f7d35db 100644 --- a/pkg/executor/infoschema_reader.go +++ b/pkg/executor/infoschema_reader.go @@ -85,6 +85,7 @@ import ( "github.com/tikv/client-go/v2/txnkv/txnlock" pd "github.com/tikv/pd/client/http" "go.uber.org/zap" + "golang.org/x/exp/maps" ) type memtableRetriever struct { @@ -591,8 +592,8 @@ func getMatchSchemas( extractor base.MemTablePredicateExtractor, is infoschema.InfoSchema, ) []model.CIStr { - ex := extractor.(plannercore.TableSchemaSelector) - if ex != nil { + ex, ok := extractor.(plannercore.TableSchemaSelector) + if ok { if schemas := ex.SelectedSchemaNames(); len(schemas) > 0 { ret := schemas[:0] for _, s := range schemas { @@ -616,12 +617,12 @@ func getMatchTableInfosForPartitions( schema model.CIStr, is infoschema.InfoSchema, ) ([]*model.TableInfo, error) { - ex := extractor.(plannercore.TableSchemaSelector) - if ex == nil || !ex.HasTables() { + ex, ok := extractor.(plannercore.TableSchemaSelector) + if !ok || !ex.HasTables() { // There is no specified table in predicate. return is.SchemaTableInfos(ctx, schema) } - tables := make([]*model.TableInfo, 0, 8) + tables := make(map[int64]*model.TableInfo, 8) // Find all table infos from predicate. for _, n := range ex.SelectedTableNames() { tbl, err := is.TableByName(ctx, schema, n) @@ -631,7 +632,8 @@ func getMatchTableInfosForPartitions( } return nil, errors.Trace(err) } - tables = append(tables, tbl.Meta()) + tblInfo := tbl.Meta() + tables[tblInfo.ID] = tblInfo } for _, pid := range ex.SelectedPartitionIDs() { tbl, db, _ := is.FindTableByPartitionID(pid) @@ -641,9 +643,10 @@ func getMatchTableInfosForPartitions( if db.Name.L != schema.L { continue } - tables = append(tables, tbl.Meta()) + tblInfo := tbl.Meta() + tables[tblInfo.ID] = tblInfo } - return deduplicateTableInfos(tables), nil + return maps.Values(tables), nil } func getMatchTableInfos( @@ -652,12 +655,12 @@ func getMatchTableInfos( schema model.CIStr, is infoschema.InfoSchema, ) ([]*model.TableInfo, error) { - ex := extractor.(plannercore.TableSchemaSelector) - if ex == nil || !ex.HasTables() { + ex, ok := extractor.(plannercore.TableSchemaSelector) + if !ok || !ex.HasTables() { // There is no specified table in predicate. return is.SchemaTableInfos(ctx, schema) } - tables := make([]*model.TableInfo, 0, 8) + tables := make(map[int64]*model.TableInfo, 8) // Find all table infos from predicate. for _, n := range ex.SelectedTableNames() { tbl, err := is.TableByName(ctx, schema, n) @@ -667,7 +670,8 @@ func getMatchTableInfos( } return nil, errors.Trace(err) } - tables = append(tables, tbl.Meta()) + tblInfo := tbl.Meta() + tables[tblInfo.ID] = tblInfo } for _, id := range ex.SelectedTableIDs() { tbl, ok := is.TableByID(id) @@ -681,22 +685,10 @@ func getMatchTableInfos( } return nil, errors.Trace(err) } - tables = append(tables, tbl.Meta()) + tblInfo := tbl.Meta() + tables[tblInfo.ID] = tblInfo } - return deduplicateTableInfos(tables), nil -} - -func deduplicateTableInfos(tables []*model.TableInfo) []*model.TableInfo { - idMap := make(map[int64]struct{}, len(tables)) - tmp := tables[:0] - for _, t := range tables { - _, found := idMap[t.ID] - if !found { - idMap[t.ID] = struct{}{} - tmp = append(tmp, t) - } - } - return tmp + return maps.Values(tables), nil } func (e *memtableRetriever) setDataFromOneTable( diff --git a/pkg/planner/core/memtable_predicate_extractor.go b/pkg/planner/core/memtable_predicate_extractor.go index b6c41a109dfca..4eb59b7bf42e5 100644 --- a/pkg/planner/core/memtable_predicate_extractor.go +++ b/pkg/planner/core/memtable_predicate_extractor.go @@ -1815,7 +1815,6 @@ type InfoSchemaTablesExtractor struct { // SkipRequest means the where clause always false, we don't need to request any component SkipRequest bool - colNames []string ColPredicates map[string]set.StringSet } @@ -1826,11 +1825,11 @@ func (e *InfoSchemaTablesExtractor) Extract(ctx base.PlanContext, predicates []expression.Expression, ) (remained []expression.Expression) { var resultSet, resultSet1 set.StringSet - e.colNames = []string{"table_schema", "constraint_schema", "table_name", "constraint_name", + colNames := []string{"table_schema", "constraint_schema", "table_name", "constraint_name", "sequence_schema", "sequence_name", "partition_name", "schema_name", "index_name", "tidb_table_id"} e.ColPredicates = make(map[string]set.StringSet) remained = predicates - for _, colName := range e.colNames { + for _, colName := range colNames { remained, e.SkipRequest, resultSet = e.extractColWithLower(ctx, schema, names, remained, colName) if e.SkipRequest { break From f292a3c2f65fa68c4fa510c63e8cdd40294ce6a5 Mon Sep 17 00:00:00 2001 From: tangenta Date: Wed, 31 Jul 2024 14:56:26 +0800 Subject: [PATCH 13/13] update bazel --- pkg/executor/BUILD.bazel | 1 + 1 file changed, 1 insertion(+) diff --git a/pkg/executor/BUILD.bazel b/pkg/executor/BUILD.bazel index b6a11412a5d39..73c60f967563f 100644 --- a/pkg/executor/BUILD.bazel +++ b/pkg/executor/BUILD.bazel @@ -288,6 +288,7 @@ go_library( "@org_golang_google_grpc//credentials", "@org_golang_google_grpc//credentials/insecure", "@org_golang_google_grpc//status", + "@org_golang_x_exp//maps", "@org_golang_x_sync//errgroup", "@org_uber_go_atomic//:atomic", "@org_uber_go_zap//:zap",