Repository: dask/dask
Pull request #8052 (2021-08-17 13:52:21)

Remove `wrap_func_like_safe`, not required with NumPy >= 1.17
Follow-up to https://github.com/dask/dask/pull/7939 where other `_like_safe` functions for compatibility with NumPy < 1.17 were removed, as they're not necessary with new NumPy minimum requirements in Dask.

pentschev: cc @jakirkham 

Patch:
diff --git a/dask/array/backends.py b/dask/array/backends.py
index 9f622850a..78ec4525f 100644
--- a/dask/array/backends.py
+++ b/dask/array/backends.py
@@ -108,8 +108,12 @@ def _tensordot(a, b, axes=2):
 def register_cupy():
     import cupy
 
+    from dask.dataframe.backends import percentile
+    from dask.dataframe.dispatch import percentile_dispatch
+
     concatenate_lookup.register(cupy.ndarray, cupy.concatenate)
     tensordot_lookup.register(cupy.ndarray, cupy.tensordot)
+    percentile_dispatch.register(cupy.ndarray, percentile)
 
     @einsum_lookup.register(cupy.ndarray)
     def _cupy_einsum(*args, **kwargs):
diff --git a/dask/array/percentile.py b/dask/array/percentile.py
index b43a14a1e..fda808782 100644
--- a/dask/array/percentile.py
+++ b/dask/array/percentile.py
@@ -12,22 +12,29 @@ from .core import Array
 
 @wraps(np.percentile)
 def _percentile(a, q, interpolation="linear"):
+    import pandas as pd
+
     n = len(a)
     if not len(a):
         return None, n
     if isinstance(q, Iterator):
         q = list(q)
     if a.dtype.name == "category":
-        result = np.percentile(a.codes, q, interpolation=interpolation)
-        import pandas as pd
+        result = np.percentile(a.cat.codes, q, interpolation=interpolation)
 
-        return pd.Categorical.from_codes(result, a.categories, a.ordered), n
-    if np.issubdtype(a.dtype, np.datetime64):
-        a2 = a.astype("i8")
-        result = np.percentile(a2, q, interpolation=interpolation).astype(a.dtype)
+        return pd.Categorical.from_codes(result, a.dtype.categories, a.dtype.ordered), n
+    if isinstance(a.dtype, pd.core.dtypes.dtypes.DatetimeTZDtype) or np.issubdtype(
+        a.dtype, np.datetime64
+    ):
+        if isinstance(a, (pd.Series, pd.Index)):
+            values = a.values
+        else:
+            values = a
+        a2 = values.view("i8")
+        result = np.percentile(a2, q, interpolation=interpolation).astype(values.dtype)
         if q[0] == 0:
             # https://github.com/dask/dask/issues/6864
-            result[0] = min(result[0], a.min())
+            result[0] = min(result[0], values.min())
         return result, n
     if not np.issubdtype(a.dtype, np.number):
         interpolation = "nearest"
