diff --git a/src/dask_awkward/lib/core.py b/src/dask_awkward/lib/core.py index 8d22b05d..40898b11 100644 --- a/src/dask_awkward/lib/core.py +++ b/src/dask_awkward/lib/core.py @@ -3,6 +3,7 @@ import keyword import logging import math +import numbers import operator import sys import warnings @@ -1425,7 +1426,7 @@ def _getitem_tuple(self, where): dtype = where[0].layout.dtype.type except AttributeError: dtype = where[0].layout.content.dtype.type - if issubclass(dtype, (np.bool_, bool, np.int64, np.int32, int)): + if issubclass(dtype, (np.bool_, numbers.Integral)): return self._getitem_outer_bool_or_int_lazy_array(where) elif where[0] is Ellipsis: @@ -1467,7 +1468,7 @@ def _getitem_single(self, where): while not hasattr(layout, "dtype"): layout = layout.content dtype = layout.dtype.type - if issubclass(dtype, (np.bool_, bool, np.int64, np.int32, int)): + if issubclass(dtype, (np.bool_, numbers.Integral)): return self._getitem_outer_bool_or_int_lazy_array(where) # a single ellipsis diff --git a/tests/test_getitem.py b/tests/test_getitem.py index bb52fb9b..e7dd3f8d 100644 --- a/tests/test_getitem.py +++ b/tests/test_getitem.py @@ -66,6 +66,21 @@ def test_single_int(daa: dak.Array, caa: ak.Array) -> None: assert caa["points"][i].tolist() == daa["points"][i].compute().tolist() +@pytest.mark.parametrize( + "dtype", + [np.int8, np.uint8, np.int16, np.uint16, np.int32, np.uint32, np.int64, np.uint64], +) +def test_multi_int(daa: dak.Array, caa: ak.Array, dtype: type) -> None: + a = daa["points"]["x"] + aix = ak.zeros_like(a, dtype=dtype) # take first elements + c = caa["points"]["x"] + cix = ak.zeros_like(c, dtype=dtype) # take first elements + assert_eq( + a[aix], + c[cix], + ) + + def test_single_ellipsis(daa: dak.Array, caa: ak.Array) -> None: assert_eq(daa[...], caa[...])