fixup

pingcap · May 6, 2024 · a3d52d5 · a3d52d5
1 parent 2cf16d8
commit a3d52d5
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 1 deletion.
diff --git a/pkg/planner/core/common_unity.go b/pkg/planner/core/common_unity.go
@@ -150,6 +150,14 @@ func collectUnityInfo(p base.LogicalPlan, o *UnityOutput) {
 			//	result[tableName].Indexes[idxName] = &UnityIndexInfo{}
 			//}
 		}
+		for _, col := range x.Schema().Columns{
+			colName := strings.ToLower(col.OrigName)
+			tName := tblName(col.OrigName)
+			if tName == "" {
+				continue
+			}
+			o.Tables[tableName].Columns[colName] = &UnityColumnInfo{}
+		}
 		for _, expr := range x.allConds {
 			collectColumnFromExpr(expr, o)
 		}
@@ -187,6 +195,10 @@ func collectUnityInfo(p base.LogicalPlan, o *UnityOutput) {
 				collectColumnFromExpr(expr, o)
 			}
 		}
+	case *LogicalProjection:
+		for _, expr := range x.Exprs {
+			collectColumnFromExpr(expr, o)
+		}
 	case *LogicalSort:
 		for _, item := range x.ByItems {
 			collectColumnFromExpr(item.Expr, o)

diff --git a/pkg/planner/core/common_unity_test.go b/pkg/planner/core/common_unity_test.go
@@ -16,9 +16,11 @@ func TestUnity(t *testing.T) {
 	tk.MustExec(`create table t1 (a int, b int, c int, key(a))`)
 	tk.MustExec(`create table t2 (a int, b int, c int, key(a))`)
 	tk.MustExec(`create table t3 (a int, b int, c int, key(a))`)
+	tk.MustExec(`create table t4 (a int, b int, c int, primary key (a))`)
+	tk.MustExec(`create table t5 (a int, b int, c int, primary key (a, b))`)
 	//formatPrint(tk, `explain format='unity' select * from t1, t2 where t1.a=t2.a`)
 	//formatPrint(tk, `explain format='unity' select 1 from t1, t2, t3 where t1.a=t2.a and t2.a=t3.a`)
-	formatPrint(tk, `explain format='unity' select max(a) as ma, count(1) from t1`)
+	formatPrint(tk, `explain format='unity' select a from t4`)
 }
 
 func formatPrint(tk *testkit.TestKit, sql string) {