No (easy) optimisation for booleans possible. Now pass all ndarray_object tests!

Author: lshaw8317

src/blosc2/__init__.py

@@ -351,6 +351,7 @@ def __array_namespace_info__() -> Info:
     get_expr_operands,
     validate_expr,
     evaluate,
+    slices_eval,
 )
 from .proxy import Proxy, ProxySource, ProxyNDSource, ProxyNDField, SimpleProxy, jit
 

src/blosc2/lazyexpr.py

@@ -1506,11 +1506,8 @@ def slices_eval(  # noqa: C901
         if out is None:
             shape_ = shape_slice if shape_slice is not None else shape
             if where is not None and len(where) < 2:
-                # The result is a linear array in the first ndims
-                try:  # call from NDArray.__getitem__
-                    shape_ = (math.prod(operands["key"].shape),) + shape[operands["key"].ndim :]
-                except KeyError:
-                    shape_ = math.prod(shape_)
+                # The result is a linear array
+                shape_ = math.prod(shape_)
             if getitem or _order:
                 out = np.empty(shape_, dtype=dtype_)
                 if _order:

src/blosc2/ndarray.py

@@ -1599,6 +1599,8 @@ def get_fselection_numpy(self, key: list | np.ndarray) -> np.ndarray:
             return out.reshape(out_shape)  # should have filled in correct order, just need to reshape
 
         # Default when there are booleans
+        # TODO: for boolean indexing could be optimised by avoiding
+        # calculating out_shape prior to loop and keeping track on-the-fly (like in LazyExpr machinery)
         return self._get_set_findex_default(_slice, out_shape)
 
     def _get_set_findex_default(self, _slice, out_shape=None, updater=None):
@@ -1718,36 +1720,36 @@ def __getitem__(  # noqa: C901
             expr = blosc2.LazyExpr._new_expr(key, self.fields, guess=False)
             return expr.where(self)
 
-        # key not iterable
-        key = key[()] if isinstance(key, NDArray) else key
+        key = key[()] if isinstance(key, NDArray) else key  # key not iterable
         key = tuple(k[()] if isinstance(k, NDArray) else k for k in key) if isinstance(key, tuple) else key
 
         # decompress NDArrays
         key_, mask = process_key(key, self.shape)  # internally handles key an integer
         key = key[()] if hasattr(key, "shape") and key.shape == () else key  # convert to scalar
+
         # fancy indexing
         if isinstance(key_, (list, np.ndarray)) or builtins.any(
             isinstance(k, (list, np.ndarray)) for k in key_
         ):
-            # check scalar booleans, which add 1 dim to beginning but which cause problems for ndindex.as_subindex
+            # check scalar booleans, which add 1 dim to beginning
             if np.issubdtype(type(key), bool) and np.isscalar(key):
                 if key:
                     _slice = ndindex.ndindex(()).expand(self.shape)  # just get whole array
                     out_shape = _slice.newshape(self.shape)
                     return np.expand_dims(self._get_set_findex_default(_slice, out_shape=out_shape), 0)
                 else:  # do nothing
                     return np.empty((0,) + self.shape, dtype=self.dtype)
-            elif hasattr(key, "dtype") and np.issubdtype(key.dtype, np.bool_):  # check ORIGINAL key
-                # This can be interpreted as a boolean expression
-                # elif key.shape == self.shape: # This is faster than the fancy indexing path
-                # expr = blosc2.lazyexpr(f"(key)")
-                # The next should be a bit faster
-                expr = blosc2.LazyExpr._new_expr("key", {"key": key}, guess=False)
+            elif (
+                hasattr(key, "dtype") and np.issubdtype(key.dtype, np.bool_) and key.shape == self.shape
+            ):  # check ORIGINAL key
+                # This can be interpreted as a boolean expression but only for key shape same as self shape
+                expr = blosc2.LazyExpr._new_expr("key", {"key": key}, guess=False).where(self)
                 # Decorate with where and force a getitem operation to return actual values.
                 # This behavior is consistent with NumPy, although different from e.g. ['expr']
                 # which returns a lazy expression.
-                return expr.where(self)[:]
-            return self.get_fselection_numpy(key_)  # fancy index
+                # This is faster than the fancy indexing path
+                return expr[:]
+            return self.get_fselection_numpy(key)  # fancy index default, can be quite slow
 
         start, stop, step, none_mask = get_ndarray_start_stop(self.ndim, key_, self.shape)
         for i, s in enumerate(step):  # (start, stop, -1) => stop < start
@@ -4096,7 +4098,7 @@ def save(array: NDArray, urlpath: str, contiguous=True, **kwargs: Any) -> None:
     array.save(urlpath, contiguous, **kwargs)
 
 
-def asarray(array: Sequence | np.ndarray | blosc2.C2Array | NDArray, copy=True, **kwargs: Any) -> NDArray:  # noqa: C901
+def asarray(array: Sequence | np.ndarray | blosc2.C2Array | NDArray, copy=True, **kwargs: Any) -> NDArray:
     """Convert the `array` to an `NDArray`.
 
     Parameters
@@ -4159,10 +4161,7 @@ def asarray(array: Sequence | np.ndarray | blosc2.C2Array | NDArray, copy=True,
             # A getitem operation should be enough to get a numpy array
             array = array[()]
 
-        if dtype != array.dtype:
-            array = array.astype(dtype=dtype, casting=casting)
-        if not array.flags.contiguous:
-            array = np.ascontiguousarray(array)
+        array = np.require(array, dtype=dtype, requirements="C")  # require contiguous array
 
         return blosc2_ext.asarray(array, chunks, blocks, **kwargs)
 
@@ -4182,7 +4181,7 @@ def asarray(array: Sequence | np.ndarray | blosc2.C2Array | NDArray, copy=True,
             for i, (c, s) in enumerate(zip(coords, chunks, strict=True))
         )
         # Ensure the array slice is contiguous and of correct dtype
-        array_slice = np.ascontiguousarray(array[slice_], dtype=dtype)
+        array_slice = np.require(array[slice_], dtype=dtype, requirements="C")
         if behaved:
             # The whole chunk is to be updated, so this fastpath is safe
             ndarr.schunk.update_data(nchunk, array_slice, copy=False)