physical expr: support placeholders

This patchs adds a support for placeholders on the physical expression
level. It allows to use generic plans and then resolve placeholders on the
execution stage.

Placeholders are resolved on the `execution(...)` phase. Each `ExecutionPlan`
is responsible for resolving placeholders in it's own expressions.

Example:
```sh
> create table a(x int);

> explain select x + $1 from a where x > $2;
+---------------+-------------------------------------------------------------------------+
| plan_type     | plan                                                                    |
+---------------+-------------------------------------------------------------------------+
| logical_plan  | Projection: a.x + $1                                                    |
|               |   Filter: a.x > $2                                                      |
|               |     TableScan: a projection=[x]                                         |
| physical_plan | ProjectionExec: expr=[x@0 + $1 as a.x + $1]                             |
|               |   RepartitionExec: partitioning=RoundRobinBatch(16), input_partitions=1 |
|               |     CoalesceBatchesExec: target_batch_size=8192                         |
|               |       FilterExec: x@0 > $2                                              |
|               |         MemoryExec: partitions=1, partition_sizes=[0]                   |
|               |                                                                         |
+---------------+-------------------------------------------------------------------------+
```

Author: askalt

datafusion-examples/examples/advanced_udwf.rs

@@ -80,6 +80,13 @@ impl WindowUDFImpl for SmoothItUdf {
     fn partition_evaluator(&self) -> Result<Box<dyn PartitionEvaluator>> {
         Ok(Box::new(MyPartitionEvaluator::new()))
     }
+
+    fn resolve_placeholders(
+        &self,
+        _param_values: &Option<datafusion_common::ParamValues>,
+    ) -> Result<Option<std::sync::Arc<dyn WindowUDFImpl>>> {
+        Ok(None)
+    }
 }
 
 /// This implements the lowest level evaluation for a window function

datafusion-examples/examples/simplify_udwf_expression.rs

@@ -84,6 +84,13 @@ impl WindowUDFImpl for SimplifySmoothItUdf {
 
         Some(Box::new(simplify))
     }
+
+    fn resolve_placeholders(
+        &self,
+        _param_values: &Option<datafusion_common::ParamValues>,
+    ) -> Result<Option<std::sync::Arc<dyn WindowUDFImpl>>> {
+        Ok(None)
+    }
 }
 
 // create local execution context with `cars.csv` registered as a table named `cars`

datafusion/core/tests/user_defined/user_defined_window_functions.rs

@@ -565,6 +565,13 @@ impl OddCounter {
             fn aliases(&self) -> &[String] {
                 &self.aliases
             }
+
+            fn resolve_placeholders(
+                &self,
+                _param_values: &Option<datafusion_common::ParamValues>,
+            ) -> Result<Option<Arc<dyn WindowUDFImpl>>> {
+                Ok(None)
+            }
         }
 
         ctx.register_udwf(WindowUDF::from(SimpleWindowUDF::new(test_state)))

datafusion/execution/src/task.rs

@@ -26,7 +26,7 @@ use crate::{
     registry::FunctionRegistry,
     runtime_env::{RuntimeEnv, RuntimeEnvBuilder},
 };
-use datafusion_common::{plan_datafusion_err, DataFusionError, Result};
+use datafusion_common::{plan_datafusion_err, DataFusionError, ParamValues, Result};
 use datafusion_expr::planner::ExprPlanner;
 use datafusion_expr::{AggregateUDF, ScalarUDF, WindowUDF};
 
@@ -53,6 +53,8 @@ pub struct TaskContext {
     window_functions: HashMap<String, Arc<WindowUDF>>,
     /// Runtime environment associated with this task context
     runtime: Arc<RuntimeEnv>,
+    /// Param values for physical placeholders.
+    param_values: Option<ParamValues>,
 }
 
 impl Default for TaskContext {
@@ -70,6 +72,7 @@ impl Default for TaskContext {
             aggregate_functions: HashMap::new(),
             window_functions: HashMap::new(),
             runtime,
+            param_values: None,
         }
     }
 }
@@ -97,6 +100,7 @@ impl TaskContext {
             aggregate_functions,
             window_functions,
             runtime,
+            param_values: None,
         }
     }
 
@@ -125,6 +129,17 @@ impl TaskContext {
         Arc::clone(&self.runtime)
     }
 
