LQRS: Learned Query Re-optimization Framework for Spark SQL

The paper proposes LQRS, a learned query re-optimization framework for Spark SQL that leverages a curriculum reinforcement learning strategy and runtime observations to dynamically refine execution plans, achieving up to a 90% reduction in end-to-end execution time compared to existing methods.

The Gist

Imagine you are the captain of a massive cargo ship (a database) trying to deliver goods (data) to a port. Your job is to figure out the fastest route to get there.

The Old Way: The Guessing Captain

Traditionally, the ship's captain (the Query Optimizer) looks at a map and tries to guess how much cargo is in each container and how fast the wind is blowing. Based on these guesses, they draw a single route on a piece of paper before the ship even leaves the dock.

The problem? The captain's guesses are often wrong. Maybe a container is actually empty, or a storm is brewing. But once the ship leaves the dock, the captain is stuck with that paper route. Even if they see a storm coming or realize a container is light, they can't change the course. They just have to power through, wasting fuel and time.

The "Smart" Way (Previous Attempts): The Pre-Flight Guru

Recently, engineers tried to replace the guessing captain with a Super-Computer trained on millions of past voyages. This computer is great at predicting the route before the ship leaves. It's much better than the old captain.

But, it still suffers from the same flaw: It can't change the plan once the ship is moving. It makes a perfect plan based on what it thinks is happening, but if reality is different, the ship is stuck with a bad route until it crashes or finishes slowly.

The New Solution: LQRS (The "Pilot Who Can Reroute")

The paper introduces LQRS, a new system that acts like a smart pilot who can reroute the ship while it's sailing.

Here is how LQRS works, using simple analogies:

1. The "Checkpoints" (Query Stages)

In the world of big data (Spark SQL), the journey isn't one long straight line. It's broken into Checkpoints. Imagine the ship drops off a batch of cargo, checks the weather, and then decides what to do next.

• Old systems only look at the map at the very beginning.
• LQRS waits for the ship to reach a checkpoint. Once the ship arrives, it knows the exact truth: "Wow, that first container was actually empty!" or "That storm is much smaller than we thought!"

2. The "Re-Optimization" (Changing the Course)

Because LQRS knows the real facts at the checkpoint, it doesn't just stick to the old paper plan. It says, "Okay, since that container is empty, let's skip the long detour and take the shortcut!"

• It can swap the order of tasks (like deciding to load the light boxes before the heavy ones).
• It can change the method of delivery (like switching from a slow truck to a fast drone).
• Crucially: It does this while the ship is moving, not just before it leaves.

3. The "Teacher and Student" (Reinforcement Learning)

LQRS learns using a method called Reinforcement Learning, which is like training a dog.

• The Student (The Actor): This is the part of the system that makes the decisions (e.g., "Let's swap these two tables").
• The Teacher (The Critic): This part watches the student. If the student makes a move that saves time, the Teacher gives a treat (positive reward). If the student makes a move that causes a delay or a crash, the Teacher gives a "no" (negative reward).
• Curriculum Learning: At first, the student is only allowed to make simple choices (like "start with this table"). As it gets smarter, the Teacher lets it try more complex moves (like "swap these two huge tables"). This prevents the student from getting overwhelmed.

4. The "Plug-and-Play" Tool

One of the coolest things about LQRS is that it doesn't require rebuilding the whole ship. It's a plug-and-play extension. Think of it like adding a high-tech GPS navigation system to an old car. The car (Spark SQL) still runs the engine, but the GPS (LQRS) can tell the driver to turn left or right instantly based on real-time traffic, overriding the old map.

Why is this a Big Deal?

The researchers tested LQRS on four different "oceans" (datasets) with thousands of complex queries.

• The Result: LQRS was up to 90% faster than other smart systems.
• The Analogy: If the old smart system took 100 minutes to deliver the cargo, LQRS did it in 10 minutes.

Summary

LQRS is a system that stops guessing and starts observing. Instead of making a rigid plan based on guesses, it waits for real data to arrive, learns from it instantly, and changes the plan on the fly. It combines the brainpower of AI with the flexibility of real-time adjustments, making database queries faster, smarter, and more efficient.

Technical Summary

Here is a detailed technical summary of the paper "LQRS: Learned Query Re-optimization Framework for Spark SQL".

1. Problem Statement

Traditional cost-based query optimizers (CBOs) rely on handcrafted cost models and estimated statistics, which often lead to suboptimal execution plans due to inaccurate cardinality predictions. While Learned Query Optimizers (LQOs) have shown promise by using machine learning to improve plan selection, they generally suffer from the "optimize-then-execute" paradigm.

• Limitation: Most LQOs (e.g., Lero, AutoSteer) make all optimization decisions before execution based on estimated statistics. They cannot leverage valuable runtime observations (e.g., actual cardinality of intermediate results) to dynamically adjust the plan once execution has started.
• Consequence: If a query encounters unexpected data distributions during execution (e.g., a table scan returns only one row instead of millions), static LQOs cannot correct the plan, leading to poor performance. Existing adaptive techniques (like Spark AQE) can adjust join strategies but lack the ability to re-order joins or make complex structural changes based on learned policies.

