Merge pull request #407 from Blosc/fixChaining

Fix chaining

Author: FrancescAlted

src/blosc2/lazyexpr.py

@@ -630,6 +630,107 @@ def visit_Call(self, node):
     return set(visitor.operands)
 
 
+def conserve_functions(  # noqa: C901
+    expression: str,
+    operands_old: dict[str, blosc2.NDArray | blosc2.LazyExpr],
+    operands_new: dict[str, blosc2.NDArray | blosc2.LazyExpr],
+) -> tuple(str, dict[str, blosc2.NDArray]):
+    """
+    Given an expression in string form, return its operands.
+
+    Parameters
+    ----------
+    expression : str
+        The expression in string form.
+
+    operands_old: dict[str : blosc2.ndarray | blosc2.LazyExpr]
+        Dict of operands from expression prior to eval.
+
+    operands_new: dict[str : blosc2.ndarray | blosc2.LazyExpr]
+        Dict of operands from expression after eval.
+    Returns
+    -------
+    newexpression
+        A modified string expression with the functions/constructors conserved and
+        true operands rebased and written in o- notation.
+    newoperands
+        Dict of the set of rebased operands.
+    """
+
+    operand_to_key = {id(v): k for k, v in operands_new.items()}
+    for k, v in operands_old.items():  # extend operands_to_key with old operands
+        if isinstance(
+            v, blosc2.LazyExpr
+        ):  # unroll operands in LazyExpr (only necessary when have reduced a lazyexpr)
+            d = v.operands
+        else:
+            d = {k: v}
+        for newk, newv in d.items():
+            try:
+                operand_to_key[id(newv)]
+            except KeyError:
+                newk = (
+                    f"o{len(operands_new)}" if newk in operands_new else newk
+                )  # possible that names coincide
+                operand_to_key[id(newv)] = newk
+                operands_new[newk] = newv
+
+    class OperandVisitor(ast.NodeVisitor):
+        def __init__(self):
+            self.operandmap = {}
+            self.operands = {}
+            self.opcounter = 0
+            self.function_names = set()
+
+        def update_func(self, localop):
+            k = operand_to_key[id(localop)]
+            if k not in self.operandmap:
+                newkey = f"o{self.opcounter}"
+                self.operands[newkey] = operands_new[k]
+                self.operandmap[k] = newkey
+                self.opcounter += 1
+                return newkey
+            else:
+                return self.operandmap[k]
+
+        def visit_Name(self, node):
+            if node.id == "np":  # Skip NumPy namespace (e.g. np.int8, which will be treated separately)
+                return
+            if node.id in self.function_names:  # Skip function names
+                return
+            elif node.id not in dtype_symbols:
+                localop = operands_old[node.id]
+                if isinstance(localop, blosc2.LazyExpr):
+                    newexpr = localop.expression
+                    for (
+                        opname,
+                        v,
+                    ) in localop.operands.items():  # expression operands already in terms of basic operands
+                        newopname = self.update_func(v)
+                        newexpr = re.sub(
+                            rf"(?<=\s){opname}|(?<=\(){opname}", newopname, newexpr
+                        )  # replace with newopname
+                    node.id = newexpr
+                else:
+                    node.id = self.update_func(localop)
+            else:
+                pass
+            self.generic_visit(node)
+
+        def visit_Call(self, node):
+            if isinstance(
+                node.func, ast.Name
+            ):  # visits Call first, then Name, so don't increment operandcounter yet
+                self.function_names.add(node.func.id)
+            self.generic_visit(node)
+
+    tree = ast.parse(expression)
+    visitor = OperandVisitor()
+    visitor.visit(tree)
+    newexpression, newoperands = ast.unparse(tree), visitor.operands
+    return newexpression, newoperands
+
+
 class TransformNumpyCalls(ast.NodeTransformer):
     def __init__(self):
         self.replacements = {}
