Repository: narwhals-dev/narwhals
Pull request #948 (2024-09-11 08:28:57)

feat: add Series.Scatter
<!--
# 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/series.md b/docs/api-reference/series.md
index b637f550..c016b566 100644
--- a/docs/api-reference/series.md
+++ b/docs/api-reference/series.md
@@ -12,6 +12,7 @@
         - any
         - arg_true
         - cast
+        - clip
         - count
         - cum_sum
         - diff
@@ -22,7 +23,6 @@
         - gather_every
         - head
         - is_between
-        - clip
         - is_duplicated
         - is_empty
         - is_first_distinct
@@ -38,12 +38,13 @@
         - min
         - mode
         - name
-        - null_count
         - n_unique
+        - null_count
         - pipe
         - quantile
         - round
         - sample
+        - scatter
         - shape
         - shift
         - sort
diff --git a/narwhals/_arrow/series.py b/narwhals/_arrow/series.py
index 50927426..8765605f 100644
--- a/narwhals/_arrow/series.py
+++ b/narwhals/_arrow/series.py
@@ -323,6 +323,23 @@ class ArrowSeries:
             return self._native_series[idx]
         return self._from_native_series(self._native_series[idx])
 
+    def scatter(self, indices: int | Sequence[int], values: Any) -> Self:
+        import numpy as np  # ignore-banned-import
+        import pyarrow as pa  # ignore-banned-import
+        import pyarrow.compute as pc  # ignore-banned-import
+
+        ca = self._native_series
+        mask = np.zeros(len(ca), dtype=bool)
+        mask[indices] = True
+        if isinstance(values, self.__class__):
+            values = validate_column_comparand(values)
+        if isinstance(values, pa.ChunkedArray):
+            values = values.combine_chunks()
+        if not isinstance(values, pa.Array):
+            values = pa.array(values)
+        result = pc.replace_with_mask(ca, mask, values.take(indices))
+        return self._from_native_series(result)
+
     def to_list(self) -> Any:
         return self._native_series.to_pylist()
 
diff --git a/narwhals/_pandas_like/series.py b/narwhals/_pandas_like/series.py
index a13f252e..8288be26 100644
--- a/narwhals/_pandas_like/series.py
+++ b/narwhals/_pandas_like/series.py
@@ -10,6 +10,7 @@ from typing import overload
 from narwhals._pandas_like.utils import int_dtype_mapper
 from narwhals._pandas_like.utils import narwhals_to_native_dtype
 from narwhals._pandas_like.utils import native_series_from_iterable
+from narwhals._pandas_like.utils import set_axis
 from narwhals._pandas_like.utils import to_datetime
 from narwhals._pandas_like.utils import translate_dtype
 from narwhals._pandas_like.utils import validate_column_comparand
@@ -167,6 +168,22 @@ class PandasLikeSeries:
     def dtype(self) -> DType:
         return translate_dtype(self._native_series)
 
+    def scatter(self, indices: int | Sequence[int], values: Any) -> Self:
+        if isinstance(values, self.__class__):
+            # .copy() is necessary in some pre-2.2 versions of pandas to avoid
+            # `values` also getting modified (!)
+            values = validate_column_comparand(self._native_series.index, values).copy()
+            values = set_axis(
+                values,
+                self._native_series.index[indices],
+                implementation=self._implementation,
+                backend_version=self._backend_version,
+            )
+        s = self._native_series
+        s.iloc[indices] = values
+        s.name = self.name
+        return self._from_native_series(s)
+
     def cast(
         self,
         dtype: Any,
diff --git a/narwhals/series.py b/narwhals/series.py
index 147838e4..9fcb07a2 100644
--- a/narwhals/series.py
+++ b/narwhals/series.py
@@ -85,6 +85,58 @@ class Series:
         ca = pa.chunked_array([self.to_arrow()])
         return ca.__arrow_c_stream__(requested_schema=requested_schema)
 
+    def scatter(self, indices: int | Sequence[int], values: Any) -> Self:
+        """
+        Set value(s) at given position(s).
+
+        Arguments:
+           indices: Position(s) to set items at.
+           values: Values to set.
+
+        Warning:
+            For some libraries (pandas, Polars), this method operates in-place,
+            whereas for others (PyArrow) it doesn't!
+            We recommend being careful with it, and not relying on the
+            in-placeness. For example, a valid use case is when updating
+            a column in an eager dataframe, see the example below.
+
+        Examples:
+            >>> import pandas as pd
+            >>> import polars as pl
+            >>> import narwhals as nw
+            >>> data = {"a": [1, 2, 3], "b": [4, 5, 6]}
+            >>> df_pd = pd.DataFrame(data)
+            >>> df_pl = pl.DataFrame(data)
+
+            We define a library agnostic function:
+
+            >>> @nw.narwhalify
+            ... def func(df):
+            ...     return df.with_columns(df["a"].scatter([0, 1], [999, 888]))
+
+            We can then pass either pandas or Polars to `func`:
+
+            >>> func(df_pd)
+                 a  b
+            0  999  4
+            1  888  5
+            2    3  6
+            >>> func(df_pl)
+            shape: (3, 2)
+            ┌─────┬─────┐
+            │ a   ┆ b   │
+            │ --- ┆ --- │
+            │ i64 ┆ i64 │
+            ╞═════╪═════╡
+            │ 999 ┆ 4   │
+            │ 888 ┆ 5   │
+            │ 3   ┆ 6   │
+            └─────┴─────┘
+        """
+        return self._from_compliant_series(
+            self._compliant_series.scatter(indices, self._extract_native(values))
+        )
+
     @property
     def shape(self) -> tuple[int]:
         """
diff --git a/utils/check_api_reference.py b/utils/check_api_reference.py
index f6e5303c..ec599def 100644
--- a/utils/check_api_reference.py
+++ b/utils/check_api_reference.py
@@ -162,6 +162,7 @@ if (
             "value_counts",
             "zip_with",
             "item",
+            "scatter",
         }
     )
 ):


Test patch:
diff --git a/tests/series_only/scatter_test.py b/tests/series_only/scatter_test.py
new file mode 100644
index 00000000..2edab2b8
--- /dev/null
+++ b/tests/series_only/scatter_test.py
@@ -0,0 +1,26 @@
+from __future__ import annotations
+
+from typing import Any
+
+import pytest
+
+import narwhals as nw
+from tests.utils import compare_dicts
+
+
+def test_scatter(constructor_eager: Any, request: pytest.FixtureRequest) -> None:
+    if "modin" in str(constructor_eager):
+        # https://github.com/modin-project/modin/issues/7392
+        request.applymarker(pytest.mark.xfail)
+    df = nw.from_native(
+        constructor_eager({"a": [1, 2, 3], "b": [142, 124, 132]}), eager_only=True
+    )
+    result = df.with_columns(
+        df["a"].scatter([0, 1], [999, 888]),
+        df["b"].scatter([0, 2, 1], df["b"]),
+    )
+    expected = {
+        "a": [999, 888, 3],
+        "b": [142, 132, 124],
+    }
+    compare_dicts(result, expected)