openGauss数据库源码解析系列文章——AI技术（二）

def main(args):
logging.basicConfig(level=logging.INFO)
# 实例化算法模型，模板化模型或DNN模型
model = SQLDiag(args.model, args.csv_file, get_config(args.config_file))
# 训练模型
if args.mode == 'train':
    # fit训练数据，提取模板或特征
        model.fit()
        # 模型保存
        model.save(args.model_path)
    # 预测
elif args.mode == 'predict':
    # 加载模型
        model.load(args.model_path)
        # 标准化预测数据，获取结果
        pred_result = model.transform()
        # 保存输出结果
        ResultSave().save(pred_result, args.predicted_file)
        logging.info('predict result in saved in {}'.format(args.predicted_file))
    # 更新模型
elif args.mode == 'finetune':
        model.fine_tune(args.model_path)
        model.save(args.model_path)

class TemplateModel(AbstractModel):
    # 初始化算法参数
    def __init__(self, params):
        super().__init__(params)
        self.bias = 1e-5
        self.__hash_table = dict(INSERT=dict(), UPDATE=dict(), DELETE=dict(), SELECT=dict(),
                                 OTHER=dict())
        self.time_list_size = params.time_list_size
        self.knn_number = params.knn_number
        self.similarity_algorithm = calc_sql_distance(params.similarity_algorithm)

def fit(self, data):
    # 对每条sql语句按照粗、细粒度进行标准化，生成模板
        for sql, duration_time in data:
            if not self.check_illegal_sql(sql):
                continue
            fine_template, rough_template = get_sql_template(sql)
            sql_prefix = fine_template.split()[0]
            if sql_prefix not in self.__hash_table:
                sql_prefix = 'OTHER'
            # 更新粗粒度模板框架
            if rough_template not in self.__hash_table[sql_prefix]:
                self.__hash_table[sql_prefix][rough_template] = dict()
                self.__hash_table[sql_prefix][rough_template]['info'] = dict()
            # 更新细粒度模板框架
            if fine_template not in self.__hash_table[sql_prefix][rough_template]['info']:
                self.__hash_table[sql_prefix][rough_template]['info'][fine_template] = \
                    dict(time_list=[], count=0, mean_time=0.0, iter_time=0.0)
            # 更新每个细粒度模板的执行时间、迭代时间、sql语句的计数。
            …

self.__hash_table[sql_prefix][rough_template]['info'][fine_template]['count'] += 1
…        
# 基于细粒度模板更新粗粒度模板信息
        for sql_prefix, sql_prefix_info in self.__hash_table.items():
            …

    def transform(self, data):
        predict_time_list = {}
        for sql in data:
            # sql语句不属于'INSERT', 'SELECT', 'UPDATE', 'DELETE', 'CREATE', 'DROP'任何一个，预测时间默认为-1
            if not self.check_illegal_sql(sql):
                predict_time_list[sql] = -1
                continue
            …
                # 若预测的sql所对应的粗粒度模板不存在，执行模板相似度计算方法获取与所有粗粒度模板的相似度
                if rough_template not in self.__hash_table[sql_prefix]:
                    for local_rough_template, local_rough_template_info in self.__hash_table[
                            sql_prefix].items():
                        similarity_info.append(
                            (self.similarity_algorithm(rough_template, local_rough_template), local_rough_template_info['mean_time']))
                # 若预测的sql所对应的细粒度模板不存在，执行模板相似度计算方法获取与所有细粒度模板的相似度
                else:
                    for local_fine_template, local_fine_template_info in \
                            self.__hash_table[sql_prefix][rough_template][
                                'info'].items():
                        similarity_info.append(
                            (self.similarity_algorithm(fine_template, local_fine_template),
                             local_fine_template_info['iter_time']))
                # 基于KNN思想计算sql执行时间
                topn_similarity_info = heapq.nlargest(self.knn_number, similarity_info)
                …

        return predict_time_list

class KerasRegression:
    # 初始化模型参数
    def __init__(self, encoding_dim=1):
        self.model = None
        self.encoding_dim = encoding_dim
# 模型定义
    @staticmethod
    def build_model(shape, encoding_dim):
        from tensorflow.keras import Input, Model
        from tensorflow.keras.layers import Dense
        inputs = Input(shape=(shape,))
        layer_dense1 = Dense(128, activation='relu', kernel_initializer='he_normal')(inputs)
        …
        model = Model(inputs=inputs, outputs=y_pred)
        # 优化器，损失函数
        model.compile(optimizer='adam', loss='mse', metrics=['mae'])
        return model
    # 模型训练
    def fit(self, features, labels, batch_size=128, epochs=300):
        …
        self.model.fit(features, labels, epochs=epochs, batch_size=batch_size, shuffle=True, verbose=2)
    # 模型预测
    def predict(self, features):
        predict_result = self.model.predict(features)
        return predict_result
    # 模型保存
    def save(self, filepath):
        self.model.save(filepath)
    # 模型读取
    def load(self, filepath):
        from tensorflow.keras.models import load_model
        self.model = load_model(filepath)

