Repository: narwhals-dev/narwhals
Pull request #489 (2024-07-11 13:57:14)

feat: Added Expr and Series Count method
<!--
# Thanks for contributing a pull request! 
## Please make sure you see our contribution guidelines: https://github.com/narwhals-dev/narwhals/blob/main/CONTRIBUTING.md
-->

## What type of PR is this? (check all applicable)

- [ ] 💾 Refactor
- [x] ✨ Feature
- [ ] 🐛 Bug Fix
- [ ] 🔧 Optimization
- [ ] 📝 Documentation
- [ ] ✅ Test
- [ ] 🐳 Other

## Related issues 

- Related issue # 
- Closes #

## Checklist

- [ ] Code follows style guide (ruff)
- [x] Tests added 
- [x] Documented the changes

## If you have comments or can explain your changes, please do so below.



MarcoGorelli: awesome @ugohuche ! now you just need `count` to `docs/api-reference/series.md` and `docs/api-reference/expr.md`
ugohuche: > awesome @ugohuche ! now you just need `count` to add `docs/api-reference/series.md` and `docs/api-reference/expr.md`

I just added it now, but I have one last failing check and I don't quite understand it 🥲

Patch:
diff --git a/docs/api-reference/expressions.md b/docs/api-reference/expressions.md
index 9068b583..461df923 100644
--- a/docs/api-reference/expressions.md
+++ b/docs/api-reference/expressions.md
@@ -9,6 +9,7 @@
         - all
         - any
         - cast
+        - count
         - cum_sum
         - diff
         - drop_nulls
diff --git a/docs/api-reference/series.md b/docs/api-reference/series.md
index 525f5af6..2170fe7c 100644
--- a/docs/api-reference/series.md
+++ b/docs/api-reference/series.md
@@ -9,6 +9,7 @@
         - all
         - any
         - cast
+        - count
         - cum_sum
         - diff
         - drop_nulls
diff --git a/narwhals/_arrow/expr.py b/narwhals/_arrow/expr.py
index 45c0c719..b450c382 100644
--- a/narwhals/_arrow/expr.py
+++ b/narwhals/_arrow/expr.py
@@ -89,6 +89,9 @@ class ArrowExpr:
     def mean(self) -> Self:
         return reuse_series_implementation(self, "mean", returns_scalar=True)
 
+    def count(self) -> Self:
+        return reuse_series_implementation(self, "count", returns_scalar=True)
+
     def std(self, ddof: int = 1) -> Self:
         return reuse_series_implementation(self, "std", ddof=ddof, returns_scalar=True)
 
diff --git a/narwhals/_arrow/series.py b/narwhals/_arrow/series.py
index 4bfa49e2..2054b783 100644
--- a/narwhals/_arrow/series.py
+++ b/narwhals/_arrow/series.py
@@ -86,6 +86,10 @@ class ArrowSeries:
         pc = get_pyarrow_compute()
         return pc.stddev(self._native_series, ddof=ddof)  # type: ignore[no-any-return]
 
+    def count(self) -> int:
+        pc = get_pyarrow_compute()
+        return pc.count(self._native_series)  # type: ignore[no-any-return]
+
     def __narwhals_namespace__(self) -> ArrowNamespace:
         return ArrowNamespace(backend_version=self._backend_version)
 
diff --git a/narwhals/_pandas_like/expr.py b/narwhals/_pandas_like/expr.py
index 30719f26..d63833f0 100644
--- a/narwhals/_pandas_like/expr.py
+++ b/narwhals/_pandas_like/expr.py
@@ -174,6 +174,9 @@ class PandasLikeExpr:
     def sum(self) -> Self:
         return reuse_series_implementation(self, "sum", returns_scalar=True)
 
+    def count(self) -> Self:
+        return reuse_series_implementation(self, "count", returns_scalar=True)
+
     def mean(self) -> Self:
         return reuse_series_implementation(self, "mean", returns_scalar=True)
 
diff --git a/narwhals/_pandas_like/series.py b/narwhals/_pandas_like/series.py
index 7b8cf0d7..0df1e08f 100644
--- a/narwhals/_pandas_like/series.py
+++ b/narwhals/_pandas_like/series.py
@@ -374,6 +374,10 @@ class PandasLikeSeries:
         ser = self._native_series
         return ser.sum()
 
+    def count(self) -> Any:
+        ser = self._native_series
+        return ser.count()
+
     def mean(self) -> Any:
         ser = self._native_series
         return ser.mean()
