Update

Author: borchero

Cargo.lock

@@ -1,8 +1,13 @@
 from typing import overload
 
+import polars as pl
+
 def format_rule_failures(
     failures: list[tuple[str, int]],
-    examples: dict[str, list[str]] | None = None,
+    *,
+    failures_from: pl.DataFrame | None,
+    examples_from: pl.DataFrame | None,
+    primary_key_columns: list[str],
 ) -> str:
     """
     Format rule failures with the same logic that produces validation errors from the
@@ -11,9 +16,13 @@ def format_rule_failures(
     Args:
         failures: The name of the failures and their counts. This should only include
             failures with a count of at least 1.
-        examples: Optional mapping from rule name to a list of example row strings.
-            When provided, up to ``len(examples[rule])`` distinct examples are included
-            in the formatted message for each rule.
+        failures_from: The data frame containing the rule columns providing the
+            failures.
+        max_examples: The maximum number of examples to include for each failure. No
+            effect if `examples_from` is not provided.
+        primary_key_columns: The primary key columns of the schema for which to format
+            rule failures. This is only relevant if `examples_from` is provided and
+            allows for better error messages for the "primary_key" rule.
 
     Returns:
         The formatted rule failures.

dataframely/_native.pyi

@@ -60,6 +60,7 @@ def all_rules_required(
     null_is_valid: bool = True,
     schema_name: str,
     data_columns: Iterable[IntoExpr] | None = None,
+    primary_key_columns: list[str] | None,
 ) -> pl.Expr:
     """Execute :mod:`~polars.all_horizontal` and `.all` for a set of rules.
 
