data-validation-engine 0.6.2__py3-none-any.whl

dve/core_engine/configuration/v1/rule_stores/models.py

@@ -0,0 +1,57 @@
+"""Models for components in the rule stores."""
+
+from typing import Any, Optional, Union
+
+from pydantic import BaseModel, Field
+from typing_extensions import Annotated, Literal
+
+from dve.core_engine.configuration.v1.filters import FilterConfigUnion
+from dve.core_engine.configuration.v1.steps import StepConfigUnion
+
+
+class BusinessSpecConfig(BaseModel):
+    """A business rule or filter within the config."""
+
+    type: str
+    """The type of business rule."""
+
+    description: Optional[str] = None
+    """A description of what the rule/filter should do."""
+    parameter_descriptions: dict[str, str] = Field(default_factory=dict)
+    """Descriptions of parameters used by the rule."""
+    parameter_defaults: dict[str, Any] = Field(default_factory=dict)
+    """Default parameters to be used by the rule if no param is passed."""
+
+
+class BusinessFilterSpecConfig(BusinessSpecConfig):
+    """A business filter within the rule store."""
+
+    type: Literal["filter"]
+
+    rule_config: FilterConfigUnion
+    """The configuration for the filter."""
+
+
+class ComplexRuleConfig(BaseModel):
+    """The rule config for a business rule."""
+
+    rules: list[StepConfigUnion] = Field(default_factory=list)
+    filters: list[FilterConfigUnion] = Field(default_factory=list)
+    post_filter_rules: list[StepConfigUnion] = Field(default_factory=list)
+
+
+class BusinessRuleSpecConfig(BusinessSpecConfig):
+    """A business rule within the rule store."""
+
+    type: Literal["complex_rule"]
+
+    rule_config: ComplexRuleConfig
+    """The configuration for the rule."""
+    dependencies: list[str] = Field(default_factory=list)
+    """The dependencies for the business rule."""
+
+
+BusinessComponentSpecConfigUnion = Annotated[
+    Union[BusinessFilterSpecConfig, BusinessRuleSpecConfig], Field(discriminator="type")
+]
+"""A union of the different business component types."""

dve/core_engine/configuration/v1/steps.py

