Repository: dask/dask
Pull request #10908 (2024-02-08 13:46:15)

Add configuration for leftsemi merges in dask-expr
- [ ] Closes #xxxx
- [ ] Tests added / passed
- [ ] Passes `pre-commit run --all-files`

precursor for the dask-expr pr


Patch:
diff --git a/dask/base.py b/dask/base.py
index ae4745e29..a82ee50dd 100644
--- a/dask/base.py
+++ b/dask/base.py
@@ -9,10 +9,11 @@ import pathlib
 import pickle
 import random
 import threading
+import types
 import uuid
 import warnings
 from collections import OrderedDict
-from collections.abc import Callable, Hashable, Iterator, Mapping
+from collections.abc import Callable, Hashable, Iterable, Iterator, Mapping
 from concurrent.futures import Executor
 from contextlib import contextmanager, suppress
 from contextvars import ContextVar
@@ -20,7 +21,6 @@ from enum import Enum
 from functools import partial
 from numbers import Integral, Number
 from operator import getitem
-from types import MappingProxyType
 from typing import Any, Literal, TypeVar
 
 from tlz import curry, groupby, identity, merge
@@ -1060,8 +1060,7 @@ def normalize_type(typ):
     return typ.__name__, typ.__module__
 
 
-@normalize_token.register(MappingProxyType)
-@normalize_token.register(dict)
+@normalize_token.register((dict, types.MappingProxyType))
 def normalize_dict(d):
     return "dict", _normalize_seq_func(sorted(d.items(), key=str))
 
@@ -1073,6 +1072,9 @@ def normalize_ordered_dict(d):
 
 @normalize_token.register(set)
 def normalize_set(s):
+    # Note: in some Python version / OS combinations, set order changes every
+    # time you recreate the set (even within the same interpreter).
+    # In most other cases, set ordering is consistent within the same interpreter.
     return "set", _normalize_seq_func(sorted(s, key=str))
 
 
@@ -1088,24 +1090,36 @@ def normalize_set(s):
 _seen: ContextVar[dict[int, tuple[int, object]]] = ContextVar("_seen")
 
 
-def _normalize_seq_func(seq):
-    seen = _seen.get()
-    out = []
+def _normalize_seq_func(seq: Iterable[object]) -> list[object]:
+    try:
+        seen = _seen.get()
+        tok = None
+    except LookupError:
+        # This is for debug only, for when normalize_token is called outside of
+        # tokenize(). It is important to reset the token on tokenize to avoid artifacts
+        # when tokenize() is called recursively.
+        seen = {}
+        tok = _seen.set(seen)
 
-    for item in seq:
-        if isinstance(item, (str, bytes, int, float, bool, type(None))):
-            # Basic data type. This is just for performance and compactness of the
-            # output. It doesn't need to be a comprehensive list.
-            norm_item = item
-        elif id(item) in seen:
-            # May or may not be a circular recursion. Maybe just a double reference.
-            seen_when, _ = seen[id(item)]
-            norm_item = "__seen", seen_when
-        else:
-            seen[id(item)] = len(seen), item
-            norm_item = normalize_token(item)
-        out.append(norm_item)
-    return out
+    try:
+        out = []
+        for item in seq:
+            if isinstance(item, (str, bytes, int, float, bool, type(None))):
+                # Basic data type. This is just for performance and compactness of the
+                # output. It doesn't need to be a comprehensive list.
+                pass
+            elif id(item) in seen:
+                # May or may not be a circular recursion. Maybe just a double reference.
+                seen_when, _ = seen[id(item)]
+                item = "__seen", seen_when
+            else:
+                seen[id(item)] = len(seen), item
+                item = normalize_token(item)
+            out.append(item)
+        return out
+    finally:
+        if tok is not None:
+            _seen.reset(tok)
 
 
 @normalize_token.register((tuple, list))
@@ -1183,9 +1197,6 @@ def normalize_callable(func: Callable) -> Callable | tuple | str | bytes:
         return _normalize_callable(func)
 
 
-normalize_function = normalize_callable
-
-
 def _normalize_callable(func: Callable) -> tuple | str | bytes:
     if isinstance(func, Compose):
         first = getattr(func, "first", None)
@@ -1366,7 +1377,7 @@ def register_numpy():
                 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)
+        return data, x.dtype, x.shape
 
     @normalize_token.register(np.matrix)
     def normalize_matrix(x):