diff --git a/narwhals/expression.py b/narwhals/expression.py
index 7703952a..79013039 100644
--- a/narwhals/expression.py
+++ b/narwhals/expression.py
@@ -497,6 +497,40 @@ class Expr:
         """
         return self.__class__(lambda plx: self._call(plx).max())
 
+    def count(self) -> Self:
+        """
+        Returns the number of non-null elements in the column.
+
+        Examples:
+            >>> import polars as pl
+            >>> import pandas as pd
+            >>> import narwhals as nw
+            >>> df_pd = pd.DataFrame({"a": [1, 2, 3], "b": [None, 4, 4]})
+            >>> df_pl = pl.DataFrame({"a": [1, 2, 3], "b": [None, 4, 4]})
+
+            Let's define a dataframe-agnostic function:
+
+            >>> @nw.narwhalify
+            ... def func(df):
+            ...     return df.select(nw.all().count())
+
+            We can then pass either pandas or Polars to `func`:
+
+            >>> func(df_pd)
+               a  b
+            0  3  2
+            >>> func(df_pl)
+            shape: (1, 2)
+            ┌─────┬─────┐
+            │ a   ┆ b   │
+            │ --- ┆ --- │
+            │ u32 ┆ u32 │
+            ╞═════╪═════╡
+            │ 3   ┆ 2   │
+            └─────┴─────┘
+        """
+        return self.__class__(lambda plx: self._call(plx).count())
+
     def n_unique(self) -> Self:
         """
          Returns count of unique values
diff --git a/narwhals/series.py b/narwhals/series.py
index b0cd69d2..c7e2c2dc 100644
--- a/narwhals/series.py
+++ b/narwhals/series.py
@@ -348,12 +348,40 @@ class Series:
             We can then pass either pandas or Polars to `func`:
 
             >>> func(s_pd)  # doctest:+SKIP
-            2.0
+            np.float64(2.0)
             >>> func(s_pl)
             2.0
         """
         return self._compliant_series.mean()
 
+    def count(self) -> Any:
+        """
+        Returns the number of non-null elements in the Series.
+
+        Examples:
+            >>> import pandas as pd
+            >>> import polars as pl
+            >>> import narwhals as nw
+            >>> s = [1, 2, 3]
+            >>> s_pd = pd.Series(s)
+            >>> s_pl = pl.Series(s)
+
+            We define a library agnostic function:
+
+            >>> @nw.narwhalify
+            ... def func(s_any):
+            ...     return s_any.count()
+
+            We can then pass either pandas or Polars to `func`:
+
+            >>> func(s_pd)  # doctest:+SKIP
+            np.int64(3)
+            >>> func(s_pl)
+            3
+
+        """
+        return self._compliant_series.count()
+
     def any(self) -> Any:
         """
         Return whether any of the values in the Series are True.
@@ -378,7 +406,7 @@ class Series:
             We can then pass either pandas or Polars to `func`:
 
             >>> func(s_pd)  # doctest:+SKIP
-            True
+            np.True_
             >>> func(s_pl)
             True
         """
@@ -405,7 +433,7 @@ class Series:
             We can then pass either pandas or Polars to `func`:
 
             >>> func(s_pd)  # doctest:+SKIP
-            False
+            np.False_
             >>> func(s_pl)
             False
 
@@ -433,7 +461,7 @@ class Series:
             We can then pass either pandas or Polars to `func`:
 
             >>> func(s_pd)  # doctest:+SKIP
-            1
+            np.int64(1)
             >>> func(s_pl)
             1
         """
@@ -460,7 +488,7 @@ class Series:
             We can then pass either pandas or Polars to `func`:
 
             >>> func(s_pd)  # doctest:+SKIP
-            3
+            np.int64(3)
             >>> func(s_pl)
             3
         """
@@ -487,7 +515,7 @@ class Series:
             We can then pass either pandas or Polars to `func`:
 
             >>> func(s_pd)  # doctest:+SKIP
-            6
+            np.int64(6)
             >>> func(s_pl)
             6
         """
@@ -518,7 +546,7 @@ class Series:
             We can then pass either pandas or Polars to `func`:
 
             >>> func(s_pd)  # doctest:+SKIP
-            1.0
+            np.float64(1.0)
             >>> func(s_pl)
             1.0
         """


Test patch:
diff --git a/tests/expr/count_test.py b/tests/expr/count_test.py
new file mode 100644
index 00000000..32f6d38a
--- /dev/null
+++ b/tests/expr/count_test.py
@@ -0,0 +1,12 @@
+from typing import Any
+
+import narwhals.stable.v1 as nw
+from tests.utils import compare_dicts
+
+
+def test_count(constructor_with_pyarrow: Any) -> None:
+    data = {"a": [1, 3, 2], "b": [4, None, 6], "z": [7.0, None, None]}
+    df = nw.from_native(constructor_with_pyarrow(data), eager_only=True)
+    result = df.select(nw.col("a", "b", "z").count())
+    expected = {"a": [3], "b": [2], "z": [1]}
+    compare_dicts(result, expected)
diff --git a/tests/series/test_common.py b/tests/series/test_common.py
index 2a573958..809bd5e6 100644
--- a/tests/series/test_common.py
+++ b/tests/series/test_common.py
@@ -116,6 +116,7 @@ def test_reductions(df_raw: Any) -> None:
     assert s.std() == 1.0
     assert s.min() == 1
     assert s.max() == 3
+    assert s.count() == 3
     assert s.sum() == 6
     assert nw.to_native(s.is_between(1, 2))[0]
     assert not nw.to_native(s.is_between(1, 2))[1]