feat: Add support for {Expr,Series}.{ceil,floor} (#3198)

m-richards · FBruzzesi · pre-commit-ci[bot] · web-flow · commit c25d3f0acf55 · 2025-10-15T11:55:24.000+01:00
---------

Co-authored-by: Francesco Bruzzesi &lt;42817048+FBruzzesi@users.noreply.github.com&gt;
Co-authored-by: pre-commit-ci[bot] &lt;66853113+pre-commit-ci[bot]@users.noreply.github.com&gt;
diff --git a/.gitignore b/.gitignore
@@ -13,6 +13,7 @@ dist/
 coverage.xml
 .hypothesis/
 .pytest_cache/
+.mypy_cache
 
 # Documentation
 site/
diff --git a/docs/api-reference/expr.md b/docs/api-reference/expr.md
@@ -9,6 +9,7 @@
         - all
         - any
         - cast
+        - ceil
         - clip
         - count
         - cum_count
@@ -24,6 +25,7 @@
         - fill_null
         - filter
         - first
+        - floor
         - is_between
         - is_close
         - is_duplicated
diff --git a/docs/api-reference/series.md b/docs/api-reference/series.md
@@ -15,6 +15,7 @@
         - arg_min
         - arg_true
         - cast
+        - ceil
         - clip
         - count
         - cum_count
@@ -31,6 +32,7 @@
         - fill_null
         - filter
         - first
+        - floor
         - from_iterable
         - from_numpy
         - gather_every
diff --git a/narwhals/_arrow/series.py b/narwhals/_arrow/series.py
@@ -516,6 +516,12 @@ def round(self, decimals: int) -> Self:
             pc.round(self.native, decimals, round_mode="half_towards_infinity")
         )
 
+    def floor(self) -> Self:
+        return self._with_native(pc.floor(self.native))
+
+    def ceil(self) -> Self:
+        return self._with_native(pc.ceil(self.native))
+
     def diff(self) -> Self:
         return self._with_native(pc.pairwise_diff(self.native.combine_chunks()))
 
diff --git a/narwhals/_compliant/column.py b/narwhals/_compliant/column.py
@@ -198,6 +198,8 @@ def rolling_var(
         self, window_size: int, *, min_samples: int, center: bool, ddof: int
     ) -> Self: ...
     def round(self, decimals: int) -> Self: ...
+    def floor(self) -> Self: ...
+    def ceil(self) -> Self: ...
     def shift(self, n: int) -> Self: ...
     def unique(self) -> Self: ...
 
diff --git a/narwhals/_compliant/expr.py b/narwhals/_compliant/expr.py
@@ -709,6 +709,12 @@ def tail(self, n: int) -> Self:
     def round(self, decimals: int) -> Self:
         return self._reuse_series("round", decimals=decimals)
 
+    def floor(self) -> Self:
+        return self._reuse_series("floor")
+
+    def ceil(self) -> Self:
+        return self._reuse_series("ceil")
+
     def len(self) -> Self:
         return self._reuse_series("len", returns_scalar=True)
 
diff --git a/narwhals/_dask/expr.py b/narwhals/_dask/expr.py
@@ -445,6 +445,16 @@ def count(self) -> Self:
     def round(self, decimals: int) -> Self:
         return self._with_callable(lambda expr: expr.round(decimals), "round")
 
+    def floor(self) -> Self:
+        import dask.array as da
+
+        return self._with_callable(da.floor, "floor")
+
+    def ceil(self) -> Self:
+        import dask.array as da
+
+        return self._with_callable(da.ceil, "ceil")
+
     def unique(self) -> Self:
         return self._with_callable(lambda expr: expr.unique(), "unique")
 
diff --git a/narwhals/_pandas_like/series.py b/narwhals/_pandas_like/series.py
@@ -788,6 +788,62 @@ def tail(self, n: int) -> Self:
     def round(self, decimals: int) -> Self:
         return self._with_native(self.native.round(decimals=decimals))
 
