spark-data-quality 1.0.0__py3-none-any.whl

spark_data_quality-1.0.0.dist-info/METADATA

@@ -0,0 +1,123 @@
+Metadata-Version: 2.4
+Name: spark-data-quality
+Version: 1.0.0
+Summary: SparkDQAgent — Data Quality validation package for K8s Spark pods
+Author-email: khailas <khailas.rangath@saal.ai>
+License-Expression: MIT
+Classifier: Programming Language :: Python :: 3
+Classifier: Operating System :: OS Independent
+Requires-Python: >=3.8
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: requests>=2.28.0
+Requires-Dist: great-expectations==0.18.12
+Requires-Dist: trino>=0.320.0
+Requires-Dist: PyYAML>=6.0
+Provides-Extra: spark
+Requires-Dist: pyspark>=3.1.1; extra == "spark"
+Provides-Extra: dev
+Requires-Dist: pytest; extra == "dev"
+Requires-Dist: ruff; extra == "dev"
+Dynamic: license-file
+
+# spark-dq
+
+Data Quality validation library for K8s Spark pods. Runs Great Expectations + Trino hybrid checks and persists results to the DQ Engine API.
+
+## Installation
+
+```bash
+pip install .            # core (GE + Trino)
+pip install ".[spark]"   # includes PySpark
+pip install -e ".[dev]"  # editable + dev tools
+```
+
+## Package structure
+
+```
+src/spark_dq/
+├── quality.py              # SparkDQAgent — main orchestrator
+├── utils.py                # JSON serialisation helpers
+├── helpers/
+│   ├── api_client.py       # DQ Engine HTTP client (GET /config, POST /save-results)
+│   ├── expectations.py     # Expectation splitting (fast/slow) and result builders
+│   ├── trino_executor.py   # Trino fast-path — batched SQL for 5 assertion types
+│   ├── ge_executor.py      # GE slow-path — parallel in-memory validators
+│   ├── sql_utils.py        # SQL injection prevention utilities
+│   └── exceptions.py       # DQEngineAPIError
+└── data_quality_helper/
+    ├── const.py             # GE datasource config templates
+    ├── datacontext.py       # In-memory GE DataContext builder
+    ├── batch_request.py     # RuntimeBatchRequest factory
+    ├── expectation_suite.py # Suite builder from assertion dicts
+    ├── ge_utils.py          # Mapping ID resolution
+    └── checkpoint_conf.py   # Checkpoint runner (legacy)
+```
+
+## Usage
+
+### Single table
+
+```python
+from spark_dq.quality import SparkDQAgent
+
+agent = SparkDQAgent(
+    catalog="my_catalog",
+    schema="public",
+    table="sales_data",
+    data_quality_url="http://dq-engine:8000/api/v1/spark",
+    catalog_type="unmanaged",
+    table_id=13,
+    test_suite_id=37,
+    trino_host="trino-host:443",
+    trino_user="user",
+    trino_pwd="pwd",
+)
+
+cfg = agent.get_config()        # single GET /config call
+results = agent.validate(df)    # hybrid Trino + GE execution
+```
+
+### Suite run (auto-discover tables)
+
+```python
+agent = SparkDQAgent(
+    catalog="", schema="", table="",
+    data_quality_url="http://dq-engine:8000/api/v1/spark",
+    catalog_type="unmanaged",
+    test_suite_id=37,
+)
+table_configs = agent.get_all_table_configs()  # returns list of per-table configs
+```
+
+## Demo
+
+```bash
+# Single table
+python demo.py \
+  --catalog khtestpostgresss --schema public --table sales_data_new \
+  --dq_url http://localhost:8000/api/v1/spark \
+  --trino_host trino-dev.digixt.ae:443 --trino_user datalake --trino_pwd <pwd> \
+  --table_id 13 --test_suite_id 37
+
+# Suite run (auto-discovers tables)
+python demo.py \
+  --test_suite_id 37 \
+  --dq_url http://localhost:8000/api/v1/spark \
+  --trino_host trino-dev.digixt.ae:443 --trino_user datalake --trino_pwd <pwd>
+```
+
+## Hybrid execution
+
+When Trino credentials are provided, the agent splits expectations into two paths:
+
+| Path | Assertion types | How |
+|------|----------------|-----|
+| **Trino fast-path** | `not_be_null`, `be_unique`, `be_between`, `match_regex`, `row_count_to_be_between` | Single batched SQL query |
+| **GE slow-path** | All others | Parallel in-memory GE validators |
+
+Both paths run **concurrently**. Results are merged and saved via `POST /save-results`.
+
+## License
+
+MIT

spark_data_quality-1.0.0.dist-info/RECORD

