Merge pull request #413 from Blosc/addLazySlice

Add slicing to string lazyexprs

Author: FrancescAlted

src/blosc2/lazyexpr.py

@@ -116,7 +116,7 @@ def ne_evaluate(expression, local_dict=None, **kwargs):
 constructors = ("arange", "linspace", "fromiter", "zeros", "ones", "empty", "full", "frombuffer")
 # Note that, as reshape is accepted as a method too, it should always come last in the list
 constructors += ("reshape",)
-reducers = ("sum", "prod", "min", "max", "std", "mean", "var", "any", "all")
+reducers = ("sum", "prod", "min", "max", "std", "mean", "var", "any", "all", "slice")
 
 functions = [
     "sin",
@@ -541,7 +541,7 @@ def compute_smaller_slice(larger_shape, smaller_shape, larger_slice):
 
 # Define the patterns for validation
 validation_patterns = [
-    r"[\;\[\:]",  # Flow control characters
+    r"[\;]",  # Flow control characters
     r"(^|[^\w])__[\w]+__($|[^\w])",  # Dunder methods
     r"\.\b(?!real|imag|(\d*[eE]?[+-]?\d+)|(\d*[eE]?[+-]?\d+j)|\d*j\b|(sum|prod|min|max|std|mean|var|any|all|where)"
     r"\s*\([^)]*\)|[a-zA-Z_]\w*\s*\([^)]*\))",  # Attribute patterns
@@ -551,7 +551,20 @@ def compute_smaller_slice(larger_shape, smaller_shape, larger_slice):
 _blacklist_re = re.compile("|".join(validation_patterns))
 
 # Define valid method names
-valid_methods = {"sum", "prod", "min", "max", "std", "mean", "var", "any", "all", "where", "reshape"}
+valid_methods = {
+    "sum",
+    "prod",
+    "min",
+    "max",
+    "std",
+    "mean",
+    "var",
+    "any",
+    "all",
+    "where",
+    "reshape",
+    "slice",
+}
 valid_methods |= {"int8", "int16", "int32", "int64", "uint8", "uint16", "uint32", "uint64"}
 valid_methods |= {"float32", "float64", "complex64", "complex128"}
 valid_methods |= {"bool", "str", "bytes"}
@@ -579,7 +592,7 @@ def validate_expr(expr: str) -> None:
         raise ValueError(f"'{expr}' is not a valid expression.")
 
     # Check for invalid characters not covered by the tokenizer
-    invalid_chars = re.compile(r"[^\w\s+\-*/%().,=<>!&|~^]")
+    invalid_chars = re.compile(r"[^\w\s+\-*/%()[].,=<>!&|~^]")
     if invalid_chars.search(skip_quotes) is not None:
         invalid_chars = invalid_chars.findall(skip_quotes)
         raise ValueError(f"Expression {expr} contains invalid characters: {invalid_chars}")
@@ -731,6 +744,41 @@ def visit_Call(self, node):
     return newexpression, newoperands
 
 
+def convert_to_slice(expression):
+    """
+    Takes expression and converts all instances of [] to .slice(....)
+
+    Parameters
+    ----------
+    expression: str
+
+    Returns
+    -------
+    new_expr : str
+    """
+
+    new_expr = ""
+    skip_to_char = 0
+    for i, expr_i in enumerate(expression):
+        if i < skip_to_char:
+            continue
+        if expr_i == "[":
+            k = expression[i:].find("]")  # start checking from after [
+            slice_convert = expression[i : i + k + 1]  # include [ and ]
+            slicer = eval(f"np.s_{slice_convert}")
+            slicer = (slicer,) if isinstance(slicer, slice) else slicer  # standardise to tuple
+            if any(isinstance(el, str) for el in slicer):  # handle fields
+                raise ValueError("Cannot handle fields for slicing lazy expressions.")
+            slicer = str(slicer)
+            # use slice so that lazyexpr uses blosc arrays internally
+            # (and doesn't decompress according to getitem syntax)
+            new_expr += ".slice(" + slicer + ")"
+            skip_to_char = i + k + 1
+        else:
+            new_expr += expr_i
+    return new_expr
+
+
 class TransformNumpyCalls(ast.NodeTransformer):
     def __init__(self):
         self.replacements = {}
@@ -2543,6 +2591,7 @@ def find_args(expr):
     def _compute_expr(self, item, kwargs):  # noqa: C901
         if any(method in self.expression for method in reducers):
             # We have reductions in the expression (probably coming from a string lazyexpr)
+            # Also includes slice
             _globals = get_expr_globals(self.expression)
             lazy_expr = eval(self.expression, _globals, self.operands)
             if not isinstance(lazy_expr, blosc2.LazyExpr):
@@ -2775,6 +2824,7 @@ def save(self, urlpath=None, **kwargs):
     def _new_expr(cls, expression, operands, guess, out=None, where=None, ne_args=None):
         # Validate the expression
         validate_expr(expression)
+        expression = convert_to_slice(expression)
         if guess:
             # The expression has been validated, so we can evaluate it
             # in guessing mode to avoid computing reductions

tests/ndarray/test_reductions.py

@@ -428,3 +428,34 @@ def test_reduction_index():
     assert arr[:10].shape == (10,)
     assert arr[0].shape == (1,)
     assert arr.shape == newarr.shape
+
+
+def test_slice_lazy():
+    shape = (20, 20)
+    a = blosc2.linspace(0, 20, num=np.prod(shape), shape=shape)
+    arr = blosc2.lazyexpr("anarr.slice(slice(10,15)) + 1", {"anarr": a})
+    newarr = arr.compute()
+    np.testing.assert_allclose(newarr[:], a.slice(slice(10, 15))[:] + 1)
+
+
+def test_slicebrackets_lazy():
+    shape = (20, 20)
+    a = blosc2.linspace(0, 20, num=np.prod(shape), shape=shape)
+    arr = blosc2.lazyexpr("anarr[10:15] + 1", {"anarr": a})
+    newarr = arr.compute()
+    np.testing.assert_allclose(newarr[:], a[10:15] + 1)
+
+    # Try with getitem
+    a = blosc2.linspace(0, 20, num=np.prod(shape), shape=shape)
+    arr = blosc2.lazyexpr("anarr[10:15] + 1", {"anarr": a})
+    newarr = arr[:3]
+    res = a[10:15] + 1
+    np.testing.assert_allclose(newarr, res[:3])
+
+    arr = blosc2.lazyexpr("anarr[10:15, 2:9] + 1", {"anarr": a})
+    newarr = arr.compute()
+    np.testing.assert_allclose(newarr[:], a[10:15, 2:9] + 1)
+
+    arr = blosc2.lazyexpr("anarr[10:15][2:9] + 1", {"anarr": a})
+    newarr = arr.compute()
+    np.testing.assert_allclose(newarr[:], a[10:15][2:9] + 1)