Repository: dask-contrib/dask-expr
Pull request #958 (2024-03-08 15:03:05)

Ensure that we respect shuffle in context manager


Patch:
diff --git a/dask_expr/_collection.py b/dask_expr/_collection.py
index 512ae05..7eb66b9 100644
--- a/dask_expr/_collection.py
+++ b/dask_expr/_collection.py
@@ -127,6 +127,7 @@ from dask_expr._util import (
     _raise_if_object_series,
     _tokenize_deterministic,
     _validate_axis,
+    get_specified_shuffle,
     is_scalar,
 )
 from dask_expr.io import FromPandasDivisions, FromScalars
@@ -803,7 +804,7 @@ Expr={expr}"""
                 on,
                 npartitions,
                 ignore_index,
-                shuffle_method,
+                get_specified_shuffle(shuffle_method),
                 options,
                 index_shuffle=on_index,
             )
@@ -3164,7 +3165,7 @@ class DataFrame(FrameBase):
                 npartitions=npartitions,
                 upsample=upsample,
                 partition_size=partition_size,
-                shuffle_method=shuffle_method,
+                shuffle_method=get_specified_shuffle(shuffle_method),
                 append=append,
                 options=options,
             )
@@ -3240,7 +3241,7 @@ class DataFrame(FrameBase):
                 sort_function_kwargs,
                 upsample,
                 ignore_index,
-                shuffle_method,
+                get_specified_shuffle(shuffle_method),
                 options=options,
             )
         )
@@ -4890,7 +4891,7 @@ def merge(
             right_index=right_index,
             suffixes=suffixes,
             indicator=indicator,
-            shuffle_method=shuffle_method,
+            shuffle_method=get_specified_shuffle(shuffle_method),
             _npartitions=npartitions,
             broadcast=broadcast,
         )
diff --git a/dask_expr/_groupby.py b/dask_expr/_groupby.py
index 4f80637..26553c6 100644
--- a/dask_expr/_groupby.py
+++ b/dask_expr/_groupby.py
@@ -70,7 +70,12 @@ from dask_expr._expr import (
 )
 from dask_expr._reductions import ApplyConcatApply, Chunk, Reduction
 from dask_expr._shuffle import RearrangeByColumn
-from dask_expr._util import PANDAS_GE_300, _convert_to_list, is_scalar
+from dask_expr._util import (
+    PANDAS_GE_300,
+    _convert_to_list,
+    get_specified_shuffle,
+    is_scalar,
+)
 
 
 def _as_dict(key, value):
@@ -1483,7 +1488,7 @@ class GroupBy:
         self.obj = obj[projection] if projection is not None else obj
         self.sort = sort
         self.observed = (
-            observed if observed is not None else True if not PANDAS_GE_300 else True
+            observed if observed is not None else False if not PANDAS_GE_300 else True
         )
         self.dropna = dropna
         self.group_keys = group_keys
@@ -1531,7 +1536,7 @@ class GroupBy:
                 split_every,
                 split_out,
                 self.sort,
-                shuffle_method,
+                get_specified_shuffle(shuffle_method),
                 *self.by,
             )
         )
@@ -1941,7 +1946,7 @@ class GroupBy:
                 meta,
                 args,
                 kwargs,
-                shuffle_method,
+                get_specified_shuffle(shuffle_method),
                 *self.by,
             )
         )
@@ -1960,7 +1965,7 @@ class GroupBy:
                 meta,
                 args,
                 kwargs,
-                shuffle_method,
+                get_specified_shuffle(shuffle_method),
                 *self.by,
             )
         )
@@ -2055,7 +2060,7 @@ class GroupBy:
                 no_default,
                 (),
                 {"numeric_only": numeric_only},
-                shuffle_method,
+                get_specified_shuffle(shuffle_method),
                 split_every,
                 *self.by,
             )
@@ -2206,7 +2211,7 @@ class SeriesGroupBy(GroupBy):
                 split_every,
                 split_out,
                 self.sort,
-                shuffle_method,
+                get_specified_shuffle(shuffle_method),
                 *self.by,
             )
         )
diff --git a/dask_expr/_util.py b/dask_expr/_util.py
index 0502016..d8c000b 100644
--- a/dask_expr/_util.py
+++ b/dask_expr/_util.py
@@ -219,3 +219,8 @@ def _is_any_real_numeric_dtype(arr_or_dtype):
             and not is_complex_dtype(arr_or_dtype)
             and not is_bool_dtype(arr_or_dtype)
         )
+
+
+def get_specified_shuffle(shuffle_method):
+    # Take the config shuffle if given, otherwise defer evaluation until optimize
+    return shuffle_method or config.get("dataframe.shuffle.method", None)


Test patch:
diff --git a/dask_expr/tests/test_groupby.py b/dask_expr/tests/test_groupby.py
index a8f78fa..f344630 100644
--- a/dask_expr/tests/test_groupby.py
+++ b/dask_expr/tests/test_groupby.py
@@ -8,7 +8,7 @@ import pytest
 from dask_expr import from_pandas
 from dask_expr._groupby import Aggregation, GroupByUDFBlockwise
 from dask_expr._reductions import TreeReduce
-from dask_expr._shuffle import Shuffle, divisions_lru
+from dask_expr._shuffle import Shuffle, TaskShuffle, divisions_lru
 from dask_expr.io import FromPandas
 from dask_expr.tests._util import _backend_library, assert_eq, xfail_gpu
 
@@ -958,6 +958,32 @@ def test_groupby_size_drop_columns(df, pdf):
     assert result.simplify()._name == df[[]].groupby("x").size().simplify()._name
 
 
+def test_groupy_respect_shuffle_context(df, pdf):
+    def _check_task_shuffle(q):
+        result = q.optimize(fuse=False)
+        assert len([x for x in result.walk() if isinstance(x, TaskShuffle)]) > 0
+
+    with dask.config.set({"dataframe.shuffle.method": "tasks"}):
+        q = df.groupby("x").apply(lambda x: x)
+    _check_task_shuffle(q)
+
+    with dask.config.set({"dataframe.shuffle.method": "tasks"}):
+        q = df.groupby("x").transform(lambda x: x)
+    _check_task_shuffle(q)
+
+    with dask.config.set({"dataframe.shuffle.method": "tasks"}):
+        q = df.groupby("x").sum(split_out=True)
+    _check_task_shuffle(q)
+
+    with dask.config.set({"dataframe.shuffle.method": "tasks"}):
+        q = df.groupby("x").y.nunique(split_out=True)
+    _check_task_shuffle(q)
+
+    with dask.config.set({"dataframe.shuffle.method": "tasks"}):
+        q = df.groupby("x").median(split_out=True)
+    _check_task_shuffle(q)
+
+
 def test_groupby_agg_meta_error(df, pdf):
     result = df.groupby(["x"]).agg({"y": ["median", "std"]})
     expected = pdf.groupby(["x"]).agg({"y": ["median", "std"]})
diff --git a/dask_expr/tests/test_shuffle.py b/dask_expr/tests/test_shuffle.py
index 8c627ee..75900f7 100644
--- a/dask_expr/tests/test_shuffle.py
+++ b/dask_expr/tests/test_shuffle.py
@@ -8,7 +8,12 @@ from dask_expr import from_pandas, new_collection
 from dask_expr._expr import Assign, Blockwise, Filter
 from dask_expr._reductions import NFirst, NLast
 from dask_expr._repartition import RepartitionToFewer
-from dask_expr._shuffle import BaseSetIndexSortValues, TaskShuffle, divisions_lru
+from dask_expr._shuffle import (
+    BaseSetIndexSortValues,
+    P2PShuffle,
+    TaskShuffle,
+    divisions_lru,
+)
 from dask_expr.io import FromPandas
 from dask_expr.tests._util import _backend_library, assert_eq, xfail_gpu
 
@@ -689,6 +694,25 @@ def test_shuffle_no_assign(df, pdf):
     assert len([x for x in q.walk() if isinstance(x, Assign)]) == 0
 
 
+@pytest.mark.parametrize("func", ["set_index", "sort_values", "shuffle"])
+def test_respect_context_shuffle(df, pdf, func):
+    with dask.config.set({"dataframe.shuffle.method": "tasks"}):
+        q = getattr(df, func)("x")
+    result = q.optimize(fuse=False)
+    assert len([x for x in result.walk() if isinstance(x, TaskShuffle)]) > 0
+
+    q = getattr(df, func)("x")
+    with dask.config.set({"dataframe.shuffle.method": "tasks"}):
+        result = q.optimize(fuse=False)
+    assert len([x for x in result.walk() if isinstance(x, TaskShuffle)]) > 0
+
+    with dask.config.set({"dataframe.shuffle.method": "p2p"}):
+        q = getattr(df, func)("x")
+    with dask.config.set({"dataframe.shuffle.method": "tasks"}):
+        result = q.optimize(fuse=False)
+    assert len([x for x in result.walk() if isinstance(x, P2PShuffle)]) > 0
+
+
 @pytest.mark.parametrize("meth", ["shuffle", "sort_values"])
 def test_shuffle_filter_pushdown(pdf, meth):
     pdf["z"] = 1