From 262327f78d2e39532a4ba9cb6eed7aa539df0990 Mon Sep 17 00:00:00 2001
From: Zhihong Yu <yuzhihong@gmail.com>
Date: Sun, 10 Sep 2023 19:48:11 -0700
Subject: [PATCH] planner: speed up PhysicalTopN#containVirtualColumn (#46812)

close pingcap/tidb#46809
---
 planner/core/task.go | 12 ++++++++----
 1 file changed, 8 insertions(+), 4 deletions(-)

diff --git a/planner/core/task.go b/planner/core/task.go
index 15b7b3a73e664..c90b52ceddee9 100644
--- a/planner/core/task.go
+++ b/planner/core/task.go
@@ -1115,14 +1115,18 @@ func (p *PhysicalTopN) canExpressionConvertedToPB(storeTp kv.StoreType) bool {
 
 // containVirtualColumn checks whether TopN.ByItems contains virtual generated columns.
 func (p *PhysicalTopN) containVirtualColumn(tCols []*expression.Column) bool {
+	tColSet := make(map[int64]struct{}, len(tCols))
+	for _, tCol := range tCols {
+		if tCol.ID > 0 && tCol.VirtualExpr != nil {
+			tColSet[tCol.ID] = struct{}{}
+		}
+	}
 	for _, by := range p.ByItems {
 		cols := expression.ExtractColumns(by.Expr)
 		for _, col := range cols {
-			for _, tCol := range tCols {
+			if _, ok := tColSet[col.ID]; ok {
 				// A column with ID > 0 indicates that the column can be resolved by data source.
-				if tCol.ID > 0 && tCol.ID == col.ID && tCol.VirtualExpr != nil {
-					return true
-				}
+				return true
 			}
 		}
 	}