feat: add suite_result and tagging

Author: ryanseq-gyg

.github/workflows/main.yaml

@@ -33,7 +33,7 @@ jobs:
           uv run python scripts/sanity_checks.py
       - name: Run tests
         run: |
-          uv run pytest tests/ --cov=dataframe_expectations
+          uv run pytest tests/ -n auto --tb=line --cov=dataframe_expectations
 
   lint:
     runs-on: ubuntu-latest

dataframe_expectations/__init__.py

@@ -9,4 +9,22 @@
     # Catch all exceptions to handle various edge cases in different environments
     __version__ = "0.0.0.dev0"
 
-__all__ = []
+from dataframe_expectations.core.suite_result import (
+    ExpectationResult,
+    SuiteExecutionResult,
+    serialize_violations,
+)
+from dataframe_expectations.suite import (
+    DataFrameExpectationsSuite,
+    DataFrameExpectationsSuiteRunner,
+    DataFrameExpectationsSuiteFailure,
+)
+
+__all__ = [
+    "ExpectationResult",
+    "SuiteExecutionResult",
+    "serialize_violations",
+    "DataFrameExpectationsSuite",
+    "DataFrameExpectationsSuiteRunner",
+    "DataFrameExpectationsSuiteFailure",
+]

dataframe_expectations/core/__init__.py

@@ -1,3 +1,15 @@
 """Core base classes and interfaces for DataFrame expectations."""
 
-__all__ = []
+from dataframe_expectations.core.suite_result import (
+    ExpectationResult,
+    ExpectationStatus,
+    SuiteExecutionResult,
+    serialize_violations,
+)
+
+__all__ = [
+    "ExpectationResult",
+    "ExpectationStatus",
+    "SuiteExecutionResult",
+    "serialize_violations",
+]

dataframe_expectations/core/aggregation_expectation.py

@@ -1,5 +1,5 @@
 from abc import abstractmethod
-from typing import List, Union
+from typing import List, Optional, Union
 
 from dataframe_expectations.core.types import DataFrameLike, DataFrameType
 from dataframe_expectations.core.expectation import DataFrameExpectation
@@ -20,6 +20,7 @@ def __init__(
         expectation_name: str,
         column_names: List[str],
         description: str,
+        tags: Optional[List[str]] = None,
     ):
         """
         Template for implementing DataFrame aggregation expectations, where data is first aggregated
@@ -28,7 +29,10 @@ def __init__(
         :param expectation_name: The name of the expectation. This will be used during logging.
         :param column_names: The list of column names to aggregate on.
         :param description: A description of the expectation used in logging.
+        :param tags: Optional tags as list of strings in "key:value" format.
+                    Example: ["priority:high", "env:test"]
         """
+        super().__init__(tags=tags)
         self.expectation_name = expectation_name
         self.column_names = column_names
         self.description = description

dataframe_expectations/core/column_expectation.py

@@ -1,4 +1,4 @@
-from typing import Callable
+from typing import Callable, List, Optional
 
 from dataframe_expectations.core.types import DataFrameLike, DataFrameType
 from dataframe_expectations.core.expectation import DataFrameExpectation
@@ -23,6 +23,7 @@ def __init__(
         fn_violations_pyspark: Callable,
         description: str,
         error_message: str,
+        tags: Optional[List[str]] = None,
     ):
         """
         Template for implementing DataFrame column expectations, where a column value is tested against a
@@ -34,7 +35,10 @@ def __init__(
         :param fn_violations_pyspark: Function to find violations in a PySpark DataFrame.
         :param description: A description of the expectation used in logging.
         :param error_message: The error message to return if the expectation fails.
+        :param tags: Optional tags as list of strings in "key:value" format.
+                    Example: ["priority:high", "env:test"]
         """
+        super().__init__(tags=tags)
         self.column_name = column_name
         self.expectation_name = expectation_name
         self.fn_violations_pandas = fn_violations_pandas

dataframe_expectations/core/expectation.py

@@ -1,5 +1,5 @@
 from abc import ABC, abstractmethod
-from typing import cast
+from typing import List, Optional, cast
 
 from pandas import DataFrame as PandasDataFrame
 from pyspark.sql import DataFrame as PySparkDataFrame
@@ -12,6 +12,7 @@
     PySparkConnectDataFrame = None  # type: ignore[misc,assignment]
 
 from dataframe_expectations.core.types import DataFrameLike, DataFrameType
+from dataframe_expectations.core.tagging import TagSet
 from dataframe_expectations.result_message import (
     DataFrameExpectationResultMessage,
 )
