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
(out_name, in_name, blockdims, index, allow_getitem_optimization)
| 226 | |
| 227 | |
| 228 | def slice_wrap_lists(out_name, in_name, blockdims, index, allow_getitem_optimization): |
| 229 | """ |
| 230 | Fancy indexing along blocked array dasks |
| 231 | |
| 232 | Handles index of type list. Calls slice_slices_and_integers for the rest |
| 233 | |
| 234 | See Also |
| 235 | -------- |
| 236 | |
| 237 | take : handle slicing with lists ("fancy" indexing) |
| 238 | slice_slices_and_integers : handle slicing with slices and integers |
| 239 | """ |
| 240 | assert all(isinstance(i, (slice, list, Integral)) or is_arraylike(i) for i in index) |
| 241 | if not len(blockdims) == len(index): |
| 242 | raise IndexError("Too many indices for array") |
| 243 | |
| 244 | # Do we have more than one list in the index? |
| 245 | where_list = [ |
| 246 | i for i, ind in enumerate(index) if is_arraylike(ind) and ind.ndim > 0 |
| 247 | ] |
| 248 | if len(where_list) > 1: |
| 249 | raise NotImplementedError("Don't yet support nd fancy indexing") |
| 250 | # Is the single list an empty list? In this case just treat it as a zero |
| 251 | # length slice |
| 252 | if where_list and not index[where_list[0]].size: |
| 253 | index = list(index) |
| 254 | index[where_list.pop()] = slice(0, 0, 1) |
| 255 | index = tuple(index) |
| 256 | |
| 257 | # No lists, hooray! just use slice_slices_and_integers |
| 258 | if not where_list: |
| 259 | return slice_slices_and_integers( |
| 260 | out_name, in_name, blockdims, index, allow_getitem_optimization |
| 261 | ) |
| 262 | |
| 263 | # Replace all lists with full slices [3, 1, 0] -> slice(None, None, None) |
| 264 | index_without_list = tuple( |
| 265 | slice(None, None, None) if is_arraylike(i) else i for i in index |
| 266 | ) |
| 267 | |
| 268 | # lists and full slices. Just use take |
| 269 | if all(is_arraylike(i) or i == slice(None, None, None) for i in index): |
| 270 | axis = where_list[0] |
| 271 | blockdims2, dsk3 = take( |
| 272 | out_name, in_name, blockdims, index[where_list[0]], axis=axis |
| 273 | ) |
| 274 | # Mixed case. Both slices/integers and lists. slice/integer then take |
| 275 | else: |
| 276 | # Do first pass without lists |
| 277 | tmp = "slice-" + tokenize((out_name, in_name, blockdims, index)) |
| 278 | dsk, blockdims2 = slice_slices_and_integers( |
| 279 | tmp, |
| 280 | in_name, |
| 281 | blockdims, |
| 282 | index_without_list, |
| 283 | allow_getitem_optimization=False, |
| 284 | ) |
| 285 |
no test coverage detected
searching dependent graphs…