Repository: dask/dask
Pull request #10925 (2024-02-15 05:42:04)

Add dataframe.query-planning-warning config
- [x] Part of #10917 
- [x] Tests added / passed
- [x] Passes `pre-commit run --all-files`


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 14m 41s :stopwatch: - 3m 36s
 13 031 tests ±0   12 100 :white_check_mark: +1     931 :zzz: ±0  0 :x:  - 1 
161 152 runs  ±0  144 615 :white_check_mark: +1  16 537 :zzz: ±0  0 :x:  - 1 

Results for commit 62089ae8. ± Comparison against base commit 71b16539.

[test-results]:data:application/gzip;base64,H4sIAOqqzWUC/02Oyw6DIBBFf8Ww7oKZkVd/piGgCamPBmFl+u9FtOrynLk5mZX1YegW9mxAPBq25JAu8jnaFOZpY5AaiirXVO/E6eTXkp3bJALnl3yHT5HmtuttGIq6Nl2MczxMzFMtSwCBB5/ltpX1pV3uZZCC1N/d0pXvZTePY0gFmESuje00GlCaPKleSPLKOQ3KkjPotEYkz74/ZAOuKhkBAAA=


Patch:
diff --git a/dask/base.py b/dask/base.py
index f32095d55..aaeec20e4 100644
--- a/dask/base.py
+++ b/dask/base.py
@@ -1005,9 +1005,13 @@ def persist(*args, traverse=True, optimize_graph=True, scheduler=None, **kwargs)
 ############
 # Tokenize #
 ############
+class TokenizationError(RuntimeError):
+    pass
 
 
-def tokenize(*args, **kwargs):
+def tokenize(
+    *args: object, ensure_deterministic: bool | None = None, **kwargs: object
+) -> str:
     """Deterministic token
 
     >>> tokenize([1, 2, '3'])
@@ -1015,19 +1019,89 @@ def tokenize(*args, **kwargs):
 
     >>> tokenize('Hello') == tokenize('Hello')
     True
+
+    Parameters
+    ----------
+    args, kwargs:
+        objects to tokenize
+    ensure_deterministic: bool, optional
+        If True, raise TokenizationError if the objects cannot be deterministically
+        tokenized, e.g. two identical objects will return different tokens.
+        Defaults to the `tokenize.ensure-deterministic` configuration parameter.
     """
-    tok = _seen.set({})
-    try:
-        token = _normalize_seq_func(args)
+    with _seen_ctx(reset=True), _ensure_deterministic_ctx(ensure_deterministic):
+        token: object = _normalize_seq_func(args)
         if kwargs:
             token = token, _normalize_seq_func(sorted(kwargs.items()))
-    finally:
-        _seen.reset(tok)
 
     # Pass `usedforsecurity=False` to support FIPS builds of Python
     return hashlib.md5(str(token).encode(), usedforsecurity=False).hexdigest()
 
 
