AQORA: A Fast Learned Adaptive Query Optimizer with Stage-Level Feedback for Spark SQL

TL;DR

AQORA introduces a learned adaptive query optimizer for Spark SQL that performs plan refinement during execution by leveraging stage-level feedback. Built on an actor–critic reinforcement learning framework with TreeCNN feature encoding, AQORA seamlessly extends Spark's AQE through a plug-in planner to adjust join orders and operator choices online, guided by a dense reward signal derived from runtime shuffles. The approach combines four strategies—runtime plan feature encoding, fine-grained feedback with interventions, automatic strategy adaptation via curriculum learning, and lightweight integration—to outperform state-of-the-art baselines, achieving up to 90% end-to-end speedups and reduced tail latency while keeping optimization overhead modest. Extensive experiments on JOB, ExtJOB, and STACK demonstrate robust improvements, with ablations highlighting the value of stage-level feedback, action-space design, and PPO-based learning for stable convergence. The work advances practical deployment by enabling continuous learning across pre-execution and in-execution phases without intrusive engine changes.

Abstract

Recent studies have identified two main approaches to improve query optimization: learned query optimization (LQO), which generates or selects better query plans before execution based on models trained in advance, and adaptive query processing (AQP), which adapts the query plan during execution based on statistical feedback collected at runtime. Although both approaches have shown promise, they also face critical limitations. LQO must commit to a fixed plan without access to actual cardinalities and typically rely on a single end-to-end feedback signal, making learning inefficient. On the other hand, AQP depends heavily on rule-based heuristics and lacks the ability to learn from experience. In this paper, we present AQORA, a learned adaptive query optimizer that continuously refines query plans through stage-level feedback, unifying pre-execution and in-execution optimization within a single learning process. AQORA addresses the above challenges through four core strategies: (1) realistic feature encoding, (2) query stage-level feedback and intervention, (3) automatic strategy adaptation, and (4) low-cost integration. Experiments show that AQORA reduces end-to-end execution time by up to 90% compared to other learned methods and by up to 70% compared to Spark SQL's default configuration with adaptive query execution.

Figures (11)

• Figure 1: AQORA is the first learned query optimizer achieving execution-time query plan optimization on Spark SQL.
• Figure 2: In a left-deep tree plan, the left child of any join operation is a subtree, and the right child must be a leaf node. In contrast, bushy tree plans allow both children of any join operation to be subtrees
• Figure 3: Execution and optimization times of selected JOB queries with and without CBO.
• Figure 4: This query performs three joins: after the first join $t1 \bowtie t2$ produces only 5 rows, AQE switches the second join from SMJ to BHJ; after the second join $(t1 \bowtie t2) \bowtie t3$ produces 100,000 rows, AQE switches the third join from BHJ to SMJ.
• Figure 5: An overview of AQORA, which consists of two main components: (1) a decision model that generates optimization actions, and (2) an AQE planner extension that applies these actions and provides feedback to the decision model.