planner: fix TRACE PLAN TARGET = 'estimation' panic when meeting part…

…ition table (#35743)

close #35117

planner/core/optimizer.go

@@ -326,14 +326,20 @@ func refineCETrace(sctx sessionctx.Context) {
 		return i.RowCount < j.RowCount
 	})
 	traceRecords := stmtCtx.OptimizerCETrace
-	is := sctx.GetInfoSchema().(infoschema.InfoSchema)
+	is := sctx.GetDomainInfoSchema().(infoschema.InfoSchema)
 	for _, rec := range traceRecords {
 		tbl, ok := is.TableByID(rec.TableID)
-		if !ok {
-			logutil.BgLogger().Warn("[OptimizerTrace] Failed to find table in infoschema",
-				zap.Int64("table id", rec.TableID))
+		if ok {
+			rec.TableName = tbl.Meta().Name.O
+			continue
+		}
+		tbl, _, _ = is.FindTableByPartitionID(rec.TableID)
+		if tbl != nil {
+			rec.TableName = tbl.Meta().Name.O
+			continue
 		}
-		rec.TableName = tbl.Meta().Name.O
+		logutil.BgLogger().Warn("[OptimizerTrace] Failed to find table in infoschema",
+			zap.Int64("table id", rec.TableID))
 	}
 }
 

statistics/trace_test.go

@@ -100,3 +100,28 @@ func TestTraceCE(t *testing.T) {
 		require.ElementsMatch(t, resultJSON, out[i].Trace)
 	}
 }
+
+func TestTraceCEPartitionTable(t *testing.T) {
+	store, _, clean := testkit.CreateMockStoreAndDomain(t)
+	defer clean()
+	tk := testkit.NewTestKit(t, store)
+	tk.MustExec("use test")
+	tk.MustExec("drop table if exists t")
+	tk.MustExec("create table t(a int, b int, d varchar(10), index idx(a, b)) PARTITION BY RANGE (a) (PARTITION p0 VALUES LESS THAN MAXVALUE);")
+	tk.MustExec(`insert into t values(1, 1, "aaa"),
+		(1, 1, "bbb"),
+		(1, 2, "ccc"),
+		(1, 2, "ddd"),
+		(2, 2, "aaa"),
+		(2, 3, "bbb")`)
+	tk.MustExec("analyze table t")
+	result := tk.MustQuery("trace plan target='estimation' select * from t where a >=1")
+	require.Len(t, result.Rows(), 1)
+	resultStr := result.Rows()[0][0].(string)
+	var resultJSON []*tracing.CETraceRecord
+	err := json.Unmarshal([]byte(resultStr), &resultJSON)
+	require.NoError(t, err)
+	for _, r := range resultJSON {
+		require.Equal(t, "t", r.TableName)
+	}
+}