@@ -22,6 +23,14 @@ class DataFrameExpectation(ABC):
     Base class for DataFrame expectations.
     """
 
+    def __init__(self, tags: Optional[List[str]] = None):
+        """
+        Initialize the base expectation with optional tags.
+        :param tags: Optional tags as list of strings in "key:value" format.
+                    Example: ["priority:high", "env:test"]
+        """
+        self.tags = TagSet(tags)
+
     def get_expectation_name(self) -> str:
         """
         Returns the class name as the expectation name.
@@ -48,29 +57,31 @@ def infer_data_frame_type(cls, data_frame: DataFrameLike) -> DataFrameType:
         """
         Infer the DataFrame type based on the provided DataFrame.
         """
-        if isinstance(data_frame, PandasDataFrame):
-            return DataFrameType.PANDAS
-        elif isinstance(data_frame, PySparkDataFrame):
-            return DataFrameType.PYSPARK
-        elif PySparkConnectDataFrame is not None and isinstance(
-            data_frame, PySparkConnectDataFrame
-        ):
-            return DataFrameType.PYSPARK
-        else:
-            raise ValueError(f"Unsupported DataFrame type: {type(data_frame)}")
+        match data_frame:
+            case PandasDataFrame():
+                return DataFrameType.PANDAS
+            case PySparkDataFrame():
+                return DataFrameType.PYSPARK
+            case _ if PySparkConnectDataFrame is not None and isinstance(
+                data_frame, PySparkConnectDataFrame
+            ):
+                return DataFrameType.PYSPARK
+            case _:
+                raise ValueError(f"Unsupported DataFrame type: {type(data_frame)}")
 
     def validate(self, data_frame: DataFrameLike, **kwargs):
         """
         Validate the DataFrame against the expectation.
         """
         data_frame_type = self.infer_data_frame_type(data_frame)
 