+    def floor(self) -> Self:
+        native = self.native
+        native_cls = type(native)
+        implementation = self._implementation
+        if get_dtype_backend(native.dtype, implementation=implementation) == "pyarrow":
+            import pyarrow.compute as pc
+
+            from narwhals._arrow.utils import native_to_narwhals_dtype
+
+            ca = native.array._pa_array
+            result_arr = cast("ChunkedArrayAny", pc.floor(ca))
+            nw_dtype = native_to_narwhals_dtype(result_arr.type, self._version)
+            out_dtype = narwhals_to_native_dtype(
+                nw_dtype, "pyarrow", self._implementation, self._version
+            )
+            result_native = native_cls(
+                result_arr, dtype=out_dtype, index=native.index, name=native.name
+            )
+        else:
+            array_funcs = self._array_funcs
+            result_arr = array_funcs.floor(self.native)
+            result_native = (
+                native_cls(result_arr, index=native.index, name=native.name)
+                if implementation.is_cudf()
+                else result_arr
+            )
+        return self._with_native(result_native)
+
+    def ceil(self) -> Self:
+        native = self.native
+        native_cls = type(native)
+        implementation = self._implementation
+        if get_dtype_backend(native.dtype, implementation=implementation) == "pyarrow":
+            import pyarrow.compute as pc
+
+            from narwhals._arrow.utils import native_to_narwhals_dtype
+
+            ca = native.array._pa_array
+            result_arr = cast("ChunkedArrayAny", pc.ceil(ca))
+            nw_dtype = native_to_narwhals_dtype(result_arr.type, self._version)
+            out_dtype = narwhals_to_native_dtype(
+                nw_dtype, "pyarrow", self._implementation, self._version
+            )
+            result_native = native_cls(
+                result_arr, dtype=out_dtype, index=native.index, name=native.name
+            )
+        else:
+            array_funcs = self._array_funcs
+            result_arr = array_funcs.ceil(self.native)
+            result_native = (
+                native_cls(result_arr, index=native.index, name=native.name)
+                if implementation.is_cudf()
+                else result_arr
+            )
+        return self._with_native(result_native)
+
     def to_dummies(self, *, separator: str, drop_first: bool) -> PandasLikeDataFrame:
         from narwhals._pandas_like.dataframe import PandasLikeDataFrame
 
diff --git a/narwhals/_polars/expr.py b/narwhals/_polars/expr.py
@@ -340,6 +340,7 @@ def struct(self) -> PolarsExprStructNamespace:
     arg_max: Method[Self]
     arg_min: Method[Self]
     arg_true: Method[Self]
+    ceil: Method[Self]
     clip: Method[Self]
     count: Method[Self]
     cum_max: Method[Self]
@@ -352,6 +353,7 @@ def struct(self) -> PolarsExprStructNamespace:
     fill_null: Method[Self]
     fill_nan: Method[Self]
     first: Method[Self]
+    floor: Method[Self]
     last: Method[Self]
     gather_every: Method[Self]
     head: Method[Self]
diff --git a/narwhals/_polars/series.py b/narwhals/_polars/series.py
@@ -81,6 +81,7 @@
         "arg_max",
         "arg_min",
         "arg_true",
+        "ceil",
         "clip",
         "count",
         "cum_max",
@@ -93,6 +94,7 @@
         "fill_null",
         "fill_nan",
         "filter",
+        "floor",
         "gather_every",
         "head",
         "is_between",
@@ -699,6 +701,7 @@ def struct(self) -> PolarsSeriesStructNamespace:
     arg_max: Method[int]
     arg_min: Method[int]
     arg_true: Method[Self]
+    ceil: Method[Self]
     clip: Method[Self]
     count: Method[int]
     cum_max: Method[Self]
@@ -711,6 +714,7 @@ def struct(self) -> PolarsSeriesStructNamespace:
     fill_null: Method[Self]
     fill_nan: Method[Self]
     filter: Method[Self]
+    floor: Method[Self]
     gather_every: Method[Self]
     head: Method[Self]
     is_between: Method[Self]
diff --git a/narwhals/_sql/expr.py b/narwhals/_sql/expr.py
@@ -569,6 +569,12 @@ def round(self, decimals: int) -> Self:
             lambda expr: self._function("round", expr, self._lit(decimals))
         )
 
