Repository: dask/dask
Pull request #10850 (2024-01-22 20:13:07)

Fix test for sizeof(pd.MultiIndex) in upstream CI
- Part of #10593 

Output is much smaller than vs. published releases, but still very accurate for non-trivial data sizes

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   3h 14m 3s :stopwatch: - 3m 3s
 12 968 tests ±0   12 036 :white_check_mark: +1     929 :zzz: ±0  3 :x:  - 1 
159 960 runs  ±0  143 376 :white_check_mark: +6  16 576 :zzz:  - 5  8 :x:  - 1 

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

Results for commit 834b40a0. ± Comparison against base commit 762388d0.

[test-results]:data:application/gzip;base64,H4sIAA3UrmUC/03NTQ6DIBCG4asY1l0AA6P0MgZQElJ/GoRV07t3qhW7fJ9JvnmxEKdxY/dG6FvDthLzVUNJNsd1+bZABUR0zftdGuzO7rfi/Y4c8MJHfBIaaSoFGyeiOtSPKa2JhJOkshyfjcGz67ICaLHisSxQX/ab7s7+X/brPMdMwTpQTnHLuZLAtQgCfQsgvdXohgBOW82NV4q9P7BV/csZAQAA


Patch:
diff --git a/dask/array/fft.py b/dask/array/fft.py
index e4cf9742f..acbb2f926 100644
--- a/dask/array/fft.py
+++ b/dask/array/fft.py
@@ -15,7 +15,7 @@ except ImportError:
 from dask.array.core import asarray
 from dask.array.core import concatenate as _concatenate
 from dask.array.creation import arange as _arange
