Fix a bug when indexing reduction expressions with ints

Author: FrancescAlted

src/blosc2/lazyexpr.py

@@ -10,6 +10,7 @@
 
 import ast
 import asyncio
+import builtins
 import concurrent.futures
 import copy
 import inspect
@@ -2601,7 +2602,7 @@ def _compute_expr(self, item, kwargs):  # noqa: C901
             _globals = get_expr_globals(self.expression)
             lazy_expr = eval(self.expression, _globals, self.operands)
             if not isinstance(lazy_expr, blosc2.LazyExpr):
-                key, _ = process_key(item, self.shape)
+                key, mask = process_key(item, self.shape)
                 # An immediate evaluation happened (e.g. all operands are numpy arrays)
                 if hasattr(self, "_where_args"):
                     # We need to apply the where() operation
@@ -2618,7 +2619,11 @@ def _compute_expr(self, item, kwargs):  # noqa: C901
                     # This is not exactly optimized, but it works for now
                     self._output[:] = lazy_expr[key]
                     return self._output
-                return lazy_expr[key]
+                arr = lazy_expr[key]
+                if builtins.sum(mask) > 0:
+                    # Correct shape to adjust to NumPy convention
+                    arr.shape = tuple(arr.shape[i] for i in range(len(mask)) if not mask[i])
+                return arr
 
             return chunked_eval(lazy_expr.expression, lazy_expr.operands, item, **kwargs)
 

tests/ndarray/test_lazyexpr.py

@@ -1451,8 +1451,6 @@ def test_to_cframe():
     N = 1_000
     dtype = "float64"
     a = blosc2.linspace(0, 1, N * N, dtype=dtype, shape=(N, N))
-    b = blosc2.linspace(1, 2, N * N, dtype=dtype, shape=(N, N))
-    c = blosc2.linspace(0, 1, N, dtype=dtype, shape=(N,))
     expr = a**3 + blosc2.sin(a**2)
     cframe = expr.to_cframe()
     assert len(cframe) > 0

tests/ndarray/test_reductions.py

@@ -426,10 +426,23 @@ def test_reduction_index():
     arr = blosc2.lazyexpr("sum(a,axis=0)", {"a": a})
     newarr = arr.compute()
     assert arr[:10].shape == (10,)
-    assert arr[0].shape == (1,)
+    assert arr[0].shape == ()
     assert arr.shape == newarr.shape
 
 
+@pytest.mark.parametrize("idx", [0, 1, (0,), slice(1, 2), (slice(0, 1),), slice(0, 4), (0, 2)])
+def test_reduction_index2(idx):
+    N = 10
+    shape = (N, N, N)
+    a = blosc2.linspace(0, 1, num=np.prod(shape), shape=(N, N, N))
+    expr = blosc2.lazyexpr("a.sum(axis=1)")
+    out = expr[idx]
+    na = blosc2.asarray(a)
+    nout = na.sum(axis=1)[idx]
+    assert out.shape == nout.shape
+    assert np.allclose(out, nout)
+
+
 def test_slice_lazy():
     shape = (20, 20)
     a = blosc2.linspace(0, 20, num=np.prod(shape), shape=shape)