class DnnModel(AbstractModel, ABC):
    # 初始化算法参数
    def __init__(self, params):
        …
        self.regression = KerasRegression(encoding_dim=1)
        self.data = None
    # 把sql语句转化为vector，如果模型不存在，则直接训练w2v模型，如果模型存在则进行增量训练
    def build_word2vector(self, data):
        self.data = list(data)
        if self.w2v.model:
            self.w2v.update(self.data)
        else:
            self.w2v.fit(self.data)

    def fit(self, data):
        self.build_word2vector(data)
        …
        # 数据归一化
        self.scaler = MinMaxScaler(feature_range=(0, 1))
        self.scaler.fit(labels)
        labels = self.scaler.transform(labels)
        self.regression.fit(features, labels, epochs=self.epoch)

# 利用回归模型预测执行时间
def transform(self, data):
…

python main.py –help

python main.py train -f train.csv --model-path test/

python main.py predict –f predict.csv –model-path test/ --predicted-file test/result.csv

python main.py finetune –f train_new.csv –model-path test/

SuggestedIndex *suggest_index(const char *query_string, _out_ int *len)
{
    ……
    // 对查询语句进行词法和语法解析，获得解析树
List *parse_tree_list = raw_parser(query_string);
…
    // 递归地搜索解析树中的SelectStmt结构
    Node *parsetree = (Node *)lfirst(list_head(parse_tree_list));
    find_select_stmt(parsetree);
   …

    // 依次解析和处理SelectStmt结构中的各个子句部分
    ListCell *item = NULL;

    foreach (item, g_stmt_list) {
        SelectStmt *stmt = (SelectStmt *)lfirst(item);
        /* 处理SelectStmt 结构体中涉及的FROM子句，提取涉及的表，解析和保存这些表中的join关系 */
        parse_from_clause(stmt->fromClause);
        …
        if (g_table_list) {
            // 处理WHERE子句，提取条件表达式中的谓词并添加候选索引，解析和保存其中的join关系
            parse_where_clause(stmt->whereClause);
            // 根据保存的join关系确定驱动表
            determine_driver_table();
            // 处理GROUP子句，如果满足条件，则将其中的谓词添加为候选索引
            if (parse_group_clause(stmt->groupClause, stmt->targetList)) {
                add_index_from_group_order(g_driver_table, stmt->groupClause, stmt->targetList, true);
            /* 处理ORDER子句，如果满足条件，则将其中的谓词添加为候选索引 */
            } else if (parse_order_clause(stmt->sortClause, stmt->targetList)) {
                add_index_from_group_order(g_driver_table, stmt->sortClause, stmt->targetList, false);
            }
            // 如果是多表查询，则根据保存的join关系为被驱动表添加候选索引
            if (g_table_list->length > 1 && g_driver_table) {
                add_index_for_drived_tables();
            }
            /* 对全局变量中的每个table依次进行处理，函数generate_final_index将前述过程生成的候选索引进行字符串拼接，并检查和已存在的索引是否重复 */
            ListCell *table_item = NULL;

            foreach (table_item, g_table_list) {
                TableCell *table = (TableCell *)lfirst(table_item);
                if (table->index != NIL) {
                    Oid table_oid = find_table_oid(query_tree->rtable, table->table_name);
                    if (table_oid == 0) {
                        continue;
                    }
                    generate_final_index(table, table_oid);
                }
            }
            g_driver_table = NULL;
        }
    }
……
    return array;
}

typedef struct hypoIndex {
    Oid oid;           /* 虚拟索引的oid，该oid是唯一的 */
    Oid relid;         /* 涉及的表的oid */
    …
    char *indexname;   /* 虚拟索引名 */

    BlockNumber pages; /* 预估索引使用的磁盘页数 */
    double tuples;     /* 预估索引所涉及的元组数目 */

    /* 索引描述信息 */
    int ncolumns;         /* 涉及的总列数 */
    int nkeycolumns;      /* 涉及的关键列数 */
    … 
    Oid relam;            /* 记录索引操作回调函数的元组的oid, 从pg_am系统表中获取的 */
    … 
} hypoIndex;

