Repository: narwhals-dev/narwhals
Pull request #568 (2024-07-21 14:51:49)

feat: add DataFrame__contains__, Expr.arg_true, and Series.arg_true
<!--
# 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
- [ ] ✨ Feature
- [ ] 🐛 Bug Fix
- [ ] 🔧 Optimization
- [ ] 📝 Documentation
- [ ] ✅ Test
- [ ] 🐳 Other

## Related issues 

- Related issue # 
- Closes #

## Checklist

- [ ] Code follows style guide (ruff)
- [ ] Tests added 
- [ ] Documented the changes

## If you have comments or can explain your changes, please do so below.




Patch:
diff --git a/docs/api-reference/expressions.md b/docs/api-reference/expressions.md
index 461df923..c70f4a5f 100644
--- a/docs/api-reference/expressions.md
+++ b/docs/api-reference/expressions.md
@@ -8,6 +8,7 @@
         - alias
         - all
         - any
+        - arg_true
         - cast
         - count
         - cum_sum
diff --git a/docs/api-reference/series.md b/docs/api-reference/series.md
index 2170fe7c..0d409408 100644
--- a/docs/api-reference/series.md
+++ b/docs/api-reference/series.md
@@ -8,6 +8,7 @@
         - alias
         - all
         - any
+        - arg_true
         - cast
         - count
         - cum_sum
diff --git a/narwhals/_arrow/expr.py b/narwhals/_arrow/expr.py
index c14dd99f..73111bff 100644
--- a/narwhals/_arrow/expr.py
+++ b/narwhals/_arrow/expr.py
@@ -216,6 +216,9 @@ class ArrowExpr:
     def is_in(self, other: ArrowExpr | Any) -> Self:
         return reuse_series_implementation(self, "is_in", other)
 
+    def arg_true(self) -> Self:
+        return reuse_series_implementation(self, "arg_true")
+
     def sample(
         self: Self,
         n: int | None = None,
diff --git a/narwhals/_arrow/series.py b/narwhals/_arrow/series.py
index 3a3ec226..cf800cbf 100644
--- a/narwhals/_arrow/series.py
+++ b/narwhals/_arrow/series.py
@@ -311,6 +311,14 @@ class ArrowSeries:
         ser = self._native_series
         return self._from_native_series(pc.is_in(ser, value_set=value_set))
 
+    def arg_true(self) -> Self:
+        np = get_numpy()
+        ser = self._native_series
+        res = np.flatnonzero(ser)
+        return self._from_iterable(
+            res, name=self.name, backend_version=self._backend_version
+        )
+
     def item(self: Self, index: int | None = None) -> Any:
         if index is None:
             if len(self) != 1:
diff --git a/narwhals/_pandas_like/expr.py b/narwhals/_pandas_like/expr.py
index b34ce5b6..2641bf7a 100644
--- a/narwhals/_pandas_like/expr.py
+++ b/narwhals/_pandas_like/expr.py
@@ -216,6 +216,9 @@ class PandasLikeExpr:
     def is_in(self, other: Any) -> Self:
         return reuse_series_implementation(self, "is_in", other=other)
 
+    def arg_true(self) -> Self:
+        return reuse_series_implementation(self, "arg_true")
+
     def filter(self, *predicates: Any) -> Self:
         from narwhals._pandas_like.namespace import PandasLikeNamespace
 
diff --git a/narwhals/_pandas_like/series.py b/narwhals/_pandas_like/series.py
index 042a5eb0..81d65a11 100644
--- a/narwhals/_pandas_like/series.py
+++ b/narwhals/_pandas_like/series.py
@@ -15,6 +15,7 @@ from narwhals._pandas_like.utils import translate_dtype
 from narwhals._pandas_like.utils import validate_column_comparand
 from narwhals.dependencies import get_cudf
 from narwhals.dependencies import get_modin
+from narwhals.dependencies import get_numpy
 from narwhals.dependencies import get_pandas
 from narwhals.dependencies import get_pyarrow_compute
 
@@ -219,6 +220,19 @@ class PandasLikeSeries:
         res = ser.isin(other)
         return self._from_native_series(res)
 
+    def arg_true(self) -> PandasLikeSeries:
+        np = get_numpy()
+        ser = self._native_series
+        res = np.flatnonzero(ser)
+        return self._from_native_series(
+            native_series_from_iterable(
+                res,
+                name=self.name,
+                index=range(len(res)),
+                implementation=self._implementation,
+            )
+        )
+
     # Binary comparisons
 
     def filter(self, other: Any) -> PandasLikeSeries:
diff --git a/narwhals/dataframe.py b/narwhals/dataframe.py
index ba3fa3de..4d1b275d 100644
--- a/narwhals/dataframe.py
+++ b/narwhals/dataframe.py
@@ -525,6 +525,9 @@ class DataFrame(BaseFrame[FrameT]):
             msg = f"Expected str or slice, got: {type(item)}"
             raise TypeError(msg)
 
+    def __contains__(self, key: str) -> bool:
+        return key in self.columns
+
     @overload
     def to_dict(self, *, as_series: Literal[True] = ...) -> dict[str, Series]: ...
     @overload
diff --git a/narwhals/expression.py b/narwhals/expression.py
index a7ca6112..5b51f5a5 100644
--- a/narwhals/expression.py
+++ b/narwhals/expression.py
@@ -1051,6 +1051,43 @@ class Expr:
         """
         return self.__class__(lambda plx: self._call(plx).is_null())
 
+    def arg_true(self) -> Self:
+        """
+        Find elements where boolean expression is True.
+
+        Examples:
+            >>> import pandas as pd
+            >>> import polars as pl
+            >>> import narwhals as nw
+            >>> data = {"a": [1, None, None, 2]}
+            >>> df_pd = pd.DataFrame(data)
+            >>> df_pl = pl.DataFrame(data)
+
+            We define a library agnostic function:
+
+            >>> @nw.narwhalify
+            ... def func(df_any):
+            ...     return df_any.select(nw.col("a").is_null().arg_true())
+
+            We can then pass either pandas or Polars to `func`:
+
+            >>> func(df_pd)
+               a
+            0  1
+            1  2
+            >>> func(df_pl)
+            shape: (2, 1)
+            ┌─────┐
+            │ a   │
+            │ --- │
+            │ u32 │
+            ╞═════╡
+            │ 1   │
+            │ 2   │
+            └─────┘
+        """
+        return self.__class__(lambda plx: self._call(plx).arg_true())
+
     def fill_null(self, value: Any) -> Self:
         """
         Fill null values with given value.
diff --git a/narwhals/series.py b/narwhals/series.py
index aa1c2524..befa0f18 100644
--- a/narwhals/series.py
+++ b/narwhals/series.py
@@ -576,6 +576,40 @@ class Series:
             self._compliant_series.is_in(self._extract_native(other))
         )
 