@@ -0,0 +1,365 @@
+"""Very basic configuration options for steps.
+
+These (mostly) only differ slightly from the metadata steps,
+but there's some repetition here to make it possible to change
+the metadata steps without altering the config classes.
+
+N.B. These are quite coarsely copied from the JSON schema.
+
+"""
+
+# pylint: disable=missing-class-docstring
+from abc import ABC, abstractmethod
+from typing import Any, Optional, Union
+
+from pydantic import BaseModel, Extra, Field, validator
+from typing_extensions import Annotated, Literal
+
+from dve.core_engine.backends.metadata.rules import (
+    AbstractStep,
+    Aggregation,
+    AntiJoin,
+    ColumnAddition,
+    ColumnRemoval,
+    ConfirmJoinHasMatch,
+    CopyEntity,
+    EntityRemoval,
+    HeaderJoin,
+    ImmediateFilter,
+    InnerJoin,
+    LeftJoin,
+    OneToOneJoin,
+    RenameEntity,
+    SelectColumns,
+    SemiJoin,
+    TableUnion,
+)
+from dve.core_engine.type_hints import MultipleExpressions
+
+
+class ConfigStep(BaseModel, ABC):
+    """The parent for the config steps."""
+
+    class Config:  # pylint: disable=too-few-public-methods
+        """Config class for dynamically generated pydantic models"""
+
+        extra = Extra.forbid
+
+    name: Optional[str] = None
+    """The 'name' of the rule. This is mapped to an ID in the entity."""
+    operation: str
+    """The operation implemented by the step."""
+
+    @abstractmethod
+    def to_step(self) -> AbstractStep:
+        """Convert the config step definition to a 'real' metadata step."""
+
+
+class AddConfig(ConfigStep):
+    """Configuration step for adding a new column"""
+
+    operation: Literal["add"]
+
+    entity: str
+    new_entity_name: Optional[str] = None
+    column_name: str
+    expression: str
+
+    def to_step(self) -> AbstractStep:
+        """Takes a config object and returns a step object"""
+        return ColumnAddition(
+            id=self.name,
+            entity_name=self.entity,
+            new_entity_name=self.new_entity_name,
+            column_name=self.column_name,
+            expression=self.expression,
+        )
+
+
+class RemoveConfig(ConfigStep):
+    """Configuration step for removing a column"""
+
+    operation: Literal["remove"]
+
+    entity: str
+    new_entity_name: Optional[str] = None
+    column_name: str
+
+    def to_step(self) -> AbstractStep:
+        """Takes a config object and returns a step object"""
+        return ColumnRemoval(
+            id=self.name,
+            entity_name=self.entity,
+            new_entity_name=self.new_entity_name,
+            column_name=self.column_name,
+        )
+
+
+class GroupByConfig(ConfigStep):
+    """Configuration step for performing a GROUP BY operation"""
+
+    operation: Literal["group_by"]
+
+    entity: str
+    new_entity_name: Optional[str] = None
+    group_by: MultipleExpressions
+    pivot_column: Optional[str] = None
+    pivot_values: Optional[list[str]] = None
+    agg_columns: MultipleExpressions
+
+    @validator("pivot_values")
+    @classmethod
+    def _ensure_no_values_if_not_column(cls, value: Optional[str], values: dict[str, Any]):
+        if value and not values["pivot_column"]:
+            raise ValueError("Cannot provide 'pivot_values' if no 'pivot_column'")
+        return value
+
+    def to_step(self) -> AbstractStep:
+        """Takes a config object and returns a step object"""
+        return Aggregation(
+            id=self.name,
+            entity_name=self.entity,
+            new_entity_name=self.new_entity_name,
+            group_by=self.group_by,
+            pivot_column=self.pivot_column,
+            pivot_values=self.pivot_values,
+            agg_columns=self.agg_columns,
+        )
+
+
+class SelectConfig(ConfigStep):
+    """Configuration step for performing a SELECT operation"""
+
+    operation: Literal["select"]
+
+    entity: str
+    new_entity_name: Optional[str] = None
+    columns: MultipleExpressions
+    distinct: bool = False
+
+    def to_step(self) -> AbstractStep:
+        """Takes a config object and returns a step object"""
+        return SelectColumns(
+            id=self.name,
+            entity_name=self.entity,
+            new_entity_name=self.new_entity_name,
+            columns=self.columns,
+            distinct=self.distinct,
+        )
+
+
+class RenameEntityConfig(ConfigStep):
+    """Configuration step for renaming an entity"""
+
+    operation: Literal["rename_entity"]
+
+    entity: str
+    new_entity_name: str
+
+    def to_step(self) -> AbstractStep:
+        """Takes a config object and returns a step object"""
+        return RenameEntity(
+            id=self.name,
+            entity_name=self.entity,
+            new_entity_name=self.new_entity_name,
+        )
+
+
+class NonNotifyingFilterConfig(ConfigStep):
+    """Configuration step for filtering out values without creating errors
+
+    mainly used on derived entities
+    """
+
+    operation: Literal["filter_without_notifying"]
+
+    entity: str
+    new_entity_name: Optional[str] = None
+    filter_rule: str
+
+    def to_step(self) -> AbstractStep:
+        """Takes a config object and returns a step object"""
+        return ImmediateFilter(
+            id=self.name,
+            entity_name=self.entity,
+            new_entity_name=self.new_entity_name,
+            expression=self.filter_rule,
+        )
+
+
+class HasMatchConfig(ConfigStep):
+    """Configuration step for checking if a value has a match in another entity"""
+
+    operation: Literal["has_match"]
+
+    entity: str
+    new_entity_name: Optional[str] = None
+    target: str
+    join_condition: str
+    column_name: str
+
+    def to_step(self) -> AbstractStep:
+        """Takes a config object and returns a step object"""
+        return ConfirmJoinHasMatch(
+            id=self.name,
+            entity_name=self.entity,
+            target_name=self.target,
+            new_entity_name=self.new_entity_name,
+            join_condition=self.join_condition,
+            column_name=self.column_name,
+        )
+
+
+class SemiOrAntiJoinConfig(ConfigStep):
+    """Configuration step for performing a SEMI or ANTI JOIN
+
+    More performant than a left or right join for checking membership in another entity
+    """
+
+    operation: Literal["semi_join", "anti_join"]
+
+    entity: str
+    new_entity_name: Optional[str] = None
+    target: str
+    join_condition: str
+
+    def to_step(self) -> AbstractStep:
+        """Takes a config object and returns a step object"""
+        type_ = AntiJoin if self.operation == "anti_join" else SemiJoin
+        return type_(
+            id=self.name,
+            entity_name=self.entity,
+            target_name=self.target,
+            new_entity_name=self.new_entity_name,
+            join_condition=self.join_condition,
+        )
+
+
+class LeftOrInnerJoinConfig(SemiOrAntiJoinConfig):
+    """Configuration step for performing a LEFT or INNER JOIN"""
+
+    operation: Literal["left_join", "inner_join"]  # type: ignore
+
+    new_columns: MultipleExpressions
+
+    def to_step(self) -> AbstractStep:
+        """Takes a config object and returns a step object"""
+        type_ = LeftJoin if self.operation == "left_join" else InnerJoin
+        return type_(
+            id=self.name,
+            entity_name=self.entity,
+            target_name=self.target,
+            new_entity_name=self.new_entity_name,
+            join_condition=self.join_condition,
+            new_columns=self.new_columns,
+        )
+
+
+class OneToOneJoinConfig(LeftOrInnerJoinConfig):
+    """Config for joining one entity to another"""
+
+    operation: Literal["join", "one_to_one_join"]  # type: ignore
+
+    perform_integrity_check: bool = True
+
+    def to_step(self) -> AbstractStep:
+        """Takes a config object and returns a step object"""
+        return OneToOneJoin(
+            id=self.name,
+            entity_name=self.entity,
+            target_name=self.target,
+            new_entity_name=self.new_entity_name,
+            join_condition=self.join_condition,
+            new_columns=self.new_columns,
+            perform_integrity_check=self.perform_integrity_check,
+        )
+
+
+class JoinHeaderConfig(ConfigStep):
+    """Config for joining a header onto another entity"""
+
+    operation: Literal["join_header"]
+
+    entity: str
+    new_entity_name: Optional[str] = None
+    target: str
+    header_column_name: str = "_Header"
+    perform_integrity_check: bool = True
+
+    def to_step(self) -> AbstractStep:
+        """Takes a config object and returns a step object"""
+        return HeaderJoin(
+            id=self.name,
+            entity_name=self.entity,
+            target_name=self.target,
+            new_entity_name=self.new_entity_name,
+            header_column_name=self.header_column_name,
+        )
+
+
+class UnionConfig(ConfigStep):
+    """Config to unioning two entities"""
+
+    operation: Literal["union"]
+
+    entity: str
+    new_entity_name: Optional[str] = None
+    target: str
+
+    def to_step(self) -> AbstractStep:
+        """Takes a config object and returns a step object"""
+        return TableUnion(
+            id=self.name,
+            entity_name=self.entity,
+            target_name=self.target,
+            new_entity_name=self.new_entity_name,
+        )
+
+
+class CopyEntityConfig(ConfigStep):
+    """Config for copying entities"""
+
+    operation: Literal["copy_entity"]
+
+    entity: str
+    new_entity_name: str
+
+    def to_step(self) -> AbstractStep:
+        """Takes a config object and returns a step object"""
+        return CopyEntity(
+            id=self.name, entity_name=self.entity, new_entity_name=self.new_entity_name
+        )
+
+
+class RemoveEntityConfig(ConfigStep):
+    """Config for removing entities"""
+
+    operation: Literal["remove_entity", "remove_entities"]
+
+    entity: Union[str, list[str]]
+
+    def to_step(self) -> AbstractStep:
+        """Takes a config object and returns a step object"""
+        return EntityRemoval(id=self.name, entity_name=self.entity)
+
+
+StepConfigUnion = Annotated[
+    Union[
+        AddConfig,
+        CopyEntityConfig,
+        GroupByConfig,
+        HasMatchConfig,
+        JoinHeaderConfig,
+        LeftOrInnerJoinConfig,
+        NonNotifyingFilterConfig,
+        OneToOneJoinConfig,
+        RemoveConfig,
+        RemoveEntityConfig,
+        RenameEntityConfig,
+        SelectConfig,
+        SemiOrAntiJoinConfig,
+        UnionConfig,
+    ],
+    Field(discriminator="operation"),
+]
+"""Pydantic configuration classes for steps."""

