feat: add ExtractColumnsFromSelectStmt function to improve SQL column extraction and enhance GetSelectivityOfSQLColumns method

Author: iwanghc

sqle/driver/mysql/mysql.go

@@ -14,6 +14,7 @@ import (
 	"github.com/actiontech/sqle/sqle/driver/mysql/plocale"
 	rulepkg "github.com/actiontech/sqle/sqle/driver/mysql/rule"
 	_ "github.com/actiontech/sqle/sqle/driver/mysql/rule/ai"
+	aiutil "github.com/actiontech/sqle/sqle/driver/mysql/rule/ai/util"
 	"github.com/actiontech/sqle/sqle/driver/mysql/session"
 	"github.com/actiontech/sqle/sqle/driver/mysql/util"
 	driverV2 "github.com/actiontech/sqle/sqle/driver/v2"
@@ -639,6 +640,83 @@ func (p *PluginProcessor) GetDriverMetas() (*driverV2.DriverMetas, error) {
 	return metas, nil
 }
 
+func (i *MysqlDriverImpl) GetSelectivityOfSQLColumns(ctx context.Context, sql string) (map[string]map[string]float32, error) {
+	node, err := util.ParseOneSql(sql)
+	if err != nil {
+		return nil, err
+	}
+
+	if _, ok := node.(*ast.SelectStmt); !ok {
+		log.NewEntry().Errorf("get selectivity of sql columns failed, sql is not a select statement, sql: %s", sql)
+		return nil, nil
+	}
+
+	selectVisitor := &util.SelectVisitor{}
+	node.Accept(selectVisitor)
+
+	result := make(map[string]map[string]float32)
+
+	for _, selectNode := range selectVisitor.SelectList {
+		if selectNode.From == nil || selectNode.From.TableRefs == nil {
+			continue
+		}
+
+		// 获取表别名映射关系
+		aliasInfo := aiutil.GetTableAliasInfoFromJoin(selectNode.From.TableRefs)
+		aliasMap := make(map[string]string)
+		allTables := make([]string, 0, len(aliasInfo))
+
+		for _, alias := range aliasInfo {
+			if alias.TableAliasName != "" {
+				aliasMap[alias.TableAliasName] = alias.TableName
+			}
+			allTables = append(allTables, alias.TableName)
+		}
+
+		// 提取列并按表分组
+		tableColumns := util.ExtractColumnsFromSelectStmt(selectNode, aliasMap, allTables)
+
+		// 遍历每个表，获取其列的选择性
+		for tableName, columnSet := range tableColumns {
+			columns := make([]string, 0, len(columnSet))
+			for colName := range columnSet {
+				columns = append(columns, colName)
+			}
+
+			if len(columns) == 0 {
+				continue
+			}
+
+			// 构造 TableName 对象
+			var schemaName string
+			for _, alias := range aliasInfo {
+				if alias.TableName == tableName {
+					schemaName = alias.SchemaName
+					break
+				}
+			}
+			tableNameObj := util.NewTableName(schemaName, tableName)
+
+			columnSelectivityMap, err := i.Ctx.GetSelectivityOfColumns(tableNameObj, columns)
+			if err != nil {
+				log.NewEntry().Errorf("get selectivity of columns failed, table: %s, columns: %v, error: %v", tableName, columns, err)
+				continue
+			}
+
+			if result[tableName] == nil {
+				result[tableName] = make(map[string]float32)
+			}
+			for columnName, selectivity := range columnSelectivityMap {
+				if selectivity > 0 {
+					result[tableName][columnName] = float32(selectivity)
+				}
+			}
+		}
+	}
+
+	return result, nil
+}
+
 func (p *PluginProcessor) Open(l *logrus.Entry, cfg *driverV2.Config) (driver.Plugin, error) {
 	return NewInspect(l, cfg)
 }

sqle/driver/mysql/util/parser_helper.go

@@ -931,3 +931,106 @@ func ConvertAliasToTable(alias string, tables []*ast.TableSource) (*ast.TableNam
 	}
 	return nil, errors.New("can not find table")
 }