@@ -0,0 +1,23 @@
+spark_data_quality-1.0.0.dist-info/licenses/LICENSE,sha256=ESYyLizI0WWtxMeS7rGVcX3ivMezm-HOd5WdeOh-9oU,1056
+spark_dq/__init__.py,sha256=udiZe4rI0Ttfpxo4KwiE-sGjTxEtO20Pa3YxzsErLVY,98
+spark_dq/quality.py,sha256=EQhg8krc6gVznGYHCa_UPcIPt8efET7ja1JwsZ5gFwY,10251
+spark_dq/utils.py,sha256=s_Swkufse9Tu08f5b0EzwKog0NbWFLbV33BamFuUiYQ,1097
+spark_dq/data_quality_helper/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+spark_dq/data_quality_helper/batch_request.py,sha256=TLnlxx7CqfH0THCCjS9O07JEgHIqWpax_dLjASOom4Y,840
+spark_dq/data_quality_helper/checkpoint_conf.py,sha256=RcTHntduI92nngS7VC7ZvTjMEvoqbeKSHK0HC2IsY8c,1213
+spark_dq/data_quality_helper/const.py,sha256=HzofT1dA8xZetqswnHzS6nOKh-cCJCa9iuRtQfiFPsM,749
+spark_dq/data_quality_helper/datacontext.py,sha256=iUwXomXjE8dh5JrZseZG6wuzDlsn7xnQpFa_z136KAA,1497
+spark_dq/data_quality_helper/expectation_suite.py,sha256=B9Jla-Jh3R-YLrLnygQxGv4e8OThzLiCljkzcEuTb_s,688
+spark_dq/data_quality_helper/ge_utils.py,sha256=w0_JMg-Gj-3StjbEasAoglIFKgU16y3tNcAn1BdmxGI,943
+spark_dq/data_quality_helper/utils.py,sha256=9LSquN7Xt0x40WNkB2dBCBub1o6Kj4M8CEfoi45pVEg,1326
+spark_dq/helpers/__init__.py,sha256=YFotyjTGDiIMFuhcwKThHdjnZ9DzA9E9uSWkirb3Ve0,660
+spark_dq/helpers/api_client.py,sha256=Shh_WeNFVd_YpeyIvRx_NvcCyuvTyEcZayFLrNdidqw,3725
+spark_dq/helpers/exceptions.py,sha256=hfhp0tuV9oFYIIukm1b81Sx67LS3imlTs8HaM8NWK9c,277
+spark_dq/helpers/expectations.py,sha256=_-L-Mfo6NwWCpI_I1HEIn2qAvwnLx5jor3REECZUdfQ,1414
+spark_dq/helpers/ge_executor.py,sha256=dudNWzRDyvHR9plBa7QXu5dlF9trFQzhjN6xuG8pDpg,5276
+spark_dq/helpers/sql_utils.py,sha256=vP62kDPYNCULsETIBoaPV_huhiWoD9k8LQc7fXN30kI,833
+spark_dq/helpers/trino_executor.py,sha256=BiBcwE9ISM7ykeamgt7kmo04ldBWUT55NDp5B6gvO_I,5697
+spark_data_quality-1.0.0.dist-info/METADATA,sha256=ewM868gPLUHWbCHsepprSx9L8rHwmdyzm6CaD6tQ81g,3937
+spark_data_quality-1.0.0.dist-info/WHEEL,sha256=K260EYznzXsJYBQGqmI8VTxEdiZYNvDZwW9cBh9-_MA,91
+spark_data_quality-1.0.0.dist-info/top_level.txt,sha256=Js6y25e8v--TS1TsPtdTCw02kmmPV0VlgBluNy0q-D0,9
+spark_data_quality-1.0.0.dist-info/RECORD,,

spark_data_quality-1.0.0.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (83.0.0)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

spark_data_quality-1.0.0.dist-info/licenses/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

spark_data_quality-1.0.0.dist-info/top_level.txt

@@ -0,0 +1 @@
+spark_dq

spark_dq/__init__.py

@@ -0,0 +1,3 @@
+"""SparkDQAgent — Data Quality validation package for K8s Spark pods."""
+
+__version__ = "1.0.0"

spark_dq/data_quality_helper/batch_request.py

@@ -0,0 +1,23 @@
+import logging
+from great_expectations.core.batch import RuntimeBatchRequest
+
+logger = logging.getLogger(__name__)
+
+
+def create_batchrequest(df, datasource_name: str, data_asset_name: str):
+    """
+    Wrap a pre-loaded Spark DataFrame in a GE RuntimeBatchRequest.
+
+    The DataFrame is already loaded by ``quality_spark.py`` (either via
+    ``spark.sql`` for Iceberg or Trino JDBC for external catalogs).
+    No data fetching happens here.
+    """
+    batch = RuntimeBatchRequest(
+        datasource_name=datasource_name,
+        data_connector_name="default_runtime_data_connector_name",
+        data_asset_name=data_asset_name,
+        runtime_parameters={"batch_data": df},
+        batch_identifiers={"default_identifier_name": "default_identifier"},
+    )
+    logger.info("Batch request created from Spark DataFrame")
+    return batch

spark_dq/data_quality_helper/checkpoint_conf.py

@@ -0,0 +1,30 @@
+from datetime import datetime
+
+from spark_dq.data_quality_helper.const import ACTION_LIST, DATETIME_FORMAT
+
+# BOOLEAN_ONLY: GE only computes pass/fail per expectation — no unexpected
+# value collection, no data movement back to the driver. This eliminates the
+# hidden per-expectation Spark action that makes validation slow at scale.
+RESULT_FORMAT = {
+    "result_format": "BOOLEAN_ONLY",
+    "include_unexpected_rows": False,
+    "partial_unexpected_count": 0,
+}
+
+
+def run_checkpoint(context, data_asset_name, expectation_suite_name, batch_request):
+    checkpoint = context.add_or_update_checkpoint(
+        name=f"{data_asset_name}.{expectation_suite_name}.chk",
+        module_name="great_expectations.checkpoint",
+        class_name="Checkpoint",
+        run_name_template=f'{datetime.now().strftime(DATETIME_FORMAT)}-my-run-name-template',
+    )
+
+    checkpoint_run_result = checkpoint.run(
+        run_name=f'{datetime.now().strftime(DATETIME_FORMAT)}-my-run-name-template',
+        batch_request=batch_request,
+        expectation_suite_name=expectation_suite_name,
+        action_list=ACTION_LIST,
+        runtime_configuration={"result_format": RESULT_FORMAT},
+    )
+    return checkpoint_run_result

spark_dq/data_quality_helper/const.py

@@ -0,0 +1,27 @@
+DATETIME_FORMAT = "%Y-%m-%d %H:%M:%S"
+
+DATASOURCE_CONFIG = """
+    name: {datasource_name}
+    class_name: Datasource
+    module_name: great_expectations.datasource
+    execution_engine:
+        class_name: SparkDFExecutionEngine
+        module_name: great_expectations.execution_engine
+        
+    data_connectors:
+        default_runtime_data_connector_name:
+            class_name: RuntimeDataConnector
+            batch_identifiers:
+            - default_identifier_name
+    """
+
+ACTION_LIST = [
+    {
+        "name": "store_validation_result",
+        "action": {"class_name": "StoreValidationResultAction"},
+    },
+    {
+        "name": "store_evaluation_params",
+        "action": {"class_name": "StoreEvaluationParametersAction"},
+    },
+]

