⭐OpenGauss Database Source Code Analysis Series Articles——AI Query Time Prediction⭐

The previous article introduced the relevant content of " 8.5 Index Collection, Prediction, and Anomaly Detection ". In this article, we introduce the relevant exciting content of "8.6 AI Query Time Prediction".

8.6 AI Query Time Prediction

The "Slow SQL Discovery" feature was introduced earlier. The typical scenario of this feature is the inspection before the launch of new services, and the input source is the SQL pipeline data collected in advance. The slow SQL discovery function is mainly used in the batch inspection of multiple SQL statements, which requires that the SQL statements have been executed before, so the results given are mainly qualitative, and in some scenarios it may be difficult to meet the user's requirements for evaluation accuracy.
Therefore, in order to make up for the shortcomings of the above scenarios, meet the user's more accurate SQL time prediction requirements, and pave the way for the AI ​​​​optimizer, the functions described in this chapter are realized.
Due to the complex nature of actual business scenarios, the existing static cost estimation models for databases are often inaccurate in statistical results, and some paths with poor execution plans are selected. Therefore, for the above complex scenarios, the cost estimation model of the database is required to have the ability of self-update. The main function of this feature is to estimate the query time and cardinality of the currently executed SQL statement based on the historical data of the query statement.

8.6.1 Usage Scenarios

The premise of AI query analysis is to obtain the execution plan. First of all, it is necessary to collect the actual query plan of complex queries (including plan structure, operator type, related data sources, filter conditions, etc.), the actual execution time of each operator node, the cost estimated by the optimizer, and the actual number of rows returned during query execution according to user needs. , the number of rows estimated by the optimizer, the number of SMP concurrent threads, and other information. Record it in the data table, and perform persistent management including regular data invalidation cleaning.
This function is mainly divided into two aspects, one is row number estimation, and the other is query prediction. The former is the prerequisite for the latter to predict whether it is good or bad. At present, openGauss estimates the result set size of each layer of the execution plan based on online learning, which is only for display and does not affect the generation of the execution plan. Follow-up can help the optimizer to estimate the result set more accurately, so as to obtain a better execution plan.
At the current stage, this demand will provide system functions for prediction, and add it to explain for actual comparison and verification.

8.6.2 Existing technologies

At present, in the field of AI4DB, there are many attempts to estimate the number of rows and predict the query delay based on machine learning.

1. Traditional method

As database optimizer expert Guy Lohman said in the blog Is query optimization a “solved” problem, the “Achilles heel” of traditional database query performance prediction is the estimation of the size of the intermediate result set. For row number estimation, the traditional method of row number estimation based on statistical information is mainly based on three types of assumptions.
(1) Data independent distribution assumption.
(2) Uniform distribution assumption.
(3) Primary and foreign key assumptions.
However, in actual scenarios, data often has certain correlation and inclination. At this time, the above assumptions may be broken, resulting in errors of several orders of magnitude in traditional database optimizers when estimating the size of intermediate result sets for multi-table joins.
Since 2000, statistical methods represented by sampling-based estimation, sampling-based kernel density function estimation, and multi-column histograms have been proposed to solve the estimation problems caused by data correlation. However, these methods all have a common problem, that is, the model cannot be incrementally maintained, and collecting these additional statistical information will increase the huge database maintenance overhead. It has greatly improved the accuracy rate, but it has not been widely adopted by major database manufacturers.
Traditional performance prediction methods mainly rely on cost models, which have obvious disadvantages in the following aspects.
(1) Accuracy: With the continuous evolution of the underlying hardware architecture and optimization technology, the complexity of the actual performance prediction model is far from being modeled by a linear model.
(2) Scalability: The development cost of the cost model is relatively high, and it cannot optimize the user's specific scenarios comprehensively.
(3) Calibrability: The flexibility of the cost model is limited to the coefficients used in the linear addition of each resource dimension, as well as some penalty costs. The flexibility is poor, and it is difficult for users to calibrate in actual use.
(4) Timeliness: The cost model relies on the collection and use of statistical information. At present, there is no incremental maintenance method, which leads to long-term failure of statistical information in scenarios with high data mobility.

2. Machine learning methods

The advantages of machine learning models in the dimensions of model complexity, calibrability, and incremental maintainability can make up for the shortcomings of traditional optimizer cost models. Machine learning-based query performance prediction has gradually become mainstream research in database academia and industry one of the directions.
In addition to the related methods introduced in the slow SQL discovery section in Section 8.3 above, the Learned Cost Estimator model of Tsinghua University is based on Multi-task Learning and the Word-Embedding method of character conditions to further improve the prediction accuracy.
So far, although the machine learning method has achieved a high accuracy rate from the experimental results, the continuous data distribution changes in real business scenarios put forward requirements for the online learning ability of the model. openGauss adopts a data-driven online learning model, continuously collects historical job performance information through the kernel, and uses the R-LSTM (recursive long short term memory, recursive long short term memory network) model on the AI ​​Engine side to determine the operator-level query delay. and intermediate result set sizes for prediction.