+
+// TableColumnMap 表示按表分组的列名集合
+type TableColumnMap map[string]map[string]struct{}
+
+// ExtractColumnsFromSelectStmt 从 SELECT 语句中提取列，并按表分组
+// 参数：
+//   - selectStmt: SELECT 语句节点
+//   - aliasMap: 表别名到实际表名的映射
+//   - allTables: 所有涉及的表名列表（用于处理无表前缀的列）
+//
+// 返回：按表名分组的列名集合
+func ExtractColumnsFromSelectStmt(selectStmt *ast.SelectStmt, aliasMap map[string]string, allTables []string) TableColumnMap {
+	tableColumns := make(TableColumnMap)
+
+	// 收集 SELECT 列表中的所有列别名
+	selectAliases := make(map[string]struct{})
+	if selectStmt.Fields != nil {
+		for _, field := range selectStmt.Fields.Fields {
+			if field.AsName.L != "" {
+				selectAliases[field.AsName.L] = struct{}{}
+			}
+		}
+	}
+
+	// 辅助函数：从表达式中提取列并按表分组
+	extractColumnsFromExpr := func(expr ast.Node, skipAliases bool) {
+		if expr == nil {
+			return
+		}
+		columnVisitor := &ColumnNameVisitor{}
+		expr.Accept(columnVisitor)
+
+		for _, colExpr := range columnVisitor.ColumnNameList {
+			if colExpr.Name == nil {
+				continue
+			}
+
+			// 如果需要跳过别名且当前列名是一个别名，则跳过
+			if skipAliases {
+				if _, isAlias := selectAliases[colExpr.Name.Name.L]; isAlias && colExpr.Name.Table.L == "" {
+					continue
+				}
+			}
+
+			var targetTableName string
+
+			// 如果列有表前缀（可能是别名或实际表名）
+			if colExpr.Name.Table.L != "" {
+				// 先尝试从别名映射中查找
+				if actualTable, exists := aliasMap[colExpr.Name.Table.L]; exists {
+					targetTableName = actualTable
+				} else {
+					// 如果不是别名，就当作实际表名
+					targetTableName = colExpr.Name.Table.L
+				}
+			}
+
+			if targetTableName != "" {
+				if tableColumns[targetTableName] == nil {
+					tableColumns[targetTableName] = make(map[string]struct{})
+				}
+				tableColumns[targetTableName][colExpr.Name.Name.L] = struct{}{}
+			} else {
+				// 没有表前缀的列，可能属于任何表
+				// 在多表查询中，尝试将该列添加到所有表
+				for _, tableName := range allTables {
+					if tableColumns[tableName] == nil {
+						tableColumns[tableName] = make(map[string]struct{})
+					}
+					tableColumns[tableName][colExpr.Name.Name.L] = struct{}{}
+				}
+			}
+		}
+	}
+
+	// 从 SELECT Fields 提取列（包括聚合函数内的列）
+	if selectStmt.Fields != nil {
+		for _, field := range selectStmt.Fields.Fields {
+			extractColumnsFromExpr(field.Expr, false)
+		}
+	}
+
+	// 从 WHERE 条件提取列
+	if selectStmt.Where != nil {
+		extractColumnsFromExpr(selectStmt.Where, false)
+	}
+
+	// 从 GROUP BY 提取列（需要跳过别名引用）
+	if selectStmt.GroupBy != nil {
+		for _, item := range selectStmt.GroupBy.Items {
+			extractColumnsFromExpr(item.Expr, true)
+		}
+	}
+
+	// 从 HAVING 提取列
+	if selectStmt.Having != nil {
+		extractColumnsFromExpr(selectStmt.Having.Expr, false)
+	}
+
+	// 注意：不从 ORDER BY 提取，因为可能包含别名引用
+
+	return tableColumns
+}