diff --git a/dask/dataframe/multi.py b/dask/dataframe/multi.py
index 0175035ef..5c5576c89 100644
--- a/dask/dataframe/multi.py
+++ b/dask/dataframe/multi.py
@@ -287,6 +287,12 @@ def merge_chunk(
                 else:
                     rhs = rhs.assign(**{col: right.astype(dtype)})
 
+    if len(args) and args[0] == "leftsemi" or kwargs.get("how", None) == "leftsemi":
+        rhs = rhs.drop_duplicates()
+        if len(args):
+            args[0] = "inner"
+        else:
+            kwargs["how"] = "inner"
     out = lhs.merge(rhs, *args, **kwargs)
 
     # Workaround for pandas bug where if the left frame of a merge operation is
diff --git a/dask/optimization.py b/dask/optimization.py
index 4ec7957d9..68944594c 100644
--- a/dask/optimization.py
+++ b/dask/optimization.py
@@ -2,10 +2,12 @@ from __future__ import annotations
 
 import math
 import numbers
-import uuid
+from collections.abc import Iterable
 from enum import Enum
+from typing import Any
 
 from dask import config, core, utils
+from dask.base import normalize_token, tokenize
 from dask.core import (
     flatten,
     get_dependencies,
@@ -15,6 +17,7 @@ from dask.core import (
     subs,
     toposort,
 )
+from dask.typing import Graph, Key
 
 
 def cull(dsk, keys):
@@ -965,20 +968,26 @@ class SubgraphCallable:
         The name to use for the function.
     """
 
-    __slots__ = ("dsk", "outkey", "inkeys", "name")
+    dsk: Graph
+    outkey: Key
+    inkeys: tuple[Key, ...]
+    name: str
+    __slots__ = tuple(__annotations__)
 
-    def __init__(self, dsk, outkey, inkeys, name=None):
+    def __init__(
+        self, dsk: Graph, outkey: Key, inkeys: Iterable[Key], name: str | None = None
+    ):
         self.dsk = dsk
         self.outkey = outkey
-        self.inkeys = inkeys
+        self.inkeys = tuple(inkeys)
         if name is None:
-            name = f"subgraph_callable-{uuid.uuid4()}"
+            name = "subgraph_callable-" + tokenize(dsk, outkey, self.inkeys)
         self.name = name
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         return self.name
 
-    def __eq__(self, other):
+    def __eq__(self, other: Any) -> bool:
         return (
             type(self) is type(other)
             and self.name == other.name
@@ -986,16 +995,22 @@ class SubgraphCallable:
             and set(self.inkeys) == set(other.inkeys)
         )
 
-    def __ne__(self, other):
-        return not (self == other)
-
-    def __call__(self, *args):
+    def __call__(self, *args: Any) -> Any:
         if not len(args) == len(self.inkeys):
             raise ValueError("Expected %d args, got %d" % (len(self.inkeys), len(args)))
         return core.get(self.dsk, self.outkey, dict(zip(self.inkeys, args)))
 
-    def __reduce__(self):
-        return (SubgraphCallable, (self.dsk, self.outkey, self.inkeys, self.name))
+    def __reduce__(self) -> tuple:
+        return SubgraphCallable, (self.dsk, self.outkey, self.inkeys, self.name)
 
-    def __hash__(self):
-        return hash(tuple((self.outkey, frozenset(self.inkeys), self.name)))
+    def __hash__(self) -> int:
+        return hash((self.outkey, frozenset(self.inkeys), self.name))
+
+    def __dask_tokenize__(self) -> object:
+        return (
+            "SubgraphCallable",
+            normalize_token(self.dsk),
+            self.outkey,
+            self.inkeys,
+            self.name,
+        )


Test patch:
diff --git a/dask/tests/test_optimization.py b/dask/tests/test_optimization.py
index 397f6fc23..925c4f93a 100644
--- a/dask/tests/test_optimization.py
+++ b/dask/tests/test_optimization.py
@@ -1228,26 +1228,53 @@ def test_SubgraphCallable_eq():
     dsk2 = {"a": (inc, 0), "b": (inc, "a"), "c": (add, "d", "e")}
     f1 = SubgraphCallable(dsk1, "c", ["d", "e"])
     f2 = SubgraphCallable(dsk2, "c", ["d", "e"])
+
     # Different graphs must compare unequal (when no name given)
     assert f1 != f2
+    assert hash(f1) != hash(f2)
+    assert tokenize(f1) != tokenize(f2)
+
+    # Identical graphs must compare as equal and generate the same name
+    dsk1b = {"a": 1, "b": 2, "c": (add, "d", "e")}
+    f1b = SubgraphCallable(dsk1b, "c", ["d", "e"])
+    assert f1b.name == f1.name
+    assert f1 == f1b
+    assert hash(f1) == hash(f1b)
+    assert tokenize(f1) == tokenize(f1b)
 
     # Different inputs must compare unequal
     f3 = SubgraphCallable(dsk2, "c", ["d", "f"], name=f1.name)
     assert f3 != f1
+    assert hash(f3) != hash(f1)
+    assert tokenize(f3) != tokenize(f1)
 
     # Different outputs must compare unequal
     f4 = SubgraphCallable(dsk2, "a", ["d", "e"], name=f1.name)
     assert f4 != f1
+    assert hash(f4) != hash(f1)
+    assert tokenize(f4) != tokenize(f1)
 
-    # Reordering the inputs must not prevent equality
+    # Reordering the inputs must not prevent equality...
     f5 = SubgraphCallable(dsk1, "c", ["e", "d"], name=f1.name)
     assert f1 == f5
     assert hash(f1) == hash(f5)
+    # ... except for tokenize, which takes into consideration the whole dsk
+    assert tokenize(f1) != tokenize(f5)
 
     # Explicitly named graphs with different names must be unequal
-    unnamed1 = SubgraphCallable(dsk1, "c", ["d", "e"], name="first")
-    unnamed2 = SubgraphCallable(dsk1, "c", ["d", "e"], name="second")
-    assert unnamed1 != unnamed2
+    f6 = SubgraphCallable(dsk1, "c", ["d", "e"], name="first")
+    f7 = SubgraphCallable(dsk1, "c", ["d", "e"], name="second")
+    assert f6 != f7
+    assert hash(f6) != hash(f7)
+    assert tokenize(f6) != tokenize(f7)
+
+    # Equality and hashing don't rely on tokenizing the full dsk,
+    # because tokenization is not guaranteed to be deterministic
+    f8 = SubgraphCallable({"a": object(), "b": "a"}, "b", [], name="n")
+    f9 = SubgraphCallable({"a": object(), "b": "a"}, "b", [], name="n")
+    assert f8 == f9
+    assert hash(f8) == hash(f9)
+    assert tokenize(f8) != tokenize(f9)
 
 
 def test_fuse_subgraphs(compare_subgraph_callables):
diff --git a/dask/tests/test_tokenize.py b/dask/tests/test_tokenize.py
index 59d9a2c6a..aff284103 100644
--- a/dask/tests/test_tokenize.py
+++ b/dask/tests/test_tokenize.py
@@ -18,7 +18,7 @@ import pytest
 from tlz import compose, curry, partial
 
 import dask
-from dask.base import function_cache, normalize_function, normalize_token, tokenize
+from dask.base import function_cache, normalize_callable, normalize_token, tokenize
 from dask.core import literal
 from dask.utils import tmpfile
 from dask.utils_test import import_or_none
@@ -31,39 +31,6 @@ pa = import_or_none("pyarrow")
 pd = import_or_none("pandas")
 
 
-def f1(a, b, c=1):
-    pass
-
-
-def f2(a, b=1, c=2):
-    pass
-
-
-def f3(a):
-    pass
-
-
-def test_normalize_function():
-    assert normalize_function(f2)
-
-    assert normalize_function(lambda a: a)
-
-    assert normalize_function(partial(f2, b=2)) == normalize_function(partial(f2, b=2))
-
-    assert normalize_function(partial(f2, b=2)) != normalize_function(partial(f2, b=3))
-
-    assert normalize_function(partial(f1, b=2)) != normalize_function(partial(f2, b=2))
-
-    assert normalize_function(compose(f2, f3)) == normalize_function(compose(f2, f3))
-
-    assert normalize_function(compose(f2, f3)) != normalize_function(compose(f2, f1))
-
-    assert normalize_function(curry(f2)) == normalize_function(curry(f2))
-    assert normalize_function(curry(f2)) != normalize_function(curry(f1))
-    assert normalize_function(curry(f2, b=1)) == normalize_function(curry(f2, b=1))
-    assert normalize_function(curry(f2, b=1)) != normalize_function(curry(f2, b=2))
-
-
 def _clear_function_cache():
     from dask.base import function_cache, function_cache_lock
 
@@ -129,7 +96,7 @@ def test_check_tokenize():
         def __init__(self):
             self.tok = random.random()
 
-        def __reduce_ex__(self, protocol):
+        def __reduce__(self):
             return A, ()
 
         def __dask_tokenize__(self):
@@ -381,12 +348,28 @@ def test_tokenize_local_functions():
     assert len(set(tokens)) == len(all_funcs)
 
 
-def my_func(a, b, c=1):
-    return a + b + c
+def f1(a, b, c=1):
+    pass
+
+
+def f2(a, b=1, c=2):
+    pass
+
+
+def f3(a):
+    pass
 
 
 def test_tokenize_callable():
-    check_tokenize(my_func)
+    assert check_tokenize(f1) != check_tokenize(f2)
+
+
+def test_tokenize_composite_functions():
+    assert check_tokenize(partial(f2, b=2)) != check_tokenize(partial(f2, b=3))
+    assert check_tokenize(partial(f1, b=2)) != check_tokenize(partial(f2, b=2))
+    assert check_tokenize(compose(f2, f3)) != check_tokenize(compose(f2, f1))
+    assert check_tokenize(curry(f2)) != check_tokenize(curry(f1))
+    assert check_tokenize(curry(f2, b=1)) != check_tokenize(curry(f2, b=2))
 
 
 @pytest.mark.skipif("not pd")
@@ -572,25 +555,35 @@ def test_staticmethods():
     assert check_tokenize(a.class_method) != check_tokenize(c.class_method)
 
 
-@pytest.mark.skipif("not np")
 def test_tokenize_sequences():
     assert check_tokenize([1]) != check_tokenize([2])
     assert check_tokenize([1]) != check_tokenize((1,))
     assert check_tokenize([1]) == check_tokenize([1])
 
-    x = np.arange(2000)  # long enough to drop information in repr
-    y = np.arange(2000)
-    y[1000] = 0  # middle isn't printed in repr
-    assert check_tokenize([x]) != check_tokenize([y])
+    # You can call normalize_token directly.
+    # Repeated objects are memoized.
+    x = (1, 2)
+    y = [x, x, [x, (2, 3)]]
+    assert normalize_token(y) == (
+        "list",
+        [
+            ("tuple", [1, 2]),
+            ("__seen", 0),
+            ("list", [("__seen", 0), ("tuple", [2, 3])]),
+        ],
+    )
 
 
 def test_tokenize_dict():
-    assert check_tokenize({"x": 1, 1: "x"}) == check_tokenize({"x": 1, 1: "x"})
+    # Insertion order is ignored. Keys can be an unsortable mix of types.
+    assert check_tokenize({"x": 1, 1: "x"}) == check_tokenize({1: "x", "x": 1})
+    assert check_tokenize({"x": 1, 1: "x"}) != check_tokenize({"x": 1, 2: "x"})
+    assert check_tokenize({"x": 1, 1: "x"}) != check_tokenize({"x": 2, 1: "x"})
 
 
 def test_tokenize_set():
     assert check_tokenize({1, 2, "x", (1, "x")}) == check_tokenize(
-        {1, 2, "x", (1, "x")}
+        {2, "x", (1, "x"), 1}
     )
 
 
@@ -685,6 +678,14 @@ def test_tokenize_range(other):
     assert check_tokenize(range(5, 10, 2)) != check_tokenize(range(*other))
 
 
+@pytest.mark.skipif("not np")
+def test_tokenize_numpy_array():
+    x = np.arange(2000)  # long enough to drop information in repr
+    y = np.arange(2000)
+    y[1000] = 0  # middle isn't printed in repr
+    assert check_tokenize([x]) != check_tokenize([y])
+
+
 @pytest.mark.skipif("not np")
 def test_tokenize_object_array_with_nans():
     a = np.array(["foo", "Jos\xe9", np.nan], dtype="O")
@@ -749,10 +750,7 @@ def test_tokenize_circular_recursion():
     b = [1, 3]
     b[0] = b
 
-    def copy(x):
-        return pickle.loads(pickle.dumps(x))
-
-    assert check_tokenize(a, copy=copy) != check_tokenize(b, copy=copy)
+    assert check_tokenize(a) != check_tokenize(b)
 
     # Different circular recursions tokenize differently
     c = [[], []]
@@ -762,12 +760,12 @@ def test_tokenize_circular_recursion():
     d = [[], []]
     d[0].append(d[1])
     d[1].append(d[0])
-    assert check_tokenize(c, copy=copy) != check_tokenize(d, copy=copy)
+    assert check_tokenize(c) != check_tokenize(d)
 
     # For dicts, the dict itself is not passed to _normalize_seq_func
     e = {}
     e[0] = e
-    check_tokenize(e, copy=copy)
+    check_tokenize(e)
 
 
 @pytest.mark.parametrize(
@@ -898,7 +896,7 @@ def test_tokenize_functions_main():
 def test_normalize_function_limited_size():
     _clear_function_cache()
     for _ in range(1000):
-        normalize_function(lambda x: x)
+        normalize_callable(lambda x: x)
     assert 50 < len(function_cache) < 600
 
 
@@ -984,13 +982,11 @@ def test_tokenize_random_functions_with_state_numpy():
 def test_tokenize_pyarrow_datatypes_simple():
     a = pa.int64()
     b = pa.float64()
-    assert tokenize(a) == tokenize(a)
-    assert tokenize(a) != tokenize(b)
+    assert check_tokenize(a) != check_tokenize(b)
 
 
 @pytest.mark.skipif("not pa")
 def test_tokenize_pyarrow_datatypes_complex():
     a = pa.struct({"x": pa.int32(), "y": pa.string()})
     b = pa.struct({"x": pa.float64(), "y": pa.int16()})
-    assert tokenize(a) == tokenize(a)
-    assert tokenize(a) != tokenize(b)
+    assert check_tokenize(a) != check_tokenize(b)