Allow to run miniexpr is more cases (specially 1-dim)

Author: FrancescAlted

CMakeLists.txt

@@ -58,7 +58,9 @@ set(MINIEXPR_BUILD_BENCH OFF CACHE BOOL "Build miniexpr benchmarks" FORCE)
 
 FetchContent_Declare(miniexpr
     GIT_REPOSITORY https://github.com/Blosc/miniexpr.git
-    GIT_TAG 3e0ad9f2800cfb46729da88553a9228845eaa731  # latest SIMD additions
+    GIT_TAG sleef  # latest SIMD additions
+    # In case you want to use a local copy of miniexpr for development, uncomment the line below
+    # SOURCE_DIR "/Users/faltet/blosc/miniexpr"
 )
 FetchContent_MakeAvailable(miniexpr)
 

src/blosc2/blosc2_ext.pyx

@@ -573,6 +573,17 @@ cdef extern from "miniexpr.h":
     int me_compile(const char *expression, const me_variable *variables,
                    int var_count, me_dtype dtype, int *error, me_expr **out)
 
+    cdef enum me_compile_status:
+        ME_COMPILE_SUCCESS
+        ME_COMPILE_ERR_OOM
+        ME_COMPILE_ERR_PARSE
+        ME_COMPILE_ERR_INVALID_ARG
+        ME_COMPILE_ERR_COMPLEX_UNSUPPORTED
+        ME_COMPILE_ERR_REDUCTION_INVALID
+        ME_COMPILE_ERR_VAR_MIXED
+        ME_COMPILE_ERR_VAR_UNSPECIFIED
+        ME_COMPILE_ERR_INVALID_ARG_TYPE
+
     int me_eval(const me_expr *expr, const void ** vars_chunk,
                 int n_vars, void *output_chunk, int chunk_nitems) nogil
 
@@ -2878,8 +2889,10 @@ cdef class NDArray:
         expression = expression.encode("utf-8") if isinstance(expression, str) else expression
         cdef me_dtype = me_dtype_from_numpy(self.dtype.num)
         cdef me_expr *out_expr
-        error = me_compile(expression, variables, n, me_dtype, &error, &out_expr)
-        if error != 0:
+        cdef int rc = me_compile(expression, variables, n, me_dtype, &error, &out_expr)
+        if rc == ME_COMPILE_ERR_INVALID_ARG_TYPE:
+            raise TypeError(f"miniexpr does not support operand or output dtype: {expression}")
+        if rc != ME_COMPILE_SUCCESS:
             raise NotImplementedError(f"Cannot compile expression: {expression}")
         udata.miniexpr_handle = out_expr
 

src/blosc2/lazyexpr.py

