Merge pull request #411 from Blosc/fixChaining

Allow use of where for string expressions and integer operands

Author: lshaw8317

src/blosc2/lazyexpr.py

@@ -2091,6 +2091,8 @@ def dtype(self):
 
         operands = {
             key: np.ones(np.ones(len(value.shape), dtype=int), dtype=value.dtype)
+            if hasattr(value, "shape")
+            else value
             for key, value in self.operands.items()
         }
 
@@ -2784,6 +2786,20 @@ def _new_expr(cls, expression, operands, guess, out=None, where=None, ne_args=No
                 expression_, operands_ = conserve_functions(
                     _expression, _operands, new_expr.operands | local_vars
                 )
+                # if new_expr has where_args, must have come from where(...) - or possibly where(where(..
+                # since 5*where, where + ...  are evaluated eagerly
+                if hasattr(new_expr, "_where_args"):
+                    st = expression_.find("where(") + len(
+                        "where("
+                    )  # expr always begins where( - should have st = 6 always
+                    finalexpr = ""
+                    counter = 0
+                    for char in expression_[st:]:  # get rid of external where(...)
+                        finalexpr += char
+                        counter += 1 * (char == "(") - 1 * (char == ")")
+                        if counter == 0 and char == ",":
+                            break
+                    expression_ = finalexpr[:-1]  # remove trailing comma
                 new_expr.expression = f"({expression_})"  # force parenthesis
                 new_expr.expression_tosave = expression
                 new_expr.operands = operands_

tests/ndarray/test_lazyexpr.py

@@ -1364,14 +1364,13 @@ def test_chain_expressions():
     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)
+    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

tests/ndarray/test_lazyexpr_fields.py

@@ -226,6 +226,47 @@ def test_where(array_fixture):
     np.testing.assert_allclose(res, nres[sl])
 
 
+# Test expressions with where() and string comps
+def test_lazy_where(array_fixture):
+    sa1, sa2, nsa1, nsa2, a1, a2, a3, a4, na1, na2, na3, na4 = array_fixture
+
+    # Test 1: where
+    # Test with string expression
+    expr = blosc2.lazyexpr("where((a1 ** 2 + a2 ** 2) > (2 * a1 * a2 + 1), 0, a1)")
+    # Test with eval
+    res = expr.compute()
+    nres = ne_evaluate("where(na1**2 + na2**2 > 2 * na1 * na2 + 1, 0, na1)")
+    np.testing.assert_allclose(res[:], nres)
+    # Test with getitem
+    sl = slice(100)
+    res = expr[sl]
+    np.testing.assert_allclose(res, nres[sl])
+
+    # Test 2: sum of wheres
+    # Test with string expression
+    expr = blosc2.lazyexpr("where(a1 < 0, 10, a1) + where(a2 < 0, 3, a2)")
+    # Test with eval
+    res = expr.compute()
+    nres = ne_evaluate("where(na1 < 0, 10, na1) + where(na2 < 0, 3, na2)")
+    np.testing.assert_allclose(res[:], nres)
+
+    # Test 3: nested wheres
+    # Test with string expression
+    expr = blosc2.lazyexpr("where(where(a2 < 0, 3, a2) > 3, 10, a1)")
+    # Test with eval
+    res = expr.compute()
+    nres = ne_evaluate("where(where(na2 < 0, 3, na2) > 3, 10, na1)")
+    np.testing.assert_allclose(res[:], nres)
+
+    # Test 4: multiplied wheres
+    # Test with string expression
+    expr = blosc2.lazyexpr("1 * where(a2 < 0, 3, a2)")
+    # Test with eval
+    res = expr.compute()
+    nres = ne_evaluate("1 * where(na2 < 0, 3, na2)")
+    np.testing.assert_allclose(res[:], nres)
+
+
 # Test where with one parameter
 def test_where_one_param(array_fixture):
     sa1, sa2, nsa1, nsa2, a1, a2, a3, a4, na1, na2, na3, na4 = array_fixture