Proposal to remove Protobuf (de)serialization

[levkk]

Hi team!

Happy New Year, hope everyone is having a good holiday break!

I had some spare time in the past few days and decided to see what would happen if I removed Protobuf (de)serialization and replaced it with raw C/Rust conversions. Basically, I had to get the AST nodes from the Postgres source code, generate Rust structs (with bindgen) and write conversion functions from C to Rust, and back (to support deparse).

The results are really great:

1. 5x performance improvement for parsing queries (string to AST)
2. 9x improvement for deparsing (AST to string)

Here's raw numbers from our benchmark:

Parse

============================================================
            parse_raw vs parse Benchmark
============================================================
Query: 2725 chars (CTEs + JOINs + subqueries + window functions)

┌─────────────────────────────────────────────────────────┐
│                    RESULTS                              │
├─────────────────────────────────────────────────────────┤
│  parse_raw (direct C struct reading):                   │
│    Iterations:          6800                            │
│    Total time:         2.03s                            │
│    Per iteration:     297.88 μs                         │
│    Throughput:          3357 queries/sec                │
├─────────────────────────────────────────────────────────┤
│  parse (protobuf serialization):                        │
│    Iterations:          1300                            │
│    Total time:         2.11s                            │
│    Per iteration:    1623.70 μs                         │
│    Throughput:           616 queries/sec                │
├─────────────────────────────────────────────────────────┤
│  COMPARISON                                             │
│    Speedup:             5.45x faster                    │
│    Time saved:       1325.82 μs per parse               │
│    Extra queries:       2741 more queries/sec           │
└─────────────────────────────────────────────────────────┘

Deparse

============================================================
           deparse_raw vs deparse Benchmark
============================================================
Query: 2725 chars (CTEs + JOINs + subqueries + window functions)

┌─────────────────────────────────────────────────────────┐
│                    RESULTS                              │
├─────────────────────────────────────────────────────────┤
│  deparse_raw (direct C struct building):                │
│    Iterations:         14700                            │
│    Total time:         2.01s                            │
│    Per iteration:     136.63 μs                         │
│    Throughput:          7319 queries/sec                │
├─────────────────────────────────────────────────────────┤
│  deparse (protobuf serialization):                      │
│    Iterations:          1600                            │
│    Total time:         2.11s                            │
│    Per iteration:    1317.66 μs                         │
│    Throughput:           759 queries/sec                │
├─────────────────────────────────────────────────────────┤
│  COMPARISON                                             │
│    Speedup:             9.64x faster                    │
│    Time saved:       1181.03 μs per deparse             │
│    Extra queries:       6560 more queries/sec           │
└─────────────────────────────────────────────────────────┘

The code is in our fork here: https://github.com/pgdogdev/pg_query.rs. We also forked libpg_query to add the necessary C wrappers: https://github.com/pgdogdev/libpg_query/.

So, what's the catch? Non-exhaustive list:

1. The node conversions are all AI-generated (Claude). This is often a deal breaker for teams, so just letting you know early on.
2. The conversions are fully recursive - it can blow the stack on large queries. We set it pretty high in PgDog to avoid this issue (32MiB right now, which is pretty high imo)
3. I blew the stack a few times because the generated code incorrectly interpreted some nodes which caused infinite recursion. Our CI caught those but we would need to basically come up with queries that use any and all Postgres SQL features - not terribly hard, but tedious and easy to make a mistake.

I also tried implementing a heap-based iterative approach, but that ended up being slower that the protobuf version (surprisingly, or not perhaps?), so I left the recursive one in place for now. It's easier to read too, tbh. We can probably add a recursion limit, just like prost does, to prevent overflows.

I know you guys maintain a large number of bindings for libpg_query and protobuf works great for making the interface between the C lib and the bindings work. Maintaining raw bindings between all languages could be difficult. Although, these days, you could probably AI-generate them in a few days/weeks, but there is a pretty big maintenance burden there for sure. Let me know if this is something that might be interesting and if yes, I'm happy to work with you to merge our fork into your code!

I'm going to be testing this with PgDog in the coming weeks and months and let you know how it goes.

Cheers!

[lfittl]

Happy New Year, hope everyone is having a good holiday break!

Happy new year @levkk!

I had some spare time in the past few days and decided to see what would happen if I removed Protobuf (de)serialization and replaced it with raw C/Rust conversions. Basically, I had to get the AST nodes from the Postgres source code, generate Rust structs (with bindgen) and write conversion functions from C to Rust, and back (to support deparse).

The results are really great:

1. 5x performance improvement for parsing queries (string to AST)

2. 9x improvement for deparsing (AST to string)

Nice, thank you for the proposal! - the performance overhead of the Protobuf serialization/deserialization is definitely a known issue, and is why we recently introduced the summary functions that avoid the Protobuf conversion for certain use cases in pganalyze (e.g. getting the referenced tables from a query, "smart" truncating it, etc.), and we saw pretty big performance benefits there.

