lakehouse_engine.algorithms.dq_validator

  1"""Module to define Data Validator class."""
  2
  3from delta.tables import DeltaTable
  4from pyspark.sql import DataFrame
  5from pyspark.sql.utils import StreamingQueryException
  6
  7from lakehouse_engine.algorithms.algorithm import Algorithm
  8from lakehouse_engine.core.definitions import DQSpec, DQValidatorSpec, InputSpec
  9from lakehouse_engine.core.exec_env import ExecEnv
 10from lakehouse_engine.dq_processors.dq_factory import DQFactory
 11from lakehouse_engine.dq_processors.exceptions import DQValidationsFailedException
 12from lakehouse_engine.io.reader_factory import ReaderFactory
 13from lakehouse_engine.utils.logging_handler import LoggingHandler
 14
 15
 16class DQValidator(Algorithm):
 17    """Validate data using an algorithm configuration (ACON represented as dict).
 18
 19    This algorithm focuses on isolate Data Quality Validations from loading,
 20    applying a set of data quality functions to a specific input dataset,
 21    without the need to define any output specification.
 22    You can use any input specification compatible with the lakehouse engine
 23    (dataframe, table, files, etc).
 24    """
 25
 26    _LOGGER = LoggingHandler(__name__).get_logger()
 27
 28    def __init__(self, acon: dict):
 29        """Construct DQValidator algorithm instances.
 30
 31        A data quality validator needs the following specifications to work
 32        properly:
 33            - input specification (mandatory): specify how and what data to
 34            read.
 35            - data quality specification (mandatory): specify how to execute
 36            the data quality process.
 37            - restore_prev_version (optional): specify if, having
 38            delta table/files as input, they should be restored to the
 39            previous version if the data quality process fails. Note: this
 40            is only considered if fail_on_error is kept as True.
 41
 42        Args:
 43            acon: algorithm configuration.
 44        """
 45        self.spec: DQValidatorSpec = DQValidatorSpec(
 46            input_spec=InputSpec(**acon["input_spec"]),
 47            dq_spec=self._get_dq_spec(acon["dq_spec"]),
 48            restore_prev_version=acon.get("restore_prev_version", None),
 49        )
 50
 51    def read(self) -> DataFrame:
 52        """Read data from an input location into a distributed dataframe.
 53
 54        Returns:
 55             Dataframe with data that was read.
 56        """
 57        current_df = ReaderFactory.get_data(self.spec.input_spec)
 58
 59        return current_df
 60
 61    def process_dq(self, data: DataFrame) -> DataFrame:
 62        """Process the data quality tasks for the data that was read.
 63
 64        It supports a single input dataframe.
 65
 66        It is possible to use data quality validators/expectations that will validate
 67        your data and fail the process in case the expectations are not met. The DQ
 68        process also generates and keeps updating a site containing the results of the
 69        expectations that were done on your data. The location of the site is
 70        configurable and can either be on file system or S3. If you define it to be
 71        stored on S3, you can even configure your S3 bucket to serve the site so that
 72        people can easily check the quality of your data. Moreover, it is also
 73        possible to store the result of the DQ process into a defined result sink.
 74
 75        Args:
 76            data: input dataframe on which to run the DQ process.
 77
 78        Returns:
 79            Validated dataframe.
 80        """
 81        return DQFactory.run_dq_process(self.spec.dq_spec, data)
 82
 83    def execute(self) -> None:
 84        """Define the algorithm execution behaviour."""
 85        self._LOGGER.info("Starting read stage...")
 86        read_df = self.read()
 87
 88        self._LOGGER.info("Starting data quality validator...")
 89        try:
 90            if read_df.isStreaming:
 91                # To handle streaming, and although we are not interested in
 92                # writing any data, we still need to start the streaming and
 93                # execute the data quality process in micro batches of data.
 94                def write_dq_validator_micro_batch(
 95                    batch_df: DataFrame, batch_id: int
 96                ) -> None:
 97                    self.process_dq(batch_df)
 98
 99                read_df.writeStream.trigger(once=True).foreachBatch(
100                    write_dq_validator_micro_batch
101                ).start().awaitTermination()
102
103            else:
104                self.process_dq(read_df)
105        except (DQValidationsFailedException, StreamingQueryException):
106            if not self.spec.input_spec.df_name and self.spec.restore_prev_version:
107                self._LOGGER.info("Restoring delta table/files to previous version...")
108
109                self._restore_prev_version()
110
111                raise DQValidationsFailedException(
112                    "Data Quality Validations Failed! The delta "
113                    "table/files were restored to the previous version!"
114                )
115
116            elif self.spec.dq_spec.fail_on_error:
117                raise DQValidationsFailedException("Data Quality Validations Failed!")
118        else:
119            self._LOGGER.info("Execution of the algorithm has finished!")
120
121    @staticmethod
122    def _get_dq_spec(input_dq_spec: dict) -> DQSpec:
123        """Get data quality specification from acon.
124
125        Args:
126            input_dq_spec: data quality specification.
127
128        Returns:
129            Data quality spec.
130        """
131        dq_spec, dq_functions, critical_functions = Algorithm.get_dq_spec(input_dq_spec)
132
133        dq_spec.dq_functions = dq_functions
134        dq_spec.critical_functions = critical_functions
135
136        return dq_spec
137
138    def _restore_prev_version(self) -> None:
139        """Restore delta table or delta files to previous version."""
140        if self.spec.input_spec.db_table:
141            delta_table = DeltaTable.forName(
142                ExecEnv.SESSION, self.spec.input_spec.db_table
143            )
144        else:
145            delta_table = DeltaTable.forPath(
146                ExecEnv.SESSION, self.spec.input_spec.location
147            )
148
149        previous_version = (
150            delta_table.history().agg({"version": "max"}).collect()[0][0] - 1
151        )
152
153        delta_table.restoreToVersion(previous_version)