data-validation-engine 0.7.4__tar.gz → 0.7.6__tar.gz

{data_validation_engine-0.7.4 → data_validation_engine-0.7.6}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: data-validation-engine
-Version: 0.7.4
+Version: 0.7.6
 Summary: `nhs data validation engine` is a framework used to validate data
 License-Expression: MIT
 License-File: LICENSE

{data_validation_engine-0.7.4 → data_validation_engine-0.7.6}/pyproject.toml

@@ -24,7 +24,7 @@ Issues = "https://github.com/NHSDigital/data-validation-engine/issues"
 Changelog = "https://github.com/NHSDigital/data-validation-engine/blob/main/CHANGELOG.md"
 
 [tool.poetry]
-version = "0.7.4"
+version = "0.7.6"
 packages = [
     { include = "dve", from = "src" },
 ]

{data_validation_engine-0.7.4 → data_validation_engine-0.7.6}/src/dve/core_engine/backends/base/backend.py

@@ -3,7 +3,7 @@
 import logging
 import warnings
 from abc import ABC, abstractmethod
-from collections.abc import Mapping, MutableMapping
+from collections.abc import MutableMapping
 from typing import Any, ClassVar, Generic, Optional
 
 from pyspark.sql import DataFrame, SparkSession
@@ -41,14 +41,12 @@ class BaseBackend(Generic[EntityType], ABC):
         self,
         contract: BaseDataContract[EntityType],
         steps: BaseStepImplementations[EntityType],
-        reference_data_loader_type: Optional[type[BaseRefDataLoader[EntityType]]],
         logger: Optional[logging.Logger] = None,
         **kwargs: Any,
     ) -> None:
         for component_name, component in (
             ("Contract", contract),
             ("Step implementation", steps),
-            ("Reference data loader", reference_data_loader_type),
         ):
             component_entity_type = getattr(component, "__entity_type__", None)
             if component_entity_type != self.__entity_type__:
@@ -61,12 +59,6 @@ class BaseBackend(Generic[EntityType], ABC):
         """The data contract implementation used by the backend."""
         self.step_implementations = steps
         """The step implementations used by the backend."""
-        self.reference_data_loader_type = reference_data_loader_type
-        """
-        The loader type to use for the reference data. If `None`, do not
-        load any reference data and error if it is provided.
-
-        """
         self.logger = logger or get_logger(type(self).__name__)
         """The `logging.Logger instance for the backend."""
 
@@ -74,29 +66,9 @@ class BaseBackend(Generic[EntityType], ABC):
         self,
         reference_entity_config: dict[EntityName, ReferenceConfigUnion],
         submission_info: Optional[SubmissionInfo],
-    ) -> Mapping[EntityName, EntityType]:
-        """Load the reference data as specified in the reference entity config."""
-        sub_info_entity: Optional[EntityType] = None
-        if submission_info:
-            sub_info_entity = self.convert_submission_info(submission_info)
-
-        if self.reference_data_loader_type is None:
-            if reference_entity_config:
-                raise ValueError(
-                    "Reference data has been specified but no reference data loader is "
-                    + "configured for this backend"
-                )
-
-            reference_data_dict = {}
-            if sub_info_entity is not None:
-                reference_data_dict["dve_submission_info"] = sub_info_entity
-            return reference_data_dict
-
-        reference_data_loader = self.reference_data_loader_type(reference_entity_config)
-        if sub_info_entity is not None:
-            reference_data_loader.entity_cache["dve_submission_info"] = sub_info_entity
-
-        return reference_data_loader
+    ) -> BaseRefDataLoader[EntityType]:
+        """Supply configured reference data loader for use with business rules"""
+        raise NotImplementedError()
 
     @abstractmethod
     def convert_submission_info(self, submission_info: SubmissionInfo) -> EntityType:

{data_validation_engine-0.7.4 → data_validation_engine-0.7.6}/src/dve/core_engine/backends/base/reference_data.py

@@ -11,6 +11,7 @@ import dve.parser.file_handling as fh
 from dve.core_engine.backends.base.core import get_entity_type
 from dve.core_engine.backends.exceptions import (
     MissingRefDataEntity,
+    NoRefDataConfigSupplied,
     RefdataLacksFileExtensionSupport,
 )
 from dve.core_engine.backends.types import EntityType
