Error when encountering an unsupported aggregate function#1020
Draft
sgrif wants to merge 5 commits into
Draft
Conversation
Contributor
Codecov Report✅ All modified and coverable lines are covered by tests. 📢 Thoughts on this report? Let us know! |
levkk
approved these changes
May 28, 2026
levkk
reviewed
May 28, 2026
This change allows the query rewriter to safely modify its behavior for direct-to-shard queries, and skip rewrite steps that are only needed for cross-shard queries. Though this commit does not take advantage of this, future changes will require this behavior and changing the cache key required enough code to justify being pulled into its own commit. The signature of Map::get and Borrow means this needs a bit more boilerplate than I would have liked. There is no way to implement Borrow<(str, _)> for (String, _) (or any other owned/ref types). Instead we have to erase the types to a trait object so we can get individual references to the fields when we need them for comparison and hashing. This trait is implemented for any combination of values/refs, and a special case for (Arc<A>, B). If in the future this trait is needed with additional pointer types, or Arc in the right position, additional manual impls will need to be added.
When handling aggregate functions in cross-shard queries, we can only return a correct result for functions that we have explicit support for. No matter what the aggregate function is doing, we need to know how to combine the results from the separate shards. Prior to this change, we would just silently do the wrong thing, treating this as any other non-aggregate expression and just returning a union of the rows from each query. We now explicitly check if an aggregate function with that name exists and error if we don't recognize it. This is future proofed against new functions in later postgres versions, as well as user defined aggregates (which we can likely never support). This will produce a false positive if a function exists and is defined as aggregate for some argument types but not others. But frankly anyone doing that is asking for trouble. This logic is objectively in the wrong place. What we are doing here is validation, not parsing. However as I'm familiarizing myself with these code paths and preparing a larger rearchitecture, I'm slowly hammering things into a shape that's easier to move around. I do not intend on leaving this logic here long term.
While we can't completely treat direct-to-shard queries the same as unsharded queries, we can at least skip certain parts of the rewriter. In particular, the aggregate rewrite does some extra work that will never be needed on direct-to-shard, and skipping it allows us to error on aggregate functions we don't support without breaking the very niche case of an unsupported aggregate function being used by a user only in direct-to-shard queries. Actually implementing this was a bit of a PITA. Since we cache the AST after the rewriter mutates it, as opposed to just the results of pg_query::parse, we needed to add whether the query is direct-to-shard to the cache key (#1027) The natural place to exit early is the same place we do the "skip this if the schema only has one shard", but that information was previously lost all the way up at the cache impl, so we needed to pass this information through quite a few levels of indirection. Ultimately I think we need to separate out parsing from rewriting more concretely, and structure rewriting much differently with a better structured context. However, as I work towards being ready to do that larger restructuring, putting this where it's most natural to make the dependency graph clear seems like the right short-term path forward.
Contributor
Author
|
(Changing base to #1027 so this PR's changes are easier to review, since that branch needs work, and this PR is reviewable separately even if it's not yet ready to merge) |
fb4cfe1 to
ed257bc
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
When handling aggregate functions in cross-shard queries, we can only return a correct result for functions that we have explicit support for. No matter what the aggregate function is doing, we need to know how to combine the results from the separate shards.
Prior to this change, we would just silently do the wrong thing, treating this as any other non-aggregate expression and just returning a union of the rows from each query. We now explicitly check if an aggregate function with that name exists and error if we don't recognize it.
This is future proofed against new functions in later postgres versions, as well as user defined aggregates (which we can likely never support). This will produce a false positive if a function exists and is defined as aggregate for some argument types but not others. But frankly anyone doing that is asking for trouble.
This logic is objectively in the wrong place. What we are doing here is validation, not parsing. However as I'm familiarizing myself with these code paths and preparing a larger rearchitecture, I'm slowly hammering things into a shape that's easier to move around. I do not intend on leaving this logic here long term.