fix failed test from #12050

Signed-off-by: jayzhan211 <[email protected]>

datafusion/expr/src/expr_schema.rs

@@ -338,10 +338,14 @@ impl ExprSchemable for Expr {
             Expr::ScalarFunction(ScalarFunction { func, args }) => {
                 Ok(func.is_nullable(args, input_schema))
             }
-            Expr::AggregateFunction(AggregateFunction { func, .. }) => {
-                Ok(func.is_nullable())
+            Expr::AggregateFunction(AggregateFunction { func, args, .. }) => {
+                let nullables = args
+                    .iter()
+                    .map(|e| e.nullable(input_schema))
+                    .collect::<Result<Vec<_>>>()?;
+                Ok(func.is_nullable(&nullables))
             }
-            Expr::WindowFunction(WindowFunction { fun, .. }) => match fun {
+            Expr::WindowFunction(WindowFunction { fun, args, .. }) => match fun {
                 WindowFunctionDefinition::BuiltInWindowFunction(func) => {
                     if func.name() == "RANK"
                         || func.name() == "NTILE"
@@ -352,7 +356,13 @@ impl ExprSchemable for Expr {
                         Ok(true)
                     }
                 }
-                WindowFunctionDefinition::AggregateUDF(func) => Ok(func.is_nullable()),
+                WindowFunctionDefinition::AggregateUDF(func) => {
+                    let nullables = args
+                        .iter()
+                        .map(|e| e.nullable(input_schema))
+                        .collect::<Result<Vec<_>>>()?;
+                    Ok(func.is_nullable(&nullables))
+                }
                 WindowFunctionDefinition::WindowUDF(udwf) => Ok(udwf.nullable()),
             },
             Expr::ScalarVariable(_, _)

datafusion/expr/src/udaf.rs

@@ -163,8 +163,8 @@ impl AggregateUDF {
         self.inner.name()
     }
 
-    pub fn is_nullable(&self) -> bool {
-        self.inner.is_nullable()
+    pub fn is_nullable(&self, nullables: &[bool]) -> bool {
+        self.inner.is_nullable(nullables)
     }
 
     /// Returns the aliases for this function.
@@ -355,8 +355,8 @@ pub trait AggregateUDFImpl: Debug + Send + Sync {
     ///
     /// Nullable means that that the function could return `null` for any inputs.
     /// For example, aggregate functions like `COUNT` always return a non null value
-    /// but others like `MIN` will return `NULL` if there is no non null input.
-    fn is_nullable(&self) -> bool {
+    /// but others like `MIN` will return `NULL` if there is nullable input.
+    fn is_nullable(&self, _nullables: &[bool]) -> bool {
         true
     }
 

datafusion/functions-aggregate/src/count.rs

@@ -121,7 +121,8 @@ impl AggregateUDFImpl for Count {
         Ok(DataType::Int64)
     }
 
-    fn is_nullable(&self) -> bool {
+    // Count is always nullable regardless of the input nullability
+    fn is_nullable(&self, _nullables: &[bool]) -> bool {
         false
     }
 

datafusion/functions/src/core/arrow_cast.rs

@@ -26,7 +26,9 @@ use datafusion_common::{
 };
 
 use datafusion_expr::simplify::{ExprSimplifyResult, SimplifyInfo};
-use datafusion_expr::{ColumnarValue, Expr, ScalarUDFImpl, Signature, Volatility};
+use datafusion_expr::{
+    ColumnarValue, Expr, ExprSchemable, ScalarUDFImpl, Signature, Volatility,
+};
 
 /// Implements casting to arbitrary arrow types (rather than SQL types)
 ///
@@ -87,6 +89,10 @@ impl ScalarUDFImpl for ArrowCastFunc {
         internal_err!("arrow_cast should return type from exprs")
     }
 
+    fn is_nullable(&self, args: &[Expr], schema: &dyn ExprSchema) -> bool {
+        args.iter().any(|e| e.nullable(schema).ok().unwrap_or(true))
+    }
+
     fn return_type_from_exprs(
         &self,
         args: &[Expr],

datafusion/physical-expr-functions-aggregate/src/aggregate.rs

@@ -56,12 +56,10 @@ pub struct AggregateExprBuilder {
     is_distinct: bool,
     /// Whether the expression is reversed
     is_reversed: bool,
-    is_nullable: bool,
 }
 
 impl AggregateExprBuilder {
     pub fn new(fun: Arc<AggregateUDF>, args: Vec<Arc<dyn PhysicalExpr>>) -> Self {
-        let is_nullable = fun.is_nullable();
         Self {
             fun,
             args,
@@ -71,7 +69,6 @@ impl AggregateExprBuilder {
             ignore_nulls: false,
             is_distinct: false,
             is_reversed: false,
-            is_nullable,
         }
     }
 
@@ -85,7 +82,6 @@ impl AggregateExprBuilder {
             ignore_nulls,
             is_distinct,
             is_reversed,
-            is_nullable,
         } = self;
         if args.is_empty() {
             return internal_err!("args should not be empty");
@@ -107,13 +103,19 @@ impl AggregateExprBuilder {
             .map(|arg| arg.data_type(&schema))
             .collect::<Result<Vec<_>>>()?;
 
+        let input_nullables = args
+            .iter()
+            .map(|arg| arg.nullable(&schema))
+            .collect::<Result<Vec<_>>>()?;
+
         check_arg_count(
             fun.name(),
             &input_exprs_types,
             &fun.signature().type_signature,
         )?;
 
         let data_type = fun.return_type(&input_exprs_types)?;
+        let is_nullable = fun.is_nullable(&input_nullables);
         let name = match alias {
             // TODO: Ideally, we should build the name from physical expressions
             None => create_function_physical_name(fun.name(), is_distinct, &[], None)?,