From e9bffaa56797bb1e535f04f556b34f92d0727d32 Mon Sep 17 00:00:00 2001 From: Dmitry Kalinkin Date: Mon, 11 Mar 2024 18:06:34 -0400 Subject: [PATCH] feat: support slicing with int16 index arrays --- src/dask_awkward/lib/core.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/dask_awkward/lib/core.py b/src/dask_awkward/lib/core.py index 8d22b05d..cda45b13 100644 --- 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