implement short_circuits function for ScalarUDFImpl trait

datafusion/expr/src/expr.rs

@@ -1266,7 +1266,7 @@ impl Expr {
     pub fn short_circuits(&self) -> bool {
         match self {
             Expr::ScalarFunction(ScalarFunction { func_def, .. }) => {
-                matches!(func_def, ScalarFunctionDefinition::UDF(fun) if fun.name().eq("coalesce"))
+                matches!(func_def, ScalarFunctionDefinition::UDF(fun) if fun.short_circuits())
             }
             Expr::BinaryExpr(BinaryExpr { op, .. }) => {
                 matches!(op, Operator::And | Operator::Or)

datafusion/expr/src/udf.rs

@@ -193,6 +193,11 @@ impl ScalarUDF {
     pub fn monotonicity(&self) -> Result<Option<FuncMonotonicity>> {
         self.inner.monotonicity()
     }
+
+    /// Get the circuits of inner implementation
+    pub fn short_circuits(&self) -> bool {
+        self.inner.short_circuits()
+    }
 }
 
 impl<F> From<F> for ScalarUDF
@@ -376,6 +381,13 @@ pub trait ScalarUDFImpl: Debug + Send + Sync {
     ) -> Result<ExprSimplifyResult> {
         Ok(ExprSimplifyResult::Original(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
+    /// Setting this to true prevents certain optimizations such as common subexpression elimination
+    fn short_circuits(&self) -> bool {
+        false
+    }
 }
 
 /// ScalarUDF that adds an alias to the underlying function. It is better to

datafusion/functions/src/math/coalesce.rs

@@ -120,6 +120,10 @@ impl ScalarUDFImpl for CoalesceFunc {
             Ok(result)
         }
     }
+
+    fn short_circuits(&self) -> bool {
+        true
+    }
 }
 
 #[cfg(test)]