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/_arrow/series.py b/src/narwhals/_arrow/series.py index 74fbac0df1..4e0883d394 100644 --- a/src/narwhals/_arrow/series.py +++ b/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) diff --git a/src/narwhals/_compliant/series.py b/src/narwhals/_compliant/series.py index ea83e6f86b..dabe1c6deb 100644 --- a/src/narwhals/_compliant/series.py +++ b/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: ... diff --git a/src/narwhals/_pandas_like/series.py b/src/narwhals/_pandas_like/series.py index 48ee244c22..28f2732733 100644 --- a/src/narwhals/_pandas_like/series.py +++ b/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: diff --git a/src/narwhals/_polars/series.py b/src/narwhals/_polars/series.py index b82a911f61..eff6a77204 100644 --- a/src/narwhals/_polars/series.py +++ b/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) diff --git a/src/narwhals/series.py b/src/narwhals/series.py index d0688cba41..c668d9aa35 100644 --- a/src/narwhals/series.py +++ b/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() + ) diff --git a/tests/series_only/factorize_test.py b/tests/series_only/factorize_test.py new file mode 100644 index 0000000000..e8a99bf6bf --- /dev/null +++ b/tests/series_only/factorize_test.py @@ -0,0 +1,169 @@ +from __future__ import annotations + +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", "null_as_value", "expected_n_unique"), + [ + ([], 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( + 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) + + has_null = any(x is None for x in values) + + df_native = constructor_eager({"a": values}) + df = nw.from_native(df_native) + 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) + 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 in input & `null_as_value=False` + assert codes.dtype.is_integer() + assert len(codes) == len(values) + + 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) + + assert codes.name == "codes" + assert uniqs.name == "uniques" + + 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"), + [ + ([], 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], +) -> 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(null_as_value=null_as_value, sort=True) + + assert_equal_series(uniqs, expected_uniqs, name="uniques") + assert_equal_series(codes, expected_codes, name="codes") + + +@pytest.mark.parametrize( + ("values", "null_as_value", "expected_unique", "expected_unique_pandas"), + [ + ([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( + 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(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) + assert_equal_series(uniqs, expected, name="uniques") diff --git a/tests/v1_test.py b/tests/v1_test.py index ac048112e5..0bce40c727 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="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 bd1a3a9036..cb21a0243f 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 = df["a"].factorize(sort=True) + + 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 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",