+    def arg_true(self) -> Self:
+        """
+        Find elements where boolean Series is True.
+
+        Examples:
+            >>> import pandas as pd
+            >>> import polars as pl
+            >>> import narwhals as nw
+            >>> data = [1, None, None, 2]
+            >>> s_pd = pd.Series(data, name="a")
+            >>> s_pl = pl.Series("a", data)
+
+            We define a library agnostic function:
+
+            >>> @nw.narwhalify
+            ... def func(s_any):
+            ...     return s_any.is_null().arg_true()
+
+            We can then pass either pandas or Polars to `func`:
+
+            >>> func(s_pd)
+            0    1
+            1    2
+            Name: a, dtype: int64
+            >>> func(s_pl)  # doctest: +NORMALIZE_WHITESPACE
+            shape: (2,)
+            Series: 'a' [u32]
+            [
+               1
+               2
+            ]
+        """
+        return self._from_compliant_series(self._compliant_series.arg_true())
+
     def drop_nulls(self) -> Self:
         """
         Drop all null values.


Test patch:
diff --git a/tests/expr_and_series/arg_true_test.py b/tests/expr_and_series/arg_true_test.py
new file mode 100644
index 00000000..11d43da9
--- /dev/null
+++ b/tests/expr_and_series/arg_true_test.py
@@ -0,0 +1,14 @@
+from typing import Any
+
+import narwhals as nw
+from tests.utils import compare_dicts
+
+
+def test_arg_true(constructor: Any) -> None:
+    df = nw.from_native(constructor({"a": [1, None, None, 3]}), eager_only=True)
+    result = df.select(nw.col("a").is_null().arg_true())
+    expected = {"a": [1, 2]}
+    compare_dicts(result, expected)
+    result = df.select(df["a"].is_null().arg_true())
+    compare_dicts(result, expected)
+    assert "a" in df  # cheeky test to hit `__contains__` method