Improve handling of known/unknown funcs

Author: lshaw8317

src/blosc2/lazyexpr.py

@@ -45,14 +45,13 @@
     _check_allowed_dtypes,
     get_chunks_idx,
     get_intersecting_chunks,
-    is_inside_new_expr,
     local_ufunc_map,
     process_key,
     ufunc_map,
     ufunc_map_1param,
 )
 
-from .shape_utils import constructors, infer_shape, lin_alg_funcs, reducers
+from .shape_utils import constructors, elementwise_funcs, infer_shape, lin_alg_attrs, lin_alg_funcs, reducers
 
 if not blosc2.IS_WASM:
     import numexpr
@@ -158,119 +157,9 @@ def ne_evaluate(expression, local_dict=None, **kwargs):
     "S": np.str_,
     "V": np.bytes_,
 }
-
-blosc2_funcs = [
-    "abs",
-    "acos",
-    "acosh",
-    "add",
-    "all",
-    "any",
-    "arange",
-    "arccos",
-    "arccosh",
-    "arcsin",
-    "arcsinh",
-    "arctan",
-    "arctan2",
-    "arctanh",
-    "asin",
-    "asinh",
-    "atan",
-    "atan2",
-    "atanh",
-    "bitwise_and",
-    "bitwise_invert",
-    "bitwise_left_shift",
-    "bitwise_or",
-    "bitwise_right_shift",
-    "bitwise_xor",
-    "broadcast_to",
-    "ceil",
-    "clip",
-    "concat",
-    "concatenate",
-    "copy",
-    "copysign",
-    "count_nonzero",
-    "divide",
-    "empty",
-    "empty_like",
-    "equal",
-    "expand_dims",
-    "expm1",
-    "eye",
-    "floor",
-    "floor_divide",
-    "frombuffer",
-    "fromiter",
-    "full",
-    "full_like",
-    "greater",
-    "greater_equal",
-    "hypot",
-    "isfinite",
-    "isinf",
-    "isnan",
-    "less_equal",
-    "less_than",
-    "linspace",
-    "log",
-    "log1p",
-    "log2",
-    "log10",
-    "logaddexp",
-    "logical_and",
-    "logical_not",
-    "logical_or",
-    "logical_xor",
-    "matmul",
-    "matrix_transpose",
-    "max",
-    "maximum",
-    "mean",
-    "meshgrid",
-    "min",
-    "minimum",
-    "multiply",
-    "nans",
-    "ndarray_from_cframe",
-    "negative",
-    "nextafter",
-    "not_equal",
-    "ones",
-    "ones_like",
-    "permute_dims",
-    "positive",
-    "pow",
-    "prod",
-    "real",
-    "reciprocal",
-    "remainder",
-    "reshape",
-    "round",
-    "sign",
-    "signbit",
-    "sort",
-    "square",
-    "squeeze",
-    "stack",
-    "sum",
-    "subtract",
-    "take",
-    "take_along_axis",
-    "tan",
-    "tanh",
-    "tensordot",
-    "transpose",
-    "trunc",
-    "var",
-    "vecdot",
-    "where",
-    "zeros",
-    "zeros_like",
-]
-
+blosc2_funcs = constructors + lin_alg_funcs + elementwise_funcs + reducers
+# functions that have to be evaluated before chunkwise lazyexpr machinery
+eager_funcs = lin_alg_funcs + reducers + ["slice"] + lin_alg_attrs
 # Gather all callable functions in numpy
 numpy_funcs = {
     name
@@ -751,12 +640,7 @@ def validate_expr(expr: str) -> None:
 
 def extract_and_replace_slices(expr, operands):
     """
-    Replaces all var.slice(...).slice(...) chains in expr with oN temporary variables.
-    Infers shapes using infer_shape and creates placeholder arrays in new_ops.
-
-    Returns:
-        new_expr: expression string with oN replacements
-        new_ops: dictionary mapping variable names (original and oN) to arrays
+    Return new expression and operands with op.slice(...) replaced by temporary operands.
     """
     # Copy shapes and operands
     shapes = {k: () if not hasattr(v, "shape") else v.shape for k, v in operands.items()}
@@ -1976,11 +1860,6 @@ def reduce_slices(  # noqa: C901
     if out is not None and reduced_shape != out.shape:
         raise ValueError("Provided output shape does not match the reduced shape.")
 
-    if is_inside_new_expr():
-        # We already have the dtype and reduced_shape, so return immediately
-        # Use a blosc2 container, as it consumes less memory in general
-        return blosc2.zeros(reduced_shape, dtype=dtype)
-
     # Choose the array with the largest shape as the reference for chunks
     # Note: we could have expr = blosc2.lazyexpr('numpy_array + 1') (i.e. no choice for chunks)
     blosc2_arrs = tuple(o for o in operands.values() if hasattr(o, "chunks"))
@@ -3023,7 +2902,7 @@ def _compute_expr(self, item, kwargs):  # noqa : C901
                 }
             )
 
-        if any(method in self.expression for method in reducers + lin_alg_funcs):
+        if any(method in self.expression for method in eager_funcs):
             # We have reductions in the expression (probably coming from a string lazyexpr)
             # Also includes slice
             _globals = get_expr_globals(self.expression)