8.6.3 Implementation Principle

Figure 8-15 Schematic diagram of AI query performance prediction architecture

In general, query performance prediction consists of two parts: the database kernel side and the AI ​​Engine side, as shown in Figure 8-15.
(1) In addition to providing basic database functions, the database kernel side also needs to collect and persist historical data, and send HTTPS requests to the AI ​​Engine side through curl.
(2) AI Engine provides interfaces such as model training, execution prediction, and model management. The server based on the Flask framework accepts HTTPS requests. The process is shown in Figure 8-16.

Figure 8-16 Schematic diagram of the relationship between the database kernel and the AI ​​Engine process

After enabling the parameters related to data collection (which may affect performance by about 5%, depending on the actual business load), historical performance data is persistently collected in the system tables of the database for model training.
Before model training, users need to configure model parameters (see 8.6.5 Usage Example for details). After the user training command is issued, the kernel process will send a configure request to the AI ​​Engine side to initialize the machine learning model. Figure 8-17 shows the sequence of the configure process.

Figure 8-17 Sequence diagram of configure process

After the model configuration is successful, the kernel process sends a train request to the AI ​​Engine side to trigger training, as shown in Figure 8-18.

Figure 8-18 Sequence diagram of train process After the model is trained, the user issues a prediction instruction, and the database will first send a setup request to the AI ​​Engine side for model loading. After the loading is successful, it will send a predict request to obtain the prediction result, as shown in Figure 8-19 Show. Figure 8-19 Sequence diagram of the complete process of model prediction, which is divided into two stages: setup and predict. This feature architecture supports multiple models. Currently, the R-LSTM model has been implemented. The model architecture is shown in Figure 8-20. In the plan, the execution order of operators also affects the performance of operators. Based on this feature, we use the LSTM neural network model to learn the meaningful dependencies between operators in the plan, and target the model structure, loss function, optimization algorithm, etc. according to the row number/time prediction scenario Optimized to improve the accuracy of learning and prediction in this scenario. Input: query plan tree, operator type on each node, corresponding table name, column name, and filter conditions. Output: number of rows, startup time, total time, Peak Memory. In the encoding stage, each plan node is encoded into a fixed length and connected into a sequence as the feature value of the input LSTM neural network. LSTM has a chain network composed of multiple repeated neural network modules, and each module has three functions to determine which information in the historical sequence will be passed to the network module of the next sequence. The output value h_t of the last module is the prediction result returned by the model.

Among them, Хt is the input of the current timing module, ht-₁ is the input of the current timing module, which is the output information of the previous timing, use the sigmoid (σ) function to get the part of Οt to be output in the current cell state; Ct represents all historical timing The retained information is processed by the tanh function and multiplied by the current state output information Οt to obtain the output ht of this state, and the prediction result of the one-dimensional vector [startup time, total time, cardinality] with three elements is compared with the real data , using ratio-error to calculate the loss function of the model.

Figure 8-20 Model architecture diagram

8.6.4 Key source code analysis

1. Project structure

The main file path involved on the AI ​​Engine side is openGauss-server/src/gausskernel/dbmind/tools/predictor, and its file structure is shown in Table 8-13.

Table 8-13 AI Engine file structure

file structure	illustrate
install	Deployment required file path
install/ca_ext.txt	certificate profile
install/requirements-gpu.txt	Use GPU (graphics processing unit, graphics processor) to train the list of dependent libraries
install/requirements.txt	List of dependent libraries used for CPU training
install/ssl.sh	Certificate generation script
python	project code path
python/certs.py	encrypted communication
python/e_log	syslog path
python/log	Model training log path
python/log.conf	configuration file
python/model.py	machine learning model
python/run.py	server main function
python/saved_models	Model training checkpoint
python/settings.py	Project configuration file
python/uploads	Curl transfer file storage path

The file path mainly involved in the kernel side is openGauss-server/src/gausskernel/optimizer/util/learn, and its file structure is shown in Table 8-14.

Table 8-14 Main file structure of the kernel side

file structure	illustrate
comm.cpp	Communication layer code implementation
encoding.cpp	data encoding
ml_model.cpp	Generic Model Call Interface
plan_tree_model.cpp	Tree Model Call Interface

2. Training process

The model training interface on the kernel side is implemented through the ModelTrainInternal function. The key parts of this function are as follows:

static void ModelTrainInternal(const char* templateName, const char* modelName, ModelAccuracy** mAcc)
{
  …
    /* 对于树形模型调用对应的训练接口 */
    char* trainResultJson = TreeModelTrain(modelinfo, labels);
    /* 解析返回结果 */
    …
    ModelTrainInfo* info = GetModelTrainInfo(jsonObj);
    cJSON_Delete(jsonObj);
    /* 更新模型信息 */
    Relation modelRel = heap_open(OptModelRelationId, RowExclusiveLock);
   …
    UpdateTrainRes(values, datumsMax, datumsAcc, nLabel, mAcc, info, labels);

    HeapTuple modelTuple = SearchSysCache1(OPTMODEL, CStringGetDatum(modelName));
   …
    HeapTuple newTuple = heap_modify_tuple(modelTuple, RelationGetDescr(modelRel), values, nulls, replaces);
    simple_heap_update(modelRel, &newTuple->t_self, newTuple);
CatalogUpdateIndexes(modelRel, newTuple);
…
}

The tree model training interface on the kernel side is implemented through the TreeModelTrain function, and the core code is as follows:

char* TreeModelTrain(Form_gs_opt_model modelinfo, char* labels)
{
    char* filename = (char*)palloc0(sizeof(char) * MAX_LEN_TEXT);
    char* buf = NULL;
    /* configure阶段 */
    ConfigureModel(modelinfo, labels, &filename);

    /* 将编码好的数据写入临时文件 */
    SaveDataToFile(filename);

    /* Train阶段 */
    buf = TrainModel(modelinfo, filename);
    return buf;
}

The URI of the web service configured on the AI ​​Engine side is /configure, and the URI of the training phase is /train. The following code snippet shows the training process.

  def fit(self, filename):
        keras.backend.clear_session()
        set_session(self.session)
        with self.graph.as_default():
            # 根据模型入参和出参维度变化情况，判断是否需要初始化模型
            feature, label, need_init = self.parse(filename) 
            os.environ['CUDA_VISIBLE_DEVICES'] = '0'
            epsilon = self.model_info.make_epsilon()
            if need_init: # 冷启动训练
                epoch_start = 0
                self.model = self._build_model(epsilon)
            else: # 增量训练
                epoch_start = int(self.model_info.last_epoch)
                ratio_error = ratio_error_loss_wrapper(epsilon)
                ratio_acc_2 = ratio_error_acc_wrapper(epsilon, 2)
                self.model = load_model(self.model_info.model_path,
                                        custom_objects={'ratio_error': ratio_error, 'ratio_acc': ratio_acc_2})
            self.model_info.last_epoch = int(self.model_info.max_epoch) + epoch_start
            self.model_info.dump_dict()
            log_path = os.path.join(settings.PATH_LOG, self.model_info.model_name + '_log.json')
            if not os.path.exists(log_path):
                os.mknod(log_path, mode=0o600)
            # 训练日志记录回调函数
            json_logging_callback = LossHistory(log_path, self.model_info.model_name, self.model_info.last_epoch)
            # 数据分割
            X_train, X_val, y_train, y_val = \
                train_test_split(feature, label, test_size=0.1)
            # 模型训练
            self.model.fit(X_train, y_train, epochs=self.model_info.last_epoch,
                           batch_size=int(self.model_info.batch_size), validation_data=(X_val, y_val),
                           verbose=0, initial_epoch=epoch_start, callbacks=[json_logging_callback])
            # 记录模型checkpoint
            self.model.save(self.model_info.model_path)
            val_pred = self.model.predict(X_val)
            val_re = get_ratio_errors_general(val_pred, y_val, epsilon)
            self.model_logger.debug(val_re)
            del self.model
            return val_re

3. Forecasting process

The model prediction process on the kernel side is mainly implemented through the ModelPredictInternal function. The tree model prediction process is implemented through the TreeModelPredict function. The tree model prediction process on the kernel side will take up some signaling to communicate with the AI ​​Engine. The communication process is as follows:

char* TreeModelPredict(const char* modelName, char* filepath, const char* ip, int port)
{
    …
    if (!TryConnectRemoteServer(conninfo, &buf)) {
        DestroyConnInfo(conninfo);
        ParseResBuf(buf, filepath, "AI engine connection failed.");
        return buf;
    }

    switch (buf[0]) {
        case '0': {
            ereport(NOTICE, (errmodule(MOD_OPT_AI), errmsg("Model setup successfully.")));
            break;
        }
        case 'M': {
            ParseResBuf(buf, filepath, "Internal error: missing compulsory key.");
            break;
        }
…
    }
    /* Predict阶段 */
    …
    if (!TryConnectRemoteServer(conninfo, &buf)) {
        ParseResBuf(buf, filepath, "AI engine connection failed.");
        return buf;
    }
    switch (buf[0]) {
        case 'M': {
            ParseResBuf(buf, filepath, "Internal error: fail to load the file to predict.");
            break;
        }
        case 'S': {
            ParseResBuf(buf, filepath, "Internal error: session is not loaded, model setup required.");
            break;
        }
        default: {
            break;
        }
    }
    return buf;
}