spark_dq/data_quality_helper/datacontext.py

@@ -0,0 +1,39 @@
+from great_expectations.data_context.types.base import DataContextConfig
+
+
+def create_datacontext():
+    """Build a GE DataContextConfig backed entirely by in-memory stores."""
+    return DataContextConfig(
+        **{
+            "config_version": 1.0,
+            "stores": {
+                "expectations_store": {
+                    "class_name": "ExpectationsStore",
+                    "store_backend": {
+                        "class_name": "InMemoryStoreBackend",
+                    },
+                },
+                "validations_store": {
+                    "class_name": "ValidationsStore",
+                    "store_backend": {
+                        "class_name": "InMemoryStoreBackend",
+                    },
+                },
+                "evaluation_parameter_store": {
+                    "class_name": "EvaluationParameterStore",
+                },
+                "checkpoint_store": {
+                    "class_name": "CheckpointStore",
+                    "store_backend": {
+                        "class_name": "InMemoryStoreBackend",
+                    },
+                },
+            },
+            "expectations_store_name": "expectations_store",
+            "validations_store_name": "validations_store",
+            "evaluation_parameter_store_name": "evaluation_parameter_store",
+            "checkpoint_store_name": "checkpoint_store",
+            "data_docs_sites": {},
+            "anonymous_usage_statistics": {"enabled": False},
+        }
+    )

spark_dq/data_quality_helper/expectation_suite.py

@@ -0,0 +1,17 @@
+from great_expectations.core.expectation_configuration import ExpectationConfiguration
+
+
+def create_expectation_function(expectation_list, suite):
+    # print(f"###inside create_expectation_function,  expectation_list:{expectation_list}")
+    for expectation in expectation_list:
+        expectation = expectation["assertion_info"]
+
+        expectation = ExpectationConfiguration(
+            expectation_type=list(expectation.keys())[0],
+            kwargs=list(expectation.values())[0],
+        )
+        # while adding, GE will overwrite the assertion arguments, if the assertion already exists .
+
+        suite.add_expectation(expectation_configuration=expectation)
+
+    return suite

spark_dq/data_quality_helper/ge_utils.py

@@ -0,0 +1,20 @@
+def get_mapping_id(actions_result: dict, expectations_list: list):
+    """Match a GE validation result back to its assertion_mapping_id."""
+    expectation_type_action = actions_result["expectation_config"]["expectation_type"]
+    expectation_type_kwargs = dict(actions_result["expectation_config"]["kwargs"])
+    expectation_type_kwargs.pop("batch_id", None)
+    expectation_type_kwargs.pop("catch_exceptions", None)
+
+    for expectation in expectations_list:
+        expectation_info = expectation["assertion_info"]
+        expectation_type_key = list(expectation_info.keys())[0]
+        expectations_type_value = dict(list(expectation_info.values())[0])
+        expectations_type_value.pop("catch_exceptions", None)
+
+        if (
+            expectation_type_key == expectation_type_action
+            and expectations_type_value == expectation_type_kwargs
+        ):
+            return expectation["assertion_mapping_id"]
+
+    return None

spark_dq/data_quality_helper/utils.py

@@ -0,0 +1,50 @@
+import base64
+import os
+import shutil
+import urllib.parse
+
+
+def convert_quality_report_to_binary(quality_report_lctn):
+    quality_report_lctn = quality_report_lctn.split("file:")[1]
+
+    with open(rf"{urllib.parse.unquote(quality_report_lctn)}", "rb") as f:
+        file_content = f.read()
+
+    # Encode the file content as base64
+    html_binary = base64.b64encode(file_content)
+
+    return html_binary
+
+
+def delete_quality_files(table_id: int, suite_id: int, quality_path: str):
+    # to delete the folder and its contents
+    folder_paths = [
+        os.path.join(
+            quality_path,
+            "expectations",
+            f"{table_id}_{suite_id}",
+        ),
+        os.path.join(
+            quality_path,
+            "uncommitted",
+            "validations",
+            f"{table_id}_{suite_id}",
+        ),
+        os.path.join(
+            quality_path,
+            "checkpoints",
+            f"{table_id}_{suite_id}",
+        ),
+        os.path.join(
+            quality_path,
+            "uncommitted",
+            "data_docs",
+            "local_site",
+            f"{table_id}_{suite_id}",
+        ),
+    ]
+
+    for folder_path in folder_paths:
+        if os.path.exists(folder_path):
+            shutil.rmtree(folder_path)
+            print(f"Folder '{folder_path}' and its contents have been deleted.")

spark_dq/helpers/__init__.py

@@ -0,0 +1,25 @@
+from spark_dq.helpers.exceptions import DQEngineAPIError
+from spark_dq.helpers.expectations import (
+    TRINO_FAST_TYPES,
+    exp_type_and_kwargs,
+    make_raw_result,
+    split_expectations,
+)
+from spark_dq.helpers.sql_utils import safe_identifier, safe_numeric, safe_regex
+from spark_dq.helpers import api_client
+from spark_dq.helpers import trino_executor
+from spark_dq.helpers import ge_executor
+
+__all__ = [
+    "DQEngineAPIError",
+    "TRINO_FAST_TYPES",
+    "exp_type_and_kwargs",
+    "make_raw_result",
+    "split_expectations",
+    "safe_identifier",
+    "safe_numeric",
+    "safe_regex",
+    "api_client",
+    "trino_executor",
+    "ge_executor",
+]

spark_dq/helpers/api_client.py