@@ -3073,7 +2952,7 @@ def _compute_expr(self, item, kwargs):  # noqa : C901
                     # Replace the constructor call by the new operand
                     newexpr = newexpr.replace(constexpr, newop)
 
-            _globals = {func: getattr(blosc2, func) for func in functions if func in newexpr}
+            _globals = get_expr_globals(newexpr)
             lazy_expr = eval(newexpr, _globals, newops)
             if isinstance(lazy_expr, blosc2.NDArray):
                 # Almost done (probably the expression is made of only constructors)

src/blosc2/shape_utils.py

@@ -1,9 +1,85 @@
 import ast
 import builtins
+import warnings
 
 from numpy import broadcast_shapes
 
-lin_alg_funcs = (
+elementwise_funcs = [
+    "abs",
+    "acos",
+    "acosh",
+    "add",
+    "arccos",
+    "arccosh",
+    "arcsin",
+    "arcsinh",
+    "arctan",
+    "arctan2",
+    "arctanh",
+    "asin",
+    "asinh",
+    "atan",
+    "atan2",
+    "atanh",
+    "bitwise_and",
+    "bitwise_invert",
+    "bitwise_left_shift",
+    "bitwise_or",
+    "bitwise_right_shift",
+    "bitwise_xor",
+    "broadcast_to",
+    "ceil",
+    "clip",
+    "copysign",
+    "divide",
+    "equal",
+    "expm1",
+    "floor",
+    "floor_divide",
+    "greater",
+    "greater_equal",
+    "hypot",
+    "isfinite",
+    "isinf",
+    "isnan",
+    "less_equal",
+    "less_than",
+    "log",
+    "log1p",
+    "log2",
+    "log10",
+    "logaddexp",
+    "logical_and",
+    "logical_not",
+    "logical_or",
+    "logical_xor",
+    "maximum",
+    "minimum",
+    "multiply",
+    "negative",
+    "nextafter",
+    "not_equal",
+    "positive",
+    "pow",
+    "real",
+    "reciprocal",
+    "remainder",
+    "round",
+    "sign",
+    "signbit",
+    "square",
+    "sum",
+    "subtract",
+    "tan",
+    "tanh",
+    "trunc",
+    "var",
+    "where",
+    "zeros",
+    "zeros_like",
+]
+
+lin_alg_funcs = [
     "concat",
     "diagonal",
     "expand_dims",
@@ -16,15 +92,13 @@
     "tensordot",
     "transpose",
     "vecdot",
-    "T",
-    "mT",
-    "take",
-    "take_along_axis",
-)
-reducers = ("sum", "prod", "min", "max", "std", "mean", "var", "any", "all", "slice", "count_nonzero")
+]
+
+lin_alg_attrs = ["T", "mT"]
+reducers = ["sum", "prod", "min", "max", "std", "mean", "var", "any", "all", "count_nonzero"]
 
 # All the available constructors and reducers necessary for the (string) expression evaluator
-constructors = (
+constructors = [
     "arange",
     "linspace",
     "fromiter",
@@ -38,9 +112,10 @@
     "ones_like",
     "empty_like",
     "eye",
-)
+    "nans",
+]
 # Note that, as reshape is accepted as a method too, it should always come last in the list
-constructors += ("reshape",)
+constructors += ["reshape"]
 
 
 # --- Shape utilities ---
@@ -336,6 +411,7 @@ def visit_Call(self, node):  # noqa : C901
                 "zeros_like",
                 "empty_like",
                 "ones_like",
+                "nans",
             ):
                 if node.args:
                     shape_arg = node.args[0]
@@ -412,6 +488,11 @@ def visit_Call(self, node):  # noqa : C901
             return FUNCTIONS[base_name](*args, **kwargs)
 
         shapes = [s for s in args if s is not None]
+        if base_name not in elementwise_funcs:
+            warnings.warn(
+                f"Function shape parser not implemented for {base_name}.", UserWarning, stacklevel=2
+            )
+        # default to elementwise but print warning that function not defined explicitly
         return elementwise(*shapes) if shapes else ()
 
     def visit_Compare(self, node):

tests/ndarray/test_reductions.py

@@ -459,12 +459,10 @@ def test_reduction_index():
     assert arr.shape == newarr.shape
 
     a = blosc2.ones(shape=(0, 0))
-    arr = blosc2.lazyexpr("sum(a, axis=(0, 1, 2))", {"a": a})
     with pytest.raises(np.exceptions.AxisError):
-        newarr = arr.compute()
-    arr = blosc2.lazyexpr("sum(a, axis=(0, 0))", {"a": a})
+        arr = blosc2.lazyexpr("sum(a, axis=(0, 1, 2))", {"a": a})
     with pytest.raises(ValueError):
-        newarr = arr.compute()
+        arr = blosc2.lazyexpr("sum(a, axis=(0, 0))", {"a": a})
 
 
 @pytest.mark.parametrize("idx", [0, 1, (0,), slice(1, 2), (slice(0, 1),), slice(0, 4), (0, 2)])