ICDE2024_GLO_Towards_Generalized_Learned_Query_Optimization_GoldenDB.pdf

迹部景吾

183

13页

4次

2024-08-01

免费下载

GLO: Towards Generalized Learned Query

Optimization

Tianyi Chen

†

, Jun Gao

†∗

, Yaofeng Tu

‡∗

, and Mo Xu

‡

†

Key Laboratory of High Conﬁdence Software Technologies, CS, Peking University, China

‡

ZTE Corporation

tianyichen@stu.pku.edu.cn, gaojun@pku.edu.cn, {tu.yaofeng, xu.mo1}@zte.com.cn

Abstract—In recent years, there has been a growing interest in

the application of deep reinforcement learning (DRL) techniques

on query execution plan generation. Although current DRL-

based query optimizers achieve competitive performance against

traditional methods on speciﬁc query workloads, these methods

encounter issues when generalizing to workloads unseen during

training. Thus, we propose GLO to address the limitations

and step towards generalized learned query optimization. First,

rather than using ungeneralizable table-speciﬁc one-hot labels in

almost all existing work, GLO relies on statistical information of

the well-established underlying DBMS along with table patterns

extracted via a clustering algorithm, enabling GLO to enhance

generalization in different scenarios. Second, GLO improves

the information capture of plans by integrating Transformer

layers into the DRL value model, empowering the model’s

capability to handle diverse queries with deeper networks and

more parameters in plan generation. In addition, GLO allows

the injection of cost estimations from the DBMS as external

knowledge for better generalization. Third, GLO recognizes and

replaces disastrously poor plans by making comparisons between

generated plans and those produced by the DBMS. We establish

our experiments on composite workloads that combine various

query sets including JOB, Extended JOB, TPC-DS, and Stack.

The results demonstrate that GLO outperforms previous state-

of-the-art learned optimizers, with a speed 1.4x faster than

LOGER and 2.1x faster than Balsa on TPC-DS when TPC-DS

queries are completely unknown during training. To the best of

our knowledge, GLO is the ﬁrst learned optimizer that directly

generates plans while possessing the preliminary generalization

ability across different query workloads.

Index Terms—learned query optimizer, generalization, deep

reinforcement learning

I. INTRODUCTION

The query optimizer is a crucial component of database

management systems, whose role is to generate plans for input

queries while seeking efﬁcient and stable execution perfor-

mance [1]. With great effort in development, traditional query

optimizers achieve stability under different circumstances.

However, these optimization algorithms face challenges in

achieving high performance on data with complicated distri-

butions and different hardware environments [2], [3]. How to

generate efﬁcient query execution plans has been an important

problem in the ﬁeld of databases for decades [4]–[6].

In order to reach higher query execution performance and

lower the development effort, a number of optimizers based

on deep reinforcement learning (DRL) have emerged in recent

* Corresponding authors

Fig. 1. The per-query latencies of Balsa, LOGER, and Lero in milliseconds

compared with PostgreSQL on TPC-DS when trained on other workloads.

years, which can be roughly categorized into two classes.

Selection-based methods like Bao [7], Lero [8], etc. select

from multiple plans generated by the underlying DBMS’s

traditional optimizer with different hints. These methods incor-

porate the development of the DBMS to produce more efﬁcient

plans, with a relatively stable but also limited performance

due to the dependence on a traditional optimizer. In contrast,

generation-based methods like RTOS [9], Balsa [10], and

LOGER [11], etc. consecutively search for intermediate plans

(subplans) from scratch until producing complete plans with

their own knowledge about the distribution. These methods do

not depend on an existing optimizer and have more potential

to achieve better performance than selection-based methods.

Although generation-based optimizers can yield impressive

performance on speciﬁc query workloads with sufﬁcient train-

ing, these methods face one severe issue that hinders them

from being practical. When the generation-based optimizers

are applied to a workload largely different from training, they

exhibit obvious performance degradation and often produce

poorer plans compared with traditional optimizers. Figure

1 illustrates the performances of state-of-the-art generation-

based optimizers Balsa [10] and LOGER [11] on the TPC-

DS [12] query workload when trained on JOB [2], Extended

JOB [13] and Stack [7], along with a selection-based method

Lero [8]. While the poor performance of generation-based

methods can be observed, the results also show that Lero

works more robustly but cannot outperform PostgreSQL ei-

ther. These results are unacceptable in real-life application

scenarios. Therefore, we investigate the causes behind the poor

generalization ability of existing generation-based methods.

First, current generation-based optimizers learn with un-

generalizable table-speciﬁc features. As far as our knowledge

4843

2024 IEEE 40th International Conference on Data Engineering (ICDE)

DOI 10.1109/ICDE60146.2024.00368

Authorized licensed use limited to: ZTE CORPORATION. Downloaded on July 31,2024 at 08:22:42 UTC from IEEE Xplore. Restrictions apply.

goes, almost all previous generation-based optimizers, includ-

ing the aforementioned Balsa and LOGER, use table-speciﬁc

one-hot label vectors or directly learned embeddings as a part

of tables’ features. These features enable these models to fully

learn implicit knowledge for each speciﬁc table seen in an

end-to-end training process. However, these features become

invalid when the query workload or underlying database

