Mask based indexing works when the mask is a list, but not a numpy array #1490
Description
Zarr version
2.16.1.dev4,2.15.0
Numcodecs version
0.10.2
Python Version
3.9.16,3.10.12
Operating System
Mac
Installation
Dev environment + pip install from pypi
Description
Indexing with a boolean mask doesn't work if the mask is a numpy array. I expect this to work because the behavior is available for python lists, and has an implementation in zarr/indexing.py
Steps to reproduce
import numpy as np
import zarr
z = zarr.array(np.arange(6).reshape((3, 2)))
z[[True, False, True], :]
# array([[0, 1],
# [4, 5]])That works, this doesn't:
z[np.array([True, False, True]), :]File ~/github/zarr-python/zarr/indexing.py:348, in BasicIndexer.__init__(self, selection, array)
345 dim_indexer = SliceDimIndexer(dim_sel, dim_len, dim_chunk_len)
347 else:
--> 348 raise IndexError(
349 "unsupported selection item for basic indexing; "
350 "expected integer or slice, got {!r}".format(type(dim_sel))
351 )
353 dim_indexers.append(dim_indexer)
355 self.dim_indexers = dim_indexers
IndexError: unsupported selection item for basic indexing; expected integer or slice, got <class 'numpy.ndarray'>Full traceback
---------------------------------------------------------------------------
IndexError Traceback (most recent call last)
Cell In[5], line 1
----> 1 z[np.array([True, False, True]), :]
File ~/github/zarr-python/zarr/core.py:844, in Array.__getitem__(self, selection)
842 result = self.get_orthogonal_selection(pure_selection, fields=fields)
843 else:
--> 844 result = self.get_basic_selection(pure_selection, fields=fields)
845 return result
File ~/github/zarr-python/zarr/core.py:970, in Array.get_basic_selection(self, selection, out, fields)
968 return self._get_basic_selection_zd(selection=selection, out=out, fields=fields)
969 else:
--> 970 return self._get_basic_selection_nd(selection=selection, out=out, fields=fields)
File ~/github/zarr-python/zarr/core.py:1010, in Array._get_basic_selection_nd(self, selection, out, fields)
1006 def _get_basic_selection_nd(self, selection, out=None, fields=None):
1007 # implementation of basic selection for array with at least one dimension
1008
1009 # setup indexer
-> 1010 indexer = BasicIndexer(selection, self)
1012 return self._get_selection(indexer=indexer, out=out, fields=fields)
File ~/github/zarr-python/zarr/indexing.py:348, in BasicIndexer.__init__(self, selection, array)
345 dim_indexer = SliceDimIndexer(dim_sel, dim_len, dim_chunk_len)
347 else:
--> 348 raise IndexError(
349 "unsupported selection item for basic indexing; "
350 "expected integer or slice, got {!r}".format(type(dim_sel))
351 )
353 dim_indexers.append(dim_indexer)
355 self.dim_indexers = dim_indexers
IndexError: unsupported selection item for basic indexing; expected integer or slice, got <class 'numpy.ndarray'>Additional output
This is also the case for 1 boolean arrays:
zarr.array(np.arange(3))[np.array([True, False, True])]
# IndexError: unsupported selection item for basic indexing; expected integer or slice, got <class 'numpy.ndarray'>Full traceback
---------------------------------------------------------------------------
IndexError Traceback (most recent call last)
Cell In[18], line 1
----> 1 zarr.array(np.arange(3))[np.array([True, False, True])]
File ~/github/zarr-python/zarr/core.py:844, in Array.__getitem__(self, selection)
842 result = self.get_orthogonal_selection(pure_selection, fields=fields)
843 else:
--> 844 result = self.get_basic_selection(pure_selection, fields=fields)
845 return result
File ~/github/zarr-python/zarr/core.py:970, in Array.get_basic_selection(self, selection, out, fields)
968 return self._get_basic_selection_zd(selection=selection, out=out, fields=fields)
969 else:
--> 970 return self._get_basic_selection_nd(selection=selection, out=out, fields=fields)
File ~/github/zarr-python/zarr/core.py:1010, in Array._get_basic_selection_nd(self, selection, out, fields)
1006 def _get_basic_selection_nd(self, selection, out=None, fields=None):
1007 # implementation of basic selection for array with at least one dimension
1008
1009 # setup indexer
-> 1010 indexer = BasicIndexer(selection, self)
1012 return self._get_selection(indexer=indexer, out=out, fields=fields)
File ~/github/zarr-python/zarr/indexing.py:348, in BasicIndexer.__init__(self, selection, array)
345 dim_indexer = SliceDimIndexer(dim_sel, dim_len, dim_chunk_len)
347 else:
--> 348 raise IndexError(
349 "unsupported selection item for basic indexing; "
350 "expected integer or slice, got {!r}".format(type(dim_sel))
351 )
353 dim_indexers.append(dim_indexer)
355 self.dim_indexers = dim_indexers
IndexError: unsupported selection item for basic indexing; expected integer or slice, got <class 'numpy.ndarray'>And I can get an error message telling me that mask based indexing is valid:
z = zarr.array(np.arange(6).reshape((3, 2)))
z[[True, False, True], np.array([1])]VindexInvalidSelectionError: unsupported selection type for vectorized indexing; only coordinate selection
(tuple of integer arrays) and mask selection (single Boolean array) are supported; got
(array([ True, False, True]), array([1]))