dve/core_engine/constants.py

@@ -0,0 +1,8 @@
+"""Constant values used in mutiple places."""
+
+ROWID_COLUMN_NAME: str = "__rowid__"
+"""The name of the column containing the row ID for each entity."""
+
+CONTRACT_ERROR_VALUE_FIELD_NAME: str = "__error_value"
+"""The name of the field that can be used to extract the field value that caused
+   a pydantic validation error"""

dve/core_engine/engine.py

@@ -0,0 +1,265 @@
+"""The core engine for the data validation engine."""
+
+import json
+import logging
+from pathlib import Path
+from types import TracebackType
+from typing import Any, Optional, Union
+
+from pydantic import BaseModel, Field, PrivateAttr, validate_arguments, validator
+from pydantic.types import FilePath
+from pyspark.sql import SparkSession
+
+from dve.core_engine.backends.base.backend import BaseBackend
+from dve.core_engine.backends.implementations.spark.backend import SparkBackend
+from dve.core_engine.backends.implementations.spark.types import SparkEntities
+from dve.core_engine.configuration.base import BaseEngineConfig
+from dve.core_engine.configuration.v1 import V1EngineConfig
+from dve.core_engine.constants import ROWID_COLUMN_NAME
+from dve.core_engine.loggers import get_child_logger, get_logger
+from dve.core_engine.models import EngineRunValidation, SubmissionInfo
+from dve.core_engine.type_hints import EntityName, JSONstring
+from dve.parser.file_handling import TemporaryPrefix, get_resource_exists, joinuri, resolve_location
+from dve.parser.type_hints import URI, Location
+
+
+class CoreEngine(BaseModel):
+    """The core engine implementation for the data validation engine."""
+
+    class Config:  # pylint: disable=too-few-public-methods
+        """`pydantic` configuration options."""
+
+        arbitrary_types_allowed = True
+        validate_assignment = True
+
+    backend_config: BaseEngineConfig
+    """The backend configuration for the given run."""
+    dataset_config_uri: URI
+    """The dischema location for the current run"""
+    output_prefix_uri: URI = Field(default_factory=lambda: Path("outputs").resolve().as_posix())
+    """The prefix for the parquet outputs."""
+    main_log: logging.Logger = Field(default_factory=lambda: get_logger("CoreEngine"))
+    """The `logging.Logger instance for the data ingest process."""
+    cache_prefix_uri: Optional[URI] = None
+    """
+    An optional cache prefix URI. If not provided, a local temporary directory will
+    be used instead (this will not play nicely in Databricks).
+
+    """
+    _cache_dir: Optional[TemporaryPrefix] = PrivateAttr(default=None)
+    """
+    The `TemporaryPrefix` indicating the cache dir.
+
+    Data will be chunked to parquet in this directory after being read,
+    and written here before filters are applied.
+
+    """
+    backend: BaseBackend = None  # type: ignore
+    """The backend to use to process the files."""
+    debug: bool = False
+    """Indication of if this run is in debug mode."""
+
+    @validator("cache_prefix_uri", "output_prefix_uri", allow_reuse=True, pre=True)
+    # pylint: disable=E0213
+    def _validate_prefix_uri(cls, location: Optional[Location]) -> Optional[URI]:
+        """Ensure we support the cache prefix scheme."""
+        if location is None:
+            return None
+        return resolve_location(location)
+
+    def __init__(self, *args, **kwargs):
+        # pylint: disable=W0235
+        super().__init__(*args, **kwargs)
+
+    @validator("backend", always=True)
+    @classmethod
+    def _ensure_backend(cls, value: Optional[BaseBackend], values: dict[str, Any]) -> BaseBackend:
+        """Ensure a default backend is created if a backend is not specified."""
+        if value is not None:
+            return value
+
+        main_logger = values.get("main_log")
+        if main_logger is None:
+            return SparkBackend(dataset_config_uri=values.get("dataset_config_uri"))
+        return SparkBackend(
+            dataset_config_uri=values.get("dataset_config_uri"),
+            logger=get_child_logger(
+                ".".join((SparkBackend.__module__, SparkBackend.__name__)), main_logger
+            ),
+        )
+
+    @classmethod
+    @validate_arguments(config={"arbitrary_types_allowed": True})
+    def build(
+        cls,
+        dataset_config_path: Union[FilePath, URI],
+        output_prefix: Location = Path("./outputs"),
+        cache_prefix: Optional[Location] = None,
+        parent_logger: Optional[logging.Logger] = None,
+        debug: Optional[bool] = False,
+        **kwargs,
+    ):
+        """Build an engine from serialised definitions.
+
+        Args:
+         - `dataset_config_path`: a URI or path indicating the location of the
+           dataset configuration.
+         - `output_prefix`: the prefix for parquet outputs (a URI or a local path).
+         - `cache_prefix`: the prefix for caching (a URI or a local path).
+         - `parent_logger`: an optional parent logger for the engine.
+         - `debug`: whether to run in debug mode (default: False).
+
+        """
+        if parent_logger:
+            main_log = get_child_logger(cls.__name__, parent_logger)
+        else:
+            main_log = get_logger(cls.__name__)
+        main_log.info("Initialising...")
+        main_log.info(f"Debug mode: {debug}")
+
+        if isinstance(dataset_config_path, Path):
+            dataset_config_uri = dataset_config_path.resolve().as_posix()
+        else:
+            dataset_config_uri = dataset_config_path
+        if isinstance(output_prefix, Path):
+            output_prefix_uri = output_prefix.resolve().as_posix()
+        else:
+            output_prefix_uri = output_prefix
+
+        backend_config = V1EngineConfig.load(dataset_config_uri)
+
+        self = cls(
+            dataset_config_uri=dataset_config_uri,
+            output_prefix_uri=output_prefix_uri,
+            main_log=main_log,
+            cache_prefix_uri=cache_prefix,
+            backend_config=backend_config,
+            debug=debug,
+            **kwargs,
+        )
+        self.main_log.info(f"Output path: {self.output_prefix_uri!r}")
+        return self
+
+    @classmethod
+    def build_from_model(cls, model_str: JSONstring):
+        """Build an engine from a serialised JSON pydantic model of definitions.
+
+        Args:
+         - `dataset_config_path`: a URI or path indicating the location of the
+           dataset configuration.
+         - `output_prefix`: the prefix for parquet outputs (a URI or a local path).
+
+        """
+        main_log = get_logger("CoreEngine")
+        main_log.info("Initalise from model...")
+        return cls.build(**EngineRunValidation(**json.loads(model_str)).dict())
+
+    def __enter__(self) -> "CoreEngine":
+        self.main_log.info("Entering pipeline context.")
+        if self._cache_dir is not None:
+            raise ValueError("Pipeline already within context")
+
+        self._cache_dir = TemporaryPrefix(self.cache_prefix_uri)
+        self._cache_dir.__enter__()
+        self.main_log.info(f"Pipeline will cache to {self.cache_prefix!r}")
+        return self
+
+    def __exit__(
+        self,
+        exc_type: Optional[type[Exception]],
+        exc_value: Optional[Exception],
+        traceback: Optional[TracebackType],
+    ) -> None:
+        self.main_log.info(f"Exiting pipeline context, clearing {self.cache_prefix!r}")
+        cache_dir = self._cache_dir
+        self._cache_dir = None
+
+        if cache_dir is not None:
+            cache_dir.__exit__(exc_type, exc_value, traceback)
+
+        self.main_log.info("Cleared cache.")
+
+    @property
+    def cache_prefix(self) -> URI:
+        """The cache directory for the pipeline run."""
+        if self._cache_dir is None:
+            raise ValueError(
+                "`cache_prefix` is undefined when the pipeline is not being used as a "
+                + "context manager"
+            )
+        return self._cache_dir.prefix
+
+    def _write_entity_outputs(self, entities: SparkEntities) -> SparkEntities:
+        """Write the final entities to the output prefix as Parquet.
+
+        This will result in a directory of files for each entity, containing
+        parquet files for each partition in the entity.
+
+        """
+        output_entities = {}
+
+        self.main_log.info(f"Writing entities to the output location: {self.output_prefix_uri}")
+        for entity_name, entity in entities.items():
+            entity = entity.drop(ROWID_COLUMN_NAME)
+
+            self.main_log.info(f"Entity: {entity_name} {type(entity)}")
+
+            output_uri = joinuri(self.output_prefix_uri, entity_name)
+            if get_resource_exists(output_uri):
+                self.main_log.info(f"{output_uri} already exists - will be overwritten")
+
+            self.main_log.info(f"+ Writing parquet output to {output_uri!r}")
+            entity.write.mode("overwrite").parquet(output_uri)
+            spark_session = SparkSession.builder.getOrCreate()
+            output_entities[entity_name] = spark_session.read.format("parquet").load(
+                output_uri, schema=entity.schema
+            )
+
+        return output_entities
+
+    def _write_outputs(self, entities: SparkEntities) -> SparkEntities:
+        """Write the outputs from the pipeline, returning the written entities
+        and messages.
+
+        """
+        entities = self._write_entity_outputs(entities)
+
+        return entities
+
+    def _show_available_entities(self, entities: SparkEntities, *, verbose: bool = False) -> None:
+        """Print current entities."""
+        self.main_log.info("Displaying available dataframes in this run:")
+
+        for entity_name, entity in entities.items():
+            # FIXME: Currently a print statement because log messages
+            #  can arrive out of sequence with the df.show()
+            if self.debug:
+                print(f"+ Entity dataframe: {entity_name} has {entity.count()} rows")
+            else:
+                print(f"+ Entity dataframe: {entity_name}")
+
+            if verbose:
+                # Cap the number of rows displayed to reduce probs with max log size on dbr
+                entity.show(n=10, truncate=False)
+
+    def run_pipeline(
+        self,
+        entity_locations: dict[EntityName, URI],
+        # pylint: disable=unused-argument
+        submission_info: Optional[SubmissionInfo] = None,
+    ) -> tuple[SparkEntities, URI]:
+        """Run the pipeline, reading in the entities and applying validation
+        and transformation rules, and then write the outputs.
+
+        The returned entities will reference the output locations, so
+        references should be valid after the pipeline context exits.
+
+        """
+        entities, errors_uri = self.backend.process_legacy(
+            self.output_prefix_uri,
+            entity_locations,
+            self.backend_config.get_contract_metadata(),
+            self.backend_config.get_rule_metadata(),
+            submission_info,
+        )
+        return self._write_outputs(entities), errors_uri