@@ -147,11 +148,11 @@ class BaseRefDataLoader(Generic[EntityType], Mapping[EntityName, EntityType], AB
     # pylint: disable=unused-argument
     def __init__(
         self,
-        reference_entity_config: dict[EntityName, ReferenceConfig],
-        dataset_config_uri: Optional[URI] = None,
+        reference_data_config: dict[EntityName, ReferenceConfig],
+        dataset_config_uri: URI,
         **kwargs,
     ) -> None:
-        self.reference_entity_config = reference_entity_config
+        self.reference_entity_config = reference_data_config
         self.dataset_config_uri = dataset_config_uri
         """
         Configuration options for the reference data. This is likely to vary
@@ -207,6 +208,8 @@ class BaseRefDataLoader(Generic[EntityType], Mapping[EntityName, EntityType], AB
             try:
                 config = self.reference_entity_config[key]
                 return self.load_entity(entity_name=key, config=config)
+            except TypeError as err:
+                raise NoRefDataConfigSupplied() from err
             except Exception as err:
                 raise MissingRefDataEntity(entity_name=key) from err
 

{data_validation_engine-0.7.4 → data_validation_engine-0.7.6}/src/dve/core_engine/backends/exceptions.py

@@ -119,6 +119,20 @@ class MissingRefDataEntity(MissingEntity, BackendErrorMixin):  # pylint: disable
         return f"Missing reference data entity {self.entity_name!r}"
 
 
+class NoRefDataConfigSupplied(BackendError):
+    """An error raised when trying to load a refdata entity when no refdata
+    config has been supplied.
+
+    """
+
+    def __init__(self, *args: object) -> None:
+        super().__init__(*args)
+
+    def get_message_preamble(self) -> EntityName:
+        """Message for logging purposes"""
+        return "Refdata loader not supplied with refdata config - unable to load refdata entities"
+
+
 class ConstraintError(ValueError, BackendErrorMixin):
     """Raised when a given constraint is violated."""
 

{data_validation_engine-0.7.4 → data_validation_engine-0.7.6}/src/dve/core_engine/backends/implementations/duckdb/duckdb_helpers.py

@@ -411,7 +411,7 @@ def get_duckdb_cast_statement_from_annotation(
             stmt = rf"CASE WHEN REGEXP_MATCHES(TRIM({quoted_name}), '{date_regex}') THEN TRY_CAST(TRIM({quoted_name}) as DATE) ELSE NULL END"  # pylint: disable=C0301
             return stmt
         if issubclass(type_, time):
-            stmt = rf"CASE WHEN REGEXP_MATCHES(TRIM({quoted_name}), '{time_regex}') THEN TRY_CAST(TRIM({quoted_name}) as TIME) ELSE NULL END" # pylint: disable=C0301
+            stmt = rf"CASE WHEN REGEXP_MATCHES(TRIM({quoted_name}), '{time_regex}') THEN TRY_CAST(TRIM({quoted_name}) as TIME) ELSE NULL END"  # pylint: disable=C0301
             return stmt
         duck_type = get_duckdb_type_from_annotation(type_)
         if duck_type:

{data_validation_engine-0.7.4 → data_validation_engine-0.7.6}/src/dve/core_engine/backends/implementations/duckdb/readers/csv.py

@@ -6,13 +6,7 @@ from typing import Any, Optional
 
 import duckdb as ddb
 import polars as pl
-from duckdb import (
-    DuckDBPyConnection,
-    DuckDBPyRelation,
-    StarExpression,
-    default_connection,
-    read_csv,
-)
+from duckdb import DuckDBPyConnection, DuckDBPyRelation, StarExpression, read_csv
 from pydantic import BaseModel
 
 from dve.core_engine.backends.base.reader import BaseFileReader, read_function
@@ -61,7 +55,7 @@ class DuckDBCSVReader(BaseFileReader):
         self.header = header
         self.delim = delim
         self.quotechar = quotechar
-        self._connection = connection if connection else default_connection
+        self._connection = connection if connection else ddb.connect(":memory:")
         self.field_check = field_check
         self.field_check_error_code = field_check_error_code
         self.field_check_error_message = field_check_error_message
@@ -181,7 +175,7 @@ class PolarsToDuckDBCSVReader(DuckDBCSVReader):
             ] + [pl.col(RECORD_INDEX_COLUMN_NAME)]
             df = df.select(pl_exprs)
 
-        return ddb.sql("SELECT * FROM df")
+        return self._connection.sql("SELECT * FROM df")
 
 
 class DuckDBCSVRepeatingHeaderReader(PolarsToDuckDBCSVReader):

{data_validation_engine-0.7.4 → data_validation_engine-0.7.6}/src/dve/core_engine/backends/implementations/duckdb/readers/json.py

@@ -4,7 +4,8 @@
 from collections.abc import Iterator
 from typing import Any, Optional
 
-from duckdb import DuckDBPyRelation, read_json
+import duckdb
+from duckdb import DuckDBPyConnection, DuckDBPyRelation
 from pydantic import BaseModel
 
 from dve.core_engine.backends.base.reader import BaseFileReader, read_function
@@ -26,9 +27,11 @@ class DuckDBJSONReader(BaseFileReader):
         self,
         *,
         json_format: Optional[str] = "array",
+        connection: Optional[DuckDBPyConnection] = None,
         **_,
     ):
         self._json_format = json_format
+        self._connection = duckdb.connect(":memory:") if not connection else connection
 
         super().__init__()
 
@@ -50,5 +53,7 @@ class DuckDBJSONReader(BaseFileReader):
         }
 
         return self.add_record_index(
-            read_json(resource, columns=ddb_schema, format=self._json_format)  # type: ignore
+            self._connection.read_json(
+                resource, columns=ddb_schema, format=self._json_format  # type: ignore
+            )
         )

{data_validation_engine-0.7.4 → data_validation_engine-0.7.6}/src/dve/core_engine/backends/implementations/duckdb/readers/xml.py

@@ -3,8 +3,9 @@
 
 from typing import Optional
 
+import duckdb
 import polars as pl
-from duckdb import DuckDBPyConnection, DuckDBPyRelation, default_connection
+from duckdb import DuckDBPyConnection, DuckDBPyRelation
 from pydantic import BaseModel
 
 from dve.core_engine.backends.base.reader import read_function
@@ -24,8 +25,8 @@ from dve.core_engine.type_hints import URI
 class DuckDBXMLStreamReader(XMLStreamReader):
     """A reader for XML files"""
 
-    def __init__(self, *, ddb_connection: Optional[DuckDBPyConnection] = None, **kwargs):
-        self.ddb_connection = ddb_connection if ddb_connection else default_connection
+    def __init__(self, *, connection: Optional[DuckDBPyConnection] = None, **kwargs):
+        self._connection = connection if connection else duckdb.connect(":memory:")
         super().__init__(**kwargs)
 
     @read_function(DuckDBPyRelation)
@@ -49,4 +50,4 @@ class DuckDBXMLStreamReader(XMLStreamReader):
                 data=self.read_to_py_iterator(resource, entity_name, schema), schema=polars_schema
             )
         )
-        return self.ddb_connection.sql("select * from _lazy_frame")
+        return self._connection.sql("select * from _lazy_frame")

{data_validation_engine-0.7.4 → data_validation_engine-0.7.6}/src/dve/core_engine/backends/implementations/duckdb/reference_data.py

@@ -1,13 +1,11 @@
 """A reference data loader for duckdb."""
 
-from typing import Optional
-
 from duckdb import DuckDBPyConnection, DuckDBPyRelation
 from pyarrow import ipc  # type: ignore
 
 from dve.core_engine.backends.base.reference_data import (
     BaseRefDataLoader,
-    ReferenceConfigUnion,
+    ReferenceConfig,
     ReferenceTable,
     mark_refdata_file_extension,
 )
@@ -19,17 +17,16 @@ from dve.parser.type_hints import URI
 class DuckDBRefDataLoader(BaseRefDataLoader[DuckDBPyRelation]):
     """A reference data loader using already existing DuckDB tables."""
 
-    connection: DuckDBPyConnection
-    """The DuckDB connection for the backend."""
-    dataset_config_uri: Optional[URI] = None
-    """The location of the dischema file"""
-
     def __init__(
         self,
-        reference_entity_config: dict[EntityName, ReferenceConfigUnion],
+        connection: DuckDBPyConnection,
+        reference_data_config: dict[EntityName, ReferenceConfig],
+        dataset_config_uri: URI,
         **kwargs,
     ) -> None:
-        super().__init__(reference_entity_config, self.dataset_config_uri, **kwargs)
+        super().__init__(reference_data_config, dataset_config_uri, **kwargs)
+
+        self.connection = connection
 
         if not self.connection:
             raise AttributeError("DuckDBConnection must be specified")

{data_validation_engine-0.7.4 → data_validation_engine-0.7.6}/src/dve/core_engine/backends/implementations/spark/backend.py

@@ -6,6 +6,7 @@ from typing import Any, Optional
 from pyspark.sql import DataFrame, SparkSession
 
 from dve.core_engine.backends.base.backend import BaseBackend
+from dve.core_engine.backends.base.reference_data import ReferenceConfigUnion
 from dve.core_engine.backends.implementations.spark.contract import SparkDataContract
 from dve.core_engine.backends.implementations.spark.reference_data import SparkRefDataLoader
 from dve.core_engine.backends.implementations.spark.rules import SparkStepImplementations
@@ -14,7 +15,7 @@ from dve.core_engine.backends.implementations.spark.types import SparkEntities
 from dve.core_engine.constants import RECORD_INDEX_COLUMN_NAME
 from dve.core_engine.loggers import get_child_logger, get_logger
 from dve.core_engine.models import SubmissionInfo
-from dve.core_engine.type_hints import URI, EntityParquetLocations
+from dve.core_engine.type_hints import URI, EntityName, EntityParquetLocations
 from dve.parser.file_handling import get_resource_exists, joinuri
 
 
@@ -26,7 +27,6 @@ class SparkBackend(BaseBackend[DataFrame]):
         dataset_config_uri: Optional[URI] = None,
         contract: Optional[SparkDataContract] = None,
         steps: Optional[SparkStepImplementations] = None,
-        reference_data_loader: Optional[type[SparkRefDataLoader]] = None,
         logger: Optional[logging.Logger] = None,
         spark_session: Optional[SparkSession] = None,
         **kwargs: Any,
@@ -36,6 +36,8 @@ class SparkBackend(BaseBackend[DataFrame]):
 
         self.spark_session = spark_session or SparkSession.builder.getOrCreate()
         """The Spark session for the backend."""
+        self.dataset_config_uri = dataset_config_uri
+        """The uri of the dischema specifying the DVE config"""
 
         if contract is None:
             contract = SparkDataContract(
@@ -46,11 +48,27 @@ class SparkBackend(BaseBackend[DataFrame]):
             steps = SparkStepImplementations.register_udfs(
                 logger=get_child_logger("SparkStepImplementations", logger)
             )
-        if reference_data_loader is None:
-            reference_data_loader = SparkRefDataLoader
-        reference_data_loader.spark = self.spark_session
-        reference_data_loader.dataset_config_uri = dataset_config_uri
-        super().__init__(contract, steps, reference_data_loader, logger, **kwargs)
+        super().__init__(contract, steps, logger, **kwargs)
+
+    def load_reference_data(
+        self,
+        reference_entity_config: dict[EntityName, ReferenceConfigUnion],
+        submission_info: Optional[SubmissionInfo],
+    ):
+        """Load the reference data as specified in the reference entity config."""
+        sub_info_entity: Optional[DataFrame] = None
+        if submission_info:
+            sub_info_entity = self.convert_submission_info(submission_info)
+
+        reference_data_loader = SparkRefDataLoader(
+            spark=self.spark_session,
+            reference_data_config=reference_entity_config,
+            dataset_config_uri=self.dataset_config_uri,  # type: ignore
+        )
+        if sub_info_entity is not None:
+            reference_data_loader.entity_cache["dve_submission_info"] = sub_info_entity
+
+        return reference_data_loader
 
     def write_entities_to_parquet(
         self, entities: SparkEntities, cache_prefix: URI

{data_validation_engine-0.7.4 → data_validation_engine-0.7.6}/src/dve/core_engine/backends/implementations/spark/reference_data.py

@@ -1,8 +1,6 @@
 # pylint: disable=no-member
 """A reference data loader for Spark."""
 
-from typing import Optional
-
 from pyspark.sql import DataFrame, SparkSession
 
 from dve.core_engine.backends.base.reference_data import (
@@ -19,17 +17,15 @@ from dve.parser.type_hints import URI
 class SparkRefDataLoader(BaseRefDataLoader[DataFrame]):
     """A reference data loader using already existing Apache Spark Tables."""
 
-    spark: SparkSession
-    """The Spark session for the backend."""
-    dataset_config_uri: Optional[URI] = None
-    """The location of the dischema file defining business rules"""
-
     def __init__(
         self,
-        reference_entity_config: dict[EntityName, ReferenceConfig],
+        spark: SparkSession,
+        reference_data_config: dict[EntityName, ReferenceConfig],
+        dataset_config_uri: URI,
         **kwargs,
     ) -> None:
-        super().__init__(reference_entity_config, self.dataset_config_uri, **kwargs)
+        super().__init__(reference_data_config, dataset_config_uri, **kwargs)
+        self.spark = spark
         if not self.spark:
             raise AttributeError("Spark session must be provided")
 

{data_validation_engine-0.7.4 → data_validation_engine-0.7.6}/src/dve/pipeline/duckdb_pipeline.py

@@ -5,10 +5,12 @@ from typing import Optional
 
 from duckdb import DuckDBPyConnection, DuckDBPyRelation
 
-from dve.core_engine.backends.base.reference_data import BaseRefDataLoader
+import dve.parser.file_handling as fh
+from dve.core_engine.backends.base.reference_data import ReferenceConfig
 from dve.core_engine.backends.implementations.duckdb.auditing import DDBAuditingManager
 from dve.core_engine.backends.implementations.duckdb.contract import DuckDBDataContract
 from dve.core_engine.backends.implementations.duckdb.duckdb_helpers import duckdb_get_entity_count
+from dve.core_engine.backends.implementations.duckdb.reference_data import DuckDBRefDataLoader
 from dve.core_engine.backends.implementations.duckdb.rules import DuckDBStepImplementations
 from dve.core_engine.models import SubmissionInfo
 from dve.core_engine.type_hints import URI
@@ -30,7 +32,6 @@ class DDBDVEPipeline(BaseDVEPipeline):
         connection: DuckDBPyConnection,
         rules_path: Optional[URI],
         submitted_files_path: Optional[URI],
-        reference_data_loader: Optional[type[BaseRefDataLoader]] = None,
         job_run_id: Optional[int] = None,
         logger: Optional[logging.Logger] = None,
     ):
@@ -42,9 +43,19 @@ class DDBDVEPipeline(BaseDVEPipeline):
             DuckDBStepImplementations.register_udfs(connection=self._connection),
             rules_path,
             submitted_files_path,
-            reference_data_loader,
             job_run_id,
             logger,
+            {"connection": self._connection},
+        )
+
+    def init_reference_data_loader(
+        self, reference_data_config: dict[str, ReferenceConfig], **kwargs
+    ) -> DuckDBRefDataLoader:
+        return DuckDBRefDataLoader(
+            connection=self._connection,
+            reference_data_config=reference_data_config,
+            dataset_config_uri=fh.get_parent(self._rules_path),  # type: ignore
+            **kwargs
         )
 
     # pylint: disable=arguments-differ

{data_validation_engine-0.7.4 → data_validation_engine-0.7.6}/src/dve/pipeline/pipeline.py

@@ -9,7 +9,7 @@ from concurrent.futures import Executor, Future, ThreadPoolExecutor
 from functools import lru_cache
 from itertools import starmap
 from threading import Lock
-from typing import Optional, Union
+from typing import Any, Optional, Union
 from uuid import uuid4
 
 import polars as pl
@@ -26,7 +26,7 @@ from dve.common.error_utils import (
 from dve.core_engine.backends.base.auditing import BaseAuditingManager
 from dve.core_engine.backends.base.contract import BaseDataContract
 from dve.core_engine.backends.base.core import EntityManager
-from dve.core_engine.backends.base.reference_data import BaseRefDataLoader
+from dve.core_engine.backends.base.reference_data import BaseRefDataLoader, ReferenceConfig
 from dve.core_engine.backends.base.rules import BaseStepImplementations
 from dve.core_engine.backends.exceptions import MessageBearingError
 from dve.core_engine.backends.readers import BaseFileReader
@@ -36,7 +36,7 @@ from dve.core_engine.exceptions import CriticalProcessingError
 from dve.core_engine.loggers import get_logger
 from dve.core_engine.message import FeedbackMessage
 from dve.core_engine.models import SubmissionInfo, SubmissionStatisticsRecord
-from dve.core_engine.type_hints import URI, DVEStageName, FileURI, InfoURI
+from dve.core_engine.type_hints import URI, DVEStageName, EntityName, FileURI, InfoURI
 from dve.parser import file_handling as fh
 from dve.parser.file_handling.implementations.file import LocalFilesystemImplementation
 from dve.parser.file_handling.service import _get_implementation
@@ -49,6 +49,7 @@ PERMISSIBLE_EXCEPTIONS: tuple[type[Exception]] = (
 )
 
 
+# pylint: disable=R0904
 class BaseDVEPipeline:
     """
     Base class for running a DVE Pipeline either by a given step or a full e2e process.
@@ -62,14 +63,13 @@ class BaseDVEPipeline:
         step_implementations: Optional[BaseStepImplementations[EntityType]],
         rules_path: Optional[URI],
         submitted_files_path: Optional[URI],
-        reference_data_loader: Optional[type[BaseRefDataLoader]] = None,
         job_run_id: Optional[int] = None,
         logger: Optional[logging.Logger] = None,
+        backend_reader_kwargs: Optional[dict[str, Any]] = None,
     ):
         self._submitted_files_path = submitted_files_path
         self._processed_files_path = processed_files_path
         self._rules_path = rules_path
