pyspark-dq-framework 0.1.3

A PySpark-based Data Quality Framework using YAML-configurable checks.

🛡️ PySpark Data Quality Framework (pyspark_dq_framework)

This framework provides a config-driven, extensible Data Quality (DQ) framework built on top of PySpark.
It allows defining validation checks in a config.yml file and running them on one or more DataFrames.

---

📁 Project Structure

pyspark_dq_framework/
├── pyspark_dq/                   # your Python package
│   ├── __init__.py
│   ├── dq_checks.py              # Core data quality functions
│   ├── config.py                 # YAML parser to config DataFrame
│   ├── run_dq.py                 # Engine to run checks with logging
│   └── validate_checks.py        # Check validation
├── README.md
├── LICENSE
├── MANIFEST.in
├── setup.py

---

⚙️ Configuration (config.yml)

return_mode: summary   # options: summary (default) | summary+clean | summary+error | all

model:
  - name: df
    columns:
      - name: ID
        checks:
          - null_check
      - name: Age
        checks:
          - null_check
          - positive_age_check:
              min_age: 21

  - name: df2
    columns:
      - name: ID
        checks:
          - null_check
      - name: Salary
        checks:
          - non_negative_check

---

🛠️ Code Example for Executing framework

If config file is in :

• local: pass file path in config_path parameter
• S3: pass s3 path
• Azure Blob Storage or (ADLS Gen2 for data lake): pass blob storage path
• Google Cloud Storage (GCS): pass GCS path

Code Example:

from pyspark.sql import SparkSession
from pyspark.sql.functions import col, lit
from pyspark_dq.run_dq import run_dq_pipeline
from pyspark_dq.dq_checks import *

# Custom check
def positive_age_check(df, row, min_age=0):
    table = row["table"]
    column = row["column"]
    check_type = "positive_age_check"

    valid_df = df.filter(col(column) >= min_age)
    invalid_df = df.filter(col(column) < min_age) \
                   .withColumn("table", lit(table)) \
                   .withColumn("column", lit(column)) \
                   .withColumn("check_type", lit(check_type))

    log_df = df.sparkSession.createDataFrame(
        [(table, column, check_type, invalid_df.count() == 0, invalid_df.count())],
        ["table", "column", "check_type", "passed", "invalid_count"]
    )
    return valid_df, invalid_df, log_df

if __name__ == "__main__":
    
    spark = SparkSession.builder.appName("pyspark_dq_framework").getOrCreate()

    # Example Data
    data1 = [(1, "Ram", 20), (None, "Shyam", 20)]
    df = spark.createDataFrame(data1, ["ID", "Name", "Age"])

    data2 = [
        (1, "Ram", "Software Engineer", 20000),
        (2, "Shyam", "Product Manager", 80000),
        (3, "Radha", "Senior Software Engineer", 50000)
    ]
    df2 = spark.createDataFrame(data2, ["ID", "Name", "Designation", "Salary"])

    df_dict = {"df": df, "df2": df2}

    # Register checks
    check_function_map = {
        "null_check": null_check,
        "unique_check": unique_check,
        "allowed_values_check": allowed_values_check,
        "range_check": range_check,
        "non_negative_check": non_negative_check,
        "regex_check": regex_check,
        "not_empty_check": not_empty_check,
        "data_type_check": data_type_check,
        "positive_age_check": positive_age_check
    }

    # Run pipeline
    results = run_dq_pipeline(
        spark,
        df_dict=df_dict,
        config_path="config.yml",
        checks=check_function_map
    )

    print(results)

---

🔄 Return Modes

The return_mode in config.yml controls what is returned:

• summary → Only summary logs
• summary+clean → Summary + cleaned DataFrames
• summary+error → Summary + error DataFrames
• all → Summary + clean + error DataFrames

---

▶️ How to Run

📦 Install PyPi package in your environment

pip install pyspark-dq-framework

📚 Import required libraries

from pyspark.sql import SparkSession
from pyspark_dq.dq_checks import *
from pyspark_dq.run_dq import run_dq_pipeline

💻 Add sample code without custom check (here I am taking file name as main.py)

/main.py

from pyspark.sql import SparkSession
from pyspark_dq.run_dq import run_dq_pipeline
from pyspark_dq.dq_checks import null_check, unique_check
    
spark = SparkSession.builder.appName("pyspark_dq_framework").getOrCreate()

# Example Data
# Dataframe 1
data1 = [
          (1, "Ram", 20),
          (None, "Shyam", 20)
        ]
df1 = spark.createDataFrame(data1, ["ID", "Name", "Age"])

# Datframe 2
data2 = [
          (1, "Ram", "Software Engineer", 20000),
          (2, "Shyam", "Product Manager", 80000),
          (3, "Radha", "Senior Software Engineer", 50000)
        ]
df2 = spark.createDataFrame(data2, ["ID", "Name", "Designation", "Salary"])

# Dataframes Dictionary
df_dict = {"df1": df1, "df2": df2}

# Register checks
check_function_map = {
        "null_check": null_check,
        "unique_check": unique_check
    }

# Run pipeline
results = run_dq_pipeline(
            spark,
            df_dict=df_dict,
            config_path="config.yml",
            checks=check_function_map
          )

print(results)

⚡ Add sample code with custom check (here I am taking file name as main.py)

/main.py

from pyspark.sql import SparkSession
from pyspark_dq.run_dq import run_dq_pipeline
from pyspark_dq.dq_checks import null_check, unique_check


