feat: support slicing with int16 index arrays

dask-contrib · Mar 11, 2024 · e9bffaa · e9bffaa
1 parent a305dd4
commit e9bffaa
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/src/dask_awkward/lib/core.py b/src/dask_awkward/lib/core.py
@@ -1425,7 +1425,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_, bool, np.int64, np.int32, np.int16, int)):
                 return self._getitem_outer_bool_or_int_lazy_array(where)
 
         elif where[0] is Ellipsis:
@@ -1467,7 +1467,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_, bool, np.int64, np.int32, np.int16, int)):
                 return self._getitem_outer_bool_or_int_lazy_array(where)
 
         # a single ellipsis