Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Allow user defined SQL planners to be registered #11208

Merged

Conversation

samuelcolvin
Copy link
Contributor

Rationale for this change

Follow up to #11180 that allows UserDefinedSQLPlanners to be registered on FunctionRegistrys, 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.

@github-actions github-actions bot added logical-expr Logical plan and expressions core Core DataFusion crate labels Jul 2, 2024
Copy link
Contributor

@alamb alamb left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 `?`
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

👍

use datafusion_expr::planner::{PlannerResult, RawBinaryExpr, UserDefinedSQLPlanner};
use datafusion_expr::BinaryExpr;

struct MyCustomPlanner;
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

❤️

expr: RawBinaryExpr,
_schema: &DFSchema,
) -> Result<PlannerResult<RawBinaryExpr>> {
match &expr.op {
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

this is awesome


struct MyCustomPlanner;

impl UserDefinedSQLPlanner for MyCustomPlanner {
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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")]
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Copy link
Contributor

@jayzhan211 jayzhan211 left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

👍
Thanks @samuelcolvin and @alamb

@jayzhan211 jayzhan211 merged commit 43ea682 into apache:main Jul 2, 2024
23 checks passed
@samuelcolvin samuelcolvin deleted the register_user_defined_sql_planners branch July 2, 2024 14:40
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
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
core Core DataFusion crate logical-expr Logical plan and expressions
Projects
None yet
Development

Successfully merging this pull request may close these issues.

3 participants