How we used pg_query to rewrite queries to fix bad query plans
Rewriting SQL queries programmatically is harder than it looks. As a human, adding an extra AND condition to a WHERE clause is simple enough. But doing the same thing in code quickly gets complicated. You might try regex, but the real difficulty is coming up with a pattern that works for every variation of a query. AI could generate plausible rewrites, but it's hard to guarantee correctness. These rewrites may look valid, but SQL has many subtle corner cases, so it's difficult to prove that the transformed query always behaves identically.
As we are developing the new Query Advisor feature in pganalyze, we need a way to take query insights one step further: not only highlight potential issues, but also suggest alternative query patterns. To do that safely, we turn to pg_query.
Using the pg_query open source library, you can parse a query into a structured parse tree, tweak it at the tree level, and then regenerate valid SQL. It ensures the output is deterministic and syntactically correct. With a recent change, it will also support pretty-printing with configurable indentation and line length, making rewrites more powerful and easier to read.
In this post, we will show a few examples of how you can use pg_query to rewrite queries, starting from a simple demonstration and then moving on to real-world patterns that benefit from rewriting.
Query rewrite 101
Let's walk through a really simple case of using pg_query to rewrite a query. Bindings are available for Ruby, Rust and Go, as well as community-maintained ports for Node.js and Python. If you want to learn more about the basics of pg_query, check out our past blog post. In this blog post, we'll use the Ruby bindings.
Let's start with a simple query:
require "pg_query"
parsed_query = PgQuery.parse("SELECT id FROM tbl1")
# => #<PgQuery::ParserResult:0x000000015da1dc50
# @aliases=nil,
# @cte_names=nil,
# @functions=nil,
# @query="SELECT id FROM tbl1",
# @tables=nil,
# @tree=
# <PgQuery::ParseResult: version: 170005, stmts: [<PgQuery::RawStmt: stmt: <PgQuery::Node: select_stmt:
# <PgQuery::SelectStmt: distinct_clause: [],
# target_list: [<PgQuery::Node: res_target: <PgQuery::ResTarget: name: "", indirection: [], val: <PgQuery::Node: column_ref: <PgQuery::ColumnRef: fields: [<PgQuery::Node: string: <PgQuery::String: sval: "id">>], location: 7>>, location: 7>>],
# from_clause: [<PgQuery::Node: range_var: <PgQuery::RangeVar: catalogname: "", schemaname: "", relname: "tbl1", inh: true, relpersistence: "p", location: 15>>],
# group_clause: [], group_distinct: false, window_clause: [], values_lists: [], sort_clause: [],
# limit_option: :LIMIT_OPTION_DEFAULT, locking_clause: [], op: :SETOP_NONE, all: false>>,
# stmt_location: 0, stmt_len: 0>]>,
# @warnings=[]>
parsed_query.tables
# => ["tbl1"]Here, parsed_query is a parse result that contains a parse tree. It also exposes useful methods, such as tables, which tells us which tables are used in the query.
The parse tree for this query looks like the following:
We can either walk the tree to visit nodes (which we'll cover later), or drill down directly to a specific node. For example, to reach the table name of the from clause:
parsed_query.tree.stmts[0].stmt.select_stmt.from_clause[0].range_var.relname
# => "tbl1"Updating this lets us change the table name. After updating the node, we can call deparse to generate SQL again:
parsed_query.tree.stmts[0].stmt.select_stmt.from_clause[0].range_var.relname = "tbl2"
parsed_query.deparse
# => "SELECT id FROM tbl2"Now that we have the basic idea of how rewriting works, let's move on to more practical examples.
Example #1 - Add +0 to ORDER BY to avoid index misuse
When you find a slow query using ORDER BY combined with LIMIT, it's important to check whether the planner is picking the right index. This is something we verify in the Query Advisor feature.
Let's start with a simple query:
SELECT * FROM items WHERE object_id = 123 LIMIT 1With this query, when the items table has an object_id index (e.g. items_object_id_idx), the planner will usually use it, and the query should finish quickly, as long as the index is selective.
Now, let's add an ORDER BY:
SELECT * FROM items WHERE object_id = 123 ORDER BY id LIMIT 1In some cases, this can cause the planner to choose a plan like "Index Scan Backward using items_pkey on items", and then filter out rows where object_id = 123. If many rows are removed by that filter, the query can become significantly slower.
A simple workaround is to add "+0" to the ORDER BY id. This prevents the planner from using the primary key index (items_pkey).
SELECT * FROM items WHERE object_id = 123 ORDER BY id + 0 LIMIT 1Let's create a parse tree from the query (without "+0") and look at the "ORDER BY id" part:
parsed_query = PgQuery.parse('SELECT * FROM items WHERE object_id = 123 ORDER BY id LIMIT 1')
parsed_query.tree.stmts[0].stmt.select_stmt.sort_clause.sort_by[0].sort_by.node
# => <PgQuery::Node: column_ref: <PgQuery::ColumnRef: fields: [<PgQuery::Node: string: <PgQuery::String: sval: "id">>], location: 51>>It's a bit hard to read, but the sort_by node here is a ColumnRef node pointing to id. To add "+0", we replace it with an A_Expr node that represents a binary expression with id on the left and 0 on the right.
In the below code, create a new A_Expr node:
sort_by_node = parsed_query.tree.stmts[0].stmt.select_stmt.sort_clause.sort_by[0].sort_by.node
new_node = PgQuery::Node.new(
a_expr: PgQuery::A_Expr.new(
kind: :AEXPR_OP,
name: [PgQuery::Node.new(string: PgQuery::String.new(sval: '+'))],
lexpr: sort_by_node.dup, # Note: to reuse existing nodes, make sure to duplicate to avoid accidentally modifying the original tree
rexpr: PgQuery::Node.new(a_const: PgQuery::A_Const.new(ival: PgQuery::Integer.new(ival: 0)))
)
)Finally, we assign the new node and deparse the query. Don't forget to use the new pretty-printing options:
parsed_query.tree.stmts[0].stmt.select_stmt.sort_clause.sort_by[0].sort_by.node = new_node
opts = PgQuery::DeparseOpts.new(pretty_print: true, indent_size: 2, trailing_newline: true)
parsed_query.deparse(opts: opts)
# => "SELECT *\nFROM items\nWHERE object_id = 123\nORDER BY id + 0\nLIMIT 1\n"For more on why this rewrite helps, see our blog post Postgres Planner Quirks: The impact of ORDER BY + LIMIT on index usage.
Example #2 - Transform multiple OR clauses to ANY
With Postgres 18, the planner can transform certain chains of OR comparisons into ANY, which can produce a better plan (https://postgr.es/c/ae4569161). This happens at the planner level, but let's take a look at how we can do the same thing explicitly using a pg_query rewrite.
Let's look at a query that compares the id column to multiple constants:
EXPLAIN SELECT id FROM items WHERE id = 41 OR id = 42 OR id = 43;
QUERY PLAN
---------------------------------------------------------------------------------
Bitmap Heap Scan on items (cost=12.89..24.37 rows=3 width=8)
Recheck Cond: ((id = 41) OR (id = 42) OR (id = 43))
-> BitmapOr (cost=12.89..12.89 rows=3 width=0)
-> Bitmap Index Scan on items_pkey (cost=0.00..4.29 rows=1 width=0)
Index Cond: (id = 41)
-> Bitmap Index Scan on items_pkey (cost=0.00..4.29 rows=1 width=0)
Index Cond: (id = 42)
-> Bitmap Index Scan on items_pkey (cost=0.00..4.29 rows=1 width=0)
Index Cond: (id = 43)
(9 rows)Now, let's rewrite it with ANY:
EXPLAIN SELECT id FROM items WHERE id = ANY('{41,42,43}');
QUERY PLAN
----------------------------------------------------------------------------------
Index Only Scan using items_pkey on items (cost=0.29..12.92 rows=3 width=8)
Index Cond: (id = ANY ('{41,42,43}'::bigint[]))
(2 rows)Notice how the cost dropped from 24.37 down to 12.92. The query returns the same results, but instead of three Bitmap Index Scans, it uses a single Index Only Scan. Let's take a closer look at the parse tree.
At the parse tree level, the first query is represented as a bool_expr (BoolExpr) with OR_EXPR, containing three equality expressions (id = 41, id = 42, id = 43).
The ANY form, on the other hand, is represented as an a_expr (A_Expr) with ANY. This corresponds to = ANY(array), with the column id on the left and an array of constants on the right ({41,42,43}).
The rewrite steps are:
- Find an OR expression made up of multiple
id = <const>comparisons - Collect all the constants
- Replace the
BoolExprwith anA_Exprrepresentingid = ANY(array)
In the example code below, to simplify the replace step, it replaces the matching args elements within the BoolExpr node with a single A_Expr node, collapsing the OR chain into = ANY(...).
In Example #1, we just drilled down to a node and swapped it. This time, let's walk the whole tree, find any matching pattern, and rewrite it.
def transform_or_to_any(query)
parsed_query = PgQuery.parse(query)
parsed_query.walk! do |node|
# Find the BoolExpr node with OR_EXPR
next unless node.is_a?(PgQuery::BoolExpr) && node.boolop == :OR_EXPR
keep_as_is = []
group_by_lexpr = {}
node.args.each do |arg|
# Note: only group when the arg is ColumnRef = A_Const (e.g. col1 = 123)
# For other cases (e.g. col1 IS TRUE, col1 != 345), leave it as is
if arg.node == :a_expr &&
arg.a_expr.name.first.node == :string &&
arg.a_expr.name.first.string.sval == '=' &&
arg.a_expr.lexpr.node == :column_ref &&
arg.a_expr.rexpr.node == :a_const
# In order to use this as a hash key, remove the location info by setting to 0
arg.a_expr.lexpr.inner.location = 0
group_by_lexpr[arg.a_expr.lexpr] ||= []
group_by_lexpr[arg.a_expr.lexpr] << arg.a_expr.rexpr.dup
else
keep_as_is << arg.dup
end
end
# No multiple ORs with the same column (lexpr)
next unless group_by_lexpr.any? { |k, v| v.length > 1 }
# Create new args with AEXPR_OP_ANY a_expr for grouped args
any_args = []
group_by_lexpr.each do |lexpr, rexprs|
if rexprs.length == 1
keep_as_is << PgQuery::Node.new(
a_expr: PgQuery::A_Expr.new(
kind: :AEXPR_OP,
name: [PgQuery::Node.new(string: PgQuery::String.new(sval: '='))],
lexpr: lexpr,
rexpr: rexprs.first
)
)
else
any_args << PgQuery::Node.new(
a_expr: PgQuery::A_Expr.new(
kind: :AEXPR_OP_ANY,
name: [PgQuery::Node.new(string: PgQuery::String.new(sval: '='))],
lexpr: lexpr,
rexpr: PgQuery::Node.new(a_const: QueryParameters.values_to_array(rexprs))
)
)
end
end
node.args.replace(any_args + keep_as_is)
end
parsed_query.deparse
endNow let's try the transform on some variations:
# Simple case
transform_or_to_any('SELECT id FROM items WHERE id = 41 OR id = 42 OR id = 43')
# => "SELECT id FROM items WHERE id = ANY('{41,42,43}')"
# With AND and OR
transform_or_to_any('SELECT id FROM items WHERE id = 41 OR id = 42 AND id = 43 OR id = 44')
# => "SELECT id FROM items WHERE id = ANY('{41,44}') OR (id = 42 AND id = 43)"
# ORs in subqueries or UNIONs
transform_or_to_any(<<~SQL)
SELECT id FROM items WHERE id IN (SELECT id FROM items2 WHERE id = 41 OR id = 42)
UNION
SELECT id FROM items3 WHERE id = 43 OR id = 44
SQL
# => "SELECT id FROM items WHERE id IN (SELECT id FROM items2 WHERE id = ANY('{41,42}')) UNION SELECT id FROM items3 WHERE id = ANY('{43,44}')"You can see that it transforms ORs properly no matter where they appear in the query.
Conclusion
The examples we looked at, such as adding an expression to influence index usage or transforming multiple OR clauses into ANY, show how query rewriting with pg_query can solve real-world problems in a safe and consistent way. These are only a few cases, and the same approach can be applied to many other kinds of transformations.
In developing the Query Advisor feature, rewriting the query using pg_query has been an essential piece. We hope that sharing these examples encourages you to explore what is possible with this library in your own projects.