feat: Add ConfigOptions to ScalarFunctionArgs #13527

Omega359 · 2024-11-22T15:28:52Z

Which issue does this PR close?

Closes #13519

Rationale for this change

Allow udf's to access df config

What changes are included in this PR?

Code.

Are these changes tested?

Existing tests.

Are there any user-facing changes?

Not specifically, this is covered with the udf signature change in #13290

Omega359 · 2024-11-22T15:46:39Z

There is a lot of file changes here but most of the important changes are in scalar_function.rs, There is a todo in expr_simplifier.rs that I would like feedback on.

alamb · 2024-11-24T22:56:58Z

I plan to review this carefully tomorrow

alamb

Thank you @Omega359 -- this is an epic plumbing exercise 🪠

The signature in ScalarFunctionArgs is 👌 very nice

This PR seems to require config_options to be cloned many times now. I wonder if it is possible to avoid that 🤔. I took a brief look and it seems to be somewhat challenging as SessionState allows mutable access to the underlying SessionConfig.

Maybe we could change the semantics so that SessionConfig has a Arc<ConfigOptions> which was cloned when it was modified (Arc::unwrap_or_clone() style) 🤔

I also think the const evaluator does need the actual correct ConfigOptions for correctness

alamb · 2024-11-25T11:31:21Z

datafusion-examples/examples/expr_api.rs

