feat(python, rust): Adds NULLIF and COALESCE SQL functions

crates/polars-sql/src/functions.rs

@@ -1,6 +1,6 @@
 use polars_core::prelude::{polars_bail, polars_err, PolarsError, PolarsResult};
 use polars_lazy::dsl::Expr;
-use polars_plan::dsl::count;
+use polars_plan::dsl::{coalesce, count, when};
 use polars_plan::logical_plan::LiteralValue;
 use polars_plan::prelude::lit;
 use sqlparser::ast::{
@@ -232,6 +232,16 @@ pub(crate) enum PolarsSqlFunctions {
     /// SELECT UPPER(column_1) from df;
     /// ```
     Upper,
+    /// SQL 'nullif' function
+    /// ```sql
+    /// SELECT NULLIF(column_1, column_2) from df;
+    /// ```
+    NullIf,
+    /// SQL 'coalesce' function
+    /// ```sql
+    /// SELECT COALESCE(column_1, ...) from df;
+    /// ```
+    Coalesce,
 
     // ----
     // Aggregate functions
@@ -389,6 +399,7 @@ impl PolarsSqlFunctions {
             "cbrt",
             "ceil",
             "ceiling",
+            "coalesce",
             "cos",
             "cosd",
             "cot",
@@ -411,6 +422,7 @@ impl PolarsSqlFunctions {
             "ltrim",
             "max",
             "min",
+            "nullif",
             "octet_length",
             "pi",
             "pow",
@@ -476,6 +488,12 @@ impl TryFrom<&'_ SQLFunction> for PolarsSqlFunctions {
             "cbrt" => Self::Cbrt,
             "round" => Self::Round,
 
+            // ----
+            // Comparison functions
+            // ----
+            "nullif" => Self::NullIf,
+            "coalesce" => Self::Coalesce,
+
             // ----
             // String functions
             // ----
@@ -580,6 +598,13 @@ impl SqlFunctionVisitor<'_> {
                     polars_bail!(InvalidOperation:"Invalid number of arguments for Round: {}",function.args.len());
                 },
             },
+
+            // ----
+            // Comparison functions
+            // ----
+            NullIf => self.visit_binary(|l, r: Expr| when(l.clone().eq(r)).then(lit(LiteralValue::Null)).otherwise(l)),
+            Coalesce => self.visit_variadic(coalesce),
+
             // ----
             // String functions
             // ----
@@ -780,6 +805,24 @@ impl SqlFunctionVisitor<'_> {
         }
     }
 
+    fn visit_variadic(&self, f: impl Fn(&[Expr]) -> Expr) -> PolarsResult<Expr> {
+        self.try_visit_variadic(|e| Ok(f(e)))
+    }
+
+    fn try_visit_variadic(&self, f: impl Fn(&[Expr]) -> PolarsResult<Expr>) -> PolarsResult<Expr> {
+        let function = self.func;
+        let args = extract_args(function);
+        let mut expr_args = vec![];
+        for arg in args {
+            if let FunctionArgExpr::Expr(sql_expr) = arg {
+                expr_args.push(parse_sql_expr(sql_expr, self.ctx)?);
+            } else {
+                return self.not_supported_error();
+            };
+        }
+        f(&expr_args)
+    }
+
     // fn visit_ternary<Arg: FromSqlExpr>(
     //     &self,
     //     f: impl Fn(Expr, Arg, Arg) -> Expr,

py-polars/tests/unit/test_sql.py

@@ -773,6 +773,35 @@ def test_sql_trim(foods_ipc_path: Path) -> None:
     }
 
 
+def test_sql_nullif_coalesce(foods_ipc_path: Path) -> None:
+    nums = pl.LazyFrame(
+        {
+            "x": [1, None, 2, 3, None, 4],
+            "y": [5, 4, None, 3, None, 2],
+            "z": [3, 4, None, 3, None, None],
+        }
+    )
+
+    res = pl.SQLContext(df=nums).execute(
+        """
+        SELECT
+        COALESCE(x,y,z) as "coal",
+        NULLIF(x,y) as "nullif x_y",
+        NULLIF(y,z) as "nullif y_z",
+        COALESCE(x, NULLIF(y,z)) as "both"
+        FROM df
+        """,
+        eager=True,
+    )
+
+    assert res.to_dict(False) == {
+        "coal": [1, 4, 2, 3, None, 4],
+        "nullif x_y": [1, None, 2, None, None, 4],
+        "nullif y_z": [5, None, None, None, None, 2],
+        "both": [1, None, 2, 3, None, 4],
+    }
+
+
 def test_sql_order_by(foods_ipc_path: Path) -> None:
     foods = pl.scan_ipc(foods_ipc_path)
     nums = pl.LazyFrame(