chore: factor out sqrt from SparkLike & DuckDB (#2876)

Author: ym-pett

narwhals/_duckdb/expr.py

@@ -303,12 +303,6 @@ def _log(expr: Expression) -> Expression:
 
         return self._with_elementwise(_log)
 
-    def sqrt(self) -> Self:
-        def _sqrt(expr: Expression) -> Expression:
-            return when(expr < lit(0), lit(float("nan"))).otherwise(F("sqrt", expr))
-
-        return self._with_elementwise(_sqrt)
-
     @property
     def str(self) -> DuckDBExprStringNamespace:
         return DuckDBExprStringNamespace(self)

narwhals/_duckdb/namespace.py

@@ -63,8 +63,15 @@ def _function(self, name: str, *args: Expression) -> Expression:  # type: ignore
     def _lit(self, value: Any) -> Expression:
         return lit(value)
 
-    def _when(self, condition: Expression, value: Expression) -> Expression:
-        return when(condition, value)
+    def _when(
+        self,
+        condition: Expression,
+        value: Expression,
+        otherwise: Expression | None = None,
+    ) -> Expression:
+        if otherwise is None:
+            return when(condition, value)
+        return when(condition, value).otherwise(otherwise)
 
     def _coalesce(self, *exprs: Expression) -> Expression:
         return CoalesceOperator(*exprs)

narwhals/_ibis/expr.py

@@ -335,12 +335,6 @@ def _log(expr: ir.NumericColumn) -> ir.Value:
 
         return self._with_callable(_log)
 
-    def sqrt(self) -> Self:
-        def _sqrt(expr: ir.NumericColumn) -> ir.Value:
-            return ibis.cases((expr < lit(0), lit(float("nan"))), else_=expr.sqrt())
-
-        return self._with_callable(_sqrt)
-
     @property
     def str(self) -> IbisExprStringNamespace:
         return IbisExprStringNamespace(self)

narwhals/_ibis/namespace.py

@@ -51,8 +51,12 @@ def _function(self, name: str, *args: ir.Value | PythonLiteral) -> ir.Value:
     def _lit(self, value: Any) -> ir.Value:
         return lit(value)
 
-    def _when(self, condition: ir.Value, value: ir.Value) -> ir.Value:
-        return ibis.cases((condition, value))
+    def _when(
+        self, condition: ir.Value, value: ir.Value, otherwise: ir.Expr | None = None
+    ) -> ir.Value:
+        if otherwise is None:
+            return ibis.cases((condition, value))
+        return ibis.cases((condition, value), else_=otherwise)  # pragma: no cover
 
     def _coalesce(self, *exprs: ir.Value) -> ir.Value:
         return ibis.coalesce(*exprs)

narwhals/_spark_like/expr.py

@@ -387,14 +387,6 @@ def _log(expr: Column) -> Column:
 
         return self._with_elementwise(_log)
 
-    def sqrt(self) -> Self:
-        def _sqrt(expr: Column) -> Column:
-            return self._F.when(expr < 0, self._F.lit(float("nan"))).otherwise(
-                self._F.sqrt(expr)
-            )
-
-        return self._with_elementwise(_sqrt)
-
     @property
     def str(self) -> SparkLikeExprStringNamespace:
         return SparkLikeExprStringNamespace(self)

narwhals/_spark_like/namespace.py

@@ -73,8 +73,12 @@ def _function(self, name: str, *args: Column | PythonLiteral) -> Column:
     def _lit(self, value: Any) -> Column:
         return self._F.lit(value)
 
-    def _when(self, condition: Column, value: Column) -> Column:
-        return self._F.when(condition, value)
+    def _when(
+        self, condition: Column, value: Column, otherwise: Column | None = None
+    ) -> Column:
+        if otherwise is None:
+            return self._F.when(condition, value)
+        return self._F.when(condition, value).otherwise(otherwise)
 
     def _coalesce(self, *exprs: Column) -> Column:
         return self._F.coalesce(*exprs)

narwhals/_sql/expr.py

@@ -181,14 +181,19 @@ def _function(self, name: str, *args: NativeExprT | PythonLiteral) -> NativeExpr
     def _lit(self, value: Any) -> NativeExprT:
         return self.__narwhals_namespace__()._lit(value)
 
-    def _when(self, condition: NativeExprT, value: NativeExprT) -> NativeExprT:
-        return self.__narwhals_namespace__()._when(condition, value)
-
     def _coalesce(self, *expr: NativeExprT) -> NativeExprT:
         return self.__narwhals_namespace__()._coalesce(*expr)
 
     def _count_star(self) -> NativeExprT: ...
 
+    def _when(
+        self,
+        condition: NativeExprT,
+        value: NativeExprT,
+        otherwise: NativeExprT | None = None,
+    ) -> NativeExprT:
+        return self.__narwhals_namespace__()._when(condition, value, otherwise)
+
     def _window_expression(
         self,
         expr: NativeExprT,
@@ -492,6 +497,16 @@ def round(self, decimals: int) -> Self:
             lambda expr: self._function("round", expr, self._lit(decimals))
         )
 
+    def sqrt(self) -> Self:
+        def _sqrt(expr: NativeExprT) -> NativeExprT:
+            return self._when(
+                expr < self._lit(0),  # type: ignore[operator]
+                self._lit(float("nan")),
+                self._function("sqrt", expr),
+            )
+
+        return self._with_elementwise(_sqrt)
+
     def exp(self) -> Self:
         return self._with_elementwise(lambda expr: self._function("exp", expr))
 

narwhals/_sql/namespace.py

@@ -20,7 +20,12 @@ class SQLNamespace(
 ):
     def _function(self, name: str, *args: NativeExprT | PythonLiteral) -> NativeExprT: ...
     def _lit(self, value: Any) -> NativeExprT: ...
-    def _when(self, condition: NativeExprT, value: NativeExprT) -> NativeExprT: ...
+    def _when(
+        self,
+        condition: NativeExprT,
+        value: NativeExprT,
+        otherwise: NativeExprT | None = None,
+    ) -> NativeExprT: ...
     def _coalesce(self, *exprs: NativeExprT) -> NativeExprT: ...
 
     # Horizontal functions