Repository: dask/dask
Pull request #10817 (2024-01-17 15:02:36)

Implement deterministic tokenization for hlg
- [ ] Closes #xxxx
- [ ] Tests added / passed
- [ ] Passes `pre-commit run --all-files`


This is something we would need for dask-expr, is the approach in this PR sensible?

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 26m 18s :stopwatch: - 15m 36s
 12 936 tests + 1   11 669 :white_check_mark: ± 0     929 :zzz: ±0  338 :x: +1 
160 272 runs  +15  143 016 :white_check_mark: +11  16 918 :zzz: +3  338 :x: +1 

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

Results for commit e9ab756d. ± Comparison against base commit 91093fdc.

[test-results]:data:application/gzip;base64,H4sIAJ/0p2UC/1XMyw6DIBCF4VcxrLsARkfoyzTCYELqpUFYNX33jrZquzzfSf6n6OMQFnGtVHOpxFJiPheV1OU4T+vW0BomfvP2awu479tSvF9RIdoT7/HBaPVJfRcHJoAjdQspzYlNsqQybW2UutXffbRrkAoP/LQVWmV2+4tv8tv28zjGzEME27m2QaqtkoC+IeMIAhmNIRgnEQM54l+83lVKCGkdAQAA


Patch:
diff --git a/dask/dataframe/__init__.py b/dask/dataframe/__init__.py
index b247c4b76..0f64a152d 100644
--- a/dask/dataframe/__init__.py
+++ b/dask/dataframe/__init__.py
@@ -14,67 +14,82 @@ def _dask_expr_enabled() -> bool:
 
 
 if _dask_expr_enabled():
-    from dask_expr import (
-        DataFrame,
-        Index,
-        Series,
-        concat,
-        from_array,
-        from_dask_array,
-        from_dask_dataframe,
-        from_delayed,
-        from_dict,
-        from_graph,
-        from_map,
-        from_pandas,
-        get_dummies,
-        isna,
-        map_overlap,
-        map_partitions,
-        merge,
-        merge_asof,
-        pivot_table,
-        read_csv,
-        read_hdf,
-        read_json,
-        read_orc,
-        read_parquet,
-        read_sql,
-        read_sql_query,
-        read_sql_table,
-        read_table,
-        repartition,
-        to_bag,
-        to_csv,
-        to_datetime,
-        to_hdf,
-        to_json,
-        to_numeric,
-        to_orc,
-        to_parquet,
-        to_records,
-        to_sql,
-        to_timedelta,
-    )
+    try:
+        from dask_expr import (
+            DataFrame,
+            Index,
+            Series,
+            concat,
+            from_array,
+            from_dask_array,
+            from_dask_dataframe,
+            from_delayed,
+            from_dict,
+            from_graph,
+            from_map,
+            from_pandas,
+            get_dummies,
+            isna,
+            map_overlap,
+            map_partitions,
+            merge,
+            merge_asof,
+            pivot_table,
+            read_csv,
+            read_hdf,
+            read_json,
+            read_orc,
+            read_parquet,
+            read_sql,
+            read_sql_query,
+            read_sql_table,
+            read_table,
+            repartition,
+            to_bag,
+            to_csv,
+            to_datetime,
+            to_hdf,
+            to_json,
+            to_numeric,
+            to_orc,
+            to_parquet,
+            to_records,
+            to_sql,
+            to_timedelta,
+        )
+
+        import dask.dataframe._pyarrow_compat
+        from dask.base import compute
+        from dask.dataframe import backends, dispatch
+        from dask.dataframe.io import demo
+        from dask.dataframe.utils import assert_eq
+
+        def raise_not_implemented_error(attr_name):
+            def inner_func(*args, **kwargs):
+                raise NotImplementedError(
+                    f"Function {attr_name} is not implemented for dask-expr."
+                )
+
+            return inner_func
 
-    import dask.dataframe._pyarrow_compat
-    from dask.base import compute
-    from dask.dataframe import backends, dispatch
-    from dask.dataframe.io import demo
-    from dask.dataframe.utils import assert_eq
+        _Frame = raise_not_implemented_error("_Frame")
+        Aggregation = raise_not_implemented_error("Aggregation")
+        read_fwf = raise_not_implemented_error("read_fwf")
+        melt = raise_not_implemented_error("melt")
 
-    def raise_not_implemented_error(attr_name):
-        def inner_func(*args, **kwargs):
-            raise NotImplementedError(
-                f"Function {attr_name} is not implemented for dask-expr."
-            )
+    # Due to the natural circular imports caused from dask-expr
+    # wanting to import things from dask.dataframe, this module's init
+    # can be run multiple times as it walks code trying to import
+    # dask-expr while dask-expr is also trying to import from dask.dataframe
+    # Each time this happens and hits a circular import, we can reload
+    # dask.dataframe to update itself until dask-expr is fully initialized.
+    # TODO: This can go away when dask-expr is merged into dask
+    except ImportError:
+        import importlib
 
-        return inner_func
+        import dask.dataframe as dd
 
-    _Frame = raise_not_implemented_error("_Frame")
-    Aggregation = raise_not_implemented_error("Aggregation")
-    read_fwf = raise_not_implemented_error("read_fwf")
-    melt = raise_not_implemented_error("melt")
+        dd = importlib.reload(dd)
 
 else:
     try:
diff --git a/dask/dataframe/methods.py b/dask/dataframe/methods.py
index 546510ee5..6b82ad912 100644
--- a/dask/dataframe/methods.py
+++ b/dask/dataframe/methods.py
@@ -32,6 +32,7 @@ from dask.dataframe.dispatch import (  # noqa: F401
     union_categoricals,
 )
 from dask.dataframe.utils import is_dataframe_like, is_index_like, is_series_like