@@ -84,6 +91,8 @@ def percentile(a, q, interpolation="linear", method="default"):
     --------
     numpy.percentile : Numpy's equivalent Percentile function
     """
+    from dask.dataframe.dispatch import _percentile
+
     from .utils import array_safe, meta_from_array
 
     if not a.ndim == 1:
diff --git a/dask/array/wrap.py b/dask/array/wrap.py
index ebccef957..866fc5846 100644
--- a/dask/array/wrap.py
+++ b/dask/array/wrap.py
@@ -101,17 +101,6 @@ def wrap_func_like(func, *args, **kwargs):
     return Array(dsk, name, chunks, meta=meta.astype(dtype))
 
 
-def wrap_func_like_safe(func, func_like, *args, **kwargs):
-    """
-    Safe implementation for wrap_func_like(), attempts to use func_like(),
-    if the shape keyword argument, falls back to func().
-    """
-    try:
-        return func_like(*args, **kwargs)
-    except TypeError:
-        return func(*args, **kwargs)
-
-
 @curry
 def wrap(wrap_func, func, **kwargs):
     func_like = kwargs.pop("func_like", None)
@@ -179,7 +168,7 @@ zeros = w(broadcast_trick(np.zeros_like), dtype="f8")
 empty = w(broadcast_trick(np.empty_like), dtype="f8")
 
 
-w_like = wrap(wrap_func_like_safe)
+w_like = wrap(wrap_func_like)
 
 
 empty_like = w_like(np.empty, func_like=np.empty_like)
diff --git a/dask/dataframe/backends.py b/dask/dataframe/backends.py
index 21572e885..48d2f5d95 100644
--- a/dask/dataframe/backends.py
+++ b/dask/dataframe/backends.py
@@ -12,6 +12,7 @@ from pandas.api.types import (
     union_categoricals,
 )
 
+from dask.array.percentile import _percentile
 from dask.sizeof import SimpleSizeof, sizeof
 
 from ..utils import is_arraylike, typename
@@ -27,6 +28,7 @@ from .dispatch import (
     make_meta_dispatch,
     make_meta_obj,
     meta_nonempty,
+    percentile_dispatch,
     tolist_dispatch,
     union_categoricals_dispatch,
 )
@@ -547,3 +549,8 @@ def is_categorical_dtype_pandas(obj):
 @make_meta_obj.register_lazy("cudf")
 def _register_cudf():
     import dask_cudf  # noqa: F401
+
+
+@percentile_dispatch.register((pd.Series, np.ndarray, pd.Index))
+def percentile(a, q, interpolation="linear"):
+    return _percentile(a, q, interpolation)
diff --git a/dask/dataframe/core.py b/dask/dataframe/core.py
index 58fa7b80c..a19ec8d80 100644
--- a/dask/dataframe/core.py
+++ b/dask/dataframe/core.py
@@ -6012,14 +6012,15 @@ def quantile(df, q, method="default"):
         }
     else:
 
-        from dask.array.percentile import _percentile, merge_percentiles
+        from dask.array.percentile import merge_percentiles
+        from dask.dataframe.dispatch import _percentile
 
         # Add 0 and 100 during calculation for more robust behavior (hopefully)
         calc_qs = np.pad(qs, 1, mode="constant")
         calc_qs[-1] = 100
         name = "quantiles-1-" + token
         val_dsk = {
-            (name, i): (_percentile, (getattr, key, "values"), calc_qs)
+            (name, i): (_percentile, key, calc_qs)
             for i, key in enumerate(df.__dask_keys__())
         }
 
diff --git a/dask/dataframe/dispatch.py b/dask/dataframe/dispatch.py
index 7698b1c66..657726f8a 100644
--- a/dask/dataframe/dispatch.py
+++ b/dask/dataframe/dispatch.py
@@ -22,6 +22,12 @@ concat_dispatch = Dispatch("concat")
 tolist_dispatch = Dispatch("tolist")
 is_categorical_dtype_dispatch = Dispatch("is_categorical_dtype")
 union_categoricals_dispatch = Dispatch("union_categoricals")
+percentile_dispatch = Dispatch("percentile_dispatch")
+
+
+def _percentile(a, q, interpolation="linear"):
+    func = percentile_dispatch.dispatch(type(a))
+    return func(a, q, interpolation)
 
 
 def concat(
diff --git a/dask/dataframe/groupby.py b/dask/dataframe/groupby.py
index 03605d454..2ff204951 100644
--- a/dask/dataframe/groupby.py
+++ b/dask/dataframe/groupby.py
@@ -968,7 +968,7 @@ def _apply_func_to_columns(df_like, prefix, func):
         columns = df_like.columns
     else:
         # handle GroupBy objects
-        columns = df_like._selected_obj.columns
+        columns = df_like.obj.columns
 
     columns = sorted(col for col in columns if col.startswith(prefix))
 
diff --git a/dask/dataframe/partitionquantiles.py b/dask/dataframe/partitionquantiles.py
index 070fe7447..c1a625657 100644
--- a/dask/dataframe/partitionquantiles.py
+++ b/dask/dataframe/partitionquantiles.py
@@ -401,21 +401,22 @@ def percentiles_summary(df, num_old, num_new, upsample, state):
         Scale factor to increase the number of percentiles calculated in
         each partition.  Use to improve accuracy.
     """
-    from dask.array.percentile import _percentile
+    from dask.dataframe.dispatch import _percentile
 
     length = len(df)
     if length == 0:
         return ()
     random_state = np.random.RandomState(state)
     qs = sample_percentiles(num_old, num_new, length, upsample, random_state)
-    data = df.values
+    data = df
     interpolation = "linear"
+
     if is_categorical_dtype(data):
