ENH PERF Speed up unique counts for strings, thereby speeding up encoding (scikit-learn#34386)

Author: itamarst

doc/whats_new/upcoming_changes/sklearn.preprocessing/34386.efficiency.rst

@@ -0,0 +1,5 @@
+- Improved fit time of :class:`preprocessing.OneHotEncoder`
+  and :class:`preprocessing.OrdinalEncoder` on object/string categorical
+  features when category counts are needed, for instance with ``min_frequency``
+  or ``max_categories``.
+  By :user:`Itamar Turner-Trauring <itamarst>`.

sklearn/utils/_encode.py

@@ -2,7 +2,8 @@
 # SPDX-License-Identifier: BSD-3-Clause
 
 from collections import Counter
-from contextlib import suppress
+from collections.abc import Iterable
+from numbers import Real
 from typing import NamedTuple
 
 import numpy as np
@@ -96,6 +97,9 @@ class MissingValues(NamedTuple):
 
     nan: bool
     none: bool
+    # All the different nan instances seen, useful because they end up merged
+    # into just np.nan:
+    all_nans: Iterable[Real] = (np.nan,)
 
     def to_list(self):
         """Convert tuple to a list where None is always first."""
@@ -130,15 +134,20 @@ def _extract_missing(values):
     if not missing_values_set:
         return values, MissingValues(nan=False, none=False)
 
+    all_nans = missing_values_set - {None}
     if None in missing_values_set:
         if len(missing_values_set) == 1:
-            output_missing_values = MissingValues(nan=False, none=True)
+            output_missing_values = MissingValues(
+                nan=False, none=True, all_nans=all_nans
+            )
         else:
             # If there is more than one missing value, then it has to be
             # float('nan') or np.nan
-            output_missing_values = MissingValues(nan=True, none=True)
+            output_missing_values = MissingValues(
+                nan=True, none=True, all_nans=all_nans
+            )
     else:
-        output_missing_values = MissingValues(nan=True, none=False)
+        output_missing_values = MissingValues(nan=True, none=False, all_nans=all_nans)
 
     # create set without the missing values
     output = values - missing_values_set
@@ -189,7 +198,7 @@ def _unique_python(values, *, return_inverse, return_counts):
         ret += (_map_to_integer(values, uniques),)
 
     if return_counts:
-        ret += (_get_counts(values, uniques),)
+        ret += (_get_counts(values, uniques, missing_values.all_nans),)
 
     return ret[0] if len(ret) == 1 else ret
 
@@ -321,40 +330,21 @@ def is_valid(value):
     return diff
 
 
-class _NaNCounter(Counter):
-    """Counter with support for nan values."""
-
-    def __init__(self, items):
-        super().__init__(self._generate_items(items))
-
-    def _generate_items(self, items):
-        """Generate items without nans. Stores the nan counts separately."""
-        for item in items:
-            if not is_scalar_nan(item):
-                yield item
-                continue
-            if not hasattr(self, "nan_count"):
-                self.nan_count = 0
-            self.nan_count += 1
-
-    def __missing__(self, key):
-        if hasattr(self, "nan_count") and is_scalar_nan(key):
-            return self.nan_count
-        raise KeyError(key)
-
-
-def _get_counts(values, uniques):
+def _get_counts(values, uniques, nan_values=(np.nan,)):
     """Get the count of each of the `uniques` in `values`.
 
-    The counts will use the order passed in by `uniques`. For non-object dtypes,
-    `uniques` is assumed to be sorted and `np.nan` is at the end.
+    The counts will use the order passed in by `uniques`.  `np.nan` is assumed
+    to be the last item in `uniques`, if it was one of the values.  For
+    non-object dtypes, `uniques` is assumed to be sorted.
     """
     if values.dtype.kind in "OU":
-        counter = _NaNCounter(values)
+        counter = Counter(values)
         output = np.zeros(len(uniques), dtype=np.int64)
         for i, item in enumerate(uniques):
-            with suppress(KeyError):
-                output[i] = counter[item]
+            output[i] = counter[item]
+        if len(uniques) > 0 and is_scalar_nan(uniques[-1]):
+            # Should be the sum of all nans:
+            output[-1] = sum(counter[nan] for nan in nan_values)
         return output
 
     unique_values, counts = _unique_np(values, return_counts=True)

sklearn/utils/_missing.py

@@ -1,9 +1,9 @@
 # Authors: The scikit-learn developers
 # SPDX-License-Identifier: BSD-3-Clause
 
-import math
-import numbers
 from contextlib import suppress
+from math import isnan
+from numbers import Real
 
 
 def is_scalar_nan(x):
@@ -37,11 +37,7 @@ def is_scalar_nan(x):
     >>> is_scalar_nan([np.nan])
     False
     """
-    return (
-        not isinstance(x, numbers.Integral)
-        and isinstance(x, numbers.Real)
-        and math.isnan(x)
-    )
+    return isinstance(x, Real) and isnan(x)
 
 
 def is_pandas_na(x):

sklearn/utils/tests/test_encode.py

@@ -233,6 +233,10 @@ def test_check_unknown_with_both_missing_values():
     assert_array_equal(valid_mask, [False, True, True, True, False, False, False])
 
 
+NAN1 = float("nan")
+NAN2 = float("nan")
+
+
 @pytest.mark.parametrize(
     "values, uniques, expected_counts",
     [
@@ -272,3 +276,25 @@ def test_check_unknown_with_both_missing_values():
 def test_get_counts(values, uniques, expected_counts):
     counts = _get_counts(values, uniques)
     assert_array_equal(counts, expected_counts)
+
+
+def test_get_counts_multiple_nans():
+    """
+    When both np.nan and float("nan") are present, they get merged into np.nan.
+    """
+    values = np.array(
+        ["a", np.nan, NAN1, np.nan, NAN2, NAN1, np.nan, "a"],
+        dtype=object,
+    )
+    uniques = np.array(["a", np.nan], dtype=object)
+    expected_counts = [2, 6]
+    assert_array_equal(
+        _get_counts(values, uniques, [np.nan, NAN1, NAN2]), expected_counts
+    )
+
+    # Now try it via _unique, to make sure this works end-to-end:
+    real_uniques, real_counts = _unique(values, return_counts=True)
+    # Comparing two arrays with nan fails cause the nans are not equal to
+    # themselves. So compare as Python lists:
+    assert list(uniques) == list(real_uniques)
+    assert_array_equal(real_counts, expected_counts)