Introduce user defined SQL planner API #11180

jayzhan211 · 2024-06-30T02:57:05Z

Which issue does this PR close?

Closes #11168
part of #10534
Closes #11137

TLDR is to make it easier for users to customize sql planning, specifically to handle operator translation

Rationale for this change

Let sql planner customizable
Follow up from #11168

What changes are included in this PR?

Array function rewrite
GetField
Array literal

TODO in next PRs

#10534, to close this, we have more functions to move.

date_part
create_struct
create_named_struct
sql_overlay_to_expr
sql_position_to_expr
sql_substring_to_expr
sql_compound_identifier_to_expr

Are these changes tested?

Are there any user-facing changes?

Signed-off-by: jayzhan211 <[email protected]>

jayzhan211 · 2024-06-30T05:47:35Z

datafusion/core/src/execution/session_state.rs

+                Arc::new(ArrayFunctionPlanner::try_new(provider).unwrap()) as _;
+
+            let field_access_planner =
+                Arc::new(FieldAccessPlanner::try_new(provider).unwrap()) as _;


get_field for struct does not need array expression, but I'm not sure should we move it out

Signed-off-by: jayzhan211 <[email protected]>

jayzhan211 · 2024-06-30T08:03:30Z

Interesting, I got stack overflow in CI but not in local.

Signed-off-by: jayzhan211 <[email protected]>

alamb

Thank you @jayzhan211 this is very cool and I think close

I left a bunch of stylistic / naming / comment suggestions -- let me know what you think

datafusion/expr/src/planner.rs

alamb · 2024-06-30T11:16:10Z

datafusion/expr/src/planner.rs