+    /// Return param values associated with thix [`TaskContext`].
+    pub fn param_values(&self) -> &Option<ParamValues> {
+        &self.param_values
+    }
+
+    /// Update the param values.
+    pub fn with_param_values(mut self, param_values: ParamValues) -> Self {
+        self.param_values = Some(param_values);
+        self
+    }
+
     /// Update the [`SessionConfig`]
     pub fn with_session_config(mut self, session_config: SessionConfig) -> Self {
         self.session_config = session_config;

datafusion/expr/src/expr_fn.rs

@@ -665,6 +665,13 @@ impl WindowUDFImpl for SimpleWindowUDF {
     fn partition_evaluator(&self) -> Result<Box<dyn crate::PartitionEvaluator>> {
         (self.partition_evaluator_factory)()
     }
+
+    fn resolve_placeholders(
+        &self,
+        _param_values: &Option<datafusion_common::ParamValues>,
+    ) -> Result<Option<Arc<dyn WindowUDFImpl>>> {
+        Ok(None)
+    }
 }
 
 pub fn interval_year_month_lit(value: &str) -> Expr {

datafusion/expr/src/udwf.rs

@@ -27,7 +27,7 @@ use std::{
     sync::Arc,
 };
 
-use datafusion_common::{not_impl_err, Result};
+use datafusion_common::{not_impl_err, ParamValues, Result};
 
 use crate::expr::WindowFunction;
 use crate::{
@@ -107,9 +107,12 @@ impl WindowUDF {
     where
         F: WindowUDFImpl + 'static,
     {
-        Self {
-            inner: Arc::new(fun),
-        }
+        Self::new_from_arc_impl(Arc::new(fun))
+    }
+
+    /// Createa a new `WindowUDF` from a dyn `[WindowUDFImpl`].
+    pub fn new_from_arc_impl(fun: Arc<dyn WindowUDFImpl>) -> Self {
+        Self { inner: fun }
     }
 
     /// Return the underlying [`WindowUDFImpl`] trait object for this function
@@ -384,6 +387,15 @@ pub trait WindowUDFImpl: Debug + Send + Sync {
     fn coerce_types(&self, _arg_types: &[DataType]) -> Result<Vec<DataType>> {
         not_impl_err!("Function {} does not implement coerce_types", self.name())
     }
+
+    /// Resolve placeholders in this expession.
+    /// Returns [`Some`] with rewrited expression if there is
+    /// at least one placeholder.
+    /// Otherwise returns [`None`].
+    fn resolve_placeholders(
+        &self,
+        param_values: &Option<ParamValues>,
+    ) -> Result<Option<Arc<dyn WindowUDFImpl>>>;
 }
 
 impl PartialEq for dyn WindowUDFImpl {
@@ -476,6 +488,21 @@ impl WindowUDFImpl for AliasedWindowUDFImpl {
     fn coerce_types(&self, arg_types: &[DataType]) -> Result<Vec<DataType>> {
         self.inner.coerce_types(arg_types)
     }
+
+    fn resolve_placeholders(
+        &self,
+        param_values: &Option<ParamValues>,
+    ) -> Result<Option<Arc<dyn WindowUDFImpl>>> {
+        let inner = self.inner.resolve_placeholders(param_values)?;
+        Ok(if let Some(inner) = inner {
+            Some(Arc::new(Self {
+                inner: inner,
+                aliases: self.aliases.clone(),
+            }))
+        } else {
+            None
+        })
+    }
 }
 
 /// Implementation of [`WindowUDFImpl`] that wraps the function style pointers
@@ -525,6 +552,13 @@ impl WindowUDFImpl for WindowUDFLegacyWrapper {
     fn partition_evaluator(&self) -> Result<Box<dyn PartitionEvaluator>> {
         (self.partition_evaluator_factory)()
     }
+
+    fn resolve_placeholders(
+        &self,
+        _param_values: &Option<ParamValues>,
+    ) -> Result<Option<Arc<dyn WindowUDFImpl>>> {
+        Ok(None)
+    }
 }
 
 #[cfg(test)]
@@ -570,6 +604,12 @@ mod test {
         fn partition_evaluator(&self) -> Result<Box<dyn PartitionEvaluator>> {
             unimplemented!()
         }
+        fn resolve_placeholders(
+            &self,
+            _param_values: &Option<datafusion_common::ParamValues>,
+        ) -> Result<Option<std::sync::Arc<dyn WindowUDFImpl>>> {
+            Ok(None)
+        }
     }
 
     #[derive(Debug, Clone)]
@@ -606,6 +646,12 @@ mod test {
         fn partition_evaluator(&self) -> Result<Box<dyn PartitionEvaluator>> {
             unimplemented!()
         }
+        fn resolve_placeholders(
+            &self,
+            _param_values: &Option<datafusion_common::ParamValues>,
+        ) -> Result<Option<std::sync::Arc<dyn WindowUDFImpl>>> {
+            Ok(None)
+        }
     }
 
     #[test]

datafusion/functions-window/src/row_number.rs

@@ -102,6 +102,13 @@ impl WindowUDFImpl for RowNumber {
             nulls_first: false,
         })
     }
+
+    fn resolve_placeholders(
+        &self,
+        _param_values: &Option<datafusion_common::ParamValues>,
+    ) -> Result<Option<std::sync::Arc<dyn WindowUDFImpl>>> {
+        Ok(None)
+    }
 }
 
 /// State for the `row_number` built-in window function.

datafusion/optimizer/src/simplify_expressions/expr_simplifier.rs

@@ -3916,5 +3916,12 @@ mod tests {
         fn partition_evaluator(&self) -> Result<Box<dyn PartitionEvaluator>> {
             unimplemented!("not needed for tests")
         }
+
+        fn resolve_placeholders(
+            &self,
+            _param_values: &Option<datafusion_common::ParamValues>,
+        ) -> Result<Option<Arc<dyn WindowUDFImpl>>> {
+            Ok(None)
+        }
     }
 }

datafusion/physical-expr/src/aggregate.rs

@@ -35,8 +35,8 @@ pub mod utils {
 }
 
 use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
-use datafusion_common::ScalarValue;
 use datafusion_common::{internal_err, not_impl_err, Result};
+use datafusion_common::{ParamValues, ScalarValue};
 use datafusion_expr::AggregateUDF;
 use datafusion_expr::ReversedUDAF;
 use datafusion_expr_common::accumulator::Accumulator;
@@ -52,6 +52,8 @@ use datafusion_expr_common::groups_accumulator::GroupsAccumulator;
 use std::fmt::Debug;
 use std::sync::Arc;
 
+use crate::expressions::resolve_placeholders_seq;
+
 /// Builder for physical [`AggregateFunctionExpr`]
 ///
 /// `AggregateFunctionExpr` contains the information necessary to call
@@ -528,6 +530,38 @@ impl AggregateFunctionExpr {
     pub fn default_value(&self, data_type: &DataType) -> Result<ScalarValue> {
         self.fun.default_value(data_type)
     }
+
+    /// Resolves placeholders.
+    /// If there are no placeholders returns [`None`].
+    /// Otherwise returns [`Some`] that contains resolved expression.
+    pub fn resolve_placeholders(
+        self: &Arc<Self>,
+        param_values: &Option<ParamValues>,
+    ) -> Result<Option<Arc<Self>>> {
+        Ok(
+            if let Some(resolved_args) =
+                resolve_placeholders_seq(&self.args, param_values)?
+            {
+                Some(Arc::new(AggregateFunctionExpr {
+                    fun: self.fun.clone(),
+                    args: resolved_args,
+                    data_type: self.data_type.clone(),
+                    name: self.name.clone(),
+                    schema: self.schema.clone(),
+                    ordering_req: self.ordering_req.clone(),
+                    ignore_nulls: self.ignore_nulls,
+                    ordering_fields: self.ordering_fields.clone(),
+                    is_distinct: self.is_distinct,
+                    is_reversed: self.is_reversed,
+                    input_types: self.input_types.clone(),
+                    is_nullable: self.is_nullable,
+                }))
+            } else {
+                // Args do not contain placeholders at all.
+                None
+            },
+        )
+    }
 }
 
 /// Stores the physical expressions used inside the `AggregateExpr`.

datafusion/physical-expr/src/expressions/mod.rs

@@ -30,6 +30,7 @@ mod literal;
 mod negative;
 mod no_op;
 mod not;
+mod placeholder;
 mod try_cast;
 mod unknown_column;
 
@@ -55,5 +56,6 @@ pub use literal::{lit, Literal};
 pub use negative::{negative, NegativeExpr};
 pub use no_op::NoOp;
 pub use not::{not, NotExpr};
+pub use placeholder::{placeholder, resolve_placeholders, resolve_placeholders_seq};
 pub use try_cast::{try_cast, TryCastExpr};
 pub use unknown_column::UnKnownColumn;