hypoIndex *hypo_index_store_parsetree(IndexStmt *node, const char *queryString)
{
……
// 获得创建索引的表的oid
    relid = RangeVarGetRelid(node->relation, AccessShareLock, false);
    ……
    // 对该创建索引的语句进行语法解析
    node = transformIndexStmt(relid, node, queryString);
    ……
    // 新建虚拟索引，该虚拟索引的结构体类型hypoIndex定于位于文件openGauss-server/src/include/dbmind/hypopg_index.h，与索引结构体IndexOptInfo类似
    entry = hypo_newIndex(relid, node->accessMethod, nkeycolumns, ninccolumns, node->options);
    // 根据语法树的解析结果为虚拟索引entry内的各个成员赋值
    PG_TRY();
{
   ……
        entry->unique = node->unique;
        entry->ncolumns = nkeycolumns + ninccolumns;
        entry->nkeycolumns = nkeycolumns;
        ……
    }
    PG_CATCH();
    {        
        hypo_index_pfree(entry);
        PG_RE_THROW();
    }
    PG_END_TRY();
    // 设置虚拟索引的名字
    hypo_set_indexname(entry, indexRelationName.data);
    // 将新建的虚拟索引entry添加到虚拟索引的全局链表hypoIndexes上，该全局变量为节点类型为hypoIndex*的List链表，记录了全部创建过的虚拟索引
    hypo_addIndex(entry);

    return entry;
}
// 该函数被赋值给全局的函数指针get_relation_info_hook，当数据库执行EXPLAIN时，会通过该函数指针跳转到本函数 
void hypo_get_relation_info_hook(PlannerInfo *root, Oid relationObjectId, bool inhparent, RelOptInfo *rel)
{
    /* 判断是否开启GUC参数enable_hypo_index，当SQL语句是EXPLAIN命令时，变量isExplain的值为真 */
    if (u_sess->attr.attr_sql.enable_hypo_index && isExplain) {
        Relation relation;

        relation = heap_open(relationObjectId, AccessShareLock);

        if (relation->rd_rel->relkind == RELKIND_RELATION) {
            ListCell *lc;
            /* 遍历全局变量链表hypoIndexes中的每个创建过的虚拟索引 */
            foreach (lc, hypoIndexes) {
                hypoIndex *entry = (hypoIndex *)lfirst(lc);
                // 判断该虚拟索引和该表是否匹配
                if (hypo_index_match_table(entry, RelationGetRelid(relation))) {
                    // 如果匹配，则将该索引加入该表的indexlist中，indexlist是节点类型为IndexOptInfo的链表，是结构体类型RelOptInfo的成员，记录了表的所有的索引
                    hypo_injectHypotheticalIndex(root, relationObjectId, inhparent, rel, relation, entry);
                }
            }
        }
        heap_close(relation, AccessShareLock);
}
……
}

opengauss=> select * from gs_index_advise('SELECT c_discount from bmsql_customer where c_w_id = 10');     table      |  column  ----------------+---------- bmsql_customer | (c_w_id)(1 row)

CREATE INDEX idx on bmsql_customer(c_w_id);

opengauss=# select * from gs_index_advise('select name, age, sex from t1 where age >= 18 and age < 35 and sex = ''f'';'); table | column-------+------------ t1    | (age, sex)(1 row)

CREATE INDEX idx1 on t1(age, sex);

opengauss=> select * from hypopg_create_index('create index on bmsql_customer(c_w_id)'); indexrelid |              indexname              ------------+-------------------------------------     329726 | <329726>btree_bmsql_customer_c_w_id(1 row)

opengauss=> set enable_hypo_index = on;SET

opengauss=> explain SELECT c_discount from bmsql_customer where c_w_id = 10;                              QUERY PLAN                              -------------------------------------------------------------------- Seq Scan on bmsql_customer  (cost=0.00..52963.06 rows=31224 width=4)   Filter: (c_w_id = 10)(2 rows)

opengauss=> explain SELECT c_discount from bmsql_customer where c_w_id = 10;                              QUERY PLAN                          -------------------------------------------------------------------- [Bypass] Index Scan using <329726>btree_bmsql_customer_c_w_id on bmsql_customer  (cost=0.00..39678.69 rows=31224 width=4)   Index Cond: (c_w_id = 10)(3 rows)

opengauss=> select * from hypopg_display_index();                 indexname                  | indexrelid |     table      |      column      --------------------------------------------+------------+----------------+------------------ <329726>btree_bmsql_customer_c_w_id        |     329726 | bmsql_customer | (c_w_id) <329729>btree_bmsql_customer_c_d_id_c_w_id |     329729 | bmsql_customer | (c_d_id, c_w_id)(2 rows)

opengauss=> select * from hypopg_estimate_size(329730); hypopg_estimate_size ----------------------             15687680(1 row)

opengauss=> select * from hypopg_drop_index(329726); hypopg_drop_index ------------------- t(1 row)

opengauss=> select * from hypopg_reset_index(); hypopg_reset_index --------------------    (1 row)

python index_advisor_workload.py [p PORT] [d DATABASE] [f FILE] [--h HOST] [-U USERNAME] [-W PASSWORD][--max_index_num MAX_INDEX_NUM] [--multi_iter_mode]

python index_advisor_workload.py 6001 opengauss tpcc_log.txt --max_index_num 10 --multi_iter_mode