# Custom check
def positive_age_check(df, row, min_age=0):
    table = row["table"]
    column = row["column"]
    check_type = "positive_age_check"

    valid_df = df.filter(col(column) >= min_age)
    invalid_df = df.filter(col(column) < min_age) \
                   .withColumn("table", lit(table)) \
                   .withColumn("column", lit(column)) \
                   .withColumn("check_type", lit(check_type))

    log_df = df.sparkSession.createDataFrame(
        [(table, column, check_type, invalid_df.count() == 0, invalid_df.count())],
        ["table", "column", "check_type", "passed", "invalid_count"]
    )
    return valid_df, invalid_df, log_df

# Inititalize spark session 
spark = SparkSession.builder.appName("pyspark_dq_framework").getOrCreate()

# Example Data
# Dataframe 1
data1 = [
          (1, "Ram", 20),
          (None, "Shyam", 20)
        ]
df1 = spark.createDataFrame(data1, ["ID", "Name", "Age"])

# Datframe 2
data2 = [
          (1, "Ram", "Software Engineer", 20000),
          (2, "Shyam", "Product Manager", 80000),
          (3, "Radha", "Senior Software Engineer", 50000)
        ]
df2 = spark.createDataFrame(data2, ["ID", "Name", "Designation", "Salary"])

# Dataframes Dictionary
df_dict = {"df1": df1, "df2": df2}

# Register checks
check_function_map = {
        "null_check": null_check,
        "unique_check": unique_check,
        "positive_age_check": positive_age_check
    }

# Run pipeline
results = run_dq_pipeline(
            spark,
            df_dict=df_dict,
            config_path="config.yml",
            checks=check_function_map
          )

print(results)

📝 General Config YAML Structure

/config.yml

return_mode: <return mode>   # options: summary (default) | summary+clean | summary+error | all

model:
  - name: <dataframe_name>
    columns:
      - name: <column_name>
        checks:
          - <check_name>
          - <check_name_with_params>:
              param1: value
              param2: value

---

💻 Run locally

python main.py

📊 Run in Databricks

dbutils.notebook.run("main", 60)

☁️ Run with S3 config

python main.py --config s3://my-bucket/config.yml

---

✅ Features

• Config-driven DQ checks
• Works with local, S3, GCP, Azure, DBFS
• Extensible with custom checks
• Flexible return_mode for different use cases

📋 Supported Data Quality Checks

• null_check
• unique_check
• allowed_values_check
• range_check
• non_negative_check
• regex_check
• not_empty_check
• data_type_check

---

Check Name	Description	Parameters
null_check	Ensures no null values.	None
unique_check	Ensures uniqueness of values.	None
allowed_values_check	Ensures value is in allowed set.	values: [list_of_values]
range_check	Ensures value is within numeric range.	min_value, max_value (optional)
non_negative_check	Ensures value >= 0.	None
regex_check	Ensures value matches regex pattern.	pattern: "<regex>"
not_empty_check	Ensures dataset/column is not empty.	None
data_type_check	Ensures value matches expected data type.	expected_type
positive_age_check	Ensures age is >= min_age.	min_age

---

📑 YAML Usage Examples for Each Check

Check Name	YAML Example
null_check	- null_check
unique_check	- unique_check
allowed_values_check	- allowed_values_check:\n values: [Male, Female, Other]
range_check	- range_check:\n min_value: 10\n max_value: 100
non_negative_check	- non_negative_check
regex_check	- regex_check:\n pattern: "^[A-Za-z0-9_]+$"
not_empty_check	- not_empty_check
data_type_check	- data_type_check:\n expected_type: integer
positive_age_check	- positive_age_check:\n min_age: 21

---

📑 Explaination of each check with proper YAML format for config.yml file

null_check

• Ensures that the column does not contain NULL values.
• Records with NULL values will go into the invalid dataset.

model:
  - name: df
    columns:
      - name: ID
        checks:
          - null_check

unique_check

• Ensures that the column contains unique values.
• Duplicates will be flagged as invalid.

model:
  - name: df
    columns:
      - name: ID
        checks:
          - unique_check

allowed_values_check

• Ensures that the column values belong to a defined set of allowed values.
• Invalid values are flagged.

Parameters

• values → list of allowed values.

model:
  - name: df
    columns:
      - name: Status
        checks:
          - allowed_values_check:
              values: ["Active", "Inactive", "Pending"]

range_check

• Ensures that the column values fall within a numeric range.

Parameters

• min_value → minimum acceptable value (optional).
• max_value → maximum acceptable value (optional).

model:
  - name: df
    columns:
      - name: Age
        checks:
          - range_check:
              min_value: 18
              max_value: 60

non_negative_check

• Ensures that the column values are greater than or equal to 0.

model:
  - name: df
    columns:
      - name: Salary
        checks:
          - non_negative_check

regex_check

• Ensures that the column values match a given regex pattern.
• Useful for emails, phone numbers, codes, etc.

Parameters

• pattern → valid regex pattern.

model:
  - name: df
    columns:
      - name: Email
        checks:
          - regex_check:
              pattern: "^[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\\.[A-Z|a-z]{2,}$"

not_empty_check

• Ensures that the column (or dataset) is not empty.
• If no records exist, the check fails.

model:
  - name: df
    columns:
      - name: ID
        checks:
          - not_empty_check

data_type_check

• Ensures that the column matches the expected data type.
• You need to provide the expected type in parameters

Parameters

• expected_type → e.g., IntegerType, StringType, DoubleType.

model:
  - name: df
    columns:
      - name: Age
        checks:
          - data_type_check:
              expected_type: IntegerType