Mask based indexing works when the mask is a list, but not a numpy array

[ivirshup]

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]))

[joshmoore]

+1: Seems like widening the isinstance in https://github.com/zarr-developers/zarr-python/blob/main/zarr/indexing.py#L33 would be straight-forward.

cc: @jni re 7c31f04

[jhamman]

This seems to have been fixed in the v3 refactor. Feel free to reopen if you hit this on 3.0.0.beta or later.