refactor coalesce to support columns as input

dask_sql/physical/rex/core/call.py

@@ -577,16 +577,18 @@ def __init__(self):
         super().__init__(self.coalesce)
 
     def coalesce(self, *operands):
+        result = None
         for operand in operands:
             if is_frame(operand):
                 # Check if frame evaluates to nan or NA
-                if not operand.isnull().all().compute():
-                    return operand
+                if len(operand) == 1 and not operand.isnull().all().compute():
+                    return operand if result is None else result.fillna(operand)
                 else:
-                    continue
+                    result = operand if result is None else result.fillna(operand)
+            elif not pd.isna(operand):
+                return operand if result is None else result.fillna(operand)
 
-            if not pd.isna(operand):
-                return operand
+        return result
 
 
 class ExtractOperation(Operation):

tests/integration/test_rex.py

@@ -386,9 +386,7 @@ def test_coalesce(c, gpu):
         }
     )
 
-    df["c2"] = df["c2"].astype("float64")
-    df["c5"] = df["c5"].astype("O")
-    assert_eq(df, expected_df)
+    assert_eq(df, expected_df, check_dtype=False)
     c.drop_table("df")