2. Methodology: LQRS Framework

The authors propose LQRS (Learned Query Re-optimization System), a framework that integrates Reinforcement Learning (RL) into the Spark SQL Adaptive Query Execution (AQE) engine. LQRS unifies pre-execution optimization and in-execution re-optimization.

Core Components

1. Planner Extension (The Interface):

   • Built on top of Spark AQE's extensibility interfaces.
   • Extraction: Intercepts the query execution at query stage boundaries (points where shuffle occurs). It extracts the partially executed plan and actual runtime statistics (cardinality, data size).
   • Injection: Receives optimization actions from the decision model and modifies the remaining unexecuted portion of the plan without restarting the query or invalidating completed stages.
   • Safety: Relies on Spark's built-in semantic equivalence checks to ensure modified plans are valid.

2. Decision Model (The Brain):

   • Architecture: Uses an Actor-Critic framework based on TreeCNN (Tree Convolutional Neural Network) to process the query plan tree structure.
   • State Representation ($s_t$): Encodes the partially executed plan as a vector tree. Features include node types, table participation, and true runtime statistics (actual cardinality and byte size) of completed stages.
   • Action Space: The model can perform five types of actions:
     • init: Select plan initialization strategy.
     • swap(i, j): Swap two leaf nodes in the join order.
     • lead(t): Move a specific table to the front of the join order.
     • broadcast(i): Force a broadcast join for a specific table (overcoming static thresholds).
     • no-op: Wait for more runtime data (crucial for stability).
   • Reward Function: Designed to minimize total execution time while penalizing unnecessary data reshuffling (intermediate shuffles). It uses a dense reward signal based on step-by-step feedback rather than just end-to-end latency.

3. Training Strategy (Curriculum Learning):

   • To ensure stability, the model is trained using a curriculum learning approach:
     • Phase 1: Only allowed to choose initialization strategies.
     • Phase 2: Allowed to perform runtime adjustments (lead/swap) based on actual statistics.
     • Phase 3: Full action space exploration.
   • This allows the model to learn simple decisions first before tackling complex, dynamic re-optimization.

3. Key Contributions

• First Learned Re-optimizer for Spark: LQRS is the first system to jointly support pre-execution and in-execution optimization on Spark SQL, effectively bridging the gap between learned optimization and adaptive query processing.
• Runtime-Driven Decision Making: Unlike previous LQOs, LQRS uses actual runtime observations (true cardinalities) to guide decisions, allowing it to correct suboptimal plans (e.g., reordering joins) that were based on poor estimates.
• Plug-and-Play Planner Extension: A novel extension module that enables online plan modification without requiring changes to the core Spark SQL engine logic, leveraging AQE's synchronization points.
• Dense Feedback Mechanism: By evaluating actions at every query stage rather than waiting for the full query to finish, the system receives feedback signals that are 3× denser than traditional end-to-end learning approaches, accelerating convergence.

4. Experimental Results

The authors evaluated LQRS on four benchmarks: JOB, Extended JOB (ExtJOB), STACK, and TPC-H, comparing it against Spark SQL (default AQE), Lero, AutoSteer, and SSA (QuerySplit).

• Performance Gains:
  • LQRS reduced end-to-end execution time by up to 90% compared to other learned optimizers and re-optimization methods.
  • Across all benchmarks, LQRS reduced total execution time by 16.3% to 73.3% compared to standard Spark SQL.
  • Case Study: In the STACK benchmark (Query q1#1), LQRS observed that a table scan produced only 1 tuple. It dynamically reordered the join plan to process this table first, reducing execution time from 21.52s (Lero) to 15.24s.
• Overhead:
  • LQRS has the lowest optimization overhead among all methods (e.g., 33.3s total overhead on JOB vs. 30,098s for Lero). This is because it avoids exhaustive plan enumeration and only intervenes when necessary (via no-op).
• Robustness:
  • LQRS demonstrated strong robustness to data distribution shifts and workload changes, outperforming baselines even when trained on different data snapshots.

5. Significance

• Paradigm Shift: LQRS challenges the static nature of modern learned optimizers by proving that deferring optimization decisions to execution time yields superior results. It demonstrates that the "unknowns" of runtime data are critical inputs for optimization.
• Practical Integration: By building on Spark AQE, LQRS provides a blueprint for integrating AI-driven optimization into production-grade distributed SQL engines without requiring a complete rewrite of the engine.
• Dynamic Adaptation: It solves the "cold start" and "distribution shift" problems better than static models by continuously adapting the plan as data flows through the system, making it highly suitable for real-world, dynamic workloads.

In summary, LQRS represents a significant advancement in database optimization by successfully merging Reinforcement Learning with Adaptive Query Execution, enabling systems to "learn while executing" and dynamically correct course based on reality rather than estimates.