Add documentation of behaviour of get_item and compute for LazyExprs

Author: Luke Shaw

examples/ndarray/lazyexpr_where_indexing.py

@@ -0,0 +1,42 @@
+# Imports
+
+import numpy as np
+
+import blosc2
+
+N = 1000
+it = ((-x + 1, x - 2, 0.1 * x) for x in range(N))
+sa = blosc2.fromiter(
+    it, dtype=[("A", "i4"), ("B", "f4"), ("C", "f8")], shape=(N,), urlpath="sa-1M.b2nd", mode="w"
+)
+expr = sa["(A < B)"]
+A = sa["A"][:]
+B = sa["B"][:]
+C = sa["C"][:]
+temp = sa[:]
+indices = A < B
+idx = np.argmax(indices)
+
+# One might think that expr[:10] gives the first 10 elements of the evaluated expression, but this is not the case.
+# It actually computes the expression on the first 10 elements of the operands; since for some elements the condition
+# is False, the result will be shorter than 10 elements.
+# Returns less than 10 elements in general
+sliced = expr.compute(slice(0, 10))
+gotitem = expr[:10]
+np.testing.assert_array_equal(sliced[:], gotitem)
+np.testing.assert_array_equal(gotitem, temp[:10][indices[:10]])  # Equivalent syntax
+# Actually this makes sense since one can understand this as a request to compute on a portion of operands.
+# If one desires a portion of the result, one should compute the whole expression and then slice it.
+
+# Get first element for which condition is true
+sliced = expr.compute(idx)
+gotitem = expr[idx]
+# Arrays of one element
+np.testing.assert_array_equal(sliced[()], gotitem)
+np.testing.assert_array_equal(gotitem, temp[idx])
+
+# Should return void arrays here.
+sliced = expr.compute(0)
+gotitem = expr[0]
+np.testing.assert_array_equal(sliced[()], gotitem)
+np.testing.assert_array_equal(gotitem, temp[0])

src/blosc2/lazyexpr.py

@@ -1829,7 +1829,8 @@ def chunked_eval(  # noqa: C901
     operands: dict
         A dictionary containing the operands for the expression.
     item: int, slice or sequence of slices, optional
-        The slice(s) to be retrieved. Note that step parameter is not honored yet.
+        The slice(s) of the operands to be used in computation. Note that step parameter is not honored yet.
+        Item is used to slice the operands PRIOR to computation.
     kwargs: Any, optional
         Additional keyword arguments supported by the :func:`empty` constructor.  In addition,
         the following keyword arguments are supported:
@@ -2688,6 +2689,21 @@ def sort(self, order: str | list[str] | None = None) -> blosc2.LazyArray:
         return lazy_expr
 
     def compute(self, item=None, **kwargs) -> blosc2.NDArray:
+        """
+        Compute the expression with the given item and kwargs.
+        Parameters
+        ----------
+        item: int, slice or sequence of slices, optional
+            The slice(s) of the operands to be used in computation. Note that step parameter is not honored yet.
+            Item is used to slice the operands PRIOR to computation.
+        kwargs
+
+        Returns:
+        blosc2.NDArray or numpy.ndarray
+        -------
+
+        """
+
         # When NumPy ufuncs are called, the user may add an `out` parameter to kwargs
         if "out" in kwargs:
             kwargs["_output"] = kwargs.pop("out")
@@ -2722,6 +2738,17 @@ def compute(self, item=None, **kwargs) -> blosc2.NDArray:
         return result
 
     def __getitem__(self, item):
+        """
+        Apply LazyExpr on a slice of the oeprands.
+        Parameters
+        ----------
+        item: int, slice or sequence of slices, optional
+            The slice(s) of the operands to be used in computation. Note that step parameter is not honored yet.
+            Item is used to slice the operands PRIOR to computation.
+        Returns:
+        numpy.ndarray
+        """
+
         kwargs = {"_getitem": True}
         return self.compute(item, **kwargs)