Support miniexpr for where inside reductions, e.g 'sum(where(a < b, b, a))'

Author: FrancescAlted

CMakeLists.txt

@@ -58,8 +58,7 @@ set(MINIEXPR_BUILD_BENCH OFF CACHE BOOL "Build miniexpr benchmarks" FORCE)
 
 FetchContent_Declare(miniexpr
     GIT_REPOSITORY https://github.com/Blosc/miniexpr.git
-    #GIT_TAG ndim  # latest me_compile_nd()/me_eval_nd() APIs
-    GIT_TAG 8c50850094e156ce568186edef667fabecbd00ff  # latest commit in ndim
+    GIT_TAG ndim  # latest me_compile_nd()/me_eval_nd() APIs
     # In case you want to use a local copy of miniexpr for development, uncomment the line below
     # SOURCE_DIR "/Users/faltet/blosc/miniexpr"
 )

src/blosc2/lazyexpr.py

@@ -1309,8 +1309,8 @@ def fast_eval(  # noqa: C901
         fp_accuracy = kwargs.pop("fp_accuracy", blosc2.FPAccuracy.DEFAULT)
         res_eval = blosc2.uninit(shape, dtype, chunks=chunks, blocks=blocks, cparams=cparams, **kwargs)
         try:
-            print("expr->miniexpr:", expression)
             res_eval._set_pref_expr(expression, operands, fp_accuracy=fp_accuracy)
+            print("expr->miniexpr:", expression)
             # Data to compress is fetched from operands, so it can be uninitialized here
             data = np.empty(res_eval.schunk.chunksize, dtype=np.uint8)
             # Exercise prefilter for each chunk
@@ -1993,7 +1993,7 @@ def reduce_slices(  # noqa: C901
         del temp
 
     # miniexpr reduction path only supported for some cases so far
-    if not (where is None and fast_path and all_ndarray and reduced_shape == ()):
+    if not (fast_path and all_ndarray and reduced_shape == ()):
         use_miniexpr = False
 
     # Some reductions are not supported yet in miniexpr
@@ -2008,6 +2008,8 @@ def reduce_slices(  # noqa: C901
         )
         if has_complex and any(tok in expression for tok in ("!=", "==", "<=", ">=", "<", ">")):
             use_miniexpr = False
+        if where is not None and len(where) != 2:
+            use_miniexpr = False
 
     if use_miniexpr:
         # Experiments say that not splitting is best (at least on Apple Silicon M4 Pro)
@@ -2037,9 +2039,12 @@ def reduce_slices(  # noqa: C901
             # For other operations, zeros should be safe
             aux_reduc = np.zeros(nblocks, dtype=dtype)
         try:
+            if where is not None:
+                expression_miniexpr = f"{reduce_op_str}(where({expression}, _where_x, _where_y))"
+            else:
+                expression_miniexpr = f"{reduce_op_str}({expression})"
+            res_eval._set_pref_expr(expression_miniexpr, operands, fp_accuracy, aux_reduc)
             print("expr->miniexpr:", expression, reduce_op)
-            expression = f"{reduce_op_str}({expression})"
-            res_eval._set_pref_expr(expression, operands, fp_accuracy, aux_reduc)
             # Data won't even try to be compressed, so buffers can be unitialized and reused
             data = np.empty(res_eval.schunk.chunksize, dtype=np.uint8)
             chunk_data = np.empty(res_eval.schunk.chunksize + blosc2.MAX_OVERHEAD, dtype=np.uint8)

tests/ndarray/test_reductions.py

@@ -65,15 +65,24 @@ def test_reduce_bool(array_fixture, reduce_op):
     np.testing.assert_allclose(res, nres, atol=tol, rtol=tol)
 
 
-def test_reduce_where(array_fixture):
+# @pytest.mark.parametrize("reduce_op", ["sum"])
+@pytest.mark.parametrize("reduce_op", ["sum", "prod", "min", "max", "any", "all", "argmax", "argmin"])
+def test_reduce_where(array_fixture, reduce_op):
     a1, a2, a3, a4, na1, na2, na3, na4 = array_fixture
-    # The next works
-    # res = blosc2.where(a1 < a2, a2, 0).sum()
-    # nres = ne_evaluate("sum(where(na1 < na2, na2, 0))")
-    # This does not work yet (it currently hangs)
-    res = blosc2.where(a1 < a2, a2, a1).sum()
-    nres = ne_evaluate("sum(where(na1 < na2, na2, na1))")
-    print("res:", res, nres)
+    if reduce_op == "prod":
+        # To avoid overflow, create a1 and a2 with small values
+        na1 = np.linspace(0, 0.1, np.prod(a1.shape), dtype=np.float32).reshape(a1.shape)
+        a1 = blosc2.asarray(na1)
+        na2 = np.linspace(0, 0.5, np.prod(a1.shape), dtype=np.float32).reshape(a1.shape)
+        a2 = blosc2.asarray(na2)
+        expr = a1 + a2 - 0.2
+        nres = eval("na1 + na2 - .2")
+    else:
+        expr = blosc2.where(a1 < a2, a2, a1)
+        nres = eval("np.where(na1 < na2, na2, na1)")
+    res = getattr(expr, reduce_op)()
+    nres = getattr(nres, reduce_op)()
+    # print("res:", res, nres, type(res), type(nres))
     tol = 1e-15 if a1.dtype == "float64" else 1e-6
     np.testing.assert_allclose(res, nres, atol=tol, rtol=tol)