alamb commented on a change in pull request #8746:
URL: https://github.com/apache/arrow/pull/8746#discussion_r529052558
##########
File path: rust/datafusion/src/optimizer/utils.rs
##########
@@ -24,7 +24,11 @@ use arrow::datatypes::{Schema, SchemaRef};
use super::optimizer::OptimizerRule;
use crate::error::{DataFusionError, Result};
use crate::logical_plan::{Expr, LogicalPlan, PlanType, StringifiedPlan};
-use crate::prelude::col;
+use crate::prelude::{col, lit};
+use crate::scalar::ScalarValue;
+
+const CASE_EXPR_MARKER: &str = "__DATAFUSION_CASE_EXPR__";
Review comment:
🤔
##########
File path: rust/datafusion/src/logical_plan/expr.rs
##########
@@ -141,6 +166,7 @@ impl Expr {
Expr::Column(name) =>
Ok(schema.field_with_name(name)?.data_type().clone()),
Expr::ScalarVariable(_) => Ok(DataType::Utf8),
Expr::Literal(l) => Ok(l.get_datatype()),
+ Expr::Case { when_then_expr, .. } =>
when_then_expr[0].1.get_type(schema),
Review comment:
I feel like there should be a check somewhere that all the `then` and
`else` clauses result in the same type -- however, I suppose we would just get
an arrow error at runtime, which while not ideal, could easily be fixed
subsequently.
##########
File path: rust/datafusion/src/logical_plan/expr.rs
##########
@@ -342,6 +386,82 @@ impl Expr {
}
}
+pub struct CaseWhenBuilder {
+ expr: Option<Box<Expr>>,
+ when_expr: Vec<Expr>,
+ then_expr: Vec<Expr>,
+}
+
+pub struct CaseThenBuilder {
+ expr: Option<Box<Expr>>,
+ when_expr: Vec<Expr>,
+ then_expr: Vec<Expr>,
+}
+
+impl CaseWhenBuilder {
+ pub fn when(&mut self, expr: Expr) -> CaseThenBuilder {
+ self.when_expr.push(expr);
+ CaseThenBuilder {
+ expr: self.expr.clone(),
+ when_expr: self.when_expr.clone(),
+ then_expr: self.then_expr.clone(),
+ }
+ }
+ pub fn or_else(&mut self, else_expr: Expr) -> Expr {
+ Expr::Case {
+ expr: self.expr.clone(),
+ when_then_expr: self
+ .when_expr
+ .iter()
+ .zip(self.then_expr.iter())
+ .map(|(w, t)| (Box::new(w.clone()), Box::new(t.clone())))
+ .collect(),
+ else_expr: Some(Box::new(else_expr)),
+ }
+ }
+ pub fn end(&mut self) -> Expr {
+ Expr::Case {
+ expr: self.expr.clone(),
+ when_then_expr: self
+ .when_expr
+ .iter()
+ .zip(self.then_expr.iter())
+ .map(|(w, t)| (Box::new(w.clone()), Box::new(t.clone())))
+ .collect(),
+ else_expr: None,
+ }
+ }
+}
Review comment:
is there a reason for a second `impl` for `CaseThenBuilder`?
Also given that `when / then` always need to come in pairs, maybe
```
pub fn when(&mut self, expr: Expr) -> CaseThenBuilder {
```
Could be translated into
```
pub fn when_then(&mut self, when_expr: Expr, then_expr) ->
CaseThenBuilder {
```
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
[email protected]