+# tokenize.ensure-deterministic flag, potentially overridden by tokenize()
+_ensure_deterministic: ContextVar[bool] = ContextVar("_ensure_deterministic")
+
+# Circular reference breaker used by _normalize_seq_func.
+# This variable is recreated anew every time you call tokenize(). Note that this means
+# that you could call tokenize() from inside tokenize() and they would be fully
+# independent.
+#
+# It is a map of {id(obj): (<first seen incremental int>, obj)} which causes an object
+# to be tokenized as ("__seen", <incremental>) the second time it's encountered while
+# traversing collections. A strong reference to the object is stored in the context to
+# prevent ids from being reused by different objects.
+_seen: ContextVar[dict[int, tuple[int, object]]] = ContextVar("_seen")
+
+
+@contextmanager
+def _ensure_deterministic_ctx(ensure_deterministic: bool | None) -> Iterator[bool]:
+    try:
+        ensure_deterministic = _ensure_deterministic.get()
+        # There's a call of tokenize() higher up in the stack
+        tok = None
+    except LookupError:
+        # Outermost tokenize(), or normalize_token() was called directly
+        if ensure_deterministic is None:
+            ensure_deterministic = config.get("tokenize.ensure-deterministic")
+        tok = _ensure_deterministic.set(ensure_deterministic)
+
+    try:
+        yield ensure_deterministic
+    finally:
+        if tok:
+            tok.var.reset(tok)
+
+
+def _maybe_raise_nondeterministic(msg: str) -> None:
+    with _ensure_deterministic_ctx(None) as ensure_deterministic:
+        if ensure_deterministic:
+            raise TokenizationError(msg)
+
+
+@contextmanager
+def _seen_ctx(reset: bool) -> Iterator[dict[int, tuple[int, object]]]:
+    if reset:
+        # It is important to reset the token on tokenize() to avoid artifacts when
+        # it is called recursively
+        seen: dict[int, tuple[int, object]] = {}
+        tok = _seen.set(seen)
+    else:
+        try:
+            seen = _seen.get()
+            tok = None
+        except LookupError:
+            # This is for debug only, for when normalize_token is called outside of
+            # tokenize()
+            seen = {}
+            tok = _seen.set(seen)
+
+    try:
+        yield seen
+    finally:
+        if tok:
+            tok.var.reset(tok)
+
+
 normalize_token = Dispatch()
 normalize_token.register(
     (
@@ -1068,30 +1142,8 @@ def normalize_set(s):
     return "set", _normalize_seq_func(sorted(s, key=str))
 
 
-# Circular reference breaker used by _normalize_seq_func.
-# This variable is recreated anew every time you call tokenize(). Note that this means
-# that you could call tokenize() from inside tokenize() and they would be fully
-# independent.
-#
-# It is a map of {id(obj): (<first seen incremental int>, obj)} which causes an object
-# to be tokenized as ("__seen", <incremental>) the second time it's encountered while
-# traversing collections. A strong reference to the object is stored in the context to
-# prevent ids from being reused by different objects.
-_seen: ContextVar[dict[int, tuple[int, object]]] = ContextVar("_seen")
-
-
 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)
-
-    try:
+    with _seen_ctx(reset=False) as seen:
         out = []
         for item in seq:
             if isinstance(item, (str, bytes, int, float, bool, type(None))):
@@ -1107,9 +1159,6 @@ def _normalize_seq_func(seq: Iterable[object]) -> list[object]:
                 item = normalize_token(item)
             out.append(item)
         return out
-    finally:
-        if tok is not None:
-            _seen.reset(tok)
 
 
 @normalize_token.register((tuple, list))
@@ -1147,9 +1196,6 @@ def normalize_builtin_function_or_method(func):
         return normalize_object(func)
 
 
-_seen_objects = set()
-
-
 @normalize_token.register(object)
 def normalize_object(o):
     method = getattr(o, "__dask_tokenize__", None)
@@ -1157,15 +1203,7 @@ def normalize_object(o):
         return method()
 
     if type(o) is object:
-        if config.get("tokenize.ensure-deterministic"):
-            raise RuntimeError(
-                "object() cannot be deterministically hashed. See "
-                "https://docs.dask.org/en/latest/custom-collections.html#implementing-deterministic-hashing "
-                "for more information."
-            )
-        # Idempotent, but not deterministic. Make sure that the id is not reused.
-        _seen_objects.add(o)
-        return "object", id(o)
+        return _normalize_pure_object(o)
 
     if dataclasses.is_dataclass(o) and not isinstance(o, type):
         return _normalize_dataclass(o)
@@ -1173,15 +1211,28 @@ def normalize_object(o):
     try:
         return _normalize_pickle(o)
     except Exception:
-        if config.get("tokenize.ensure-deterministic"):
-            raise RuntimeError(
-                f"Object {o!r} cannot be deterministically hashed. See "
-                "https://docs.dask.org/en/latest/custom-collections.html#implementing-deterministic-hashing "
-                "for more information."
-            )
+        _maybe_raise_nondeterministic(
+            f"Object {o!r} cannot be deterministically hashed. See "
+            "https://docs.dask.org/en/latest/custom-collections.html#implementing-deterministic-hashing "
+            "for more information."
+        )
         return uuid.uuid4().hex
 
 
+_seen_objects = set()
+
+
+def _normalize_pure_object(o: object) -> tuple[str, int]:
+    _maybe_raise_nondeterministic(
+        "object() cannot be deterministically hashed. See "
+        "https://docs.dask.org/en/latest/custom-collections.html#implementing-deterministic-hashing "
+        "for more information."
+    )
+    # Idempotent, but not deterministic. Make sure that the id is not reused.
+    _seen_objects.add(o)
+    return "object", id(o)
+
+
 def _normalize_pickle(o: object) -> tuple:
     buffers: list[pickle.PickleBuffer] = []
     pik: bytes | None
@@ -1303,12 +1354,11 @@ def register_numpy():
         try:
             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"
-                )
+            _maybe_raise_nondeterministic(
+                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
 
 
diff --git a/dask/dask-schema.yaml b/dask/dask-schema.yaml
index a9ce92e80..20412201c 100644
--- a/dask/dask-schema.yaml
+++ b/dask/dask-schema.yaml
@@ -93,6 +93,10 @@ properties:
         type: [boolean]
         description: |
           Whether to use query planning.
+      query-planning-warning:
+        type: [boolean]
+        description: |
+          Whether to warn about future use of query planning.
 
   array:
     type: object
diff --git a/dask/dask.yaml b/dask/dask.yaml
index ac6030061..3b264f152 100644
--- a/dask/dask.yaml
+++ b/dask/dask.yaml
@@ -16,6 +16,7 @@ dataframe:
     metadata-task-size-remote: 1  # Number of files per remote metadata-processing task
   convert-string: null  # Whether to convert string-like data to pyarrow strings
   query-planning: false  # Whether to use dask-expr
+  query-planning-warning: true  # Warn about future query optimization change
 
 array:
   backend: "numpy"  # Backend array library for input IO and data creation
diff --git a/dask/dataframe/__init__.py b/dask/dataframe/__init__.py
index fca52eba9..b68398ab8 100644
--- a/dask/dataframe/__init__.py
+++ b/dask/dataframe/__init__.py
@@ -159,9 +159,11 @@ else:
             '  python -m pip install "dask[dataframe]" --upgrade  # or python -m pip install'
         )
         raise ImportError(msg) from e