@@ -0,0 +1,121 @@
+import json
+import logging
+
+import requests
+
+from spark_dq.helpers.exceptions import DQEngineAPIError
+from spark_dq.data_quality_helper.ge_utils import get_mapping_id
+from spark_dq.utils import jsonify
+
+logger = logging.getLogger(__name__)
+
+API_TIMEOUT = 30
+
+
+def api_call(base_url: str, method: str, endpoint: str, body=None):
+    """Make an HTTP request to the DQ Engine API."""
+    url = base_url + endpoint
+    headers = {
+        "Accept": "application/json, text/plain, */*",
+        "Content-Type": "application/json; charset=UTF-8",
+    }
+    if body is None:
+        body = {}
+    if method.upper() == "GET":
+        resp = requests.request(
+            method, url, params=body, headers=headers, timeout=API_TIMEOUT,
+        )
+    else:
+        resp = requests.request(
+            method, url, data=body, headers=headers, timeout=API_TIMEOUT,
+        )
+
+    if resp.status_code in (200, 201):
+        return resp.json().get("data")
+
+    logger.error(
+        f"DQ Engine API error — {method} {url}\n"
+        f"  Status : {resp.status_code}\n"
+        f"  Body   : {body}\n"
+        f"  Response: {resp.content}"
+    )
+    raise DQEngineAPIError(resp.status_code, resp.text)
+
+
+def fetch_config(base_url: str, params: dict) -> dict | list:
+    """GET /config — works for both single-table and suite-run modes.
+
+    - Single-table (params has catalog_name/schema_name/table_name):
+        returns a single config dict.
+    - Suite-run (params has only test_suite_id):
+        returns a list of per-table config dicts.
+    """
+    data = api_call(base_url, "GET", "/config", body=params)
+    if isinstance(data, dict) and "tables" in data:
+        return data["tables"]
+    return data
+
+
+def save_results(
+    base_url: str,
+    *,
+    table_id: int,
+    schema_id: int,
+    catalog_id: int,
+    suite_id: int,
+    execution_type: str,
+    raw_results: list,
+    expectations_list: list,
+) -> dict:
+    """POST /save-results from a list of normalised raw-result dicts."""
+    total = len(raw_results)
+    passed = sum(1 for r in raw_results if r.get("success"))
+    stats = {
+        "evaluated_expectations": total,
+        "successful_expectations": passed,
+        "unsuccessful_expectations": total - passed,
+        "success_percent": round(passed / total * 100, 4) if total else 100.0,
+    }
+
+    results_payload = []
+    for ar in raw_results:
+        ec = ar.get("expectation_config", {})
+        exp_type = ec.get("expectation_type", "unknown")
+        exp_kwargs = ec.get("kwargs", {})
+        exception_info = ar.get("exception_info", {})
+
+        if isinstance(exception_info, dict) and exception_info.get("raised_exception"):
+            status = "SKIPPED"
+        elif ar.get("success"):
+            status = "SUCCESS"
+        else:
+            status = "FAILURE"
+
+        mapping_id = get_mapping_id(
+            actions_result=ar,
+            expectations_list=expectations_list,
+        )
+
+        run_event = {
+            "expectation_config": {"expectation_type": exp_type, "kwargs": exp_kwargs},
+            "result": ar.get("result", {}),
+            "success": ar.get("success", False),
+            "exception_info": exception_info,
+            "statistics": stats,
+        }
+        results_payload.append({
+            "assertion_mapping_id": mapping_id,
+            "assertion_name": exp_type,
+            "status": status,
+            "run_event": jsonify(run_event),
+        })
+
+    body = json.dumps({
+        "table_id": table_id,
+        "schema_id": schema_id,
+        "catalog_id": catalog_id,
+        "test_suite_id": suite_id,
+        "execution_type": execution_type,
+        "results": results_payload,
+    })
+    return api_call(base_url, "POST", "/save-results", body=body)

spark_dq/helpers/exceptions.py

@@ -0,0 +1,6 @@
+class DQEngineAPIError(Exception):
+    """Raised when the DQ Engine API returns a non-success response."""
+
+    def __init__(self, status_code: int, detail: str):
+        self.status_code = status_code
+        super().__init__(f"DQ Engine API error ({status_code}): {detail}")

spark_dq/helpers/expectations.py

@@ -0,0 +1,43 @@
+from typing import Tuple
+
+TRINO_FAST_TYPES = frozenset({
+    "expect_column_values_to_not_be_null",
+    "expect_column_values_to_be_unique",
+    "expect_column_values_to_be_between",
+    "expect_column_values_to_match_regex",
+    "expect_table_row_count_to_be_between",
+})
+
+
+def exp_type_and_kwargs(exp: dict) -> Tuple[str, dict]:
+    """Extract (expectation_type, kwargs) from an assertion_info dict."""
+    info = exp["assertion_info"]
+    exp_type = next(iter(info))
+    kwargs = dict(info[exp_type])
+    return exp_type, kwargs
+
+
+def make_raw_result(exp_type: str, kwargs: dict, success: bool) -> dict:
+    """Build a normalised raw-result dict for save-results payload."""
+    return {
+        "expectation_config": {
+            "expectation_type": exp_type,
+            "kwargs": {k: v for k, v in kwargs.items() if k != "catch_exceptions"},
+        },
+        "success": success,
+        "result": {},
+        "exception_info": {
+            "raised_exception": False,
+            "exception_message": None,
+            "exception_traceback": None,
+        },
+    }
+
+
+def split_expectations(expectations_list: list) -> Tuple[list, list]:
+    """Return (fast_exps, slow_exps) split by TRINO_FAST_TYPES membership."""
+    fast, slow = [], []
+    for exp in expectations_list:
+        etype, _ = exp_type_and_kwargs(exp)
+        (fast if etype in TRINO_FAST_TYPES else slow).append(exp)
+    return fast, slow

spark_dq/helpers/ge_executor.py

