feat: add nw.factorize

docs/api-reference/narwhals.md

@@ -16,6 +16,7 @@ Here are the top-level functions available in Narwhals.
         - corr
         - cov
         - exclude
+        - factorize
         - format
         - from_arrow
         - from_dict

src/narwhals/__init__.py

@@ -57,6 +57,7 @@
     corr,
     cov,
     exclude,
+    factorize,
     format,
     from_arrow,
     from_dict,
@@ -145,6 +146,7 @@
     "dtypes",
     "exceptions",
     "exclude",
+    "factorize",
     "format",
     "from_arrow",
     "from_dict",

src/narwhals/functions.py

@@ -30,6 +30,7 @@
     is_numpy_array_2d,
     is_pyarrow_table,
 )
+from narwhals.dtypes import Int32
 from narwhals.exceptions import InvalidOperationError
 from narwhals.expr import Expr
 from narwhals.schema import Schema
@@ -54,6 +55,7 @@
         IntoDType,
         IntoExpr,
         IntoSchema,
+        IntoSeriesT,
         NonNestedLiteral,
         PythonLiteral,
         _2DArray,
@@ -2021,3 +2023,55 @@ def list_(*exprs: IntoExpr | Sequence[IntoExpr]) -> Expr:
     return Expr(
         ExprNode(ExprKind.ELEMENTWISE, "list", exprs=flat_exprs, allow_multi_output=True)
     )
+
+
+def factorize(
+    values: Series[IntoSeriesT], *, sort: bool = False
+) -> tuple[Series[IntoSeriesT], Series[IntoSeriesT]]:
+    """Encode values as integer codes and unique values.
+
+    Arguments:
+        values: A series to factorize.
+        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.factorize(series["groups"], sort=True)
+        >>> codes
+        ┌─────┐
+        |  a  |
+        |-----|
+        | i32 |
+        |-----|
+        |   0 |
+        |   1 |
+        |   0 |
+        |  -1 |
+        └─────┘
+        >>> uniques
+        ┌─────┐
+        |  a  |
+        |-----|
+        | str |
+        |-----|
+        | b   |
+        | a   |
+        └─────┘
+    """
+    uniques = values.unique().drop_nulls()
+    if sort:
+        uniques = uniques.sort()
+
+    codes = values.replace_strict(
+        uniques.to_list(), [*range(len(uniques))], default=-1, return_dtype=Int32()
+    )
+    return codes, uniques

src/narwhals/stable/v2/__init__.py

@@ -1171,6 +1171,25 @@ def struct(*exprs: IntoExpr | Sequence[IntoExpr], **named_exprs: IntoExpr) -> Ex
     return _stableify(nw_f.struct(*exprs, **named_exprs))
 
 
+def factorize(
+    values: Series[IntoSeriesT], *, sort: bool = False
+) -> tuple[Series[IntoSeriesT], Series[IntoSeriesT]]:
+    """Encode values as integer codes and unique values.
+
+    Arguments:
+        values: A series to factorize.
+        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.
+    """
+    codes, uniques = nw_f.factorize(values, sort=sort)
+    return _stableify(codes), _stableify(uniques)
+
+
 __all__ = [
     "Array",
     "Binary",
@@ -1221,6 +1240,7 @@ def struct(*exprs: IntoExpr | Sequence[IntoExpr], **named_exprs: IntoExpr) -> Ex
     "dtypes",
     "exceptions",
     "exclude",
+    "factorize",
     "format",
     "from_arrow",
     "from_dict",

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 = nw.factorize(df["a"])
+
+    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 = nw.factorize(df["a"], 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 = nw.factorize(df["a"])
+
+    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/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 = nw_v2.factorize(df["a"], 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