CTable.__getitem__ now supports boolean blosc2.LazyExpr / boolean blosc2.NDArray

Author: FrancescAlted

doc/reference/ctable.rst

@@ -42,6 +42,9 @@ explicitly call :meth:`~blosc2.CTable.to_arrow` or iterate with
        * ``str`` — column name returns a :class:`Column`; any other string
          is interpreted as a boolean expression and behaves like
          :meth:`where`.
+       * boolean :class:`~blosc2.LazyExpr` / :class:`~blosc2.NDArray` —
+         filtered row view, same as :meth:`where`, e.g.
+         ``t[t.temperature_f > 70]``.
        * ``int`` — single row as a namedtuple-like object.
        * ``slice`` — row-range view.
        * ``list[int]`` / ``ndarray[int]`` — gathered-row view.
@@ -190,11 +193,14 @@ CTable indexing is type-driven::
     t[3:8]                      # row view
     t[[1, 4, 7]]                # gathered-row view
     t[mask]                     # filtered row view
+    t[t.amount > 100]           # LazyExpr filtered row view, like where()
     t[["region", "amount"]]   # projected column view
 
 String keys first try exact column-name lookup.  If the string is not a
 column name, it is interpreted as a boolean expression and behaves like
-:meth:`CTable.where`.
+:meth:`CTable.where`.  Boolean :class:`~blosc2.LazyExpr` and boolean
+:class:`~blosc2.NDArray` keys also behave like :meth:`CTable.where`, so computed
+column predicates such as ``t[t.temperature_f > 70]`` are supported.
 
 For explicit filtered projection, use::
 

src/blosc2/ctable.py

@@ -4356,11 +4356,47 @@ def __array__(self, dtype=None, copy=None):
         return arr.copy() if copy else arr
 
     def __getitem__(self, key):
-        """Type-driven indexing: column name, boolean expression, row int/slice, mask, or column list."""
+        """Type-driven indexing for columns, rows, projections, and filters.
+
+        Supported keys are:
+
+        - ``str``: return a :class:`Column` when it matches a stored or computed
+          column name; otherwise evaluate it as a boolean expression via
+          :meth:`where`.
+        - boolean :class:`blosc2.LazyExpr` or :class:`blosc2.NDArray`: return the
+          same filtered view as :meth:`where`, e.g. ``t[t.temperature_f > 70]``.
+        - ``int``: return one live row as a namedtuple-like object.
+        - ``slice``: return a row-range view.
+        - integer array/list: return a gathered-row view.
+        - boolean NumPy array/list: return a boolean-mask filtered view.
+        - string list: return a column-projection view, equivalent to
+          :meth:`select`.
+
+        Examples
+        --------
+        Access columns and rows::
+
+            temps = t["temperature"]
+            first = t[0]
+            view = t[10:20]
+
+        Filter rows with a string expression, a stored-column expression, or a
+        computed-column expression::
+
+            warm = t["temperature > 20"]
+            warm_active = t[(t.temperature > 20) & t.active]
+            hot_fahrenheit = t[t.temperature_f > 70]
+
+        Project columns::
+
+            slim = t[["sensor_id", "temperature_f"]]
+        """
         if isinstance(key, str):
             if key in self._cols or key in self._computed_cols:
                 return Column(self, key)
             return self.where(key)
+        if isinstance(key, (blosc2.NDArray, blosc2.LazyExpr)) and getattr(key, "dtype", None) == np.bool_:
+            return self.where(key)
         if isinstance(key, tuple):
             raise TypeError("Tuple indexing is not supported for CTable in V1")
         return self._getitem_row_selector(key)

tests/ctable/test_ctable_computed_cols.py

@@ -163,6 +163,16 @@ def test_computed_column_where_via_col():
     assert len(view) == 3  # 9, 16, 25
 
 
+def test_getitem_boolean_lazyexpr_matches_where_for_computed_column():
+    t = _make_invoice_table(5)
+    t.add_computed_column("total", lambda cols: cols["price"] * cols["qty"])
+    expr = t.total >= 9
+    view_getitem = t[expr]
+    view_where = t.where(expr)
+    assert len(view_getitem) == len(view_where) == 3
+    np.testing.assert_array_equal(view_getitem["price"][:], view_where["price"][:])
+
+
 # ---------------------------------------------------------------------------
 # 6. Expression composability
 # ---------------------------------------------------------------------------