+}
+
+/// This trait allows users to customize the behavior of the SQL planner
+pub trait UserDefinedPlanner {


I was thinking about different names for this struct that might be more specific to SQL and thus make it easier to find

How about UserDefinedSQLPlanner? or SQLPlannerExtensions?

alamb · 2024-06-30T11:20:00Z

datafusion/expr/src/planner.rs

+    }
+}
+
+pub struct BinaryExpr {


I think it may be confusing if this struct is also called BinaryExpr given https://docs.rs/datafusion/latest/datafusion/logical_expr/struct.BinaryExpr.html. What do you think about callng this RawBinaryExpr or maybe SqlBinaryExpr (and then rename FieldAccessExpr similarly)?

Here is some suggested comments to explain the structure (it would be good to do the same for FieldAccessExpr too):

Suggested change

pub struct BinaryExpr {

/// An operator with two arguments to plan

///

/// Note `left` and `right` are DataFusion [`Expr`]s but the `op` is the SQL AST operator.

/// This structure is used by [`UserDefinedPlanner`] to plan operators with custom expressions.

pub struct BinaryExpr {

I also struggle with this naming too 😆 . I will take RawBinaryExpr

alamb · 2024-06-30T11:21:49Z

datafusion/expr/src/planner.rs

+pub enum PlannerSimplifyResult {
+    /// The function call was simplified to an entirely new Expr
+    Simplified(Expr),
+    /// the function call could not be simplified, and the arguments
+    /// are return unmodified.
+    OriginalBinaryExpr(BinaryExpr),
+    OriginalFieldAccessExpr(FieldAccessExpr),
+    OriginalArray(Vec<Expr>),
+}


I think this might be easier to use if it were generic so that the callsites would know exactly the type of the returns Original

SO something like

Suggested change

pub enum PlannerSimplifyResult {

/// The function call was simplified to an entirely new Expr

Simplified(Expr),

/// the function call could not be simplified, and the arguments

/// are return unmodified.

OriginalBinaryExpr(BinaryExpr),

OriginalFieldAccessExpr(FieldAccessExpr),

OriginalArray(Vec<Expr>),

}

pub enum PlannerSimplifyResult<T> {

/// The structure was planned as an entirely new Expr by the planner

Simplified(Expr),

/// the planner did not handle planning the structure, and it is returned

/// unmodified.

OriginalBinaryExpr(T),

}

alamb · 2024-06-30T11:25:31Z

datafusion/sql/src/expr/mod.rs

+                            field_access_expr = expr;
+                        }
+                        _ => {
+                            return exec_err!("Unexpected result encountered. Did you expect an OriginalFieldAccessExpr?")


We can avoid this type of check by using a generic in PlannerSimplifyResult

alamb · 2024-06-30T11:26:07Z

datafusion/expr/src/planner.rs

+    pub expr: Expr,
+}
+
+pub enum PlannerSimplifyResult {


Since this isn't really "simplifying" anything maybe we could call it PlannerResult instead of PlannerSimplifyResult

alamb · 2024-06-30T11:26:46Z

datafusion/sql/src/planner.rs

-    /// Get all user defined window function names
-    fn udwf_names(&self) -> Vec<String>;
-}
+pub use datafusion_expr::planner::ContextProvider;


👍 for backwards compatibility

alamb · 2024-06-30T11:28:04Z

datafusion/sql/src/planner.rs

+    /// user defined planner extensions
+    pub(crate) planners: Vec<Arc<dyn UserDefinedPlanner>>,


What do you think about using a single planner

Suggested change

/// user defined planner extensions

pub(crate) planners: Vec<Arc<dyn UserDefinedPlanner>>,

/// user defined planner extensions

pub planner: Option<Arc<dyn UserDefinedPlanner>>>,

This would match the basic pattern we have elsewhere (like proto extension codec) and people could implement composable planners like @lewiszlw did in https://github.com/apache/datafusion/blob/main/datafusion-examples/examples/composed_extension_codec.rs

I'm not sure about this

If we only accept single planner, does that mean we need to implement all the trait functions in one struct? We can't have array function planner optional in this case.

If the user expect their own planner + our builtin planner, how can they enable both? With current design, they just chain them with with_user_defined_planner.

query .with_user_defined_planner(array_planner) .with_user_defined_planner(field_access_planner) .with_user_defined_planner(user's planner) ...

The example given https://github.com/apache/datafusion/blob/main/datafusion-examples/examples/composed_extension_codec.rs seems to have multiple codecs too 🤔

/// A PhysicalExtensionCodec that tries one of multiple inner codecs /// until one works #[derive(Debug)] struct ComposedPhysicalExtensionCodec { codecs: Vec<Arc<dyn PhysicalExtensionCodec>>, }

alamb · 2024-06-30T11:30:08Z

datafusion/sqllogictest/test_files/array.slt

@@ -4910,10 +4910,11 @@ select array_ndims(arrow_cast([null], 'List(List(List(Int64)))'));
 3

 # array_ndims scalar function #2
+# TODO: dimensions 20 is the maximum without stack overflow, find a way to enable deep nested array


This is perhaps because the stackframe got larger in debug builds with this change -- one way to workaround it might be to pull the code for planning binary_exprs and the code for planning array access into their own functions rather than leave them inline in plan_expr

Signed-off-by: jayzhan211 <[email protected]>

alamb

Thank you @jayzhan211 -- I think this looks great. I pushed two commits to this branch to update comments and rename the variant to avoid having to go back and forth again. Hopefully that is ok

@samuelcolvin what do you think of this PR ? Is it enough to work with instead of #11137?

alamb · 2024-07-01T11:04:18Z

datafusion/core/src/execution/session_state.rs

+        // register crate of array expressions (if enabled)
+        #[cfg(feature = "array_expressions")]
+        {
+            let array_planner =


samuelcolvin · 2024-07-01T11:40:40Z

In general this looks great. I'm going to try to use it on datafusion-functions-json and let you know how I get on.

Should be able to get this done today, or tomorrow if life get's in the way.

jayzhan211 · 2024-07-01T12:39:24Z

datafusion/expr/src/planner.rs

+///
+/// This structure is used by [`UserDefinedSQLPlanner`] to plan operators with
+/// custom expressions.
+#[derive(Debug, Clone)]


I think we don't need Clone 🤔 Same as RawFieldAccessExpr and PlannerResult . Others looks good to me

I agree there isn't a critical usecase for it now, but I figured it didn't hurt. If you feel strongly I will remove Clone

I think it is not a critical issue.

jayzhan211 · 2024-07-02T00:00:28Z

Thanks @alamb

samuelcolvin · 2024-07-02T10:28:55Z

I'm a bit surprised this got merged without waiting for me to try it on datafusion-functions-json!

I said just above that I would try to do that yesterday or today.

As far as I can see, this PR provides no way to register external custom query planners. I guess I'll work on a follow up...

alamb · 2024-07-02T10:47:41Z

Thanks @samuelcolvin I agree iterating with some other PRs is a good idea.

alamb · 2024-07-02T10:53:56Z

I wrote up #11207 to track moving the remaining SQL planning features to this API

rtyler · 2024-07-07T16:53:02Z

datafusion/sql/src/expr/mod.rs

+                    }
+                }
+
+                not_impl_err!("GetFieldAccess not supported by UserDefinedExtensionPlanners: {field_access_expr:?}")


👋 I'm curious about this default branch because I'm seeing this error when doing integration testing with the delta-rs package. To the best of my knowledge we're not using any extension planners, but now this is failing with the latest datafusion. I'm uncertain whether we're doing something wrong, or if the default behavior of falling through to not_impl_err is causing us grief

Hi @rtyler -- looks to me like the delta.rs code is managing its own SessionState -- like in https://github.com/delta-io/delta-rs/blob/main/crates/core/src/delta_datafusion/mod.rs#L1687 (BTW the Delta API is really nicely thought out)

So I think you'll need to register the same planners in your SessionContext 🤔

Helpfully I think @Omega359 just made a PR to make this easier: #11296

I feel in general DataFusion is hard to use / configure correctly if you are using a custom SessionState / configuration -- one potential thing we were discussing is #11182 (comment) -- I'll file a ticket to make this a thing.

I think I had a similar issue you you @rtyler. Found a workaround described here: #11477

Proposed fix is here: #11485

lewiszlw · 2024-07-08T08:45:52Z

Sorry, I didn't follow up this topic before. I'm thinking if we can have one UserDefinedSQLPlanner trait which all methods have default official implementation, users just need to impl method they want to customize.
We have to loop UserDefinedSQLPlanner instances in current design, it seems not very efficient and has much boilerplate code. And these UserDefinedSQLPlanner instaces might have conflicts.

alamb · 2024-07-08T10:02:54Z

We have to loop UserDefinedSQLPlanner instances in current design, it seems not very efficient and has much boilerplate code. And these UserDefinedSQLPlanner instaces might have conflicts.

I agree the looping is not ideal.

In terms of conflicts, I think the semantics are pretty well defined now (the planners are ordered and whichever successfully plans the option first will be chosen).

* Prototype user defined sql planner might look like * at arrow Signed-off-by: jayzhan211 <[email protected]> * get field Signed-off-by: jayzhan211 <[email protected]> * cleanup Signed-off-by: jayzhan211 <[email protected]> * plan array literal Signed-off-by: jayzhan211 <[email protected]> * move to functions-array Signed-off-by: jayzhan211 <[email protected]> * license Signed-off-by: jayzhan211 <[email protected]> * cleanup Signed-off-by: jayzhan211 <[email protected]> * fix Signed-off-by: jayzhan211 <[email protected]> * change nested array test Signed-off-by: jayzhan211 <[email protected]> * address comment Signed-off-by: jayzhan211 <[email protected]> * fix stack overflow issue Signed-off-by: jayzhan211 <[email protected]> * upd cli Signed-off-by: jayzhan211 <[email protected]> * fix doc Signed-off-by: jayzhan211 <[email protected]> * fix doc Signed-off-by: jayzhan211 <[email protected]> * Rename PlannerResult::Simplified to PlannerResult::Planned * Update comments and add Debug/Clone impls --------- Signed-off-by: jayzhan211 <[email protected]> Co-authored-by: Andrew Lamb <[email protected]>

alamb and others added 6 commits June 28, 2024 13:29

Prototype user defined sql planner might look like

a7c9e7f

merge main

3c6801c

Signed-off-by: jayzhan211 <[email protected]>

at arrow

3362985

Signed-off-by: jayzhan211 <[email protected]>

get field

42f6414

Signed-off-by: jayzhan211 <[email protected]>

cleanup

8ece394

Signed-off-by: jayzhan211 <[email protected]>

plan array literal

d5d3189

Signed-off-by: jayzhan211 <[email protected]>

github-actions bot added the sql SQL Planner label Jun 30, 2024

move to functions-array

cdf5342

Signed-off-by: jayzhan211 <[email protected]>

github-actions bot added logical-expr Logical plan and expressions core Core DataFusion crate labels Jun 30, 2024

license

d386702

Signed-off-by: jayzhan211 <[email protected]>

jayzhan211 commented Jun 30, 2024

View reviewed changes

jayzhan211 added 2 commits June 30, 2024 14:00

cleanup

24372e4

Signed-off-by: jayzhan211 <[email protected]>

fix

e889009

Signed-off-by: jayzhan211 <[email protected]>

change nested array test

10477a8

Signed-off-by: jayzhan211 <[email protected]>

github-actions bot added the sqllogictest SQL Logic Tests (.slt) label Jun 30, 2024

jayzhan211 marked this pull request as ready for review June 30, 2024 09:21

alamb reviewed Jun 30, 2024

View reviewed changes

alamb changed the title ~~User defined sql planner~~ Introduce user defined SQL planner API Jun 30, 2024

alamb mentioned this pull request Jun 30, 2024

Discussion: make it easier for specify SQL --> function translation #10534

Closed

jayzhan211 added 2 commits June 30, 2024 20:44

address comment

b662b1e

Signed-off-by: jayzhan211 <[email protected]>

fix stack overflow issue

160e032

Signed-off-by: jayzhan211 <[email protected]>

github-actions bot removed the sqllogictest SQL Logic Tests (.slt) label Jun 30, 2024

jayzhan211 added 3 commits July 1, 2024 08:35

upd cli

c4e2b33

Signed-off-by: jayzhan211 <[email protected]>

fix doc

dafd53e

Signed-off-by: jayzhan211 <[email protected]>

fix doc

5e9af66

Signed-off-by: jayzhan211 <[email protected]>

alamb mentioned this pull request Jul 1, 2024

DataFusion weekly project plan (Andrew Lamb) - July 1, 2024 #11190

Closed

10 tasks

alamb added 2 commits July 1, 2024 07:06

Merge remote-tracking branch 'apache/main' into udplan

a0786eb

Rename PlannerResult::Simplified to PlannerResult::Planned

5963626

Update comments and add Debug/Clone impls

f1f1b4c

alamb approved these changes Jul 1, 2024

View reviewed changes

alamb mentioned this pull request Jul 1, 2024

Custom operator support #11137

Closed

6 tasks

jayzhan211 commented Jul 1, 2024

View reviewed changes

jayzhan211 merged commit 48a1754 into apache:main Jul 2, 2024
24 checks passed

jayzhan211 deleted the udplan branch July 2, 2024 00:00

samuelcolvin mentioned this pull request Jul 2, 2024

Allow user defined SQL planners to be registered #11208

Merged

jayzhan211 mentioned this pull request Jul 3, 2024

feat: support COUNT() #11229

Merged

rtyler reviewed Jul 7, 2024

View reviewed changes

alamb mentioned this pull request Jul 7, 2024

Extract registering default features from SessionState and into its own function #11320

Closed

This was referenced Jul 8, 2024

DataFusion weekly project plan (Andrew Lamb) - July 8, 2024 #11334

Closed

[Epic] Complete pulling out special SQL planning from the Sql Parser #11207

Closed

Revert "remove derive(Copy) from Operator (#11132)" #11341

Merged

tshauck mentioned this pull request Jul 9, 2024

feat: add raw aggregate udf planner #11371

Closed

This was referenced Jul 16, 2024

Get expr planners when creating new planner #11485

Merged

ExprPlanner not propagated to SqlToRel #11477

Closed

This pull request was closed.

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Introduce user defined SQL planner API #11180

Introduce user defined SQL planner API #11180

jayzhan211 commented Jun 30, 2024 •

edited by alamb

Loading

jayzhan211 Jun 30, 2024

jayzhan211 commented Jun 30, 2024

alamb left a comment

alamb Jun 30, 2024

alamb Jun 30, 2024

jayzhan211 Jun 30, 2024

alamb Jun 30, 2024

alamb Jun 30, 2024

alamb Jun 30, 2024

alamb Jun 30, 2024

alamb Jun 30, 2024

jayzhan211 Jun 30, 2024 •

edited

Loading

alamb Jun 30, 2024

alamb left a comment

alamb Jul 1, 2024

samuelcolvin commented Jul 1, 2024

jayzhan211 Jul 1, 2024

alamb Jul 1, 2024

jayzhan211 Jul 2, 2024

jayzhan211 commented Jul 2, 2024

samuelcolvin commented Jul 2, 2024

alamb commented Jul 2, 2024

alamb commented Jul 2, 2024

rtyler Jul 7, 2024

alamb Jul 7, 2024

wjones127 Jul 15, 2024

alamb Jul 16, 2024

lewiszlw commented Jul 8, 2024 •

edited

Loading

alamb commented Jul 8, 2024

-pub struct BinaryExpr {
+/// An operator with two arguments to plan
+///
+/// Note `left` and `right` are DataFusion [`Expr`]s but the `op` is the SQL AST operator.
+/// This structure is used by [`UserDefinedPlanner`] to plan operators with custom expressions.
+pub struct BinaryExpr {

		/// user defined planner extensions
		pub(crate) planners: Vec<Arc<dyn UserDefinedPlanner>>,

Introduce user defined SQL planner API #11180

Introduce user defined SQL planner API #11180

Conversation

jayzhan211 commented Jun 30, 2024 • edited by alamb Loading

Which issue does this PR close?

Rationale for this change

What changes are included in this PR?

Are these changes tested?

Are there any user-facing changes?

Choose a reason for hiding this comment

jayzhan211 commented Jun 30, 2024

alamb left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

jayzhan211 Jun 30, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

alamb left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

samuelcolvin commented Jul 1, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

jayzhan211 commented Jul 2, 2024

samuelcolvin commented Jul 2, 2024

alamb commented Jul 2, 2024

alamb commented Jul 2, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

lewiszlw commented Jul 8, 2024 • edited Loading

alamb commented Jul 8, 2024

jayzhan211 commented Jun 30, 2024 •

edited by alamb

Loading

jayzhan211 Jun 30, 2024 •

edited

Loading

lewiszlw commented Jul 8, 2024 •

edited

Loading