ddl: pre-split region for partitioned table

ddl/table.go

@@ -72,7 +72,15 @@ func onCreateTable(d *ddlCtx, t *meta.Meta, job *model.Job) (ver int64, _ error)
 		}
 		if atomic.LoadUint32(&EnableSplitTableRegion) != 0 {
 			// TODO: Add restrictions to this operation.
-			go splitTableRegion(d.store, tbInfo.ID)
+			pi := tbInfo.GetPartitionInfo()
+			if pi != nil {
+				// Max partition count is 4096, should we sample and just choose some of the partition to split?
+				for _, def := range pi.Definitions {
+					go splitTableRegion(d.store, def.ID)
+				}
+			} else {
+				go splitTableRegion(d.store, tbInfo.ID)
+			}
 		}
 		// Finish this job.
 		job.FinishTableJob(model.JobStateDone, model.StatePublic, ver, tbInfo)

ddl/table_split_test.go

@@ -26,6 +26,7 @@ import (
 	"github.com/pingcap/tidb/store/mockstore"
 	"github.com/pingcap/tidb/store/tikv"
 	"github.com/pingcap/tidb/tablecodec"
+	"github.com/pingcap/tidb/util/testkit"
 )
 
 type testDDLTableSplitSuite struct{}
@@ -41,21 +42,39 @@ func (s *testDDLTableSplitSuite) TestTableSplit(c *C) {
 	atomic.StoreUint32(&ddl.EnableSplitTableRegion, 1)
 	dom, err := session.BootstrapSession(store)
 	c.Assert(err, IsNil)
+	tk := testkit.NewTestKit(c, store)
+	tk.MustExec("use test")
+	tk.MustExec(`create table t_part (a int key) partition by range(a) (
+		partition p0 values less than (10),
+		partition p1 values less than (20)
+	)`)
 	defer dom.Close()
 	atomic.StoreUint32(&ddl.EnableSplitTableRegion, 0)
 	infoSchema := dom.InfoSchema()
 	c.Assert(infoSchema, NotNil)
 	t, err := infoSchema.TableByName(model.NewCIStr("mysql"), model.NewCIStr("tidb"))
 	c.Assert(err, IsNil)
-	regionStartKey := tablecodec.EncodeTablePrefix(t.Meta().ID)
+	checkRegionStartWithTableID(c, t.Meta().ID, store.(kvStore))
 
-	type kvStore interface {
-		GetRegionCache() *tikv.RegionCache
+	t, err = infoSchema.TableByName(model.NewCIStr("test"), model.NewCIStr("t_part"))
+	c.Assert(err, IsNil)
+	pi := t.Meta().GetPartitionInfo()
+	c.Assert(pi, NotNil)
+	for _, def := range pi.Definitions {
+		checkRegionStartWithTableID(c, def.ID, store.(kvStore))
 	}
+}
+
+type kvStore interface {
+	GetRegionCache() *tikv.RegionCache
+}
+
+func checkRegionStartWithTableID(c *C, id int64, store kvStore) {
+	regionStartKey := tablecodec.EncodeTablePrefix(id)
 	var loc *tikv.KeyLocation
 	for i := 0; i < 10; i++ {
-		cache := store.(kvStore).GetRegionCache()
-		loc, err = cache.LocateKey(tikv.NewBackoffer(context.Background(), 5000), regionStartKey)
+		cache := store.GetRegionCache()
+		loc, err := cache.LocateKey(tikv.NewBackoffer(context.Background(), 5000), regionStartKey)
 		c.Assert(err, IsNil)
 
 		// Region cache may be out of date, so we need to drop this expired region and load it again.