fix: Remove supported coalesce types

viirya · Mar 5, 2024 · ef4760e · ef4760e
1 parent 2a490e4
commit ef4760e
Show file tree

Hide file tree

Showing 3 changed files with 58 additions and 26 deletions.
diff --git a/datafusion/expr/src/built_in_function.rs b/datafusion/expr/src/built_in_function.rs
@@ -26,9 +26,7 @@ use std::sync::{Arc, OnceLock};
 use crate::signature::TIMEZONE_WILDCARD;
 use crate::type_coercion::binary::get_wider_type;
 use crate::type_coercion::functions::data_types;
-use crate::{
-    conditional_expressions, FuncMonotonicity, Signature, TypeSignature, Volatility,
-};
+use crate::{FuncMonotonicity, Signature, TypeSignature, Volatility};
 
 use arrow::datatypes::{DataType, Field, Fields, IntervalUnit, TimeUnit};
 use datafusion_common::{exec_err, plan_err, DataFusionError, Result};
@@ -941,10 +939,9 @@ impl BuiltinScalarFunction {
             | BuiltinScalarFunction::ConcatWithSeparator => {
                 Signature::variadic(vec![Utf8], self.volatility())
             }
-            BuiltinScalarFunction::Coalesce => Signature::variadic(
-                conditional_expressions::SUPPORTED_COALESCE_TYPES.to_vec(),
-                self.volatility(),
-            ),
+            BuiltinScalarFunction::Coalesce => {
+                Signature::variadic_equal(self.volatility())
+            }
             BuiltinScalarFunction::SHA224
             | BuiltinScalarFunction::SHA256
             | BuiltinScalarFunction::SHA384
@@ -1643,4 +1640,13 @@ mod tests {
             assert_eq!(func_from_str, *func_original);
         }
     }
+
+    #[test]
+    fn test_coalesce_return_types() {
+        let coalesce = BuiltinScalarFunction::Coalesce;
+        let return_type = coalesce
+            .return_type(&[DataType::Date32, DataType::Date32])
+            .unwrap();
+        assert_eq!(return_type, DataType::Date32);
+    }
 }
diff --git a/datafusion/expr/src/conditional_expressions.rs b/datafusion/expr/src/conditional_expressions.rs
@@ -22,25 +22,6 @@ use arrow::datatypes::DataType;
 use datafusion_common::{plan_err, DFSchema, Result};
 use std::collections::HashSet;
 
-/// Currently supported types by the coalesce function.
-/// The order of these types correspond to the order on which coercion applies
-/// This should thus be from least informative to most informative
-pub static SUPPORTED_COALESCE_TYPES: &[DataType] = &[
-    DataType::Boolean,
-    DataType::UInt8,
-    DataType::UInt16,
-    DataType::UInt32,
-    DataType::UInt64,
-    DataType::Int8,
-    DataType::Int16,
-    DataType::Int32,
-    DataType::Int64,
-    DataType::Float32,
-    DataType::Float64,
-    DataType::Utf8,
-    DataType::LargeUtf8,
-];
-
 /// Helper struct for building [Expr::Case]
 pub struct CaseBuilder {
     expr: Option<Box<Expr>>,

diff --git a/datafusion/sqllogictest/test_files/scalar.slt b/datafusion/sqllogictest/test_files/scalar.slt
@@ -1822,6 +1822,51 @@ SELECT COALESCE(c1 * c2, 0) FROM test
 statement ok
 drop table test
 
+# coalesce date32
+
+statement ok
+CREATE TABLE test(
+  d1_date DATE,
+  d2_date DATE,
+  d3_date DATE
+) as VALUES
+  ('2022-12-12','2022-12-12','2022-12-12'),
+  (NULL,'2022-12-11','2022-12-12'),
+  ('2022-12-12','2022-12-10','2022-12-12'),
+  ('2022-12-12',NULL,'2022-12-12'),
+  ('2022-12-12','2022-12-8','2022-12-12'),
+  ('2022-12-12','2022-12-7',NULL),
+  ('2022-12-12',NULL,'2022-12-12'),
+  (NULL,'2022-12-5','2022-12-12')
+;
+
+query D
+SELECT COALESCE(d1_date, d2_date, d3_date) FROM test
+----
+2022-12-12
+2022-12-11
+2022-12-12
+2022-12-12
+2022-12-12
+2022-12-12
+2022-12-12
+2022-12-05
+
+query T
+SELECT arrow_typeof(COALESCE(d1_date, d2_date, d3_date)) FROM test
+----
+Date32
+Date32
+Date32
+Date32
+Date32
+Date32
+Date32
+Date32
+
+statement ok
+drop table test
+
 statement ok
 CREATE TABLE test(
   i32 INT,