@@ -2678,10 +2779,14 @@ def _new_expr(cls, expression, operands, guess, out=None, where=None, ne_args=No
             _dtype = new_expr.dtype
             _shape = new_expr.shape
             if isinstance(new_expr, blosc2.LazyExpr):
-                # Restore the original expression and operands
-                new_expr.expression = f"({_expression})"  # forcibly add parenthesis
-                new_expr.expression_tosave = _expression
-                new_expr.operands = _operands
+                # DO NOT restore the original expression and operands
+                # Instead rebase operands and restore only constructors
+                expression_, operands_ = conserve_functions(
+                    _expression, _operands, new_expr.operands | local_vars
+                )
+                new_expr.expression = f"({expression_})"  # force parenthesis
+                new_expr.expression_tosave = expression
+                new_expr.operands = operands_
                 new_expr.operands_tosave = operands
             else:
                 # An immediate evaluation happened (e.g. all operands are numpy arrays)

tests/ndarray/test_lazyexpr.py

@@ -1354,45 +1354,52 @@ def test_chain_expressions():
     le3_ = blosc2.lazyexpr("(le2 & (b < 0))", {"le2": le2_, "b": b})
     assert (le3_[:] == le3[:]).all()
 
-    # TODO: This test should pass eventually
-    # le1 = a ** 3 + blosc2.sin(a ** 2)
-    # le2 = le1 < c
-    # le3 = (b < 0)
-    # le4 = le2 & le3
-    # le1_ = blosc2.lazyexpr("a ** 3 + sin(a ** 2)", {"a": a})
-    # le2_ = blosc2.lazyexpr("(le1 < c)", {"le1": le1_, "c": c})
-    # le3_ = blosc2.lazyexpr("(b < 0)", {"b": b})
-    # le4_ = blosc2.lazyexpr("(le2 & le3)", {"le2": le2_, "le3": le3_})
-    # assert (le4_[:] == le4[:]).all()
-
-
-# TODO: Test the chaining of multiple persistent lazy expressions
-# def test_chain_persistentexpressions():
-#     N = 1_000
-#     dtype = "float64"
-#     a = blosc2.linspace(0, 1, N * N, dtype=dtype, shape=(N, N), urlpath="a.b2nd", mode="w")
-#     b = blosc2.linspace(1, 2, N * N, dtype=dtype, shape=(N, N), urlpath="b.b2nd", mode="w")
-#     c = blosc2.linspace(0, 1, N, dtype=dtype, shape=(N,), urlpath="c.b2nd", mode="w")
-#
-#     le1 = a ** 3 + blosc2.sin(a ** 2)
-#     le2 = le1 < c
-#     le3 = (b < 0)
-#     le4 = le2 & le3
-#
-#     le1_ = blosc2.lazyexpr("a ** 3 + sin(a ** 2)", {"a": a})
-#     le1_.save("expr1.b2nd", mode="w")
-#     myle1 = blosc2.open("expr1.b2nd")
-#
-#     le2_ = blosc2.lazyexpr("(le1 < c)", {"le1": myle1, "c": c})
-#     le2_.save("expr2.b2nd", mode="w")
-#     myle2 = blosc2.open("expr2.b2nd")
-#
-#     le3_ = blosc2.lazyexpr("(b < 0)", {"b": b})
-#     le3_.save("expr3.b2nd", mode="w")
-#     myle3 = blosc2.open("expr3.b2nd")
-#
-#     le4_ = blosc2.lazyexpr("(le2 & le3)", {"le2": myle2, "le3": myle3})
-#     le4_.save("expr4.b2nd", mode="w")
-#     myle4 = blosc2.open("expr4.b2nd")
-#     print((myle4[:] == le4[:]).all())
-#
+    le1 = a**3 + blosc2.sin(a**2)
+    le2 = le1 < c
+    le3 = b < 0
+    le4 = le2 & le3
+    le1_ = blosc2.lazyexpr("a ** 3 + sin(a ** 2)", {"a": a})
+    le2_ = blosc2.lazyexpr("(le1 < c)", {"le1": le1_, "c": c})
+    le3_ = blosc2.lazyexpr("(b < 0)", {"b": b})
+    le4_ = blosc2.lazyexpr("(le2 & le3)", {"le2": le2_, "le3": le3_})
+    assert (le4_[:] == le4[:]).all()
+
+    # TODO: Eventually this test should pass
+    # expr1 = blosc2.lazyexpr("arange(N) + b")
+    # expr2 = blosc2.lazyexpr("a * b + 1")
+    # expr = blosc2.lazyexpr("expr1 - expr2")
+    # expr_final = blosc2.lazyexpr("expr * expr")
+    # nres = (expr * expr)[:]
+    # res = expr_final.compute()
+    # np.testing.assert_allclose(res[:], nres)
+
+
+# Test the chaining of multiple persistent lazy expressions
+def test_chain_persistentexpressions():
+    N = 1_000
+    dtype = "float64"
+    a = blosc2.linspace(0, 1, N * N, dtype=dtype, shape=(N, N), urlpath="a.b2nd", mode="w")
+    b = blosc2.linspace(1, 2, N * N, dtype=dtype, shape=(N, N), urlpath="b.b2nd", mode="w")
+    c = blosc2.linspace(0, 1, N, dtype=dtype, shape=(N,), urlpath="c.b2nd", mode="w")
+
+    le1 = a**3 + blosc2.sin(a**2)
+    le2 = le1 < c
+    le3 = b < 0
+    le4 = le2 & le3
+
+    le1_ = blosc2.lazyexpr("a ** 3 + sin(a ** 2)", {"a": a})
+    le1_.save("expr1.b2nd", mode="w")
+    myle1 = blosc2.open("expr1.b2nd")
+
+    le2_ = blosc2.lazyexpr("(le1 < c)", {"le1": myle1, "c": c})
+    le2_.save("expr2.b2nd", mode="w")
+    myle2 = blosc2.open("expr2.b2nd")
+
+    le3_ = blosc2.lazyexpr("(b < 0)", {"b": b})
+    le3_.save("expr3.b2nd", mode="w")
+    myle3 = blosc2.open("expr3.b2nd")
+
+    le4_ = blosc2.lazyexpr("(le2 & le3)", {"le2": myle2, "le3": myle3})
+    le4_.save("expr4.b2nd", mode="w")
+    myle4 = blosc2.open("expr4.b2nd")
+    assert (myle4[:] == le4[:]).all()