Validation Engine¶

def validate(
    df: pl.DataFrame | pl.LazyFrame,
    schema: SchemaModel,
    registry: FunctionRegistry,
    *,
    lazy: bool | None = None,
    coerce_strategy: Literal["strict", "null_on_failure"] = "strict",
    error_report: ErrorReportConfig | None = None,
) -> ValidationResult:
    """Validate a frame according to the schema and registry.

    Pipeline order (NEW):
    1. Column resolution (synonyms)
    2. Count original nulls
    3. Frame parsers
    4. Column parsers
    5. COERCE (moved before checks)
    6. Frame checks
    7. Column checks (with auto-injected non_null)
    8. Build error report (before nullification)
    9. Apply lenient behavior (nullify failures)
    10. Final nullable check (safety)
    11. Build validation report
    12. Return

    Args:
        df: Input DataFrame or LazyFrame.
        schema: Schema model.
        registry: Function registry holding parsers/checks.
        lazy: Optional override for lazy execution. Defaults to schema.lazy.
        coerce_strategy: How to handle coercion failures.
            - "strict": Raises on coercion failure.
            - "null_on_failure": Casts with strict=False (nulls on failure).
        error_report: Error reporting configuration. Defaults to summary mode.

    Returns:
        ValidationResult with validated data, errors DataFrame, and validation report.

    Examples:
        >>> # Summary mode (default) - just counts
        >>> result = validate(df, schema, registry)
        >>> result.errors
        shape: (2, 3)
        ┌──────────┬───────────┬───────┐
        │ column   │ check     │ count │
        │ ---      │ ---       │ ---   │
        │ str      │ str       │ u32   │
        ╞══════════╪═══════════╪═══════╡
        │ age      │ positive  │ 5     │
        │ name     │ non_null  │ 12    │
        └──────────┴───────────┴───────┘

        >>> # Rows mode - counts + indices
        >>> config = ErrorReportConfig(mode="rows", limit=100)
        >>> result = validate(df, schema, registry, error_report=config)
        >>> result.errors
        shape: (2, 3)
        ┌──────────┬───────────┬───────┬──────────────┐
        │ column   │ check     │ count │ row_indices  │
        │ ---      │ ---       │ ---   │ ---          │
        │ str      │ str       │ u32   │ list[u32]    │
        ╞══════════╪═══════════╪═══════╪══════════════╡
        │ age      │ positive  │ 5     │ [0, 3, 5,…]  │
        │ name     │ non_null  │ 12    │ [1, 2, 7,…]  │
        └──────────┴───────────┴───────┴──────────────┘

        >>> # Cells mode - individual rows with values
        >>> config = ErrorReportConfig(mode="cells", include_values=True, limit=10)
        >>> result = validate(df, schema, registry, error_report=config)
        >>> result.errors
        shape: (17, 4)
        ┌──────────┬───────────┬───────────┬────────┐
        │ column   │ check     │ row_index │ value  │
        │ ---      │ ---       │ ---       │ ---    │
        │ str      │ str       │ u32       │ i64    │
        ╞══════════╪═══════════╪═══════════╪════════╡
        │ age      │ positive  │ 0         │ -5     │
        │ age      │ positive  │ 3         │ 0      │
        │ name     │ non_null  │ 1         │ null   │
        │ …        │ …         │ …         │ …      │
        └──────────┴───────────┴───────────┴────────┘
    """
    if error_report is None:
        error_report = ErrorReportConfig(mode="summary")

    use_lazy = schema.lazy if lazy is None else lazy
    lf: pl.LazyFrame = df.lazy() if isinstance(df, pl.DataFrame) else df
    lf = lf.with_row_index("__row_index__")

    total_rows = int(lf.select(pl.len()).collect().item())

    # Phase 1: Column resolution
    lf = resolve_column_names(schema, lf)

    # Phase 2: Count original nulls (BEFORE any transformations)
    original_nulls = _count_original_nulls(lf, schema)

    # Phase 3: Frame parsers
    lf = _apply_frame_parsers(lf, schema, registry)

    # Phase 4: Column parsers (string transformations)
    lf = _apply_column_parsers(lf, schema, registry)

    # Phase 5: COERCE (MOVED BEFORE CHECKS)
    # Track coercion failures when using lenient strategy
    coercion_failures = {}
    if schema.coerce:
        if coerce_strategy == "null_on_failure":
            # Count nulls before coercion
            try:
                before_nulls = {
                    col: int(lf.select(pl.col(col).is_null().sum()).collect().item()) for col in schema.columns.keys()
                }
            except pl.exceptions.SchemaError as e:
                # Provide helpful error message for dtype mismatches
                current_dtypes = lf.collect_schema()
                schema_dtypes = {name: _resolve_dtype(col.dtype) for name, col in schema.columns.items()}

                error_details = []
                for col_name in schema.columns.keys():
                    current = current_dtypes.get(col_name)
                    expected = schema_dtypes.get(col_name)
                    if current and expected and current != expected:
                        error_details.append(f"  - {col_name}: current type is {current}, schema expects {expected}")

                raise ValueError(
                    "Schema dtype mismatch after parsers. This often happens when:\n"
                    "1. Column parsers (e.g., to_int, to_float) already convert the dtype\n"
                    "2. The schema's dtype doesn't match what the parsers produce\n\n"
                    "Mismatched columns:\n" + "\n".join(error_details) + "\n\n"
                    f"Solutions:\n"
                    f"  - Remove parsers and rely on coercion alone, OR\n"
                    f"  - Update schema dtypes to match parser outputs (Int64 -> i64, Float64 -> f64), OR\n"
                    f"  - Set coerce: false if parsers handle all type conversions\n\n"
                    f"Original error: {e}"
                ) from e

            lf = _apply_coercion(lf, schema, coerce_strategy)
            for col in schema.columns.keys():
                after = int(lf.select(pl.col(col).is_null().sum()).collect().item())
                coercion_failures[col] = after - before_nulls[col]
        else:
            try:
                lf = _apply_coercion(lf, schema, coerce_strategy)
            except pl.exceptions.SchemaError as e:
                # Provide helpful error message for dtype mismatches
                current_dtypes = lf.collect_schema()
                schema_dtypes = {name: _resolve_dtype(col.dtype) for name, col in schema.columns.items()}

                error_details = []
                for col_name in schema.columns.keys():
                    current = current_dtypes.get(col_name)
                    expected = schema_dtypes.get(col_name)
                    if current and expected and current != expected:
                        error_details.append(f"  - {col_name}: current type is {current}, schema expects {expected}")

                raise ValueError(
                    "Schema dtype mismatch after parsers. This often happens when:\n"
                    "1. Column parsers (e.g., to_int, to_float) already convert the dtype\n"
                    "2. The schema's dtype doesn't match what the parsers produce\n\n"
                    "Mismatched columns:\n" + "\n".join(error_details) + "\n\n"
                    f"Solutions:\n"
                    f"  - Remove parsers and rely on coercion alone, OR\n"
                    f"  - Update schema dtypes to match parser outputs (Int64 -> i64, Float64 -> f64), OR\n"
                    f"  - Set coerce: false if parsers handle all type conversions\n\n"
                    f"Original error: {e}"
                ) from e

    # Phase 6: Frame checks
    lf = _apply_frame_checks(lf, schema, registry)

    # Keep copy before checks for error value extraction
    lf_before_checks = lf

    # Phase 7: Column checks (with auto-injected non_null for nullable=False)
    lf, check_exprs = _collect_column_checks(lf, schema, registry)

    # Phase 8: Build error report BEFORE nullification (captures ALL failures)
    errors_df = _build_error_report(lf, check_exprs, error_report, lf_before_checks)

    # Phase 9: Apply lenient behavior (nullify where on_failure='null')
    lf, nullified_counts = _apply_lenient_checks(lf, check_exprs, schema)

    # Phase 10: Final nullable check (safety assertion)
    _check_final_nullable(lf, schema)

    # Phase 11: Count final nulls
    final_nulls = {col: int(lf.select(pl.col(col).is_null().sum()).collect().item()) for col in schema.columns.keys()}

    # Phase 12: Build validation report
    column_stats = {}
    for col_name in schema.columns.keys():
        # Count check failures from error report
        if not errors_df.is_empty():
            col_check_fails = errors_df.filter(pl.col("column") == col_name)
            check_fail_count = int(col_check_fails.select(pl.col("count").sum()).item() or 0)
        else:
            check_fail_count = 0

        column_stats[col_name] = ColumnValidationStats(
            column_name=col_name,
            parse_failures=0,  # Not tracked yet
            coercion_failures=coercion_failures.get(col_name, 0),
            check_failures=check_fail_count,
            nullified=nullified_counts.get(col_name, 0),
            final_null_count=final_nulls.get(col_name, 0),
            original_null_count=original_nulls.get(col_name, 0),
        )

    # Calculate valid rows
    if errors_df.is_empty():
        valid_rows = total_rows
    # Conservative: total - unique failing rows
    elif error_report.mode == "cells":
        failed_row_count = len(errors_df.select(pl.col("row_index").unique()).collect())
        valid_rows = total_rows - failed_row_count
    else:
        valid_rows = max(0, total_rows - int(errors_df.select(pl.col("count").sum()).item() or 0))

    validation_report = ValidationReport(
        rows_processed=total_rows,
        rows_valid=valid_rows,
        on_failure=schema.on_failure,
        columns=column_stats,
    )

    # Phase 13: Clean up and return
    lf = lf.drop("__row_index__")
    data_out: pl.DataFrame | pl.LazyFrame = lf if use_lazy else lf.collect()

    return ValidationResult(
        data=data_out,
        errors=errors_df,
        report=validation_report,
    )

