From d38d4cbd3b5c7e3083c9c82a37b4dbd38b1ad51c Mon Sep 17 00:00:00 2001 From: Cameron Riddell Date: Mon, 20 Jul 2026 14:49:05 -0700 Subject: [PATCH 01/12] feat: add nw.factorize factorize produces unique values and value mappings back to the locations of those values. --- docs/api-reference/narwhals.md | 1 + src/narwhals/__init__.py | 2 + src/narwhals/functions.py | 54 ++++++++++++++ src/narwhals/stable/v2/__init__.py | 20 ++++++ tests/series_only/factorize_test.py | 107 ++++++++++++++++++++++++++++ 5 files changed, 184 insertions(+) create mode 100644 tests/series_only/factorize_test.py diff --git a/docs/api-reference/narwhals.md b/docs/api-reference/narwhals.md index 15314410c7..a089e0c85c 100644 --- a/docs/api-reference/narwhals.md +++ b/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 diff --git a/src/narwhals/__init__.py b/src/narwhals/__init__.py index afd37b8870..5539257c0e 100644 --- a/src/narwhals/__init__.py +++ b/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", diff --git a/src/narwhals/functions.py b/src/narwhals/functions.py index a576c6b33d..a124b09448 100644 --- a/src/narwhals/functions.py +++ b/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 diff --git a/src/narwhals/stable/v2/__init__.py b/src/narwhals/stable/v2/__init__.py index 04ec4d6c36..6533722cdf 100644 --- a/src/narwhals/stable/v2/__init__.py +++ b/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", diff --git a/tests/series_only/factorize_test.py b/tests/series_only/factorize_test.py new file mode 100644 index 0000000000..f5c38b5ab3 --- /dev/null +++ b/tests/series_only/factorize_test.py @@ -0,0 +1,107 @@ +from __future__ import annotations + +from math import isnan +from typing import Any + +import pytest + +import narwhals as nw +from tests.utils import ConstructorEager, assert_equal_data, assert_equal_series + + +@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: + 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: + 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: + is_pandas_backend = "pandas" in str(constructor_eager) + + 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) From dfcc905281af522d7e158b3ca117b62f3541c0e5 Mon Sep 17 00:00:00 2001 From: Cameron Riddell Date: Mon, 20 Jul 2026 15:56:44 -0700 Subject: [PATCH 02/12] factorize test skip Polars < 1.0 due to lack of replace_strict --- tests/series_only/factorize_test.py | 21 +++++++++++++++++++-- 1 file changed, 19 insertions(+), 2 deletions(-) diff --git a/tests/series_only/factorize_test.py b/tests/series_only/factorize_test.py index f5c38b5ab3..c90ef74d53 100644 --- a/tests/series_only/factorize_test.py +++ b/tests/series_only/factorize_test.py @@ -6,7 +6,15 @@ import pytest import narwhals as nw -from tests.utils import ConstructorEager, assert_equal_data, assert_equal_series +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( @@ -23,6 +31,9 @@ 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}) @@ -60,6 +71,9 @@ def test_factorize_sort( 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) @@ -79,7 +93,10 @@ def test_factorize_sort( def test_factorize_nan_semantics( values: list[float], constructor_eager: ConstructorEager ) -> None: - is_pandas_backend = "pandas" in str(constructor_eager) + 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) From 9a5dacfea57f5c46c97cec30fac94af3d5ae0f28 Mon Sep 17 00:00:00 2001 From: Cameron Riddell Date: Mon, 20 Jul 2026 15:57:04 -0700 Subject: [PATCH 03/12] add factorize test for nw V2 --- tests/v2_test.py | 30 ++++++++++++++++++++++++++++++ 1 file changed, 30 insertions(+) diff --git a/tests/v2_test.py b/tests/v2_test.py index bd1a3a9036..d39c13454e 100644 --- a/tests/v2_test.py +++ b/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 From 7d8459a44bc9081ceae460a85bb3014602b25122 Mon Sep 17 00:00:00 2001 From: Cameron Riddell Date: Wed, 22 Jul 2026 07:36:10 -0700 Subject: [PATCH 04/12] fix: nw.factorize doctest & docstring returns --- src/narwhals/functions.py | 46 +++++++++++++++++++++------------------ 1 file changed, 25 insertions(+), 21 deletions(-) diff --git a/src/narwhals/functions.py b/src/narwhals/functions.py index a124b09448..dfc9df22c9 100644 --- a/src/narwhals/functions.py +++ b/src/narwhals/functions.py @@ -2035,37 +2035,41 @@ def factorize( sort: Whether to sort the unique values before assigning codes. Returns: - - codes: An integer series where each value represents the index + 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. + 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, uniques = nw.factorize(nw_df["groups"], sort=True) >>> codes - ┌─────┐ - | a | - |-----| - | i32 | - |-----| - | 0 | - | 1 | - | 0 | - | -1 | - └─────┘ + ┌──────────────────────┐ + | Narwhals Series | + |----------------------| + |shape: (4,) | + |Series: 'groups' [i32]| + |[ | + | 0 | + | 1 | + | 0 | + | -1 | + |] | + └──────────────────────┘ >>> uniques - ┌─────┐ - | a | - |-----| - | str | - |-----| - | b | - | a | - └─────┘ + ┌──────────────────────┐ + | Narwhals Series | + |----------------------| + |shape: (2,) | + |Series: 'groups' [str]| + |[ | + | "a" | + | "b" | + |] | + └──────────────────────┘ """ uniques = values.unique().drop_nulls() if sort: From 8b798aafa964ccfbbbad3a08219b3f7ec8030045 Mon Sep 17 00:00:00 2001 From: Cameron Riddell Date: Thu, 23 Jul 2026 12:26:20 -0700 Subject: [PATCH 05/12] ref: move factorize to Series method & make specific for each backend - `nw.factorize` -> `nw.Series.factorize` - factorize code now lives within each backend, allowing for fastpaths like pandas.factorize instead of the generic unique -> replace_strict logic - added v1 test --- docs/api-reference/narwhals.md | 1 - docs/api-reference/series.md | 1 + src/narwhals/__init__.py | 2 - src/narwhals/_arrow/series.py | 13 +++++++ src/narwhals/_compliant/series.py | 1 + src/narwhals/_pandas_like/series.py | 9 +++++ src/narwhals/_polars/series.py | 12 ++++++ src/narwhals/functions.py | 58 ----------------------------- src/narwhals/series.py | 46 +++++++++++++++++++++++ src/narwhals/stable/v2/__init__.py | 20 ---------- tests/series_only/factorize_test.py | 6 +-- tests/v1_test.py | 30 +++++++++++++++ tests/v2_test.py | 2 +- utils/check_api_reference.py | 1 + 14 files changed, 117 insertions(+), 85 deletions(-) diff --git a/docs/api-reference/narwhals.md b/docs/api-reference/narwhals.md index a089e0c85c..15314410c7 100644 --- a/docs/api-reference/narwhals.md +++ b/docs/api-reference/narwhals.md @@ -16,7 +16,6 @@ Here are the top-level functions available in Narwhals. - corr - cov - exclude - - factorize - format - from_arrow - from_dict diff --git a/docs/api-reference/series.md b/docs/api-reference/series.md index 91c4874fe9..073ec116ae 100644 --- a/docs/api-reference/series.md +++ b/docs/api-reference/series.md @@ -30,6 +30,7 @@ - dtype - ewm_mean - exp + - factorize - fill_nan - fill_null - filter diff --git a/src/narwhals/__init__.py b/src/narwhals/__init__.py index 5539257c0e..afd37b8870 100644 --- a/src/narwhals/__init__.py +++ b/src/narwhals/__init__.py @@ -57,7 +57,6 @@ corr, cov, exclude, - factorize, format, from_arrow, from_dict, @@ -146,7 +145,6 @@ "dtypes", "exceptions", "exclude", - "factorize", "format", "from_arrow", "from_dict", diff --git a/src/narwhals/_arrow/series.py b/src/narwhals/_arrow/series.py index 74fbac0df1..e732473828 100644 --- a/src/narwhals/_arrow/series.py +++ b/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) diff --git a/src/narwhals/_compliant/series.py b/src/narwhals/_compliant/series.py index ea83e6f86b..22e8f742cc 100644 --- a/src/narwhals/_compliant/series.py +++ b/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: ... diff --git a/src/narwhals/_pandas_like/series.py b/src/narwhals/_pandas_like/series.py index 48ee244c22..84e93fb0e4 100644 --- a/src/narwhals/_pandas_like/series.py +++ b/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: diff --git a/src/narwhals/_polars/series.py b/src/narwhals/_polars/series.py index b82a911f61..fe759f1142 100644 --- a/src/narwhals/_polars/series.py +++ b/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) diff --git a/src/narwhals/functions.py b/src/narwhals/functions.py index dfc9df22c9..a576c6b33d 100644 --- a/src/narwhals/functions.py +++ b/src/narwhals/functions.py @@ -30,7 +30,6 @@ 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 @@ -55,7 +54,6 @@ IntoDType, IntoExpr, IntoSchema, - IntoSeriesT, NonNestedLiteral, PythonLiteral, _2DArray, @@ -2023,59 +2021,3 @@ 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(nw_df["groups"], sort=True) - >>> codes - ┌──────────────────────┐ - | Narwhals Series | - |----------------------| - |shape: (4,) | - |Series: 'groups' [i32]| - |[ | - | 0 | - | 1 | - | 0 | - | -1 | - |] | - └──────────────────────┘ - >>> uniques - ┌──────────────────────┐ - | Narwhals Series | - |----------------------| - |shape: (2,) | - |Series: 'groups' [str]| - |[ | - | "a" | - | "b" | - |] | - └──────────────────────┘ - """ - 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 diff --git a/src/narwhals/series.py b/src/narwhals/series.py index d0688cba41..1f96d0a431 100644 --- a/src/narwhals/series.py +++ b/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. diff --git a/src/narwhals/stable/v2/__init__.py b/src/narwhals/stable/v2/__init__.py index 6533722cdf..04ec4d6c36 100644 --- a/src/narwhals/stable/v2/__init__.py +++ b/src/narwhals/stable/v2/__init__.py @@ -1171,25 +1171,6 @@ 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", @@ -1240,7 +1221,6 @@ def factorize( "dtypes", "exceptions", "exclude", - "factorize", "format", "from_arrow", "from_dict", diff --git a/tests/series_only/factorize_test.py b/tests/series_only/factorize_test.py index c90ef74d53..44121769a2 100644 --- a/tests/series_only/factorize_test.py +++ b/tests/series_only/factorize_test.py @@ -38,7 +38,7 @@ def test_factorize_invariants( df_native = constructor_eager({"a": values}) df = nw.from_native(df_native) - codes, uniqs = nw.factorize(df["a"]) + 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) @@ -76,7 +76,7 @@ def test_factorize_sort( df_native = constructor_eager({"a": values}) df = nw.from_native(df_native) - codes, uniqs = nw.factorize(df["a"], sort=True) + codes, uniqs = df["a"].factorize(sort=True) assert_equal_series(uniqs, expected_uniqs, name="a") assert_equal_series(codes, expected_codes, name="a") @@ -100,7 +100,7 @@ def test_factorize_nan_semantics( df_native = constructor_eager({"a": values}) df = nw.from_native(df_native) - codes, uniqs = nw.factorize(df["a"]) + 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) diff --git a/tests/v1_test.py b/tests/v1_test.py index ac048112e5..268decb53c 100644 --- a/tests/v1_test.py +++ b/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 diff --git a/tests/v2_test.py b/tests/v2_test.py index d39c13454e..cc3edacc61 100644 --- a/tests/v2_test.py +++ b/tests/v2_test.py @@ -609,7 +609,7 @@ def test_factorize( df_native = constructor_eager({"a": values}) df = nw_v2.from_native(df_native) - codes, uniqs = nw_v2.factorize(df["a"], sort=True) + codes, uniqs = df["a"].factorize(sort=True) assert_equal_series(uniqs, expected_uniqs, name="a") assert_equal_series(codes, expected_codes, name="a") diff --git a/utils/check_api_reference.py b/utils/check_api_reference.py index db0af805d3..83d1b2cd79 100644 --- a/utils/check_api_reference.py +++ b/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", From e3b4cb4cfecf3df747eed7ac436f63ba896855fb Mon Sep 17 00:00:00 2001 From: Cameron Riddell Date: Thu, 23 Jul 2026 14:39:54 -0700 Subject: [PATCH 06/12] perf: pyarrow factorize fastpath --- src/narwhals/_arrow/series.py | 32 ++++++++++++++++++++++---------- 1 file changed, 22 insertions(+), 10 deletions(-) diff --git a/src/narwhals/_arrow/series.py b/src/narwhals/_arrow/series.py index e732473828..48b92993cd 100644 --- a/src/narwhals/_arrow/series.py +++ b/src/narwhals/_arrow/series.py @@ -1039,17 +1039,29 @@ def hist_from_bin_count( ) 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(), + 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)) + + native_dict_encoded = pc.dictionary_encode(self.native).unify_dictionaries() + uniques = native_dict_encoded.chunks[0].dictionary + codes = pa.chunked_array( + [chunk.indices for chunk in native_dict_encoded.iterchunks()] + ) + if not sort: + return ( + self._with_native(pc.fill_null(codes, -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, -1)), + self._with_native(sorted_uniques), ) - return codes, uniques def __iter__(self) -> Iterator[Any]: for x in self.native: From e8d73e9dd3055c7f828c4f151d03adb8bb749aa6 Mon Sep 17 00:00:00 2001 From: Cameron Riddell Date: Fri, 24 Jul 2026 06:04:58 -0700 Subject: [PATCH 07/12] perf: polars factorize pass native series instead of .to_list --- src/narwhals/_polars/series.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/narwhals/_polars/series.py b/src/narwhals/_polars/series.py index fe759f1142..0feeee93c7 100644 --- a/src/narwhals/_polars/series.py +++ b/src/narwhals/_polars/series.py @@ -675,8 +675,8 @@ def factorize(self, *, sort: bool = False) -> tuple[Self, Self]: if sort: uniques = uniques.sort(descending=False, nulls_last=True) codes = self.native.replace_strict( - old=uniques.to_list(), - new=[*range(len(uniques))], + old=uniques.native, + new=range(len(uniques)), default=-1, return_dtype=pl.Int32(), ) From 2abe03db3d02c94f535a978c942a4c43ee34fc1f Mon Sep 17 00:00:00 2001 From: Cameron Riddell Date: Fri, 24 Jul 2026 13:02:48 -0700 Subject: [PATCH 08/12] feat: add null_as_value to factorize --- src/narwhals/_arrow/series.py | 42 +++++++-- src/narwhals/_compliant/series.py | 4 +- src/narwhals/_pandas_like/series.py | 27 +++++- src/narwhals/_polars/series.py | 27 ++++-- src/narwhals/series.py | 21 ++++- tests/series_only/factorize_test.py | 132 ++++++++++++++++++---------- 6 files changed, 182 insertions(+), 71 deletions(-) diff --git a/src/narwhals/_arrow/series.py b/src/narwhals/_arrow/series.py index 48b92993cd..4e0883d394 100644 --- a/src/narwhals/_arrow/series.py +++ b/src/narwhals/_arrow/series.py @@ -1038,28 +1038,54 @@ def hist_from_bin_count( .to_frame() ) - def factorize(self, *, sort: bool = False) -> tuple[Self, Self]: + 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)) - native_dict_encoded = pc.dictionary_encode(self.native).unify_dictionaries() - uniques = native_dict_encoded.chunks[0].dictionary - codes = pa.chunked_array( - [chunk.indices for chunk in native_dict_encoded.iterchunks()] - ) + # 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, -1)), + 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, -1)), + self._with_native(pc.fill_null(new_codes, pa.scalar(-1))), self._with_native(sorted_uniques), ) diff --git a/src/narwhals/_compliant/series.py b/src/narwhals/_compliant/series.py index 22e8f742cc..dabe1c6deb 100644 --- a/src/narwhals/_compliant/series.py +++ b/src/narwhals/_compliant/series.py @@ -133,7 +133,9 @@ 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 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: ... diff --git a/src/narwhals/_pandas_like/series.py b/src/narwhals/_pandas_like/series.py index 84e93fb0e4..2c79679e07 100644 --- a/src/narwhals/_pandas_like/series.py +++ b/src/narwhals/_pandas_like/series.py @@ -1163,10 +1163,33 @@ 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]: + def factorize( + self, *, null_as_value: bool = False, sort: bool = False + ) -> tuple[Self, Self]: pdx = self.__native_namespace__() name = self.native.name - codes, uniques = self.native.factorize(sort=sort) + + # 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.rename(name)), + self._with_native(uniques.rename(name)), + ) + + codes, uniques = self.native.factorize( + sort=sort, use_na_sentinel=not null_as_value + ) return ( self._with_native(pdx.Series(codes, name=name)), self._with_native(pdx.Series(uniques, name=name)), diff --git a/src/narwhals/_polars/series.py b/src/narwhals/_polars/series.py index 0feeee93c7..eff6a77204 100644 --- a/src/narwhals/_polars/series.py +++ b/src/narwhals/_polars/series.py @@ -670,17 +670,26 @@ 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() + 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) - codes = self.native.replace_strict( - old=uniques.native, - new=range(len(uniques)), - default=-1, - return_dtype=pl.Int32(), - ) - return self._with_native(codes.alias(self.name)), uniques + + 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: diff --git a/src/narwhals/series.py b/src/narwhals/series.py index 1f96d0a431..378ce4b8cd 100644 --- a/src/narwhals/series.py +++ b/src/narwhals/series.py @@ -2905,16 +2905,27 @@ 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]: + def factorize( + self, *, null_as_value: bool = False, sort: bool = False + ) -> tuple[Self, Self]: """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. + of the corresponding value in `uniques`. Null values are encoded + as -1. uniques: A series containing the unique non-null values. Examples: @@ -2948,7 +2959,9 @@ def factorize(self, *, sort: bool = False) -> tuple[Self, Self]: |] | └──────────────────────┘ """ - codes, uniques = self._compliant_series.factorize(sort=sort) + codes, uniques = self._compliant_series.factorize( + null_as_value=null_as_value, sort=sort + ) return self._with_compliant(codes), self._with_compliant(uniques) @unstable diff --git a/tests/series_only/factorize_test.py b/tests/series_only/factorize_test.py index 44121769a2..69f83fb28f 100644 --- a/tests/series_only/factorize_test.py +++ b/tests/series_only/factorize_test.py @@ -1,6 +1,5 @@ from __future__ import annotations -from math import isnan from typing import Any import pytest @@ -18,18 +17,28 @@ @pytest.mark.parametrize( - ("values", "expected_n_unique"), + ("values", "null_as_value", "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), + ([], False, 0), + ([], True, 0), + ([*"abcabc"], False, 3), + ([*"abcabc"], True, 3), + ([1, 2, 3, 2], False, 3), + ([1, 2, 3, 2], True, 3), + ([1.1, 2.2, 3.3, 2.2], False, 3), + ([1.1, 2.2, 3.3, 2.2], True, 3), + ([*"abc", None], False, 3), + ([*"abc", None], True, 4), + ([*"aaabbbccc", None], False, 3), + ([*"aaabbbccc", None], True, 4), ], ) def test_factorize_invariants( - values: list[Any], expected_n_unique: int, constructor_eager: ConstructorEager + constructor_eager: ConstructorEager, + *, + values: list[Any], + null_as_value: bool, + expected_n_unique: int, ) -> None: if "polars" in str(constructor_eager) and polars_lt_v1: pytest.skip(reason=pl_skip_reason) @@ -38,87 +47,116 @@ def test_factorize_invariants( df_native = constructor_eager({"a": values}) df = nw.from_native(df_native) - codes, uniqs = df["a"].factorize() + codes, uniqs = df["a"].factorize(null_as_value=null_as_value) 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 + # codes should be integer, preserve length, and + # only contain -1 in the presence of nulls in input & `null_as_value=False` 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()) + min_code = 0 if null_as_value else -1 + assert (codes >= min_code).all() + assert (codes == -1).any() == (has_null and not null_as_value) @pytest.mark.parametrize( - ("values", "expected_uniqs", "expected_codes"), + ("values", "null_as_value", "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]), + ([], False, [], []), + ([None], False, [], [-1]), + ([*"abc"], False, [*"abc"], [0, 1, 2]), + ([*"abcabc"], False, [*"abc"], [0, 1, 2, 0, 1, 2]), + ([*"abcabc", None], False, [*"abc"], [0, 1, 2, 0, 1, 2, -1]), + ([*"aaabbbccc"], False, [*"abc"], [0, 0, 0, 1, 1, 1, 2, 2, 2]), + ([10, 11, 12], False, [10, 11, 12], [0, 1, 2]), + ([10, 11, 12, 10, 11, 12], False, [10, 11, 12], [0, 1, 2, 0, 1, 2]), + ([10, 10, 11, 11, 12, 12], False, [10, 11, 12], [0, 0, 1, 1, 2, 2]), + ([10, 11, 12, None], False, [10, 11, 12], [0, 1, 2, -1]), + ([], True, [], []), + ([None], True, [None], [0]), + ([*"abc"], True, [*"abc"], [0, 1, 2]), + ([*"abcabc"], True, [*"abc"], [0, 1, 2, 0, 1, 2]), + ([*"abcabc", None], True, [*"abc", None], [0, 1, 2, 0, 1, 2, 3]), + ([*"aaabbbccc"], True, [*"abc"], [0, 0, 0, 1, 1, 1, 2, 2, 2]), + ([10, 11, 12], True, [10, 11, 12], [0, 1, 2]), + ([10, 11, 12, 10, 11, 12], True, [10, 11, 12], [0, 1, 2, 0, 1, 2]), + ([10, 10, 11, 11, 12, 12], True, [10, 11, 12], [0, 0, 1, 1, 2, 2]), + ([10, 11, 12, None], True, [10, 11, 12, None], [0, 1, 2, 3]), ], ) def test_factorize_sort( + constructor_eager: ConstructorEager, + *, values: list[Any], + null_as_value: bool, 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) + codes, uniqs = df["a"].factorize(null_as_value=null_as_value, sort=True) assert_equal_series(uniqs, expected_uniqs, name="a") assert_equal_series(codes, expected_codes, name="a") @pytest.mark.parametrize( - "values", + ("values", "null_as_value", "expected_unique", "expected_unique_pandas"), [ - [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")], + ([1.1, 2.2, 1.1, float("nan")], False, [1.1, 2.2, float("nan")], [1.1, 2.2]), + ( + [1.1, 2.2, 1.1, float("nan"), float("nan")], + False, + [1.1, 2.2, float("nan")], + [1.1, 2.2], + ), + ( + [1.1, 2.2, 1.1, None, float("nan")], + False, + [1.1, 2.2, float("nan")], + [1.1, 2.2], + ), + ([1.1, 2.2, 1.1, float("nan")], True, [1.1, 2.2, float("nan")], [1.1, 2.2, None]), + ( + [1.1, 2.2, 1.1, float("nan"), float("nan")], + True, + [1.1, 2.2, float("nan")], + [1.1, 2.2, None], + ), + ( + [1.1, 2.2, 1.1, None, float("nan")], + True, + [1.1, 2.2, float("nan"), None], + [1.1, 2.2, None], + ), ], ) def test_factorize_nan_semantics( - values: list[float], constructor_eager: ConstructorEager + constructor_eager: ConstructorEager, + *, + values: list[float | None], + null_as_value: bool, + expected_unique: list[Any], + expected_unique_pandas: list[Any], ) -> 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")) + expected = expected_unique_pandas if is_pandas_backend else expected_unique df_native = constructor_eager({"a": values}) df = nw.from_native(df_native) - codes, uniqs = df["a"].factorize() + codes, uniqs = df["a"].factorize(null_as_value=null_as_value, sort=True) 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) + assert_equal_series(uniqs, expected, name="a") From 9d7a132029241d8eb7dfd71769ff045552daf850 Mon Sep 17 00:00:00 2001 From: Cameron Riddell Date: Fri, 24 Jul 2026 13:14:47 -0700 Subject: [PATCH 09/12] ref: factorize overwrite result names to "codes" and "uniques" --- src/narwhals/_pandas_like/series.py | 10 +++------- src/narwhals/series.py | 5 ++++- tests/series_only/factorize_test.py | 9 ++++++--- 3 files changed, 13 insertions(+), 11 deletions(-) diff --git a/src/narwhals/_pandas_like/series.py b/src/narwhals/_pandas_like/series.py index 2c79679e07..28f2732733 100644 --- a/src/narwhals/_pandas_like/series.py +++ b/src/narwhals/_pandas_like/series.py @@ -1167,7 +1167,6 @@ def factorize( self, *, null_as_value: bool = False, sort: bool = False ) -> tuple[Self, Self]: pdx = self.__native_namespace__() - name = self.native.name # https://github.com/apache/arrow/issues/33297; input pa.NullArray's don't dictionary_encode properly if self.native.dtype == "null[pyarrow]": @@ -1182,17 +1181,14 @@ def factorize( pdx.Series([], dtype=self.native.dtype), ) - return ( - self._with_native(codes.rename(name)), - self._with_native(uniques.rename(name)), - ) + 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, name=name)), - self._with_native(pdx.Series(uniques, name=name)), + self._with_native(pdx.Series(codes)), + self._with_native(pdx.Series(uniques)), ) def _apply_pyarrow_compute_func( diff --git a/src/narwhals/series.py b/src/narwhals/series.py index 378ce4b8cd..3510375b5a 100644 --- a/src/narwhals/series.py +++ b/src/narwhals/series.py @@ -2962,7 +2962,10 @@ def factorize( codes, uniques = self._compliant_series.factorize( null_as_value=null_as_value, sort=sort ) - return self._with_compliant(codes), self._with_compliant(uniques) + return ( + self._with_compliant(codes).alias("codes"), + self._with_compliant(uniques).alias("uniques"), + ) @unstable def any_value(self, *, ignore_nulls: bool = False) -> PythonLiteral: diff --git a/tests/series_only/factorize_test.py b/tests/series_only/factorize_test.py index 69f83fb28f..f084853753 100644 --- a/tests/series_only/factorize_test.py +++ b/tests/series_only/factorize_test.py @@ -63,6 +63,9 @@ def test_factorize_invariants( assert (codes >= min_code).all() assert (codes == -1).any() == (has_null and not null_as_value) + assert codes.name == "codes" + assert uniqs.name == "uniques" + @pytest.mark.parametrize( ("values", "null_as_value", "expected_uniqs", "expected_codes"), @@ -104,8 +107,8 @@ def test_factorize_sort( df = nw.from_native(df_native) codes, uniqs = df["a"].factorize(null_as_value=null_as_value, sort=True) - assert_equal_series(uniqs, expected_uniqs, name="a") - assert_equal_series(codes, expected_codes, name="a") + assert_equal_series(uniqs, expected_uniqs, name="uniques") + assert_equal_series(codes, expected_codes, name="codes") @pytest.mark.parametrize( @@ -159,4 +162,4 @@ def test_factorize_nan_semantics( reconstructed_values = {"a": [uniqs[i] if i >= 0 else None for i in codes]} assert_equal_data(df, reconstructed_values) - assert_equal_series(uniqs, expected, name="a") + assert_equal_series(uniqs, expected, name="uniques") From 44c3b556694f54e7c9d611e6ffed1310cdebec6d Mon Sep 17 00:00:00 2001 From: Cameron Riddell Date: Fri, 24 Jul 2026 13:50:20 -0700 Subject: [PATCH 10/12] ref: factorize to produce Namedtuple result --- src/narwhals/series.py | 33 ++++++++++++++++++++++++++--- tests/series_only/factorize_test.py | 7 +++++- 2 files changed, 36 insertions(+), 4 deletions(-) diff --git a/src/narwhals/series.py b/src/narwhals/series.py index 3510375b5a..c4c0afbc54 100644 --- a/src/narwhals/series.py +++ b/src/narwhals/series.py @@ -3,7 +3,16 @@ import math from collections.abc import Callable, Iterable, Iterator, Mapping, Sequence from functools import partial -from typing import TYPE_CHECKING, Any, ClassVar, Generic, Literal, cast, overload +from typing import ( + TYPE_CHECKING, + Any, + ClassVar, + Generic, + Literal, + NamedTuple, + cast, + overload, +) from narwhals._expression_parsing import ExprKind, ExprNode from narwhals._utils import ( @@ -2907,7 +2916,7 @@ def is_close( def factorize( self, *, null_as_value: bool = False, sort: bool = False - ) -> tuple[Self, Self]: + ) -> Encoded[IntoSeriesT]: """Encode values as integer codes and unique values. The integer codes are index locations that map the unique values back to their @@ -2962,7 +2971,7 @@ def factorize( codes, uniques = self._compliant_series.factorize( null_as_value=null_as_value, sort=sort ) - return ( + return Encoded( self._with_compliant(codes).alias("codes"), self._with_compliant(uniques).alias("uniques"), ) @@ -3008,3 +3017,21 @@ def list(self) -> SeriesListNamespace[Self]: @property def struct(self) -> SeriesStructNamespace[Self]: return SeriesStructNamespace(self) + + +class Encoded(NamedTuple, Generic[IntoSeriesT]): + """Result of `factorize`. Unpacks as `(codes, uniques)` like pandas.""" + + codes: Series[IntoSeriesT] + uniques: Series[IntoSeriesT] + + @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() + ) diff --git a/tests/series_only/factorize_test.py b/tests/series_only/factorize_test.py index f084853753..9e56f63d42 100644 --- a/tests/series_only/factorize_test.py +++ b/tests/series_only/factorize_test.py @@ -47,7 +47,8 @@ def test_factorize_invariants( df_native = constructor_eager({"a": values}) df = nw.from_native(df_native) - codes, uniqs = df["a"].factorize(null_as_value=null_as_value) + encoded_result = df["a"].factorize(null_as_value=null_as_value) + codes, uniqs = encoded_result reconstructed_values = {"a": [uniqs[i] if i >= 0 else None for i in codes]} assert_equal_data(df, reconstructed_values) @@ -66,6 +67,10 @@ def test_factorize_invariants( assert codes.name == "codes" assert uniqs.name == "uniques" + assert isinstance(encoded_result, nw.series.Encoded) + assert [*encoded_result.mapping.keys()] == [*encoded_result.uniques] + assert [*encoded_result.mapping.values()] == [*range(len(encoded_result.uniques))] + @pytest.mark.parametrize( ("values", "null_as_value", "expected_uniqs", "expected_codes"), From f6510a07df141b90ef715c6b87cc4bc83d49d1d2 Mon Sep 17 00:00:00 2001 From: Cameron Riddell Date: Fri, 24 Jul 2026 14:14:10 -0700 Subject: [PATCH 11/12] ref: factorize: namedtuple -> dataclass to avoid multiple inheritance --- src/narwhals/series.py | 19 ++++++++----------- tests/series_only/factorize_test.py | 1 - 2 files changed, 8 insertions(+), 12 deletions(-) diff --git a/src/narwhals/series.py b/src/narwhals/series.py index c4c0afbc54..c668d9aa35 100644 --- a/src/narwhals/series.py +++ b/src/narwhals/series.py @@ -2,17 +2,9 @@ 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, - NamedTuple, - cast, - overload, -) +from typing import TYPE_CHECKING, Any, ClassVar, Generic, Literal, cast, overload from narwhals._expression_parsing import ExprKind, ExprNode from narwhals._utils import ( @@ -3019,12 +3011,17 @@ def struct(self) -> SeriesStructNamespace[Self]: return SeriesStructNamespace(self) -class Encoded(NamedTuple, Generic[IntoSeriesT]): +@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.""" diff --git a/tests/series_only/factorize_test.py b/tests/series_only/factorize_test.py index 9e56f63d42..e8a99bf6bf 100644 --- a/tests/series_only/factorize_test.py +++ b/tests/series_only/factorize_test.py @@ -67,7 +67,6 @@ def test_factorize_invariants( assert codes.name == "codes" assert uniqs.name == "uniques" - assert isinstance(encoded_result, nw.series.Encoded) assert [*encoded_result.mapping.keys()] == [*encoded_result.uniques] assert [*encoded_result.mapping.values()] == [*range(len(encoded_result.uniques))] From fa085fd38bbd29a51183810d02bed6d0fef466e2 Mon Sep 17 00:00:00 2001 From: Cameron Riddell Date: Fri, 24 Jul 2026 14:20:24 -0700 Subject: [PATCH 12/12] fix v1/v2 factorize tests with new column names --- tests/v1_test.py | 4 ++-- tests/v2_test.py | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/tests/v1_test.py b/tests/v1_test.py index 268decb53c..0bce40c727 100644 --- a/tests/v1_test.py +++ b/tests/v1_test.py @@ -1271,8 +1271,8 @@ def test_factorize( 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_equal_series(uniqs, expected_uniqs, name="uniques") + assert_equal_series(codes, expected_codes, name="codes") assert codes._version is Version.V1 assert uniqs._version is Version.V1 diff --git a/tests/v2_test.py b/tests/v2_test.py index cc3edacc61..cb21a0243f 100644 --- a/tests/v2_test.py +++ b/tests/v2_test.py @@ -611,8 +611,8 @@ def test_factorize( 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_equal_series(uniqs, expected_uniqs, name="uniques") + assert_equal_series(codes, expected_codes, name="codes") assert codes._version is Version.V2 assert uniqs._version is Version.V2