@@ -0,0 +1,145 @@
+import logging
+import yaml
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from typing import List
+
+from great_expectations.core import RunIdentifier
+from great_expectations.data_context import BaseDataContext
+
+from spark_dq.data_quality_helper import (
+    batch_request,
+    datacontext,
+    expectation_suite,
+)
+from spark_dq.data_quality_helper.const import DATASOURCE_CONFIG
+
+logger = logging.getLogger(__name__)
+
+
+def run_ge(table_name: str, expectations_list: list, df) -> dict:
+    """Build in-memory GE context, add expectations, run validator on *df*."""
+    env = "prod"
+    suite_name = f"quality_{env}_suite_{table_name}"
+    asset_name = f"quality_{env}_data_asset_{table_name}"
+    ds_name = f"quality_{env}_datasource_{table_name}"
+
+    data_context_config = datacontext.create_datacontext()
+    context = BaseDataContext(project_config=data_context_config)
+
+    ds_config_dict = yaml.safe_load(
+        DATASOURCE_CONFIG.format(datasource_name=ds_name)
+    )
+    context.add_datasource(**ds_config_dict)
+
+    context.add_or_update_expectation_suite(expectation_suite_name=suite_name)
+    suite_obj = context.get_expectation_suite(expectation_suite_name=suite_name)
+    suite_with_validations = expectation_suite.create_expectation_function(
+        expectation_list=expectations_list, suite=suite_obj
+    )
+    context.save_expectation_suite(
+        expectation_suite=suite_with_validations,
+        expectation_suite_name=suite_name,
+        overwrite_existing=True,
+    )
+
+    batch_req = batch_request.create_batchrequest(
+        df=df, datasource_name=ds_name, data_asset_name=asset_name
+    )
+    validator = context.get_validator(
+        batch_request=batch_req, expectation_suite=suite_with_validations
+    )
+    validator.save_expectation_suite(discard_failed_expectations=False)
+
+    validation_result = validator.validate(
+        result_format={
+            "result_format": "BOOLEAN_ONLY",
+            "include_unexpected_rows": False,
+            "partial_unexpected_count": 0,
+        },
+    )
+
+    run_id = RunIdentifier(run_name=f"validator-run-{table_name}")
+    return {
+        "run_results": {
+            run_id: {
+                "validation_result": (
+                    validation_result.to_json_dict()
+                    if hasattr(validation_result, "to_json_dict")
+                    else dict(validation_result)
+                )
+            }
+        }
+    }
+
+
+def extract_raw_from_checkpoint(checkpoint_run_result: dict) -> list:
+    """Unpack a GE validator/checkpoint result into a list of raw dicts."""
+    run_results = checkpoint_run_result.get("run_results", {})
+    if not run_results:
+        return []
+    assert_key = next(iter(run_results))
+    validation_result = run_results[assert_key]["validation_result"]
+
+    raw = []
+    for ar_obj in validation_result.get("results", []):
+        ar = ar_obj if isinstance(ar_obj, dict) else ar_obj.__dict__
+        ec_raw = ar.get("expectation_config", {})
+
+        exp_type, exp_kwargs = "unknown", {}
+        if hasattr(ec_raw, "expectation_type"):
+            exp_type = ec_raw.expectation_type
+            exp_kwargs = dict(ec_raw.kwargs)
+        else:
+            ec = ec_raw.__dict__ if hasattr(ec_raw, "__dict__") else ec_raw
+            if isinstance(ec, dict):
+                exp_type = (
+                    ec.get("expectation_type") or ec.get("_expectation_type", "unknown")
+                )
+                exp_kwargs = dict(ec.get("kwargs") or ec.get("_kwargs") or {})
+
+        raw.append({
+            "expectation_config": {"expectation_type": exp_type, "kwargs": exp_kwargs},
+            "success": ar.get("success", False),
+            "result": ar.get("result", {}),
+            "exception_info": ar.get("exception_info", {}),
+        })
+    return raw
+
+
+def run_ge_chunk(table_name: str, expectations_list: list, df, chunk_idx: int) -> List[dict]:
+    """Run a single GE validator for one chunk with unique context names."""
+    checkpoint = run_ge(f"{table_name}_p{chunk_idx}", expectations_list, df)
+    return extract_raw_from_checkpoint(checkpoint)
+
+
+def run_ge_parallel(
+    table_name: str, expectations_list: list, df, workers: int
+) -> List[dict]:
+    """Split expectations into chunks and run each in its own thread."""
+    n = min(workers, len(expectations_list))
+    if n <= 1:
+        return run_ge_chunk(table_name, expectations_list, df, 0)
+
+    chunk_size = -(-len(expectations_list) // n)
+    chunks = [
+        expectations_list[i : i + chunk_size]
+        for i in range(0, len(expectations_list), chunk_size)
+    ]
+    logger.info(
+        f"GE parallel: {len(expectations_list)} exps → "
+        f"{len(chunks)} chunks of ~{chunk_size}"
+    )
+
+    all_raw: List[dict] = []
+    with ThreadPoolExecutor(max_workers=len(chunks)) as pool:
+        futures = {
+            pool.submit(run_ge_chunk, table_name, chunk, df, idx): idx
+            for idx, chunk in enumerate(chunks)
+        }
+        for future in as_completed(futures):
+            idx = futures[future]
+            chunk_raw = future.result()
+            passed = sum(1 for r in chunk_raw if r["success"])
+            logger.info(f"  GE chunk {idx} done — {passed}/{len(chunk_raw)} passed")
+            all_raw.extend(chunk_raw)
+    return all_raw

spark_dq/helpers/sql_utils.py

@@ -0,0 +1,26 @@
+import re
+
+_IDENTIFIER_RE = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$")
+
+
+def safe_identifier(name: str) -> str:
+    """Validate and quote a SQL identifier to prevent injection."""
+    if not name or not _IDENTIFIER_RE.match(name):
+        raise ValueError(f"Invalid SQL identifier: {name!r}")
+    return f'"{name}"'
+
+
+def safe_numeric(value) -> str:
+    """Ensure a value is a valid number literal for SQL."""
+    try:
+        float(value)
+    except (TypeError, ValueError):
+        raise ValueError(f"Expected numeric value, got: {value!r}")
+    return str(value)
+
+
+def safe_regex(pattern: str) -> str:
+    """Escape single quotes in a regex pattern for SQL string literal."""
+    if not isinstance(pattern, str):
+        raise ValueError(f"Expected regex string, got: {type(pattern).__name__}")
+    return pattern.replace("'", "''")

spark_dq/helpers/trino_executor.py

@@ -0,0 +1,171 @@
+import logging
+from typing import Dict
+
+from spark_dq.helpers.expectations import (
+    exp_type_and_kwargs,
+    make_raw_result,
+)
+from spark_dq.helpers.sql_utils import (
+    safe_identifier,
+    safe_numeric,
+    safe_regex,
+)
+
+logger = logging.getLogger(__name__)
+
+
+def run_trino_fast_path(
+    *,
+    trino_host: str,
+    trino_user: str,
+    trino_pwd: str,
+    catalog: str,
+    schema: str,
+    table: str,
+    row_limit: int | None,
+    fast_expectations: list,
+) -> list:
+    """
+    Execute fast-path expectations as ONE batched SQL query on Trino.
+
+    Uses validated/quoted identifiers and numeric literals to prevent
+    SQL injection. The connection is guaranteed to be closed via finally.
+    """
+    import trino as trino_mod
+
+    host_port = trino_host.rsplit(":", 1)
+    host = host_port[0]
+    port = int(host_port[1]) if len(host_port) > 1 else 443
+
+    q_catalog = safe_identifier(catalog)
+    q_schema = safe_identifier(schema)
+    q_table = safe_identifier(table)
+    fq_table = f"{q_catalog}.{q_schema}.{q_table}"
+
+    if row_limit:
+        table_ref = f"(SELECT * FROM {fq_table} LIMIT {safe_numeric(row_limit)})"
+    else:
+        table_ref = fq_table
+
+    selects = ["COUNT(*) AS __total"]
+    nonnull_selects: Dict[str, str] = {}
+    index_map: Dict[str, tuple] = {}
+    rowcount_exps = []
+
+    for i, exp in enumerate(fast_expectations):
+        etype, kwargs = exp_type_and_kwargs(exp)
+        col = kwargs.get("column", "")
+        qcol = safe_identifier(col) if col else '""'
+
+        if etype == "expect_column_values_to_not_be_null":
+            alias = f"__n_{i}"
+            selects.append(
+                f"SUM(CASE WHEN {qcol} IS NULL THEN 1 ELSE 0 END) AS {alias}"
+            )
+            index_map[alias] = (exp, etype, kwargs, "null")
+
+        elif etype == "expect_column_values_to_be_unique":
+            alias = f"__u_{i}"
+            selects.append(
+                f"(COUNT({qcol}) - COUNT(DISTINCT {qcol})) AS {alias}"
+            )
+            index_map[alias] = (exp, etype, kwargs, "unique")
+
+        elif etype == "expect_column_values_to_be_between":
+            alias = f"__b_{i}"
+            min_v = kwargs.get("min_value")
+            max_v = kwargs.get("max_value")
+            parts = [f"{qcol} IS NOT NULL"]
+            if min_v is not None and max_v is not None:
+                parts.append(
+                    f"NOT ({qcol} BETWEEN {safe_numeric(min_v)} AND {safe_numeric(max_v)})"
+                )
+            elif min_v is not None:
+                parts.append(f"{qcol} < {safe_numeric(min_v)}")
+            elif max_v is not None:
+                parts.append(f"{qcol} > {safe_numeric(max_v)}")
+            selects.append(
+                f"SUM(CASE WHEN {' AND '.join(parts)} THEN 1 ELSE 0 END) AS {alias}"
+            )
+            index_map[alias] = (exp, etype, kwargs, "between")
+
+        elif etype == "expect_column_values_to_match_regex":
+            alias = f"__r_{i}"
+            regex = safe_regex(kwargs.get("regex", ".*"))
+            selects.append(
+                f"SUM(CASE WHEN {qcol} IS NOT NULL"
+                f" AND NOT REGEXP_LIKE(CAST({qcol} AS VARCHAR), '{regex}')"
+                f" THEN 1 ELSE 0 END) AS {alias}"
+            )
+            index_map[alias] = (exp, etype, kwargs, "regex")
+
+        elif etype == "expect_table_row_count_to_be_between":
+            rowcount_exps.append((exp, etype, kwargs))
+
+        if col and etype != "expect_table_row_count_to_be_between":
+            nn_alias = f"__nn_{col}"
+            if nn_alias not in nonnull_selects:
+                nonnull_selects[nn_alias] = f"COUNT({qcol}) AS {safe_identifier(nn_alias)}"
+
+    selects.extend(nonnull_selects.values())
+
+    sql = f"SELECT {', '.join(selects)} FROM {table_ref}"
+    logger.info(
+        f"Trino fast-path: {len(fast_expectations)} expectations → 1 SQL query | "
+        f"SQL: {sql[:500]}"
+    )
+
+    conn = trino_mod.dbapi.connect(
+        host=host,
+        port=port,
+        user=trino_user,
+        http_scheme="https",
+        auth=trino_mod.auth.BasicAuthentication(trino_user, trino_pwd),
+    )
+    try:
+        cursor = conn.cursor()
+        cursor.execute(sql)
+        row = cursor.fetchone()
+        col_names = [d[0] for d in cursor.description]
+        cursor.close()
+    finally:
+        conn.close()
+
+    result_map = dict(zip(col_names, row))
+    total_rows = int(result_map.get("__total") or 0)
+    raw_results = []
+
+    for alias, (exp, etype, kwargs, check_type) in index_map.items():
+        mostly = float(kwargs.get("mostly") or 1.0)
+        unexpected = int(result_map.get(alias) or 0)
+
+        col = kwargs.get("column", "")
+        nonnull_count = int(result_map.get(f"__nn_{col}") or total_rows)
+        denominator = nonnull_count if check_type != "null" else total_rows
+
+        if mostly >= 1.0:
+            success = unexpected == 0
+        else:
+            success = (
+                (unexpected / denominator) <= (1.0 - mostly)
+                if denominator > 0
+                else True
+            )
+        raw_results.append(make_raw_result(etype, kwargs, success))
+
+    for exp, etype, kwargs in rowcount_exps:
+        min_v = kwargs.get("min_value")
+        max_v = kwargs.get("max_value")
+        success = True
+        if min_v is not None and total_rows < min_v:
+            success = False
+        if max_v is not None and total_rows > max_v:
+            success = False
+        raw_results.append(make_raw_result(etype, kwargs, success))
+
+    passed = sum(1 for r in raw_results if r["success"])
+    logger.info(
+        f"Trino fast-path complete: {passed}/{len(raw_results)} passed "
+        f"(total_rows={total_rows:,})"
+    )
+    return raw_results

spark_dq/quality.py

@@ -0,0 +1,282 @@
+import logging
+import os
+import time
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from typing import Dict, Any, List
+
+from spark_dq.helpers.expectations import split_expectations
+from spark_dq.helpers.api_client import (
+    fetch_config,
+    save_results,
+)
+from spark_dq.helpers.trino_executor import run_trino_fast_path
+from spark_dq.helpers.ge_executor import run_ge_parallel
+
+logger = logging.getLogger(__name__)
+
+_MAX_SUITE_WORKERS = 8
+
+
+def _fmt(seconds: float) -> str:
+    if seconds < 60:
+        return f"{seconds:.2f}s"
+    m, s = divmod(seconds, 60)
+    return f"{int(m)}m {s:.2f}s"
+
+
+class SparkDQAgent:
+    """
+    Data Quality Agent designed for K8s Spark pods.
+
+    Receives a pre-loaded Spark DataFrame, runs validations, and persists
+    results to the Data Quality Engine service via its HTTP API.
+
+    **Hybrid execution** (active when Trino credentials are supplied):
+      - TRINO_FAST_TYPES  → single batched SQL query on Trino per suite
+      - All other types   → Great Expectations on the Spark DataFrame
+
+    When every expectation in a suite belongs to TRINO_FAST_TYPES the
+    caller may pass ``df=None`` — Spark is not touched at all.
+
+    Remote calls per suite:
+      1. ``GET  /config``       — fetch query, table type, suites + expectations
+      2. ``POST /save-results`` — persist results
+    """
+
+    def __init__(
+        self,
+        catalog: str,
+        schema: str,
+        table: str,
+        data_quality_url: str,
+        catalog_type: str,
+        log_level: str = "INFO",
+        test_suite_id: int = None,
+        table_id: int = None,
+        trino_host: str = "",
+        trino_user: str = "",
+        trino_pwd: str = "",
+        is_suite_run: bool = False,
+        config: dict = None,
+    ):
+        self.catalog = catalog
+        self.schema = schema
+        self.table = table
+        self.data_quality_url = data_quality_url.rstrip("/")
+        self.catalog_type = catalog_type
+        self.log_level = log_level
+        self.test_suite_id = test_suite_id
+        self.table_id = table_id
+        self.trino_host = trino_host or ""
+        self.trino_user = trino_user or ""
+        self.trino_pwd = trino_pwd or ""
+        self.is_suite_run = is_suite_run
+        self.row_limit = None
+        self._config = config
+
+    # ------------------------------------------------------------------
+    # Public API
+    # ------------------------------------------------------------------
+
+    def get_config(self) -> dict:
+        """Fetch config for this table from the API (or return cached).
+
+        Returns a dict with quality_query, table_type, suites, scan_limit, etc.
+        """
+        if self._config is None:
+            params = {
+                "catalog_name": self.catalog,
+                "schema_name": self.schema,
+                "table_name": self.table,
+                "catalog_type": self.catalog_type,
+            }
+            if self.table_id is not None:
+                params["table_id"] = self.table_id
+            if self.test_suite_id is not None:
+                params["test_suite_id"] = self.test_suite_id
+            self._config = fetch_config(self.data_quality_url, params)
+        return self._config
+
+    def get_all_table_configs(self) -> list:
+        """Fetch configs for ALL tables in the test suite from the API.
+
+        Returns a list of config dicts, one per table.
+        """
+        params = {"test_suite_id": self.test_suite_id, "catalog_type": self.catalog_type}
+        return fetch_config(self.data_quality_url, params)
+
+    def validate(self, df) -> Dict[str, Any]:
+        """
+        Run validations against the provided Spark DataFrame (or ``None``).
+
+        All suites are processed **in parallel** via ThreadPoolExecutor.
+        Within each suite, Trino fast-path and GE slow-path also run concurrently.
+        """
+        fqn = f"{self.catalog}.{self.schema}.{self.table}"
+        logger.info(f"Starting validation for {fqn}")
+
+        cfg = self.get_config()
+        suites = cfg.get("suites", [])
+        schema_id = cfg.get("schema_id")
+        catalog_id = cfg.get("catalog_id")
+
+        config_limit = cfg.get("scan_limit")
+        if config_limit is not None:
+            self.row_limit = int(config_limit)
+
+        if not suites:
+            logger.warning(f"No suites returned for {fqn}")
+            return {}
+
+        use_hybrid = bool(self.trino_host and self.trino_user)
+        all_results: Dict[str, Any] = {}
+
+        workers = min(len(suites), _MAX_SUITE_WORKERS)
+        with ThreadPoolExecutor(max_workers=workers) as pool:
+            futures = {
+                pool.submit(
+                    self._validate_suite, suite, df, use_hybrid,
+                    schema_id=schema_id, catalog_id=catalog_id,
+                ): suite
+                for suite in suites
+            }
+            for future in  (futures):
+                suite = futures[future]
+                suite_id = suite["suite_id"]
+                table_name = suite.get("table_name", self.table)
+                result_key = f"{table_name}_{suite_id}"
+                try:
+                    all_results[result_key] = future.result()
+                except Exception as exc:
+                    logger.error(f"Validation failed for {result_key}: {exc}")
+                    all_results[result_key] = {"error": str(exc)}
+
+        logger.info(f"Validation complete for {fqn}")
+        return all_results
+
+    # ------------------------------------------------------------------
+    # Internal
+    # ------------------------------------------------------------------
+
+    def _validate_suite(
+        self, suite: dict, df, use_hybrid: bool,
+        schema_id: int = None, catalog_id: int = None,
+    ) -> Dict[str, Any]:
+        """Validate a single suite — called in parallel from validate()."""
+        suite_id = suite["suite_id"]
+        table_id = suite.get("table_id") or self.table_id
+        table_name = suite.get("table_name", self.table)
+        expectations_list = suite.get("expectations", [])
+        ge_workers = min(4, os.cpu_count() or 4)
+
+        t_suite = time.perf_counter()
+
+        if use_hybrid:
+            all_raw = self._run_hybrid(
+                suite_id, table_name, expectations_list, df, ge_workers, t_suite,
+            )
+        else:
+            logger.info(
+                f"[suite {suite_id}] GE-only — "
+                f"{len(expectations_list)} expectations, workers: {ge_workers}"
+            )
+            t0 = time.perf_counter()
+            all_raw = run_ge_parallel(table_name, expectations_list, df, ge_workers)
+            logger.info(
+                f"[suite {suite_id}][timer] GE parallel "
+                f"({len(expectations_list)} exps, {ge_workers} workers): "
+                f"{_fmt(time.perf_counter() - t0)}"
+            )
+
+        t0 = time.perf_counter()
+        try:
+            save_results(
+                self.data_quality_url,
+                table_id=table_id,
+                schema_id=schema_id,
+                catalog_id=catalog_id,
+                suite_id=suite_id,
+                execution_type="SCHEDULED" if self.is_suite_run else "MANUAL",
+                raw_results=all_raw,
+                expectations_list=expectations_list,
+            )
+        except Exception as exc:
+            logger.error(
+                f"[suite {suite_id}] POST /save-results failed: {exc}. "
+                f"Computed {len(all_raw)} results (not persisted)."
+            )
+            raise
+        logger.info(
+            f"[suite {suite_id}][timer] POST /save-results "
+            f"({len(all_raw)} results): {_fmt(time.perf_counter() - t0)}"
+        )
+
+        passed = sum(1 for r in all_raw if r.get("success"))
+        total = len(all_raw)
+        logger.info(
+            f"[suite {suite_id}][timer] Suite total: {_fmt(time.perf_counter() - t_suite)} "
+            f"| passed {passed}/{total}"
+        )
+        return {
+            "evaluated_expectations": total,
+            "successful_expectations": passed,
+            "unsuccessful_expectations": total - passed,
+            "success_percent": round(passed / total * 100, 4) if total else 100.0,
+        }
+
+    def _run_hybrid(
+        self,
+        suite_id: int,
+        table_name: str,
+        expectations_list: list,
+        df,
+        ge_workers: int,
+        t_suite: float,
+    ) -> List[dict]:
+        """Run Trino fast-path and GE slow-path concurrently."""
+        fast_exps, slow_exps = split_expectations(expectations_list)
+        logger.info(
+            f"[suite {suite_id}] Hybrid — "
+            f"Trino fast: {len(fast_exps)}, GE slow: {len(slow_exps)}, "
+            f"workers: {ge_workers}"
+        )
+
+        fast_raw: List[dict] = []
+        ge_raw: List[dict] = []
+
+        with ThreadPoolExecutor(max_workers=2) as pool:
+            futures = {}
+            if fast_exps:
+                futures[pool.submit(
+                    run_trino_fast_path,
+                    trino_host=self.trino_host,
+                    trino_user=self.trino_user,
+                    trino_pwd=self.trino_pwd,
+                    catalog=self.catalog,
+                    schema=self.schema,
+                    table=self.table,
+                    row_limit=self.row_limit,
+                    fast_expectations=fast_exps,
+                )] = "trino"
+            if slow_exps:
+                futures[pool.submit(
+                    run_ge_parallel, table_name, slow_exps, df, ge_workers,
+                )] = "ge"
+
+            for future in as_completed(futures):
+                label = futures[future]
+                if label == "trino":
+                    fast_raw = future.result()
+                    logger.info(
+                        f"[suite {suite_id}][timer] Trino fast-path "
+                        f"({len(fast_exps)} exps): {_fmt(time.perf_counter() - t_suite)}"
+                    )
+                else:
+                    ge_raw = future.result()
+                    logger.info(
+                        f"[suite {suite_id}][timer] GE parallel "
+                        f"({len(slow_exps)} exps, {ge_workers} workers): "
+                        f"{_fmt(time.perf_counter() - t_suite)}"
+                    )
+
+        return fast_raw + ge_raw

spark_dq/utils.py

@@ -0,0 +1,32 @@
+from datetime import date, datetime, timedelta
+from decimal import Decimal
+from enum import Enum
+
+from great_expectations.core.expectation_configuration import ExpectationConfiguration
+from great_expectations.core.expectation_validation_result import (
+    ExpectationValidationResult,
+)
+
+datetime_format = "%Y-%m-%d %H:%M:%S"
+
+
+def jsonify(data):
+    if isinstance(data, dict):
+        return {key: jsonify(value) for key, value in data.items()}
+    elif isinstance(data, list):
+        return [jsonify(list_obj) for list_obj in data]
+    elif isinstance(data, datetime):
+        return data.strftime(datetime_format)
+    elif isinstance(data, date):
+        return data.isoformat()
+    elif isinstance(data, timedelta):
+        return data.total_seconds()
+    elif isinstance(data, (ExpectationConfiguration, ExpectationValidationResult)):
+        return {key: jsonify(value) for key, value in dict(data).items()}
+    elif isinstance(data, Decimal):
+        return float(data)
+    elif isinstance(data, Enum):
+        return data.value
+    elif data is None:
+        return data
+    return data