def resolve_column_names(schema: SchemaModel, df: pl.DataFrame | pl.LazyFrame) -> pl.DataFrame | pl.LazyFrame:
    """Rename columns using canonical names and synonyms.

    Args:
        schema: SchemaModel defining columns and synonyms.
        df: Input frame.

    Returns:
        Frame with columns renamed to canonical names where possible.

    Raises:
        SchemaResolutionError: If required columns are missing or ambiguous.
    """
    columns = set(df.collect_schema().names() if isinstance(df, pl.LazyFrame) else df.columns)
    mapping: dict[str, str] = {}
    used: set[str] = set()

    for canonical, col_schema in schema.columns.items():
        candidates = {canonical} | set(col_schema.synonyms)
        found = [c for c in columns if c in candidates]
        if not found:
            if col_schema.required:
                raise SchemaResolutionError(
                    f"Required column '{canonical}' is missing (synonyms: {col_schema.synonyms})"
                )
            continue
        if len(found) > 1:
            raise SchemaResolutionError(
                f"Ambiguous columns for '{canonical}': {found}. Only one canonical/synonym is allowed."
            )
        physical = found[0]
        if physical in used:
            raise SchemaResolutionError(f"Column '{physical}' is mapped multiple times.")
        used.add(physical)
        if physical != canonical:
            mapping[physical] = canonical

    if not mapping:
        return df
    return df.rename(mapping)