-    let physical_expr =
-        datafusion_physical_expr::create_physical_expr(&expr, &df_schema, &props)?;
+    let config_options = Arc::new(ConfigOptions::default());
+    let physical_expr = datafusion_physical_expr::create_physical_expr(


It seems somewhat inevitable that creating a physical expr will require the config options

However, I also think threading through the config options down through to the physical creation will (finally) permit people to pass things from the session down to function implementations (I think @cisaacson also was trying to do this in the past)

alamb · 2024-11-25T11:32:21Z

datafusion/core/src/datasource/listing/helpers.rs

@@ -283,10 +284,16 @@ async fn prune_partitions(

    // TODO: Plumb this down


This todo may have now be complete

alamb · 2024-11-25T11:36:27Z

datafusion/expr/src/udf.rs

@@ -336,6 +337,8 @@ pub struct ScalarFunctionArgs<'a> {
    // The return type of the scalar function returned (from `return_type` or `return_type_from_exprs`)
    // when creating the physical expression from the logical expression
    pub return_type: &'a DataType,
+    // The config options which can be used to lookup configuration properties
+    pub config_options: Arc<ConfigOptions>,


alamb · 2024-11-25T11:38:45Z

datafusion/optimizer/src/simplify_expressions/expr_simplifier.rs

-                Ok(e) => e,
-                Err(err) => return ConstSimplifyResult::SimplifyRuntimeError(err, expr),
-            };
+        // todo - should the config options be the actual options here or is this sufficient?


I think the actual configuration options are needed here. Otherwise what will happen is that any function whose behavior relies on the ConfigOptions may have different behavior on columns and constants (or other expressions that can be constant folded)

Omega359 · 2024-11-25T14:38:55Z

This PR seems to require config_options to be cloned many times now. I wonder if it is possible to avoid that 🤔. I took a brief look and it seems to be somewhat challenging as SessionState allows mutable access to the underlying SessionConfig.

Yes, it's a bit annoying. I was tempted to see if I could switch to &'a ConfigOptions everywhere. There is at least one 'real' (vs Arc::clone) clone for every query, possibly more as I haven't checked.

Maybe we could change the semantics so that SessionConfig has a Arc<ConfigOptions> which was cloned when it was modified (Arc::unwrap_or_clone() style) 🤔

Certainly possible, I can attempt that.

I also think the const evaluator does need the actual correct ConfigOptions for correctness

I was afraid of that. I was avoiding it because of the signature changes it would required just about everywhere which would cause even more headaches for those systems trying to upgrade.

alamb · 2024-11-25T21:05:57Z

Yeah, it is a tricky one for sure

alamb · 2024-11-27T19:24:57Z

Marking as draft as I think this PR is no longer waiting on feedback. Please mark it as ready for review when it is ready for another look

…rgs_session_config

Omega359 · 2024-11-28T21:33:32Z

Maybe we could change the semantics so that SessionConfig has a Arc<ConfigOptions> which was cloned when it was modified (Arc::unwrap_or_clone() style) 🤔

Certainly possible, I can attempt that.

@alamb I did a quick attempt at implementing that however it breaks a commonly used method - SessionConfig.options_mut(). Not having that available breaks a bunch of stuff and while switching to SessionConfig.set(..) is quite possible it's not as clean.

Trying with &ConfigOptions in ScalarFunctionExpr leads to lifetime hell in areas I have no idea how to overcome right now.

As much as I want this feature I'm going to put it aside for now

…config # Conflicts: # datafusion/optimizer/src/simplify_expressions/expr_simplifier.rs # datafusion/proto/src/physical_plan/from_proto.rs

…Expr. This means no cloning of ConfigOptions except for ScalarFunctions.

Omega359 · 2024-12-23T20:29:31Z

I think this may be ready for review again. For this round I refactored the code to use &ConfigOptions everywhere except for ScalarFunctionExpr so the cost for cloning ConfigOptions is only incurred when creating a scalar UDF.

…_args_session_config # Conflicts: # datafusion/physical-expr/src/equivalence/properties.rs

Omega359 · 2024-12-28T23:59:14Z

Examples check failure is transient I believe.

alamb · 2024-12-29T13:17:56Z

Examples check failure is transient I believe.

I restarted the checks

alamb

Thank you @Omega359 -- I think this PR is a major step forward and we could merge it as is.

However, I feel strongly there are two things that should be improved soon (if not this PR then follow on ones). I left specific comments about each

Don't copy ConfigOptions in every call to ScalarFunction::create_physical_expr
Add convenience methods to get default &ConfigOptions and Arc<ConfigOptions> which I think will help people upgrading to the next version of DataFusion quickly migrate their code

alamb · 2024-12-29T13:25:31Z

datafusion-examples/examples/composed_extension_codec.rs

    let result_exec_plan: Arc<dyn ExecutionPlan> = proto
-        .try_into_physical_plan(&ctx, runtime.deref(), &composed_codec)
+        .try_into_physical_plan(&ctx, config_options, runtime.deref(), &composed_codec)


This is an API change, but I think it is required to thread the config options through as each argument is specific

We could potentially improve the try_into_physical_plan API (as a follow on PR) to make it easier to update the API in the future using a trait or something like

https://github.com/apache/datafusion/blob/e99e02b9b9093ceb0c13a2dd32a2a89beba47930/datafusion/expr/src/expr_schema.rs#L39-L38

So this would look something like

pub trait ProtobufContext { /// return a function registry fn function_registry(&self) -> &dyn FunctionRegistry; /// return the runtime env fn runtime_env(&self) -> &RuntimeEnv; /// return the config options fn config_options(&self) -> &ConfigOptions; /// return extension codec fn extension_codec(&self) -> &dyn PhysicalExtensionCodec; }

impl AsExecutionPlan for protobuf::PhysicalPlanNode { ... fn try_into_physical_plan( &self, registry: &dyn FunctionRegistry, config_options: &ConfigOptions, runtime: &RuntimeEnv, extension_codec: &dyn PhysicalExtensionCodec, ) -> Result<Arc<dyn ExecutionPlan>> {

alamb · 2024-12-29T13:29:21Z

datafusion-examples/examples/pruning.rs

@@ -187,7 +188,9 @@ impl PruningStatistics for MyCatalog {
 fn create_pruning_predicate(expr: Expr, schema: &SchemaRef) -> PruningPredicate {
    let df_schema = DFSchema::try_from(schema.as_ref().clone()).unwrap();
    let props = ExecutionProps::new();
-    let physical_expr = create_physical_expr(&expr, &df_schema, &props).unwrap();
+    let config_options = ConfigOptions::default();


Something we could potentially do to make this API slightly easier to use might be to create a static default ConfigOptions

Something like

impl ConfigOptions { /// returns a reference to default ConfigOptions pub fn default_singleton() -> &'static ConfigOptions }

This would then make it easier to return &ConfigOptions in various places when only the default was needed

For example, then in LocalCsvTableFunc you could avoid having to thread the ConfigOptions through as in that example having the actual config options isn't important

alamb · 2024-12-29T13:34:22Z

datafusion/core/src/datasource/listing/helpers.rs

@@ -241,6 +241,7 @@ async fn prune_partitions(
    partitions: Vec<Partition>,
    filters: &[Expr],
    partition_cols: &[(String, DataType)],
+    ctx: &SessionState,


I recommend passing &ConfigOptions directly here as the function only needs the config, not the entire session state

alamb · 2024-12-29T13:42:33Z

datafusion/physical-expr/src/scalar_function.rs

@@ -243,6 +292,7 @@ pub fn create_physical_expr(
            Arc::new(fun.clone()),
            input_phy_exprs.to_vec(),
            return_type,
+            Arc::new(config_options.clone()),


This line worries me -- it means each distinct scalar function in the plan will get an entirely new copy of the ConfigOptions.

I think the Arc should be passed in as the argument like this (I realize this will be a significant code change) so that the config options are copied at most once per plan

/// Create a physical expression for the UDF. pub fn create_physical_expr( fun: &ScalarUDF, input_phy_exprs: &[Arc<dyn PhysicalExpr>], input_schema: &Schema, args: &[Expr], input_dfschema: &DFSchema, config_options: &Arc<ConfigOptions>, // <--- I think this should be an `Arc` ) -> Result<Arc<dyn PhysicalExpr>> {

I think that would also make it clearer that physical planning makes a read only copy of the configuration (Arc<ConfigOptions>) that is then unchanged during execution

That was my main concern as well. I am unsure that this change should merged in as is to be honest as any fix will be just as disruptive as this PR is api wise.

Pushing the clone higher up the stack is possible but I did run into major issues trying to push it all the way up including all kinds of disruptive changes like changing OptimizerConfig from &ConfigOptions to the Arc version. Doing the opposite - pushing &ConfigOptions all the way down ran into issues with ScalarFunctionExpr and DynEq/DynHash.

Speaking for myself I prefer the latter as I feel it's less disruptive overall ... I just couldn't get it to work before. Changing the signature of DynHash and DynEq traits may allow it to work but I haven't tried that yet. Maybe with some more thought in January I can get it to work but I expect my January time to be pretty limited.

Pushing the clone higher up the stack is possible but I did run into major issues trying to push it all the way up including all kinds of disruptive changes like changing OptimizerConfig from &ConfigOptions to the Arc version. Doing the opposite - pushing &ConfigOptions all the way down ran into issues with ScalarFunctionExpr and DynEq/DynHash.

If you have an Arc<ConfigOptions> you can always get a &ConfigOptions by calling options.as_ref()

So in other words, I don't know if you need to push the Arc through everywhere (at least at first) -- we could just change the create physical expr. I would love to help with this -- perhaps I can find time later in the week (though I have a few other things going on too)

Another option to avoid cloning ConfigOptions for each instance of ScalarFunction might be to add some sort of API / way for the function to communicate "I need the config options"

Omega359 added 3 commits November 20, 2024 18:00

Fix md5 return_type to only return Utf8 as per current code impl.

a18c82f

Merge remote-tracking branch 'origin/main'

ad50ae6

Add SessionConfig reference to ScalarFunctionArgs

42cd8dc

github-actions bot added logical-expr Logical plan and expressions physical-expr Physical Expressions optimizer Optimizer rules core Core DataFusion crate common Related to common crate proto Related to proto crate functions labels Nov 22, 2024

Omega359 changed the title ~~Feature/scalar func args session config~~ feat: Add ConfigOptions to ScalarFuntionArgs Nov 22, 2024

Omega359 changed the title ~~feat: Add ConfigOptions to ScalarFuntionArgs~~ feat: Add ConfigOptions to ScalarFunctionArgs Nov 22, 2024

Omega359 marked this pull request as ready for review November 22, 2024 16:42

alamb reviewed Nov 25, 2024

View reviewed changes

alamb marked this pull request as draft November 27, 2024 19:24

Merge remote-tracking branch 'origin/main' into feature/scalar_func_a…

e84ca6d

…rgs_session_config

Omega359 added 7 commits November 28, 2024 21:37

Merge remote-tracking branch 'origin/main'

1e87f7e

Merge remote-tracking branch 'origin/main'

cc1f540

Merge remote-tracking branch 'origin/main'

a733c8f

Merge remote-tracking branch 'origin/main'

87b388e

Merge branch 'refs/heads/main' into feature/scalar_func_args_session_…

6a6d86c

…config # Conflicts: # datafusion/optimizer/src/simplify_expressions/expr_simplifier.rs # datafusion/proto/src/physical_plan/from_proto.rs

Refactored to use &ConfigOptions everywhere except for ScalarFunction…

b8bf209

…Expr. This means no cloning of ConfigOptions except for ScalarFunctions.

Cleanup updates.

dc9f5db

Omega359 marked this pull request as ready for review December 23, 2024 20:29

Omega359 added 2 commits December 28, 2024 23:21

Merge remote-tracking branch 'upstream/main' into feature/scalar_func…

1041525

…_args_session_config # Conflicts: # datafusion/physical-expr/src/equivalence/properties.rs

Fixed merge conflict

5557f64

alamb added the api change Changes the API exposed to users of the crate label Dec 29, 2024

alamb approved these changes Dec 29, 2024

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

feat: Add ConfigOptions to ScalarFunctionArgs #13527

feat: Add ConfigOptions to ScalarFunctionArgs #13527

Omega359 commented Nov 22, 2024

Omega359 commented Nov 22, 2024

alamb commented Nov 24, 2024

alamb left a comment

alamb Nov 25, 2024

alamb Nov 25, 2024

alamb Nov 25, 2024

alamb Nov 25, 2024

Omega359 commented Nov 25, 2024

alamb commented Nov 25, 2024

alamb commented Nov 27, 2024

Omega359 commented Nov 28, 2024

Omega359 commented Dec 23, 2024

Omega359 commented Dec 28, 2024

alamb commented Dec 29, 2024

alamb left a comment

alamb Dec 29, 2024

alamb Dec 29, 2024

alamb Dec 29, 2024

alamb Dec 29, 2024

Omega359 Dec 29, 2024

alamb Dec 30, 2024

		@@ -283,10 +284,16 @@ async fn prune_partitions(

		// TODO: Plumb this down

feat: Add ConfigOptions to ScalarFunctionArgs #13527

Are you sure you want to change the base?

feat: Add ConfigOptions to ScalarFunctionArgs #13527

Conversation

Omega359 commented Nov 22, 2024

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?

Omega359 commented Nov 22, 2024

alamb commented Nov 24, 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

Omega359 commented Nov 25, 2024

alamb commented Nov 25, 2024

alamb commented Nov 27, 2024

Omega359 commented Nov 28, 2024

Omega359 commented Dec 23, 2024

Omega359 commented Dec 28, 2024

alamb commented Dec 29, 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