Training a 4B model to produce 81% faster query plans than Postgres
rohanbansal.com533 points by polyphilz 14 hours ago
533 points by polyphilz 14 hours ago
“81% faster query plans than Postgres”…on an 8 GB dataset that fits entirely in memory, with shared_buffers constrained to a fraction of that, queries warmed before measuring, and read-only SELECTs.
I would be cautious about over fitting, it’s tough to say if those query plans would really be more optimal than Postgres heuristics at scale and with a bit more realistic OLTP workloads.
In any case, such is life with profile guided optimization. Many of us appreciate how database workloads can drift over time and with scale.
Kudos to the author for getting their hands dirty and writing up their experiments.
With a 4B parameter model that probably ran through 8GBs of RAM multiple times to run.
At a certain point we should seriously talk about CUDA accelerating Postgres instead.
I would think it's possible to make it so that the 4B model only needs to be called during an initial phase, and then the same queries it constructed can just be re-used with values replaced, unless you're generating a lot of unique on-the-fly query shapes.
pardon but aren't disks usually the bottleneck? im all for CUDA acceleration and CUDA accelerating culture
Parent post was talking about an 8GB dataset.
8GB isn't even CPU RAM these days. That's GPU super-mega-awesome ram. Ordinary Server CPUs are regularly pushing 2TB capacities.
GPUs are in the 8GB to 32GB typically, at least for smaller and more regular GPUs. This GPU RAM is also well known to be at least 10x the bandwidth of CPU RAM.
Yeah, I have a GPU from almost six years ago in my desktop that has twice that much VRAM. Less than a year ago my wife got a 5070 Ti with the same for around $750 without needing to wait for it to be in stock or anything. I'm inclined to think that for a server that needs a GPU, even 32 GB would probably be considered small.
not necessarily, no. With SSDs you get much better IOPS for cold data, and many datasets fit in RAM. So a lot of (OLTP/HTAP) workloads can become CPU-bound due to sorting/hashing - bread and butter of joins.
Though many use network block storage (e.g. EBS), which is significantly slower than an SSD.
Which is why a good query plan is so important, so that as much disk I/O can be avoided as possible (predicate push down, index elimination, join ordering, partition/scan pruning). Like the old CTE optimization fence problem.
What would you accelerate? Is there a lot of linear algebra you could throw cuda at in Postgres?
You know that GPUs are more flexible than just linear algebra, right?
GPUs are simply faster at fundamental algorithms like sorting (which has huge parallelism), and hashing. This is because both sorting and hashing benefit from endless growth of parallelism, offering enough "work" for these 10,000 SIMD-core systems to crunch work upon.
And because of modern algorithms/libraries with 'Mergepath sort' (a GPU-SIMD parallel sorting algorithm), its not even that difficult to implement anymore.
Naturally, this then leads to parallel Sort Merge Join, as well as parallel Hash-Join (two ways to implement left or right joins in a GPU that benefit from significant parallelism).
So yeah, Joins. https://www.kenchoi.dev/papers/gpu-joins.pdf (This paper also has a description of "Mergepath sort", a GPU parallel way of sorting)
---------
Even if GPUs weren't fundamentally faster at these kinds of operations... the RAM is simply 10x higher bandwidth and we all know its a RAM-constrained problem.
Your typical SQL query is going to need multiple joins, probably a sort and possibly some "group" operations. As long as you have more than 10,000 elements or so (IE: can saturate all 10,000+ SIMD-units of a GPU), you'll be able to at least benefit from the faster RAM.
If you have a LOT of joins (a recursive join or some other kind of deeply nested computationally complex query), you probably benefit even more from the greater compute-power offered by GPUs. These operations (joins really) are nominally over the entire set of data, and cleanly break down into obvious parallelism.
The GPU isn’t connected to the disk though. Usually. So you’d still have to load from disk, to ram, then from ram to the GPU.
I think in principle you could clone your database in prod and at least test to see if your most difficult + common queries are indeed faster after running through the LLM optimizer?
Frankly, just a general extension to feed a query log to a batch job to do offline optimisation of common actual reoccurring query shapes based on a query log might well be worth it.
> I think in principle you could clone your database in prod and at least test to see if your most difficult + common queries are indeed faster after running through the LLM optimizer?
That is the responsibility of whoever thought it would be a good idea to write this article. It's their responsibility to show that their idea has merit, and that their results are significant. I mean, don't they have a vested interest in manipulating and cherry-picking their results to inflate their relevance?
This is why academic papers are peer reviewed.
Optimal plan construction is math-heavy, algorithm-heavy and vary even by workload. There are options like creating just-in-time indexes, so solution space grows even faster than article presents. Sometimes it is the query planner which is the slow part of total execution time.
LLM is kind of blunt weapon to use here. I am waiting rather for alphago style neural net heuristic.
I also wondered why an LLM would be the right starting point. Why would Balzac or billions of lines of rwir code or reddit be relevant to mapping this smallish, well-defined language (SQL) to this other tiny constrained specification language (the query plan suggestions)? You could make a (relatively) tiny network and then actually pass it some relevant features of the actual data, like as numbers, not just as text returned from a tool call.
It’s ultimately based on a lot of hand-written heuristics. Google has some non-LLM based machine learning technique to guide optimization heuristics in LLVM; that would be closer to what you are looking for.
What if we use a hybrid model of using both query optimizer and LLM? Whichever produces better result, the database can use?
- a question from someone with lack of DB depth, me.
This is about to bake your noodle:
GEQO is not to get a better plan than the traditional optimizer, it is to be able to get a plan at all when the query is large. And it's widely known for creating poor plans.
Yes, but it's exactly the kind of hybrid between a regular planner and something generative (writ broadly) that they were asking about. Practically speaking if you're hitting the GEQO you've already failed as a query writer unless it's a purely OLAP on a dedicated beefy machine.
I think calling GEQO generative is a bit of a stretch; it's just a different way of searching through the same space with the same cost model. More or less devolving to “let's take a bunch of randomized join orders and see which one is best” :-)
And yes, large joins is definitely for OLAP use. If you have 20-way joins for OLTP, you're either crazy or you're using an ORM.
The immediate problem: How do you know which one is better without running them?
You create formulas to estimate the cost of running a given query plan. Use statistics collected about the tables (e.g. how many rows) to try to be accurate. The topic is "Cost Based Optimization".
If you have formulas that actually match reality, what do you need the LLM for? An optimizer is perfectly capable of finding the optimal plan if it has a perfect estimator. In fact, if you could only estimate the number of rows in each subplan perfectly, you have as good as solved the problem already.
Could you A/B at random, use that to collect data and eventually feed that back in to prefer A or B depending on the shape of the query?
There are papers and Postgres projects that attempt this kind of learning-based optimization, with some success. None are in widespread use. (One part, but certainly not the entirety, of the problem is that it's not just A/B, it's an exponential number of options that all could seem close to each other.)
[flagged]
This immediately halves your throughput.
Only in the worst case when the plans are equivalent: If one plan is significantly faster, then it'll finish first, and the loser can get canceled before it finishes.
Good and bad plans can have orders of magnitude performance difference. The bad one can easily do enough damage cutting the performance in half before it is canceled.
Engineer: "HELP, our production DB is frozen on this query that worked fine before!"
Infra: "Hmm, let's check... Well would you look at that, it seems like your LLM query planner usually works and produces fast queries, but this time when you changed a variable name to trigger query rebuild, it happened to hallucinate and miss an index, would you mind re-running the LLM a few times until you get a faster query?"
Funnily enough, you could replace "LLM query planner" with just "query planner" and this comment would still hold true
That bug is fixable and verifiable. The LLM you cross your fingers till the next time the same thing happens.
> fixable and verifiable
By people with a specific skill set. LLMs generation can also be fixed and verified by people with a certain skill set, and non-deterministic computing doesn't automatically mean unpredictable. When people say that the LLMs are a black box, it means unpredictability in unknown situations.
You do structured output, input validation, output validation, lower temperature, limit decisions, RL, etc. to increase predictability to near certainty. It's just statistics after all. Or you can as well generate the code to do the job.
It's just that the required skill set is a different one to do those things, and unusual in the context of DB administration.
Only if someone is planning on running a pinned self hosted version of an LLM alongside the DB to fix the problem. The developer can change the binary easy enough and test it but the LLM approach just seems either theoretical or bending ourselves in knots to justify using an LLM.
I didn't argue that it'd make sense, I just said that it could be reasonably fixable when problems occur and verifyable that the fix works. Even if shipping and RLing an LLM were easy tasks in terms of software distribution (they are not) we'd still hit the skill mismatch, as I said in my previous comment.
I just find the "all llms are non dererministic and therefore unreliable" narrative a bit backwards. All software that has more than 0 users needs to deal with non-determinism anyway :)
This exact situation has happened to me with Postgres' stock query planner when an automatic analyze got a bad sample of a large table.
I’ve seen this happen to SQL Server many times. Every time the solution is a stored proc with the recompile option enabled.
The answer to every problem is to rebuild statistics. (And never use stored procedures. It's just a shitty API layer in the worst language imaginable, sitting outside of source control. If you need an API layer, write it in a real language, ideally the one you're already using.)
Our database 12.34 was working great, but 12.35 deployment had some optimizer changes that had regression on exact scenario that you have in your statistics. Shit happens, sorry. Use this hint.
> Frontier intelligence is extremely powerful; the distillation I did off Astra trajectories is proof enough that large models are not going anywhere
Wouldn't admitting this invite trouble due to accusations of distillation flying around between closed and open models.
On the flip side, its important to know just how much we would be sacrificing if big frontier gets their way in convincing the courts that distillation is a bad thing
Yeah don't steal my stolen stuff.
Like this scene from Pirates of Silicon Valley between Bill Gates and Steve Jobs: https://www.youtube.com/watch?v=CBri-xgYvHQ Bill was saying that Steve was stealing from Xerox so it's okay if Bill stole from Xerox first.
I don't think anybody's denying that small models are doing distillation, the issue at hand is frontier lab vs frontier lab when it comes to big models.
Frontier model trainers stole almost all the data they've trained on (the whole internet, all copyrighted).
It's very hard for them to claim the moral high ground here.
It's like stealing an apple from the British Colonial Empire.
Having the moral high ground matters less than having a big warchest of money to spend on lawyers.
Does it matter if they companies doing are not in the jurisdiction or even if they are, maybe the can't prove it?
They can't prove it but can force you to spend time and money to disprove it.
The only thing they're going to try to do is ban Chinese models, and that's not going to fly outside the US, so Americans are going be the only losers.
You'd be shocked at some of the punishments handed out from various empires over history for offences as minor as stealing an apple. Whether they can claim the moral high ground is not question of ethics but a question of power.
Unless they’ve been elided, there were no indices other than the PK on any table, and no additional statistics. There are correlated columns here: a given country may have produced more movies in a given range of years, as its movie industry built up; a given country may produce more TV series than movies, etc.
Nearly every time I’ve seen someone resorting to hints for a query, it’s because their statistics are incorrect. Adding hints is papering over the problem, and can backfire later if the data shape changes.
> I paid ~$800 to rent a 2x H100 SXM node from Lambda for ~95 hours, and ~$400 in OpenAI API fees to generate the Astra trajectory demonstrations.
> a tiny 4B model went from not being able to understand the harness it was wrapped in, to achieving a 1.81x geometric mean speedup and a summed latency decrease of 44.7% across a workload of join-heavy SQL queries
I can’t find it in the article (may have skimmed it too much), but I suspect they didn’t include those ~95 hours in the benchmark numbers.
I think all database vendors know their query optimizers could do much better if they could afford to spend lots of time to derive query plans.
⇒ this may be useful for some workloads, but even then, can you afford to spend hours every now and then to update your 4B model to ensure it still picks a good query plan?
> ⇒ this may be useful for some workloads, but even then, can you afford to spend hours every now and then to update your 4B model to ensure it still picks a good query plan?
I think this would be likely comparable to a scheduled backup, so I think it would be an acceptable maintenance window. However, deterministic algorithms would likely beat re-training (or re-fine-tuning) the model. For example, one could analyze actual distributions or whatever (instead of assuming uniform), and then some plans would automatically be eliminated.
Imo a good thought experiment is to look at places that are hyper-optimized, like compilers. Would LLMs bring anything to the table (architecturally or performance-wise) to a piece of software that has been carefully crafted for decades? (Methinks no.)
The Postgres query planner has had to operate, for those same decades, in a much more realtime-sensitive and restricted environment than compilers. It can only draw its conclusions from summary statistics on tables in isolation, not on their relationships with each other (and even less so when filters are involved). For many cases this is fine! For many others it isn't.
There's a good number of heuristic choices in compilation where, maybe, you could get more optimal outcomes with machine learning - but at the cost of compilation resources, both time and space, and possibly determinism too.
As an example, register allocation is graph colouring, and thus NP complete; a model for producing an allocation plan is learning heuristics that might look at more features in combination than the ones hand-crafted into the compiler. An LLM for the job might do better than a more focused model like a GNN, due to sheer size, the effectiveness of transformers, or magic. But it probably won't do an overall better job than the handcrafted heuristics, because those handcrafted heuristics also tend to compile very, very fast with a small memory footprint, and can be debugged (more) easily when they go wrong.
Could you stitch a specific query plan to your view?
‘create plan llm_optimized …’
‘create view foo (select x, y, z from table bar) with plan llm_optimized’
I think the idea would be making a frontier model that does this. One that is trained on multiple queries
the currency of the counts also, right.. actually recounting table contents is done from time to time
This was a thoroughly enjoyable read, both the writing and presentation. I really liked the level of writing as it's basically introducing a whole lot of advanced topics but at just the right level for a non-AI researcher type of engineer like myself to be able to understand what's going on, and I felt it made some elements of LLMs actually something I could understand rather than wizardry done by maths PhDs. Probably because it's more like applied engineering rather than hard mathematics here. Thank you for a delightful post.
I have a theory that soon enough every code library will ship a CLI and a very tiny finetune for that specific lib alongside it
81%is nothing. It is not hard to be more than 3x better than Postgres [1]. And you don’t need a model to do that, let alone a 4B model. How much additional compute is needed to just run that model?
[1] https://github.com/datalevin/datalevin/tree/master/benchmark...
It's certainly a good idea to train a neural network to find good query plans but... an LLM??
An 8GB dataset? That’s literally a few seconds worth of records generated in my world, and any speedup at that scale is completely meaningless.
Let’s talk when you are looking at double digit TB at a minimum.
Funny enough I was thinking about something very similar to this based on the Jev model posted yesterday.
I’ve played with it already. I don’t think this is the use case. I think Jev’s use case is fast, cheap and somewhat easy classification. It’s not trainable in the way you would want here. Even though it’s fast it wont be faster than pgs query optimizer.
At least as I understand things.
How did you plan to use Jev for query optimization?
I am still struggling to understand a use-case for Jev. Isn't what was explained in this article a classification problem? I.e. find and aggregate data?
It’s classifying faster and cheaper. A lot of immediate ideas are better solved by pre-classifying + embedding, but their doom example or the wikipedia runs are one where you can’t preclassify.
The number of options has to be small and bounded. The query planning is more of a search/optimization problem than a classification problem since the number of options increases wildly based on query size.
Suppose you run a platform, and you run a couple thousand different queries of different types throughout the day. It would make sense to have an auto-optimizer that would read long queries, ponder over them with an LLM, come up with some good plans, and store them as hints. This seems like quite a good idea? Is there a product for this?
Mentioned elsewhere but classic ml seems the right tool for this problem
such a (visually) beautiful blogpost.
> How hard can it be?
> As it turns out: enormously hard.
This exactly tracks me learning everything
Really creative use of distillation here.
Nice article about how to train/fine-tune a language model.
However, it misses the whole point of database query planning. You can't just ignore the planning time itself, as if the database query were a static entity to be optimized once at a leisurely pace.
The real constraint on live query planners is quite different: they must improve the combined time - planning + query - based on live database statistics. You can amortize the planning with prepared statements, but that too is fraught since optimal plans can change quite frequently and based on input parameters. "Live" and "faster than the queries themselves" are the hard requirements to be considered a viable database query planner. This project does neither.
Aren't optimizations suppose to be deterministic?
They are not. Choice among several query plans depends on various summary statistics about the data, which might not be the most recent.
Including the input parameters.
It's not unusual for us to end up with bad query plans because the shape of our data can vary pretty greatly. In many cases, a Foo has 1 Bar. But in some cases, a Foo has a million Bars. That can cause the query optimizer to treat lookups on the bar table as if there are few elements there (causing a scan instead of a seek).
For the general case, the optimizer gets it right. However, the fringe case is one that causes the entire system to crash. It's a bit akin to how an insertion sort can be faster than quick sort when n is small. The optimizer might make a bad assumption about the size of n which makes it pick an expensive n lookup when log(n) is available (but slower for small n).
I'd like to contribute my amateur hour entry into this thread, although I did administer and develop mssql stuff for awhile.
sure optimizations based on stats, but the stats are the wildcard, in my experience query plans can change suddenly.
Queries are translated into plans according to statistics. However the transforms will be deterministic and should only change one valid plan to another. I could very easily see a neural network manipulate transforms the same way the current programming does, its just that the neural networks are by nature really nicely suitable because the "decisions" are based on training, and this training can be closed world type things like the ai assists that chess engines are now getting. Obviously ai still can't play chess but apparently its very good at ranking board positions just by developing that much statistical info because its training comes not from reading the web, but playing a gazzilian games against itself in a "closed" chess world of its own.
I'm thinking that the ai does "this legal transform of the query plan should be applied to this pattern of data (statistics, cardinality, etc)" simply because the ai encountered it in closed world training, much like the chess thing.
Just a theory tho feel free to correct!
> Obviously ai still can't play chess
I believe you are wrong on that. Do you mean large language models can’t play chess?
I don't believe that a pure neural network can currently abide by the rules of chess, even with unreasonable amounts of training. Do you have a counter example? I never mind having beliefs challenged with facts lol
edit: I think you could provide an AI with a service or skill that asks "is this move legal" but given all the overhead for llms or whatever to call a "legal move" service external to its process, well then you aren't really searching the tree very efficiently lol.
However if you just let a neural network score boards and the neural network is in the same process well then I think thats the working solution for using neural networks in chess. The net does not need to score all boards either, simple value based heuristics can obviously provide a preliminary list of good boards (moves) at a certain depth or ply and then select the move that produces the board that the neural net scores highest. I kinda sorta think thats whats done today but as usual I could be full of it lol
There's quite a big difference between "AI" and "pure neural networks." No, frontier chess (Alpha zero etc) is not "pure," because there's Monte Carlo tree search and a hard-coded game rules implementation.
But how will you know that the query plan actually does what your query asked for?
I would like to think that pg_hint_plan is designed in such a way that any hint it accepts must be a valid plan for the query. I’m quite confident that schemes with this property that can also express high quality plans are possible and not even excessively complicated.
This is not to say that it’s possible to genetically verify that a proposed algorithm does what you want it to — that would be undecidable or NP-hard or co-NP-hard depending on how you formulate the question.
I wouldn't be very excited about adding a 4B param model to my database deployment, but using this kind of approach while testing an app to identify query plans where Postgres is leaving performance on the table seems valuable without much risk.
Given the approach from the article, you can commit the hints to git and run tests for verification. The model would be used during coding.
If your statistics or workload change, this approach is useless. The hints are generated being generated ahead of time, taking 95hrs to do so.
`pg_hint_plan` has a debug log so you can verify Postgres actually used the hint or not! Used this during evaluations
You're misunderstanding the setup here. The LLM doesn't modify the query, just some details about how to choose between different ways to break the query into basic operations on the tables. The SQL doesn't change. It's still up to postgres to guarantee that the results match the query. If the proposed plan were nonsense that didn't amount to carrying out the query, postgres would ignore it.
I imagine you can perform operations on query plans to transform them and determine equivalence?
you'll have to prove equivalence through some Lean4 code perhaps? or some weird clause tree comparisons... good question indeed.
Because P!=NP (very probably IMHO) there's a huge class of problem for which LLMs are useful on the expensive and heuristic-y generation side because the verification is relatively inexpensive.
[flagged]
[flagged]
[flagged]
[dead]
[dead]
[dead]
[dead]
TL:DR; for people. Index your data properly.
It’s 8G of data in memory. You can raw dog that no problem.
(Indexes would be nice though)
This is about getting the planner to more effectively use your 'proper indexes', concerned with the edge cases where it currently does not.
why is this write-up so long?
Need 5 days just to go through it.
Think of it like a paper. You wouldn't ask why a paper was so long.
Also you can now ask AI to summarise it for you and even probe with questions pertaining to your specific interests.
Why wouldn't you ask that? That's exactly why papers have supplementary material
> You wouldn't ask why a paper was so long.
I've asked myself that question a whole bunch of times