Allow user defined SQL planners to be registered #11208
Merged
Conversation
This file contains 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
github-actions
bot
added
logical-expr
13 tasks
alamb
approved these changes
Jul 2, 2024
There was a problem hiding this comment.
Thank you @samuelcolvin -- this looks great to me
cc @jayzhan211
| @@ -99,6 +100,8 @@ pub struct SessionState { | |||
| session_id: String, | |||
| /// Responsible for analyzing and rewrite a logical plan before optimization | |||
| analyzer: Analyzer, | |||
| /// Provides support for customising the SQL planner, e.g. to add support for custom operators like `->>` or `?` | |||
| use datafusion_expr::planner::{PlannerResult, RawBinaryExpr, UserDefinedSQLPlanner}; | ||
| use datafusion_expr::BinaryExpr; | ||
|
|
||
| struct MyCustomPlanner; |
| expr: RawBinaryExpr, | ||
| _schema: &DFSchema, | ||
| ) -> Result<PlannerResult<RawBinaryExpr>> { | ||
| match &expr.op { |
|
|
||
| struct MyCustomPlanner; | ||
|
|
||
| impl UserDefinedSQLPlanner for MyCustomPlanner { |
There was a problem hiding this comment.
Suggested change
| impl UserDefinedSQLPlanner for MyCustomPlanner { | |
| /// Planner that plans `->` and `-->` specially | |
| impl UserDefinedSQLPlanner for MyCustomPlanner { |
| // custom planners are registered first, so they're run first and take precedence over built-in planners | ||
| for planner in self.user_defined_sql_planners.iter() { | ||
| query = query.with_user_defined_planner(planner.clone()); | ||
| } | ||
|
|
||
| // register crate of array expressions (if enabled) | ||
| #[cfg(feature = "array_expressions")] |
There was a problem hiding this comment.
It might make sense to move the registering of these planners into the constructor of SessionState rather than here (so users could potentially more carefully control the planner content). But we can do that as a follow on PR
jayzhan211
approved these changes
Jul 2, 2024
findepi
pushed a commit
to findepi/datafusion
that referenced
this pull request
Jul 16, 2024
* Allow user defined SQL planners to be registered * fix clippy, remove unused Default * format
This pull request was closed.
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.
Rationale for this change
Follow up to #11180 that allows
UserDefinedSQLPlanners to be registered onFunctionRegistrys, mostly copied from #11137.Are these changes tested?
Yes, I've added a test
Are there any user-facing changes?
Adds
FunctionRegistry::register_user_defined_sql_planner.