-        self._reference_data_loader = reference_data_loader
         self._job_run_id = job_run_id
         self._audit_tables = audit_tables
         self._data_contract = data_contract
@@ -78,6 +78,7 @@ class BaseDVEPipeline:
         self._summary_lock = Lock()
         self._rec_tracking_lock = Lock()
         self._aggregates_lock = Lock()
+        self._backend_reader_kwargs = backend_reader_kwargs
 
         if self._data_contract:
             self._data_contract.logger = self._logger
@@ -109,11 +110,23 @@ class BaseDVEPipeline:
         """The step implementations to apply the business rules to a given dataset"""
         return self._step_implementations
 
+    @property
+    def backend_reader_kwargs(self) -> dict[str, Any] | None:
+        """Important required arguments for all readers related to the specific backend
+        that can't be specified at time of writing config eg. duckdb connection"""
+        return self._backend_reader_kwargs
+
     @staticmethod
     def get_entity_count(entity: EntityType) -> int:
         """Get a row count of an entity stored as parquet"""
         raise NotImplementedError()
 
+    def init_reference_data_loader(
+        self, reference_data_config: dict[EntityName, ReferenceConfig], **kwargs
+    ) -> BaseRefDataLoader:
+        """Get reference data loader if required for business rules"""
+        raise NotImplementedError()
+
     def get_submission_status(
         self, step_name: DVEStageName, submission_id: str
     ) -> SubmissionStatus:
@@ -199,7 +212,9 @@ class BaseDVEPipeline:
 
         for model_name, model in models.items():
             self._logger.info(f"Transforming {model_name} to stringified parquet")
-            reader: BaseFileReader = load_reader(dataset, model_name, ext)
+            reader: BaseFileReader = load_reader(
+                dataset, model_name, ext, self.backend_reader_kwargs
+            )
             try:
                 if not entity_type:
                     reader.write_parquet(
@@ -527,7 +542,7 @@ class BaseDVEPipeline:
 
         return processed_files, failed_processing
 
-    def apply_business_rules(   # pylint: disable=R0914
+    def apply_business_rules(  # pylint: disable=R0914
         self, submission_info: SubmissionInfo, submission_status: Optional[SubmissionStatus] = None
     ) -> tuple[SubmissionInfo, SubmissionStatus]:
         """Apply the business rules to a given submission, the submission may have failed at the
@@ -542,9 +557,6 @@ class BaseDVEPipeline:
         if not self.rules_path:
             raise AttributeError("business rules path not provided.")
 
-        if not self._reference_data_loader:
-            raise AttributeError("reference data loader not provided.")
-
         if not self.processed_files_path:
             raise AttributeError("processed files path has not been provided.")
 
@@ -556,8 +568,10 @@ class BaseDVEPipeline:
             self._processed_files_path, submission_info.submission_id
         )
         ref_data = config.get_reference_data_config()
+        reference_data: BaseRefDataLoader = self.init_reference_data_loader(
+            reference_data_config=ref_data
+        )
         rules = config.get_rule_metadata()
-        reference_data = self._reference_data_loader(ref_data)  # type: ignore
         entities = {}
         contract = fh.joinuri(
             self.processed_files_path, submission_info.submission_id, "data_contract"
@@ -582,10 +596,7 @@ class BaseDVEPipeline:
         key_fields = {model: conf.reporting_fields for model, conf in model_config.items()}
 
         _errors_uri, rules_success = self.step_implementations.apply_rules(  # type: ignore
-            working_directory,
-            entity_manager,
-            rules,
-            key_fields
+            working_directory, entity_manager, rules, key_fields
         )
 
         rule_messages = load_feedback_messages(

{data_validation_engine-0.7.4 → data_validation_engine-0.7.6}/src/dve/pipeline/spark_pipeline.py

@@ -6,9 +6,11 @@ from typing import Optional
 
 from pyspark.sql import DataFrame, SparkSession
 
-from dve.core_engine.backends.base.reference_data import BaseRefDataLoader
+import dve.parser.file_handling as fh
+from dve.core_engine.backends.base.reference_data import ReferenceConfig
 from dve.core_engine.backends.implementations.spark.auditing import SparkAuditingManager
 from dve.core_engine.backends.implementations.spark.contract import SparkDataContract
+from dve.core_engine.backends.implementations.spark.reference_data import SparkRefDataLoader
 from dve.core_engine.backends.implementations.spark.rules import SparkStepImplementations
 from dve.core_engine.backends.implementations.spark.spark_helpers import spark_get_entity_count
 from dve.core_engine.models import SubmissionInfo
@@ -31,7 +33,6 @@ class SparkDVEPipeline(BaseDVEPipeline):
         audit_tables: SparkAuditingManager,
         rules_path: Optional[URI],
         submitted_files_path: Optional[URI],
-        reference_data_loader: Optional[type[BaseRefDataLoader]] = None,
         spark: Optional[SparkSession] = None,
         job_run_id: Optional[int] = None,
         logger: Optional[logging.Logger] = None,
@@ -44,11 +45,20 @@ class SparkDVEPipeline(BaseDVEPipeline):
             SparkStepImplementations.register_udfs(self._spark),
             rules_path,
             submitted_files_path,
-            reference_data_loader,
             job_run_id,
             logger,
         )
 
+    def init_reference_data_loader(
+        self, reference_data_config: dict[str, ReferenceConfig], **kwargs
+    ) -> SparkRefDataLoader:
+        return SparkRefDataLoader(
+            spark=self._spark,
+            reference_data_config=reference_data_config,
+            dataset_config_uri=fh.get_parent(self._rules_path),  # type: ignore
+            **kwargs
+        )
+
     # pylint: disable=arguments-differ
     def write_file_to_parquet(  # type: ignore
         self, submission_file_uri: URI, submission_info: SubmissionInfo, output: URI

{data_validation_engine-0.7.4 → data_validation_engine-0.7.6}/src/dve/pipeline/utils.py

@@ -3,7 +3,7 @@
 
 import json
 from threading import Lock
-from typing import Optional
+from typing import Any, Optional
 
 from pydantic.main import ModelMetaclass
 from pyspark.sql import SparkSession
@@ -45,10 +45,17 @@ def load_config(
     return models, config, dataset
 
 
-def load_reader(dataset: Dataset, model_name: str, file_extension: str):
+def load_reader(
+    dataset: Dataset,
+    model_name: str,
+    file_extension: str,
+    backend_reader_kwargs: Optional[dict[str, Any]] = None,
+):
     """Loads the readers for the diven feed, model name and file extension"""
     reader_config = dataset[model_name].reader_config[f".{file_extension.lower()}"]
-    reader = _READER_REGISTRY[reader_config.reader](**reader_config.kwargs_)
+    reader = _READER_REGISTRY[reader_config.reader](
+        **reader_config.kwargs_, **backend_reader_kwargs if backend_reader_kwargs else {}
+    )
     return reader