Fix CTable.where() applying a short mask to the wrong rows. Closes #607.

Author: FrancescAlted

src/blosc2/ctable.py

@@ -13439,10 +13439,21 @@ def where(  # noqa: C901
 
         filter_len = len(filter)
         if filter_len != target_len:
-            if filter_len == self.nrows:
-                physical = blosc2.zeros(target_len, dtype=np.bool_)
-                physical[self._valid_rows] = filter[:]
-                filter = physical
+            n_live = self.nrows
+            if filter_len <= n_live:
+                # A mask no longer than the live-row count is logical: entry i
+                # selects the i-th live row of this view.  A short one simply
+                # leaves the trailing live rows unselected.  Padding it out to
+                # the physical length instead would align it with the
+                # underlying column, so it would pick up rows outside the view.
+                if filter_len < n_live:
+                    logical = blosc2.zeros(n_live, dtype=np.bool_)
+                    logical[:filter_len] = filter[:]
+                    filter = logical
+                if n_live != target_len:
+                    physical = blosc2.zeros(target_len, dtype=np.bool_)
+                    physical[self._valid_rows] = filter[:]
+                    filter = physical
                 filter_intersected = True
             elif filter_len > target_len:
                 filter = filter[:target_len]

tests/ctable/test_where_expressions.py

@@ -73,3 +73,63 @@ def test_where_string_expression_must_be_boolean():
 
     with pytest.raises(TypeError, match="Expected boolean"):
         t.where("value * category")
+
+
+@dataclass
+class IdRow:
+    id: int = blosc2.field(blosc2.int64(), default=0)
+
+
+def _id_table(n=6):
+    t = blosc2.CTable(IdRow, expected_size=n)
+    arr = np.empty(n, dtype=[("id", "<i8")])
+    arr["id"] = np.arange(n)
+    t.extend(arr, validate=False)
+    return t
+
+
+def test_where_short_mask_is_relative_to_the_view():
+    """A mask shorter than the view selects the view's own rows.
+
+    It used to be padded out to the physical column length, which aligned it
+    with the underlying column and let rows outside the view through.
+    """
+    view = _id_table()[1:4]  # ids [1, 2, 3]
+
+    # Predicate built from a slice of the view's column: [1, 2] > 1 -> [F, T].
+    result = view.where(view["id"][0:2] > 1)
+
+    np.testing.assert_array_equal(result.id[:], np.array([2], dtype=np.int64))
+
+
+def test_where_short_mask_skips_deleted_rows():
+    t = _id_table()
+    t.delete([0, 2])  # live ids [1, 3, 4, 5]
+
+    result = t.where(np.array([True, False]))
+
+    np.testing.assert_array_equal(result.id[:], np.array([1], dtype=np.int64))
+
+
+@pytest.mark.parametrize(
+    ("mask", "expected"),
+    [
+        (np.array([], dtype=bool), []),
+        (np.array([True, True]), [1, 2]),
+        (np.array([False, True, True]), [2, 3]),
+    ],
+)
+def test_where_mask_lengths_on_a_sliced_view(mask, expected):
+    view = _id_table()[1:4]  # ids [1, 2, 3]
+
+    result = view.where(mask)
+
+    np.testing.assert_array_equal(result.id[:], np.array(expected, dtype=np.int64))
+
+
+def test_where_short_mask_on_a_sorted_view_follows_sort_order():
+    ordered = _id_table().sort_by("id", ascending=False)  # ids [5, 4, 3, 2, 1, 0]
+
+    result = ordered.where(np.array([True, False]))
+
+    np.testing.assert_array_equal(result.id[:], np.array([5], dtype=np.int64))