[branch-51] Revert rewrite for coalesce, nvl and nvl2 simplification (#18567)

Note this targets the branch-51 release branch

## Which issue does this PR close?

- part of #17558
- resolves #17801 in the 51
release branch

## Rationale for this change

- We merged some clever rewrites for `coalesce` and `nvl2` to use `CASE`
which are faster and more correct (👏 @chenkovsky @kosiew )
- However, these rewrites cause subtle schema mismatches in some cases
planning (b/c the CASE simplification nullability logic can't determine
the correct nullability in some cases - see
#17801)
- @pepijnve has some heroic efforts to fix the schema mismatch in
#17813 (comment),
but it is non trivial and I am worried about merging it so close to the
51 release and introducing new edge cases

## What changes are included in this PR?

1. Revert #17357 /
e5dcc8c
3. Revert #17991 /
ea83c26
2. Revert #18191 /
22c4214
2. Cherry-pick 6202254, a test that
reproduces the schema mismatch issue (from
#18536)
3. Cherry-pick 735cacf, a fix for the
benchmarks that regressed due to the revert (from
#17833)
4. Update datafusion-testing (see separate PR here) for extended tests
(see apache/datafusion-testing#15)

## Are these changes tested?

Yes I added a new test

## Are there any user-facing changes?

<!--
If there are user-facing changes then we may require documentation to be
updated before approving the PR.
-->

<!--
If there are any breaking changes to public APIs, please add the `api
change` label.
-->

Author: alamb

datafusion-testing

@@ -477,9 +477,6 @@ fn criterion_benchmark(c: &mut Criterion) {
     };
 
     let raw_tpcds_sql_queries = (1..100)
-        // skip query 75 until it is fixed
-        // https://github.com/apache/datafusion/issues/17801
-        .filter(|q| *q != 75)
         .map(|q| std::fs::read_to_string(format!("{tests_path}tpc-ds/{q}.sql")).unwrap())
         .collect::<Vec<_>>();
 

datafusion/core/benches/sql_planner.rs

@@ -274,33 +274,6 @@ async fn test_nvl2() -> Result<()> {
 
     Ok(())
 }
-
-#[tokio::test]
-async fn test_nvl2_short_circuit() -> Result<()> {
-    let expr = nvl2(
-        col("a"),
-        arrow_cast(lit("1"), lit("Int32")),
-        arrow_cast(col("a"), lit("Int32")),
-    );
-
-    let batches = get_batches(expr).await?;
-
-    assert_snapshot!(
-        batches_to_string(&batches),
-        @r#"
-    +-----------------------------------------------------------------------------------+
-    | nvl2(test.a,arrow_cast(Utf8("1"),Utf8("Int32")),arrow_cast(test.a,Utf8("Int32"))) |
-    +-----------------------------------------------------------------------------------+
-    | 1                                                                                 |
-    | 1                                                                                 |
-    | 1                                                                                 |
-    | 1                                                                                 |
-    +-----------------------------------------------------------------------------------+
-    "#
-    );
-
-    Ok(())
-}
 #[tokio::test]
 async fn test_fn_arrow_typeof() -> Result<()> {
     let expr = arrow_typeof(col("l"));

datafusion/core/tests/dataframe/dataframe_functions.rs

@@ -320,26 +320,6 @@ async fn test_create_physical_expr() {
     create_simplified_expr_test(lit(1i32) + lit(2i32), "3");
 }
 
-#[test]
-fn test_create_physical_expr_nvl2() {
-    let batch = &TEST_BATCH;
-    let df_schema = DFSchema::try_from(batch.schema()).unwrap();
-    let ctx = SessionContext::new();
-
-    let expect_err = |expr| {
-        let physical_expr = ctx.create_physical_expr(expr, &df_schema).unwrap();
-        let err = physical_expr.evaluate(batch).unwrap_err();
-        assert!(
-            err.to_string()
-                .contains("nvl2 should have been simplified to case"),
-            "unexpected error: {err:?}"
-        );
-    };
-
-    expect_err(nvl2(col("i"), lit(1i64), lit(0i64)));
-    expect_err(nvl2(lit(1i64), col("i"), lit(0i64)));
-}
-
 #[tokio::test]
 async fn test_create_physical_expr_coercion() {
     // create_physical_expr does apply type coercion and unwrapping in cast

datafusion/core/tests/expr_api/mod.rs

@@ -1051,10 +1051,13 @@ async fn regression_test(query_no: u8, create_physical: bool) -> Result<()> {
 
     for sql in &sql {
         let df = ctx.sql(sql).await?;
-        let (state, plan) = df.into_parts();
-        let plan = state.optimize(&plan)?;
-        if create_physical {
-            let _ = state.create_physical_plan(&plan).await?;
+        // attempt to mimic planning steps
+        if !create_physical {
+            let (state, plan) = df.into_parts();
+            let _ = state.optimize(&plan)?;
+        } else {
+            // this is what df.execute() does internally
+            let _ = df.create_physical_plan().await?;
         }
     }
 

datafusion/core/tests/tpcds_planning.rs

@@ -252,21 +252,7 @@ impl ScalarUDF {
         Ok(result)
     }
 
-    /// Determines which of the arguments passed to this function are evaluated eagerly
-    /// and which may be evaluated lazily.
-    ///
-    /// See [ScalarUDFImpl::conditional_arguments] for more information.
-    pub fn conditional_arguments<'a>(
-        &self,
-        args: &'a [Expr],
-    ) -> Option<(Vec<&'a Expr>, Vec<&'a Expr>)> {
-        self.inner.conditional_arguments(args)
-    }
-
-    /// Returns true if some of this `exprs` subexpressions may not be evaluated
-    /// and thus any side effects (like divide by zero) may not be encountered.
-    ///
-    /// See [ScalarUDFImpl::short_circuits] for more information.
+    /// Get the circuits of inner implementation
     pub fn short_circuits(&self) -> bool {
         self.inner.short_circuits()
     }
@@ -696,42 +682,10 @@ pub trait ScalarUDFImpl: Debug + DynEq + DynHash + Send + Sync {
     ///
     /// Setting this to true prevents certain optimizations such as common
     /// subexpression elimination
-    ///
-    /// When overriding this function to return `true`, [ScalarUDFImpl::conditional_arguments] can also be
-    /// overridden to report more accurately which arguments are eagerly evaluated and which ones
-    /// lazily.
     fn short_circuits(&self) -> bool {
         false
     }
 
-    /// Determines which of the arguments passed to this function are evaluated eagerly
-    /// and which may be evaluated lazily.
-    ///
-    /// If this function returns `None`, all arguments are eagerly evaluated.
-    /// Returning `None` is a micro optimization that saves a needless `Vec`
-    /// allocation.
-    ///
-    /// If the function returns `Some`, returns (`eager`, `lazy`) where `eager`
-    /// are the arguments that are always evaluated, and `lazy` are the
-    /// arguments that may be evaluated lazily (i.e. may not be evaluated at all
-    /// in some cases).
-    ///
-    /// Implementations must ensure that the two returned `Vec`s are disjunct,
-    /// and that each argument from `args` is present in one the two `Vec`s.
-    ///
-    /// When overriding this function, [ScalarUDFImpl::short_circuits] must
-    /// be overridden to return `true`.
-    fn conditional_arguments<'a>(
-        &self,
-        args: &'a [Expr],
-    ) -> Option<(Vec<&'a Expr>, Vec<&'a Expr>)> {
-        if self.short_circuits() {
-            Some((vec![], args.iter().collect()))
-        } else {
-            None
-        }
-    }
-
     /// Computes the output [`Interval`] for a [`ScalarUDFImpl`], given the input
     /// intervals.
     ///
@@ -921,13 +875,6 @@ impl ScalarUDFImpl for AliasedScalarUDFImpl {
         self.inner.simplify(args, info)
     }
 
-    fn conditional_arguments<'a>(
-        &self,
-        args: &'a [Expr],
-    ) -> Option<(Vec<&'a Expr>, Vec<&'a Expr>)> {
-        self.inner.conditional_arguments(args)
-    }
-
     fn short_circuits(&self) -> bool {
         self.inner.short_circuits()
     }

datafusion/expr/src/udf.rs

@@ -15,13 +15,14 @@
 // specific language governing permissions and limitations
 // under the License.
 
+use arrow::array::{new_null_array, BooleanArray};
+use arrow::compute::kernels::zip::zip;
+use arrow::compute::{and, is_not_null, is_null};
 use arrow::datatypes::{DataType, Field, FieldRef};
-use datafusion_common::{exec_err, internal_err, plan_err, Result};
+use datafusion_common::{exec_err, internal_err, Result};
 use datafusion_expr::binary::try_type_union_resolution;
-use datafusion_expr::conditional_expressions::CaseBuilder;
-use datafusion_expr::simplify::{ExprSimplifyResult, SimplifyInfo};
 use datafusion_expr::{
-    ColumnarValue, Documentation, Expr, ReturnFieldArgs, ScalarFunctionArgs,
+    ColumnarValue, Documentation, ReturnFieldArgs, ScalarFunctionArgs,
 };
 use datafusion_expr::{ScalarUDFImpl, Signature, Volatility};
 use datafusion_macros::user_doc;
@@ -47,7 +48,7 @@ use std::any::Any;
 )]
 #[derive(Debug, PartialEq, Eq, Hash)]
 pub struct CoalesceFunc {
-    pub(super) signature: Signature,
+    signature: Signature,
 }
 
 impl Default for CoalesceFunc {
@@ -94,45 +95,61 @@ impl ScalarUDFImpl for CoalesceFunc {
         Ok(Field::new(self.name(), return_type, nullable).into())
     }
 
-    fn simplify(
-        &self,
-        args: Vec<Expr>,
-        _info: &dyn SimplifyInfo,
-    ) -> Result<ExprSimplifyResult> {
+    /// 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.
         if args.is_empty() {
-            return plan_err!("coalesce must have at least one argument");
-        }
-        if args.len() == 1 {
-            return Ok(ExprSimplifyResult::Simplified(
-                args.into_iter().next().unwrap(),
-            ));
+            return exec_err!(
+                "coalesce was called with {} arguments. It requires at least 1.",
+                args.len()
+            );
         }
 
-        let n = args.len();
-        let (init, last_elem) = args.split_at(n - 1);
-        let whens = init
-            .iter()
-            .map(|x| x.clone().is_not_null())
-            .collect::<Vec<_>>();
-        let cases = init.to_vec();
-        Ok(ExprSimplifyResult::Simplified(
-            CaseBuilder::new(None, whens, cases, Some(Box::new(last_elem[0].clone())))
-                .end()?,
-        ))
-    }
-
-    /// coalesce evaluates to the first value which is not NULL
-    fn invoke_with_args(&self, _args: ScalarFunctionArgs) -> Result<ColumnarValue> {
-        internal_err!("coalesce should have been simplified to case")
-    }
-
-    fn conditional_arguments<'a>(
-        &self,
-        args: &'a [Expr],
-    ) -> Option<(Vec<&'a Expr>, Vec<&'a Expr>)> {
-        let eager = vec![&args[0]];
-        let lazy = args[1..].iter().collect();
-        Some((eager, lazy))
+        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, &current_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, &current_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)
+        }
     }
 
     fn short_circuits(&self) -> bool {