-    warnings.warn(
-        """The current Dask DataFrame implementation is deprecated. 
-In a future release, Dask DataFrame will use new implementation that
+
+    if dask.config.get("dataframe.query-planning-warning"):
+        warnings.warn(
+            """The current Dask DataFrame implementation is deprecated. 
+In a future release, Dask DataFrame will use a new implementation that
 contains several improvements including a logical query planning.
 The user-facing DataFrame API will remain unchanged.
 
@@ -181,10 +183,14 @@ https://docs.dask.org/en/stable/dask-expr-api.html
 
 Any feedback can be reported on the Dask issue tracker
 https://github.com/dask/dask/issues 
+
+To disable this warning in the future, set dask config:
+
+    >>> dask.config.set({'dataframe.query-planning-warning': False})
 """,
-        DeprecationWarning,
-        stacklevel=2,
-    )
+            DeprecationWarning,
+            stacklevel=2,
+        )
 
 
 from dask.dataframe._testing import test_dataframe


Test patch:
diff --git a/dask/tests/test_tokenize.py b/dask/tests/test_tokenize.py
index adce0573a..a7f6060bd 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 normalize_token, tokenize
+from dask.base import TokenizationError, normalize_token, tokenize
 from dask.core import literal
 from dask.utils import tmpfile
 from dask.utils_test import import_or_none
@@ -32,6 +32,24 @@ pd = import_or_none("pandas")
 numba = import_or_none("numba")
 
 
+@pytest.fixture(autouse=True)
+def check_contextvars():
+    """Test that tokenize() and normalize_token() properly clean up context
+    variables at all times
+    """
+    from dask.base import _ensure_deterministic, _seen
+
+    with pytest.raises(LookupError):
+        _ensure_deterministic.get()
+    with pytest.raises(LookupError):
+        _seen.get()
+    yield
+    with pytest.raises(LookupError):
+        _ensure_deterministic.get()
+    with pytest.raises(LookupError):
+        _seen.get()
+
+
 def check_tokenize(*args, **kwargs):
     with dask.config.set({"tokenize.ensure-deterministic": True}):
         before = tokenize(*args, **kwargs)
@@ -264,13 +282,24 @@ def test_normalize_numpy_ufunc_unserializable():
     # Make a ufunc that isn't in the numpy namespace and can't be serialized
     inc = np.frompyfunc(lambda x: x + 1, 1, 1)
     with dask.config.set({"tokenize.ensure-deterministic": False}):
+        # Not idempotent
         assert tokenize(inc) != tokenize(inc)
+        # You can call normalize_token directly
+        assert normalize_token(inc) != normalize_token(inc)
+
     with dask.config.set({"tokenize.ensure-deterministic": True}):
         with pytest.raises(
-            RuntimeError, match=r"Cannot tokenize.*dask\.array\.ufunc.*instead"
+            TokenizationError, match=r"Cannot tokenize.*dask\.array\.ufunc.*instead"
         ):
             tokenize(inc)
 
+    # Test env override
+    assert tokenize(inc, ensure_deterministic=False) != tokenize(
+        inc, ensure_deterministic=False
+    )
+    with pytest.raises(TokenizationError, match="Cannot tokenize"):
+        tokenize(inc, ensure_deterministic=True)
+
 
 def test_normalize_object_unserializable():
     class C:
@@ -280,11 +309,24 @@ def test_normalize_object_unserializable():
     c = C()
 
     with dask.config.set({"tokenize.ensure-deterministic": False}):
+        # Not idempotent
         assert tokenize(c) != tokenize(c)
+        # You can call normalize_token directly
+        assert normalize_token(c) != normalize_token(c)
+
     with dask.config.set({"tokenize.ensure-deterministic": True}):
-        with pytest.raises(RuntimeError, match="cannot be deterministically hashed"):
+        with pytest.raises(
+            TokenizationError, match="cannot be deterministically hashed"
+        ):
             tokenize(c)
 
+    # Test env override
+    assert tokenize(c, ensure_deterministic=False) != tokenize(
+        c, ensure_deterministic=False
+    )
+    with pytest.raises(TokenizationError, match="cannot be deterministically hashed"):
+        tokenize(c, ensure_deterministic=True)
+
 
 def test_tokenize_partial_func_args_kwargs_consistent():
     f = partial(f3, f2, c=f1)
@@ -308,18 +350,50 @@ def test_normalize_base():
 
 
 def test_tokenize_object():
-    # object() tokenization is idempotent
-    o = object()
-    assert tokenize(o) == tokenize(o)
+    with dask.config.set({"tokenize.ensure-deterministic": False}):
+        # object() tokenization is idempotent...
+        o = object()
+        assert tokenize(o) == tokenize(o)
+        # ...but not deterministic
+        assert tokenize(object()) != tokenize(object())
+
+        # Two objects don't tokenize to the same token even if their pickle output is
+        # identical. Stress id reuse by creating and dereferencing many objects in quick
+        # succession.
+        assert len({tokenize(object()) for _ in range(100)}) == 100
+
+        # You can call normalize_token even if the _ensure_deterministic context
+        # variable hasn't been set
+        assert normalize_token(o) == normalize_token(o)
 
     with dask.config.set({"tokenize.ensure-deterministic": True}):
-        with pytest.raises(RuntimeError, match="deterministic"):
+        with pytest.raises(TokenizationError, match="deterministic"):
             tokenize(o)
+        with pytest.raises(TokenizationError, match="deterministic"):
+            normalize_token(o)
+
+    # Test env override
+    assert tokenize(o, ensure_deterministic=False) == tokenize(
+        o, ensure_deterministic=False
+    )
+    with pytest.raises(TokenizationError, match="deterministic"):
+        tokenize(o, ensure_deterministic=True)
 
-    # Two objects don't tokenize to the same token even if their pickle output is
-    # identical. Stress id reuse by creating and dereferencing many objects in quick
-    # succession.
-    assert len({tokenize(object()) for _ in range(100)}) == 100
+
+def nested_tokenize_ensure_deterministic():
+    """Test that the ensure_deterministic override is not lost if tokenize() is
+    called recursively
+    """
+
+    class C:
+        def __dask_tokenize__(self):
+            return tokenize(object())
+
+    assert tokenize(C(), ensure_deterministic=False) != tokenize(
+        C(), ensure_deterministic=False
+    )
+    with pytest.raises(TokenizationError):
+        tokenize(C())
 
 
 _GLOBAL = 1
@@ -757,6 +831,28 @@ def test_tokenize_sequences():
     )
 
 
+def test_nested_tokenize_seen():
+    """Test that calling tokenize() recursively doesn't alter the output due to
+    memoization of already-seen objects
+    """
+    o = [1, 2, 3]
+
+    class C:
+        def __init__(self, x):
+            self.x = x
+            self.tok = None
+
+        def __dask_tokenize__(self):
+            if not self.tok:
+                self.tok = tokenize(self.x)
+            return self.tok
+
+    c1, c2 = C(o), C(o)
+    check_tokenize(o, c1, o)
+    assert c1.tok
+    assert check_tokenize(c1) == check_tokenize(c2)
+
+
 def test_tokenize_dict():
     # 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})