Repository: dask/dask

Issue:
Array slicing with boolean mask and `np.newaxis` stopped working at some point
<!-- Please include a self-contained copy-pastable example that generates the issue if possible.

Please be concise with code posted. See guidelines below on how to provide a good bug report:

- Craft Minimal Bug Reports http://matthewrocklin.com/blog/work/2018/02/28/minimal-bug-reports
- Minimal Complete Verifiable Examples https://stackoverflow.com/help/mcve

Bug reports that follow these guidelines are easier to diagnose, and so are often handled much more quickly.
-->

**Describe the issue**:
Complex array slicing operations that use both boolean masks and `np.newaxis` used to work (as of Dask 2024.5.0), but in newer versions of Dask (tested with 2024.12.1) they fail with an error. Either type of operation works on its own, it's only when they are combined that the exception appears. I discovered this because some older code I had written broke with the new version.

**Minimal Complete Verifiable Example**:

```python
import numpy as np
import dask.array as da

arr = da.array([0, 1, 2])
new = arr[[True, True, False], np.newaxis]
```
In Dask 2024.5.0, this succeeds, producing an array with shape `(2, 1)`, but in Dask 2024.12.1, it leads to a `TypeError` with this traceback:
```
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
Cell In[5], line 1
----> 1 arr[[True, True, False], np.newaxis]

File ~/mambaforge/envs/dask-test/lib/python3.12/site-packages/dask/array/core.py:2006, in Array.__getitem__(self, index)
   2003     return self
   2005 out = "getitem-" + tokenize(self, index2)
-> 2006 dsk, chunks = slice_array(out, self.name, self.chunks, index2, self.itemsize)
   2008 graph = HighLevelGraph.from_collections(out, dsk, dependencies=[self])
   2010 meta = meta_from_array(self._meta, ndim=len(chunks))

File ~/mambaforge/envs/dask-test/lib/python3.12/site-packages/dask/array/slicing.py:177, in slice_array(out_name, in_name, blockdims, index, itemsize)
    174 index += (slice(None, None, None),) * missing
    176 # Pass down to next function
--> 177 dsk_out, bd_out = slice_with_newaxes(out_name, in_name, blockdims, index)
    179 bd_out = tuple(map(tuple, bd_out))
    180 return dsk_out, bd_out

File ~/mambaforge/envs/dask-test/lib/python3.12/site-packages/dask/array/slicing.py:212, in slice_with_newaxes(out_name, in_name, blockdims, index)
    210 if k[0] == out_name:
    211     k2 = (out_name,) + expand(k[1:], 0)
--> 212     dsk2[k2] = Task(k2, v.func, v.args[0], expand_orig(v.args[1], None))
    213 else:
    214     dsk2[k] = v

File <dynamic>:2, in expand(seq, val)

TypeError: 'Alias' object is not subscriptable
```

**Anything else we need to know?**:

**Environment**:

- Dask version: 2024.12.1
- Python version: 3.12.8
- Operating System: Ubuntu 24.04
- Install method (conda, pip, source): Conda


Discussion:
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 26m 49s ⏱️ + 6m 44s
 17 142 tests + 1   15 962 ✅ + 1   1 180 💤 ±0  0 ❌ ±0 
211 225 runs  +13  194 202 ✅ +15  17 023 💤  - 2  0 ❌ ±0 

Results for commit 3b25bb64. ± Comparison against base commit d8eb38e9.

[test-results]:data:application/gzip;base64,H4sIANKJk2cC/02NSQ7DIAxFrxKx7sKYQOJepmKUUDNUBFZV714SZVq+9+XnLwtx8At7Nlw+GraUmC9yJekc52llBUBV1TVve8dbPPi1FGu3I1I3+Y6fVfIeThd0HKq7hE9pTrtJZVrTyDmi3PlMU4uAp9zTHaA43C298b1s53GMuQITBqUxqrXG+ED1EYFDEEEjKUEe+s6hoN6y3x9wTPXdGgEAAA==


Patch:
diff --git a/dask/array/slicing.py b/dask/array/slicing.py
index 0927e2613..473c42646 100644
--- a/dask/array/slicing.py
+++ b/dask/array/slicing.py
@@ -12,7 +12,7 @@ import numpy as np
 from tlz import concat, memoize, merge, pluck
 
 from dask import core
-from dask._task_spec import Alias, Task, TaskRef
+from dask._task_spec import Alias, DataNode, Task, TaskRef
 from dask.array.chunk import getitem
 from dask.base import is_dask_collection, tokenize
 from dask.highlevelgraph import HighLevelGraph
@@ -209,7 +209,16 @@ def slice_with_newaxes(out_name, in_name, blockdims, index):
         for k, v in dsk.items():
             if k[0] == out_name:
                 k2 = (out_name,) + expand(k[1:], 0)
