Merge pull request #414 from posit-dev/fix-segments-lazyframe-error

fix: segments-based LazyFrame error

Author: rich-iannone

pointblank/validate.py

@@ -21113,6 +21113,8 @@ def _evaluate_segments(self, validation_info):
                 continue
 
             # Evaluate the segments expression
+            seg_tuples = []
+
             try:
                 # Get the table for this step, it can either be:
                 # 1. the target table itself
@@ -21158,6 +21160,8 @@ def _evaluate_segments(self, validation_info):
 
             except Exception:  # pragma: no cover
                 validation.eval_error = True
+                expanded_validation_info.append(validation)
+                continue
 
             # For each segmentation resolved, create a new validation step and add it to the list of
             # expanded validation steps
@@ -22651,6 +22655,10 @@ def _seg_expr_from_string(data_tbl: Any, segments_expr: str) -> tuple[str, str]:
         data_nw = nw.from_native(data_tbl)
         unique_vals = data_nw.select(nw.col(segments_expr)).unique()
 
+        # LazyFrames must be collected before item indexing
+        if is_narwhals_lazyframe(unique_vals):
+            unique_vals = unique_vals.collect()
+
         # Convert to list of values
         seg_categories = unique_vals[segments_expr].to_list()
     elif tbl_type in IBIS_BACKENDS:

tests/test_segments.py

@@ -211,6 +211,20 @@ def test_segments_with_multiple_seg_groups(tbl_type):
         )
         .interrogate()
     )
+
     assert validation.n_passed(i=1, scalar=True) == 7
     assert validation.n_passed(i=2, scalar=True) == 3
     assert validation.n_passed(i=3, scalar=True) == 2
+
+
+def test_segments_str_lazyframe():
+    df = pl.DataFrame({"x": [1.0, -1.0, 2.0], "region": ["US", "US", "EU"]})
+    validation_lazy = Validate(data=df.lazy()).col_vals_gt("x", 0, segments="region").interrogate()
+    validation_eager = Validate(data=df).col_vals_gt("x", 0, segments="region").interrogate()
+
+    # Both should produce the same number of expanded steps
+    assert len(validation_lazy.validation_info) == len(validation_eager.validation_info)
+
+    # EU: 1 row, 1 passes; US: 2 rows, 1 passes
+    assert validation_lazy.n_passed(i=1, scalar=True) == validation_eager.n_passed(i=1, scalar=True)
+    assert validation_lazy.n_passed(i=2, scalar=True) == validation_eager.n_passed(i=2, scalar=True)