Repository: dask/dask
Pull request #11661 (2025-01-10 12:33:45)

Add redirects for groupby docs after dask-expr merge
- [ ] Closes #xxxx
- [ ] Tests added / passed
- [ ] 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   4h 32m 25s ⏱️ + 3m 52s
 17 153 tests ±0   15 958 ✅ ±0   1 195 💤 ±0  0 ❌ ±0 
211 329 runs  ±0  194 152 ✅ +1  17 177 💤  - 1  0 ❌ ±0 

Results for commit b9afcdfc. ± Comparison against base commit 2e9e2477.

[test-results]:data:application/gzip;base64,H4sIABMcgWcC/02MSw6EIBAFr2JYz8KmRei5jEGEhIyfCcLKzN1HjKLLqpdXG3N+tCt7VyBeFVuTjzcNKejolzlzi01W+xqPXYLAi7s1GXOcSKhbfvw3S6By7Jz24+7qImwISzhNSHNOcwDkdHJJUwOCF3mmJUh5uUf64GfZLNPk4w6sJ+3M4IysLTagFA1oaytMKxRBT4CYmSP7/QHTApI5GgEAAA==


Patch:
diff --git a/dask/array/core.py b/dask/array/core.py
index 85fe6b1b3..f3185b48c 100644
--- a/dask/array/core.py
+++ b/dask/array/core.py
@@ -3881,7 +3881,7 @@ def to_zarr(
 
     if storage_options:
         if _zarr_v3():
-            store = zarr.storage.RemoteStore(
+            store = zarr.storage.FsspecStore(
                 url, mode=kwargs.pop("mode", "a"), **storage_options
             )
         else:
diff --git a/dask/array/optimization.py b/dask/array/optimization.py
index 20893e30a..6cf9ae70a 100644
--- a/dask/array/optimization.py
+++ b/dask/array/optimization.py
@@ -1,15 +1,38 @@
 from __future__ import annotations
 
+from collections.abc import Callable
 from itertools import zip_longest
 from numbers import Integral
+from typing import Any
+
+import numpy as np
 
 from dask import config
-from dask._task_spec import convert_legacy_graph, fuse_linear_task_spec
+from dask._task_spec import (
+    Alias,
+    DataNode,
+    GraphNode,
+    Task,
+    _execute_subgraph,
+    convert_legacy_graph,
+    fuse_linear_task_spec,
+)
+from dask.array.chunk import getitem
+from dask.array.core import getter, getter_inline, getter_nofancy
 from dask.blockwise import fuse_roots, optimize_blockwise
 from dask.core import flatten
 from dask.highlevelgraph import HighLevelGraph
 from dask.utils import ensure_dict
 
+# All get* functions the optimizations know about
+
+
+GETTERS = (getter, getter_nofancy, getter_inline, getitem)
+# These get* functions aren't ever completely removed from the graph,
+# even if the index should be a no-op by numpy semantics. Some array-like's
+# don't completely follow semantics, making indexing always necessary.
+GETNOREMOVE = (getter, getter_nofancy)
+
 
 def optimize(dsk, keys, **kwargs):
     """Optimize dask for array computation
@@ -37,6 +60,129 @@ def optimize(dsk, keys, **kwargs):
 
     dsk = convert_legacy_graph(dsk)
     dsk = fuse_linear_task_spec(dsk, keys=keys)
+    dsk = optimize_slices(dsk)
+
+    return dsk
+
+
+def _is_getter_task(
+    value: GraphNode,
+) -> tuple[Callable, Any, Any, bool, bool | None] | None:
+    """Check if a value in a Dask graph looks like a getter.
+    1. Is it a tuple with the first element a known getter.
+    If a getter is found, it returns a tuple with (getter, array, index, asarray, lock).
+    Otherwise it returns ``None``.
+    """
+    if not isinstance(value, Task):
+        return None
+    func = value.func
+    if func in GETTERS:
+        get = func
+    else:
+        return None
+
+    length = len(value.args)
+    if length == 2:
+        # getter defaults to asarray=True, getitem is semantically False
+        return get, value.args[0], value.args[1], get is not getitem, None
+    elif length == 4:
+        return get, *list(value.args)
+
+    return None
+
+
+def optimize_slices(dsk):
+    """Optimize slices
+    1.  Fuse repeated slices, like x[5:][2:6] -> x[7:11]
+
+    This is generally not very important since we are fusing those tasks anyway. There
+    is one specific exception to how xarray implements opening netcdf files and subsequent
+    slices. Not merging them together can cause reading the whole netcdf file before
+    we drop the unnecessary data. Fusing slices avoids that pattern.
+
+    See https://github.com/pydata/xarray/issues/9926
+    """
+    fancy_ind_types = (list, np.ndarray)
+    dsk = dsk.copy()
+    for _, v in dsk.items():
+        if not (isinstance(v, Task) and v.func is _execute_subgraph):
+            continue
+
+        inner_graph: dict = v.args[0]
+
+        seen = set()
+
+        for inner_key, inner_value in inner_graph.items():
+            if inner_key in seen:
+                continue
+            else:
+                seen.add(inner_key)
+
+            if a_task := _is_getter_task(inner_value):
+                temp_key = inner_key
+                get, a, a_index, a_asarray, a_lock = a_task
+                fused = False
+
+                while a.key in inner_graph and (
+                    b_task := _is_getter_task(inner_graph[a.key])
+                ):
+                    seen.add(a.key)
+                    f2, b, b_index, b_asarray, b_lock = b_task
+
+                    if isinstance(a_index, DataNode):
+                        a_index = a_index.value
+                    if isinstance(b_index, DataNode):
+                        b_index = b_index.value
+
+                    if a_lock and a_lock is not b_lock:
+                        break
+                    if (type(a_index) is tuple) != (type(b_index) is tuple):
+                        break
+                    if type(a_index) is tuple:
+                        indices = b_index + a_index
+                        if len(a_index) != len(b_index) and any(
+                            i is None for i in indices
+                        ):
+                            break
+                        if f2 is getter_nofancy and not any(
+                            isinstance(i, fancy_ind_types) for i in indices
+                        ):
+                            break
+                    elif f2 is getter_nofancy and (
+                        type(a_index) in fancy_ind_types
+                        or type(b_index) in fancy_ind_types
+                    ):
+                        break
+                    try:
+                        c_index = fuse_slice(b_index, a_index)
+                        # rely on fact that nested gets never decrease in
+                        # strictness e.g. `(getter_nofancy, (getter, ...))` never
+                        # happens
+                        get = getter if f2 is getter_inline else f2
+
+                        inner_graph[temp_key] = Alias(temp_key, a.key)
+                        temp_key = a.key
+                        fused = True
+                    except NotImplementedError:
+                        break
+                    a, a_index, a_lock = b, c_index, b_lock
+                    a_asarray |= b_asarray
+
+                    if not isinstance(a, Task):
+                        break
+
+                if not fused:
+                    pass
+                elif get is getitem or (a_asarray and not a_lock):
+                    # default settings are fine, drop the extra parameters Since we
+                    # always fallback to inner `get` functions, `get is getitem`
+                    # can only occur if all gets are getitem, meaning all
+                    # parameters must be getitem defaults.
+                    inner_graph[temp_key] = Task(temp_key, get, a, a_index)
+                else:
+                    inner_graph[temp_key] = Task(
+                        temp_key, get, a, a_index, a_asarray, a_lock
+                    )
 
     return dsk
 
diff --git a/dask/dataframe/__init__.py b/dask/dataframe/__init__.py
index 03bf50a94..468c72e1b 100644
--- a/dask/dataframe/__init__.py
+++ b/dask/dataframe/__init__.py
@@ -69,3 +69,54 @@ from dask.dataframe.io import (
     to_sql,
 )
 from dask.dataframe.utils import assert_eq
+
+__all__ = [
+    "DataFrame",
+    "Index",
+    "Scalar",
+    "Series",
+    "Aggregation",
+    "backends",
+    "compute",
+    "concat",
+    "dispatch",
+    "from_array",
+    "from_dask_array",
+    "from_delayed",
+    "from_dict",
+    "from_graph",
+    "from_map",
+    "from_pandas",
+    "get_collection_type",
+    "get_dummies",
+    "isna",
+    "map_overlap",
+    "map_partitions",
+    "melt",
+    "merge",
+    "merge_asof",
+    "pivot_table",
+    "read_csv",
+    "read_fwf",
+    "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",
+    "assert_eq",
+]
diff --git a/dask/dataframe/api.py b/dask/dataframe/api.py
index a39b62d38..315eb00c6 100644
--- a/dask/dataframe/api.py
+++ b/dask/dataframe/api.py
@@ -1,4 +1,10 @@
 from __future__ import annotations
 
+from dask.dataframe.core import get_parallel_type  # noqa: F401
+from dask.dataframe.dask_expr import FrameBase, elemwise, new_collection  # noqa: F401
+from dask.dataframe.dask_expr._expr import are_co_aligned, emulate  # noqa: F401
 from dask.dataframe.dask_expr._groupby import GroupBy, SeriesGroupBy  # noqa: F401
+from dask.dataframe.dask_expr._reductions import ApplyConcatApply  # noqa: F401
 from dask.dataframe.dask_expr._rolling import Rolling  # noqa: F401
+from dask.dataframe.extensions import make_array_nonempty, make_scalar  # noqa: F401
+from dask.dataframe.utils import meta_nonempty  # noqa: F401
diff --git a/dask/dataframe/dask_expr/_collection.py b/dask/dataframe/dask_expr/_collection.py
index b4bb73e57..3d4994825 100644
--- a/dask/dataframe/dask_expr/_collection.py
+++ b/dask/dataframe/dask_expr/_collection.py
@@ -3101,7 +3101,7 @@ class DataFrame(FrameBase):
                 "installed."
             )
         if meta is None:
-            meta = expr._emulate(M.map, self, func, na_action=na_action, udf=True)
+            meta = expr.emulate(M.map, self, func, na_action=na_action, udf=True)
             warnings.warn(meta_warning(meta))
         return new_collection(expr.Map(self, arg=func, na_action=na_action, meta=meta))
 
@@ -3233,7 +3233,7 @@ class DataFrame(FrameBase):
             )
             raise NotImplementedError(msg)
         if meta is no_default:
-            meta = expr._emulate(
+            meta = expr.emulate(
                 M.apply, self, function, args=args, udf=True, axis=axis, **kwargs
             )
             warnings.warn(meta_warning(meta))
@@ -4200,7 +4200,7 @@ class Series(FrameBase):
                     expr.MapAlign(self, arg, op=None, na_action=na_action, meta=meta)
                 )
         if meta is None:
-            meta = expr._emulate(M.map, self, arg, na_action=na_action, udf=True)
+            meta = expr.emulate(M.map, self, arg, na_action=na_action, udf=True)
             warnings.warn(meta_warning(meta))
         return new_collection(expr.Map(self, arg=arg, na_action=na_action, meta=meta))
 
@@ -4384,7 +4384,7 @@ class Series(FrameBase):
         """
         self._validate_axis(axis)
         if meta is no_default:
-            meta = expr._emulate(M.apply, self, function, args=args, udf=True, **kwargs)
+            meta = expr.emulate(M.apply, self, function, args=args, udf=True, **kwargs)
             warnings.warn(meta_warning(meta))
         return new_collection(self.expr.apply(function, *args, meta=meta, **kwargs))
 
@@ -4727,7 +4727,7 @@ class Index(Series):
                     expr.MapIndexAlign(self, arg, na_action, meta, is_monotonic)
                 )
         if meta is None:
-            meta = expr._emulate(M.map, self, arg, na_action=na_action, udf=True)
+            meta = expr.emulate(M.map, self, arg, na_action=na_action, udf=True)
             warnings.warn(meta_warning(meta))
         return new_collection(
             expr.Map(
diff --git a/dask/dataframe/dask_expr/_expr.py b/dask/dataframe/dask_expr/_expr.py
index 4c852e5df..dac7d0615 100644
--- a/dask/dataframe/dask_expr/_expr.py
+++ b/dask/dataframe/dask_expr/_expr.py
@@ -4033,7 +4033,7 @@ def _extract_meta(x, nonempty=False):
         return x
 
 
-def _emulate(func, *args, udf=False, **kwargs):
+def emulate(func, *args, udf=False, **kwargs):
     """
     Apply a function using args / kwargs. If arguments contain dd.DataFrame /
     dd.Series, using internal cache (``_meta``) for calculation
@@ -4053,7 +4053,7 @@ def _get_meta_map_partitions(args, dfs, func, kwargs, meta, parent_meta):
         # Use non-normalized kwargs here, as we want the real values (not
         # delayed values)
         a = [meta_nonempty(arg._meta) if isinstance(arg, Expr) else arg for arg in args]
-        meta = _emulate(func, *a, udf=True, **kwargs)
+        meta = emulate(func, *a, udf=True, **kwargs)
         meta_is_emulated = True
     else:
         meta = make_meta(meta, index=meta_index, parent_meta=parent_meta)
diff --git a/docs/source/conf.py b/docs/source/conf.py
index 903d533f1..b68225158 100644
--- a/docs/source/conf.py
+++ b/docs/source/conf.py
@@ -415,6 +415,46 @@ redirect_files = [
     ("how-to/setup-prometheus.html", "prometheus.html"),
 ]
 
+for func in [
+    "aggregate",
+    "apply",
+    "bfill",
+    "count",
+    "cumcount",
+    "cumprod",
+    "cumsum",
+    "fillna",
+    "ffill",
+    "get_group",
+    "max",
+    "mean",
+    "min",
+    "size",
+    "std",
+    "sum",
+    "var",
+    "cov",
+    "corr",
+    "first",
+    "last",
+    "idxmin",
+    "idxmax",
+    "rolling",
+    "transform",
+]:
+    redirect_files.append(
+        (
+            f"generated/dask.dataframe.groupby.DataFrameGroupBy.{func}.html",
+            f"generated/dask.dataframe.api.GroupBy.{func}.html",
+        )
+    )
+    redirect_files.append(
+        (
+            f"generated/dask.dataframe.groupby.SeriesGroupBy.{func}.html",
+            f"generated/dask.dataframe.api.SeriesGroupBy.{func}.html",
+        )
+    )
+
 
 redirect_template = """\
 <html>


Test patch:
diff --git a/dask/array/tests/test_creation.py b/dask/array/tests/test_creation.py
index 7b6f947c8..d7687e43f 100644
--- a/dask/array/tests/test_creation.py
+++ b/dask/array/tests/test_creation.py
@@ -988,6 +988,22 @@ def test_nan_zeros_ones_like(fn, shape_chunks, dtype):
     )
 
 
+def test_from_array_getitem_fused():
+    arr = np.arange(100).reshape(10, 10)
+    darr = da.from_array(arr, chunks=(5, 5))
+    result = darr[slice(1, 5), :][slice(1, 3), :]
+    dsk = collections_to_dsk([result])
+    # Ensure that slices are merged properly
+    key = [k for k in dsk if "array-getitem" in k[0]][0]
+    key_2 = [
+        k
+        for k, v in dsk[key].args[0].items()
+        if "getitem" in k[0] and not isinstance(v, Alias)
+    ][0]
+    assert dsk[key].args[0][key_2].args[1] == ((slice(2, 4), slice(0, None)))
+    assert_eq(result, arr[slice(1, 5), :][slice(1, 3), :])
+
+
 @pytest.mark.parametrize("shape_chunks", [((50, 4), (10, 2)), ((50,), (10,))])
 @pytest.mark.parametrize("dtype", ["u4", np.float32, None, np.int64])
 def test_nan_empty_like(shape_chunks, dtype):