-        if data_frame_type == DataFrameType.PANDAS:
-            return self.validate_pandas(data_frame=data_frame, **kwargs)
-        elif data_frame_type == DataFrameType.PYSPARK:
-            return self.validate_pyspark(data_frame=data_frame, **kwargs)
-        else:
-            raise ValueError(f"Unsupported DataFrame type: {data_frame_type}")
+        match data_frame_type:
+            case DataFrameType.PANDAS:
+                return self.validate_pandas(data_frame=data_frame, **kwargs)
+            case DataFrameType.PYSPARK:
+                return self.validate_pyspark(data_frame=data_frame, **kwargs)
+            case _:
+                raise ValueError(f"Unsupported DataFrame type: {data_frame_type}")
 
     @abstractmethod
     def validate_pandas(

dataframe_expectations/core/suite_result.py

@@ -0,0 +1,174 @@
+"""Suite execution result models for capturing validation outcomes."""
+
+from datetime import datetime
+from typing import Any, Dict, List, Literal, Optional
+
+from pydantic import BaseModel, Field, computed_field
+
+from dataframe_expectations.core.types import DataFrameType, DataFrameLike
+from dataframe_expectations.core.tagging import TagSet
+
+
+from enum import Enum
+
+
+class ExpectationStatus(str, Enum):
+    PASSED = "passed"
+    FAILED = "failed"
+    SKIPPED = "skipped"
+
+
+class ExpectationResult(BaseModel):
+    """
+    Representation of a single expectation result within a suite execution.
+    Captures the outcome (passed, failed, skipped) using status.
+    Does not store raw dataframes, only serialized violation samples.
+    """
+
+    expectation_name: str = Field(..., description="Name of the expectation class")
+    description: str = Field(..., description="Human-readable description of the expectation")
+    status: ExpectationStatus = Field(..., description="Outcome status: passed, failed, or skipped")
+    tags: Optional[TagSet] = Field(
+        default=None, description="User-defined tags for this specific expectation"
+    )
+    error_message: Optional[str] = Field(
+        default=None, description="Error message if expectation failed"
+    )
+    violation_count: Optional[int] = Field(
+        default=None, description="Total count of violations (if applicable)"
+    )
+    violation_sample: Optional[List[Dict[str, Any]]] = Field(
+        default=None,
+        description="Sample of violations as list of dicts (limited by violation_sample_limit)",
+    )
+
+    model_config = {"frozen": True, "arbitrary_types_allowed": True}  # Make immutable, allow TagSet
+
+
+class SuiteExecutionResult(BaseModel):
+    """Result of a complete suite execution.
+    Captures all metadata about the suite run including timing, dataframe info,
+    and individual expectation results. Does not store raw dataframes.
+    """
+
+    suite_name: Optional[str] = Field(default=None, description="Optional name for the suite")
+    context: Dict[str, Any] = Field(
+        default_factory=dict, description="Additional runtime metadata (e.g., job_id, environment)"
+    )
+    applied_filters: TagSet = Field(
+        default_factory=TagSet, description="Tag filters that were applied to select expectations"
+    )
+    tag_match_mode: Optional[Literal["any", "all"]] = Field(
+        default=None, description="How tags were matched: 'any' (OR) or 'all' (AND)"
+    )
+    results: List[ExpectationResult] = Field(
+        ..., description="Results for each expectation in execution order (including skipped)"
+    )
+    start_time: datetime = Field(..., description="Suite execution start timestamp")
+    end_time: datetime = Field(..., description="Suite execution end timestamp")
+    dataframe_type: DataFrameType = Field(..., description="Type of dataframe validated")
+    dataframe_row_count: int = Field(..., description="Number of rows in validated dataframe")
+    dataframe_was_cached: bool = Field(
+        default=False, description="Whether PySpark dataframe was cached during execution"
+    )
+
+    model_config = {"frozen": True, "arbitrary_types_allowed": True}  # Make immutable, allow TagSet
+
+    @computed_field  # type: ignore[misc]
+    @property
+    def total_duration_seconds(self) -> float:
+        """Total execution time in seconds."""
+        return (self.end_time - self.start_time).total_seconds()
+
+    @computed_field  # type: ignore[misc]
+    @property
+    def total_expectations(self) -> int:
+        """Total number of expectations in the suite (including skipped)."""
+        return len(self.results)
+
+    @computed_field  # type: ignore[misc]
+    @property
+    def total_passed(self) -> int:
+        """Number of expectations that passed."""
+        return sum(1 for r in self.results if r.status == ExpectationStatus.PASSED)
+
+    @computed_field  # type: ignore[misc]
+    @property
+    def total_failed(self) -> int:
+        """Number of expectations that failed."""
+        return sum(1 for r in self.results if r.status == ExpectationStatus.FAILED)
+
+    @computed_field  # type: ignore[misc]
+    @property
+    def total_skipped(self) -> int:
+        """Number of expectations that were skipped due to tag filtering."""
+        return sum(1 for r in self.results if r.status == ExpectationStatus.SKIPPED)
+
+    @computed_field  # type: ignore[misc]
+    @property
+    def pass_rate(self) -> float:
+        """Percentage of expectations that passed (0.0 to 1.0)."""
+        executed = self.total_passed + self.total_failed
+        if executed == 0:
+            return 1.0
+        return self.total_passed / executed
+
+    @computed_field  # type: ignore[misc]
+    @property
+    def success(self) -> bool:
+        """Whether all executed expectations passed (ignores skipped)."""
+        return self.total_failed == 0
+
+    @computed_field  # type: ignore[misc]
+    @property
+    def passed_expectations(self) -> List[ExpectationResult]:
+        """List of expectations that passed."""
+        return [r for r in self.results if r.status == ExpectationStatus.PASSED]
+
+    @computed_field  # type: ignore[misc]
+    @property
+    def failed_expectations(self) -> List[ExpectationResult]:
+        """List of expectations that failed."""
+        return [r for r in self.results if r.status == ExpectationStatus.FAILED]
+
+    @computed_field  # type: ignore[misc]
+    @property
+    def skipped_expectations(self) -> List[ExpectationResult]:
+        """List of expectations that were skipped due to tag filtering."""
+        return [r for r in self.results if r.status == ExpectationStatus.SKIPPED]
+
+
+def serialize_violations(
+    violations_df: Optional[DataFrameLike],
+    df_type: DataFrameType,
+    limit: int = 5,
+) -> tuple[Optional[int], Optional[List[Dict[str, Any]]]]:
+    """Serialize violation dataframe to count and sample for storage.
+
+    Converts dataframes to JSON-serializable format without storing raw dataframes.
+
+    :param violations_df: DataFrame containing violations (pandas or PySpark).
+    :param df_type: Type of the violations dataframe.
+    :param limit: Maximum number of violation rows to include in sample.
+    :return: Tuple of (total_count, sample_as_list_of_dicts).
+    """
+    if violations_df is None:
+        return None, None
+
+    count: Optional[int] = None
+    sample: Optional[list[dict[str, Any]]] = None
+
+    try:
+        if df_type == DataFrameType.PANDAS:
+            pandas_df = violations_df  # type: ignore[assignment]
+            count = len(pandas_df)  # type: ignore[arg-type]
+            sample = pandas_df.head(limit).to_dict("records")  # type: ignore[assignment,union-attr]
+        elif df_type == DataFrameType.PYSPARK:
+            pyspark_df = violations_df  # type: ignore[assignment]
+            count = pyspark_df.count()  # type: ignore[assignment]
+            sample = pyspark_df.limit(limit).toPandas().to_dict("records")  # type: ignore[assignment,operator]
+
+        return count, sample
+    except Exception:
+        # If serialization fails, return None to avoid breaking the suite
+        return None, None