-                dsk2[k2] = Task(k2, v.func, v.args[0], expand_orig(v.args[1], None))
+                if isinstance(v.args[1], Alias):
+                    # positional indexing with newaxis
+                    indexer = expand_orig(dsk[v.args[1].key].value[1], None)
+                    tok = "shuffle-taker-" + tokenize(indexer)
+                    dsk2[tok] = DataNode(tok, (1, indexer))
+                    arg = TaskRef(tok)
+                else:
+                    arg = expand_orig(v.args[1], None)
+                # raise NotImplementedError
+                dsk2[k2] = Task(k2, v.func, v.args[0], arg)
             else:
                 dsk2[k] = v
 
diff --git a/dask/dataframe/dask_expr/_reductions.py b/dask/dataframe/dask_expr/_reductions.py
index 5b5f248dd..5ba667896 100644
--- a/dask/dataframe/dask_expr/_reductions.py
+++ b/dask/dataframe/dask_expr/_reductions.py
@@ -1427,6 +1427,7 @@ class ValueCounts(ReductionConstantDim):
     reduction_chunk = M.value_counts
     reduction_aggregate = methods.value_counts_aggregate
     reduction_combine = methods.value_counts_combine
+    split_by = None
 
     @functools.cached_property
     def _meta(self):
@@ -1441,8 +1442,8 @@ class ValueCounts(ReductionConstantDim):
             return func(_concat(inputs), observed=True, **kwargs)
 
     @property
-    def split_by(self):
-        return self.frame._meta.name
+    def shuffle_by_index(self):
+        return True
 
     @property
     def chunk_kwargs(self):


Test patch:
diff --git a/dask/array/tests/test_array_core.py b/dask/array/tests/test_array_core.py
index e72ec9d30..bf2171d00 100644
--- a/dask/array/tests/test_array_core.py
+++ b/dask/array/tests/test_array_core.py
@@ -5145,9 +5145,16 @@ def test_dask_array_holds_scipy_sparse_containers():
 @pytest.mark.parametrize("sparse_module_path", ["scipy.sparse", "cupyx.scipy.sparse"])
 def test_scipy_sparse_indexing(index, sparse_module_path):
     sp = pytest.importorskip(sparse_module_path)
-    x = da.random.default_rng().random((1000, 10), chunks=(100, 10))
-    x[x < 0.9] = 0
-    y = x.map_blocks(sp.csr_matrix)
+
+    if sparse_module_path == "cupyx.scipy.sparse":
+        backend = "cupy"
+    else:
+        backend = "numpy"
+
+    with dask.config.set({"array.backend": backend}):
+        x = da.random.default_rng().random((1000, 10), chunks=(100, 10))
+        x[x < 0.9] = 0
+        y = x.map_blocks(sp.csr_matrix)
 
     assert not (
         x[index, :].compute(scheduler="single-threaded")
diff --git a/dask/array/tests/test_slicing.py b/dask/array/tests/test_slicing.py
index 46366c2a8..28eef7841 100644
--- a/dask/array/tests/test_slicing.py
+++ b/dask/array/tests/test_slicing.py
@@ -1131,3 +1131,9 @@ def test_vindex_with_dask_array():
 
     with pytest.raises(IndexError, match=msg):
         darr.vindex[(dindexer, None)]
+
+
+def test_positional_indexer_newaxis():
+    arr = da.array([0, 1, 2])
+    new = arr[[True, True, False], np.newaxis]
+    assert_eq(new, arr.compute()[[True, True, False], np.newaxis])
diff --git a/dask/dataframe/dask_expr/tests/test_reductions.py b/dask/dataframe/dask_expr/tests/test_reductions.py
index b10030ac0..744e648cb 100644
--- a/dask/dataframe/dask_expr/tests/test_reductions.py
+++ b/dask/dataframe/dask_expr/tests/test_reductions.py
@@ -558,3 +558,16 @@ def test_reductions_timestamps_display():
     data = pd.to_datetime(["2024-10-02 12:00:00", "2024-10-02 14:00:00"])
     df = from_pandas(pd.DataFrame({"valid_time": data}))
     assert df["valid_time"].min().__repr__()
+
+
+def test_value_counts_shuffle_properly():
+    pdf = pd.DataFrame(
+        {
+            "A": [1, 2, 3, 4, 5, 6, 7, 8, 9, 10],
+            "B": [11, 12, 13, 4, 15, 6, 17, 8, 19, 10],
+        }
+    )
+    df = from_pandas(pdf, npartitions=2)
+    result = (df["A"] == df["B"]).value_counts()
+    expected = (pdf["A"] == pdf["B"]).value_counts()
+    assert_eq(result, expected)