@@ -1238,6 +1238,9 @@ def fast_eval(  # noqa: C901
         ne_args = {}
     dtype = kwargs.pop("dtype", None)
     where: dict | None = kwargs.pop("_where_args", None)
+    if where is not None:
+        # miniexpr does not support where(); use the regular path.
+        use_miniexpr = False
     if isinstance(out, blosc2.NDArray):
         # If 'out' has been passed, and is a NDArray, use it as the base array
         basearr = out
@@ -1290,15 +1293,17 @@ def fast_eval(  # noqa: C901
         use_miniexpr = False
 
     if use_miniexpr:
+        op_dtypes = {op.dtype for op in operands.values() if isinstance(op, blosc2.NDArray)}
+        if len(op_dtypes) > 1:
+            use_miniexpr = False
+        # Avoid padding issues except for 1D arrays (contiguous along the only axis).
+        if len(shape) != 1 and builtins.any(s % c != 0 for s, c in zip(shape, chunks, strict=True)):
+            use_miniexpr = False
         for op in operands.values():
             # Only NDArray in-memory operands
             if not (isinstance(op, blosc2.NDArray) and op.urlpath is None and out is None):
                 use_miniexpr = False
                 break
-            # Check that partitions are well-behaved (no padding)
-            if not blosc2.are_partitions_behaved(op.shape, op.chunks, op.blocks):
-                use_miniexpr = False
-                break
             # Ensure blocks fit exactly in chunks
             blocks_fit = builtins.all(c % b == 0 for c, b in zip(op.chunks, op.blocks, strict=True))
             if not blocks_fit:
@@ -1310,7 +1315,7 @@ def fast_eval(  # noqa: C901
         # All values will be overwritten, so we can use an uninitialized array
         res_eval = blosc2.uninit(shape, dtype, chunks=chunks, blocks=blocks, cparams=cparams, **kwargs)
         try:
-            # print("expr->miniexpr:", expression)
+            print("expr->miniexpr:", expression)
             res_eval._set_pref_expr(expression, operands)
             # Data to compress is fetched from operands, so it can be uninitialized here
             data = np.empty(res_eval.schunk.chunksize, dtype=np.uint8)
@@ -2001,6 +2006,16 @@ def reduce_slices(  # noqa: C901
 
     # Only behaved partitions are supported in miniexpr reductions
     if use_miniexpr:
+        # Avoid padding issues except for 1D arrays (contiguous along the only axis).
+        if len(shape) != 1 and builtins.any(s % c != 0 for s, c in zip(shape, chunks, strict=True)):
+            use_miniexpr = False
+        if use_miniexpr and isinstance(expression, str):
+            has_complex = any(
+                isinstance(op, blosc2.NDArray) and blosc2.isdtype(op.dtype, "complex floating")
+                for op in operands.values()
+            )
+            if has_complex and any(tok in expression for tok in ("!=", "==", "<=", ">=", "<", ">")):
+                use_miniexpr = False
         for op in operands.values():
             # Check that chunksize is multiple of blocksize and blocks fit exactly in chunks
             blocks_fit = builtins.all(c % b == 0 for c, b in zip(op.chunks, op.blocks, strict=True))

tests/ndarray/test_lazyexpr.py

@@ -180,7 +180,10 @@ def test_simple_expression(array_fixture):
     expr = a1 + a2 - a3 * a4
     nres = ne_evaluate("na1 + na2 - na3 * na4")
     res = expr.compute(cparams=blosc2.CParams())
-    np.testing.assert_allclose(res[:], nres)
+    if na1.dtype == np.float32:
+        np.testing.assert_allclose(res[:], nres, rtol=1e-6, atol=1e-6)
+    else:
+        np.testing.assert_allclose(res[:], nres)
 
 
 # Mix Proxy and NDArray operands
@@ -205,10 +208,16 @@ def test_iXXX(array_fixture):
         expr **= 2.3  # __ipow__
     res = expr.compute()
     if not blosc2.IS_WASM:
-        nres = ne_evaluate("(((((na1 ** 3 + na2 ** 2 + na3 ** 3 - na4 + 3) + 5) - 15) * 2) / 7) ** 2.3")
+        expr_str = "(((((na1 ** 3 + na2 ** 2 + na3 ** 3 - na4 + 3) + 5) - 15) * 2) / 7) ** 2.3"
     else:
-        nres = ne_evaluate("(((((na1 ** 3 + na2 ** 2 + na3 ** 3 - na4 + 3) + 5) - 15) * 2) / 7)")
-    np.testing.assert_allclose(res[:], nres)
+        expr_str = "(((((na1 ** 3 + na2 ** 2 + na3 ** 3 - na4 + 3) + 5) - 15) * 2) / 7)"
+    if na1.dtype == np.float32:
+        with np.errstate(invalid="ignore"):
+            nres = eval(expr_str, {"np": np}, {"na1": na1, "na2": na2, "na3": na3, "na4": na4})
+        np.testing.assert_allclose(res[:], nres, rtol=1e-5, atol=1e-6)
+    else:
+        nres = ne_evaluate(expr_str)
+        np.testing.assert_allclose(res[:], nres)
 
 
 def test_complex_evaluate(array_fixture):
@@ -253,7 +262,10 @@ def test_expression_with_constants(array_fixture):
     # Test with operands with same chunks and blocks
     expr = a1 + 2 - a3 * 3.14
     nres = ne_evaluate("na1 + 2 - na3 * 3.14")
-    np.testing.assert_allclose(expr[:], nres)
+    if na1.dtype == np.float32:
+        np.testing.assert_allclose(expr[:], nres, rtol=1e-6)
+    else:
+        np.testing.assert_allclose(expr[:], nres)
 
 
 @pytest.mark.parametrize("compare_expressions", [True, False])