+    def floor(self) -> Self:
+        return self._with_elementwise(lambda expr: self._function("floor", expr))
+
+    def ceil(self) -> Self:
+        return self._with_elementwise(lambda expr: self._function("ceil", expr))
+
     def sqrt(self) -> Self:
         def _sqrt(expr: NativeExprT) -> NativeExprT:
             return self._when(
diff --git a/narwhals/expr.py b/narwhals/expr.py
@@ -1536,6 +1536,52 @@ def round(self, decimals: int = 0) -> Self:
             lambda plx: self._to_compliant_expr(plx).round(decimals)
         )
 
+    def floor(self) -> Self:
+        r"""Compute the numerical floor.
+
+        Examples:
+            >>> import pyarrow as pa
+            >>> import narwhals as nw
+            >>> df_native = pa.table({"values": [1.1, 4.3, -1.3]})
+            >>> df = nw.from_native(df_native)
+            >>> result = df.with_columns(floor=nw.col("values").floor())
+            >>> result
+            ┌────────────────────────┐
+            |   Narwhals DataFrame   |
+            |------------------------|
+            |pyarrow.Table           |
+            |values: double          |
+            |floor: double           |
+            |----                    |
+            |values: [[1.1,4.3,-1.3]]|
+            |floor: [[1,4,-2]]       |
+            └────────────────────────┘
+        """
+        return self._with_elementwise(lambda plx: self._to_compliant_expr(plx).floor())
+
+    def ceil(self) -> Self:
+        r"""Compute the numerical ceiling.
+
+        Examples:
+            >>> import pyarrow as pa
+            >>> import narwhals as nw
+            >>> df_native = pa.table({"values": [1.1, 4.3, -1.3]})
+            >>> df = nw.from_native(df_native)
+            >>> result = df.with_columns(ceil=nw.col("values").ceil())
+            >>> result
+            ┌────────────────────────┐
+            |   Narwhals DataFrame   |
+            |------------------------|
+            |pyarrow.Table           |
+            |values: double          |
+            |ceil: double            |
+            |----                    |
+            |values: [[1.1,4.3,-1.3]]|
+            |ceil: [[2,5,-1]]        |
+            └────────────────────────┘
+        """
+        return self._with_elementwise(lambda plx: self._to_compliant_expr(plx).ceil())
+
     def len(self) -> Self:
         r"""Return the number of elements in the column.
 
diff --git a/narwhals/series.py b/narwhals/series.py
@@ -2045,6 +2045,46 @@ def round(self, decimals: int = 0) -> Self:
         """
         return self._with_compliant(self._compliant_series.round(decimals))
 
+    def floor(self) -> Self:
+        r"""Compute the numerical floor.
+
+        Examples:
+            >>> import pyarrow as pa
+            >>> import narwhals as nw
+            >>> s_native = pa.chunked_array([[1.1, 4.3, -1.3]])
+            >>> s = nw.from_native(s_native, series_only=True)
+            >>> s.floor().to_native()  # doctest:+ELLIPSIS
+            <pyarrow.lib.ChunkedArray object at ...>
+            [
+              [
+                1,
+                4,
+                -2
+              ]
+            ]
+        """
+        return self._with_compliant(self._compliant_series.floor())
+
+    def ceil(self) -> Self:
+        r"""Compute the numerical ceiling.
+
+        Examples:
+            >>> import pyarrow as pa
+            >>> import narwhals as nw
+            >>> s_native = pa.chunked_array([[1.1, 4.3, -1.3]])
+            >>> s = nw.from_native(s_native, series_only=True)
+            >>> s.ceil().to_native()  # doctest:+ELLIPSIS
+            <pyarrow.lib.ChunkedArray object at ...>
+            [
+              [
+                2,
+                5,
+                -1
+              ]
+            ]
+        """
+        return self._with_compliant(self._compliant_series.ceil())
+
     def to_dummies(
         self, *, separator: str = "_", drop_first: bool = False
     ) -> DataFrame[Any]:
diff --git a/tests/expr_and_series/floor_ceil_test.py b/tests/expr_and_series/floor_ceil_test.py
diff --git a/tests/expr_and_series/log_test.py b/tests/expr_and_series/log_test.py
