|
13 | 13 |
|
14 | 14 |
|
15 | 15 | if TYPE_CHECKING: |
16 | | - from typing import TypeAlias |
| 16 | + from collections.abc import Iterable |
| 17 | + from typing import Literal, TypeAlias |
17 | 18 |
|
18 | 19 | from fast_array_utils.typing import CpuArray, DiskArray, GpuArray |
19 | 20 | from testing.fast_array_utils import ArrayType |
20 | 21 |
|
21 | 22 | Array: TypeAlias = CpuArray | GpuArray | DiskArray | types.CSDataset | types.DaskArray |
| 23 | + ExtendedArray: TypeAlias = Array | types.COOBase | types.CupyCOOMatrix |
22 | 24 |
|
23 | 25 |
|
24 | 26 | WARNS_NUMBA = pytest.warns(RuntimeWarning, match="numba is not installed; falling back to slow conversion") |
25 | 27 |
|
26 | 28 |
|
27 | 29 | @pytest.mark.parametrize("to_cpu_memory", [True, False], ids=["to_cpu_memory", "not_to_cpu_memory"]) |
28 | | -def test_to_dense(array_type: ArrayType[Array], *, to_cpu_memory: bool) -> None: |
| 30 | +@pytest.mark.parametrize("order", argvalues=["K", "C", "F"]) # “A” behaves like “K” |
| 31 | +def test_to_dense(array_type: ArrayType[Array], *, order: Literal["K", "C", "F"], to_cpu_memory: bool) -> None: |
29 | 32 | x = array_type([[1, 2, 3], [4, 5, 6]], dtype=np.float32) |
30 | 33 | if not to_cpu_memory and array_type.cls in {types.CSCDataset, types.CSRDataset}: |
31 | 34 | with pytest.raises(ValueError, match="to_cpu_memory must be True if x is an CS{R,C}Dataset"): |
32 | | - to_dense(x, to_cpu_memory=to_cpu_memory) |
| 35 | + to_dense(x, order=order, to_cpu_memory=to_cpu_memory) |
33 | 36 | return |
34 | 37 |
|
35 | | - with WARNS_NUMBA if issubclass(array_type.cls, types.CSBase) and not find_spec("numba") else nullcontext(): |
36 | | - arr = to_dense(x, to_cpu_memory=to_cpu_memory) |
| 38 | + with ( |
| 39 | + pytest.warns(RuntimeWarning, match="Dask can not be made to emit F-contiguous arrays") |
| 40 | + if (order == "F" and array_type.cls is types.DaskArray) |
| 41 | + else nullcontext(), |
| 42 | + WARNS_NUMBA if issubclass(array_type.cls, types.CSBase) and not find_spec("numba") else nullcontext(), |
| 43 | + ): |
| 44 | + arr = to_dense(x, order=order, to_cpu_memory=to_cpu_memory) |
| 45 | + |
37 | 46 | assert_expected_cls(x, arr, to_cpu_memory=to_cpu_memory) |
38 | 47 | assert arr.shape == (2, 3) |
| 48 | + # Dask is unreliable: for explicit “F”, we emit a warning (tested above), for “K” we just ignore the result |
| 49 | + if not (array_type.cls is types.DaskArray and order in {"F", "K"}): |
| 50 | + assert_expected_order(x, arr, order=order) |
39 | 51 |
|
40 | 52 |
|
41 | 53 | @pytest.mark.parametrize("to_cpu_memory", [True, False], ids=["to_cpu_memory", "not_to_cpu_memory"]) |
42 | | -def test_to_dense_extra(coo_matrix_type: ArrayType[Array], *, to_cpu_memory: bool) -> None: |
| 54 | +@pytest.mark.parametrize("order", argvalues=["K", "C", "F"]) # “A” behaves like “K” |
| 55 | +def test_to_dense_extra(coo_matrix_type: ArrayType[types.COOBase | types.CupyCOOMatrix], *, order: Literal["K", "C", "F"], to_cpu_memory: bool) -> None: |
43 | 56 | src_mtx = coo_matrix_type([[1, 2, 3], [4, 5, 6]], dtype=np.float32) |
| 57 | + |
44 | 58 | with WARNS_NUMBA if not find_spec("numba") else nullcontext(): |
45 | | - arr = to_dense(src_mtx, to_cpu_memory=to_cpu_memory) |
| 59 | + arr = to_dense(src_mtx, order=order, to_cpu_memory=to_cpu_memory) |
| 60 | + |
46 | 61 | assert_expected_cls(src_mtx, arr, to_cpu_memory=to_cpu_memory) |
47 | 62 | assert arr.shape == (2, 3) |
| 63 | + assert_expected_order(src_mtx, arr, order=order) |
48 | 64 |
|
49 | 65 |
|
50 | | -def assert_expected_cls(orig: Array, converted: Array, *, to_cpu_memory: bool) -> None: |
| 66 | +def assert_expected_cls(orig: ExtendedArray, converted: Array, *, to_cpu_memory: bool) -> None: |
51 | 67 | match (to_cpu_memory, orig): |
52 | 68 | case False, types.DaskArray(): |
53 | 69 | assert isinstance(converted, types.DaskArray) |
54 | | - assert_expected_cls(orig._meta, converted._meta, to_cpu_memory=to_cpu_memory) # noqa: SLF001 |
| 70 | + assert_expected_cls(orig.compute(), converted.compute(), to_cpu_memory=to_cpu_memory) |
55 | 71 | case False, types.CupyArray() | types.CupySpMatrix(): |
56 | 72 | assert isinstance(converted, types.CupyArray) |
57 | 73 | case _: |
58 | 74 | assert isinstance(converted, np.ndarray) |
| 75 | + |
| 76 | + |
| 77 | +def assert_expected_order(orig: ExtendedArray, converted: Array, *, order: Literal["K", "C", "F"]) -> None: |
| 78 | + match converted: |
| 79 | + case types.CupyArray() | np.ndarray(): |
| 80 | + orders = {order_exp: converted.flags[f"{order_exp}_CONTIGUOUS"] for order_exp in (get_orders(orig) if order == "K" else {order})} # type: ignore[index] |
| 81 | + assert any(orders.values()), orders |
| 82 | + case types.DaskArray(): |
| 83 | + assert_expected_order(orig, converted.compute(), order=order) |
| 84 | + case _: |
| 85 | + pytest.fail(f"Unsupported array type: {type(converted)}") |
| 86 | + |
| 87 | + |
| 88 | +def get_orders(orig: ExtendedArray) -> Iterable[Literal["C", "F"]]: |
| 89 | + """Get the orders of an array. |
| 90 | +
|
| 91 | + Numpy arrays with at most one axis of a length >1 are valid in both orders. |
| 92 | + So are COO sparse matrices/arrays. |
| 93 | + """ |
| 94 | + match orig: |
| 95 | + case np.ndarray() | types.CupyArray(): |
| 96 | + if orig.flags.c_contiguous: |
| 97 | + yield "C" |
| 98 | + if orig.flags.f_contiguous: |
| 99 | + yield "F" |
| 100 | + case _ if isinstance(orig, types.CSBase | types.COOBase | types.CupyCSMatrix | types.CupyCOOMatrix | types.CSDataset): |
| 101 | + if orig.format in {"csr", "coo"}: |
| 102 | + yield "C" |
| 103 | + if orig.format in {"csc", "coo"}: |
| 104 | + yield "F" |
| 105 | + case types.DaskArray(): |
| 106 | + yield from get_orders(orig.compute()) |
| 107 | + case types.ZarrArray() | types.H5Dataset(): |
| 108 | + yield "C" |
| 109 | + case _: |
| 110 | + pytest.fail(f"Unsupported array type: {type(orig)}") |
0 commit comments