+from dask.utils import _deprecated_kwarg
 
 # cuDF may try to import old dispatch functions
 hash_df = hash_object_dispatch
@@ -56,6 +57,7 @@ def iloc(df, cindexer=None):
     return df.iloc[:, cindexer]
 
 
+@_deprecated_kwarg("convert_dtype", None)
 def apply(df, *args, **kwargs):
     with check_convert_dtype_deprecation():
         with check_apply_dataframe_deprecation():
diff --git a/dask/highlevelgraph.py b/dask/highlevelgraph.py
index 0f13212ce..5448019ee 100644
--- a/dask/highlevelgraph.py
+++ b/dask/highlevelgraph.py
@@ -21,7 +21,7 @@ import tlz as toolz
 
 import dask
 from dask import config
-from dask.base import clone_key, flatten, is_dask_collection
+from dask.base import clone_key, flatten, is_dask_collection, normalize_token
 from dask.core import keys_in_tasks, reverse_dict
 from dask.typing import DaskCollection, Graph, Key
 from dask.utils import ensure_dict, import_required, key_split
@@ -993,3 +993,8 @@ def _get_some_layer_name(collection) -> str:
         # collection does not define the optional __dask_layers__ method
         # or it spuriously returns more than one layer
         return str(id(collection))
+
+
+@normalize_token.register(HighLevelGraph)
+def register_highlevelgraph(hlg):
+    return normalize_token(list(hlg.layers.keys()))


Test patch:
diff --git a/.github/workflows/publish-test-results.yaml b/.github/workflows/publish-test-results.yaml
index fb5dfe194..2acb14465 100644
--- a/.github/workflows/publish-test-results.yaml
+++ b/.github/workflows/publish-test-results.yaml
@@ -47,3 +47,4 @@ jobs:
           event_file: artifacts/Event File/event.json
           event_name: ${{ github.event.workflow_run.event }}
           junit_files: artifacts/**/*.xml
+          fail_on: "errors"
diff --git a/dask/dataframe/io/tests/test_parquet.py b/dask/dataframe/io/tests/test_parquet.py
index 6d423b1ba..cfd3b91ae 100644
--- a/dask/dataframe/io/tests/test_parquet.py
+++ b/dask/dataframe/io/tests/test_parquet.py
@@ -2267,7 +2267,6 @@ def test_writing_parquet_with_unknown_kwargs(tmpdir, engine):
         ddf.to_parquet(fn, engine=engine, unknown_key="unknown_value")
 
 
-@pytest.mark.skipif(DASK_EXPR_ENABLED, reason="circular import")
 def test_to_parquet_with_get(tmpdir, engine):
     from dask.multiprocessing import get as mp_get
 
diff --git a/dask/dataframe/tests/test_dataframe.py b/dask/dataframe/tests/test_dataframe.py
index 8ac825885..470ea0088 100644
--- a/dask/dataframe/tests/test_dataframe.py
+++ b/dask/dataframe/tests/test_dataframe.py
@@ -3411,7 +3411,7 @@ def test_apply_convert_dtype(convert_dtype):
     df = pd.DataFrame({"x": [2, 3, 4, 5], "y": [10, 20, 30, 40]})
     ddf = dd.from_pandas(df, npartitions=2)
     kwargs = {} if convert_dtype is None else {"convert_dtype": convert_dtype}
-    should_warn = PANDAS_GE_210 and convert_dtype is not None
+    pd_should_warn = PANDAS_GE_210 and convert_dtype is not None
 
     meta_val = ddf.x._meta_nonempty.iloc[0]
 
@@ -3420,11 +3420,17 @@ def test_apply_convert_dtype(convert_dtype):
         assert x != meta_val
         return x + 1
 
-    with _check_warning(should_warn, FutureWarning, "the convert_dtype parameter"):
+    with _check_warning(pd_should_warn, FutureWarning, "the convert_dtype parameter"):
         expected = df.x.apply(func, **kwargs)
-    with _check_warning(should_warn, FutureWarning, "the convert_dtype parameter"):
+    with _check_warning(pd_should_warn, FutureWarning, "the convert_dtype parameter"):
         result = ddf.x.apply(func, **kwargs, meta=expected)
-    assert_eq(result, expected)
+
+    dask_should_warn = pytest.warns(FutureWarning, match="the 'convert_dtype' keyword")
+    if convert_dtype is None:
+        dask_should_warn = contextlib.nullcontext()
+
+    with dask_should_warn:
+        assert_eq(result, expected)
 
 
 def test_apply_warns():
diff --git a/dask/tests/test_highgraph.py b/dask/tests/test_highgraph.py
index 87283c233..03a21c763 100644
--- a/dask/tests/test_highgraph.py
+++ b/dask/tests/test_highgraph.py
@@ -9,6 +9,7 @@ from concurrent.futures import ThreadPoolExecutor
 import pytest
 
 import dask
+from dask.base import normalize_token
 from dask.blockwise import Blockwise, blockwise_token
 from dask.highlevelgraph import HighLevelGraph, Layer, MaterializedLayer, to_graphviz
 from dask.utils_test import inc
@@ -300,3 +301,13 @@ def test_node_tooltips_exist():
             end = layer.find('"', start)
             tooltip = layer[start:end]
             assert len(tooltip) > 0
+
+
+def test_tokenize_hlg():
+    import dask.bag as db
+
+    a = db.from_sequence(list(range(10)), npartitions=2).max()
+    b = db.from_sequence(list(range(10)), npartitions=2).max()
+    c = db.from_sequence(list(range(10)), npartitions=3).max()
+    assert normalize_token(a.dask) == normalize_token(b.dask)
+    assert normalize_token(a.dask) != normalize_token(c.dask)