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,57 @@ def hist_from_bin_count(
             .to_frame()
         )
 
+    def factorize(
+        self, *, null_as_value: bool = False, sort: bool = False
+    ) -> tuple[Self, Self]:
+        if len(self.native) == 0:
+            codes = pa.chunked_array([[]], type=pa.int32())
+            uniques = pa.chunked_array([[]], type=self.native.type)
+            return (self._with_native(codes), self._with_native(uniques))
+
+        # https://github.com/apache/arrow/issues/33297; input pa.NullArray's don't dictionary_encode properly
+        if pa.types.is_null(self.native.type):
+            if null_as_value:
+                codes, uniques = (
+                    pa.repeat(0, len(self.native)),
+                    pa.nulls(1, type=self.native.type),
+                )
+            else:
+                codes, uniques = (
+                    pa.repeat(-1, len(self.native)),
+                    pa.nulls(0, type=self.native.type),
+                )
+
+            return (self._with_native(codes.cast(pa.int32())), self._with_native(uniques))
+
+        native = self.native
+        if pa.types.is_dictionary(native.type):
+            # re-encode if already dict encoded; can't be certain how the original dict encoding was done
+            native = native.cast(native.type.value_type)
+
+        null_encoding: Literal["encode", "mask"] = "encode" if null_as_value else "mask"
+        encoded = pc.dictionary_encode(
+            native, null_encoding=null_encoding
+        ).unify_dictionaries()
+        uniques = encoded.chunk(0).dictionary  # type: ignore[attr-defined]
+        codes = pa.chunked_array([c.indices for c in encoded.chunks], type=pa.int32())  # type: ignore[attr-defined]
+        codes = cast("pa.ChunkedArray[pa.Int32Scalar]", codes)
+
+        if not sort:
+            return (
+                self._with_native(pc.fill_null(codes, pa.scalar(-1))),
+                self._with_native(uniques),
+            )
+
+        sorted_uniques = pc.take(uniques, pc.sort_indices(uniques))
+        new_mapping = pc.index_in(uniques, value_set=sorted_uniques)
+        new_codes = pc.take(new_mapping, codes)
+
+        return (
+            self._with_native(pc.fill_null(new_codes, pa.scalar(-1))),
+            self._with_native(sorted_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,9 @@ def arg_max(self) -> int: ...
     def arg_min(self) -> int: ...
     def arg_true(self) -> Self: ...
     def count(self) -> int: ...
+    def factorize(
+        self, *, null_as_value: bool, 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,34 @@ 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, *, null_as_value: bool = False, sort: bool = False
+    ) -> tuple[Self, Self]:
+        pdx = self.__native_namespace__()
+
+        # https://github.com/apache/arrow/issues/33297; input pa.NullArray's don't dictionary_encode properly
+        if self.native.dtype == "null[pyarrow]":
+            if null_as_value:
+                codes, uniques = (
+                    pdx.Series(0, index=self.native.index),
+                    pdx.Series([None], dtype=self.native.dtype),
+                )
+            else:
+                codes, uniques = (
+                    pdx.Series(-1, index=self.native.index),
+                    pdx.Series([], dtype=self.native.dtype),
+                )
+
+            return (self._with_native(codes), self._with_native(uniques))
+
+        codes, uniques = self.native.factorize(
+            sort=sort, use_na_sentinel=not null_as_value
+        )
+        return (
+            self._with_native(pdx.Series(codes)),
+            self._with_native(pdx.Series(uniques)),
+        )
+
     def _apply_pyarrow_compute_func(
         self, native: NativeSeriesT, pc_func: Callable[[ChunkedArrayAny], ChunkedArrayAny]
     ) -> NativeSeriesT:

src/narwhals/_polars/series.py

@@ -670,6 +670,27 @@ 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, *, null_as_value: bool = False, sort: bool = False
+    ) -> tuple[Self, Self]:
+        uniques = self.unique() if null_as_value else self.unique().drop_nulls()
+        if sort:
+            uniques = uniques.sort(descending=False, nulls_last=True)
+
+        if null_as_value:
+            codes = self.native.replace_strict(
+                old=uniques.native, new=range(len(uniques)), return_dtype=pl.Int32()
+            )
+        else:
+            codes = self.native.replace_strict(
+                old=uniques.native,
+                new=range(len(uniques)),
+                default=-1,
+                return_dtype=pl.Int32(),
+            )
+
+        return self._with_native(codes.cast(pl.Int32())), uniques
+
     @property
     def dt(self) -> PolarsSeriesDateTimeNamespace:
         return PolarsSeriesDateTimeNamespace(self)

src/narwhals/series.py

@@ -2,6 +2,7 @@
 
 import math
 from collections.abc import Callable, Iterable, Iterator, Mapping, Sequence
+from dataclasses import dataclass
 from functools import partial
 from typing import TYPE_CHECKING, Any, ClassVar, Generic, Literal, cast, overload
 
@@ -2905,6 +2906,68 @@ def is_close(
         result = result.rename(orig_name) if name_is_none else result
         return cast("Self", result)
 
+    def factorize(
+        self, *, null_as_value: bool = False, sort: bool = False
+    ) -> Encoded[IntoSeriesT]:
+        """Encode values as integer codes and unique values.
+
+        The integer codes are index locations that map the unique values back to their
+        positions within the original array.
+
+        Arguments:
+            null_as_value: Whether to treat null as a regular value. When False,
+                nulls are removed from the returned unique values and the code -1 is
+                used to indicate the location of null values in the original array.
+                When True, nulls are preserved in the returned unique values and a
+                positive integer is used to indicate their location in the original
+                array.
+            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(
+            null_as_value=null_as_value, sort=sort
+        )
+        return Encoded(
+            self._with_compliant(codes).alias("codes"),
+            self._with_compliant(uniques).alias("uniques"),
+        )
+
     @unstable
     def any_value(self, *, ignore_nulls: bool = False) -> PythonLiteral:
         """Get a random value from the column.
@@ -2946,3 +3009,26 @@ def list(self) -> SeriesListNamespace[Self]:
     @property
     def struct(self) -> SeriesStructNamespace[Self]:
         return SeriesStructNamespace(self)
+
+
+@dataclass(frozen=True)
+class Encoded(Generic[IntoSeriesT]):
+    """Result of `factorize`. Unpacks as `(codes, uniques)` like pandas."""
+
+    codes: Series[IntoSeriesT]
+    uniques: Series[IntoSeriesT]
+
+    def __iter__(self) -> Iterator[Series[IntoSeriesT]]:
+        yield self.codes
+        yield self.uniques
+
+    @property
+    def mapping(self) -> Mapping[Any, int]:
+        """Forward map as a joinable ``(value, code)`` frame; works for any dtype."""
+        name = self.uniques.name
+        return dict(
+            self.uniques.to_frame()
+            .with_row_index("code")
+            .select(name, "code")
+            .iter_rows()
+        )