-from dask.array.numpy_compat import _numpy_200
+from dask.array.numpy_compat import NUMPY_GE_200
 from dask.utils import derived_from, skip_doctest
 
 chunk_error = (
@@ -167,7 +167,7 @@ def fft_wrap(fft_func, kind=None, dtype=None):
                 if s is None:
                     axes = tuple(range(a.ndim))
                 else:
-                    if _numpy_200:
+                    if NUMPY_GE_200:
                         # Match deprecation in numpy
                         warnings.warn(
                             "DeprecationWarning: `axes` should not be `None` "
diff --git a/dask/array/numpy_compat.py b/dask/array/numpy_compat.py
index 05e92dd08..5489263ab 100644
--- a/dask/array/numpy_compat.py
+++ b/dask/array/numpy_compat.py
@@ -8,16 +8,18 @@ from packaging.version import parse as parse_version
 from dask.utils import derived_from
 
 _np_version = parse_version(np.__version__)
-_numpy_122 = _np_version >= parse_version("1.22.0")
-_numpy_123 = _np_version >= parse_version("1.23.0")
-_numpy_124 = _np_version >= parse_version("1.24.0")
-_numpy_125 = _np_version.release >= (1, 25, 0)
-_numpy_200 = _np_version.release >= (2, 0, 0)
+NUMPY_GE_122 = _np_version.release >= (1, 22)
+NUMPY_GE_123 = _np_version.release >= (1, 23)
+NUMPY_GE_124 = _np_version.release >= (1, 24)
+NUMPY_GE_125 = _np_version.release >= (1, 25)
+NUMPY_GE_200 = _np_version.release >= (2, 0)
 
 
-if _numpy_200:
+if NUMPY_GE_200:
+    from numpy.exceptions import AxisError, ComplexWarning  # noqa: F401
     from numpy.lib.array_utils import normalize_axis_index, normalize_axis_tuple
 else:
+    from numpy import AxisError, ComplexWarning  # noqa: F401
     from numpy.core.numeric import normalize_axis_index  # type: ignore[attr-defined]
     from numpy.core.numeric import normalize_axis_tuple  # type: ignore[attr-defined]
 
@@ -183,11 +185,7 @@ def rollaxis(a, axis, start=0):
 
 # kwarg is renamed in numpy 1.22.0
 def percentile(a, q, method="linear"):
-    if _numpy_122:
+    if NUMPY_GE_122:
         return np.percentile(a, q, method=method)
     else:
         return np.percentile(a, q, interpolation=method)
-
-
-ComplexWarning = np.exceptions.ComplexWarning if _numpy_200 else np.ComplexWarning
-AxisError = np.exceptions.AxisError if _numpy_200 else np.AxisError
diff --git a/dask/array/percentile.py b/dask/array/percentile.py
index e4ee436a2..a7f347446 100644
--- a/dask/array/percentile.py
+++ b/dask/array/percentile.py
@@ -9,7 +9,7 @@ import numpy as np
 from tlz import merge
 
 from dask.array.core import Array
-from dask.array.numpy_compat import _numpy_122
+from dask.array.numpy_compat import NUMPY_GE_122
 from dask.array.numpy_compat import percentile as np_percentile
 from dask.base import tokenize
 from dask.highlevelgraph import HighLevelGraph
@@ -126,7 +126,7 @@ def percentile(a, q, method="linear", internal_method="default", **kwargs):
         internal_method = method
 
     if "interpolation" in kwargs:
-        if _numpy_122:
+        if NUMPY_GE_122:
             warnings.warn(
                 "In Dask 2022.1.0, the `interpolation=` argument to percentile was renamed to "
                 "`method= ` ",
diff --git a/dask/array/routines.py b/dask/array/routines.py
index f089df2f8..8a5c0f922 100644
--- a/dask/array/routines.py
+++ b/dask/array/routines.py
@@ -29,7 +29,7 @@ from dask.array.core import (
 )
 from dask.array.creation import arange, diag, empty, indices, tri
 from dask.array.einsumfuncs import einsum  # noqa
-from dask.array.numpy_compat import _numpy_200
+from dask.array.numpy_compat import NUMPY_GE_200
 from dask.array.reductions import reduction
 from dask.array.ufunc import multiply, sqrt
 from dask.array.utils import (
@@ -1801,7 +1801,7 @@ def unique(ar, return_index=False, return_inverse=False, return_counts=False):
         # mapping of the original values.
         matches = (ar[:, None] == out["values"][None, :]).astype(np.intp)
         inverse = (matches * out["inverse"]).sum(axis=1)
-        if _numpy_200:
+        if NUMPY_GE_200:
             inverse = inverse.reshape(orig_shape)
         result.append(inverse)
     if return_counts:
diff --git a/dask/array/slicing.py b/dask/array/slicing.py
index e2a6b34a0..20df6fe02 100644
--- a/dask/array/slicing.py
+++ b/dask/array/slicing.py
@@ -559,9 +559,14 @@ def slicing_plan(chunks, index):
 
     if not is_arraylike(index):
         index = np.asanyarray(index)
-    cum_chunks = cached_cumsum(chunks)
 
-    cum_chunks = asarray_safe(cum_chunks, like=index)
+    cum_chunks_tup = cached_cumsum(chunks)
+    cum_chunks = asarray_safe(cum_chunks_tup, like=index)
+    if cum_chunks.dtype.kind != "f":  # Don't cast NaN chunks to int
+        # This is important when index.dtype=uint64 (or uint32 on 32-bit hosts) to
+        # prevent accidental automatic casting during `index - cum_chunks` below
+        cum_chunks = cum_chunks.astype(index.dtype)
+
     # this dispactches to the array library
     chunk_locations = np.searchsorted(cum_chunks, index, side="right")
 
@@ -634,8 +639,6 @@ def take(outname, inname, chunks, index, itemsize, axis=0):
             PerformanceWarning,
             stacklevel=6,
         )
-    if not is_arraylike(index):
-        index = np.asarray(index)
 
     # Check for chunks from the plan that would violate the user's
     # configured chunk size.
diff --git a/dask/dataframe/backends.py b/dask/dataframe/backends.py
index ae629539b..88b334786 100644
--- a/dask/dataframe/backends.py
+++ b/dask/dataframe/backends.py
@@ -11,7 +11,7 @@ from dask.array.core import Array
 from dask.array.dispatch import percentile_lookup
 from dask.array.percentile import _percentile
 from dask.backends import CreationDispatch, DaskBackendEntrypoint
-from dask.dataframe._compat import is_any_real_numeric_dtype
+from dask.dataframe._compat import PANDAS_GE_220, is_any_real_numeric_dtype
 from dask.dataframe.core import DataFrame, Index, Scalar, Series, _Frame
 from dask.dataframe.dispatch import (
     categorical_dtype_dispatch,
@@ -623,11 +623,11 @@ def concat_pandas(
         if uniform
         else any(isinstance(df, pd.DataFrame) for df in dfs2)
     ):
-        if uniform:
+        if uniform or PANDAS_GE_220:
             dfs3 = dfs2
             cat_mask = dfs2[0].dtypes == "category"
         else:
-            # When concatenating mixed dataframes and series on axis 1, Pandas
+            # When concatenating mixed dataframes and series on axis 1, Pandas <2.2
             # converts series to dataframes with a single column named 0, then
             # concatenates.
             dfs3 = [
@@ -647,7 +647,7 @@ def concat_pandas(
                     **kwargs,
                 ).any()
 
-        if cat_mask.any():
+        if isinstance(cat_mask, pd.Series) and cat_mask.any():
             not_cat = cat_mask[~cat_mask].index
             # this should be aligned, so no need to filter warning
             out = pd.concat(
diff --git a/dask/sizeof.py b/dask/sizeof.py
index 5a98d9465..9ad81b2a1 100644
--- a/dask/sizeof.py
+++ b/dask/sizeof.py
@@ -213,10 +213,7 @@ def register_pandas():
 
     @sizeof.register(pd.MultiIndex)
     def sizeof_pandas_multiindex(i):
-        p = sum(sizeof(lev) for lev in i.levels)
-        for c in i.codes:
-            p += c.nbytes
-        return p
+        return sum(sizeof(l) for l in i.levels) + sum(c.nbytes for c in i.codes)
 
 
 @sizeof.register_lazy("scipy")


Test patch:
diff --git a/.github/workflows/test-report.yaml b/.github/workflows/test-report.yaml
index 0583aec7b..86d341d61 100644
--- a/.github/workflows/test-report.yaml
+++ b/.github/workflows/test-report.yaml
@@ -42,7 +42,7 @@ jobs:
       - name: mamba list
         run: mamba list
 
-      - uses: actions/cache@v3
+      - uses: actions/cache@v4
         id: cache
         with:
           # Suffix is depending on the backend / OS. Let's be agnostic here
diff --git a/dask/array/tests/test_array_core.py b/dask/array/tests/test_array_core.py
index d87283b35..f4b190f2f 100644
--- a/dask/array/tests/test_array_core.py
+++ b/dask/array/tests/test_array_core.py
@@ -51,7 +51,7 @@ from dask.array.core import (
     stack,
     store,
 )
-from dask.array.numpy_compat import _numpy_200
+from dask.array.numpy_compat import NUMPY_GE_200
 from dask.array.reshape import _not_implemented_message
 from dask.array.tests.test_dispatch import EncapsulateNDArray
 from dask.array.utils import assert_eq, same_keys
@@ -888,7 +888,7 @@ def test_elemwise_on_scalars():
     dy = from_array(ny, chunks=(5,))
     dz = dx.sum() * dy
 
-    if _numpy_200:
+    if NUMPY_GE_200:
         assert_eq(dz, nz)
     else:
         # Dask 0-d arrays do not behave like numpy scalars for type promotion
@@ -3112,25 +3112,19 @@ def test_slice_with_floats():
         d[[1, 1.5]]
 
 
-def test_slice_with_integer_types():
+@pytest.mark.parametrize("dtype", [np.int32, np.int64, np.uint32, np.uint64])
+def test_slice_with_integer_types(dtype):
     x = np.arange(10)
     dx = da.from_array(x, chunks=5)
-    inds = np.array([0, 3, 6], dtype="u8")
+    inds = np.array([0, 3, 6], dtype=dtype)
     assert_eq(dx[inds], x[inds])
-    assert_eq(dx[inds.astype("u4")], x[inds.astype("u4")])
-
-    inds = np.array([0, 3, 6], dtype=np.int64)
-    assert_eq(dx[inds], x[inds])
-    assert_eq(dx[inds.astype("u4")], x[inds.astype("u4")])
 
 
-def test_index_with_integer_types():
+@pytest.mark.parametrize("cls", [int, np.int32, np.int64, np.uint32, np.uint64])
+def test_index_with_integer_types(cls):
     x = np.arange(10)
     dx = da.from_array(x, chunks=5)
-    inds = int(3)
-    assert_eq(dx[inds], x[inds])
-
-    inds = np.int64(3)
+    inds = cls(3)
     assert_eq(dx[inds], x[inds])
 
 
diff --git a/dask/array/tests/test_fft.py b/dask/array/tests/test_fft.py
index 59d9a3e17..6b07b16c1 100644
--- a/dask/array/tests/test_fft.py
+++ b/dask/array/tests/test_fft.py
@@ -10,7 +10,7 @@ import dask.array as da
 import dask.array.fft
 from dask.array.core import normalize_chunks
 from dask.array.fft import fft_wrap
-from dask.array.numpy_compat import _numpy_200
+from dask.array.numpy_compat import NUMPY_GE_200
 from dask.array.utils import assert_eq, same_keys
 
 all_1d_funcnames = ["fft", "ifft", "rfft", "irfft", "hfft", "ihfft"]
@@ -60,7 +60,7 @@ def test_fft2n_shapes(funcname):
         da_fft(darr3, (12, 11), axes=(1, 0)), np_fft(nparr, (12, 11), axes=(1, 0))
     )
 
-    if _numpy_200 and funcname.endswith("fftn"):
+    if NUMPY_GE_200 and funcname.endswith("fftn"):
         ctx = pytest.warns(
             DeprecationWarning,
             match="`axes` should not be `None` if `s` is not `None`",
diff --git a/dask/array/tests/test_masked.py b/dask/array/tests/test_masked.py
index 851e0eab5..4f45ea9db 100644
--- a/dask/array/tests/test_masked.py
+++ b/dask/array/tests/test_masked.py
@@ -9,7 +9,7 @@ import numpy as np
 import pytest
 
 import dask.array as da
-from dask.array.numpy_compat import ComplexWarning, _numpy_123
+from dask.array.numpy_compat import NUMPY_GE_123, ComplexWarning
 from dask.array.utils import assert_eq
 from dask.base import tokenize
 from dask.utils import typename
@@ -387,7 +387,7 @@ def test_average_weights_with_masked_array(keepdims):
 
     da_avg = da.ma.average(d_a, weights=d_weights, axis=1, keepdims=keepdims)
 
-    if _numpy_123:
+    if NUMPY_GE_123:
         assert_eq(da_avg, np.ma.average(a, weights=weights, axis=1, keepdims=keepdims))
     elif not keepdims:
         assert_eq(da_avg, np.ma.average(a, weights=weights, axis=1))
diff --git a/dask/array/tests/test_reductions.py b/dask/array/tests/test_reductions.py
index 0a2c1a894..2af27e089 100644
--- a/dask/array/tests/test_reductions.py
+++ b/dask/array/tests/test_reductions.py
@@ -13,7 +13,7 @@ import itertools
 
 import dask.array as da
 import dask.config as config
-from dask.array.numpy_compat import ComplexWarning, _numpy_122
+from dask.array.numpy_compat import NUMPY_GE_122, ComplexWarning
 from dask.array.utils import assert_eq, same_keys
 from dask.core import get_deps
 
@@ -261,7 +261,7 @@ def test_arg_reductions(dfunc, func):
         assert_eq(dfunc(a, 0), func(x, 0))
         assert_eq(dfunc(a, 1), func(x, 1))
         assert_eq(dfunc(a, 2), func(x, 2))
-    if _numpy_122:
+    if NUMPY_GE_122:
         assert_eq(dfunc(a, keepdims=True), func(x, keepdims=True))
 
     pytest.raises(ValueError, lambda: dfunc(a, 3))
diff --git a/dask/array/tests/test_routines.py b/dask/array/tests/test_routines.py
index a1d6f5528..5a0e6f285 100644
--- a/dask/array/tests/test_routines.py
+++ b/dask/array/tests/test_routines.py
@@ -14,7 +14,7 @@ from dask.delayed import delayed
 np = pytest.importorskip("numpy")
 
 import dask.array as da
-from dask.array.numpy_compat import AxisError, _numpy_123, _numpy_200
+from dask.array.numpy_compat import NUMPY_GE_123, NUMPY_GE_200, AxisError
 from dask.array.utils import assert_eq, same_keys
 
 
@@ -2351,7 +2351,7 @@ def test_result_type():
     # Effect of scalars depends on their value
     assert da.result_type(1, b) == np.int16
     assert da.result_type(1.0, a) == np.float32
-    if _numpy_200:
+    if NUMPY_GE_200:
         assert da.result_type(np.int64(1), b) == np.int64
         assert da.result_type(np.ones((), np.int64), b) == np.int64
         assert da.result_type(1e200, a) == np.float32
@@ -2594,7 +2594,7 @@ def test_average_keepdims(a):
 
     da_avg = da.average(d_a, keepdims=True)
 
-    if _numpy_123:
+    if NUMPY_GE_123:
         np_avg = np.average(a, keepdims=True)
         assert_eq(np_avg, da_avg)
 
@@ -2609,7 +2609,7 @@ def test_average_weights(keepdims):
 
     da_avg = da.average(d_a, weights=d_weights, axis=1, keepdims=keepdims)
 
-    if _numpy_123:
+    if NUMPY_GE_123:
         assert_eq(da_avg, np.average(a, weights=weights, axis=1, keepdims=keepdims))
     elif not keepdims:
         assert_eq(da_avg, np.average(a, weights=weights, axis=1))
diff --git a/dask/dataframe/io/tests/test_parquet.py b/dask/dataframe/io/tests/test_parquet.py
index cfd3b91ae..861b411fc 100644
--- a/dask/dataframe/io/tests/test_parquet.py
+++ b/dask/dataframe/io/tests/test_parquet.py
@@ -17,7 +17,7 @@ from packaging.version import parse as parse_version
 import dask
 import dask.dataframe as dd
 import dask.multiprocessing
-from dask.array.numpy_compat import _numpy_124
+from dask.array.numpy_compat import NUMPY_GE_124
 from dask.blockwise import Blockwise, optimize_blockwise
 from dask.dataframe._compat import PANDAS_GE_150, PANDAS_GE_200, PANDAS_GE_202
 from dask.dataframe.io.parquet.core import get_engine
@@ -3311,7 +3311,7 @@ def test_pandas_timestamp_overflow_pyarrow(tmpdir):
     info = np.iinfo(np.dtype("int64"))
     # In `numpy=1.24.0` NumPy warns when an overflow is encountered when casting from float to int
     # https://numpy.org/doc/stable/release/1.24.0-notes.html#numpy-now-gives-floating-point-errors-in-casts
-    if _numpy_124:
+    if NUMPY_GE_124:
         ctx = pytest.warns(RuntimeWarning, match="invalid value encountered in cast")
     else:
         ctx = contextlib.nullcontext()
diff --git a/dask/dataframe/tests/test_arithmetics_reduction.py b/dask/dataframe/tests/test_arithmetics_reduction.py
index a3a22b4ad..78c9a90c5 100644
--- a/dask/dataframe/tests/test_arithmetics_reduction.py
+++ b/dask/dataframe/tests/test_arithmetics_reduction.py
@@ -10,7 +10,7 @@ import pytest
 from pandas.api.types import is_scalar
 
 import dask.dataframe as dd
-from dask.array.numpy_compat import _numpy_125
+from dask.array.numpy_compat import NUMPY_GE_125
 from dask.dataframe._compat import (
     PANDAS_GE_140,
     PANDAS_GE_150,
@@ -889,7 +889,7 @@ def test_reductions_out(axis, redfunc):
         # explicitly when calling np.var(dask)
         with ctx:
             np_redfunc(dsk_in, axis=axis, ddof=1, out=out)
-    elif _numpy_125 and redfunc == "product" and out is None:
+    elif NUMPY_GE_125 and redfunc == "product" and out is None:
         with pytest.warns(DeprecationWarning, match="`product` is deprecated"):
             np_redfunc(dsk_in, axis=axis, out=out)
     else:
@@ -923,7 +923,7 @@ def test_reductions_numpy_dispatch(axis, redfunc):
         # explicitly when calling np.var(dask)
         expect = np_redfunc(pdf, axis=axis, ddof=1)
         actual = np_redfunc(df, axis=axis, ddof=1)
-    elif _numpy_125 and redfunc == "product":
+    elif NUMPY_GE_125 and redfunc == "product":
         expect = np_redfunc(pdf, axis=axis)
         with pytest.warns(DeprecationWarning, match="`product` is deprecated"):
             actual = np_redfunc(df, axis=axis)
diff --git a/dask/tests/test_sizeof.py b/dask/tests/test_sizeof.py
index 9b43a7ae4..ef3d6ba37 100644
--- a/dask/tests/test_sizeof.py
+++ b/dask/tests/test_sizeof.py
@@ -78,7 +78,7 @@ def test_pandas_contiguous_dtypes():
 
 @requires_pandas
 def test_pandas_multiindex():
-    index = pd.MultiIndex.from_product([range(5), ["a", "b", "c", "d", "e"]])
+    index = pd.MultiIndex.from_product([range(50), list("abcdefghilmnopqrstuvwxyz")])
     actual_size = sys.getsizeof(index)
 
     assert 0.5 * actual_size < sizeof(index) < 2 * actual_size