alamb commented on code in PR #17357: URL: https://github.com/apache/datafusion/pull/17357#discussion_r2325943637
########## datafusion/sqllogictest/test_files/select.slt: ########## @@ -1656,10 +1656,10 @@ query TT explain select coalesce(1, y/x), coalesce(2, y/x) from t; ---- logical_plan -01)Projection: coalesce(Int64(1), CAST(t.y / t.x AS Int64)), coalesce(Int64(2), CAST(t.y / t.x AS Int64)) +01)Projection: CASE WHEN Boolean(true) THEN Int64(1) ELSE CAST(t.y / t.x AS Int64) END AS coalesce(Int64(1),t.y / t.x), CASE WHEN Boolean(true) THEN Int64(2) ELSE CAST(t.y / t.x AS Int64) END AS coalesce(Int64(2),t.y / t.x) 02)--TableScan: t projection=[x, y] physical_plan -01)ProjectionExec: expr=[coalesce(1, CAST(y@1 / x@0 AS Int64)) as coalesce(Int64(1),t.y / t.x), coalesce(2, CAST(y@1 / x@0 AS Int64)) as coalesce(Int64(2),t.y / t.x)] +01)ProjectionExec: expr=[CASE WHEN true THEN 1 ELSE CAST(y@1 / x@0 AS Int64) END as coalesce(Int64(1),t.y / t.x), CASE WHEN true THEN 2 ELSE CAST(y@1 / x@0 AS Int64) END as coalesce(Int64(2),t.y / t.x)] Review Comment: It seems like this should be able to be simplifed even more -- `CASE WHEN true THEN ...` should go to `1` I filed a ticket to track this idea: - https://github.com/apache/datafusion/issues/17448 ########## datafusion/functions/src/core/coalesce.rs: ########## @@ -95,61 +94,36 @@ impl ScalarUDFImpl for CoalesceFunc { Ok(Field::new(self.name(), return_type, nullable).into()) } - /// coalesce evaluates to the first value which is not NULL - fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> { - let args = args.args; - // do not accept 0 arguments. + fn simplify( + &self, + args: Vec<Expr>, + _info: &dyn SimplifyInfo, + ) -> Result<ExprSimplifyResult> { if args.is_empty() { - return exec_err!( - "coalesce was called with {} arguments. It requires at least 1.", - args.len() - ); + return plan_err!("coalesce must have at least one argument"); } - - let return_type = args[0].data_type(); - let mut return_array = args.iter().filter_map(|x| match x { - ColumnarValue::Array(array) => Some(array.len()), - _ => None, - }); - - if let Some(size) = return_array.next() { - // start with nulls as default output - let mut current_value = new_null_array(&return_type, size); - let mut remainder = BooleanArray::from(vec![true; size]); - - for arg in args { - match arg { - ColumnarValue::Array(ref array) => { - let to_apply = and(&remainder, &is_not_null(array.as_ref())?)?; - current_value = zip(&to_apply, array, ¤t_value)?; - remainder = and(&remainder, &is_null(array)?)?; - } - ColumnarValue::Scalar(value) => { - if value.is_null() { - continue; - } else { - let last_value = value.to_scalar()?; - current_value = zip(&remainder, &last_value, ¤t_value)?; - break; - } - } - } - if remainder.iter().all(|x| x == Some(false)) { - break; - } - } - Ok(ColumnarValue::Array(current_value)) - } else { - let result = args - .iter() - .filter_map(|x| match x { - ColumnarValue::Scalar(s) if !s.is_null() => Some(x.clone()), - _ => None, - }) - .next() - .unwrap_or_else(|| args[0].clone()); - Ok(result) + if args.len() == 1 { + return Ok(ExprSimplifyResult::Simplified( + args.into_iter().next().unwrap(), + )); } + + let n = args.len(); Review Comment: This is quite a clever implementation. However, I worry it may cause problems for comet which uses physical evaluation directly @comphead or @andygrove do you know if comet uses the `COALESCE` implementation directly? -- 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...@datafusion.apache.org For queries about this service, please contact Infrastructure at: us...@infra.apache.org --------------------------------------------------------------------- To unsubscribe, e-mail: github-unsubscr...@datafusion.apache.org For additional commands, e-mail: github-h...@datafusion.apache.org