But from having looked at how you use pg_query.rs in pgDog before, I know that using summary functions won't really work all that well, and it doesn't seem feasible to handle all the use cases. So I think its worthwhile to consider how we can integrate your proposal.

The code is in our fork here: https://github.com/pgdogdev/pg_query.rs. We also forked libpg_query to add the necessary C wrappers: https://github.com/pgdogdev/libpg_query/.

FWIW, if you wanted to reduce the diff to our upstream to maintain here, I'd be open to getting the libpg_query part here integrated earlier - that seems more broadly applicable to any caller that wants to work with the AST in C.

So, what's the catch? Non-exhaustive list:

1. The node conversions are all AI-generated (Claude). This is often a deal breaker for teams, so just letting you know early on.

I think this is a deal-breaker from a maintenance perspective - I don't necessarily mind using Claude for generating code that generates code (as long as the generation code itself is human reviewed) - but generating the struct translation code with Claude on each release seems brittle and prone to have undiscovered errors.

Here two different ideas on how to take this forward:

1. If staying with the current approach (translating the C structs to the Protobuf Rust structs), I would rework it so that the generation happens based on the node definitions, so they are guaranteed to match.

This is what we do to generate the Protobuf definitions on each release (https://github.com/pganalyze/libpg_query/blob/17-latest/scripts/generate_protobuf_and_funcs.rb), and you can use that same data from the libpg_query srcdata folder to generate all the functions that you have in src/raw_deparse.rs.

2. It might be worth considering not using the Protobuf structs, and instead just using Rust bindgen on the original node definitions and working with what was generated.

This is probably a bit less ergonomic (since you'll run into more Postgres C concepts like bitmapsets, lists, etc), but would presumably be noticeably faster, since you don't have a translation step. This is closer to what pgrx does, and what we also do internally in pganalyze for what we call pg_plan (library for running the planner like pg_query runs the parser). If this is of interest we could probably open-source a few of the wrapper helpers around pg_list and such.

I know you guys maintain a large number of bindings for libpg_query and protobuf works great for making the interface between the C lib and the bindings work. Maintaining raw bindings between all languages could be difficult. Although, these days, you could probably AI-generate them in a few days/weeks, but there is a pretty big maintenance burden there for sure.

Yeah, I'm sure we could generate them, but I don't think the maintenance trade-off is worth it - and I think its potentially okay for the Rust implementation to work a bit differently than the other bindings (i.e. do something fast here, but don't change the other ones).

I'm going to be testing this with PgDog in the coming weeks and months and let you know how it goes.

Great, excited to see how this works for you!

[levkk]

FWIW, if you wanted to reduce the diff to our upstream to maintain here, I'd be open to getting the libpg_query part here integrated earlier - that seems more broadly applicable to any caller that wants to work with the AST in C.

Sounds good, I'll put up a PR for that one. There is a bit of copy/pasta going on there because I wanted to leave the protobuf paths untouched, but maybe we can merge the two.

I don't necessarily mind using Claude for generating code that generates code (as long as the generation code itself is human reviewed)

Absolutely. My original plan was to use macros to generate conversion code for each field, but there is still the giant case/match statement for each node.

It might be worth considering not using the Protobuf structs, and instead just using Rust bindgen on the original node definitions and working with what was generated.

The protobuf structs are convenient because they are all using Rust-native data types, while the bindgen-generated ones are basically just mutable pointers to C memory wrapped in a Rust struct with repr(C) on it, so a conversion to safe Rust types is still needed. We could probably use macros to generate those, thought, but that's pretty much what Prost does.

presumably be noticeably faster, since you don't have a translation step.

Might be a bit - no additional memory allocation is great, but we would need to do quite a bit of pointer arithmetic and casting to get to the data inside of each node.

This is probably a bit less ergonomic (since you'll run into more Postgres C concepts like bitmapsets, lists, etc), but would presumably be noticeably faster, since you don't have a translation step. This is closer to what pgrx does, and what we also do internally in pganalyze for what we call pg_plan (library for running the planner like pg_query runs the parser). If this is of interest we could probably open-source a few of the wrapper helpers around pg_list and such.

Yeah, that would be good I think. It feels like the right solution. If you guys could open source those wrappers when you have a moment, we could give it a look.

Yeah, I'm sure we could generate them, but I don't think the maintenance trade-off is worth it - and I think its potentially okay for the Rust implementation to work a bit differently than the other bindings (i.e. do something fast here, but don't change the other ones).

Agreed. The Postgres API will keep us honest here. In dynamic languages like Ruby, Python, etc, the performance benefits won't be noticeable anyway. Added benefit of keeping protobuf around is we have the same contract between all libpg_query wrappers, we're just substituting Prost with a bit of C/Rust glue.