Repository: dask/dask
Pull request #10922 (2024-02-14 10:09:35)

Make nunique faster again
- [ ] Closes #xxxx
- [ ] Tests added / passed
- [ ] Passes `pre-commit run --all-files`

The current solution is still super slow, this will get a 10x speed improvement for high cardinality cases and will most likely be faster for low cardinality cases as well, but not so significantly

github-actions[bot]: ## Unit Test Results
_See [test report](https://dask.github.io/dask/test_report.html) for an extended history of previous test failures. This is useful for diagnosing flaky tests._

     15 files  ±0       15 suites  ±0   3h 19m 29s :stopwatch: + 1m 42s
 13 021 tests ±0   12 090 :white_check_mark:  - 1     930 :zzz: ±0  1 :x: +1 
161 002 runs  ±0  144 509 :white_check_mark:  - 3  16 492 :zzz: +2  1 :x: +1 

For more details on these failures, see [this check](https://github.com/dask/dask/runs/21560535131).

Results for commit a5bd8688. ± Comparison against base commit 07099e5c.

[test-results]:data:application/gzip;base64,H4sIAAyZzGUC/02Myw7CIBQFf8WwdnF5luvPGCiQEPswFFbGfxeq0i5nzsm8SIiT38jtQuX1QrYS80GuJJPjujSmqLCquuZ958Don+9bGccmGSAc8hGfVSI/VDBxar8ufEprqqZdUln2sqIA7Me9LIQE7PJbpkpgP57SO5/L4zrPMVcgRlqnldbovAhG6SFwIcNgLGc2cGsogrcDdeT9Af7y1wUZAQAA


Patch:
diff --git a/dask/base.py b/dask/base.py
index 795ebd536..f32095d55 100644
--- a/dask/base.py
+++ b/dask/base.py
@@ -7,7 +7,6 @@ import hashlib
 import inspect
 import pathlib
 import pickle
-import random
 import types
 import uuid
 import warnings
@@ -16,7 +15,6 @@ from collections.abc import Hashable, Iterable, Iterator, Mapping
 from concurrent.futures import Executor
 from contextlib import contextmanager, suppress
 from contextvars import ContextVar
-from enum import Enum
 from functools import partial
 from numbers import Integral, Number
 from operator import getitem
@@ -1052,17 +1050,9 @@ normalize_token.register(
 )
 
 
-@normalize_token.register(type)
-def normalize_type(typ):
-    # This allows for some ambiguity in case a type is redefined in a local
-    # scope of a module however, defining this further breaks tokenization for
-    # local objects entirely
-    return typ.__name__, typ.__module__
-
-
-@normalize_token.register((dict, types.MappingProxyType))
+@normalize_token.register((types.MappingProxyType, dict))
 def normalize_dict(d):
-    return "dict", _normalize_seq_func(sorted(d.items(), key=str))
+    return "dict", _normalize_seq_func(sorted(d.items(), key=lambda kv: str(kv[0])))
 
 
 @normalize_token.register(OrderedDict)
@@ -1132,21 +1122,6 @@ def normalize_literal(lit):
     return "literal", normalize_token(lit())
 
 
-@normalize_token.register(range)
-def normalize_range(r):
-    return "range", _normalize_seq_func((r.start, r.stop, r.step))
-
-
-@normalize_token.register(Enum)
-def normalize_enum(e):
-    return _normalize_seq_func((type(e), e.name, e.value))
-
-
-@normalize_token.register(random.Random)
-def normalize_random_state(state):
-    return normalize_token(state.getstate())
-
-
 @normalize_token.register(Compose)
 def normalize_compose(func):
     return _normalize_seq_func((func.first,) + func.funcs)
@@ -1178,7 +1153,7 @@ _seen_objects = set()
 @normalize_token.register(object)
 def normalize_object(o):
     method = getattr(o, "__dask_tokenize__", None)
-    if method is not None:
+    if method is not None and not isinstance(o, type):
         return method()
 
     if type(o) is object:
@@ -1323,90 +1298,18 @@ def register_pyarrow():
 def register_numpy():
     import numpy as np
 
-    @normalize_token.register(np.generic)
-    @normalize_token.register(np.ndarray)
-    def normalize_array(x):
-        if not x.shape:
-            # Note: this path is used for both scalar and empty arrays
-            return (x.item(), x.dtype)
-
-        if x.size == 0:
-            return (None, x.dtype, x.shape)
-
-        if x.dtype.hasobject:
-            try:
-                try:
-                    # string fast-path
-                    data = hash_buffer_hex(
-                        "-".join(x.flat).encode(
-                            encoding="utf-8", errors="surrogatepass"
-                        )
-                    )
-                except UnicodeDecodeError:
-                    # bytes fast-path
-                    data = hash_buffer_hex(b"-".join(x.flat))
-            except (TypeError, UnicodeDecodeError):
-                data = _normalize_pickle(x)
-        else:
-            try:
-                data = hash_buffer_hex(x.ravel(order="K").view("i1"))
-            except (BufferError, AttributeError, ValueError):
-                data = hash_buffer_hex(x.copy().ravel(order="K").view("i1"))
-        return data, x.dtype, x.shape
-
-    @normalize_token.register(np.matrix)
-    def normalize_matrix(x):
-        return _normalize_seq_func((type(x), x.view(type=np.ndarray)))
-
-    normalize_token.register(np.dtype, repr)
-
     @normalize_token.register(np.ufunc)
-    def normalize_ufunc(x):
+    def normalize_ufunc(func):
         try:
-            name = x.__name__
-            if getattr(np, name) is x:
-                return "np." + name
-        except AttributeError:
-            try:
-                return _normalize_pickle(x)
-            except Exception:
-                if config.get("tokenize.ensure-deterministic"):
-                    raise RuntimeError(
-                        f"Cannot tokenize numpy ufunc {x!r}. Please use functions "
-                        "of the dask.array.ufunc module instead. See also "
-                        "https://docs.dask.org/en/latest/array-numpy-compatibility.html"
-                    )
-                return uuid.uuid4().hex
-
-    @normalize_token.register(np.random.RandomState)
-    def normalize_np_random_state(state):
-        return normalize_token(state.get_state())
-
-    @normalize_token.register(np.random.BitGenerator)
-    def normalize_bit_generator(bg):
-        return normalize_token(bg.state)
-
-
-@normalize_token.register_lazy("scipy")
-def register_scipy():
-    import scipy.sparse as sp
-
-    def normalize_sparse_matrix(x, attrs):
-        return _normalize_seq_func((type(x), [getattr(x, key) for key in attrs]))
-
-    for cls, attrs in [
-        (sp.dia_matrix, ("data", "offsets", "shape")),
-        (sp.bsr_matrix, ("data", "indices", "indptr", "blocksize", "shape")),
-        (sp.coo_matrix, ("data", "row", "col", "shape")),
-        (sp.csr_matrix, ("data", "indices", "indptr", "shape")),
-        (sp.csc_matrix, ("data", "indices", "indptr", "shape")),
-        (sp.lil_matrix, ("data", "rows", "shape")),
-    ]:
-        normalize_token.register(cls, partial(normalize_sparse_matrix, attrs=attrs))
-
-    @normalize_token.register(sp.dok_matrix)
-    def normalize_dok_matrix(x):
-        return _normalize_seq_func((type(x), x.tocoo()))
+            return _normalize_pickle(func)
+        except Exception:
+            if config.get("tokenize.ensure-deterministic"):
+                raise RuntimeError(
+                    f"Cannot tokenize numpy ufunc {func!r}. Please use functions "
+                    "of the dask.array.ufunc module instead. See also "
+                    "https://docs.dask.org/en/latest/array-numpy-compatibility.html"
+                )
+            return uuid.uuid4().hex
 
 
 def _colorize(t):
diff --git a/dask/dataframe/groupby.py b/dask/dataframe/groupby.py
index b85102336..6dc737e0c 100644
--- a/dask/dataframe/groupby.py
+++ b/dask/dataframe/groupby.py
@@ -745,6 +745,12 @@ def _cov_agg(_t, levels, ddof, std=False, sort=False):
 
 def _nunique_df_chunk(df, *by, **kwargs):
     name = kwargs.pop("name")
+    try:
+        # This is a lot faster but kind of a pain to implement when by
+        # has a boolean series in it.
+        return df.drop_duplicates(subset=list(by) + [name]).set_index(list(by))
+    except Exception:
+        pass
     group_keys = {}
     if PANDAS_GE_150:
         group_keys["group_keys"] = True
diff --git a/pyproject.toml b/pyproject.toml
index 7d83259ab..b6ad64c9b 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -129,8 +129,6 @@ markers = [
 ]
 addopts = "-v -rsxfE --durations=10 --color=yes --cov-config=pyproject.toml"
 filterwarnings = [
-    # From Cython-1753
-    "ignore:can't resolve:ImportWarning",
     "error:::dask[.*]",
     "error:::pandas[.*]",
     "error:::numpy[.*]",


Test patch:
diff --git a/dask/tests/test_tokenize.py b/dask/tests/test_tokenize.py
index f418e4543..adce0573a 100644
--- a/dask/tests/test_tokenize.py
+++ b/dask/tests/test_tokenize.py
@@ -29,6 +29,7 @@ np = import_or_none("numpy")
 sp = import_or_none("scipy.sparse")
 pa = import_or_none("pyarrow")
 pd = import_or_none("pandas")
+numba = import_or_none("numba")
 
 
 def check_tokenize(*args, **kwargs):
@@ -242,8 +243,7 @@ def test_tokenize_numpy_memmap_no_filename():
 
 @pytest.mark.skipif("not np")
 def test_tokenize_numpy_ufunc():
-    assert check_tokenize(np.sin) == check_tokenize("np.sin")
-    assert check_tokenize(np.cos) == check_tokenize("np.cos")
+    assert check_tokenize(np.sin) != check_tokenize(np.cos)
 
     np_ufunc = np.sin
     np_ufunc2 = np.cos
@@ -437,23 +437,14 @@ def test_tokenize_functions_unique_token():
     assert len(set(tokens)) == len(tokens)
 
 
-def test_tokenize_local_classes_from_different_contexts():
-    def f():
-        class C:
-            pass
-
-        return C
-
-    assert check_tokenize(f()) == check_tokenize(f())
-
-
 @pytest.mark.xfail(reason="https://github.com/cloudpipe/cloudpickle/issues/453")
-def test_tokenize_local_instances_from_different_contexts():
+@pytest.mark.parametrize("instance", [False, True])
+def test_tokenize_local_classes_from_different_contexts(instance):
     def f():
         class C:
             pass
 
-        return C()
+        return C() if instance else C
 
     assert check_tokenize(f()) == check_tokenize(f())
 
@@ -790,6 +781,42 @@ def test_tokenize_ordered_dict():
     assert check_tokenize(a) != check_tokenize(c)
 
 
+def test_tokenize_dict_doesnt_call_str_on_values():
+    class C:
+        def __dask_tokenize__(self):
+            return "C"
+
+        def __repr__(self):
+            assert False
+
+    check_tokenize({1: C(), "2": C()})
+
+
+def test_tokenize_sorts_dict_before_seen_map():
+    """When sequence values are repeated, the 2nd+ entry is tokenized as (__seen, 0).
+    This makes it important to ensure that dicts are sorted *before* you call
+    normalize_token() on their elements.
+    """
+    v = (1, 2, 3)
+    d1 = {1: v, 2: v}
+    d2 = {2: v, 1: v}
+    assert "__seen" in str(normalize_token(d1))
+    assert check_tokenize(d1) == check_tokenize(d2)
+
+
+def test_tokenize_sorts_set_before_seen_map():
+    """Same as test_tokenize_sorts_dict_before_seen_map, but for sets.
+
+    Note that this test is only meaningful if set insertion order impacts iteration
+    order, which is an implementation detail of the Python interpreter.
+    """
+    v = (1, 2, 3)
+    s1 = {(i, v) for i in range(100)}
+    s2 = {(i, v) for i in reversed(range(100))}
+    assert "__seen" in str(normalize_token(s1))
+    assert check_tokenize(s1) == check_tokenize(s2)
+
+
 def test_tokenize_timedelta():
     assert check_tokenize(datetime.timedelta(days=1)) == check_tokenize(
         datetime.timedelta(days=1)
@@ -1201,3 +1228,100 @@ def test_tokenize_opaque_object_with_buffers():
             self.x = np.array(x)
 
     assert check_tokenize(C([1, 2])) != check_tokenize(C([1, 3]))
+
+
+if not numba:
+
+    class NumbaDummy:
+        def __bool__(self):
+            return False
+
+        def _dummy_decorator(self, *args, **kwargs):
+            def wrapper(func):
+                return func
+
+            return wrapper
+
+        jit = vectorize = guvectorize = _dummy_decorator
+
+    numba = NumbaDummy()
+
+
+@numba.jit(nopython=True)
+def numba_jit(x, y):
+    return x + y
+
+
+@numba.jit("f8(f8, f8)", nopython=True)
+def numba_jit_with_signature(x, y):
+    return x + y
+
+
+@numba.vectorize(nopython=True)
+def numba_vectorize(x, y):
+    return x + y
+
+
+@numba.vectorize("f8(f8, f8)", nopython=True)
+def numba_vectorize_with_signature(x, y):
+    return x + y
+
+
+@numba.guvectorize(["f8,f8,f8[:]"], "(),()->()")
+def numba_guvectorize(x, y, out):
+    out[0] = x + y
+
+
+all_numba_funcs = [
+    numba_jit,
+    numba_jit_with_signature,
+    numba_vectorize,
+    numba_vectorize_with_signature,
+    numba_guvectorize,
+]
+
+
+@pytest.mark.skipif("not numba")
+@pytest.mark.parametrize("func", all_numba_funcs)
+def test_tokenize_numba(func):
+    assert func(1, 2) == 3
+    check_tokenize(func)
+
+
+@pytest.mark.skipif("not numba")
+def test_tokenize_numba_unique_token():
+    tokens = [check_tokenize(func) for func in all_numba_funcs]
+    assert len(tokens) == len(set(tokens))
+
+
+@pytest.mark.skipif("not numba")
+def test_numba_local():
+    @numba.jit(nopython=True)
+    def local_jit(x, y):
+        return x + y
+
+    @numba.jit("f8(f8, f8)", nopython=True)
+    def local_jit_with_signature(x, y):
+        return x + y
+
+    @numba.vectorize(nopython=True)
+    def local_vectorize(x, y):
+        return x + y
+
+    @numba.vectorize("f8(f8, f8)", nopython=True)
+    def local_vectorize_with_signature(x, y):
+        return x + y
+
+    @numba.guvectorize(["f8,f8,f8[:]"], "(),()->()")
+    def local_guvectorize(x, y, out):
+        out[0] = x + y
+
+    all_funcs = [
+        local_jit,
+        local_jit_with_signature,
+        local_vectorize,
+        local_vectorize_with_signature,
+        local_guvectorize,
+    ]
+    tokens = [check_tokenize(func) for func in all_funcs]
+    assert len(tokens) == len(set(tokens))