alamb commented on code in PR #4701:
URL: https://github.com/apache/arrow-datafusion/pull/4701#discussion_r1069123054
##########
datafusion/expr/src/expr_schema.rs:
##########
@@ -128,7 +128,14 @@ impl ExprSchemable for Expr {
Expr::Like { .. } | Expr::ILike { .. } | Expr::SimilarTo { .. } =>
{
Ok(DataType::Boolean)
}
- Expr::Placeholder { data_type, .. } => Ok(data_type.clone()),
+ Expr::Placeholder { data_type, .. } => {
+ let data_type = data_type.clone().ok_or_else(|| {
+ DataFusionError::Internal(
+ "Placeholder type could not be resolved".to_owned(),
+ )
+ })?;
+ Ok(data_type)
Review Comment:
```suggestion
data_type.clone().ok_or_else(|| {
DataFusionError::Internal(
"Placeholder type could not be resolved".to_owned(),
)
})
```
##########
datafusion/expr/src/logical_plan/plan.rs:
##########
@@ -603,6 +604,73 @@ impl LogicalPlan {
Ok(new_plan)
}
+ /// Walk the logical plan, find any `PlaceHolder` tokens, and return a map
of their IDs and DataTypes
+ pub fn get_parameter_types(
+ &self,
+ ) -> Result<HashMap<String, Option<DataType>>, DataFusionError> {
+ struct ParamTypeVisitor {
+ param_types: HashMap<String, Option<DataType>>,
+ }
+
+ struct ExprParamTypeVisitor {
+ param_types: HashMap<String, Option<DataType>>,
+ }
+
+ impl ExpressionVisitor for ExprParamTypeVisitor {
+ fn pre_visit(
+ mut self,
+ expr: &Expr,
+ ) -> datafusion_common::Result<Recursion<Self>>
+ where
+ Self: ExpressionVisitor,
+ {
+ if let Expr::Placeholder { id, data_type } = expr {
+ let _ = self.param_types.insert(id.clone(),
data_type.clone());
+ }
+ Ok(Recursion::Continue(self))
+ }
+ }
+
+ impl PlanVisitor for ParamTypeVisitor {
Review Comment:
I suspect you could `impl PlanVisitor` and `impl ExprVisitor` for
`ParamTypeVisitor` rather than having multiple visitors and having to copy
items from the hash maps
##########
datafusion/expr/src/logical_plan/plan.rs:
##########
@@ -603,6 +604,73 @@ impl LogicalPlan {
Ok(new_plan)
}
+ /// Walk the logical plan, find any `PlaceHolder` tokens, and return a map
of their IDs and DataTypes
+ pub fn get_parameter_types(
+ &self,
+ ) -> Result<HashMap<String, Option<DataType>>, DataFusionError> {
+ struct ParamTypeVisitor {
+ param_types: HashMap<String, Option<DataType>>,
+ }
+
+ struct ExprParamTypeVisitor {
+ param_types: HashMap<String, Option<DataType>>,
+ }
+
+ impl ExpressionVisitor for ExprParamTypeVisitor {
+ fn pre_visit(
+ mut self,
+ expr: &Expr,
+ ) -> datafusion_common::Result<Recursion<Self>>
+ where
+ Self: ExpressionVisitor,
+ {
+ if let Expr::Placeholder { id, data_type } = expr {
+ let _ = self.param_types.insert(id.clone(),
data_type.clone());
+ }
+ Ok(Recursion::Continue(self))
+ }
+ }
+
+ impl PlanVisitor for ParamTypeVisitor {
+ type Error = DataFusionError;
+
+ fn pre_visit(&mut self, plan: &LogicalPlan) -> Result<bool,
Self::Error> {
+ let mut visitor = ExprParamTypeVisitor {
+ param_types: Default::default(),
+ };
+ match plan {
Review Comment:
There are probably more plan types (like Aggregates for example) that could
have references to parameter values
I poked around for a "visit all `Expr` in a plan" preexisting code and could
not find one. So longer term I think it would be good to move this code
somewhere more general (maybe near `ExpressionVisitor`):
```rust
/// visit all expressions in all nodes of the LogicalPlan
fn visit_plan_exprs(plan: &LogicalPlan, visitor: ExpressionVisitor) {
...
}
However think it is ok to start with this approach of something very
specialized for parameter type inference and then refactor into something more
general as a follow on PR. I also think we can track adding support for
parameters in other plan types with a ticket as well if you don't want to
support the entire thing in the first PR
##########
datafusion/expr/src/logical_plan/plan.rs:
##########
@@ -603,6 +604,73 @@ impl LogicalPlan {
Ok(new_plan)
}
+ /// Walk the logical plan, find any `PlaceHolder` tokens, and return a map
of their IDs and DataTypes
+ pub fn get_parameter_types(
+ &self,
+ ) -> Result<HashMap<String, Option<DataType>>, DataFusionError> {
+ struct ParamTypeVisitor {
+ param_types: HashMap<String, Option<DataType>>,
+ }
+
+ struct ExprParamTypeVisitor {
+ param_types: HashMap<String, Option<DataType>>,
+ }
+
+ impl ExpressionVisitor for ExprParamTypeVisitor {
+ fn pre_visit(
+ mut self,
+ expr: &Expr,
+ ) -> datafusion_common::Result<Recursion<Self>>
+ where
+ Self: ExpressionVisitor,
+ {
+ if let Expr::Placeholder { id, data_type } = expr {
+ let _ = self.param_types.insert(id.clone(),
data_type.clone());
Review Comment:
Should we error here if the parameter was previously referenced with a
different datatype 🤔 (aka don't ignore `_`?)
##########
datafusion/expr/src/expr_schema.rs:
##########
@@ -128,7 +128,14 @@ impl ExprSchemable for Expr {
Expr::Like { .. } | Expr::ILike { .. } | Expr::SimilarTo { .. } =>
{
Ok(DataType::Boolean)
}
- Expr::Placeholder { data_type, .. } => Ok(data_type.clone()),
+ Expr::Placeholder { data_type, .. } => {
+ let data_type = data_type.clone().ok_or_else(|| {
+ DataFusionError::Internal(
+ "Placeholder type could not be resolved".to_owned(),
+ )
+ })?;
+ Ok(data_type)
Review Comment:
very minor style suggestion
##########
datafusion/expr/src/expr_schema.rs:
##########
@@ -128,7 +128,14 @@ impl ExprSchemable for Expr {
Expr::Like { .. } | Expr::ILike { .. } | Expr::SimilarTo { .. } =>
{
Ok(DataType::Boolean)
}
- Expr::Placeholder { data_type, .. } => Ok(data_type.clone()),
+ Expr::Placeholder { data_type, .. } => {
+ let data_type = data_type.clone().ok_or_else(|| {
+ DataFusionError::Internal(
+ "Placeholder type could not be resolved".to_owned(),
+ )
+ })?;
+ Ok(data_type)
Review Comment:
I also think this error might be better "Plan" or "Unsupported" -- it isn't
really an internal (unexpected) error type
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
