Repository: dask/dask

Issue:
Dask 2024.8 started failing when indexing result of numpy.flatnonzero
<!-- 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**:

scikit-image's CI started failing with Dask 2024.8 when it worked fine before with 2024.7. As described in https://github.com/scikit-image/scikit-image/issues/7491, we traced the failure 
```plain
ValueError: Array chunk size or shape is unknown. 
```
to a line that tries to index the result of a `numpy.flatnonzero` operation. See the MCVE below which fails for 2024.8 but works fine for 2024.7. 

**Minimal Complete Verifiable Example**:

```python
import numpy as np
import dask.array as da

hist = da.from_array(np.arange(256, dtype=int), chunks=(256,))
result = np.flatnonzero(hist)
result[[0, -1]]  # ValueError: Array chunk size or shape is unknown. 
```

**Anything else we need to know?**:

From reading the [docs that are pointed out by the error message](https://docs.dask.org/en/latest/array-chunks.html#unknown-chunks), it seems this behavior might not be entirely unexpected and this might not actually be a bug. However, I'm unsure why it worked before with 2024.7. And since I currently don't really now how to address this on scikit-image's side without adding an extra code path to support Dask 2024.8, I'd be grateful for feedback and context.

Sorry if I missed any existing issues or docs that already describe or explain this.

**Environment**:

- Dask version: 2024.8.0
- Python version: 3.12.0
- Operating System: Linux 6.10.3-arch1-1 (64-bit)
- Install method (conda, pip, source): pip


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._

     13 files  ± 0       13 suites  ±0   3h 1m 30s :stopwatch: + 5m 2s
 13 132 tests + 1   12 063 :white_check_mark: ±0   1 068 :zzz: ±0  1 :x: +1 
136 838 runs  +11  117 770 :white_check_mark: +9  19 067 :zzz: +1  1 :x: +1 

For more details on these failures, see [this check](https://github.com/dask/dask/runs/28645524931).

Results for commit fd9ec834. ± Comparison against base commit 1ba9c5bf.

[test-results]:data:application/gzip;base64,H4sIAKPguWYC/02Myw6DIBQFf8Ww7oILeMH+TIMICamPBmHV9N8LPtDlzMmZL3F+tCt5NsAfDVmTjxcNKejol7kwVR3NKq9x34Gzk19rMqZIRpFf8u0/2xNVdU77sbgqbAhLyKa0Q5r3NCquDq5pkFLSKo90R1Ge7pbe+F42yzT5mIG4obNGcTG0wvagWq2RMWFAoqVMSXAgNWDfk98fxzxUAxoBAAA=


Patch:
diff --git a/dask/array/_shuffle.py b/dask/array/_shuffle.py
index 37aacb4bb..89f435d7b 100644
--- a/dask/array/_shuffle.py
+++ b/dask/array/_shuffle.py
@@ -8,6 +8,7 @@ import numpy as np
 from dask import config
 from dask.array.chunk import getitem
 from dask.array.core import Array, unknown_chunk_message
+from dask.array.dispatch import concatenate_lookup
 from dask.base import tokenize
 from dask.highlevelgraph import HighLevelGraph
 
@@ -187,7 +188,8 @@ def _shuffle(chunks, indexer, axis, in_name, out_name, token):
 
 
 def concatenate_arrays(arrs, sorter, axis):
-    return np.take(np.concatenate(arrs, axis=axis), np.argsort(sorter), axis=axis)
+    concatenate = concatenate_lookup.dispatch(type(arrs[0]))
+    return np.take(concatenate(arrs, axis=axis), np.argsort(sorter), axis=axis)
 
 
 def convert_key(key, chunk, axis):
diff --git a/dask/array/core.py b/dask/array/core.py
index b710fe491..3009dcb11 100644
--- a/dask/array/core.py
+++ b/dask/array/core.py
@@ -4524,6 +4524,13 @@ def retrieve_from_ooc(
     return load_dsk
 
 
+def _as_dtype(a, dtype):
+    if dtype is None:
+        return a
+    else:
+        return a.astype(dtype)
+
+
 def asarray(
     a, allow_unknown_chunksizes=False, dtype=None, order=None, *, like=None, **kwargs
 ):
@@ -4573,16 +4580,22 @@ def asarray(
     >>> y = [[1, 2, 3], [4, 5, 6]]
     >>> da.asarray(y)
     dask.array<array, shape=(2, 3), dtype=int64, chunksize=(2, 3), chunktype=numpy.ndarray>
+
+    .. warning::
+        `order` is ignored if `a` is an `Array`, has the attribute ``to_dask_array``,
+        or is a list or tuple of `Array`'s.
     """
     if like is None:
         if isinstance(a, Array):
-            return a
+            return _as_dtype(a, dtype)
         elif hasattr(a, "to_dask_array"):
-            return a.to_dask_array()
+            return _as_dtype(a.to_dask_array(), dtype)
         elif type(a).__module__.split(".")[0] == "xarray" and hasattr(a, "data"):
-            return asarray(a.data)
+            return _as_dtype(asarray(a.data, order=order), dtype)
         elif isinstance(a, (list, tuple)) and any(isinstance(i, Array) for i in a):
-            return stack(a, allow_unknown_chunksizes=allow_unknown_chunksizes)
+            return _as_dtype(
+                stack(a, allow_unknown_chunksizes=allow_unknown_chunksizes), dtype
+            )
         elif not isinstance(getattr(a, "shape", None), Iterable):
             a = np.asarray(a, dtype=dtype, order=order)
     else:
@@ -4641,16 +4654,20 @@ def asanyarray(a, dtype=None, order=None, *, like=None, inline_array=False):
     >>> y = [[1, 2, 3], [4, 5, 6]]
     >>> da.asanyarray(y)
     dask.array<array, shape=(2, 3), dtype=int64, chunksize=(2, 3), chunktype=numpy.ndarray>
+
+    .. warning::
+        `order` is ignored if `a` is an `Array`, has the attribute ``to_dask_array``,
+        or is a list or tuple of `Array`'s.
     """
     if like is None:
         if isinstance(a, Array):
-            return a
+            return _as_dtype(a, dtype)
         elif hasattr(a, "to_dask_array"):
-            return a.to_dask_array()
+            return _as_dtype(a.to_dask_array(), dtype)
         elif type(a).__module__.split(".")[0] == "xarray" and hasattr(a, "data"):
-            return asanyarray(a.data)
+            return _as_dtype(asarray(a.data, order=order), dtype)
         elif isinstance(a, (list, tuple)) and any(isinstance(i, Array) for i in a):
-            return stack(a)
+            return _as_dtype(stack(a), dtype)
         elif not isinstance(getattr(a, "shape", None), Iterable):
             a = np.asanyarray(a, dtype=dtype, order=order)
     else:
diff --git a/dask/array/slicing.py b/dask/array/slicing.py
index f9add043a..3f345cfb8 100644
--- a/dask/array/slicing.py
+++ b/dask/array/slicing.py
@@ -580,6 +580,15 @@ def take(outname, inname, chunks, index, axis=0):
         )
         chunks, graph = _shuffle(chunks, indexer, axis, inname, outname, token)
         return chunks, graph
+    elif len(chunks[axis]) == 1:
+        slices = [slice(None)] * len(chunks)
+        slices[axis] = list(index)
+        slices = tuple(slices)
+        chunk_tuples = list(product(*(range(len(c)) for i, c in enumerate(chunks))))
+        dsk = {
+            (outname,) + ct: (getitem, (inname,) + ct, slices) for ct in chunk_tuples
+        }
+        return chunks, dsk
     else:
         from dask.array.core import unknown_chunk_message
 
diff --git a/docs/source/spark.rst b/docs/source/spark.rst
index 2eabde8ad..c6ee3e240 100644
--- a/docs/source/spark.rst
+++ b/docs/source/spark.rst
@@ -19,6 +19,7 @@ Generally Dask is smaller and lighter weight than Spark.  This means that it
 has fewer features and, instead, is used in conjunction with other libraries,
 particularly those in the numeric Python ecosystem.  It couples with libraries
 like Pandas or Scikit-Learn to achieve high-level functionality.
+Additionally, Dask is often faster and more robustly performant on `standard benchmarks <https://docs.coiled.io/blog/tpch.html?utm_source=dask-docs&utm_medium=spark-vs-dask#dask-vs-spark>`_ than Spark.
 
 Language
 ~~~~~~~~


Test patch:
diff --git a/dask/array/tests/test_array_core.py b/dask/array/tests/test_array_core.py
index 70eeb0a9f..d975cfced 100644
--- a/dask/array/tests/test_array_core.py
+++ b/dask/array/tests/test_array_core.py
@@ -2704,6 +2704,18 @@ def test_asarray(asarray):
     assert_eq(asarray(y), x)
 
 
+@pytest.mark.parametrize("asarray", [da.asarray, da.asanyarray])
+def test_asarray_array_dtype(asarray):
+    # test array input
+    x = asarray([1, 2])
+    assert_eq(asarray(x, dtype=da.float32), np.asarray(x, dtype=np.float32))
+
+    x = asarray(x, dtype=da.float64)
+    assert x.dtype == da.float64
+    x = asarray(x, dtype=da.int32)
+    assert x.dtype == da.int32
+
+
 @pytest.mark.parametrize("asarray", [da.asarray, da.asanyarray])
 def test_asarray_dask_dataframe(asarray):
     # https://github.com/dask/dask/issues/3885
diff --git a/dask/array/tests/test_slicing.py b/dask/array/tests/test_slicing.py
index ed8a99e0c..0e9cfce28 100644
--- a/dask/array/tests/test_slicing.py
+++ b/dask/array/tests/test_slicing.py
@@ -1049,6 +1049,29 @@ def test_shuffle_slice(size, chunks):
     assert_eq(a, b)
 
 
+def test_unknown_chunks_length_one():
+    a = np.arange(256, dtype=int)
+    arr = da.from_array(a, chunks=(256,))
+    # np.flatnonzero dispatches
+    result = np.flatnonzero(arr)
+    assert_eq(result[[0, -1]], np.flatnonzero(a)[[0, -1]])
+
+    result = da.flatnonzero(arr)
+    assert_eq(result[[0, -1]], np.flatnonzero(a)[[0, -1]])
+
+    a = a.reshape(16, 16)
+    arr = da.from_array(a, chunks=(8, 16))
+    arr._chunks = ((8, 8), (np.nan,))
+    result = arr[:, [0, -1]]
+    expected = a[:, [0, -1]]
+    assert_eq(result, expected)
+
+    arr = da.from_array(a, chunks=(8, 8))
+    arr._chunks = ((8, 8), (np.nan, np.nan))
+    with pytest.raises(ValueError, match="Array chunk size or shape"):
+        arr[:, [0, -1]]
+
+
 @pytest.mark.parametrize("lock", [True, False])
 @pytest.mark.parametrize("asarray", [True, False])
 @pytest.mark.parametrize("fancy", [True, False])
@@ -1069,6 +1092,12 @@ def test_slice_array_3d_with_bool_numpy_array():
     assert_eq(actual, expected)
 
 
+def test_slice_masked_arrays():
+    arr = np.ma.array(range(8), mask=[0, 0, 1, 0, 0, 1, 0, 1])
+    darr = da.from_array(arr, chunks=(4, 4))
+    assert_eq(darr[[2, 6]], arr[[2, 6]])
+
+
 def test_slice_array_null_dimension():
     array = da.from_array(np.zeros((3, 0)))
     expected = np.zeros((3, 0))[[0]]