rdettai commented on a change in pull request #1128: URL: https://github.com/apache/arrow-datafusion/pull/1128#discussion_r730659006
########## File path: datafusion/src/optimizer/utils.rs ########## @@ -468,6 +473,117 @@ pub fn rewrite_expression(expr: &Expr, expressions: &[Expr]) -> Result<Expr> { } } +///Tests an expression to see if it contains only literal expressions such as 3+4.5 and immutable scalar builtins or UDFs. +pub fn expr_is_const(expr: &Expr) -> bool { + match expr { + Expr::Column(_) + | Expr::ScalarVariable(_) + | Expr::AggregateFunction { .. } + | Expr::WindowFunction { .. } + | Expr::AggregateUDF { .. } + | Expr::Sort { .. } + | Expr::Wildcard => false, + + Expr::Literal(_) => true, + + Expr::Alias(child, _) + | Expr::Not(child) + | Expr::IsNotNull(child) + | Expr::IsNull(child) + | Expr::Negative(child) + | Expr::Cast { expr: child, .. } + | Expr::TryCast { expr: child, .. } => expr_is_const(child), + + Expr::ScalarFunction { fun, args } => match fun.volatility() { + Volatility::Immutable => args.iter().all(|arg| expr_is_const(arg)), + Volatility::Stable | Volatility::Volatile => false, + }, + Expr::BinaryExpr { left, right, .. } => { + expr_is_const(left) && expr_is_const(right) + } + Expr::Between { + expr, low, high, .. + } => expr_is_const(expr) && expr_is_const(low) && expr_is_const(high), + Expr::Case { + expr, + when_then_expr, + else_expr, + } => { + let expr_constant = expr.as_ref().map(|e| expr_is_const(e)).unwrap_or(true); + let else_constant = + else_expr.as_ref().map(|e| expr_is_const(e)).unwrap_or(true); + let when_then_constant = when_then_expr + .iter() + .all(|(w, th)| expr_is_const(w) && expr_is_const(th)); + expr_constant && else_constant && when_then_constant + } + + Expr::ScalarUDF { fun, args } => match fun.volatility() { + Volatility::Immutable => args.iter().all(|arg| expr_is_const(arg)), + Volatility::Stable | Volatility::Volatile => false, + }, + + Expr::InList { expr, list, .. } => { + expr_is_const(expr) && list.iter().all(|e| expr_is_const(e)) + } + } +} + +///Evaluates an expression if it only contains literal expressions. If a non-literal expression or non-immutable function is found then it returns an error. +pub fn evalute_const_expr(expr: &Expr) -> Result<ScalarValue> { Review comment: Do you plan to use this method in a PR soon? Otherwise don't expose this function if it is never used (YAGNI). Same would apply to `expr_is_const` that is only used here 😃. ########## File path: datafusion/src/optimizer/utils.rs ########## @@ -468,6 +473,117 @@ pub fn rewrite_expression(expr: &Expr, expressions: &[Expr]) -> Result<Expr> { } } +///Tests an expression to see if it contains only literal expressions such as 3+4.5 and immutable scalar builtins or UDFs. +pub fn expr_is_const(expr: &Expr) -> bool { + match expr { + Expr::Column(_) + | Expr::ScalarVariable(_) + | Expr::AggregateFunction { .. } + | Expr::WindowFunction { .. } + | Expr::AggregateUDF { .. } + | Expr::Sort { .. } + | Expr::Wildcard => false, + + Expr::Literal(_) => true, + + Expr::Alias(child, _) + | Expr::Not(child) + | Expr::IsNotNull(child) + | Expr::IsNull(child) + | Expr::Negative(child) + | Expr::Cast { expr: child, .. } + | Expr::TryCast { expr: child, .. } => expr_is_const(child), + + Expr::ScalarFunction { fun, args } => match fun.volatility() { + Volatility::Immutable => args.iter().all(|arg| expr_is_const(arg)), + Volatility::Stable | Volatility::Volatile => false, + }, + Expr::BinaryExpr { left, right, .. } => { + expr_is_const(left) && expr_is_const(right) + } + Expr::Between { + expr, low, high, .. + } => expr_is_const(expr) && expr_is_const(low) && expr_is_const(high), + Expr::Case { + expr, + when_then_expr, + else_expr, + } => { + let expr_constant = expr.as_ref().map(|e| expr_is_const(e)).unwrap_or(true); + let else_constant = + else_expr.as_ref().map(|e| expr_is_const(e)).unwrap_or(true); + let when_then_constant = when_then_expr + .iter() + .all(|(w, th)| expr_is_const(w) && expr_is_const(th)); + expr_constant && else_constant && when_then_constant + } + + Expr::ScalarUDF { fun, args } => match fun.volatility() { + Volatility::Immutable => args.iter().all(|arg| expr_is_const(arg)), + Volatility::Stable | Volatility::Volatile => false, + }, + + Expr::InList { expr, list, .. } => { + expr_is_const(expr) && list.iter().all(|e| expr_is_const(e)) + } + } +} + +///Evaluates an expression if it only contains literal expressions. If a non-literal expression or non-immutable function is found then it returns an error. +pub fn evalute_const_expr(expr: &Expr) -> Result<ScalarValue> { + if !expr_is_const(expr) { + return Err(DataFusionError::Execution("The expression was not contsant and could not be evaluated. This means it contained stable or volatile functions, aggregates, or column expressions.".to_owned())); + } + evaluate_const_expr_unchecked(expr) +} + +///Evaluates an expression. Note that this will return incorrect results if the expression contains function calls which change return values call to call, such as Now(). +pub fn evaluate_const_expr_unchecked(expr: &Expr) -> Result<ScalarValue> { Review comment: I understand this is somewhat covered by the tests in `constant_folding.rs`, but if it is exposed publicly it should have its own set of tests. -- 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: github-unsubscr...@arrow.apache.org For queries about this service, please contact Infrastructure at: us...@infra.apache.org