Repository: dask/dask (python)

Issue:
`optimize` not lowering fully
```python
import pandas as pd
from dask.dataframe import from_pandas
df = pd.DataFrame(
    {
        "A": [1, 2, 3, 4, 5, 6, 7, 8, 9],
        "B": [9, 8, 7, 6, 5, 4, 3, 2, 1],
        "C": [True, False, True] * 3,
    },
    columns=list("ABC"),
)
ddf = from_pandas(df, 2)
expr = ddf[ddf.C.repartition([0, 2, 5, 8])].expr
expr = expr.optimize()
expr.__dask_graph__()
```

The graph cannot be generated and is running into a `NotImplementedError: Expressions should define either _layer (full dictionary) or _task (single task).  This expression <class 'dask.dataframe.dask_expr._expr.FilterAlign'> defines neither` exception.

The problem arises after the second `simplify` in `optimize_until` which generates another expression that is not fully lowered

Pull request:
Fix expression getitem to avoid alignment
- [ ] Closes #11759
- [ ] Tests added / passed
- [ ] Passes `pre-commit run --all-files`


this was a mess 


Patch:
diff --git a/dask/array/creation.py b/dask/array/creation.py
index 7a1edb2a0..b906bc405 100644
--- a/dask/array/creation.py
+++ b/dask/array/creation.py
@@ -452,7 +452,7 @@ def arange(
         blockstop = start + ((elem_count + bs) * step)
         task = Task(
             (name, i),
-            partial(chunk.arange, like=like),
+            partial(chunk.arange, like=meta),
             blockstart,
             blockstop,
             step,
diff --git a/dask/dataframe/dask_expr/_collection.py b/dask/dataframe/dask_expr/_collection.py
index 8335aba4d..4a2656bc5 100644
--- a/dask/dataframe/dask_expr/_collection.py
+++ b/dask/dataframe/dask_expr/_collection.py
@@ -397,6 +397,8 @@ class FrameBase(DaskMethodsMixin):
 
     def __getitem__(self, other):
         if isinstance(other, FrameBase):
+            if not expr.are_co_aligned(self.expr, other.expr):
+                return new_collection(expr.FilterAlign(self, other))
             return new_collection(self.expr.__getitem__(other.expr))
         elif isinstance(other, slice):
             from pandas.api.types import is_float_dtype
diff --git a/dask/dataframe/dask_expr/_expr.py b/dask/dataframe/dask_expr/_expr.py
index 7093e7235..763f07a1b 100644
--- a/dask/dataframe/dask_expr/_expr.py
+++ b/dask/dataframe/dask_expr/_expr.py
@@ -178,8 +178,6 @@ class Expr(core.Expr):
 
     def __getitem__(self, other):
         if isinstance(other, Expr):
-            if not are_co_aligned(self, other):
-                return FilterAlign(self, other)
             return Filter(self, other)
         else:
             return Projection(self, other)  # df[["a", "b", "c"]]
@@ -522,7 +520,9 @@ class Expr(core.Expr):
     def _filter_simplification(self, parent, predicate=None):
         if predicate is None:
             predicate = parent.predicate.substitute(self, self.frame)
-        return type(self)(self.frame[predicate], *self.operands[1:])
+        if are_co_aligned(self.frame, predicate):
+            # Only do this if we are aligned
+            return type(self)(self.frame[predicate], *self.operands[1:])
 
     def fuse(self):
         return optimize_blockwise_fusion(self)
@@ -1631,14 +1631,26 @@ class Where(Elemwise):
 
 
 def _check_divisions(df, i, division_min, division_max, last):
+    if not len(df):
+        return df
+    if is_index_like(df):
+        index = df
+    else:
+        try:
+            index = df.index.get_level_values(0)
+        except AttributeError:
+            index = df.index
     # Check divisions
-    real_min = df.index.min()
-    real_max = df.index.max()
+    real_min = index.min()
+    real_max = index.max()
     # Upper division of the last partition is often set to
     # the max value. For all other partitions, the upper
     # division should be greater than the maximum value.
-    valid_min = real_min >= division_min
-    valid_max = (real_max <= division_max) if last else (real_max < division_max)
+    valid_min = valid_max = True
+    if not pd.isna(division_min):
+        valid_min = real_min >= division_min
+    if not pd.isna(division_max):
+        valid_max = (real_max <= division_max) if last else (real_max < division_max)
     if not (valid_min and valid_max):
         raise RuntimeError(
             f"`enforce_runtime_divisions` failed for partition {i}."
@@ -1665,7 +1677,7 @@ class EnforceRuntimeDivisions(Blockwise):
             self.divisions[index + 1],
             index == (self.npartitions - 1),
         ]
-        return (self.operation,) + tuple(args)  # type: ignore
+        return Task(name, self.operation, *args)
 
 
 class Abs(Elemwise):
diff --git a/dask/dataframe/dask_expr/_repartition.py b/dask/dataframe/dask_expr/_repartition.py
index fe13dd362..d5a9dd5d7 100644
--- a/dask/dataframe/dask_expr/_repartition.py
+++ b/dask/dataframe/dask_expr/_repartition.py
@@ -149,6 +149,10 @@ class Repartition(Expr):
         if isinstance(parent, Filter) and self._filter_passthrough_available(
             parent, dependents
         ):
+            if self._name == parent.predicate._name:
+                # We shouldn't push through the predicate, these pushdowns should
+                # always come from frame.
+                return
             return self._filter_simplification(parent)
         if isinstance(parent, Projection):
             return plain_column_projection(self, parent, dependents)
diff --git a/dask/dataframe/utils.py b/dask/dataframe/utils.py
index 9e156c41c..7260ae0d6 100644
--- a/dask/dataframe/utils.py
+++ b/dask/dataframe/utils.py
@@ -15,7 +15,7 @@ import pandas as pd
 from pandas.api.types import is_dtype_equal
 
 import dask
-from dask.base import get_scheduler, is_dask_collection
+from dask.base import is_dask_collection
 from dask.core import get_deps
 from dask.dataframe._compat import tm  # noqa: F401
 from dask.dataframe.dispatch import (  # noqa : F401
@@ -603,24 +603,7 @@ def assert_divisions(ddf, scheduler=None):
     if not getattr(ddf, "known_divisions", False):
         return
 
-    def index(x):
-        if is_index_like(x):
-            return x
-        try:
-            return x.index.get_level_values(0)
-        except AttributeError:
-            return x.index
-
-    get = get_scheduler(scheduler=scheduler, collections=[type(ddf)])
-    results = get(ddf.dask, ddf.__dask_keys__())
-    for i, df in enumerate(results[:-1]):
-        if len(df):
-            assert index(df).min() >= ddf.divisions[i]
-            assert index(df).max() < ddf.divisions[i + 1]
-
-    if len(results[-1]):
-        assert index(results[-1]).min() >= ddf.divisions[-2]
-        assert index(results[-1]).max() <= ddf.divisions[-1]
+    ddf.enforce_runtime_divisions().compute(scheduler=scheduler)
 
 
 def assert_sane_keynames(ddf):


Test patch:
diff --git a/dask/array/tests/test_creation.py b/dask/array/tests/test_creation.py
index 9887886c9..fba157f8a 100644
--- a/dask/array/tests/test_creation.py
+++ b/dask/array/tests/test_creation.py
@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+import pickle
+
 import pytest
 
 from dask._task_spec import Alias
@@ -1093,3 +1095,18 @@ def test_nan_full_like(val, shape_chunks, dtype):
         da.full_like(y1, val, dtype=dtype),
         np.full_like(y2, val, dtype=dtype),
     )
+
+
+@pytest.mark.parametrize(
+    "func", [da.array, da.asarray, da.asanyarray, da.arange, da.tri]
+)
+def test_like_forgets_graph(func):
+    """Test that array creation functions with like=x do not
+    internally store the graph of x
+    """
+    x = da.arange(3).map_blocks(lambda x: x)
+    with pytest.raises(Exception, match="local object"):
+        pickle.dumps(x)
+
+    a = func(1, like=x)
+    pickle.dumps(a)
diff --git a/dask/dataframe/dask_expr/tests/test_collection.py b/dask/dataframe/dask_expr/tests/test_collection.py
index af5bba927..6f241ff38 100644
--- a/dask/dataframe/dask_expr/tests/test_collection.py
+++ b/dask/dataframe/dask_expr/tests/test_collection.py
@@ -2786,6 +2786,22 @@ def test_to_backend_simplify():
         assert str(df3.expr) == str(df[["y"]].expr)
 
 
+def test_getitem_triggering_unnecessary_alignment():
+    df = pd.DataFrame(
+        {
+            "A": [1, 2, 3, 4, 5, 6, 7, 8, 9],
+            "B": [9, 8, 7, 6, 5, 4, 3, 2, 1],
+            "C": [True, False, True] * 3,
+        },
+        columns=list("ABC"),
+    )
+    ddf = from_pandas(df, 2)
+    result = ddf[ddf.C]
+    expr = result.expr.optimize()
+    assert expr.__dask_graph__()
+    assert_eq(result, df[df.C])
+
+
 def test_projection_on_series():
     pdf = pd.DataFrame(data={"a": [1, 3, 2]}).a
     df = from_pandas(pdf, npartitions=1)