Wrap a dask array to support explicit indexing.
| 1820 | |
| 1821 | |
| 1822 | class DaskIndexingAdapter(IndexingAdapter): |
| 1823 | """Wrap a dask array to support explicit indexing.""" |
| 1824 | |
| 1825 | __slots__ = ("array",) |
| 1826 | |
| 1827 | def __init__(self, array): |
| 1828 | """This adapter is created in Variable.__getitem__ in |
| 1829 | Variable._broadcast_indexes. |
| 1830 | """ |
| 1831 | self.array = array |
| 1832 | |
| 1833 | def _oindex_get(self, indexer: OuterIndexer): |
| 1834 | key = indexer.tuple |
| 1835 | try: |
| 1836 | return self.array[key] |
| 1837 | except NotImplementedError: |
| 1838 | # manual orthogonal indexing |
| 1839 | value = self.array |
| 1840 | for axis, subkey in reversed(list(enumerate(key))): |
| 1841 | value = value[(slice(None),) * axis + (subkey,)] |
| 1842 | return value |
| 1843 | |
| 1844 | def _vindex_get(self, indexer: VectorizedIndexer): |
| 1845 | try: |
| 1846 | return self.array.vindex[indexer.tuple] |
| 1847 | except IndexError as e: |
| 1848 | # TODO: upstream to dask |
| 1849 | has_dask = any(is_duck_dask_array(i) for i in indexer.tuple) |
| 1850 | # this only works for "small" 1d coordinate arrays with one chunk |
| 1851 | # it is intended for idxmin, idxmax, and allows indexing with |
| 1852 | # the nD array output of argmin, argmax |
| 1853 | if ( |
| 1854 | not has_dask |
| 1855 | or len(indexer.tuple) > 1 |
| 1856 | or math.prod(self.array.numblocks) > 1 |
| 1857 | or self.array.ndim > 1 |
| 1858 | ): |
| 1859 | raise e |
| 1860 | (idxr,) = indexer.tuple |
| 1861 | if idxr.ndim == 0: |
| 1862 | return self.array[idxr.data] |
| 1863 | else: |
| 1864 | import dask.array |
| 1865 | |
| 1866 | return dask.array.map_blocks( |
| 1867 | _apply_vectorized_indexer_dask_wrapper, |
| 1868 | idxr[..., np.newaxis], |
| 1869 | self.array, |
| 1870 | chunks=idxr.chunks, |
| 1871 | drop_axis=-1, |
| 1872 | dtype=self.array.dtype, |
| 1873 | ) |
| 1874 | |
| 1875 | def __getitem__(self, indexer: ExplicitIndexer): |
| 1876 | self._check_and_raise_if_non_basic_indexer(indexer) |
| 1877 | return self.array[indexer.tuple] |
| 1878 | |
| 1879 | def _oindex_set(self, indexer: OuterIndexer, value: Any) -> None: |
no outgoing calls
searching dependent graphs…