feat: add nw.factorize - #3809
Conversation
factorize produces unique values and value mappings back to the locations of those values.
6d1f882 to
9a5dacf
Compare
|
Will this call
Given "codes" are supposed to be non-negative integers, I'd say this is good. Or at least, it's much better than other options that comes to my mind: NaN, MaxInt32, n_uniques + 1, ...? |
FBruzzesi
left a comment
FBruzzesi
left a comment
There was a problem hiding this comment.
Thanks @camriddell - I left a couple of inline comments.
On top of those:
- I think the concern from @cakedev0 on performances is totally valid
- Is there a reason to not expose the feature in
stable.v1? - I didn't check the tests 🙈
- `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
|
There was a problem hiding this comment.
Thanks again @camriddell 🙏🏼
Left a few inline suggestions. On top of those:
-
I like the possibility in pyarrow to treat null as any other value:
dictionary_encode(array, null_encoding='mask')-> keeps null's in the indicesdictionary_encode(array, null_encoding='encode')-> encodes null values as any other value, not with a sentinel value.
I think this is preferable, a user can always control the sentinel value to use bynull_encoding='mask'+fill_null(<whatever value they want>).
On the variable name (null_encoding), and its values ("mask"and"encode"), we can discuss better options
-
As of #3809 (comment), I understand the appeal of being able to do
uniques[codes] == values. However, since this is greenfield and very unlikely that polars will have the same method, I would like to have some creative freedom and propose a way to satisfy both returning uniques and a mapping.
We can return a tuple of (codes, uniques), however, we can do so in aNamedTuple, so that we can enrich its behavior (draft version):from narwhals import Series from typing import Any, NamedTuple, Generic from collections.abc import Mapping from narwhals.typing import IntoSeriesT 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() )
In this way, one can both:
codes, uniques = series.factorize() encoded = series.factorize() encoded.mapping
Let me know your thoughts - In general I am not the biggest fan of NamedTuple's for various reasons, yet this seem like a good case to have one ;)
…feat-factorize
|
@FBruzzesi I think that wraps up all of your requests. Seems like I'll need to add a back up path for an old pandas version, then we should hopefully be squared away. |
FBruzzesi
left a comment
FBruzzesi
left a comment
There was a problem hiding this comment.
Thanks @camriddell - I did a quick iteration for the top-level behavior, I will wait for the CI to be green for all the nitty gritty details
| @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() | ||
| ) |
There was a problem hiding this comment.
I am basically self-reviewing this bit - I ended up suggesting:
- A wrong docstring (from another version I was writing)
- possibly overlapping name ("code")
| @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() | |
| ) | |
| @property | |
| def mapping(self) -> Mapping[Any, int]: | |
| """Forward map `value -> code` as a dict (handy for small cardinality).""" | |
| uniques = self.uniques | |
| return dict(zip(uniques, range(len(self.uniques)), strict=True)) |
| 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. |
There was a problem hiding this comment.
I would still hold against having a sentinel value unless provided. I understand the ergonomic, but in polars spirit, I think there should be a way to preserve null values. null_as_value is probably not a good name for it, but we could enrich what pyarrow is doing:
"mask"-> preserve null values -> I would like this to be the default behavior"encode"-> same asnull_as_value = Truei.e. null's get encoded as any other value<value>-> encode with such sentinel value, equivalent tocodes, _ = series.factorize("mask") codes.fill_null(<value)>
I am curious to know how others feel about this
There was a problem hiding this comment.
Beyond ergonomics, preserving the null value in the returned codes breaks base pandas as the returned NaN upcasts the integer codes to become floats. I'd prefer not to have the default case create a surprise for the backend where this function is likely to be used most often.
I am fine with the remaining semantics, just would advocate to make the "encode" behavior the default.
There was a problem hiding this comment.
Beyond ergonomics, preserving the null value in the returned codes breaks ... as the returned NaN upcasts the integer codes to become floats.
Is it not possible to use pd.Int64Dtype?
I'd prefer not to have the default case create a surprise for the backend where this function is likely to be used most often.
I understand this concern, but I agree with @FBruzzesi that the default behavior should feel like polars.
Third option
What if we held back on nw.Series.factorize for now, and instead introduced how to do it into our docs?
This feels pretty related to discord thread and would slide in perfectly into a whole section about transitioning from pd.Categorical 🙂
@camriddell, correct me if I'm wrong, but that feels like something you'd have a lot to speak on 😅
There was a problem hiding this comment.
Is it not possible to use pd.Int64Dtype?
@dangotbanned yeah we can return an Int64 dtype. Originally I had thought that we supported a pandas version that didn't have nullable types but I am no longer certain that is the case.
What if we held back on nw.Series.factorize for now, and instead introduced how to do it into our docs?
I am biased as I have already sunk a few hours into this, but I would like to see it land. This is also a fairly common operation for our array-based libraries downstream users, and the backend implementations are all quite different from one another.
@FBruzzesi & @dangotbanned how do you feel about this as an API? I would love it if the SENTINEL policy was its own struct, so one could embed the sentinel value inside of it but we're using the wrong language for that.
class NullEncoding(StrEnum):
PRESERVE = auto() # keep the Nulls in the resultant codes
ENCODE = auto() # encode Nulls as one would like any other value
SENTINEL = auto() # replace Nulls with a particular `sentinel` value
class Series:
def factorize(
self, *, sort: bool = False, null_policy: NullEncoding = NullEncoding.PRESERVE, sentinel: Any | NoDefault = NO_DEFAULT,
) -> Encoded[IntoSeriesT]:
if null_policy is NullEncoding.SENTINEL:
if sentinel is NO_DEFAULT:
msg = "Must supply `sentinel` when null_policy=NullEncoding.SENTINEL"
raise TypeError(msg)
elif sentinel is not NO_DEFAULT:
msg = f"Argument `sentinel` is ignored when null_policy={null_policy}"
raise TypeError(msg)
...There was a problem hiding this comment.
Hey @camriddell sorry for the late feedback, I completely lost this thread!
I like having null_policy (and such name) as well as using sentinel. I would disagree on creating a custom Enum to expose to the user. I think Literal[...] is just fine.
On the larger picture of having the feature vs documenting it, I would be in favor of having it with the unstable tag. As Cam mentioned, this topic came up quite a lot recently but not only, and there is real need for it.
On the Int64 dtype for pandas: I am quite ok with that, keep in mind that one day pyarrow might be a default/requirement if I remember correctly. I am sure Marco knows much more about it.
| of the corresponding value in `uniques`. Null values are encoded | ||
| as -1. |
There was a problem hiding this comment.
Out of sync with the new behaviour
dangotbanned
left a comment
dangotbanned
left a comment
There was a problem hiding this comment.
Thanks @camriddell, hope some of this can be helpful
| result = result.rename(orig_name) if name_is_none else result | ||
| return cast("Self", result) | ||
|
|
||
| def factorize( |
There was a problem hiding this comment.
Could we please mark this as unstable?
| @dataclass(frozen=True) | ||
| class Encoded(Generic[IntoSeriesT]): |
There was a problem hiding this comment.
Could this also be marked as unstable and/or private please?
It is mostly an implementation detail
| assert schema._version is Version.V2 | ||
|
|
||
|
|
||
| @pytest.mark.parametrize( |
There was a problem hiding this comment.
Per the other comments, if factorize is unstable, would we still have these v1, v2 tests?
| 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), | ||
| ) |
There was a problem hiding this comment.
As somebody who has battled against pyarrow's frustrating API a lot - I feel like this is doing too many steps 🤔
I'm not 100% sure what is going on here, so I don't a concrete suggestion.
But I do have some code that may be helpful to look at. If only inspire an oh that's how these bits could fit together moment 😉
Show (maybe related) pyarrow functions
Definitions
- https://github.com/narwhals-dev/narwhals/blob/c179a78e4f3d6c95530aea10a7b2f917ec54ac21/src/narwhals/_plan/arrow/functions/_categorical.py
- https://github.com/narwhals-dev/narwhals/blob/c179a78e4f3d6c95530aea10a7b2f917ec54ac21/src/narwhals/_plan/arrow/functions/_multiplex.py
- https://github.com/narwhals-dev/narwhals/blob/c179a78e4f3d6c95530aea10a7b2f917ec54ac21/src/narwhals/_plan/arrow/functions/_sort.py
Some usage
narwhals/src/narwhals/_plan/arrow/dataframe.py
Lines 301 to 313 in c179a78
narwhals/src/narwhals/_plan/arrow/functions/_aggregation.py
Lines 152 to 159 in c179a78
narwhals/src/narwhals/_plan/arrow/series.py
Lines 183 to 191 in c179a78
narwhals/src/narwhals/_plan/arrow/functions/_lists.py
Lines 352 to 365 in c179a78
There was a problem hiding this comment.
I tested pc.dictionary_encoding vs uniques -> replace and found that the former was faster at large dataset sizes.
Timing code
import pyarrow as pa
import pyarrow.compute as pc
def factorize_via_replace(arr, *, sort=False):
uniques = pc.unique(arr)
if sort:
uniques = pc.take(uniques, pc.sort_indices(uniques))
return pc.index_in(arr, value_set=uniques), uniques
def factorize_via_dictionary(arr, sort=False):
encoded = pc.dictionary_encode(arr, null_encoding='mask').combine_chunks()
uniques = encoded.dictionary
codes = pa.chunked_array(encoded.indices)
if not sort:
return codes, 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 new_codes, sorted_uniques
import numpy as np
from numpy.random import default_rng
from string import ascii_lowercase
from timeit import timeit
from itertools import product
rng = default_rng(0)
n = 10
for sz, n_chunks, sort in product([10_000, 1_000_000], [5, 100, 10_000], [True, False]):
n_chunks = min(n_chunks, sz)
arr = rng.choice([*ascii_lowercase], size=(sz, 4)).view('<U4').ravel()
arr = pa.chunked_array(np.array_split(arr, n_chunks), type=pa.string())
c1, u1 = factorize_via_replace(arr, sort=sort)
c2, u2 = factorize_via_dictionary(arr, sort=sort)
assert u1.equals(u2)
assert c1.equals(c2)
print(f'N={sz:<10,}| Chunks={n_chunks:<10,}| sort={sort!r:<5}| Unique={len(u1):,}')
print(f"{timeit(lambda: factorize_via_replace(arr, sort=sort), number=n) = :.6f}")
print(f"{timeit(lambda: factorize_via_dictionary(arr, sort=sort), number=n) = :.6f}")
print()Timing results, stepping through dictionary_encoding is *generally* faster, but not in every scenario. ~2× faster for the largest array size (1M values, 10k chunks) when `sort=False`.
N=10,000 | Chunks=5 | sort=True | Unique=9,874
timeit(lambda: factorize_via_replace(arr, sort=sort), number=n) = 0.021563
timeit(lambda: factorize_via_dictionary(arr, sort=sort), number=n) = 0.028813
N=10,000 | Chunks=5 | sort=False| Unique=9,910
timeit(lambda: factorize_via_replace(arr, sort=sort), number=n) = 0.007005
timeit(lambda: factorize_via_dictionary(arr, sort=sort), number=n) = 0.003449
N=10,000 | Chunks=100 | sort=True | Unique=9,894
timeit(lambda: factorize_via_replace(arr, sort=sort), number=n) = 0.028647
timeit(lambda: factorize_via_dictionary(arr, sort=sort), number=n) = 0.029794
N=10,000 | Chunks=100 | sort=False| Unique=9,893
timeit(lambda: factorize_via_replace(arr, sort=sort), number=n) = 0.015438
timeit(lambda: factorize_via_dictionary(arr, sort=sort), number=n) = 0.006601
N=10,000 | Chunks=10,000 | sort=True | Unique=9,905
timeit(lambda: factorize_via_replace(arr, sort=sort), number=n) = 0.067098
timeit(lambda: factorize_via_dictionary(arr, sort=sort), number=n) = 0.241810
N=10,000 | Chunks=10,000 | sort=False| Unique=9,894
timeit(lambda: factorize_via_replace(arr, sort=sort), number=n) = 0.049481
timeit(lambda: factorize_via_dictionary(arr, sort=sort), number=n) = 0.222008
N=1,000,000 | Chunks=5 | sort=True | Unique=405,646
timeit(lambda: factorize_via_replace(arr, sort=sort), number=n) = 3.767795
timeit(lambda: factorize_via_dictionary(arr, sort=sort), number=n) = 3.885547
N=1,000,000 | Chunks=5 | sort=False| Unique=405,821
timeit(lambda: factorize_via_replace(arr, sort=sort), number=n) = 2.810773
timeit(lambda: factorize_via_dictionary(arr, sort=sort), number=n) = 1.740194
N=1,000,000 | Chunks=100 | sort=True | Unique=406,059
timeit(lambda: factorize_via_replace(arr, sort=sort), number=n) = 3.903205
timeit(lambda: factorize_via_dictionary(arr, sort=sort), number=n) = 3.496898
N=1,000,000 | Chunks=100 | sort=False| Unique=405,580
timeit(lambda: factorize_via_replace(arr, sort=sort), number=n) = 2.629999
timeit(lambda: factorize_via_dictionary(arr, sort=sort), number=n) = 1.398467
N=1,000,000 | Chunks=10,000 | sort=True | Unique=405,871
timeit(lambda: factorize_via_replace(arr, sort=sort), number=n) = 3.914798
timeit(lambda: factorize_via_dictionary(arr, sort=sort), number=n) = 3.816666
N=1,000,000 | Chunks=10,000 | sort=False| Unique=405,657
timeit(lambda: factorize_via_replace(arr, sort=sort), number=n) = 2.829571
timeit(lambda: factorize_via_dictionary(arr, sort=sort), number=n) = 1.563878
There was a problem hiding this comment.
(do note that in my timing I used combine_chunks() as that ended up being MUCH faster than the unify_dictionaries() approach. So I'll need to swap that out regardless.
| 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) |
There was a problem hiding this comment.
I originally had a version that dispatched based on the presense of nulls, but it felt a bit more like an implementation detail that produced a surprising behavior. Was planning to open an upstream issue.
>>> import pyarrow as pa
>>> arr = pa.array([*'ab', None])
>>> arr.dictionary_encode("encode").is_null()
<pyarrow.lib.BooleanArray object at 0x7f76b91b56c0>
[
false,
false,
false
]
>>> arr.dictionary_encode("mask").is_null()
<pyarrow.lib.BooleanArray object at 0x7f76b91b5c00>
[
false,
false,
true
]| def factorize( | ||
| self, *, null_as_value: bool = False, sort: bool = False | ||
| ) -> tuple[Self, Self]: |
There was a problem hiding this comment.
Is polars.Series.to_physical no good for this?
The example says:
Replicating the pandas pd.Series.factorize method
There was a problem hiding this comment.
factorize needs to guarantee that the returned unique values and the codes align. Since to_physical requires a Categorical/Enum input, the global categorical cache makes the physical codes surprising. However, we could reasonably make a fast path for an Enum type:
pl.Categorical
>>> s1 = pl.Series("X", [*"abc"]).cast(pl.Categorical)
>>> s2 = pl.Series("X", [*"abd"]).cast(pl.Categorical)
>>> s1.to_physical()
shape: (3,)
Series: 'X' [u32]
[
0
1
2
]
>>> s2.to_physical()
shape: (3,)
Series: 'X' [u32]
[
0
1
3
]pl.Enum
>>> import polars as pl
>>> e1 = pl.Enum(['a', 'b', 'c'])
>>> e2 = pl.Enum(['a', 'b', 'd'])
>>> s1 = pl.Series("X", [*"abc"]).cast(e1)
>>> s2 = pl.Series("X", [*"abd"]).cast(e2)
>>> s1.to_physical()
shape: (3,)
Series: 'X' [u8]
[
0
1
2
]
>>> s2.to_physical()
shape: (3,)
Series: 'X' [u8]
[
0
1
2
]| 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() | ||
| ) |
There was a problem hiding this comment.
Have you considered starting this with a self.native.to_frame() and then doing all the steps with expressions?
Show pattern
This is based on the magic polars does for most Series methods:
narwhals/src/narwhals/_polars/series.py
Lines 808 to 811 in 05ce363
Chaining Series methods has to do all these steps on every intermediate call. And this is just the python part 😅
For example, what if native was this?
import polars as pl
N = 10_000_000
N_UNIQUE = 100
native = (
pl.select(big_boi=pl.int_range(N), eager=False)
.gather((pl.col.big_boi % N_UNIQUE).shuffle(1))
.collect()
.to_series()
)
nativeShow output
shape: (10_000_000,)
Series: 'big_boi' [i64]
[
28
88
5
1
71
…
58
1
23
92
27
]
Would there be a benefit to not materializing until as late as possible?
uniques = native.unique().drop_nulls().sort(nulls_last=True)
codes = native.replace_strict(uniques, range(len(uniques)), return_dtype=pl.Int32()).cast(
pl.Int32()
)There was a problem hiding this comment.
I had intentionally re-used our PolarsSeries.sort as it implements a fallback for a previous version of Polars where the nulls_last argument did not exist and I didn't want to duplicate that code.
Would there be a benefit to not materializing until as late as possible?
Probably, but "as late as possible" comes pretty early as we must materialize before calling replace (otherwise range(len(uniques)) does not work). So pretty much by using expressions we can push the unique() and possibly drop_nulls & sort into the same operation.
4 participants
factorize produces unique values and value mappings back to the locations of those values. Inspired by pandas.factorize
Implementation Details
Value Semantics
uniquesis the set of distinct non-null values.codes[i]is the index into uniques forvalues[i].Missing Value Semantics:
-1.To Do/Discuss
Pure Narwhals implementation: is it tolerable to only have this implemented at the
narwhals.functionslevel and NOT within each individual backend? This function can be composed purely of pieces that already existMissing Values:
use_na_sentinelfrom pandas.factorize?drop_nullsornull_as_value)-1an appropriate "missing" mapping value? Currently when one factorizes a Series with a Null (or NaN in pandas), the Null is automatically dropped from the unique value set and the corresponding mapping code becomes-1.Description
What type of PR is this? (check all applicable)
Related issues
AI assistance
Checklist
Code follows style guide (ruff)
Tests added
Documented the changes
If this is your first PR to narwhals, attach a screenshot of
pytestpassing locally (not CI):