feat: add nw.factorize

docs/api-reference/series.md

@@ -30,6 +30,7 @@
         - dtype
         - ewm_mean
         - exp
+        - factorize
         - fill_nan
         - fill_null
         - filter

src/narwhals/_arrow/series.py

@@ -1038,6 +1038,19 @@ def hist_from_bin_count(
             .to_frame()
         )
 
+    def factorize(self, *, sort: bool = False) -> tuple[Self, Self]:
+        dtypes = self._version.dtypes
+        uniques = self.unique().drop_nulls()
+        if sort:
+            uniques = uniques.sort(descending=False, nulls_last=True)
+        codes = self.replace_strict(
+            old=uniques.to_list(),
+            new=[*range(len(uniques))],
+            default=-1,
+            return_dtype=dtypes.Int32(),
+        )
+        return codes, uniques
+
     def __iter__(self) -> Iterator[Any]:
         for x in self.native:
             yield maybe_extract_py_scalar(x, return_py_scalar=True)

src/narwhals/_compliant/series.py

@@ -133,6 +133,7 @@ def arg_max(self) -> int: ...
     def arg_min(self) -> int: ...
     def arg_true(self) -> Self: ...
     def count(self) -> int: ...
+    def factorize(self, *, sort: bool = False) -> tuple[Self, Self]: ...
     def filter(self, predicate: Any) -> Self: ...
     def first(self) -> PythonLiteral: ...
     def last(self) -> PythonLiteral: ...

src/narwhals/_pandas_like/series.py

@@ -1163,6 +1163,15 @@ def is_native_dtype_pyarrow(self, native_dtype: Any) -> bool:
         impl = self._implementation
         return get_dtype_backend(native_dtype, implementation=impl) == "pyarrow"
 
+    def factorize(self, *, sort: bool = False) -> tuple[Self, Self]:
+        pdx = self.__native_namespace__()
+        name = self.native.name
+        codes, uniques = self.native.factorize(sort=sort)
+        return (
+            self._with_native(pdx.Series(codes, name=name)),
+            self._with_native(pdx.Series(uniques, name=name)),
+        )
+
     def _apply_pyarrow_compute_func(
         self, native: NativeSeriesT, pc_func: Callable[[ChunkedArrayAny], ChunkedArrayAny]
     ) -> NativeSeriesT:

src/narwhals/_polars/series.py

@@ -670,6 +670,18 @@ def last(self) -> PythonLiteral:
     def any_value(self, *, ignore_nulls: bool) -> PythonLiteral:
         return self.drop_nulls().first() if ignore_nulls else self.first()
 
+    def factorize(self, *, sort: bool = False) -> tuple[Self, Self]:
+        uniques = self.unique().drop_nulls()
+        if sort:
+            uniques = uniques.sort(descending=False, nulls_last=True)
+        codes = self.native.replace_strict(
+            old=uniques.to_list(),
+            new=[*range(len(uniques))],
+            default=-1,
+            return_dtype=pl.Int32(),
+        )
+        return self._with_native(codes.alias(self.name)), uniques
+
     @property
     def dt(self) -> PolarsSeriesDateTimeNamespace:
         return PolarsSeriesDateTimeNamespace(self)

src/narwhals/series.py