schema is changed, eventually leading to poor performance.

Second, the model sizes of existing methods are not ad-

equate to handle diverse patterns of workloads and datasets.

Modern generation-based optimizers use Tree-CNN [14] or

Tree-LSTM [15] to capture information from the input sub-

plans. The delicately designed structures and the parameter-

sharing strategy enable these models to effectively handle tree-

structured subplans of arbitrary heights with low model sizes.

However, their simple network structures and few parameters

also incur limited model capabilities. A more capable model

with a deeper network and more parameters is essential for

better generalization, as the patterns of queries vary between

workloads, and even a change of a single predicate can lead

to an entirely different optimal join order for a query.

Third, existing generation-based optimizers lack the ability

to prevent producing disastrously poor plans especially when

adapted to new workloads and database schemas. It’s almost

inevitable for a generation-based optimizer to produce disas-

trous plans for some queries, as the training workload size

is always limited and the cases of queries cannot cover all

patterns of future inputs. A mechanism to remedy the issue

by recognizing and erasing poor plans is therefore crucial for

the generalization ability of a generation-based optimizer.

In order to overcome the aforementioned limitations, we

propose GLO, a generation-based method that steps towards

generalized learned query optimization, aiming to achieve both

high performance and generalization ability on unfamiliar data

and workload distributions. As a preliminary method towards

generalization, GLO currently selects only join orders. We will

extend it to support physical operator selection in the future.

GLO’s source code is available on GitHub.

GLO attempts to leverage information from the underly-

ing DBMS to handle the ﬁrst limitation. Recognizing the

generalization ability of traditional optimizers that delicately

utilize statistics, GLO discards the ungeneralizable features

and instead relies on statistical information from the DBMS to

enhance generalization across diverse workloads and datasets.

In addition, noticing that tables from different schemas may

exhibit similar patterns, GLO categorizes tables and captures

the similarities by employing the K-means algorithm. Further-

more, GLO allows the injection of cost estimations from a

cost estimator into GLO’s value prediction process, providing

insight from a different perspective for plan generation.

To face the second challenge, we draw inspiration from lan-

guage models and integrate the Transformer architecture [16]

into GLO’s value model to improve the model capability.

Beneﬁting from a deeper network structure and the atten-

https://github.com/TianyiChen0316/GLO

tion mechanism, Transformer-based models show superior

performance compared with CNNs and RNNs [17], [18],

and a high generalization potential on various tasks [19].

Therefore, we leverage Transformer layers in the value model

to predict the minimal latency for subplans. A similar idea

has been proposed in the work of QueryFormer [20], which

focuses on yielding plan representations to provide inputs for

various downstream tasks like cost estimation. In contrast,

GLO’s Transformer layers take subplan forests as inputs for

value predictions with more sophisticated features, aiming to

reinforce the generalization ability.

To further improve the tail performance and avoid disastrous

plans, we borrow the idea from selection-based optimizers [7],

[8] and propose a comparison mechanism, which effectively

selects the better one from the GLO’s generated plan and

the DBMS optimizer’s plan. Instead of introducing a separate

comparator model, GLO efﬁciently reuses the knowledge

learned by the value model and directly compares the predicted

values of generated plans and the DBMS’s plans. GLO applies

a comparator loss to enable correct comparison with the

value model apart from latency prediction. Training with the

comparator loss also improves the value model’s sensitivity to

relative differences and eventually beneﬁts the performance.

Additionally, we implement a rule-based augmentation strat-

egy to collect extra samples for the comparator loss without

any plan execution latency overhead.

We employ several experiment settings that contain various

query workloads, including JOB, Extended JOB, Stack, and

TPC-DS, to assess GLO’s performance and generalization

ability. The experiment results illustrate that GLO reaches

a performance better than previous state-of-the-art learned

optimizers. To demonstrate the generalization capability, we

evaluate GLO’s performance on the TPC-DS workload while

training it with all other three workloads. On the completely

unseen TPC-DS workload, GLO achieves a speed 1.4x faster

than LOGER and 2.1x faster than Balsa. Furthermore, by

introducing extra knowledge from Stack and TPC-DS apart

from JOB queries, GLO achieves a speedup ratio of 1.48

on Extended JOB. To the best of our knowledge, this is

the highest performance of existing learned optimizers and

even outperforms the reported 1.2x speedup of Balsa-8x

which requires an extremely expensive overhead of training 8

agents [10]. Results of more workload splits and experiment

settings are demonstrated in Section VI.

II. T

HE FRAMEWORK OF GLO

Figure 2 illustrates the method framework of GLO, from

which we can see that GLO’s plan generation process consists

of three steps. For each input query, GLO ﬁrst vectorizes both

table-level and column-level statistical information from the

database with a statistics extractor. The statistic vectors are

then integrated into table embeddings along with the informa-

tion of the input query. Next, GLO’s plan generator generates

plans step-by-step through a value-based DRL process, in

which GLO uses a value model to evaluate subplans, and

continuously selects the next subplan until a complete plan

4844

Authorized licensed use limited to: ZTE CORPORATION. Downloaded on July 31,2024 at 08:22:42 UTC from IEEE Xplore. Restrictions apply.

of 13

免费下载

关注

评论