Fancy indexing along blocked array dasks Handles index of type list. Calls slice_slices_and_integers for the rest See Also -------- take : handle slicing with lists ("fancy" indexing) slice_slices_and_integers : handle slicing with slices and integers
(x, index, allow_getitem_optimization)
| 268 | |
| 269 | |
| 270 | def slice_wrap_lists(x, index, allow_getitem_optimization): |
| 271 | """ |
| 272 | Fancy indexing along blocked array dasks |
| 273 | |
| 274 | Handles index of type list. Calls slice_slices_and_integers for the rest |
| 275 | |
| 276 | See Also |
| 277 | -------- |
| 278 | |
| 279 | take : handle slicing with lists ("fancy" indexing) |
| 280 | slice_slices_and_integers : handle slicing with slices and integers |
| 281 | """ |
| 282 | assert all(isinstance(i, (slice, list, Integral)) or is_arraylike(i) for i in index) |
| 283 | if not len(x.chunks) == len(index): |
| 284 | raise IndexError("Too many indices for array") |
| 285 | |
| 286 | # Do we have more than one list in the index? |
| 287 | where_list = [ |
| 288 | i for i, ind in enumerate(index) if is_arraylike(ind) and ind.ndim > 0 |
| 289 | ] |
| 290 | if len(where_list) > 1: |
| 291 | raise NotImplementedError("Don't yet support nd fancy indexing") |
| 292 | # Is the single list an empty list? In this case just treat it as a zero |
| 293 | # length slice |
| 294 | if where_list and not index[where_list[0]].size: |
| 295 | index = list(index) |
| 296 | index[where_list.pop()] = slice(0, 0, 1) |
| 297 | index = tuple(index) |
| 298 | |
| 299 | # No lists, hooray! just use slice_slices_and_integers |
| 300 | if not where_list: |
| 301 | return slice_slices_and_integers(x, index, allow_getitem_optimization) |
| 302 | |
| 303 | # Replace all lists with full slices [3, 1, 0] -> slice(None, None, None) |
| 304 | index_without_list = tuple( |
| 305 | slice(None, None, None) if is_arraylike(i) else i for i in index |
| 306 | ) |
| 307 | |
| 308 | # lists and full slices. Just use take |
| 309 | if all(is_arraylike(i) or i == slice(None, None, None) for i in index): |
| 310 | axis = where_list[0] |
| 311 | x = take(x, index[where_list[0]], axis=axis) |
| 312 | # Mixed case. Both slices/integers and lists. slice/integer then take |
| 313 | else: |
| 314 | x = slice_slices_and_integers( |
| 315 | x, |
| 316 | index_without_list, |
| 317 | allow_getitem_optimization=False, |
| 318 | ) |
| 319 | axis = where_list[0] |
| 320 | axis2 = axis - sum( |
| 321 | 1 for i, ind in enumerate(index) if i < axis and isinstance(ind, Integral) |
| 322 | ) |
| 323 | x = take(x, index[axis], axis=axis2) |
| 324 | |
| 325 | return x |
| 326 | |
| 327 |
no test coverage detected
searching dependent graphs…