@@ -80,6 +81,8 @@ def all_rules_required(
         data_columns: Optional data columns to include for generating example rows in
             error messages. If provided, up to 5 distinct example rows are included
             for each failing rule.
+        primary_key_columns: Optional list of primary key columns which are used for
+            better error messages if data columns are provided.
 
     Returns:
         A scalar boolean expression.
@@ -95,6 +98,7 @@ def all_rules_required(
             "null_is_valid": null_is_valid,
             "schema_name": schema_name,
             "num_rule_columns": num_rule_columns,
+            "primary_key_columns": primary_key_columns or [],
         },
         use_abs_path=True,
         is_elementwise=True,

dataframely/_plugin.py

@@ -409,13 +409,19 @@ def validate(
             # information to properly construct a useful error message.
             filtered, failures = cls.filter(data, cast=cast, eager=True)
             if any(len(failure) > 0 for failure in failures.values()):
-                errors = {
-                    member: format_rule_failures(
-                        list(failure.counts().items()), failure.examples()
+                errors: dict[str, str] = {}
+                for member, failure in failures.items():
+                    if len(failure) == 0:
+                        continue
+
+                    counts = failure.counts()
+                    errors[member] = format_rule_failures(
+                        list(counts.items()),
+                        failures_from=failure._df.select(counts.keys()),
+                        examples_from=failure.invalid(),
+                        primary_key_columns=cls.member_schemas()[member].primary_key(),
                     )
-                    for member, failure in failures.items()
-                    if len(failure) > 0
-                }
+
                 details = [
                     f" > Member '{member}' failed validation:\n"
                     + textwrap.indent(error, "   ")
@@ -453,7 +459,11 @@ def validate(
                         )
                         .filter(
                             all_rules_required(
-                                filter_names, null_is_valid=False, schema_name=name
+                                filter_names,
+                                null_is_valid=False,
+                                schema_name=name,
+                                data_columns=cls.common_primary_key(),
+                                primary_key_columns=cls.common_primary_key(),
                             )
                         )
                         .drop(filter_names)

dataframely/collection/collection.py

@@ -146,22 +146,6 @@ def counts(self) -> dict[str, int]:
         """
         return _compute_counts(self._df, self._rule_columns)
 
-    def examples(self, max_examples: int = 5) -> dict[str, list[str]]:
-        """Example rows for each failing rule.
-
-        For each rule that has at least one failure, returns up to `max_examples`
-        distinct example rows (as formatted strings) from the original data columns.
-
-        Args:
-            max_examples: The maximum number of distinct example rows to return per
-                rule.
-
-        Returns:
-            A mapping from rule name to a list of example row strings. Rules with no
-            failures are not included.
-        """
-        return _compute_examples(self._df, self._rule_columns, max_examples)
-
     def cooccurrence_counts(self) -> dict[frozenset[str], int]:
         """The number of validation failures per co-occurring rule validation failure.
 
@@ -425,28 +409,6 @@ def _compute_counts(df: pl.DataFrame, rule_columns: list[str]) -> dict[str, int]
     }
 
 
-def _compute_examples(
-    df: pl.DataFrame, rule_columns: list[str], max_examples: int
-) -> dict[str, list[str]]:
-    if len(rule_columns) == 0:
-        return {}
-
-    data_columns = [c for c in df.columns if c not in rule_columns]
-    if not data_columns:
-        return {}
-
-    result = {}
-    for rule_name in rule_columns:
-        failing = df.filter(pl.col(rule_name).not_())
-        if len(failing) == 0:
-            continue
-        examples_df = (
-            failing.select(data_columns).unique(maintain_order=True).head(max_examples)
-        )
-        result[rule_name] = [str(row) for row in examples_df.to_dicts()]
-    return result
-
-
 def _compute_cooccurrence_counts(
     df: pl.DataFrame, rule_columns: list[str]
 ) -> dict[frozenset[str], int]:

dataframely/filter_result.py

@@ -576,9 +576,13 @@ def validate(
         if eager:
             out, failure = cls.filter(df, cast=cast, eager=True)
             if len(failure) > 0:
+                counts = failure.counts()
                 raise ValidationError(
                     format_rule_failures(
-                        list(failure.counts().items()), failure.examples()
+                        list(counts.items()),
+                        failures_from=failure._df.select(counts.keys()),
+                        examples_from=failure.invalid(),
+                        primary_key_columns=cls.primary_key(),
                     )
                 )
             return out
@@ -594,6 +598,7 @@ def validate(
                             rules.keys(),
                             schema_name=cls.__name__,
                             data_columns=cls.column_names(),
+                            primary_key_columns=cls.primary_key(),
                         )
                     )
                     .drop(rules.keys())

dataframely/schema.py

@@ -2,8 +2,6 @@ mod rule_failure;
 mod utils;
 mod validation_error;
 
-use std::collections::{HashMap, HashSet};
-
 use polars::prelude::*;
 use polars_core::POOL;
 use pyo3_polars::derive::polars_expr;
@@ -64,55 +62,11 @@ pub fn all_rules(inputs: &[Series]) -> PolarsResult<Series> {
 struct RequiredValidationKwargs {
     schema_name: String,
     null_is_valid: bool,
+    primary_key_columns: Option<Vec<String>>,
     #[serde(default)]
     num_rule_columns: Option<usize>,
 }
 
-/// The maximum number of distinct example rows included in validation error messages.
-const MAX_EXAMPLES: usize = 5;
-
-/// Format a single data row (at `row_idx`) from the given data series as a Python-like dict string.
-fn format_example_row(data_series: &[Series], row_idx: usize) -> String {
-    let kvs: Vec<String> = data_series
-        .iter()
-        .map(|s| {
-            let val = s.get(row_idx).unwrap_or(AnyValue::Null);
-            format!("'{}': {}", s.name(), val)
-        })
-        .collect();
-    format!("{{{}}}", kvs.join(", "))
-}
-
-/// Compute up to `max_examples` distinct example rows for a failing rule.
-fn compute_examples(
-    bool_ca: &BooleanChunked,
-    null_is_valid: bool,
-    data_series: &[Series],
-    max_examples: usize,
-) -> Vec<String> {
-    let mut seen: HashSet<String> = HashSet::new();
-    let mut examples: Vec<String> = Vec::new();
-
-    for (i, val) in bool_ca.iter().enumerate() {
-        let is_failure = match val {
-            Some(false) => true,
-            None => !null_is_valid,
-            _ => false,
-        };
-        if is_failure {
-            let row_str = format_example_row(data_series, i);
-            if seen.insert(row_str.clone()) {
-                examples.push(row_str);
-                if examples.len() >= max_examples {
-                    break;
-                }
-            }
-        }
-    }
-
-    examples
-}
-
 /// Reduce a set of boolean columns into a single boolean scalar, AND-ing all values.
 /// Null values are treated as `true`.
 /// In contrast to `all_rules`, this function raises an error if the returned value would be
@@ -140,26 +94,26 @@ pub fn all_rules_required(
         return Ok(column.take_materialized_series());
     }
 
-    // Compute examples for each failing rule using the data columns.
-    let examples: HashMap<String, Vec<String>> = if data_inputs.is_empty() {
-        HashMap::new()
-    } else {
-        failures
+    // Aggregate failures into a validation error
+    let failures_from = DataFrame::new(
+        rule_inputs[0].len(),
+        rule_inputs
             .iter()
-            .map(|failure| {
-                let rule_series = rule_inputs
-                    .iter()
-                    .find(|s| s.name().as_str() == failure.rule)
-                    .expect("failing rule not found in inputs");
-                let bool_ca = as_bool(rule_series)?;
-                let examples =
-                    compute_examples(bool_ca, kwargs.null_is_valid, data_inputs, MAX_EXAMPLES);
-                Ok((failure.rule.to_string(), examples))
-            })
-            .collect::<PolarsResult<HashMap<_, _>>>()?
-    };
-
-    // Aggregate failure counts into a validation error.
-    let error = RuleValidationError::new(failures);
-    Err(polars_err!(ComputeError: format!("\n{}", error.to_string(Some(&kwargs.schema_name), Some(&examples)))))
+            .map(|s| s.clone().into_column())
+            .collect(),
+    )?;
+    let examples_from = DataFrame::new(
+        data_inputs[0].len(),
+        data_inputs
+            .iter()
+            .map(|s| s.clone().into_column())
+            .collect(),
+    )?;
+    let error = RuleValidationError::new(
+        failures,
+        Some(failures_from),
+        Some(examples_from),
+        kwargs.primary_key_columns.unwrap_or_default(),
+    );
+    Err(polars_err!(ComputeError: format!("\n{}", error.to_string(Some(&kwargs.schema_name)))))
 }