Merge pull request #483 from veprbl/patch-2

feat: support slicing with int16/int8 index arrays

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

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[...])