feat[expr-common]: support REE in coalesce

We were missing a couple of branches to unwrap REE in
type_union_resolution_coercion.

Signed-off-by: Alfonso Subiotto Marques <alfonso.subiotto@polarsignals.com>

Author: asubiotto

datafusion/expr-common/src/type_coercion/binary.rs

@@ -23,19 +23,19 @@ use std::sync::Arc;
 
 use crate::operator::Operator;
 
-use arrow::array::{Array, new_empty_array};
+use arrow::array::{new_empty_array, Array};
 use arrow::compute::can_cast_types;
 use arrow::datatypes::IntervalUnit::MonthDayNano;
 use arrow::datatypes::TimeUnit::*;
 use arrow::datatypes::{
-    DECIMAL32_MAX_PRECISION, DECIMAL32_MAX_SCALE, DECIMAL64_MAX_PRECISION,
-    DECIMAL64_MAX_SCALE, DECIMAL128_MAX_PRECISION, DECIMAL128_MAX_SCALE,
-    DECIMAL256_MAX_PRECISION, DECIMAL256_MAX_SCALE, DataType, Field, FieldRef, Fields,
-    TimeUnit,
+    DataType, Field, FieldRef,
+    Fields, TimeUnit, DECIMAL128_MAX_PRECISION,
+    DECIMAL128_MAX_SCALE, DECIMAL256_MAX_PRECISION, DECIMAL256_MAX_SCALE, DECIMAL32_MAX_PRECISION, DECIMAL32_MAX_SCALE, DECIMAL64_MAX_PRECISION,
+    DECIMAL64_MAX_SCALE,
 };
 use datafusion_common::{
-    Diagnostic, Result, Span, Spans, exec_err, internal_err, not_impl_err,
-    plan_datafusion_err, plan_err,
+    exec_err, internal_err, not_impl_err, plan_datafusion_err, plan_err, Diagnostic, Result,
+    Span, Spans,
 };
 use itertools::Itertools;
 
@@ -528,11 +528,12 @@ enum TypeCategory {
 impl From<&DataType> for TypeCategory {
     fn from(data_type: &DataType) -> Self {
         match data_type {
-            // Dict is a special type in arrow, we check the value type
+            // Dict and REE are special types in arrow, we check the value type.
             DataType::Dictionary(_, v) => {
                 let v = v.as_ref();
                 TypeCategory::from(v)
             }
+            DataType::RunEndEncoded(_, v) => TypeCategory::from(v.data_type()),
             _ => {
                 if data_type.is_numeric() {
                     return TypeCategory::Numeric;
@@ -709,6 +710,27 @@ fn type_union_resolution_coercion(
                 None => None,
             }
         }
+        (
+            DataType::RunEndEncoded(lhs_run, lhs_val),
+            DataType::RunEndEncoded(rhs_run, rhs_val),
+        ) => {
+            let new_run =
+                type_union_resolution_coercion(lhs_run.data_type(), rhs_run.data_type())?;
+            let new_val =
+                type_union_resolution_coercion(lhs_val.data_type(), rhs_val.data_type())?;
+            Some(DataType::RunEndEncoded(
+                Arc::new(lhs_run.as_ref().clone().with_data_type(new_run)),
+                Arc::new(lhs_val.as_ref().clone().with_data_type(new_val)),
+            ))
+        }
+        (DataType::RunEndEncoded(run, val), other)
+        | (other, DataType::RunEndEncoded(run, val)) => {
+            let new_val = type_union_resolution_coercion(val.data_type(), other)?;
+            Some(DataType::RunEndEncoded(
+                Arc::clone(run),
+                Arc::new(val.as_ref().clone().with_data_type(new_val)),
+            ))
+        }
         (DataType::Struct(lhs), DataType::Struct(rhs)) => {
             if lhs.len() != rhs.len() {
                 return None;

datafusion/sqllogictest/test_files/coalesce.slt

@@ -443,3 +443,19 @@ query T
 select coalesce(arrow_cast('', 'Utf8View'), arrow_cast('', 'Dictionary(UInt32, Utf8)'));
 ----
 (empty)
+
+# RunEndEncoded column coalesced with a string literal
+statement ok
+create table ree_t as
+select arrow_cast(c, 'RunEndEncoded("run_ends": non-null Int32, "values": Utf8)') as ree_col
+from (values ('hello'), (NULL), ('world')) as t(c);
+
+query ?
+select coalesce(ree_col, '__null__') from ree_t;
+----
+hello
+__null__
+world
+
+statement ok
+drop table ree_t;