Merge pull request #18 from getyourguide/add-expectation-suite-runner

feat!: implement builder pattern for expectation suite runner

Author: ryanseq-gyg

README.md

@@ -21,6 +21,14 @@ standardized solution for this use case. As a result, any contributions made her
 pip install dataframe-expectations
 ```
 
+### Requirements
+
+* Python 3.10+
+* pandas >= 1.5.0
+* pydantic >= 2.12.4
+* pyspark >= 3.3.0
+* tabulate >= 0.8.9
+
 ### Development setup
 
 To set up the development environment:
@@ -45,67 +53,134 @@ uv run pytest tests/ --cov=dataframe_expectations
 
 ### Using the library
 
-**Pandas example:**
+**Basic usage with Pandas:**
 ```python
-from dataframe_expectations.expectations_suite import DataFameExpectationsSuite
+from dataframe_expectations.expectations_suite import DataFrameExpectationsSuite
+import pandas as pd
 
+# Build a suite with expectations
 suite = (
     DataFrameExpectationsSuite()
-    .expect_value_greater_than("age", 18)
-    .expect_value_less_than("age", 10)
+    .expect_min_rows(min_rows=3)
+    .expect_max_rows(max_rows=10)
+    .expect_value_greater_than(column_name="age", value=18)
+    .expect_value_less_than(column_name="salary", value=100000)
+    .expect_value_not_null(column_name="name")
 )
 
-# Create a Pandas DataFrame
-import pandas as pd
-test_pandas_df = pd.DataFrame({"age": [20, 15, 30], "name": ["Alice", "Bob", "Charlie"]})
-
-suite.run(test_pandas_df)
+# Create a runner
+runner = suite.build()
 
+# Validate a DataFrame
+df = pd.DataFrame({
+    "age": [25, 15, 45, 22],
+    "name": ["Alice", "Bob", "Charlie", "Diana"],
+    "salary": [50000, 60000, 80000, 45000]
+})
+runner.run(df)
 ```
 
-
 **PySpark example:**
 ```python
 from dataframe_expectations.expectations_suite import DataFrameExpectationsSuite
+from pyspark.sql import SparkSession
 
+# Initialize Spark session
+spark = SparkSession.builder.appName("example").getOrCreate()
+
+# Build a validation suite (same API as Pandas!)
 suite = (
     DataFrameExpectationsSuite()
-    .expect_value_greater_than("age", 18)
-    .expect_value_less_than("age", 40)
+    .expect_min_rows(min_rows=3)
+    .expect_max_rows(max_rows=10)
+    .expect_value_greater_than(column_name="age", value=18)
+    .expect_value_less_than(column_name="salary", value=100000)
+    .expect_value_not_null(column_name="name")
 )
 
+# Build the runner
+runner = suite.build()
+
 # Create a PySpark DataFrame
-test_spark_df = spark.createDataFrame(
-    [
-        {"name": "Alice", "age": 20},
-        {"name": "Bob", "age": 15},
-        {"name": "Charlie", "age": 30},
-    ]
-)
+data = [
+    {"age": 25, "name": "Alice", "salary": 50000},
+    {"age": 15, "name": "Bob", "salary": 60000},
+    {"age": 45, "name": "Charlie", "salary": 80000},
+    {"age": 22, "name": "Diana", "salary": 45000}
+]
+df = spark.createDataFrame(data)
+
+# Validate
+runner.run(df)
+```
+
+**Decorator pattern for automatic validation:**
+```python
+from dataframe_expectations.expectations_suite import DataFrameExpectationsSuite
+from pyspark.sql import SparkSession
 
-suite.run(test_spark_df)
+# Initialize Spark session
+spark = SparkSession.builder.appName("example").getOrCreate()
+
+suite = (
+    DataFrameExpectationsSuite()
+    .expect_min_rows(min_rows=3)
+    .expect_max_rows(max_rows=10)
+    .expect_value_greater_than(column_name="age", value=18)
+    .expect_value_less_than(column_name="salary", value=100000)
+    .expect_value_not_null(column_name="name")
+)
 
+# Build the runner
+runner = suite.build()
+
+# Apply decorator to automatically validate function output
+@runner.validate
+def load_employee_data():
+    """Load and return employee data - automatically validated."""
+    return spark.createDataFrame(
+        [
+            {"age": 25, "name": "Alice", "salary": 50000},
+            {"age": 15, "name": "Bob", "salary": 60000},
+            {"age": 45, "name": "Charlie", "salary": 80000},
+            {"age": 22, "name": "Diana", "salary": 45000}
+        ]
+    )
+
+# Function execution automatically validates the returned DataFrame
+df = load_employee_data()  # Raises DataFrameExpectationsSuiteFailure if validation fails
+
+# Allow functions that may return None
+@runner.validate(allow_none=True)
+def conditional_load(should_load: bool):
+    """Conditionally load data - validation only runs when DataFrame is returned."""
+    if should_load:
+        return spark.createDataFrame([{"age": 25, "name": "Alice", "salary": 50000}])
+    return None  # No validation when None is returned
 ```
 
 **Output:**
 ```python
 ========================== Running expectations suite ==========================
-ExpectationValueGreaterThan ('age' greater than 18) ... FAIL
-ExpectationValueLessThan ('age' less than 40) ... OK
-============================ 1 success, 1 failures =============================
+ExpectationMinRows (DataFrame contains at least 3 rows) ... OK
+ExpectationMaxRows (DataFrame contains at most 10 rows) ... OK
+ExpectationValueGreaterThan ('age' is greater than 18) ... FAIL
+ExpectationValueLessThan ('salary' is less than 100000) ... OK
+ExpectationValueNotNull ('name' is not null) ... OK
+============================ 4 success, 1 failures =============================
 
-ExpectationSuiteFailure: (1/2) expectations failed.
+ExpectationSuiteFailure: (1/5) expectations failed.
 
 ================================================================================
 List of violations:
 --------------------------------------------------------------------------------
-[Failed 1/1] ExpectationValueGreaterThan ('age' greater than 18): Found 1 row(s) where 'age' is not greater than 18.
+[Failed 1/1] ExpectationValueGreaterThan ('age' is greater than 18): Found 1 row(s) where 'age' is not greater than 18.
 Some examples of violations:
-+-----+------+
-| age | name |
-+-----+------+
-| 15  | Bob  |
-+-----+------+
++-----+------+--------+
+| age | name | salary |
++-----+------+--------+
+| 15  | Bob  | 60000  |
++-----+------+--------+
 ================================================================================
 
 ```