Make sure 'order' kwarg will not crash the 'astype' method in dask (#9317)

GenevieveBuckley · web-flow · commit 377940b64914 · 2022-08-04T09:47:35.000-04:00
Allows the `order` kwarg to be passed in to the dask `astype` method without triggering an error. Our friends over at zarr have found a small bug in dask. While the numpy `astype` method allows the user to use an `order` keyword argument ([docs here](https://numpy.org/doc/stable/reference/generated/numpy.ndarray.astype.html)), the corresponding dask `astype` method produces an error Why we should do this: 1. It's not always as obvious as telling the user to edit the line in their code that uses `astype`, since it's often used very indirectly. Here's one example: zarr-developers/zarr-python#962 (comment) 2. It reflects better on dask not to have odd errors popping up, even if this PR won't completely solve the issue being discussed over at zarr zarr-developers/zarr-python#962
diff --git a/dask/array/core.py b/dask/array/core.py
@@ -2233,7 +2233,14 @@ def astype(self, dtype, **kwargs):
             By default, astype always returns a newly allocated array. If this
             is set to False and the `dtype` requirement is satisfied, the input
             array is returned instead of a copy.
+
+            .. note::
+
+                Dask does not respect the contiguous memory layout of the array,
+                and will ignore the ``order`` keyword argument.
+                The default order is 'C' contiguous.
         """
+        kwargs.pop("order", None)  # `order` is not respected, so we remove this kwarg
         # Scalars don't take `casting` or `copy` kwargs - as such we only pass
         # them to `map_blocks` if specified by user (different than defaults).
         extra = set(kwargs) - {"casting", "copy"}
diff --git a/dask/array/tests/test_array_core.py b/dask/array/tests/test_array_core.py
@@ -2286,6 +2286,34 @@ def test_astype():
     assert d.astype("f8") is d
 
 
+def test_astype_gh1151():
+    a = np.arange(5).astype(np.int32)
+    b = da.from_array(a, (1,))
+    assert_eq(a.astype(np.int16), b.astype(np.int16))
+
+
+def test_astype_gh9318():
+    # `order`` kwarg in `astype` should not cause an error
+    a = np.array([[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12]], order="C")
+    b = da.from_array(a, chunks=(2, 2))
+    result_a = a.astype(float, order="F")
+    result_b = b.astype(float, order="F")  # if no error at this line, pytest passes
+    assert_eq(result_a, result_b)  # won't check the order matches, but checks results
+
+
+@pytest.mark.xfail(reason="Github issue https://github.com/dask/dask/issues/9316")
+def test_astype_gh9316():
+    # Issue https://github.com/dask/dask/issues/9316
+    # Can be combined with test_astype_gh9318 above when XFAIL marker is removed
+    a = np.array([[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12]], order="C")
+    b = da.from_array(a, chunks=(2, 2))
+    result_a = a.astype(float, order="F")
+    result_b = b.astype(float, order="F")
+    result_b = result_b.compute()
+    assert result_a.flags.c_contiguous == result_b.flags.c_contiguous
+    assert result_a.flags.f_contiguous == result_b.flags.f_contiguous
+
+
 def test_arithmetic():
     x = np.arange(5).astype("f4") + 2
     y = np.arange(5).astype("i8") + 2
@@ -3569,12 +3597,6 @@ def test_npartitions():
     assert da.ones((5, 5), chunks=(2, 3)).npartitions == 6
 
 
-def test_astype_gh1151():
-    a = np.arange(5).astype(np.int32)
-    b = da.from_array(a, (1,))
-    assert_eq(a.astype(np.int16), b.astype(np.int16))
-
-
 def test_elemwise_name():
     assert (da.ones(5, chunks=2) + 1).name.startswith("add-")
 
