br: backup/restore statistic data in the partition dimension

br/cmd/br/debug.go

@@ -80,7 +80,7 @@ func newCheckSumCommand() *cobra.Command {
 			}
 
 			reader := metautil.NewMetaReader(backupMeta, s, &cfg.CipherInfo)
-			dbs, err := utils.LoadBackupTables(ctx, reader)
+			dbs, err := metautil.LoadBackupTables(ctx, reader)
 			if err != nil {
 				return errors.Trace(err)
 			}
@@ -173,7 +173,7 @@ func newBackupMetaValidateCommand() *cobra.Command {
 				return errors.Trace(err)
 			}
 			reader := metautil.NewMetaReader(backupMeta, s, &cfg.CipherInfo)
-			dbs, err := utils.LoadBackupTables(ctx, reader)
+			dbs, err := metautil.LoadBackupTables(ctx, reader)
 			if err != nil {
 				log.Error("load tables failed", zap.Error(err))
 				return errors.Trace(err)

br/pkg/backup/schema.go

@@ -40,6 +40,7 @@ type schemaInfo struct {
 	totalKvs   uint64
 	totalBytes uint64
 	stats      *util.JSONTable
+	statsIndex []*backuppb.StatsFileIndex
 }
 
 type iterFuncTp func(kv.Storage, func(*model.DBInfo, *model.TableInfo)) error
@@ -147,7 +148,8 @@ func (ss *Schemas) BackupSchemas(
 					}
 				}
 				if statsHandle != nil {
-					if err := schema.dumpStatsToJSON(statsHandle, backupTS); err != nil {
+					statsWriter := metaWriter.NewStatsWriter()
+					if err := schema.dumpStatsToJSON(ctx, statsWriter, statsHandle, backupTS); err != nil {
 						logger.Error("dump table stats failed", logutil.ShortError(err))
 					}
 				}
@@ -209,15 +211,19 @@ func (s *schemaInfo) calculateChecksum(
 	return nil
 }
 
-func (s *schemaInfo) dumpStatsToJSON(statsHandle *handle.Handle, backupTS uint64) error {
+func (s *schemaInfo) dumpStatsToJSON(ctx context.Context, statsWriter *metautil.StatsWriter, statsHandle *handle.Handle, backupTS uint64) error {
 	log.Info("dump stats to json", zap.Stringer("db", s.dbInfo.Name), zap.Stringer("table", s.tableInfo.Name))
-	jsonTable, err := statsHandle.DumpStatsToJSONBySnapshot(
-		s.dbInfo.Name.String(), s.tableInfo, backupTS, true)
-	if err != nil {
+	if err := statsHandle.DumpStatsToJSONBySnapshotFunc(
+		ctx, s.dbInfo.Name.String(), s.tableInfo, backupTS, statsWriter.BackupStats,
+	); err != nil {
 		return errors.Trace(err)
 	}
 
-	s.stats = jsonTable
+	statsFileIndexes, err := statsWriter.BackupStatsDone(ctx)
+	if err != nil {
+		return errors.Trace(err)
+	}
+	s.statsIndex = statsFileIndexes
 	return nil
 }
 

br/pkg/metautil/BUILD.bazel

@@ -2,15 +2,21 @@ load("@io_bazel_rules_go//go:def.bzl", "go_library", "go_test")
 
 go_library(
     name = "metautil",
-    srcs = ["metafile.go"],
+    srcs = [
+        "load.go",
+        "metafile.go",
+        "statsfile.go",
+    ],
     importpath = "github.com/pingcap/tidb/br/pkg/metautil",
     visibility = ["//visibility:public"],
     deps = [
         "//br/pkg/errors",
         "//br/pkg/logutil",
         "//br/pkg/storage",
         "//br/pkg/summary",
+        "//br/pkg/utils",
         "//pkg/parser/model",
+        "//pkg/statistics/handle/types",
         "//pkg/statistics/handle/util",
         "//pkg/tablecodec",
         "//pkg/util/encrypt",
@@ -21,6 +27,7 @@ go_library(
         "@com_github_pingcap_kvproto//pkg/brpb",
         "@com_github_pingcap_kvproto//pkg/encryptionpb",
         "@com_github_pingcap_log//:log",
+        "@org_golang_x_sync//errgroup",
         "@org_uber_go_zap//:zap",
     ],
 )
@@ -29,15 +36,21 @@ go_test(
     name = "metautil_test",
     timeout = "short",
     srcs = [
+        "load_test.go",
         "main_test.go",
         "metafile_test.go",
     ],
     embed = [":metautil"],
     flaky = True,
-    shard_count = 6,
+    shard_count = 8,
     deps = [
         "//br/pkg/mock/storage",
+        "//br/pkg/storage",
+        "//pkg/parser/model",
+        "//pkg/statistics/handle/util",
+        "//pkg/tablecodec",
         "//pkg/testkit/testsetup",
+        "@com_github_golang_protobuf//proto",
         "@com_github_pingcap_kvproto//pkg/brpb",
         "@com_github_pingcap_kvproto//pkg/encryptionpb",
         "@com_github_stretchr_testify//require",

br/pkg/metautil/load.go

@@ -0,0 +1,75 @@
+// 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 metautil
+
+import (
+	"context"
+
+	"github.com/pingcap/errors"
+	"github.com/pingcap/tidb/pkg/parser/model"
+)
+
+// Database wraps the schema and tables of a database.
+type Database struct {
+	Info   *model.DBInfo
+	Tables []*Table
+}
+
+// GetTable returns a table of the database by name.
+func (db *Database) GetTable(name string) *Table {
+	for _, table := range db.Tables {
+		if table.Info.Name.String() == name {
+			return table
+		}
+	}
+	return nil
+}
+
+// LoadBackupTables loads schemas from BackupMeta.
+func LoadBackupTables(ctx context.Context, reader *MetaReader) (map[string]*Database, error) {
+	ch := make(chan *Table)
+	errCh := make(chan error)
+	go func() {
+		if err := reader.ReadSchemasFiles(ctx, ch); err != nil {
+			errCh <- errors.Trace(err)
+		}
+		close(ch)
+	}()
+
+	databases := make(map[string]*Database)
+	for {
+		select {
+		case <-ctx.Done():
+			return nil, ctx.Err()
+		case err := <-errCh:
+			return nil, errors.Trace(err)
+		case table, ok := <-ch:
+			if !ok {
+				close(errCh)
+				return databases, nil
+			}
+			dbName := table.DB.Name.String()
+			db, ok := databases[dbName]
+			if !ok {
+				db = &Database{
+					Info:   table.DB,
+					Tables: make([]*Table, 0),
+				}
+				databases[dbName] = db
+			}
+			db.Tables = append(db.Tables, table)
+		}
+	}
+}

br/pkg/utils/schema_test.go → br/pkg/metautil/load_test.go

@@ -1,6 +1,17 @@
-// Copyright 2020 PingCAP, Inc. Licensed under Apache-2.0.
-
-package utils
+// 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 metautil
 
 import (
 	"context"
@@ -11,7 +22,7 @@ import (
 	"github.com/golang/protobuf/proto"
 	backuppb "github.com/pingcap/kvproto/pkg/brpb"
 	"github.com/pingcap/kvproto/pkg/encryptionpb"
-	"github.com/pingcap/tidb/br/pkg/metautil"
+
 	"github.com/pingcap/tidb/br/pkg/storage"
 	"github.com/pingcap/tidb/pkg/parser/model"
 	"github.com/pingcap/tidb/pkg/statistics/handle/util"
@@ -84,12 +95,12 @@ func TestLoadBackupMeta(t *testing.T) {
 	require.NoError(t, err)
 
 	ctx := context.Background()
-	err = store.WriteFile(ctx, metautil.MetaFile, data)
+	err = store.WriteFile(ctx, MetaFile, data)
 	require.NoError(t, err)
 
 	dbs, err := LoadBackupTables(
 		ctx,
-		metautil.NewMetaReader(
+		NewMetaReader(
 			meta,
 			store,
 			&backuppb.CipherInfo{
@@ -179,12 +190,12 @@ func TestLoadBackupMetaPartionTable(t *testing.T) {
 	require.NoError(t, err)
 
 	ctx := context.Background()
-	err = store.WriteFile(ctx, metautil.MetaFile, data)
+	err = store.WriteFile(ctx, MetaFile, data)
 	require.NoError(t, err)
 
 	dbs, err := LoadBackupTables(
 		ctx,
-		metautil.NewMetaReader(
+		NewMetaReader(
 			meta,
 			store,
 			&backuppb.CipherInfo{
@@ -265,12 +276,12 @@ func BenchmarkLoadBackupMeta64(b *testing.B) {
 		require.NoError(b, err)
 
 		ctx := context.Background()
-		err = store.WriteFile(ctx, metautil.MetaFile, data)
+		err = store.WriteFile(ctx, MetaFile, data)
 		require.NoError(b, err)
 
 		dbs, err := LoadBackupTables(
 			ctx,
-			metautil.NewMetaReader(
+			NewMetaReader(
 				meta,
 				store,
 				&backuppb.CipherInfo{
@@ -297,12 +308,12 @@ func BenchmarkLoadBackupMeta1024(b *testing.B) {
 		require.NoError(b, err)
 
 		ctx := context.Background()
-		err = store.WriteFile(ctx, metautil.MetaFile, data)
+		err = store.WriteFile(ctx, MetaFile, data)
 		require.NoError(b, err)
 
 		dbs, err := LoadBackupTables(
 			ctx,
-			metautil.NewMetaReader(
+			NewMetaReader(
 				meta,
 				store,
 				&backuppb.CipherInfo{
@@ -329,12 +340,12 @@ func BenchmarkLoadBackupMeta10240(b *testing.B) {
 		require.NoError(b, err)
 
 		ctx := context.Background()
-		err = store.WriteFile(ctx, metautil.MetaFile, data)
+		err = store.WriteFile(ctx, MetaFile, data)
 		require.NoError(b, err)
 
 		dbs, err := LoadBackupTables(
 			ctx,
-			metautil.NewMetaReader(
+			NewMetaReader(
 				meta,
 				store,
 				&backuppb.CipherInfo{