The web interface of the Setup process on the AI ​​Engine side is /model_setup, and the web interface of the prediction phase is /predict, and their protocols are both Post.

4. Data encoding

Data encoding is divided into the following two dimensions.
(1) Operator dimension: includes the attributes of each execution plan operator, as shown in Table 8-15.

Table 8-15 Operator dimensions

attribute name	meaning	coding strategy
Optname	operator type	One-hot
Orientation	return tuple storage format	One-hot
Strategy	logical attribute	One-hot
Options	physical properties	One-hot
which ones	predicate	hash
Projection	return projected column	hash

(2) Planning dimension.
For each operator, in addition to its inherent attributes, openGauss also records the query id, plan node id, and parent node id. In the training/prediction phase, use these information to reconstruct the operator information into a tree plan structure, And the sub-plan tree can be constructed recursively for data augmentation, thereby improving the generalization ability of the model. The tree data structure is shown in Figure 8-21.

Figure 8-21 Schematic diagram of tree data structure The tree data encoding on the kernel side is realized through the GetOPTEncoding function.

5. Model structure

For model analysis, training and prediction of AI Engine, see chapter 8.6.4. The following code shows the structure of the model.

class RnnModel():
    def _build_model(self, epsilon):
        model = Sequential()
        model.add(LSTM(units=int(self.model_info.hidden_units), return_sequences=True, input_shape=(None, int(self.model_info.feature_length))))
        model.add(LSTM(units=int(self.model_info.hidden_units), return_sequences=False))
        model.add(Dense(units=int(self.model_info.hidden_units), activation='relu'))
        model.add(Dense(units=int(self.model_info.hidden_units), activation='relu'))
        model.add(Dense(units=int(self.model_info.label_length), activation='sigmoid'))
        optimizer = keras.optimizers.Adadelta(lr=float(self.model_info.learning_rate), rho=0.95)
        ratio_error = ratio_error_loss_wrapper(epsilon)
        ratio_acc_2 = ratio_error_acc_wrapper(epsilon, 2)
        model.compile(loss=ratio_error, metrics=[ratio_acc_2], optimizer=optimizer)
        return model

The loss function of AI Engine uses ratio error (qerror is used in some literatures). Compared with MRE and MSE, the advantage of this loss function is that it can equally punish overestimation and underestimation. The formula is: ε is declared as

performance The infinitesimal value of the predicted value to prevent the denominator from being 0.

8.6.5 Example of use

An example of using the AI ​​query time prediction function is as follows.
① Define the performance prediction model, the code is as follows:

INSERT INTO gs_opt_model VALUES(‘rlstm’, ‘model_name’, ‘host_ip’, ‘port’);

② Start data collection through GUC parameters, configure the parameter list, the code is as follows:

enable_resource_track = on;
enable_resource_record = on;

③ Encode the training data, the code is as follows:

SELECT gather_encoding_info('db_name');

④ Calibration model, the code is as follows:

SELECT model_train_opt('template_name', 'model_name');

⑤ Monitor the training status, the code is as follows:

SELECT track_train_process('host_ip', 'port');

⑥ Predict the performance of SQL query through explain + SQL statement, the code is as follows:

EXPLAIN (..., predictor 'model_name') SELECT ...

Get the results, where the "p-time" column is the label prediction.

Row Adapter  (cost=110481.35..110481.35 rows=100 p-time=99..182 width=100) (actual time=375.158..375.160 rows=2 loops=1)

8.6.6 Evolution Route

The generalization ability of the current model is insufficient, relying on external AI Engine components, and the deep learning network is relatively heavy, which will cause difficulties for deployment; the model needs data for training, and the connection in the cold start stage is not smooth enough. Follow-up from the following aspects evolve.
(1) Add models of different complexity and support multi-model fusion analysis to provide more robust model prediction results and confidence.
(2) AI Engine is considering joining the task queue. Currently, it only supports single concurrent prediction/training, and multiple servers can be considered for concurrent business.
(3) Based on the enhancement of online learning/transfer learning, consider adding an anchor penalty cost to the loss function to avoid the problem of disaster forgetting. At the same time, optimize the data management mode, consider the data score mechanism, and assign weights according to the timeliness of the data.
(4) Integrate this function with the optimizer to explore AI-based path selection methods.

Thank you for learning the wonderful content of "8.6 AI Query Time Prediction" in Chapter 8 AI Technology. In the next article, we will introduce the relevant content of "8.7 DeepSQL".
Stay tuned.