@@ -2905,6 +2905,52 @@ def is_close(
         result = result.rename(orig_name) if name_is_none else result
         return cast("Self", result)
 
+    def factorize(self, *, sort: bool = False) -> tuple[Self, Self]:
+        """Encode values as integer codes and unique values.
+
+        Arguments:
+            sort: Whether to sort the unique values before assigning codes.
+
+        Returns:
+            codes: An integer series where each value represents the index
+              of the corresponding value in `uniques`. Null values are encoded
+              as -1.
+            uniques: A series containing the unique non-null values.
+
+        Examples:
+            >>> import polars as pl
+            >>> import narwhals as nw
+            >>> df = pl.DataFrame({"groups": ["a", "b", "a", None]})
+            >>> nw_df = nw.from_native(df)
+            >>> codes, uniques = nw_df["groups"].factorize(sort=True)
+            >>> codes
+            ┌──────────────────────┐
+            |   Narwhals Series    |
+            |----------------------|
+            |shape: (4,)           |
+            |Series: 'groups' [i32]|
+            |[                     |
+            |        0             |
+            |        1             |
+            |        0             |
+            |        -1            |
+            |]                     |
+            └──────────────────────┘
+            >>> uniques
+            ┌──────────────────────┐
+            |   Narwhals Series    |
+            |----------------------|
+            |shape: (2,)           |
+            |Series: 'groups' [str]|
+            |[                     |
+            |        "a"           |
+            |        "b"           |
+            |]                     |
+            └──────────────────────┘
+        """
+        codes, uniques = self._compliant_series.factorize(sort=sort)
+        return self._with_compliant(codes), self._with_compliant(uniques)
+
     @unstable
     def any_value(self, *, ignore_nulls: bool = False) -> PythonLiteral:
         """Get a random value from the column.

tests/series_only/factorize_test.py

@@ -0,0 +1,124 @@
+from __future__ import annotations
+
+from math import isnan
+from typing import Any
+
+import pytest
+
+import narwhals as nw
+from tests.utils import (
+    POLARS_VERSION,
+    ConstructorEager,
+    assert_equal_data,
+    assert_equal_series,
+)
+
+polars_lt_v1 = POLARS_VERSION < (1, 0, 0)
+pl_skip_reason = "replace_strict only available after 1.0"
+
+
+@pytest.mark.parametrize(
+    ("values", "expected_n_unique"),
+    [
+        ([], 0),
+        ([*"abcabc"], 3),
+        ([1, 2, 3, 2], 3),
+        ([1.1, 2.2, 3.3, 2.2], 3),
+        ([*"abc", None], 3),
+        ([*"aaabbbccc", None], 3),
+    ],
+)
+def test_factorize_invariants(
+    values: list[Any], expected_n_unique: int, constructor_eager: ConstructorEager
+) -> None:
+    if "polars" in str(constructor_eager) and polars_lt_v1:
+        pytest.skip(reason=pl_skip_reason)
+
+    has_null = any(x is None for x in values)
+
+    df_native = constructor_eager({"a": values})
+    df = nw.from_native(df_native)
+    codes, uniqs = df["a"].factorize()
+
+    reconstructed_values = {"a": [uniqs[i] if i >= 0 else None for i in codes]}
+    assert_equal_data(df, reconstructed_values)
+    assert uniqs.dtype == df["a"].dtype
+    assert len(uniqs) == expected_n_unique
+
+    # codes should be integer, preserve length, and only contain -1 in the presence of nulls
+    assert codes.dtype.is_integer()
+    assert len(codes) == len(values)
+    assert (codes >= -1).all()
+    assert (codes == -1).any() == has_null
+
+    # Null values should always be dropped out from the unique returned values
+    assert not (uniqs.is_null().any())
+
+
+@pytest.mark.parametrize(
+    ("values", "expected_uniqs", "expected_codes"),
+    [
+        ([], [], []),
+        ([*"abc"], [*"abc"], [0, 1, 2]),
+        ([*"abcabc"], [*"abc"], [0, 1, 2, 0, 1, 2]),
+        ([*"aaabbbccc"], [*"abc"], [0, 0, 0, 1, 1, 1, 2, 2, 2]),
+        ([*"abcabc", None], [*"abc"], [0, 1, 2, 0, 1, 2, -1]),
+    ],
+)
+def test_factorize_sort(
+    values: list[Any],
+    expected_uniqs: list[Any],
+    expected_codes: list[int],
+    constructor_eager: ConstructorEager,
+) -> None:
+    if "polars" in str(constructor_eager) and polars_lt_v1:
+        pytest.skip(reason=pl_skip_reason)
+
+    df_native = constructor_eager({"a": values})
+    df = nw.from_native(df_native)
+    codes, uniqs = df["a"].factorize(sort=True)
+
+    assert_equal_series(uniqs, expected_uniqs, name="a")
+    assert_equal_series(codes, expected_codes, name="a")
+
+
+@pytest.mark.parametrize(
+    "values",
+    [
+        [1.1, 2.2, 1.1, float("nan")],
+        [1.1, 2.2, 1.1, float("nan"), float("nan")],
+        [1.1, 2.2, 1.1, None, float("nan")],
+    ],
+)
+def test_factorize_nan_semantics(
+    values: list[float], constructor_eager: ConstructorEager
+) -> None:
+    if "polars" in str(constructor_eager) and polars_lt_v1:
+        pytest.skip(reason=pl_skip_reason)
+
+    is_pandas_backend = any(x in str(constructor_eager) for x in ("pandas", "modin"))
+
+    df_native = constructor_eager({"a": values})
+    df = nw.from_native(df_native)
+    codes, uniqs = df["a"].factorize()
+
+    reconstructed_values = {"a": [uniqs[i] if i >= 0 else None for i in codes]}
+    assert_equal_data(df, reconstructed_values)
+
+    if is_pandas_backend:
+        # pandas treats NaN as missing, so NaN is not retained as a unique value.
+        assert len(uniqs) == 2
+        assert (codes == -1).any()
+        assert not uniqs.is_null().any()
+    else:
+        # Other backends treat NaN as a value, not as null.
+        assert len(uniqs) == 3
+
+        # The NaN should round-trip through codes -> uniques.
+        nan_index = (
+            i
+            for i, value in enumerate(values)
+            if isinstance(value, float) and isnan(value)
+        )
+        nan_codes = (codes[nan_i] for nan_i in nan_index)
+        assert all(isnan(uniqs[nan_c]) for nan_c in nan_codes)

tests/v1_test.py

@@ -1246,3 +1246,33 @@ def test_schema_from_generator() -> None:
     )
     assert schema == nw_v1.Schema({"a": nw_v1.Int64(), "b": nw_v1.String()})
     assert schema._version is Version.V1
+
+
+@pytest.mark.parametrize(
+    ("values", "expected_uniqs", "expected_codes"),
+    [
+        ([], [], []),
+        ([*"abc"], [*"abc"], [0, 1, 2]),
+        ([*"abcabc"], [*"abc"], [0, 1, 2, 0, 1, 2]),
+        ([*"aaabbbccc"], [*"abc"], [0, 0, 0, 1, 1, 1, 2, 2, 2]),
+        ([*"abcabc", None], [*"abc"], [0, 1, 2, 0, 1, 2, -1]),
+    ],
+)
+def test_factorize(
+    values: list[Any],
+    expected_uniqs: list[Any],
+    expected_codes: list[int],
+    constructor_eager: ConstructorEager,
+) -> None:
+    if "polars" in str(constructor_eager) and (POLARS_VERSION < (1, 0, 0)):
+        pytest.skip(reason="replace_strict only available after 1.0")
+
+    df_native = constructor_eager({"a": values})
+    df = nw_v1.from_native(df_native)
+    codes, uniqs = df["a"].factorize(sort=True)
+
+    assert_equal_series(uniqs, expected_uniqs, name="a")
+    assert_equal_series(codes, expected_codes, name="a")
+
+    assert codes._version is Version.V1
+    assert uniqs._version is Version.V1

tests/v2_test.py

@@ -586,3 +586,33 @@ def test_schema_from_generator() -> None:
     )
     assert schema == nw_v2.Schema({"a": nw_v2.Int64(), "b": nw_v2.String()})
     assert schema._version is Version.V2
+
+
+@pytest.mark.parametrize(
+    ("values", "expected_uniqs", "expected_codes"),
+    [
+        ([], [], []),
+        ([*"abc"], [*"abc"], [0, 1, 2]),
+        ([*"abcabc"], [*"abc"], [0, 1, 2, 0, 1, 2]),
+        ([*"aaabbbccc"], [*"abc"], [0, 0, 0, 1, 1, 1, 2, 2, 2]),
+        ([*"abcabc", None], [*"abc"], [0, 1, 2, 0, 1, 2, -1]),
+    ],
+)
+def test_factorize(
+    values: list[Any],
+    expected_uniqs: list[Any],
+    expected_codes: list[int],
+    constructor_eager: ConstructorEager,
+) -> None:
+    if "polars" in str(constructor_eager) and (POLARS_VERSION < (1, 0, 0)):
+        pytest.skip(reason="replace_strict only available after 1.0")
+
+    df_native = constructor_eager({"a": values})
+    df = nw_v2.from_native(df_native)
+    codes, uniqs = df["a"].factorize(sort=True)
+
+    assert_equal_series(uniqs, expected_uniqs, name="a")
+    assert_equal_series(codes, expected_codes, name="a")
+
+    assert codes._version is Version.V2
+    assert uniqs._version is Version.V2

utils/check_api_reference.py

@@ -80,6 +80,7 @@ def read_documented_members(source: str | Path) -> list[str]:
     "arg_min",
     "arg_true",
     "dtype",
+    "factorize",
     "from_iterable",
     "from_numpy",
     "gather_every",