$30 off During Our Annual Pro Sale. View Details »

Balsa: Learning a Query Optimizer Without Expert Demonstrations

Balsa: Learning a Query Optimizer Without Expert Demonstrations

SIGMOD 2022
Talk video: https://www.youtube.com/watch?v=2Y88teiluGo
Paper: https://zongheng.me/pubs/balsa-sigmod2022.pdf
Code: https://github.com/balsa-project/balsa

Balsa is a reinforcement learning agent that learns to optimize SQL queries.

Abstract:

Query optimizers are a performance-critical component in every database system. Due to their complexity, optimizers take experts months to write and years to refine. In this work, we demonstrate for the first time that learning to optimize queries without learn- ing from an expert optimizer is both possible and efficient. We present Balsa, a query optimizer built by deep reinforcement learn- ing. Balsa first learns basic knowledge from a simple, environment- agnostic simulator, followed by safe learning in real execution. On the Join Order Benchmark, Balsa matches the performance of two expert query optimizers, both open-source and commercial, with two hours of learning, and outperforms them by up to 2.8× in workload runtime after a few more hours. Balsa thus opens the possibility of automatically learning to optimize in future compute environments where expert-designed optimizers do not exist.

Presenter: Zongheng Yang (https://zongheng.me), @zongheng_yang.

Zongheng Yang

June 19, 2022
Tweet

More Decks by Zongheng Yang

Other Decks in Research

Transcript

  1. Balsa: Learning a Query Optimizer
    Without Expert Demonstrations
    Zongheng Yang, Wei-Lin Chiang*, Frank Luan*,
    Gautam Mittal, Michael Luo, Ion Stoica
    SIGMOD 2022
    bit.ly/balsa-pdf
    github.com/balsa-project

    View Slide

  2. SELECT ...
    FROM A, B, C
    WHERE A.age > 21
    AND C.city IN (SFO,HNL)
    Optimizers are essential for data systems
    Query
    Optimizer
    Execution Plan
    C
    Hash Join
    Loop Join
    A B
    2

    View Slide

  3. SELECT ...
    FROM A, B, C
    WHERE A.age > 21
    AND C.city IN (SFO,HNL)
    Optimizers are essential for data systems
    Query
    Optimizer
    Execution Plan
    C
    Hash Join
    Loop Join
    A B
    Critical for performance:
    pick join orders, physical operators
    Widely applicable: used by
    many kinds of data systems
    2

    View Slide

  4. Optimizers take years to develop
    First optimizer before 2000s;
    still seeing changes to date
    Built cost-based
    optimizer 3 years after
    first release
    Shipped first optimizer by
    a team after
    “9 months of intense effort” [1]
    [1] https://www.cockroachlabs.com/blog/building-cost-based-sql-optimizer

    View Slide

  5. New systems: even higher dev cost
    New computation models, optimization objectives, or engines:
    no expert optimizer to build off from.

    View Slide

  6. Can we learn to optimize queries
    without learning from
    an expert optimizer?

    View Slide

  7. Balsa: a Learned Query Optimizer
    Environment
    Training Workload
    Execution Engine
    Learned
    Optimizer
    Agent
    Balsa

    View Slide

  8. Balsa: a Learned Query Optimizer
    Environment
    Training Workload
    Execution Engine
    Learned
    Optimizer
    Agent
    Balsa
    Query

    View Slide

  9. Balsa: a Learned Query Optimizer
    Environment
    Training Workload
    Execution Engine
    Learned
    Optimizer
    Agent
    Balsa
    Query
    Plan

    View Slide

  10. Balsa: a Learned Query Optimizer
    Environment
    Training Workload
    Execution Engine
    Learned
    Optimizer
    Agent
    Balsa
    Query
    Plan
    Latency

    View Slide

  11. Balsa: a Learned Query Optimizer
    Environment
    Training Workload
    Execution Engine
    Learned
    Optimizer
    Agent
    Balsa
    Query
    Plan
    Latency
    Update policy

    View Slide

  12. Balsa: a Learned Query Optimizer
    Environment
    Training Workload
    Execution Engine
    Learned
    Optimizer
    Agent
    Balsa
    Query
    Plan
    Latency
    Update policy
    Without learning from an expert system,
    Balsa can outperform open-source & commercial optimizers

    View Slide

  13. Balsa: a Learned Query Optimizer
    Environment
    Training Workload
    Execution Engine
    Learned
    Optimizer
    Agent
    Balsa
    Query
    Plan
    Latency
    Update policy
    7

    View Slide

  14. Planning: value network
    Value: (query, partial plan) → final latency
    Query
    (tables to join)
    }
    A B C
    {
    8

    View Slide

  15. Planning: value network
    Loop Join
    A B
    Hash Join
    C B
    Value: (query, partial plan) → final latency
    Query
    (tables to join)
    }
    A B C
    {
    pick join order,
    physical op
    8

    View Slide

  16. Planning: value network
    “Leads to a final plan that
    runs in time T”
    Loop Join
    A B
    Value ( )
    Hash Join
    C B
    Value ( )
    Analogy: final win prob.
    Value: (query, partial plan) → final latency
    Query
    (tables to join)
    }
    A B C
    {
    pick join order,
    physical op
    10 secs
    1 min.
    8

    View Slide

  17. Value network & tree search
    Loop Join
    A B
    Value ( )
    Hash Join
    C B
    Value ( )
    Value: (query, partial plan) → final latency
    Query
    (tables to join)
    }
    A B C
    {
    Run tree search to get
    a complete plan
    ML advance (AlphaGo)
    Used in optimizing Halide,
    theorem proving, etc.
    10 secs
    1 min.
    9
    pick join order,
    physical op

    View Slide

  18. Balsa: a Learned Query Optimizer
    Environment
    Training Workload
    Execution Engine
    Balsa
    Query
    Plan
    Latency
    Learned
    Value Network Tree Search
    train loop:
    plan()
    execute()
    update()
    10

    View Slide

  19. Balsa: a Learned Query Optimizer
    Environment
    Training Workload
    Execution Engine
    Balsa
    Query
    Plan
    Latency
    Learned
    Value Network Tree Search
    train loop:
    plan()
    execute()
    update()
    Must execute plans to get rewards
    No expert data → can easily run into
    bad plans, which are slow
    11

    View Slide

  20. Key challenge: Slow plans are everywhere!
    Expensive plans
    Good plans
    1-hour plan
    1-sec plan
    5-sec plan
    Plan space for a given query (schematic*)
    12
    * Based on our experiments on Join Order Benchmark, and "How Good Are Query Optimizers, Really?”, Leis et al., VLDB’15.
    1-day plan

    View Slide

  21. Key challenge: Slow plans are everywhere!
    Expensive plans
    Good plans
    1-hour plan
    1-sec plan
    5-sec plan
    Agent actions
    ? ?
    ?
    ?
    ?
    Plan space for a given query (schematic*)
    12
    * Based on our experiments on Join Order Benchmark, and "How Good Are Query Optimizers, Really?”, Leis et al., VLDB’15.
    1-day plan

    View Slide

  22. Sim-to-real learning
    Environment
    Training Workload
    Execution Engine
    Balsa
    Learned
    Value Network Tree Search
    Query
    Plan
    Latency
    13
    Pretrain in a minimal, logical-only “simulator”
    Instead of executing plans, quickly estimate their costs

    View Slide

  23. Sim-to-real learning
    Environment
    Training Workload
    Execution Engine
    Balsa
    Learned
    Value Network Tree Search
    14
    Pretrain in a minimal, logical-only “simulator”
    Instead of executing plans, quickly estimate their costs

    View Slide

  24. Sim-to-real learning
    Environment
    Training Workload
    Execution Engine
    Balsa
    Learned
    Value Network Tree Search
    Simulation
    Environment
    Training Workload
    Two-liner
    Cost Model
    Query
    Plan
    Cost
    14
    Pretrain in a minimal, logical-only “simulator”
    Instead of executing plans, quickly estimate their costs

    View Slide

  25. Sim-to-real learning
    Environment
    Training Workload
    Execution Engine
    Balsa
    Learned
    Value Network Tree Search
    Simulation
    Environment
    Training Workload
    Two-liner
    Cost Model
    Query
    Plan
    Cost
    Simulation enables the agent to efficiently
    learn to avoid the most disastrous (high-cost) plans.
    14
    Pretrain in a minimal, logical-only “simulator”
    Instead of executing plans, quickly estimate their costs

    View Slide

  26. Simulator: a minimal cost model
    15
    Cout(T) =
    (
    |T| B7 T Bb i#H2fb2H2+iBQM
    |T| + Cout(T1) + Cout(T2) B7 T = T1 ./ T2
    AAACfnicbVFNb9NAEF2br2K+Ahy5DCRUDRWtHSHRS6WKXjgWKWkrxZG13ozTVde71u4YiFz/DP4YN34LFzaphUjLSCu9fTPzZvdNXinpKI5/BeGdu/fuP9h6GD16/OTps97zF6fO1FbgRBhl7HnOHSqpcUKSFJ5XFnmZKzzLL49X+bOvaJ00ekzLCmclX2hZSMHJU1nvR3ScNaamdmc8hMMozXEhdSO8omujq/EVbKeE36mRBQzGA5AOOBD36vt+JoqVSAtpui7dhb9aWTLcuI6GG0JwCL4E0tx8I4kejwZtlKKed5OzXj/ei9cBt0HSgT7r4iTr/UznRtQlahKKOzdN4opmDbckhUKvXTusuLjkC5x6qHmJbtas7WvhrWfmUBjrjyZYs/92NLx0blnmvrLkdOFu5lbk/3LTmoqDWSN1VRNqcT2oqBWQgdUuYC6tt1AtPeDCSv9WEBfcckF+Y5E3Ibn55dvgdLSXePzlQ//oU2fHFnvF3rAdlrCP7Ih9ZidswgT7HbwO3gW7IQu3w/fh/nVpGHQ9L9lGhAd/ANtVvBw=
    AAACfnicbVFNb9NAEF2br2K+Ahy5DCRUDRWtHSHRS6WKXjgWKWkrxZG13ozTVde71u4YiFz/DP4YN34LFzaphUjLSCu9fTPzZvdNXinpKI5/BeGdu/fuP9h6GD16/OTps97zF6fO1FbgRBhl7HnOHSqpcUKSFJ5XFnmZKzzLL49X+bOvaJ00ekzLCmclX2hZSMHJU1nvR3ScNaamdmc8hMMozXEhdSO8omujq/EVbKeE36mRBQzGA5AOOBD36vt+JoqVSAtpui7dhb9aWTLcuI6GG0JwCL4E0tx8I4kejwZtlKKed5OzXj/ei9cBt0HSgT7r4iTr/UznRtQlahKKOzdN4opmDbckhUKvXTusuLjkC5x6qHmJbtas7WvhrWfmUBjrjyZYs/92NLx0blnmvrLkdOFu5lbk/3LTmoqDWSN1VRNqcT2oqBWQgdUuYC6tt1AtPeDCSv9WEBfcckF+Y5E3Ibn55dvgdLSXePzlQ//oU2fHFnvF3rAdlrCP7Ih9ZidswgT7HbwO3gW7IQu3w/fh/nVpGHQ9L9lGhAd/ANtVvBw=
    AAACfnicbVFNb9NAEF2br2K+Ahy5DCRUDRWtHSHRS6WKXjgWKWkrxZG13ozTVde71u4YiFz/DP4YN34LFzaphUjLSCu9fTPzZvdNXinpKI5/BeGdu/fuP9h6GD16/OTps97zF6fO1FbgRBhl7HnOHSqpcUKSFJ5XFnmZKzzLL49X+bOvaJ00ekzLCmclX2hZSMHJU1nvR3ScNaamdmc8hMMozXEhdSO8omujq/EVbKeE36mRBQzGA5AOOBD36vt+JoqVSAtpui7dhb9aWTLcuI6GG0JwCL4E0tx8I4kejwZtlKKed5OzXj/ei9cBt0HSgT7r4iTr/UznRtQlahKKOzdN4opmDbckhUKvXTusuLjkC5x6qHmJbtas7WvhrWfmUBjrjyZYs/92NLx0blnmvrLkdOFu5lbk/3LTmoqDWSN1VRNqcT2oqBWQgdUuYC6tt1AtPeDCSv9WEBfcckF+Y5E3Ibn55dvgdLSXePzlQ//oU2fHFnvF3rAdlrCP7Ih9ZidswgT7HbwO3gW7IQu3w/fh/nVpGHQ9L9lGhAd/ANtVvBw=
    AAACfnicbVFNb9NAEF2br2K+Ahy5DCRUDRWtHSHRS6WKXjgWKWkrxZG13ozTVde71u4YiFz/DP4YN34LFzaphUjLSCu9fTPzZvdNXinpKI5/BeGdu/fuP9h6GD16/OTps97zF6fO1FbgRBhl7HnOHSqpcUKSFJ5XFnmZKzzLL49X+bOvaJ00ekzLCmclX2hZSMHJU1nvR3ScNaamdmc8hMMozXEhdSO8omujq/EVbKeE36mRBQzGA5AOOBD36vt+JoqVSAtpui7dhb9aWTLcuI6GG0JwCL4E0tx8I4kejwZtlKKed5OzXj/ei9cBt0HSgT7r4iTr/UznRtQlahKKOzdN4opmDbckhUKvXTusuLjkC5x6qHmJbtas7WvhrWfmUBjrjyZYs/92NLx0blnmvrLkdOFu5lbk/3LTmoqDWSN1VRNqcT2oqBWQgdUuYC6tt1AtPeDCSv9WEBfcckF+Y5E3Ibn55dvgdLSXePzlQ//oU2fHFnvF3rAdlrCP7Ih9ZidswgT7HbwO3gW7IQu3w/fh/nVpGHQ9L9lGhAd/ANtVvBw=
    Cout [1]: cost = sum(estimated cardinalities of all operators)
    ✓ generic (logical-only): no assumption on physical environment
    ✓ correlates with execution speed
    [1] “On the complexity of generating optimal left-deep processing trees with cross products”, Cluet and Moerkotte, ICDT’95.

    View Slide

  27. Safety mechanisms
    Executing a suboptimal plan of a training query
    In real execution, two ways to hit slow plans
    16
    Exploration: occasionally execute unfamiliar plans

    View Slide

  28. Safety mechanisms
    Executing a suboptimal plan of a training query
    In real execution, two ways to hit slow plans
    16
    Exploration: occasionally execute unfamiliar plans
    Safe execution: time out & set its unknown latency to LARGE_NUM

    View Slide

  29. Safety mechanisms
    Executing a suboptimal plan of a training query
    In real execution, two ways to hit slow plans
    16
    Exploration: occasionally execute unfamiliar plans
    Safe execution: time out & set its unknown latency to LARGE_NUM
    Safe exploration: only explore “probably good” plans
    Use k-best plans from tree search

    View Slide

  30. Balsa: the complete architecture
    Balsa
    Learned
    Value Network Tree Search
    Minimal
    Simulator
    Safely
    Execute, Explore
    17

    View Slide

  31. Bonus: Diversifying experiences
    18
    Exp.
    Buffer
    Trained Agent 1
    Exp.
    Buffer
    Trained Agent N
    Observation: agent experiences differ a lot
    Ex: Some prefers loop join, vs. merge join

    View Slide

  32. Bonus: Diversifying experiences
    18
    Exp.
    Buffer
    Trained Agent 1
    Exp.
    Buffer
    Trained Agent N
    Merged
    experiences
    Diversified agent
    Offline training
    Observation: agent experiences differ a lot
    Ex: Some prefers loop join, vs. merge join
    Idea: diversified experiences
    Offline retrain an agent on it
    Significantly enhances both training &
    unseen test query performance (see paper)

    View Slide

  33. Prior art, ML for QO
    Learned optimizers

    DQ [Krishnan et al.]

    Learns from expert cost models
    Neo [VLDB ’19]

    Learns from expert plans (demonstrations)
    19
    Optimizer assistant

    Bao [SIGMOD ’21]
    Doesn’t produce plans; sets flags for expert optimizer
    Intra-query learning: SkinnerDB [SIGMOD ’19]

    During a query’s execution, explores many join orders
    Limited to picking left-deep join orders

    View Slide

  34. Prior art, ML for QO
    Learned optimizers

    DQ [Krishnan et al.]

    Learns from expert cost models
    Neo [VLDB ’19]

    Learns from expert plans (demonstrations)
    19
    Optimizer assistant

    Bao [SIGMOD ’21]
    Doesn’t produce plans; sets flags for expert optimizer
    Intra-query learning: SkinnerDB [SIGMOD ’19]

    During a query’s execution, explores many join orders
    Limited to picking left-deep join orders
    For the first time, Balsa removes
    the need to learn from expert
    optimizers.

    View Slide

  35. Prior art, ML for QO
    Learned optimizers

    DQ [Krishnan et al.]

    Learns from expert cost models
    Neo [VLDB ’19]

    Learns from expert plans (demonstrations)
    19
    Optimizer assistant

    Bao [SIGMOD ’21]
    Doesn’t produce plans; sets flags for expert optimizer
    Intra-query learning: SkinnerDB [SIGMOD ’19]

    During a query’s execution, explores many join orders
    Limited to picking left-deep join orders
    For the first time, Balsa removes
    the need to learn from expert
    optimizers.
    Balsa learns to produce

    plans with physical operators

    & bushy join orders.

    View Slide

  36. Workload speedup
    PostgreSQL v12
    20

    View Slide

  37. Workload speedup
    PostgreSQL v12
    Workload Speedup
    JOB JOB-Slow TPC-H
    train test
    1x
    20

    View Slide

  38. Workload speedup
    PostgreSQL v12
    Workload Speedup
    JOB JOB-Slow TPC-H
    train test
    1.2
    1.3
    1.7
    1.1
    1.3
    2.1
    1x
    20

    View Slide

  39. Workload speedup
    PostgreSQL v12 Commercial DBMS
    Workload Speedup
    JOB JOB-Slow TPC-H
    train test
    1.2
    1.3
    1.7
    1.1
    1.3
    2.1
    Workload Speedup
    JOB JOB-Slow TPC-H
    1.0
    1.5
    1.9
    1.1
    2.4
    2.8
    train test
    1x 1x
    Without learning from expert optimizers, Balsa outperforms them
    on both train (up to 2.8x) & unseen (up to 1.9x) queries 20

    View Slide

  40. Training Balsa is efficient
    0 5 10 15
    Elasped Time (hours)
    20
    10
    5
    1
    0.5
    Normalized Runtime [log]
    JOB JOB, Slow TPC-H
    Faster than expert plans
    Slower than expert plans

    View Slide

  41. Training Balsa is efficient
    Training Balsa is efficient: a few hours
    to nightly maintenance range
    0 5 10 15
    Elasped Time (hours)
    20
    10
    5
    1
    0.5
    Normalized Runtime [log]
    JOB JOB, Slow TPC-H
    Faster than expert plans
    Slower than expert plans

    View Slide

  42. Balsa vs. Expert Demonstrations (Neo)
    Test set performance
    22
    0 2 4 6 8 10
    Elasped Time (hours)
    50
    10
    5
    1
    0.5
    Normalized Runtime [log]
    Neo-impl Balsa

    View Slide

  43. Balsa vs. Expert Demonstrations (Neo)
    Test set performance
    22
    0 2 4 6 8 10
    Elasped Time (hours)
    50
    10
    5
    1
    0.5
    Normalized Runtime [log]
    Neo-impl Balsa
    Balsa significantly improves generalization to unseen queries
    (more robust; can surpass expert)
    Failed to generalize to
    unseen queries without
    simulation learning

    View Slide

  44. Balsa
    Learn to optimize, without learning from expert optimizers

    Key challenge: mitigate slow plans during learning
    Paper: more ideas & evaluation
    Open sourced on GitHub
    For researchers: reproduce & build on
    For users: find best plans for your queries/hardware
    {zongheng, weichiang, lsf, gbm, michael.luo, istoica}
    @berkeley.edu
    bit.ly/balsa-pdf
    github.com/balsa-project

    View Slide