-        data = data.codes
+        data = data.cat.codes
         interpolation = "nearest"
-    elif np.issubdtype(data.dtype, np.integer) and not is_cupy_type(data):
-        # CuPy doesn't currently support "nearest" interpolation,
-        # so it's special cased in the condition above.
+    elif isinstance(data.dtype, pd.core.dtypes.dtypes.DatetimeTZDtype) or np.issubdtype(
+        data.dtype, np.integer
+    ):
         interpolation = "nearest"
     vals, n = _percentile(data, qs, interpolation=interpolation)
     if (
diff --git a/dask/local.py b/dask/local.py
index 9f035ef27..664ae93ab 100644
--- a/dask/local.py
+++ b/dask/local.py
@@ -181,7 +181,7 @@ def start_state_from_dask(dsk, cache=None, sortkey=None):
     waiting_data = dict((k, v.copy()) for k, v in dependents.items() if v)
 
     ready_set = set([k for k, v in waiting.items() if not v])
-    ready = sorted(ready_set, key=sortkey)
+    ready = sorted(ready_set, key=sortkey, reverse=True)
     waiting = dict((k, v) for k, v in waiting.items() if v)
 
     state = {
@@ -262,7 +262,7 @@ def finish_task(
 
     Mutates.  This should run atomically (with a lock).
     """
-    for dep in sorted(state["dependents"][key], key=sortkey):
+    for dep in sorted(state["dependents"][key], key=sortkey, reverse=True):
         s = state["waiting"][dep]
         s.remove(key)
         if not s:
@@ -463,11 +463,15 @@ def get_async(
                     ntasks = nready
                     chunksize = -(ntasks // -num_workers)
                 else:
-                    ntasks = min(nready, chunksize * num_workers)
+                    used_workers = -(len(state["running"]) // -chunksize)
+                    avail_workers = max(num_workers - used_workers, 0)
+                    ntasks = min(nready, chunksize * avail_workers)
 
                 # Prep all ready tasks for submission
                 args = []
-                for i, key in zip(range(ntasks), state["ready"]):
+                for _ in range(ntasks):
+                    # Get the next task to compute (most recently added)
+                    key = state["ready"].pop()
                     # Notify task is running
                     state["running"].add(key)
                     for f in pretask_cbs:
@@ -487,7 +491,6 @@ def get_async(
                             pack_exception,
                         )
                     )
-                del state["ready"][:ntasks]
 
                 # Batch submit
                 for i in range(-(len(args) // -chunksize)):


Test patch:
diff --git a/dask/tests/test_local.py b/dask/tests/test_local.py
index 8830a216e..a45aab7b9 100644
--- a/dask/tests/test_local.py
+++ b/dask/tests/test_local.py
@@ -1,3 +1,5 @@
+import pytest
+
 import dask
 from dask.local import finish_task, get_sync, sortkey, start_state_from_dask
 from dask.order import order
@@ -170,3 +172,22 @@ def test_ordering():
     get_sync(dsk, "y")
 
     assert L == sorted(L)
+
+
+def test_complex_ordering():
+    da = pytest.importorskip("dask.array")
+    from dask.diagnostics import Callback
+
+    actual_order = []
+
+    def track_order(key, dask, state):
+        actual_order.append(key)
+
+    x = da.random.normal(size=(20, 20), chunks=(-1, -1))
+    res = (x.dot(x.T) - x.mean(axis=0)).std()
+    dsk = dict(res.__dask_graph__())
+    exp_order_dict = order(dsk)
+    exp_order = sorted(exp_order_dict.keys(), key=exp_order_dict.get)
+    with Callback(pretask=track_order):
+        get_sync(dsk, exp_order[-1])
+    assert actual_order == exp_order
diff --git a/dask/tests/test_order.py b/dask/tests/test_order.py
index 1d90bece5..d58179d66 100644
--- a/dask/tests/test_order.py
+++ b/dask/tests/test_order.py
@@ -351,7 +351,7 @@ def test_local_parents_of_reduction(abcde):
     b1, b2, b3 = [b + i for i in "123"]
     c1, c2, c3 = [c + i for i in "123"]
 
-    expected = [a3, b3, c3, a2, a1, b2, b1, c2, c1]
+    expected = [a3, a2, a1, b3, b2, b1, c3, c2, c1]
 
     log = []