Use prepared statements for primary key lookups

Problem

PgBouncer in transaction pooling mode does not support prepared statements. This means that each database query we send to PG will have to be analyzed and planned by PG. The overhead can be significant for very simple statements, such as primary key lookups:

SELECT * FROM issues WHERE id = 1 LIMIT 5;
Limit  (cost=0.43..3.45 rows=1 width=1520) (actual time=0.033..0.034 rows=1 loops=1)
  Buffers: shared hit=4
  ->  Index Scan using index_issues_on_id_and_weight on issues  (cost=0.43..3.45 rows=1 width=1520) (actual time=0.032..0.032 rows=1 loops=1)
        Index Cond: (id = 13991538)
        Buffers: shared hit=4
Planning Time: 0.371 ms
Execution Time: 0.083 ms

As you can see the planning time is 4 times slower than the execution time.

Idea

Pick the most frequently executed primary key lookup queries and turn them into function calls. SQL queries in functions are using prepared statements implicitly so we will be to leverage prepared statements. This technique was also mentioned at PGConf 2022.

This is an example function for locating one issue:

CREATE OR REPLACE FUNCTION issueById(issueId int)
  RETURNS issues AS $$
BEGIN
  return (SELECT issues FROM issues WHERE id = issueId LIMIT 1);
END;
$$ LANGUAGE plpgsql STABLE PARALLEL SAFE LEAKPROOF COST 1;

The function returns one row, which can be used in DB queries the following way:

select * from issueById(4);

Note: the function needs a bit of tweaking because, for a non-existing id, it returns an empty row.

Planning is still required, however, planning a function call seems to be much cheaper:

 Function Scan on issuebyid  (cost=0.00..0.01 rows=1 width=820) (actual time=0.075..0.076 rows=1 loops=1)
   Buffers: shared hit=2
 Planning Time: 0.040 ms
 Execution Time: 0.099 ms

Local perf test

I used PGBench to run 10K transactions of both versions of the queries:

  • SQL query TPS: 8711
  • Function call TPS: 15074

I see 30-40% speed improvement by running the test locally.

Implementation

  1. Collect a list of primary key lookup queries which are executed frequently on PRD.
  2. Measure the queries on a production-like system to make sure we would actually improve things.
  3. Review the function definition with PG experts.
  4. Implement the change behind a feature flag.
  5. Document the change.

We can rewrite DB queries by patching ActiveRecord. Pseudo code:

if issues_primary_key_lookup_query && rewrite_feature_enabled
  primary_key_value = # get it from the binding
  
  %{select * from issueById(#{primary_key_value})}
else
  super
end

Another way, write the actual query using AR:

Issue.from("issueById(4) as issues").to_a.first
Edited by Adam Hegyi